Thread safe In Memory Connector

procrastinate/testing.py

@@ -1,7 +1,9 @@
 from __future__ import annotations
 
+import asyncio
 import datetime
 import json
+import threading
 from collections import Counter
 from collections.abc import Iterable
 from itertools import count
@@ -83,6 +85,12 @@ def open(self, pool: connector.Pool | None = None) -> None:
         self.states.append("open")
 
     async def open_async(self, pool: connector.Pool | None = None) -> None:
+        """
+        Save the current event loop and its thread id so that later notifications
+        can be scheduled on this loop.
+        """
+        self._loop = asyncio.get_running_loop()
+        self._loop_thread_id = threading.get_ident()
         self.states.append("open_async")
 
     def close(self) -> None:
@@ -182,20 +190,31 @@ def finished_jobs(self) -> list[JobRow]:
             if job["status"] in {"failed", "succeeded"}
         ]
 
-    async def _notify(self, queue_name: str, notification: jobs.Notification):
+    async def _notify(self, queue_name: str, notification: jobs.Notification) -> None:
+        """
+        Instead of directly awaiting on_notification, we check the current thread.
+        If we are not on the same thread as the one where the loop was saved,
+        we schedule the notification on the correct loop.
+        """
         if not self.on_notification:
             return
 
         destination_channels = {
             "procrastinate_any_queue_v1",
             f"procrastinate_queue_v1#{queue_name}",
         }
-
         for channel in set(self.notify_channels).intersection(destination_channels):
-            await self.on_notification(
-                channel=channel,
-                payload=json.dumps(notification),
+            coro = self.on_notification(
+                channel=channel, payload=json.dumps(notification)
             )
+            if threading.get_ident() == self._loop_thread_id:
+                # Already on the right thread: just await.
+                await coro
+            else:
+                # Not on the correct thread: schedule the coroutine on the saved loop.
+                future = asyncio.run_coroutine_threadsafe(coro, self._loop)
+                # Wrap the concurrent.futures.Future so we can await it.
+                await asyncio.wrap_future(future)
 
     async def fetch_job_one(self, queues: Iterable[str] | None) -> dict:
         # Creating a copy of the iterable so that we can modify it while we iterate