diff --git a/awslambdaric/lambda_multi_concurrent_utils.py b/awslambdaric/lambda_multi_concurrent_utils.py index 053e624..c18ad9a 100644 --- a/awslambdaric/lambda_multi_concurrent_utils.py +++ b/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), ) diff --git a/tests/test_concurrency.py b/tests/test_concurrency.py index 61ceeca..781d523 100644 --- a/tests/test_concurrency.py +++ b/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() diff --git a/tests/test_multi_concurrent_runner.py b/tests/test_multi_concurrent_runner.py index 5cc1371..d4a7cbb 100644 --- a/tests/test_multi_concurrent_runner.py +++ b/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"