Skip to content

Commit c81ba62

Browse files
committed
perf: reduce the overhead of wait_async()
Block on a plain Future, resolved by the fd callbacks or by a timer, instead of an Event awaited through wait_for(), which wraps every wait in a Task. Also look up the event loop only after the first yield: during copy() the generator usually completes without blocking at all. copy() is 10-70% faster across the platforms and Python versions we support. Query round trips gain around 10% up to Python 3.11 and little after it, where wait_for() stopped creating a Task. Diagnosis and first implementation by @DylanYoung in #1331, which polled the socket with `await sleep(0)`, burning a core while idle. wait_conn_async() gets the same structure, and now raises OperationalError instead of OSError.
1 parent d2bbfe4 commit c81ba62

2 files changed

Lines changed: 60 additions & 29 deletions

File tree

‎docs/news.rst‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,8 @@ Psycopg 3.3.6 (unreleased)
2424
- Fix dumping of nested subclasses of lists as arrays (:ticket:`#1398`).
2525
- Discard prepared statements upon :sql:`DEALLOCATE ALL` (:ticket:`#1408`).
2626
- Better guards dumping large Python `!int` to binary numeric (:ticket:`#1414`).
27+
- Improve performance of async queries by reducing the overhead of the
28+
`!wait_async()` function (:ticket:`#1331`).
2729

2830

2931
Current release

‎psycopg/psycopg/waiting.py‎

Lines changed: 58 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,7 @@
1616
import logging
1717
import selectors
1818
from time import monotonic
19-
from asyncio import Event, TimeoutError, get_running_loop, wait_for
19+
from asyncio import Future, get_running_loop
2020
from selectors import DefaultSelector
2121

2222
from . import errors as e
@@ -175,39 +175,54 @@ async def wait_async(
175175
"""
176176
if interval is None:
177177
raise ValueError("indefinite wait not supported anymore")
178-
deadline = monotonic() + timeout if timeout is not None else None
179178

180-
# Use an event to block and restart after the fd state changes.
181-
# Not sure this is the best implementation but it's a start.
182-
ev = Event()
179+
# Don't touch the event loop before we know that we will have to wait: a
180+
# generator completing without blocking (e.g. every copy() round) pays
181+
# nothing more than this.
182+
try:
183+
s = next(gen)
184+
except StopIteration as ex:
185+
rv: RV = ex.value
186+
return rv
187+
188+
deadline = monotonic() + timeout if timeout is not None else None
183189
loop = get_running_loop()
184-
ready: int
185-
s: Wait
190+
ready: int = 0
191+
fut: Future[None] | None = None
186192

193+
# Block on a future and resolve it from the fd callbacks, or from a timer
194+
# if the interval expires first. A future is cheaper than an Event plus
195+
# wait_for(), which would wrap every single wait in a Task.
187196
def wakeup(state: Ready) -> None:
188197
nonlocal ready
189198
ready |= state
190-
ev.set()
199+
if fut is not None and not fut.done():
200+
fut.set_result(None)
201+
202+
def timedout() -> None:
203+
if fut is not None and not fut.done():
204+
fut.set_result(None)
191205

192206
try:
193-
s = next(gen)
194207
while True:
195208
reader = s & WAIT_R
196209
writer = s & WAIT_W
197210
if not (reader or writer):
198211
raise e.InternalError(f"bad poll status: {s}")
199-
ev.clear()
200212
ready = 0
201213
if reader:
202214
loop.add_reader(fileno, wakeup, READY_R)
203215
if writer:
204216
loop.add_writer(fileno, wakeup, READY_W)
205217
try:
206218
t = interval if deadline is None else _wait_time(interval, deadline)
219+
fut = loop.create_future()
220+
timer = loop.call_later(t, timedout)
207221
try:
208-
await wait_for(ev.wait(), t)
209-
except TimeoutError:
210-
pass
222+
await fut
223+
finally:
224+
timer.cancel()
225+
fut = None
211226
finally:
212227
if reader:
213228
loop.remove_reader(fileno)
@@ -222,7 +237,7 @@ def wakeup(state: Ready) -> None:
222237
# Assume the connection was closed
223238
raise e.OperationalError("connection socket closed") from ex
224239
except StopIteration as ex:
225-
rv: RV = ex.value
240+
rv = ex.value
226241
return rv
227242

228243

@@ -246,39 +261,50 @@ async def wait_conn_async(
246261
"""
247262
if interval is None:
248263
raise ValueError("indefinite wait not supported anymore")
249-
deadline = monotonic() + timeout if timeout is not None else None
250264

251-
# Use an event to block and restart after the fd state changes.
252-
# Not sure this is the best implementation but it's a start.
253-
ev = Event()
265+
# Don't touch the event loop before we know that we will have to wait.
266+
# See `wait_async()` for the rationale of this function's structure.
267+
try:
268+
fileno, s = next(gen)
269+
except StopIteration as ex:
270+
rv: RV = ex.value
271+
return rv
272+
273+
deadline = monotonic() + timeout if timeout is not None else None
254274
loop = get_running_loop()
255-
ready: Ready
256-
s: Wait
275+
ready: int = 0
276+
fut: Future[None] | None = None
257277

258278
def wakeup(state: Ready) -> None:
259279
nonlocal ready
260-
ready = state
261-
ev.set()
280+
ready |= state
281+
if fut is not None and not fut.done():
282+
fut.set_result(None)
283+
284+
def timedout() -> None:
285+
if fut is not None and not fut.done():
286+
fut.set_result(None)
262287

263288
try:
264-
fileno, s = next(gen)
265289
while True:
266290
reader = s & WAIT_R
267291
writer = s & WAIT_W
268292
if not (reader or writer):
269293
raise e.InternalError(f"bad poll status: {s}")
270-
ev.clear()
271-
ready = 0 # type: ignore[assignment]
294+
ready = 0
272295
if reader:
273296
loop.add_reader(fileno, wakeup, READY_R)
274297
if writer:
275298
loop.add_writer(fileno, wakeup, READY_W)
276299
try:
277300
t = interval if deadline is None else _wait_time(interval, deadline)
301+
fut = loop.create_future()
302+
timer = loop.call_later(t, timedout)
278303
try:
279-
await wait_for(ev.wait(), t)
280-
except TimeoutError:
281-
pass
304+
await fut
305+
finally:
306+
timer.cancel()
307+
fut = None
282308
finally:
283309
if reader:
284310
loop.remove_reader(fileno)
@@ -289,8 +315,11 @@ def wakeup(state: Ready) -> None:
289315
if deadline is not None and monotonic() >= deadline:
290316
raise e._WaitTimeout("wait timeout expired")
291317

318+
except OSError as ex:
319+
# Assume the connection was closed
320+
raise e.OperationalError("connection socket closed") from ex
292321
except StopIteration as ex:
293-
rv: RV = ex.value
322+
rv = ex.value
294323
return rv
295324

296325

0 commit comments

Comments
 (0)