Start on synchronous testing

Close queues, hopefully fixed race condition

Author: AlanCoding

dispatcher/service/pool.py

@@ -90,6 +90,7 @@ async def stop(self) -> None:
         except asyncio.TimeoutError:
             logger.error(f'Worker {self.worker_id} pid={self.process.pid} failed to send exit message in 3 seconds')
             self.status = 'error'  # can signal for result task to exit, since no longer waiting for it here
+            self.process.message_queue.close()
 
         await self.join()  # If worker fails to exit, this returns control without raising an exception
 
@@ -461,6 +462,18 @@ async def shutdown(self) -> None:
             except asyncio.CancelledError:
                 logger.info('The finished task was canceled, but we are shutting down so that is alright')
 
+        if self.management_task:
+            logger.info('Canceling worker management task')
+            self.management_task.cancel()
+            try:
+                await asyncio.wait_for(self.management_task, timeout=self.shutdown_timeout)
+            except asyncio.TimeoutError:
+                logger.error('The scaleup task failed to shut down')
+            except asyncio.CancelledError:
+                pass  # intended
+
+        self.process_manager.shutdown()
+
         logger.info('Pool is shut down')
 
     def active_task_ct(self) -> int:
@@ -576,6 +589,7 @@ async def read_results_forever(self) -> None:
                 async with self.workers.management_lock:
                     worker.status = 'exited'
                     worker.exit_msg_event.set()
+                    worker.process.message_queue.close()
 
                 if self.shutting_down:
                     if all(worker.inactive for worker in self.workers):

dispatcher/service/process.py

@@ -99,6 +99,9 @@ async def read_finished(self) -> dict[str, Union[str, int]]:
         message = await self.get_event_loop().run_in_executor(None, self.finished_queue.get)
         return message
 
+    def shutdown(self):
+        self.finished_queue.close()
+
 
 class ForkServerManager(ProcessManager):
     mp_context = 'forkserver'

dispatcher/testing/__init__.py

@@ -0,0 +1,3 @@
+from .subprocess import adispatcher_service, dispatcher_service
+
+__all__ = ['adispatcher_service', 'dispatcher_service']

dispatcher/testing/subprocess.py

@@ -0,0 +1,136 @@
+import asyncio
+import contextlib
+import logging
+import multiprocessing
+import sys
+from multiprocessing.context import BaseContext
+from types import ModuleType
+from typing import Any, AsyncGenerator, Union
+
+from ..config import DispatcherSettings
+from ..factories import from_settings
+from ..service.main import DispatcherMain
+
+logger = logging.getLogger(__name__)
+
+
+class CommunicationItems:
+    """Various things used for communication between the parent process and the subprocess service
+
+    This will be passed in the call to the subprocess.
+    """
+
+    def __init__(self, main_events: tuple[str], pool_events: tuple[str], context: Union[BaseContext, ModuleType]) -> None:
+        self.q_in: multiprocessing.Queue = context.Queue()
+        self.q_out: multiprocessing.Queue = context.Queue()
+        self.main_events = main_events
+        self.pool_events = pool_events
+
+
+@contextlib.asynccontextmanager
+async def adispatcher_service(config: dict) -> AsyncGenerator[DispatcherMain, Any]:
+    dispatcher = None
+    try:
+        settings = DispatcherSettings(config)
+        dispatcher = from_settings(settings=settings)  # type: ignore[arg-type]
+
+        await dispatcher.connect_signals()
+        await dispatcher.start_working()
+        await dispatcher.wait_for_producers_ready()
+        await dispatcher.pool.events.workers_ready.wait()
+
+        assert dispatcher.pool.finished_count == 0  # sanity
+        assert dispatcher.control_count == 0
+
+        yield dispatcher
+    finally:
+        if dispatcher:
+            try:
+                await dispatcher.shutdown()
+                await dispatcher.cancel_tasks()
+            except Exception:
+                logger.exception('shutdown had error')
+
+
+async def asyncio_target(config: dict, comms: CommunicationItems) -> None:
+    loop = asyncio.get_event_loop()
+    async with adispatcher_service(config) as dispatcher:
+        comms.q_out.put('ready')
+
+        events: dict[str, asyncio.Event] = {}
+        for event_name in comms.main_events:
+            events[event_name] = getattr(dispatcher.events, event_name)
+        for event_name in comms.pool_events:
+            events[event_name] = getattr(dispatcher.pool.events, event_name)
+
+        event_tasks: dict[str, asyncio.Task] = {}
+        for event_name, event in events.items():
+            event_tasks[event_name] = asyncio.create_task(event.wait(), name=f'waiting_for_{event_name}')
+
+        new_message_task = None
+
+        while True:
+            if new_message_task is None:
+                new_message_task = loop.run_in_executor(None, comms.q_in.get)
+
+            all_tasks = list(event_tasks.values()) + [new_message_task]
+            await asyncio.wait(all_tasks, return_when=asyncio.FIRST_COMPLETED)
+
+            # Update our parent process with any events they requested from us
+            for event_name, event in events.items():
+                if event.is_set():
+                    comms.q_out.put(event_name)
+                    # await loop.run_in_executor(None, comms.q_out.put, event_name)
+                event.clear()
+                event_tasks[event_name] = asyncio.create_task(event.wait())
+
+            # If no no instructions came from parent then work is done, continue loop
+            if not new_message_task.done():
+                continue
+
+            message = new_message_task.result()
+            new_message_task = None
+
+            if message == 'stop':
+                print('shutting down pool server')
+                for event in events.values():
+                    event.set()  # close out other tasks
+                await dispatcher.shutdown()
+                break
+            else:
+                eval(message)
+
+
+def subprocess_main(config, comms):
+    loop = asyncio.new_event_loop()
+    try:
+        loop.run_until_complete(asyncio_target(config, comms))
+    except Exception:
+        # The main process is very likely waiting for message of an event
+        # and exceptions may not automatically halt the test, so give a value
+        comms.q_out.put('error')
+        raise
+    finally:
+        loop.close()
+
+
+@contextlib.contextmanager
+def dispatcher_service(config, main_events=(), pool_events=()):
+    ctx = multiprocessing.get_context('spawn')
+    comms = CommunicationItems(main_events=main_events, pool_events=pool_events, context=ctx)
+    process = multiprocessing.Process(target=subprocess_main, args=(config, comms))
+    try:
+        process.start()
+        ready_msg = comms.q_out.get()
+        if ready_msg != 'ready':
+            raise RuntimeError(f'Never got "ready" message from server, got {ready_msg}')
+        yield comms
+    finally:
+        comms.q_in.put('stop')
+        process.join(timeout=1)
+        if process.is_alive():
+            process.terminate()  # SIGTERM
+        comms.q_in.close()
+        comms.q_out.close()
+        sys.stdout.flush()
+        sys.stderr.flush()

tests/conftest.py

@@ -83,14 +83,6 @@ def conn_config():
     return {'conninfo': CONNECTION_STRING}
 
 
-@pytest.fixture
-def pg_dispatcher() -> DispatcherMain:
-    # We can not reuse the connection between tests
-    config = BASIC_CONFIG.copy()
-    config['brokers']['pg_notify'].pop('async_connection_factory')
-    return DispatcherMain(config)
-
-
 @pytest.fixture
 def test_settings():
     return DispatcherSettings(BASIC_CONFIG)

tests/integration/conftest.py

@@ -0,0 +1,32 @@
+import pytest
+
+from dispatcher.testing import dispatcher_service
+from dispatcher.factories import get_publisher_from_settings
+from dispatcher.config import DispatcherSettings
+
+from tests.conftest import CONNECTION_STRING
+
+
+BASIC_CONFIG = {
+    "version": 2,
+    "brokers": {
+        "pg_notify": {
+            "channels": ['test_channel', 'test_channel2', 'test_channel3'],
+            "config": {'conninfo': CONNECTION_STRING},
+            "sync_connection_factory": "dispatcher.brokers.pg_notify.connection_saver",
+            "default_publish_channel": "test_channel"
+        }
+    }
+}
+
+
+@pytest.fixture
+def pg_dispatcher():
+    with dispatcher_service(BASIC_CONFIG, pool_events=('work_cleared',)) as comms:
+        yield comms
+
+
+@pytest.fixture()
+def pg_broker():
+    settings = DispatcherSettings(BASIC_CONFIG)
+    return get_publisher_from_settings(settings=settings)

tests/integration/test_main.py

@@ -23,13 +23,10 @@ async def wait_to_receive(dispatcher, ct, timeout=5.0, interval=0.05):
         raise RuntimeError(f'Failed to receive expected {ct} messages {dispatcher.pool.received_count}')
 
 
-@pytest.mark.asyncio
-async def test_run_lambda_function(apg_dispatcher, pg_message):
-    clearing_task = asyncio.create_task(apg_dispatcher.pool.events.work_cleared.wait(), name='test_lambda_clear_wait')
-    await pg_message('lambda: "This worked!"')
-    await asyncio.wait_for(clearing_task, timeout=3)
-
-    assert apg_dispatcher.pool.finished_count == 1
+def test_run_lambda_function(pg_dispatcher, pg_broker):
+    pg_broker.publish_message(message='lambda: "This worked!"')
+    message = pg_dispatcher.q_out.get()
+    assert message == 'work_cleared'
 
 
 @pytest.mark.asyncio