Merge pull request #233 from aws/anzhey/pin-fork

fix: pin multi-concurrent worker start method to fork

Author: anzheyazzz

awslambdaric/lambda_multi_concurrent_utils.py

@@ -67,9 +67,10 @@ def run_concurrent(
     ):
         cls._emit_worker_pool_event(max_concurrency)
 
+        ctx = multiprocessing.get_context("fork")
         processes = []
         for _ in range(max_concurrency):
-            p = multiprocessing.Process(
+            p = ctx.Process(
                 target=cls.run_single,
                 args=(handler, api_addr, use_thread, socket_path),
             )

tests/test_concurrency.py

@@ -3,6 +3,10 @@
 """
 
 import multiprocessing
+import os
+import sys
+import threading
+import types
 import unittest
 from unittest.mock import patch, MagicMock
 
@@ -52,6 +56,50 @@ def fake_bootstrap_run(handler, lambda_runtime_client):
         self.assertEqual(success_counter.value, 6)
         self.assertEqual(fail_counter.value, 2)
 
+    def test_workers_are_forked_children_regardless_of_default_start_method(self):
+        previous = multiprocessing.get_start_method(allow_none=True)
+        multiprocessing.set_start_method("forkserver", force=True)
+        try:
+            parent_pid = os.getpid()
+            parent_threads_at_fork = threading.active_count()
+            marker = "_fork_pin_marker_module"
+            sys.modules[marker] = types.ModuleType(marker)
+            n = 3
+            ppids = multiprocessing.Array("i", n)
+            inherited = multiprocessing.Array("i", n)
+            slot = multiprocessing.Value("i", 0)
+
+            def fake_bootstrap_run(handler, lambda_runtime_client):
+                with slot.get_lock():
+                    i = slot.value
+                    slot.value += 1
+                ppids[i] = os.getppid()
+                inherited[i] = int(marker in sys.modules)
+
+            with patch(
+                "awslambdaric.lambda_multi_concurrent_utils.MultiConcurrentRunner._redirect_output"
+            ), patch(
+                "awslambdaric.lambda_multi_concurrent_utils.MultiConcurrentRunner._emit_worker_pool_event"
+            ), patch(
+                "awslambdaric.lambda_multi_concurrent_utils.bootstrap.run",
+                side_effect=fake_bootstrap_run,
+            ):
+                MultiConcurrentRunner.run_concurrent(
+                    self.handler,
+                    self.addr,
+                    self.use_thread,
+                    self.socket,
+                    max_concurrency=n,
+                )
+
+            self.assertEqual(multiprocessing.get_start_method(), "forkserver")
+            self.assertEqual(parent_threads_at_fork, 1)
+            self.assertEqual(list(ppids), [parent_pid] * n)
+            self.assertEqual(list(inherited), [1] * n)
+        finally:
+            sys.modules.pop(marker, None)
+            multiprocessing.set_start_method(previous, force=True)
+
 
 if __name__ == "__main__":
     unittest.main()

tests/test_multi_concurrent_runner.py

@@ -59,32 +59,48 @@ def test_run_single_creates_client_and_calls_bootstrap(
         mock_bootstrap.run.assert_called_once_with("h.fn", mock_client)
 
     @patch.object(MultiConcurrentRunner, "_emit_worker_pool_event")
-    @patch("multiprocessing.Process")
-    def test_run_concurrent_spawns_and_joins(self, mock_process, mock_emit):
+    @patch("multiprocessing.get_context")
+    def test_run_concurrent_spawns_and_joins(self, mock_get_context, mock_emit):
         fake_proc = MagicMock()
-        mock_process.return_value = fake_proc
+        mock_ctx = MagicMock()
+        mock_ctx.Process.return_value = fake_proc
+        mock_get_context.return_value = mock_ctx
 
         MultiConcurrentRunner.run_concurrent(
             "h", "a", False, "/sock", max_concurrency=3
         )
 
-        self.assertEqual(mock_process.call_count, 3)
+        self.assertEqual(mock_ctx.Process.call_count, 3)
         self.assertEqual(fake_proc.start.call_count, 3)
         self.assertEqual(fake_proc.join.call_count, 3)
 
-        for call_args in mock_process.call_args_list:
+        for call_args in mock_ctx.Process.call_args_list:
             target = call_args.kwargs.get("target") or call_args[1].get("target")
             args = call_args.kwargs.get("args") or call_args[1].get("args")
             self.assertEqual(target, MultiConcurrentRunner.run_single)
             self.assertEqual(args, ("h", "a", False, "/sock"))
 
-    @patch("multiprocessing.Process")
+    @patch.object(MultiConcurrentRunner, "_emit_worker_pool_event")
+    @patch("multiprocessing.get_context")
+    def test_run_concurrent_pins_fork_start_method(self, mock_get_context, mock_emit):
+        mock_ctx = MagicMock()
+        mock_get_context.return_value = mock_ctx
+
+        MultiConcurrentRunner.run_concurrent(
+            "h", "a", False, "/sock", max_concurrency=2
+        )
+
+        mock_get_context.assert_called_once_with("fork")
+        self.assertEqual(mock_ctx.Process.call_count, 2)
+
+    @patch("multiprocessing.get_context")
     def test_run_concurrent_emits_worker_pool_event_once_before_spawning(
-        self, mock_process
+        self, mock_get_context
     ):
-        mock_process.return_value = MagicMock()
+        mock_ctx = MagicMock()
+        mock_get_context.return_value = mock_ctx
         order_tracker = MagicMock()
-        order_tracker.attach_mock(mock_process, "process")
+        order_tracker.attach_mock(mock_ctx.Process, "process")
 
         with patch.object(
             MultiConcurrentRunner, "_emit_worker_pool_event"