-
Notifications
You must be signed in to change notification settings - Fork 462
fix: clear 0.4.3 blockers with cancel-safe teardown #1172
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Changes from all commits
f438c4d
63f4355
30a23fb
5d177e3
366941d
3b06765
f1069e8
55255ab
2575ab4
97b7688
157e125
da95efb
de95e72
5de7a2d
d97979c
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -255,6 +255,15 @@ async def _best_effort_close(ws: ClientConnection) -> None: | |
| pass # Best effort | ||
|
|
||
|
|
||
| async def _best_effort_graceful_close(ws: ClientConnection) -> None: | ||
| """Notify the server, then close a socket without propagating failures.""" | ||
| try: | ||
| await ws.send(json.dumps({"type": "close"})) | ||
| except (Exception, asyncio.CancelledError): | ||
| pass # Best effort | ||
| await _best_effort_close(ws) | ||
|
|
||
|
|
||
| class EnvClient(ABC, Generic[ActT, ObsT, StateT]): | ||
| """ | ||
| Async environment client for persistent sessions. | ||
|
|
@@ -538,6 +547,13 @@ async def _connect_async(self) -> "EnvClient": | |
| self._ws = None | ||
| self._ws_loop = None | ||
|
|
||
| # A timed-out request drops its socket immediately but closes it in the | ||
| # background so the timeout itself remains prompt. Wait for that close | ||
| # before opening a replacement: the old server-side session continues | ||
| # occupying a capacity slot until the close handshake finishes, and | ||
| # many environments allow only one session. | ||
| await self._drain_pending_close_tasks() | ||
|
Contributor
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Key correctness fix for single-session environments. A timed-out |
||
|
|
||
| try: | ||
| self._start_provider_if_needed() | ||
| except Exception: | ||
|
|
@@ -573,24 +589,51 @@ async def _connect_async(self) -> "EnvClient": | |
| def disconnect(self) -> Any: | ||
| return self._dispatch(self._disconnect_async) | ||
|
|
||
| def _schedule_socket_close( | ||
| self, ws: ClientConnection, *, notify_server: bool = False | ||
| ) -> asyncio.Task[None]: | ||
| """Schedule and track a socket close until its handshake finishes.""" | ||
| close_coro = ( | ||
| _best_effort_graceful_close(ws) if notify_server else _best_effort_close(ws) | ||
| ) | ||
| close_task = asyncio.create_task(close_coro) | ||
| self._pending_close_tasks.add(close_task) | ||
| close_task.add_done_callback(self._pending_close_tasks.discard) | ||
| return close_task | ||
|
|
||
| async def _disconnect_async(self) -> None: | ||
| """Close the WebSocket connection.""" | ||
| if self._ws is not None: | ||
| ws = self._ws | ||
| ws_loop = self._ws_loop | ||
| same_loop = ws_loop is asyncio.get_running_loop() | ||
| try: | ||
| if same_loop: | ||
| await ws.send(json.dumps({"type": "close"})) | ||
| except Exception: | ||
| pass # Best effort | ||
| try: | ||
| if same_loop: | ||
| await ws.close() | ||
| except Exception: | ||
| pass | ||
| # Detach first so cancellation during the close handshake cannot | ||
| # leave a stale socket cached for a later operation. | ||
| self._ws = None | ||
| self._ws_loop = None | ||
| same_loop = ws_loop is asyncio.get_running_loop() | ||
| if same_loop: | ||
| # Keep ownership of the handshake if this caller is cancelled. | ||
| # A later close/reconnect drains the task before tearing down | ||
| # the loop or opening a replacement server session. | ||
| close_task = self._schedule_socket_close(ws, notify_server=True) | ||
| await asyncio.shield(close_task) | ||
|
|
||
| async def _drain_pending_close_tasks(self) -> None: | ||
| """Wait for background socket closes owned by the current event loop. | ||
|
|
||
| Shielding keeps cancellation of the caller from cancelling the close | ||
| tasks themselves. This matters both before reconnecting, when the old | ||
| server session must release its capacity slot, and during explicit | ||
| client shutdown. | ||
| """ | ||
| loop = asyncio.get_running_loop() | ||
| tasks = [ | ||
| task | ||
| for task in tuple(self._pending_close_tasks) | ||
| if not task.done() and task.get_loop() is loop | ||
| ] | ||
| if tasks: | ||
| await asyncio.shield(asyncio.gather(*tasks, return_exceptions=True)) | ||
|
|
||
| async def _ensure_connected(self) -> None: | ||
| """Ensure WebSocket connection is established on the current loop. | ||
|
|
@@ -641,9 +684,7 @@ async def _receive(self) -> Dict[str, Any]: | |
| # would actually block for up to 10s before its deadline was | ||
| # honored. Scheduling it lets the exception propagate | ||
| # immediately while the close still happens in the background. | ||
| close_task = asyncio.ensure_future(_best_effort_close(ws)) | ||
| self._pending_close_tasks.add(close_task) | ||
| close_task.add_done_callback(self._pending_close_tasks.discard) | ||
| self._schedule_socket_close(ws) | ||
| raise | ||
| return json.loads(raw) | ||
|
|
||
|
|
@@ -950,29 +991,73 @@ async def _state_async(self) -> StateT: | |
| def close(self) -> Any: | ||
| return self._dispatch(self._close_async) | ||
|
|
||
| async def _close_child_clients(self) -> asyncio.CancelledError | None: | ||
| """Close every captured child while deferring caller cancellation.""" | ||
| children = list(self._child_clients) | ||
| self._child_clients.clear() | ||
| close_tasks = [] | ||
| deferred_cancellation: asyncio.CancelledError | None = None | ||
|
|
||
| for child in children: | ||
| try: | ||
| close_tasks.append(asyncio.ensure_future(child.close())) | ||
| except asyncio.CancelledError as exc: | ||
| if deferred_cancellation is None: | ||
| deferred_cancellation = exc | ||
| except Exception: | ||
| pass # Best effort | ||
|
|
||
| if close_tasks: | ||
| close_group = asyncio.gather(*close_tasks, return_exceptions=True) | ||
| while not close_group.done(): | ||
| try: | ||
| await asyncio.shield(close_group) | ||
| except asyncio.CancelledError as exc: | ||
| if deferred_cancellation is None: | ||
| deferred_cancellation = exc | ||
|
|
||
| for result in close_group.result(): | ||
| if ( | ||
| isinstance(result, asyncio.CancelledError) | ||
| and deferred_cancellation is None | ||
| ): | ||
| deferred_cancellation = result | ||
|
|
||
| return deferred_cancellation | ||
|
|
||
| async def _close_async(self) -> None: | ||
| """ | ||
| Close the WebSocket connection and clean up resources. | ||
|
|
||
| If this client was created via from_docker_image() or from_env(), | ||
| this will also stop and remove the associated container/process. | ||
| """ | ||
| for child in list(self._child_clients): | ||
| with suppress(Exception): | ||
| await child.close() | ||
| self._child_clients.clear() | ||
|
|
||
| deferred_cancellation: asyncio.CancelledError | None = None | ||
| try: | ||
| # Wait out any backgrounded closes from a dropped socket (see | ||
| # `_receive()` / `_best_effort_close`) so a real close() call still | ||
| # sees the handshake through. SyncEnvClient.close() waits for | ||
| # `_close_async()` before stopping its loop, so the relevant risk | ||
| # is async-context cancellation of close itself — not `_stop_loop()`. | ||
| # Keep this gather inside the provider-teardown try/finally so a | ||
| # cancelled close cannot skip container/process cleanup. | ||
| if self._pending_close_tasks: | ||
| await asyncio.gather(*self._pending_close_tasks, return_exceptions=True) | ||
| await self._disconnect_async() | ||
| try: | ||
| try: | ||
| child_cancellation = await self._close_child_clients() | ||
| except asyncio.CancelledError as exc: | ||
| child_cancellation = exc | ||
| if child_cancellation is not None: | ||
| deferred_cancellation = child_cancellation | ||
|
|
||
| # A real close waits out backgrounded closes, while shielding | ||
| # their socket handshakes from cancellation. | ||
| try: | ||
| await self._drain_pending_close_tasks() | ||
| except asyncio.CancelledError as exc: | ||
| if deferred_cancellation is None: | ||
| deferred_cancellation = exc | ||
| finally: | ||
| # Run even when child or pending-close cleanup is cancelled. A | ||
| # client may already have reconnected, and that current socket | ||
| # must not remain cached or open during teardown. | ||
| try: | ||
| await self._disconnect_async() | ||
| except asyncio.CancelledError as exc: | ||
| if deferred_cancellation is None: | ||
| deferred_cancellation = exc | ||
| finally: | ||
| try: | ||
| if self._provider is not None: | ||
|
|
@@ -986,6 +1071,9 @@ async def _close_async(self) -> None: | |
| self._base_url = None | ||
| self._ws_url = None | ||
|
|
||
| if deferred_cancellation is not None: | ||
| raise deferred_cancellation | ||
|
Contributor
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. The deferred-cancellation teardown is the crux of this PR and looks correct: child-close → pending-close drain → Optional (non-blocking): |
||
|
|
||
| def _stop_provider_best_effort(self) -> None: | ||
| """Stop the underlying provider directly, ignoring any errors. | ||
|
|
||
|
|
||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -53,7 +53,12 @@ def relative_path(value: str) -> str: | |
| return value | ||
| if ( | ||
| not value | ||
| or "\x00" in value | ||
| or any( | ||
| ord(character) < 0x20 | ||
| or 0x7F <= ord(character) <= 0x9F | ||
| or ord(character) in (0x2028, 0x2029) | ||
| for character in value | ||
| ) | ||
| or "\\" in value | ||
| or PurePosixPath(value).is_absolute() | ||
| or any(part in ("", ".", "..") for part in value.split("/")) | ||
|
|
@@ -64,8 +69,12 @@ def relative_path(value: str) -> str: | |
|
|
||
| # Positive components exclude "." and ".." without lookaround, which some | ||
| # JSON Schema regex engines do not support. | ||
| _CONTROL_CHARACTER_PATTERN = r"[\x00-\x1f\x7f-\x9f\u2028\u2029]" | ||
| _PATH_COMPONENT_PATTERN = ( | ||
| r"(?:[^./\\\x00]|\.[^./\\\x00]|\.\.[^./\\\x00]|\.\.\.)[^/\\\x00]*" | ||
| r"(?:[^./\\\x00-\x1f\x7f-\x9f\u2028\u2029]" | ||
| r"|\.[^./\\\x00-\x1f\x7f-\x9f\u2028\u2029]" | ||
| r"|\.\.[^./\\\x00-\x1f\x7f-\x9f\u2028\u2029]" | ||
| r"|\.\.\.)[^/\\\x00-\x1f\x7f-\x9f\u2028\u2029]*" | ||
| ) | ||
| RelativePath = Annotated[ | ||
| NonEmpty, | ||
|
|
@@ -79,6 +88,7 @@ def relative_path(value: str) -> str: | |
| rf"(?:/{_PATH_COMPONENT_PATTERN})*)$" | ||
| ), | ||
| }, | ||
| {"not": {"pattern": _CONTROL_CHARACTER_PATTERN}}, | ||
|
Contributor
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. This separate unanchored negative clause (not just widening the char classes inside the main anchored pattern) is the load-bearing part. Model ( |
||
| ] | ||
| } | ||
| ), | ||
|
|
||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Implements the XDG-spec hardening flagged during the #1167 review: a relative
XDG_CACHE_HOMEis now ignored (rather than resolved CWD-relative, which an untrusted checkout could exploit to plant a cache), and the~/.cachefallback fails closed withRuntimeErrorwhenPath.home()isn't absolute. Good defense-in-depth on top of #1167's trusted-open/atomic-save. Covered bytest_relative_xdg_cache_home_cannot_redirect_into_working_treeandtest_relative_home_cannot_restore_a_relative_cache_path.