Repository navigation
Expand file tree
/
Copy pathtest_concurrency.py
More file actions
105 lines (91 loc) · 3.88 KB
/
Copy pathtest_concurrency.py
File metadata and controls
105 lines (91 loc) · 3.88 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
"""
Copyright 2025 Amazon.com, Inc. or its affiliates. All Rights Reserved.
"""
import multiprocessing
import os
import sys
import threading
import types
import unittest
from unittest.mock import patch, MagicMock
from awslambdaric.lambda_multi_concurrent_utils import MultiConcurrentRunner
class LambdaRuntimeConcurrencyTest(unittest.TestCase):
def setUp(self):
# common args
self.handler = "h.fn"
self.addr = "addr"
self.use_thread = False
self.socket = "/tmp/sock"
def test_success_and_failure_isolation(self):
multiprocessing.set_start_method("fork", force=True)
success_counter = multiprocessing.Value("i", 0)
fail_counter = multiprocessing.Value("i", 0)
process_index = multiprocessing.Value("i", 0)
def fake_bootstrap_run(handler, lambda_runtime_client):
with process_index.get_lock():
idx = process_index.value
process_index.value += 1
if idx % 2 == 0:
for _ in range(3):
with success_counter.get_lock():
success_counter.value += 1
else:
with fail_counter.get_lock():
fail_counter.value += 1
raise RuntimeError("Simulated failure")
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,
):
# spawn 4 multi-concurrent processes
MultiConcurrentRunner.run_concurrent(
self.handler, self.addr, self.use_thread, self.socket, max_concurrency=4
)
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()