| #!/usr/bin/env python |
| |
| # |
| # Licensed to the Apache Software Foundation (ASF) under one |
| # or more contributor license agreements. See the NOTICE file |
| # distributed with this work for additional information |
| # regarding copyright ownership. The ASF licenses this file |
| # to you under the Apache License, Version 2.0 (the |
| # "License"); you may not use this file except in compliance |
| # with the License. You may obtain a copy of the License at |
| # |
| # http://www.apache.org/licenses/LICENSE-2.0 |
| # |
| # Unless required by applicable law or agreed to in writing, |
| # software distributed under the License is distributed on an |
| # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY |
| # KIND, either express or implied. See the License for the |
| # specific language governing permissions and limitations |
| # under the License. |
| # |
| |
| from __future__ import unicode_literals |
| from __future__ import division |
| from __future__ import absolute_import |
| from __future__ import print_function |
| |
| import argparse |
| import os |
| import selectors |
| import signal |
| import socket |
| import sys |
| from threading import Thread |
| import time |
| import traceback |
| import types |
| |
| from system_test import Logger |
| from system_test import TIMEOUT |
| |
| |
| class GracefulExitSignaler: |
| kill_now = False |
| |
| def __init__(self): |
| signal.signal(signal.SIGINT, self.exit_gracefully) |
| signal.signal(signal.SIGTERM, self.exit_gracefully) |
| |
| def exit_gracefully(self, signum, frame): |
| self.kill_now = True |
| |
| |
| def split_chunk_for_display(raw_bytes): |
| """ |
| Given some raw bytes, return a display string |
| Only show the beginning and end of largish (2xMAGIC_SIZE) arrays. |
| :param raw_bytes: |
| :return: display string |
| """ |
| MAGIC_SIZE = 50 # Content repeats after chunks this big - used by echo client, too |
| if len(raw_bytes) > 2 * MAGIC_SIZE: |
| result = repr(raw_bytes[:MAGIC_SIZE]) + " ... " + repr(raw_bytes[-MAGIC_SIZE:]) |
| else: |
| result = repr(raw_bytes) |
| return result |
| |
| |
| class TcpEchoClient: |
| |
| def __init__(self, prefix, host, port, size, count, timeout, logger): |
| """ |
| :param host: connect to this host |
| :param port: connect to this port |
| :param size: size of individual payload chunks in bytes |
| :param count: number of payload chunks |
| :param strategy: "1" Send one payload; # TODO more strategies |
| Recv one payload |
| :param logger: Logger() object |
| :return: |
| """ |
| # Start up |
| self.sock = None |
| self.prefix = prefix |
| self.host = host |
| self.port = int(port) |
| self.size = size |
| self.count = count |
| self.timeout = timeout |
| self.logger = logger |
| self.keep_running = True |
| self.is_running = False |
| self.exit_status = None |
| self.error = None |
| self._thread = Thread(target=self.run) |
| self._thread.daemon = True |
| self._thread.start() |
| |
| def run(self): |
| self.logger.log("%s Client is starting up" % self.prefix) |
| try: |
| start_time = time.time() |
| self.is_running = True |
| self.logger.log('%s Connecting to host:%s, port:%d, size:%d, count:%d' % |
| (self.prefix, self.host, self.port, self.size, self.count)) |
| total_sent = 0 |
| total_rcvd = 0 |
| |
| if self.count > 0 and self.size > 0: |
| # outbound payload only if count and size both greater than zero |
| payload_out = [] |
| out_list_idx = 0 # current _out array being sent |
| out_byte_idx = 0 # next-to-send in current array |
| out_ready_to_send = True |
| # Generate unique content for each message so you can tell where the message |
| # or fragment belongs in the whole stream. Chunks look like: |
| # b'[localhost:33333:6:0]ggggggggggggggggggggggggggggg' |
| # host: localhost |
| # port: 33333 |
| # index: 6 |
| # offset into message: 0 |
| CONTENT_CHUNK_SIZE = 50 # Content repeats after chunks this big - used by echo server, too |
| for idx in range(self.count): |
| body_msg = "" |
| padchar = "abcdefghijklmnopqrstuvwxyz@#$%"[idx % 30] |
| while len(body_msg) < self.size: |
| chunk = "[%s:%d:%d:%d]" % (self.host, self.port, idx, len(body_msg)) |
| padlen = CONTENT_CHUNK_SIZE - len(chunk) |
| chunk += padchar * padlen |
| body_msg += chunk |
| if len(body_msg) > self.size: |
| body_msg = body_msg[:self.size] |
| payload_out.append(bytearray(body_msg.encode())) |
| # incoming payloads |
| payload_in = [] |
| in_list_idx = 0 # current _in array being received |
| for i in range(self.count): |
| payload_in.append(bytearray()) |
| else: |
| # when count or size .LE. zero then just connect-disconnect |
| self.keep_running = False |
| |
| # set up connection |
| host_address = (self.host, self.port) |
| self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) |
| self.sock.connect(host_address) |
| self.sock.setblocking(False) |
| |
| # set up selector |
| sel = selectors.DefaultSelector() |
| sel.register(self.sock, |
| selectors.EVENT_READ | selectors.EVENT_WRITE) |
| |
| # event loop |
| while self.keep_running: |
| if self.timeout > 0.0: |
| elapsed = time.time() - start_time |
| if elapsed > self.timeout: |
| self.exit_status = "%s Exiting due to timeout. Total sent= %d, total rcvd= %d" % \ |
| (self.prefix, total_sent, total_rcvd) |
| break |
| for key, mask in sel.select(timeout=0.1): |
| sock = key.fileobj |
| if mask & selectors.EVENT_READ: |
| recv_data = sock.recv(1024) |
| if recv_data: |
| total_rcvd = len(recv_data) |
| payload_in[in_list_idx].extend(recv_data) |
| if len(payload_in[in_list_idx]) == self.size: |
| self.logger.log("%s Rcvd message %d" % (self.prefix, in_list_idx)) |
| in_list_idx += 1 |
| if in_list_idx == self.count: |
| # Received all bytes of all chunks - done. |
| self.keep_running = False |
| # Verify the received data |
| if not payload_in == payload_out: |
| for idxc in range(self.count): |
| if not payload_in[idxc] == payload_out[idxc]: |
| for idxs in range(self.size): |
| ob = payload_out[idxc][idxs] |
| ib = payload_in[idxc][idxs] |
| if ob != ib: |
| self.error = "%s ERROR Rcvd message verify fail. row:%d, col:%d, " \ |
| "expected:%s, actual:%s" \ |
| % (self.prefix, idxc, idxs, repr(ob), repr(ib)) |
| break |
| else: |
| out_ready_to_send = True |
| sel.modify(sock, selectors.EVENT_READ | selectors.EVENT_WRITE) |
| elif len(payload_in[in_list_idx]) > self.size: |
| self.error = "ERROR Received message too big. Expected:%d, actual:%d" % \ |
| (self.size, len(payload_in[in_list_idx])) |
| break |
| else: |
| pass # still accumulating a message |
| else: |
| # socket closed |
| self.keep_running = False |
| if not in_list_idx == self.count: |
| self.error = "ERROR server closed. Echoed %d of %d messages." % (in_list_idx, self.count) |
| if self.keep_running and mask & selectors.EVENT_WRITE: |
| if out_ready_to_send: |
| n_sent = self.sock.send(payload_out[out_list_idx][out_byte_idx:]) |
| total_sent += n_sent |
| out_byte_idx += n_sent |
| if out_byte_idx == self.size: |
| self.logger.log("%s Sent message %d" % (self.prefix, out_list_idx)) |
| out_byte_idx = 0 |
| out_list_idx += 1 |
| sel.modify(self.sock, selectors.EVENT_READ) # turn off write events |
| out_ready_to_send = False # turn on when rcvr receives |
| else: |
| pass # logger.log("DEBUG: ignoring EVENT_WRITE") |
| |
| # shut down |
| sel.unregister(self.sock) |
| self.sock.close() |
| |
| except Exception: |
| self.error = "ERROR: exception : '%s'" % traceback.format_exc() |
| |
| self.is_running = False |
| |
| def wait(self, timeout=TIMEOUT): |
| self.logger.log("%s Client is shutting down" % self.prefix) |
| self.keep_running = False |
| self._thread.join(timeout) |
| |
| |
| def main(argv): |
| retval = 0 |
| # parse args |
| p = argparse.ArgumentParser() |
| p.add_argument('--host', '-b', |
| help='Required target host') |
| p.add_argument('--port', '-p', type=int, |
| help='Required target port number') |
| p.add_argument('--size', '-s', type=int, default=100, const=1, nargs='?', |
| help='Size of payload in bytes must be >= 0. Size of zero connects and disconnects with no data traffic.') |
| p.add_argument('--count', '-c', type=int, default=1, const=1, nargs='?', |
| help='Number of payloads to process must be >= 0. Count of zero connects and disconnects with no data traffic.') |
| p.add_argument('--name', |
| help='Optional logger prefix') |
| p.add_argument('--timeout', '-t', type=float, default=0.0, const=1, nargs="?", |
| help='Timeout in seconds. Default value "0" disables timeouts') |
| p.add_argument('--log', '-l', |
| action='store_true', |
| help='Write activity log to console') |
| del argv[0] |
| args = p.parse_args(argv) |
| |
| # host |
| if args.host is None: |
| raise Exception("User must specify a host") |
| host = args.host |
| |
| # port |
| if args.port is None: |
| raise Exception("User must specify a port number") |
| port = args.port |
| |
| # size |
| if args.size < 0: |
| raise Exception("Size must be greater than or equal to zero") |
| size = args.size |
| |
| # count |
| if args.count < 0: |
| raise Exception("Count must be greater than or equal to zero") |
| count = args.count |
| |
| # name / prefix |
| prefix = args.name if args.name is not None else "ECHO_CLIENT (%d_%d_%d)" % \ |
| (port, size, count) |
| |
| # timeout |
| if args.timeout < 0.0: |
| raise Exception("Timeout must be greater than or equal to zero") |
| |
| signaller = GracefulExitSignaler() |
| logger = None |
| |
| try: |
| # logging |
| logger = Logger(title="%s host:%s port %d size:%d count:%d" % (prefix, host, port, size, count), |
| print_to_console=args.log, |
| save_for_dump=False) |
| |
| client = TcpEchoClient(prefix, host, port, size, count, args.timeout, logger) |
| |
| keep_running = True |
| while keep_running: |
| time.sleep(0.1) |
| if client.error is not None: |
| logger.log("%s Client stopped with error: %s" % (prefix, client.error)) |
| keep_running = False |
| retval = 1 |
| if client.exit_status is not None: |
| logger.log("%s Client stopped with status: %s" % (prefix, client.exit_status)) |
| keep_running = False |
| if signaller.kill_now: |
| logger.log("%s Process killed with signal" % prefix) |
| keep_running = False |
| if keep_running and not client.is_running: |
| logger.log("%s Client stopped with no error or status" % prefix) |
| keep_running = False |
| |
| except Exception: |
| client.error = "ERROR: exception : '%s'" % traceback.format_exc() |
| if logger is not None: |
| logger.log("%s Exception: %s" % (prefix, traceback.format_exc())) |
| retval = 1 |
| |
| if client.error is not None: |
| # write client errors to stderr |
| def eprint(*args, **kwargs): |
| print(*args, file=sys.stderr, **kwargs) |
| |
| elines = client.error.split("\n") |
| for line in elines: |
| eprint("ERROR:", prefix, line) |
| |
| return retval |
| |
| |
| if __name__ == "__main__": |
| sys.exit(main(sys.argv)) |
