From 565062068b20bc99edc34ccaa3e264ca4e16dbee Mon Sep 17 00:00:00 2001 From: Xero <258684383+xero-aiescape@users.noreply.github.com> Date: Fri, 10 Jul 2026 10:19:32 -0500 Subject: [PATCH] fix(socket-mode): settle aiohttp tasks on close Stop connect attempts once client shutdown begins, discard sessions that complete after close, and await cancellation of SDK-owned tasks before closing the aiohttp session.\n\nAdds lifecycle regressions for closed clients, close races, current-task ownership, backoff, and task/session cleanup.\n\nRefs #1913 --- slack_sdk/socket_mode/aiohttp/__init__.py | 33 ++- .../socket_mode/test_aiohttp.py | 230 ++++++++++++++++++ 2 files changed, 255 insertions(+), 8 deletions(-) diff --git a/slack_sdk/socket_mode/aiohttp/__init__.py b/slack_sdk/socket_mode/aiohttp/__init__.py index 761683020..17899c14a 100644 --- a/slack_sdk/socket_mode/aiohttp/__init__.py +++ b/slack_sdk/socket_mode/aiohttp/__init__.py @@ -349,7 +349,7 @@ async def connect(self): # a new monitor and a new message receiver are also created. # If a new session is created but we failed to create the new # monitor or the new message, we should try it. - while True: + while not self.closed: try: old_session: Optional[ClientWebSocketResponse] = ( None if self.current_session is None else self.current_session @@ -369,18 +369,28 @@ async def connect(self): except Exception as e: self.logger.exception(f"Failed to close the old session : {e}") + if self.closed: + break + if self.wss_uri is None: # If the underlying WSS URL does not exist, # acquiring a new active WSS URL from the server-side first self.wss_uri = await self.issue_new_wss_url() - self.current_session = await self.aiohttp_client_session.ws_connect( + if self.closed: + break + + new_session = await self.aiohttp_client_session.ws_connect( self.wss_uri, autoping=False, heartbeat=self.ping_interval, proxy=self.proxy, ssl=self.web_client.ssl if self.web_client.ssl is not None else True, ) + if self.closed: + await new_session.close() + break + self.current_session = new_session session_id: str = await self.session_id() self.auto_reconnect_enabled = self.default_auto_reconnect_enabled self.stale = False @@ -405,6 +415,8 @@ async def connect(self): self.logger.debug(f"A new receive_messages() executor has been recreated for {session_id}") break except Exception as e: + if self.closed: + break self.logger.exception(f"Failed to connect (error: {e}); Retrying...") await asyncio.sleep(self.ping_interval) @@ -447,12 +459,17 @@ async def close(self): self.closed = True self.auto_reconnect_enabled = False await self.disconnect() - if self.message_processor is not None: - self.message_processor.cancel() - if self.current_session_monitor is not None: - self.current_session_monitor.cancel() - if self.message_receiver is not None: - self.message_receiver.cancel() + + current_task = asyncio.current_task() + owned_tasks = [] + for task in (self.message_processor, self.current_session_monitor, self.message_receiver): + if task is not None and task is not current_task and task not in owned_tasks: + if not task.done(): + task.cancel() + owned_tasks.append(task) + if owned_tasks: + await asyncio.gather(*owned_tasks, return_exceptions=True) + if self.aiohttp_client_session is not None: await self.aiohttp_client_session.close() diff --git a/tests/slack_sdk_async/socket_mode/test_aiohttp.py b/tests/slack_sdk_async/socket_mode/test_aiohttp.py index eda1fecdd..99b8cc24f 100644 --- a/tests/slack_sdk_async/socket_mode/test_aiohttp.py +++ b/tests/slack_sdk_async/socket_mode/test_aiohttp.py @@ -1,4 +1,5 @@ import asyncio +import time import unittest from slack_sdk.socket_mode.aiohttp import SocketModeClient @@ -8,6 +9,74 @@ from tests.mock_web_api_server import setup_mock_web_api_server_async, cleanup_mock_web_api_server_async +class ControlledClientSession: + def __init__(self): + self.closed = False + self.attempts = 0 + self.attempt_started = asyncio.Event() + self.closed_event = asyncio.Event() + + async def ws_connect(self, *args, **kwargs): + self.attempts += 1 + self.attempt_started.set() + await self.closed_event.wait() + raise RuntimeError("Session is closed") + + async def close(self): + self.closed = True + self.closed_event.set() + + +class CompletingClientSession: + def __init__(self): + self.closed = False + self.attempts = 0 + self.attempt_started = asyncio.Event() + self.closed_event = asyncio.Event() + self.websocket = ConnectedWebSocket() + + async def ws_connect(self, *args, **kwargs): + self.attempts += 1 + self.attempt_started.set() + await self.closed_event.wait() + return self.websocket + + async def close(self): + self.closed = True + self.closed_event.set() + + +class ConnectedWebSocket: + def __init__(self): + self.closed = False + self.receive_forever = asyncio.Event() + + async def close(self): + self.closed = True + + async def ping(self, data): + pass + + async def receive(self): + await self.receive_forever.wait() + + +class RetryingClientSession: + def __init__(self): + self.closed = False + self.attempt_times = [] + self.websocket = ConnectedWebSocket() + + async def ws_connect(self, *args, **kwargs): + self.attempt_times.append(time.monotonic()) + if len(self.attempt_times) == 1: + raise ConnectionError("temporary failure") + return self.websocket + + async def close(self): + self.closed = True + + class TestAiohttp(unittest.TestCase): def setUp(self): setup_mock_web_api_server_async(self, MockHandler) @@ -44,6 +113,167 @@ async def test_init_with_loop(self): finally: await client.close() + @async_test + async def test_connect_returns_when_client_is_already_closed(self): + client = SocketModeClient( + app_token="xapp-A111-222-xyz", + web_client=self.web_client, + auto_reconnect_enabled=False, + ping_interval=0.01, + ) + client.wss_uri = "ws://127.0.0.1:1/link" + + await client.close() + await asyncio.wait_for(client.connect(), timeout=0.05) + + self.assertTrue(client.closed) + self.assertTrue(client.aiohttp_client_session.closed) + + @async_test + async def test_connect_stops_when_close_lands_during_an_attempt(self): + client = SocketModeClient( + app_token="xapp-A111-222-xyz", + web_client=self.web_client, + auto_reconnect_enabled=False, + ping_interval=0.01, + ) + await client.aiohttp_client_session.close() + controlled_session = ControlledClientSession() + client.aiohttp_client_session = controlled_session + client.wss_uri = "ws://127.0.0.1:1/link" + + connect_task = asyncio.create_task(client.connect()) + await controlled_session.attempt_started.wait() + await client.close() + + await asyncio.wait_for(asyncio.shield(connect_task), timeout=0.05) + self.assertTrue(connect_task.done()) + self.assertEqual(1, controlled_session.attempts) + self.assertTrue(controlled_session.closed) + + @async_test + async def test_connect_discards_session_completed_after_close(self): + client = SocketModeClient( + app_token="xapp-A111-222-xyz", + web_client=self.web_client, + auto_reconnect_enabled=False, + ping_interval=0.01, + ) + await client.aiohttp_client_session.close() + completing_session = CompletingClientSession() + client.aiohttp_client_session = completing_session + client.wss_uri = "ws://127.0.0.1:1/link" + + connect_task = asyncio.create_task(client.connect()) + await completing_session.attempt_started.wait() + await client.close() + await asyncio.wait_for(connect_task, timeout=0.05) + + self.assertTrue(completing_session.websocket.closed) + self.assertIsNone(client.current_session) + self.assertIsNone(client.current_session_monitor) + self.assertIsNone(client.message_receiver) + + @async_test + async def test_connect_preserves_ping_interval_backoff_for_live_client(self): + client = SocketModeClient( + app_token="xapp-A111-222-xyz", + web_client=self.web_client, + auto_reconnect_enabled=False, + ping_interval=0.03, + ) + await client.aiohttp_client_session.close() + retrying_session = RetryingClientSession() + client.aiohttp_client_session = retrying_session + client.wss_uri = "ws://127.0.0.1:1/link" + + try: + await client.connect() + self.assertEqual(2, len(retrying_session.attempt_times)) + self.assertGreaterEqual( + retrying_session.attempt_times[1] - retrying_session.attempt_times[0], + client.ping_interval, + ) + finally: + await client.close() + + @async_test + async def test_close_awaits_owned_task_settlement_before_session_close(self): + client = SocketModeClient( + app_token="xapp-A111-222-xyz", + web_client=self.web_client, + auto_reconnect_enabled=False, + ) + client.message_processor.cancel() + await asyncio.gather(client.message_processor, return_exceptions=True) + await client.aiohttp_client_session.close() + controlled_session = ControlledClientSession() + client.aiohttp_client_session = controlled_session + + started = [asyncio.Event() for _ in range(3)] + cancelling = [asyncio.Event() for _ in range(3)] + release_cleanup = asyncio.Event() + + async def owned_task(index): + started[index].set() + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + cancelling[index].set() + await release_cleanup.wait() + raise + + owned_tasks = [asyncio.create_task(owned_task(index)) for index in range(3)] + client.message_processor, client.current_session_monitor, client.message_receiver = owned_tasks + for event in started: + await event.wait() + + close_task = asyncio.create_task(client.close()) + for event in cancelling: + await event.wait() + + self.assertFalse(close_task.done()) + self.assertFalse(controlled_session.closed) + + release_cleanup.set() + await asyncio.wait_for(close_task, timeout=0.05) + + self.assertTrue(controlled_session.closed) + self.assertTrue(all(task.done() for task in owned_tasks)) + self.assertTrue(all(task.cancelled() for task in owned_tasks)) + + @async_test + async def test_close_does_not_cancel_or_await_current_reconnect_owner(self): + client = SocketModeClient( + app_token="xapp-A111-222-xyz", + web_client=self.web_client, + auto_reconnect_enabled=False, + ) + client.message_processor.cancel() + await asyncio.gather(client.message_processor, return_exceptions=True) + await client.aiohttp_client_session.close() + controlled_session = ControlledClientSession() + client.aiohttp_client_session = controlled_session + + async def settled_task(): + pass + + settled = asyncio.create_task(settled_task()) + await settled + client.message_processor = settled + client.message_receiver = settled + + async def reconnect_owner(): + client.current_session_monitor = asyncio.current_task() + await client.close() + + owner_task = asyncio.create_task(reconnect_owner()) + await asyncio.wait_for(asyncio.shield(owner_task), timeout=0.05) + + self.assertTrue(owner_task.done()) + self.assertFalse(owner_task.cancelled()) + self.assertTrue(controlled_session.closed) + @async_test async def test_issue_new_wss_url(self): client = SocketModeClient(