Hold messages some time, allow multiple clients

Author: SemaiCZE

monitor/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "0.1.0"
+__version__ = "0.2.0"
 
 __all__ = ['websocket_connections', 'zeromq_connection', 'config_manager', 'main']
 

monitor/main.py

@@ -29,7 +29,7 @@ def main():
     # get logger
     logger = init_logger(*config.get_logger_settings())
     # here we'll store all active connections
-    connections = ClientConnections(logger)
+    connections = ClientConnections(logger, loop)
 
     websock_server = None
     try:

monitor/test/test_ClientConnections.py

@@ -9,23 +9,34 @@
 class TestClientConnections(unittest.TestCase):
     def setUp(self):
         logger = MagicMock()
-        self._connections = ClientConnections(logger)
+        loop = MagicMock
+        self._connections = ClientConnections(logger, loop)
 
     def test_add_client(self):
         queue = self._connections.add_client("1234")
         self.assertNotEqual(queue, None, "No queue returned")
         self.assertIsInstance(queue, asyncio.Queue, "Wrong type of returned value")
-        self.assertEqual(queue, self._connections._clients["1234"],
+        self.assertEqual(queue, self._connections._clients["1234"][0],
                          "Queue not present at requested id")
         queue2 = self._connections.add_client("456777")
         self.assertNotEqual(queue, queue2, "Different clients have same future")
+        self._connections.add_client("1234")
+        self.assertEqual(len(self._connections._clients["1234"]), 2, "No concurrent clients for same id")
+
+    def test_remove_channel(self):
+        self._connections.add_client("1234")
+        self._connections.remove_channel("4567")
+        self.assertEqual(len(self._connections._clients), 1, "Nonexisting channel removed")
+        self._connections.remove_channel("1234")
+        self.assertEqual(len(self._connections._clients), 0, "Existing channel not removed")
 
     def test_remove_client(self):
-        fut = self._connections.add_client("1234")
-        self._connections.remove_client("4567")
-        self.assertEqual(len(self._connections._clients), 1, "Nonexisting client removed")
-        self._connections.remove_client("1234")
-        self.assertEqual(len(self._connections._clients), 0, "Existing client not removed")
+        queue1 = self._connections.add_client("1234")
+        queue2 = self._connections.add_client("1234")
+        self._connections.remove_client("4567", queue1)
+        self.assertEqual(len(self._connections._clients["1234"]), 2, "Nonexisting client removed")
+        self._connections.remove_client("1234", queue2)
+        self.assertEqual(len(self._connections._clients["1234"]), 1, "Existing client not removed")
 
     def test_remove_all_clients(self):
         self._connections.add_client("1234")
@@ -37,11 +48,11 @@ def test_remove_all_clients(self):
 
     def test_send_message(self):
         queue = self._connections.add_client("1234")
-        ret = self._connections.send_message("1234", "testing message")
-        self.assertTrue(ret, "Message not sent")
+        self._connections.send_message("1234", "testing message")
         self.assertEqual(queue.get_nowait(), "testing message", "Result from queue differs")
-        ret2 = self._connections.send_message("1234", "msg two")
-        self.assertTrue(ret2, "Add second message to queue failed")
+        self._connections.send_message("1234", "msg two")
         self.assertEqual(queue.get_nowait(), "msg two", "Queue has not second message")
-        ret3 = self._connections.send_message("51236", "random string")
-        self.assertFalse(ret3, "Message sent to nonexisting host id")
+
+        self._connections.send_message("51236", "random string")
+        queue2 = self._connections.add_client("51236")
+        self.assertEqual(queue2.get_nowait(), "random string", "Result from queue differs")

monitor/test/test_websock_server.py

@@ -48,7 +48,7 @@ def test_connection_handler(self):
         websocket_mock.recv.assert_called_once_with()
         websocket_mock.send.assert_called_once_with('result text')
         connection_mock.add_client.assert_called_once_with("1234")
-        connection_mock.remove_client.assert_called_once_with("1234")
+        connection_mock.remove_client.assert_called_once_with("1234", queue)
 
     @patch('asyncio.set_event_loop')
     def test_run(self, mock_set_loop):

monitor/websocket_connections.py

@@ -16,43 +16,79 @@ class ClientConnections:
     thread safe.
     """
 
-    def __init__(self, logger):
+    def __init__(self, logger, loop):
         """
         Initialize empty client dictionary
+
+        :param logger: System logger.
+        :param loop: Main asyncio event loop.
         """
         self._clients = dict()
+        self._saved_messages = dict()
         self._logger = logger
+        self._loop = loop
 
     def add_client(self, id):
         """
         Register new client which wants to receive messages to identifier 'id'.
-        Only one subscriber per stream is allowed. Latter overrides previous one.
+        If there are any such messages, they are sent immediately. There can be
+        more subscribers per stream.
 
         :param id: Identifier of required stream of messages
         :return: Returns new asyncio.Queue on which can be wait for by
             'yield from' command.
         """
         new_queue = asyncio.Queue()
-        self._clients[id] = new_queue
+        if id not in self._clients.keys():
+            self._clients[id] = []
+        self._clients[id].append(new_queue)
+
+        # if there are already any messages, send them
+        if id in self._saved_messages.keys():
+            for msg in self._saved_messages[id]:
+                new_queue.put_nowait(msg)
+
         self._logger.debug("client connection: new client '{}' registered".format(id))
         return new_queue
 
-    def remove_client(self, id):
+    def remove_channel(self, id):
         """
-        Remove client listening on 'id' message stream. This means removing associated
-        queue and deleting the entry from internal dictionary.
-        If no such client exists, nothing is done.
+        Remove all clients listening on 'id' channel. This means removing all associated
+        queues and received messages. If no such channel exists, nothing is done.
+        This method is called 5 minutes after last message of each channel.
 
         :param id: Identifier of required stream of messages
         :return: Nothing
         """
+        if id in self._saved_messages.keys():
+            del self._saved_messages[id]
+
         if id in self._clients.keys():
             del self._clients[id]
+            self._logger.debug("client connection: channel '{}' removed".format(id))
+        else:
+            self._logger.debug("client connection: channel '{}' removing failed - "
+                               " not present".format(id))
+
+    def remove_client(self, id, queue):
+        """
+        Remove client listening on 'id' message stream with queue 'queue'.
+        This means removing associated queue and deleting the entry from internal dictionary.
+        If no such client exists, nothing is done.
+
+        :param id: Identifier of required stream of messages
+        :param queue: Queue associated with client to be removed
+        :return: Nothing
+        """
+        if id in self._clients.keys():
+            clients = self._clients[id]
+            clients.remove(queue)
             self._logger.debug("client connection: client '{}' removed".format(id))
         else:
             self._logger.debug("client connection: client '{}' removing failed - "
                                " not present".format(id))
 
+
     def remove_all_clients(self):
         """
         Remove all registered clients. This method could be called on app
@@ -61,26 +97,31 @@ def remove_all_clients(self):
         :return: Nothing
         """
         self._clients.clear()
+        self._saved_messages.clear()
         self._logger.debug("client connection: all clients removed")
 
     def send_message(self, id, message):
         """
         Send 'message' to client listening on stream with 'id'. If 'id' is not
-        known, the message is silently dropped. The message is put into queue,
-        so no message will get lost.
+        known, the message is saved for latter use. Messages for connected
+        clients are put into queues, so no message will get lost.
 
         :param id: Identifier of required stream of messages
         :param message: String containing text to be sent
-        :return: Returns True if message was sent, False otherwise
+        :return: Nothing
         """
+
+        if id not in self._saved_messages.keys():
+            self._saved_messages[id] = []
+        self._saved_messages[id].append(message)
+
         if id in self._clients.keys():
-            queue = self._clients[id]
-            queue.put_nowait(message)
-            return True
-        else:
-            self._logger.warning("client connection: Dropping message '{}' for "
-                                 "non-existing stream '{}'".format(message, id))
-        return False
+            for queue in self._clients[id]:
+                queue.put_nowait(message)
+
+        # on last message schedule removing whole channel after 5 minute wait
+        if message is None:
+            self._loop.call_later(5*60, self.remove_channel, id)
 
 
 class WebsocketServer(threading.Thread):
@@ -118,8 +159,8 @@ def connection_handler(self, websocket, path):
 
         :param websocket: Socket with request
         :param path: Requested path of socket (not used)
-        :return: Returns when socket is closed or future from ClientConnections
-            is cancelled.
+        :return: Returns when socket is closed or poison pill is found in message queue
+            from ClientConnections.
         """
         wanted_id = None
         try:
@@ -138,7 +179,8 @@ def connection_handler(self, websocket, path):
         except websockets.ConnectionClosed:
             self._logger.info("websocket server: connection closed for channel '{}'". format(wanted_id))
         finally:
-            self._connections.remove_client(wanted_id)
+            self._connections.remove_client(wanted_id, queue)
+
 
     def run(self):
         """

monitor/zeromq_connection.py

@@ -70,4 +70,9 @@ def start(self, message_callback):
                 break
             # call registered callback with given data
             message_callback(client_id, data)
+
+            # after last message (command FINISHED) send also poison pill
+            # to close listening sockets
+            if decoded_message[1] == "FINISHED":
+                message_callback(client_id, None)
         return True