Allow using forkserver

Author: AlanCoding

dispatcher/pool.py

@@ -4,7 +4,7 @@
 from asyncio import Task
 from typing import Iterator, Optional
 
-from dispatcher.process import ProcessManager, ProcessProxy
+from dispatcher.process import ForkServerManager, ProcessProxy
 from dispatcher.utils import DuplicateBehavior, MessageAction
 
 logger = logging.getLogger(__name__)
@@ -98,7 +98,7 @@ def __init__(self, num_workers: int, fd_lock: Optional[asyncio.Lock] = None):
         self.num_workers = num_workers
         self.workers: dict[int, PoolWorker] = {}
         self.next_worker_id = 0
-        self.process_manager = ProcessManager()
+        self.process_manager = ForkServerManager()
         self.queued_messages: list[dict] = []  # TODO: use deque, invent new kinds of logging anxiety
         self.read_results_task: Optional[Task] = None
         self.start_worker_task: Optional[Task] = None

dispatcher/process.py

@@ -1,14 +1,15 @@
 import asyncio
 import multiprocessing
-from typing import Callable, Iterable, Optional, Union
+from multiprocessing.context import BaseContext
+from typing import Callable, Iterable, Optional, Union, Sized
 
 from dispatcher.worker.task import work_loop
 
 
 class ProcessProxy:
-    def __init__(self, args: Iterable, finished_queue: multiprocessing.Queue, target: Callable = work_loop) -> None:
-        self.message_queue: multiprocessing.Queue = multiprocessing.Queue()
-        self._process = multiprocessing.Process(target=target, args=tuple(args) + (self.message_queue, finished_queue))
+    def __init__(self, args: Iterable, finished_queue: multiprocessing.Queue, target: Callable = work_loop, ctx: BaseContext = multiprocessing) -> None:
+        self.message_queue: multiprocessing.Queue = ctx.Queue()
+        self._process = ctx.Process(target=target, args=tuple(args) + (self.message_queue, finished_queue))
 
     def start(self) -> None:
         self._process.start()
@@ -37,8 +38,11 @@ def terminate(self) -> None:
 
 
 class ProcessManager:
+    mp_context = 'fork'
+
     def __init__(self) -> None:
-        self.finished_queue: multiprocessing.Queue = multiprocessing.Queue()
+        self.ctx = multiprocessing.get_context(self.mp_context)
+        self.finished_queue: multiprocessing.Queue = self.ctx.Queue()
         self._loop = None
 
     def get_event_loop(self):
@@ -47,8 +51,16 @@ def get_event_loop(self):
         return self._loop
 
     def create_process(self, args: Iterable[int | str], **kwargs) -> ProcessProxy:
-        return ProcessProxy(args, self.finished_queue, **kwargs)
+        return ProcessProxy(args, self.finished_queue, ctx=self.ctx, **kwargs)
 
     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
+
+
+class ForkServerManager(ProcessManager):
+    mp_context = 'forkserver'
+
+    def __init__(self, preload_modules: Sized = ()):
+        super().__init__()
+        self.ctx.set_forkserver_preload(preload_modules)

tests/unit/service/test_process.py

@@ -1,6 +1,8 @@
 from multiprocessing import Queue
 
-from dispatcher.process import ProcessManager, ProcessProxy
+import pytest
+
+from dispatcher.process import ProcessManager, ForkServerManager, ProcessProxy
 
 
 def test_pass_messages_to_worker():
@@ -17,13 +19,19 @@ def work_loop(a, b, c, in_q, out_q):
     assert msg == 'done 1 2 3 start'
 
 
-def test_pass_messages_via_process_manager():
-    def work_loop(var, in_q, out_q):
-        has_read = in_q.get()
-        out_q.put(f'done {var} {has_read}')
+def work_loop2(var, in_q, out_q):
+    """
+    Due to the mechanics of forkserver, this can not be defined in local variables,
+    it has to be importable, but this _is_ importable from the test module.
+    """
+    has_read = in_q.get()
+    out_q.put(f'done {var} {has_read}')
+
 
-    process_manager = ProcessManager()
-    process = process_manager.create_process(('value',), target=work_loop)
+@pytest.mark.parametrize('manager_cls', [ProcessManager, ForkServerManager])
+def test_pass_messages_via_process_manager(manager_cls):
+    process_manager = manager_cls()
+    process = process_manager.create_process(('value',), target=work_loop2)
     process.start()
 
     process.message_queue.put('msg1')