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()
class MasterApplication(Actionable): """ This class represents the master controlling server. It tracks slaves, sinks, and clients. """ slave_registry = {} sink_registry = {} next_slave_id = 0 next_client_id = 0 client_registry = {} clients = Sender(None) removed_slaves = set() def __init__(self, context): log.info("Master initialized") # wsgi context, unused: self.context = context # Slaves send data in via this socket self.slave_in = context.socket(zmq.PULL) self.slave_in.bind("tcp://*:{}".format(SLAVE_MASTER_PORT)) # listen_to_slaves handles any incoming messages gevent.spawn(self.listen_to_slaves) # This broadcasts to all slaves simultaneously. # TODO: Selective slave message sending self.slave_out = context.socket(zmq.PUB) self.slave_out.bind("tcp://*:{}".format(SLAVE_PUB_PORT)) # What a clever boy I am (clever code is bad code) # Allow self.slaves.<action>(data) magic self.slaves = Sender(self.broadcast_data_to_slaves) # Same for clients (where client is a browser connected via websocket) self.clients = Sender(self.broadcast_data_to_clients) def broadcast_data_to_slaves(self, data): """ Send data to all slaves as a bson dump :param data: Data to serialize """ self.slave_out.send(bson.dumps(data)) def broadcast_data_to_clients(self, data): """ Send data to all web clients (in serial) as a json dump :param data: Data to serialize """ # TODO: parallelize for client in self.client_registry.values(): client.send_data(data) def __call__(self, environ, start_response): """ This class is called whenever a browser tries to connect. Note: This is green-threaded via gevent :param environ: The wsgi environment :param start_response: A helper to return http responses :return: """ # Respond to any websocket messages ws = environ["wsgi.websocket"] self.listen_to_websocket(ws) def listen_to_websocket(self, ws): """ Listen to incoming messages on the websocket and delegate them to a message handler. :param ws: Websocket to listen to """ client = self.add_client(ws) log.info("Websocket {} connected!".format(client.id)) # Listen for messages until they disconnect while True: try: msg = ws.receive() client.handle_websocket_message(msg) except WebSocketError as e: self.disconnect_client(client) break def get_next_client_id(self): """ :return: an incrementing id to identify clients """ self.next_client_id += 1 return self.next_client_id - 1 def add_client(self, ws): """ Create a new Client instance and store it in the client registry :param ws: Websocket to associate with the client so the instance can send messages :return: The client instance :rtype Client: """ id = self.get_next_client_id() client = Client(self, ws, id) self.client_registry[id] = client return client def disconnect_client(self, client): """ Close the connection to this specific client and deregister them :param client: Client to disconnect """ if client.id in self.client_registry: client.websocket.close() del self.client_registry[client.id] def listen_to_slaves(self): """ Delegate any messages received from the slaves to the proper handler :return: """ log.info("Listening to slaves") while True: self.handle_slave_message(self.slave_in.recv()) def handle_slave_message(self, msg): """ Slaves send in heartbeats and test results. Handle finding the associated local slave instance and interpreting the message. :param msg: A bson encoded dict with message data. action: A string representing the action being taken. Must have an associated function on the local slave class with the @action decorator slave_id: The id assigned to the slave. If this is not present the slave doesn't have one and should be sent one <message specific data> """ # TODO: Think through the security implications of letting slaves report their own ids and uuids data = bson.loads(msg) if "action" in data: # Slaves without a valid slave_id are unregistered and need to be created locally # TODO: Move slave creation into load_slave if "slave_id" not in data or data["slave_id"] not in self.slave_registry: if data["action"] == "heartbeat": # Heartbeats are typically the way slaves are discovered slave = self.new_slave(data) # TODO: handle slave creation failure else: log.error("Got non-heartbeat message from unknown slave") return else: # Load the local instance slave = self.load_slave(data) # TODO: handle save loading failure assert slave # run_action looks up the @action method on the slave and runs it with data slave.run_action(data) def load_slave(self, data): """ Inspect the message data and retrieve the associated slave :param data: Message data :return: The loaded slave :rtype Slave: """ if "slave_id" not in data: log.warn("load_slave without slave_id") return slave_id = data["slave_id"] if slave_id not in self.slave_registry: log.warn("load_slave couldn't find slave") return return self.slave_registry[slave_id] def new_slave(self, heartbeat_data): """ Create a slave from its heartbeat :param heartbeat_data: :return: The created Slave :rtype Slave: """ slave = Slave(self, self.get_next_slave_id()) # Modify the data, providing a slave_id to run_action if it needs it heartbeat_data["slave_id"] = slave.id # Kinda a hack, we want to get the sink_host before the action is run, # and that requires heartbeat data, so we store it if we have it slave.last_beat = heartbeat_data # Broadcast a message to all slaves associating that uuid with the local slave_id self.slaves.set_id({"new_slave_id": slave.id, "slave_uuid": heartbeat_data["slave_uuid"]}) # Slaves send test data to the sink, and the sink collates it before sending it to the server. # There is one sink per unique hostname sink = self.pick_sink(slave) # TODO: handle if no sink assert sink # Tell the new slave what sink to use self.slaves.set_sink({"sink_id": sink.id, "sink_host": sink.get_sink_host()}) # Register the slave so it can be found again self.slave_registry[slave.id] = slave return slave def get_next_slave_id(self): """ :return: An auto incrementing id to identify slaves """ self.next_slave_id += 1 return self.next_slave_id def pick_sink(self, slave): """ Pick a sink on the same host as the given slave :param slave: Slave to find a sink for :return: """ hostname = slave.get_hostname() if hostname: # If it is a new host then it is the sink! if hostname not in self.sink_registry: return slave else: # otherwise look up the hostname return self.sink_registry[hostname] def set_sink(self, slave_id, sink_id, client=None): """ Forcefully set the sink for all slaves on a host. :param slave_id: Unused :param sink_id: The slave_id to promote to a sink :param client: The client that requested the sink change """ # TODO: Move client error reporting into the client class if sink_id not in self.slave_registry: log.warn("set_sink to unknown sink") if client: client.send.error({"error": "Sink({}) not found".format(sink_id)}) return # Look up the hostname for the promoted sink sink_host = self.slave_registry[sink_id].get_sink_host() if not sink_host: log.warn("Could sink_host not found") if client: client.send.error({"error": "Sink({}) does not have a host yet".format(sink_id)}) return data = {"sink_host": sink_host, "sink_id": sink_id} # Tell all slaves with the host <sink_host> to connect to the given sink self.slaves.set_sink(data) def register_sink(self, sink): """ Record when a slave connects to a sink :param sink: Sink that has been connected to """ self.sink_registry[sink.get_hostname()] = sink def remove_slave(self, id): """ Remove, disconnect, and shut down the given slave id :param id: Slave to remove """ if id in self.slave_registry: log.info("Killing slave {}".format(id)) del self.slave_registry[id] # Any slave that reports to that id will quit # TODO: switch to slave.broadcast_to.quit() self.slaves.quit({"slave_id": id}) self.clients.slave_disconnected({"slave_id": id}) def check_slaves(self): """ Iterate through all slaves and check the last time they heartbeat If it is too long ago remove them :return: """ the_time = int(time()) # Since we modify the dict (and that messes with iteration) first get a static list of all of them slave_ids = self.slave_registry.keys() for slave_id in slave_ids: slave = self.slave_registry[slave_id] if the_time - slave.last_beat.get("generated", 0) > HEARTBEAT_PERIOD * BEATS_TO_KILL: self.remove_slave(slave_id) def watch_slaves(self): """ This will periodically check slaves :return: """ while True: self.check_slaves() gevent.sleep(HEARTBEAT_PERIOD)