class Slave(Actionable): """ This class manages connections to the master and handles client-server logistics. Actions on it are run by master. """ id = None last_bw_out = None bw_out = 0 total_bw_out = 0 last_bw_in = None bw_in = 0 total_bw_in = 0 running_tests = [] test_runner = None sink_obj = None sink = None client_tests = {} message_lock = BoundedSemaphore(1) def __init__(self, context): self.context = context # This socket sends messages to master self.socket_out = context.socket(zmq.PUSH) self.socket_out.connect("tcp://localhost:{}".format(SLAVE_MASTER_PORT)) # Syntax allowing self.master.<action>(data) self.master = Sender(self.send_to_master) # This socket receives broadcasts from master self.socket_in = context.socket(zmq.SUB) self.socket_in.setsockopt(zmq.SUBSCRIBE, "") self.socket_in.connect("tcp://localhost:{}".format(SLAVE_PUB_PORT)) self.sink_host = '{}:{}'.format(self.hostname(), SLAVE_SINK_PORT) self.uuid = uuid4().hex def send_to_master(self, data): """ Sends a bson serialized message to master, inserting slave specific information such as id :param data: Message data to send action: (required) RPC to call on serverside slave instance """ if self.id: data['slave_id'] = self.id if data['action'] != 'heartbeat': log.debug('MASTER->{}'.format(data['action'])) self.socket_out.send(bson.dumps(data)) if data['action'] != 'heartbeat': log.debug('MASTER->{} end'.format(data['action'])) def send_to_sink(self, data): """ Send bson encoded data to the sink for collation :param data: Message data to send """ if self.id: data['slave_id'] = self.id if self.sink_obj: self.sink_obj.send_to_sink(data) else: log.warn('{}->SINK failed, no sink'.format(data['action'])) def listen_to_master(self): """ Wait for messages from the server and delegate them """ log.info("Listening to master") while True: self.handle_master_message(self.socket_in.recv()) def handle_master_message(self, msg): """ Decode a message from the server, if it matches our id or has no id decode and run it. :param msg: BSON encoded message from the server action: Local function to run slave_id: If present and not our id ignore message """ data = bson.loads(msg) if 'action' in data: if 'slave_id' not in data or data['slave_id'] == self.id: log.debug('{}<-MASTER'.format(data['action'])) self.run_action(data) log.debug('{}<-MASTER end'.format(data['action'])) else: log.warn('Server sent message with no action') def heartbeat(self): """ Generate a status report and send it to master """ # physical_memory = psutil.phymem_usage() virtual_memory = psutil.virtual_memory() swap_memory = psutil.swap_memory() load = psutil.cpu_percent(interval=1) network = psutil.net_io_counters() hostname = self.hostname() if not self.last_bw_out: self.last_bw_out = network.bytes_sent self.bw_out = (network.bytes_sent - self.last_bw_out)/HEARTBEAT_PERIOD self.total_bw_out += self.bw_out * HEARTBEAT_PERIOD self.last_bw_out = network.bytes_sent if not self.last_bw_in: self.last_bw_in = network.bytes_recv self.bw_in = (network.bytes_recv - self.last_bw_in)/HEARTBEAT_PERIOD self.total_bw_in += self.bw_in * HEARTBEAT_PERIOD self.last_bw_in = network.bytes_recv self.master.heartbeat({ 'hostname': hostname, 'memory': { # 'physical': (physical_memory.used, physical_memory.total), 'physical': (0, 0), 'virtual': (virtual_memory.used, virtual_memory.total), 'swap': (swap_memory.used, swap_memory.total) }, 'load': load, 'bandwidth': { 'in': (self.bw_in, self.total_bw_in), 'out': (self.bw_out, self.total_bw_out) }, 'generated': int(time()), 'sink_id': self.sink_obj.id if self.sink_obj else None, 'slave_uuid': self.uuid }) def heartbeat_forever(self): log.info('Beating heart forever') while True: self.heartbeat() gevent.sleep(HEARTBEAT_PERIOD) @action def set_id(self, data): """ Set id if it matches our uuid :param data: Message specific data: slave_uuid: UUID to check against ours new_slave_id: Id to set if match """ if 'slave_uuid' not in data or 'new_slave_id' not in data: log.warn('receiving id without slave_uuid/new_slave_id') return if data['slave_uuid'] == self.uuid: self.id = data['new_slave_id'] log.info('Received id: {}'.format(self.id)) @action def set_sink(self, data): """ Set sink :param data: Message specific data: sink_id: Sink id to connect to sink_host: Host to connect to :return: """ if 'sink_id' in data and 'sink_host' in data and data['sink_host'] == self.sink_host: self.connect_to_sink(data['sink_id'], data['sink_host']) else: log.warn('Incomplete sink call') @staticmethod def hostname(): return socket.gethostname() def connect_to_sink(self, sink_id, host): """ Attempt to connect to the given sink, and create a local sink instance :param sink_id: ID to assign to the sink :param host: host to connect to """ log.info('Attempting to connect to sink({}) - {}'.format(sink_id, host)) if self.sink_obj: log.info('Closing previous sink') self.sink_obj.close() # Create our sink self.sink_obj = Sink(id=sink_id, host=host, slave=self) # Initialize it, and tell master if successful if self.sink_obj.setup(): log.info('Ok!') # RPC for sink self.sink = Sender(self.send_to_sink) self.master.connected_to_sink({ 'sink_id': self.sink_obj.id }) @action def quit(self, data): """ If this RPC is run master told this slave to quit """ log.info('Master told us to quit, quitting.') sys.exit(0) @action def run_test(self, data): """ Receive a test, instantiate it and run :param data: Message specific data: client_id: Client that started the test test: Test data to hand off to the test :return: """ log.info('slave({}).run_test client({})'.format(self.id, data['client_id'])) if not self.sink_obj: log.warn('Error, no sink') return if 'test' in data: client_id = data['client_id'] test = Test(self.sink, data) # Only allow one test at a time if client_id in self.client_tests: log.info('Stopping previous test') self.client_tests[client_id].stop() self.client_tests[data['client_id']] = test test.run() else: log.warn('No test given') @action def stop_test(self, data): """ Stop any tests from that client :param data: Message specific data: client_id: Client to stop tests for """ client_id = data['client_id'] if client_id in self.client_tests: log.info('slave({}).stop_test client({})'.format(self.id, client_id)) self.client_tests[client_id].stop() del self.client_tests[client_id] else: self.master.test_stopped()
class Client(Actionable): """ This class is the local representation of a remote client (browser) When a websocket message is received it is routed into the proper local instance When the master sends a message to a client it goes through this class """ def __init__(self, master, websocket, id): self.master = master self.websocket = websocket self.id = id # This is syntax for sending message through a websocket. # client.send.<action>(data) self.send = Sender(self.send_data) def send_data(self, data): self.websocket.send(json.dumps(data)) def handle_websocket_message(self, msg): """ This handles any message that comes in by looking at the action key and seeing if it is a local method, then calling it :param msg: JSON encoded data from the websocket """ if msg: data = json.loads(msg) if "action" in data: self.run_action(data) @action def get_id(self, data): self.send.set_id({"client_id": self.id}) @action def quit(self, data): """ Websocket telling a slave to quit :param data: Message specific data: slave_id: slave to tell to quit """ if "slave_id" in data: self.master.remove_slave(data["slave_id"]) else: self.send.error({"error": "Id not specified in data"}) @action def request_slaves(self, data): """ Websocket requesting slave data """ data = {"slaves": dict((slave.id, slave.last_beat) for slave in self.master.slave_registry.values())} # Send a message with action=receive_slaves to the browser self.send.receive_slaves(data) @action def set_sink(self, data): """ Set the sink on a particular slave, which sets the sink for all slaves on that host :param data: message specific data: slave_id: Slave to set the sink for (unused) sink_id: Slave to set the sink to """ if "sink_id" not in data or "slave_id" not in data: raise Exception("Missing sink_id or slave_id") # Let the master handle actually doing this self.master.set_sink(data["slave_id"], data["sink_id"], self) @action def request_available_tests(self, data): """ Client asking to load tests :param data: """ tests_glob = os.path.join(os.path.dirname(__file__), TEST_DIR, "*.json") files = glob.glob(tests_glob) files_stripped = [os.path.basename(filename)[:-5] for filename in files] # Send list of files to the client self.send.receive_available_tests({"tests": files_stripped}) @action def request_test(self, data): """ Client asking to load a test :param data: Message specific data: name: filename to load """ filename = os.path.join(os.path.dirname(__file__), TEST_DIR, data["name"] + ".json") # TODO: Prevent path traversal if os.path.exists(filename): with open(filename) as file: return self.send.receive_test({"test": json.loads(file.read())}) return self.send.error({"error": "Test not found"}) @action def save_test(self, data): """ Save a test to a file :param data: Message specific data: name: Filename to save to test: data representing a test """ if "test" in data: test = data["test"] if "name" in test: # TODO: Prevent path traversal filename = os.path.join(os.path.dirname(__file__), TEST_DIR, test["name"] + ".json") with open(filename, "w") as file: file.write(json.dumps(test)) self.send.save_successful() @action def delete_test(self, data): """ Delete a test :param data: Message specific data: test_name: Name of test to delete """ if "test_name" in data: # TODO: Prevent path traversal filename = os.path.join(os.path.dirname(__file__), TEST_DIR, data["test_name"] + ".json") if os.path.exists(filename): os.remove(filename) else: return self.send.error({"error": "Test not found"}) else: return self.send.error({"error": "Cannot delete test, not found"}) @action def run_test(self, data): """ Send a test to all slaves to run :param data: Message specific data: test: Test data to run """ if "test" in data: data["client_id"] = self.id if "runs" in data["test"]: try: # Each slave gets a portion of the total runs data["runs"] = int(data["test"]["runs"]) / len(self.master.slave_registry) except ValueError: data["runs"] = 1 else: data["runs"] = 1 self.master.slaves.run_test(data) @action def stop_test(self, data): """ Halt all tests run by this client :param data: """ data["client_id"] = self.id if len(self.master.slave_registry) != 0: self.master.slaves.stop_test(data) else: self.send.test_stopped()