1616import logging
1717import selectors
1818from time import monotonic
19- from asyncio import Event , TimeoutError , get_running_loop , wait_for
19+ from asyncio import Future , get_running_loop
2020from selectors import DefaultSelector
2121
2222from . 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