Skip to content

Commit 69d657e

Browse files
NoahStappNoahStapp
committed
PYTHON-4864 - Create async version of SpecRunnerThread
1 parent b4e32a1 commit 69d657e

File tree

5 files changed

+177
-75
lines changed

5 files changed

+177
-75
lines changed

‎test/asynchronous/unified_format.py

Lines changed: 15 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -35,6 +35,7 @@
3535
client_knobs,
3636
unittest,
3737
)
38+
from test.asynchronous.utils_spec_runner import SpecRunnerTask
3839
from test.unified_format_shared import (
3940
KMS_TLS_OPTS,
4041
PLACEHOLDER_MAP,
@@ -58,7 +59,6 @@
5859
snake_to_camel,
5960
wait_until,
6061
)
61-
from test.utils_spec_runner import SpecRunnerThread
6262
from test.version import Version
6363
from typing import Any, Dict, List, Mapping, Optional
6464

@@ -382,8 +382,8 @@ async def drop(self: AsyncGridFSBucket, *args: Any, **kwargs: Any) -> None:
382382
return
383383
elif entity_type == "thread":
384384
name = spec["id"]
385-
thread = SpecRunnerThread(name)
386-
thread.start()
385+
thread = SpecRunnerTask(name)
386+
await thread.start()
387387
self[name] = thread
388388
return
389389

@@ -1177,16 +1177,23 @@ def primary_changed() -> bool:
11771177

11781178
wait_until(primary_changed, "change primary", timeout=timeout)
11791179

1180-
def _testOperation_runOnThread(self, spec):
1180+
async def _testOperation_runOnThread(self, spec):
11811181
"""Run the 'runOnThread' operation."""
11821182
thread = self.entity_map[spec["thread"]]
1183-
thread.schedule(lambda: self.run_entity_operation(spec["operation"]))
1183+
if _IS_SYNC:
1184+
await thread.schedule(lambda: self.run_entity_operation(spec["operation"]))
1185+
else:
1186+
1187+
async def op():
1188+
await self.run_entity_operation(spec["operation"])
1189+
1190+
await thread.schedule(op)
11841191

1185-
def _testOperation_waitForThread(self, spec):
1192+
async def _testOperation_waitForThread(self, spec):
11861193
"""Run the 'waitForThread' operation."""
11871194
thread = self.entity_map[spec["thread"]]
1188-
thread.stop()
1189-
thread.join(10)
1195+
await thread.stop()
1196+
await thread.join(10)
11901197
if thread.exc:
11911198
raise thread.exc
11921199
self.assertFalse(thread.is_alive(), "Thread {} is still running".format(spec["thread"]))

‎test/asynchronous/utils_spec_runner.py

Lines changed: 76 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -54,39 +54,82 @@
5454

5555
_IS_SYNC = False
5656

57-
58-
class SpecRunnerThread(threading.Thread):
59-
def __init__(self, name):
60-
super().__init__()
61-
self.name = name
62-
self.exc = None
63-
self.daemon = True
64-
self.cond = threading.Condition()
65-
self.ops = []
66-
self.stopped = False
67-
68-
def schedule(self, work):
69-
self.ops.append(work)
70-
with self.cond:
71-
self.cond.notify()
72-
73-
def stop(self):
74-
self.stopped = True
75-
with self.cond:
76-
self.cond.notify()
77-
78-
def run(self):
79-
while not self.stopped or self.ops:
80-
if not self.ops:
81-
with self.cond:
82-
self.cond.wait(10)
83-
if self.ops:
84-
try:
85-
work = self.ops.pop(0)
86-
work()
87-
except Exception as exc:
88-
self.exc = exc
89-
self.stop()
57+
if _IS_SYNC:
58+
59+
class SpecRunnerThread(threading.Thread):
60+
def __init__(self, name):
61+
super().__init__()
62+
self.name = name
63+
self.exc = None
64+
self.daemon = True
65+
self.cond = threading.Condition()
66+
self.ops = []
67+
self.stopped = False
68+
69+
def schedule(self, work):
70+
self.ops.append(work)
71+
with self.cond:
72+
self.cond.notify()
73+
74+
def stop(self):
75+
self.stopped = True
76+
with self.cond:
77+
self.cond.notify()
78+
79+
def run(self):
80+
while not self.stopped or self.ops:
81+
if not self.ops:
82+
with self.cond:
83+
self.cond.wait(10)
84+
if self.ops:
85+
try:
86+
work = self.ops.pop(0)
87+
work()
88+
except Exception as exc:
89+
self.exc = exc
90+
self.stop()
91+
else:
92+
93+
class SpecRunnerTask:
94+
def __init__(self, name):
95+
self.name = name
96+
self.exc = None
97+
self.cond = asyncio.Condition()
98+
self.ops = []
99+
self.stopped = False
100+
self.task = None
101+
102+
async def schedule(self, work):
103+
self.ops.append(work)
104+
async with self.cond:
105+
self.cond.notify()
106+
107+
async def stop(self):
108+
self.stopped = True
109+
async with self.cond:
110+
self.cond.notify()
111+
112+
async def start(self):
113+
self.task = asyncio.create_task(self.run(), name=self.name)
114+
115+
async def join(self, timeout: int = 0):
116+
await asyncio.wait([self.task], timeout=timeout)
117+
118+
def is_alive(self):
119+
return not self.stopped
120+
121+
async def run(self):
122+
while not self.stopped or self.ops:
123+
if not self.ops:
124+
async with self.cond:
125+
await asyncio.wait_for(self.cond.wait(), timeout=10)
126+
if self.ops:
127+
try:
128+
work = self.ops.pop(0)
129+
await work()
130+
except Exception as exc:
131+
self.exc = exc
132+
await self.stop()
90133

91134

92135
class AsyncSpecTestCreator:

‎test/unified_format.py

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1167,7 +1167,14 @@ def primary_changed() -> bool:
11671167
def _testOperation_runOnThread(self, spec):
11681168
"""Run the 'runOnThread' operation."""
11691169
thread = self.entity_map[spec["thread"]]
1170-
thread.schedule(lambda: self.run_entity_operation(spec["operation"]))
1170+
if _IS_SYNC:
1171+
thread.schedule(lambda: self.run_entity_operation(spec["operation"]))
1172+
else:
1173+
1174+
def op():
1175+
self.run_entity_operation(spec["operation"])
1176+
1177+
thread.schedule(op)
11711178

11721179
def _testOperation_waitForThread(self, spec):
11731180
"""Run the 'waitForThread' operation."""

‎test/utils_spec_runner.py

Lines changed: 76 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -54,39 +54,82 @@
5454

5555
_IS_SYNC = True
5656

57-
58-
class SpecRunnerThread(threading.Thread):
59-
def __init__(self, name):
60-
super().__init__()
61-
self.name = name
62-
self.exc = None
63-
self.daemon = True
64-
self.cond = threading.Condition()
65-
self.ops = []
66-
self.stopped = False
67-
68-
def schedule(self, work):
69-
self.ops.append(work)
70-
with self.cond:
71-
self.cond.notify()
72-
73-
def stop(self):
74-
self.stopped = True
75-
with self.cond:
76-
self.cond.notify()
77-
78-
def run(self):
79-
while not self.stopped or self.ops:
80-
if not self.ops:
81-
with self.cond:
82-
self.cond.wait(10)
83-
if self.ops:
84-
try:
85-
work = self.ops.pop(0)
86-
work()
87-
except Exception as exc:
88-
self.exc = exc
89-
self.stop()
57+
if _IS_SYNC:
58+
59+
class SpecRunnerThread(threading.Thread):
60+
def __init__(self, name):
61+
super().__init__()
62+
self.name = name
63+
self.exc = None
64+
self.daemon = True
65+
self.cond = threading.Condition()
66+
self.ops = []
67+
self.stopped = False
68+
69+
def schedule(self, work):
70+
self.ops.append(work)
71+
with self.cond:
72+
self.cond.notify()
73+
74+
def stop(self):
75+
self.stopped = True
76+
with self.cond:
77+
self.cond.notify()
78+
79+
def run(self):
80+
while not self.stopped or self.ops:
81+
if not self.ops:
82+
with self.cond:
83+
self.cond.wait(10)
84+
if self.ops:
85+
try:
86+
work = self.ops.pop(0)
87+
work()
88+
except Exception as exc:
89+
self.exc = exc
90+
self.stop()
91+
else:
92+
93+
class SpecRunnerThread:
94+
def __init__(self, name):
95+
self.name = name
96+
self.exc = None
97+
self.cond = asyncio.Condition()
98+
self.ops = []
99+
self.stopped = False
100+
self.task = None
101+
102+
def schedule(self, work):
103+
self.ops.append(work)
104+
with self.cond:
105+
self.cond.notify()
106+
107+
def stop(self):
108+
self.stopped = True
109+
with self.cond:
110+
self.cond.notify()
111+
112+
def start(self):
113+
self.task = asyncio.create_task(self.run(), name=self.name)
114+
115+
def join(self, timeout: int = 0):
116+
asyncio.wait([self.task], timeout=timeout)
117+
118+
def is_alive(self):
119+
return not self.stopped
120+
121+
def run(self):
122+
while not self.stopped or self.ops:
123+
if not self.ops:
124+
with self.cond:
125+
asyncio.wait_for(self.cond.wait(), timeout=10)
126+
if self.ops:
127+
try:
128+
work = self.ops.pop(0)
129+
work()
130+
except Exception as exc:
131+
self.exc = exc
132+
self.stop()
90133

91134

92135
class SpecTestCreator:

‎tools/synchro.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -119,6 +119,8 @@
119119
"_async_create_lock": "_create_lock",
120120
"_async_create_condition": "_create_condition",
121121
"_async_cond_wait": "_cond_wait",
122+
"AsyncDummyMonitor": "DummyMonitor",
123+
"SpecRunnerTask": "SpecRunnerThread",
122124
}
123125

124126
docstring_replacements: dict[tuple[str, str], str] = {

0 commit comments

Comments
 (0)