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
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
|
# --------------------------------------------------------------------------------------------
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License. See License.txt in the project root for license information.
# --------------------------------------------------------------------------------------------
import asyncio
import threading
from concurrent.futures import ThreadPoolExecutor, as_completed
import time
from devtools_testutils.perfstress_tests import EventPerfTest
class MockEventProcessor:
def __init__(self, partitions, process_event, process_error, error_after=None, max_events_per_second=None):
self.partitions = partitions
self.error_after = error_after
self.max_events_per_second = max_events_per_second
self.shutdown = False
self._error_raised = False
self._error_lock = threading.Lock()
self._event_args = [{"partition": i, "data": "hello"} for i in range(self.partitions)]
self._events_raised = [0] * self.partitions
self._starttime = None
self._process_event = process_event
self._process_error = process_error
self._executor = ThreadPoolExecutor(max_workers=self.partitions)
def _process_error_after(self, partition):
with self._error_lock:
if not self._error_raised:
error = Exception("Raising error after {} seconds.".format(self.error_after))
self._process_error(partition=partition, error=error)
self._error_raised = True
def start_processing(self):
self._starttime = time.time()
with self._executor as ex:
futures = [ex.submit(self._process, p) for p in range(self.partitions)]
for future in as_completed(futures):
future.result()
def _process(self, partition):
event_args = self._event_args[partition]
if self.max_events_per_second:
max_events_per_partition = float(self.max_events_per_second) / self.partitions
while not self.shutdown:
elapsed = time.time() - self._starttime
if self.error_after and not self._error_raised and elapsed > self.error_after:
self._process_error_after(partition)
else:
target_events = elapsed * max_events_per_partition
if self._events_raised[partition] < target_events:
self._process_event(**event_args)
self._events_raised[partition] += 1
else:
time.sleep(1 / max_events_per_partition)
else:
while not self.shutdown:
elapsed = time.time() - self._starttime
if self.error_after and not self._error_raised and elapsed > self.error_after:
self._process_error_after(partition)
else:
self._process_event(**event_args)
class AsyncMockEventProcessor:
def __init__(self, partitions, process_event, process_error, error_after=None, max_events_per_second=None):
self.partitions = partitions
self.error_after = error_after
self.max_events_per_second = max_events_per_second
self.shutdown = False
self._error_raised = False
self._error_lock = asyncio.Lock()
self._starttime = None
self._process_event = process_event
self._process_error = process_error
self._event_args = [{"partition": i, "data": "hello"} for i in range(self.partitions)]
self._events_raised = [0] * self.partitions
async def _process_error_after(self, partition):
async with self._error_lock:
if not self._error_raised:
error = Exception("Raising error after {} seconds.".format(self.error_after))
await self._process_error(partition=partition, error=error)
self._error_raised = True
async def start_processing(self):
self._starttime = time.time()
await asyncio.gather(*[self._process(i) for i in range(self.partitions)])
async def _process(self, partition):
event_args = self._event_args[partition]
sleep_count = 0
if self.max_events_per_second:
max_events_per_partition = float(self.max_events_per_second) / self.partitions
while not self.shutdown:
sleep_count += 1
elapsed = time.time() - self._starttime
if self.error_after and not self._error_raised and elapsed > self.error_after:
await self._process_error_after(partition)
else:
target_events = elapsed * max_events_per_partition
if self._events_raised[partition] < target_events:
await self._process_event(**event_args)
self._events_raised[partition] += 1
else:
await asyncio.sleep(1 / max_events_per_partition)
await asyncio.sleep(0)
else:
while not self.shutdown:
sleep_count += 1
elapsed = time.time() - self._starttime
if self.error_after and not self._error_raised and elapsed > self.error_after:
await self._process_error_after(partition)
else:
await self._process_event(**event_args)
if sleep_count % 500 == 0:
await asyncio.sleep(0)
class SampleEventTest(EventPerfTest):
def __init__(self, arguments):
super().__init__(arguments)
# Setup service clients
self.event_processor = MockEventProcessor(
self.args.partitions,
self.process_event_sync,
self.process_error_sync,
error_after=self.args.error_after_seconds,
max_events_per_second=self.args.max_events_per_second,
)
self.async_event_processor = AsyncMockEventProcessor(
self.args.partitions,
self.process_event_async,
self.process_error_async,
error_after=self.args.error_after_seconds,
max_events_per_second=self.args.max_events_per_second,
)
def process_event_sync(self, **kwargs):
self.event_raised_sync()
def process_error_sync(self, error=None, **kwargs):
self.error_raised_sync(error)
async def process_event_async(self, **kwargs):
await self.event_raised_async()
async def process_error_async(self, error=None, **kwargs):
await self.error_raised_async(error)
def start_events_sync(self) -> None:
"""
Start the process for receiving events.
"""
self.event_processor.start_processing()
def stop_events_sync(self) -> None:
"""
Stop the process for receiving events.
"""
self.event_processor.shutdown = True
async def start_events_async(self) -> None:
"""
Start the process for receiving events.
"""
await self.async_event_processor.start_processing()
async def stop_events_async(self) -> None:
"""
Stop the process for receiving events.
"""
self.async_event_processor.shutdown = True
@staticmethod
def add_arguments(parser):
super(SampleEventTest, SampleEventTest).add_arguments(parser)
parser.add_argument(
"--error-after-seconds", nargs="?", type=int, help="Raise error after this number of seconds."
)
parser.add_argument(
"--max-events-per-second", nargs="?", type=int, help="Maximum events per second across all partitions."
)
parser.add_argument("--partitions", nargs="?", type=int, help="Number of partitions. Default is 8.", default=8)
|
