Send reply messages from the main process (#88)

* Send reply messages directly

* Always forever autocommit

* Clean up SQL

* Log cleanup

* Remove tasks file no longer needed

* Add type hinting and rename from review

Author: AlanCoding

dispatcher/brokers/pg_notify.py

@@ -18,14 +18,21 @@
 
 async def acreate_connection(**config) -> psycopg.AsyncConnection:
     "Create a new asyncio connection"
-    return await psycopg.AsyncConnection.connect(**config)
+    connection = await psycopg.AsyncConnection.connect(**config)
+    if not connection.autocommit:
+        await connection.set_autocommit(True)
+    return connection
 
 
 def create_connection(**config) -> psycopg.Connection:
-    return psycopg.Connection.connect(**config)
+    connection = psycopg.Connection.connect(**config)
+    if not connection.autocommit:
+        connection.set_autocommit(True)
+    return connection
 
 
 class Broker:
+    NOTIFY_QUERY_TEMPLATE = 'SELECT pg_notify(%s, %s);'
 
     def __init__(
         self,
@@ -63,13 +70,18 @@ def __init__(
 
         if config:
             self._config: dict = config.copy()
-            self._config['autocommit'] = True
         else:
             self._config = {}
 
         self.channels = channels
         self.default_publish_channel = default_publish_channel
 
+        # If we are in the notification loop (receiving messages),
+        # then we have to break out before sending messages
+        # These variables track things so that we can exit, send, and re-enter
+        self.notify_loop_active: bool = False
+        self.notify_queue: list = []
+
     def get_publish_channel(self, channel: Optional[str] = None) -> str:
         "Handle default for the publishing channel for calls to publish_message, shared sync and async"
         if channel is not None:
@@ -112,23 +124,35 @@ async def aprocess_notify(self, connected_callback: Optional[Callable] = None) -
 
             while True:
                 logger.debug('Starting listening for pg_notify notifications')
+                self.notify_loop_active = True
                 async for notify in connection.notifies():
                     yield notify.channel, notify.payload
+                    if self.notify_queue:
+                        break
+                self.notify_loop_active = False
+                for reply_to, reply_message in self.notify_queue:
+                    await self.apublish_message_from_cursor(cur, channel=reply_to, message=reply_message)
+                self.notify_queue = []
+
+    async def apublish_message_from_cursor(self, cursor: psycopg.AsyncCursor, channel: Optional[str] = None, message: str = '') -> None:
+        """The inner logic of async message publishing where we already have a cursor"""
+        await cursor.execute(self.NOTIFY_QUERY_TEMPLATE, (channel, message))
 
     async def apublish_message(self, channel: Optional[str] = None, message: str = '') -> None:  # public
         """asyncio way to publish a message, used to send control in control-and-reply
 
         Not strictly necessary for the service itself if it sends replies in the workers,
         but this may change in the future.
         """
+        if self.notify_loop_active:
+            self.notify_queue.append((channel, message))
+            return
+
         connection = await self.aget_connection()
         channel = self.get_publish_channel(channel)
 
         async with connection.cursor() as cur:
-            if not message:
-                await cur.execute(f'NOTIFY {channel};')
-            else:
-                await cur.execute(f"NOTIFY {channel}, '{message}';")
+            await self.apublish_message_from_cursor(cur, channel=channel, message=message)
 
         logger.debug(f'Sent pg_notify message of {len(message)} chars to {channel}')
 
@@ -159,10 +183,7 @@ def publish_message(self, channel: Optional[str] = None, message: str = '') -> N
         channel = self.get_publish_channel(channel)
 
         with connection.cursor() as cur:
-            if message:
-                cur.execute('SELECT pg_notify(%s, %s);', (channel, message))
-            else:
-                cur.execute(f'NOTIFY {channel};')
+            cur.execute(self.NOTIFY_QUERY_TEMPLATE, (channel, message))
 
         logger.debug(f'Sent pg_notify message of {len(message)} chars to {channel}')
 
@@ -189,7 +210,6 @@ def connection_saver(**config) -> psycopg.Connection:
     Dispatcher does not manage connections, so this a simulation of that.
     """
     if connection_save._connection is None:
-        config['autocommit'] = True
         connection_save._connection = create_connection(**config)
     return connection_save._connection
 
@@ -202,6 +222,5 @@ async def async_connection_saver(**config) -> psycopg.AsyncConnection:
     Dispatcher does not manage connections, so this a simulation of that.
     """
     if connection_save._async_connection is None:
-        config['autocommit'] = True
         connection_save._async_connection = await acreate_connection(**config)
     return connection_save._async_connection

dispatcher/control.py

@@ -22,26 +22,29 @@ class ControlCallbacks:
     it exists to interact with producers, using variables relevant to the particular
     control message being sent"""
 
-    def __init__(self, queuename, send_data, expected_replies):
+    def __init__(self, queuename, send_data, expected_replies) -> None:
         self.queuename = queuename
         self.send_data = send_data
         self.expected_replies = expected_replies
 
-        self.received_replies = []
+        # received_replies only tracks the reply message, not the channel name
+        # because they come via a temporary reply_to channel and that is not user-facing
+        self.received_replies: list[str] = []
         self.events = ControlEvents()
         self.shutting_down = False
 
-    async def process_message(self, payload, producer=None, channel=None):
+    async def process_message(self, payload, producer=None, channel=None) -> tuple[Optional[str], Optional[str]]:
         self.received_replies.append(payload)
         if self.expected_replies and (len(self.received_replies) >= self.expected_replies):
             self.events.exit_event.set()
+        return (None, None)
 
     async def connected_callback(self, producer) -> None:
         payload = json.dumps(self.send_data)
         await producer.notify(channel=self.queuename, message=payload)
         logger.info('Sent control message, expecting replies soon')
 
-    def fatal_error_callback(self, *args):
+    def fatal_error_callback(self, *args) -> None:
         if self.shutting_down:
             return
 

dispatcher/producers/brokered.py

@@ -36,7 +36,9 @@ async def produce_forever(self, dispatcher) -> None:
         self.dispatcher = dispatcher
         async for channel, payload in self.broker.aprocess_notify(connected_callback=self.connected_callback):
             self.produced_count += 1
-            await dispatcher.process_message(payload, producer=self, channel=channel)
+            reply_to, reply_payload = await dispatcher.process_message(payload, producer=self, channel=channel)
+            if reply_to:
+                await self.notify(channel=reply_to, message=reply_payload)
 
     async def notify(self, channel: Optional[str] = None, message: str = '') -> None:
         await self.broker.apublish_message(channel=channel, message=message)

dispatcher/service/main.py

@@ -170,8 +170,16 @@ def create_delayed_task(self, message: dict) -> None:
         capsule.task = new_task
         self.delayed_messages.append(capsule)
 
-    async def process_message(self, payload: dict, producer: Optional[BaseProducer] = None, channel: Optional[str] = None) -> None:
-        # Convert payload from client into python dict
+    async def process_message(
+        self, payload: dict, producer: Optional[BaseProducer] = None, channel: Optional[str] = None
+    ) -> tuple[Optional[str], Optional[str]]:
+        """Called by producers to trigger a new task
+
+        Convert payload from producer into python dict
+        Process uuid default
+        Delay tasks when applicable
+        Send to next layer of internal processing
+        """
         # TODO: more structured validation of the incoming payload from publishers
         if isinstance(payload, str):
             try:
@@ -182,7 +190,7 @@ async def process_message(self, payload: dict, producer: Optional[BaseProducer]
             message = payload
         else:
             logger.error(f'Received unprocessable type {type(payload)}')
-            return
+            return (None, None)
 
         # A client may provide a task uuid (hope they do it correctly), if not add it
         if 'uuid' not in message:
@@ -195,28 +203,24 @@ async def process_message(self, payload: dict, producer: Optional[BaseProducer]
             # NOTE: control messages with reply should never be delayed, document this for users
             self.create_delayed_task(message)
         else:
-            await self.process_message_internal(message, producer=producer)
+            return await self.process_message_internal(message, producer=producer)
+        return (None, None)
 
-    async def process_message_internal(self, message: dict, producer=None) -> None:
+    async def process_message_internal(self, message: dict, producer=None) -> tuple[Optional[str], Optional[str]]:
+        """Route message based on needed action - delay for later, return reply, or dispatch to worker"""
         if 'control' in message:
             method = getattr(self.ctl_tasks, message['control'])
             control_data = message.get('control_data', {})
             returned = await method(self, **control_data)
             if 'reply_to' in message:
-                logger.info(f"Control action {message['control']} returned {returned}, sending via worker")
+                logger.info(f"Control action {message['control']} returned {returned}, sending back reply")
                 self.control_count += 1
-                await self.pool.dispatch_task(
-                    {
-                        'task': 'dispatcher.service.tasks.reply_to_control',
-                        'args': [message['reply_to'], json.dumps(returned)],
-                        'uuid': f'control-{self.control_count}',
-                        'control': 'reply',  # for record keeping
-                    }
-                )
+                return (message['reply_to'], json.dumps(returned))
             else:
                 logger.info(f"Control action {message['control']} returned {returned}, done")
         else:
             await self.pool.dispatch_task(message)
+        return (None, None)
 
     async def start_working(self) -> None:
         logger.debug('Filling the worker pool')

dispatcher/service/pool.py

@@ -107,7 +107,6 @@ def __init__(self, max_workers: int, process_manager: ProcessManager, settings:
         self.start_worker_task: Optional[Task] = None
         self.shutting_down = False
         self.finished_count: int = 0
-        self.control_count: int = 0
         self.canceled_count: int = 0
         self.discard_count: int = 0
         self.shutdown_timeout = 3
@@ -369,8 +368,6 @@ async def process_finished(self, worker, message) -> None:
         async with self.management_lock:
             if worker.is_active_cancel and result == '<cancel>':
                 self.canceled_count += 1
-            elif 'control' in worker.current_task:
-                self.control_count += 1
             else:
                 self.finished_count += 1
             worker.mark_finished_task()

dispatcher/service/tasks.py

@@ -100,7 +100,7 @@ async def test_cancel_task(apg_dispatcher, pg_message, pg_control):
     await asyncio.wait_for(clearing_task, timeout=3)
 
     pool = apg_dispatcher.pool
-    assert [pool.finished_count, pool.canceled_count, pool.control_count] == [0, 1, 1], 'cts: [finished, canceled, control]'
+    assert [pool.finished_count, pool.canceled_count, apg_dispatcher.control_count] == [0, 1, 1], 'cts: [finished, canceled, control]'
 
 
 @pytest.mark.asyncio
@@ -117,13 +117,10 @@ async def test_message_with_delay(apg_dispatcher, pg_message, pg_control):
     assert running_job['uuid'] == 'delay_task'
     await asyncio.wait_for(apg_dispatcher.pool.events.work_cleared.wait(), timeout=3)
     pool = apg_dispatcher.pool
-    assert [pool.finished_count, pool.canceled_count, pool.control_count] == [0, 0, 1], 'cts: [finished, canceled, control]'
-    # Completing the reply itself will be a work_cleared event, so we have to clear the event
-    apg_dispatcher.pool.events.work_cleared.clear()
 
     # Wait for task to finish, assertions after completion
     await asyncio.wait_for(apg_dispatcher.pool.events.work_cleared.wait(), timeout=3)
-    assert [pool.finished_count, pool.canceled_count, pool.control_count] == [1, 0, 1], 'cts: [finished, canceled, control]'
+    assert [pool.finished_count, pool.canceled_count, apg_dispatcher.control_count] == [1, 0, 1], 'cts: [finished, canceled, control]'
 
 
 @pytest.mark.asyncio

tests/integration/test_main.py

@@ -4,8 +4,7 @@
 import os
 import sys
 
-from dispatcher.factories import get_publisher_from_settings
-from dispatcher.control import Control
+from dispatcher.factories import get_publisher_from_settings, get_control_from_settings
 from dispatcher.utils import MODULE_METHOD_DELIMITER
 from dispatcher.config import setup
 
@@ -49,7 +48,7 @@ def main():
     print('performing a task cancel')
     # submit a task we will "find" two different ways
     broker.publish_message(message=json.dumps({'task': 'lambda: __import__("time").sleep(3.1415)', 'uuid': 'foobar'}))
-    ctl = Control('test_channel')
+    ctl = get_control_from_settings()
     canceled_jobs = ctl.control_with_reply('cancel', data={'uuid': 'foobar'})
     print(json.dumps(canceled_jobs, indent=2))