Skip to content

Commit 86c6d27

Browse files
[bug] Prevent CancelledError from escaping (#158)
1 parent d12efb8 commit 86c6d27

File tree

1 file changed

+16
-7
lines changed

1 file changed

+16
-7
lines changed

src/replit_river/v2/session.py

Lines changed: 16 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -287,7 +287,7 @@ def unbind_connecting_task() -> None:
287287
else:
288288
logger.debug("unbind_connecting_task failed, id did not match")
289289

290-
if not self._connecting_task:
290+
if not self._connecting_task or self._connecting_task.done():
291291
self._connecting_task = asyncio.create_task(
292292
_do_ensure_connected(
293293
transport_options=self._transport_options,
@@ -308,9 +308,16 @@ def unbind_connecting_task() -> None:
308308
)
309309
)
310310

311-
await self._connecting_task
311+
try:
312+
await self._connecting_task
313+
except asyncio.CancelledError:
314+
pass
315+
312316
if self._terminating_task:
313-
await self._terminating_task
317+
try:
318+
await self._terminating_task
319+
except asyncio.CancelledError:
320+
pass
314321

315322
def is_terminal(self) -> bool:
316323
"""
@@ -403,7 +410,10 @@ async def close(
403410
"seconds to close, leaking",
404411
)
405412
return
406-
await self._close_internal(reason)
413+
try:
414+
await self._close_internal(reason)
415+
except asyncio.CancelledError:
416+
pass
407417

408418
def _close_internal_nowait(self, reason: Exception | None = None) -> None:
409419
"""
@@ -501,10 +511,9 @@ async def do_close() -> None:
501511
# This will get us GC'd, so this should be the last thing.
502512
self._close_session_callback(self)
503513

504-
if self._terminating_task:
505-
return self._terminating_task
514+
if not self._terminating_task:
515+
self._terminating_task = asyncio.create_task(do_close())
506516

507-
self._terminating_task = asyncio.create_task(do_close())
508517
return self._terminating_task
509518

510519
def _start_buffered_message_sender(

0 commit comments

Comments
 (0)