Skip to content

Commit 21e1e09

Browse files
jsonbaileyjsonbailey
committed
fix: Close the owned streaming session if _run fails
The async streaming processor builds its own aiohttp session in _run when no factory is injected, but only closed it on a normal loop exit or via stop(). If a later step raised — create(), the SSE iterator, or interrupt() — the exception escaped _run with no finally and the session leaked. Wrap the stream body in try/finally so the owned session is always closed.
1 parent c7f3f28 commit 21e1e09

2 files changed

Lines changed: 71 additions & 46 deletions

File tree

‎ldclient/impl/datasource/async_streaming.py‎

Lines changed: 49 additions & 46 deletions
Original file line numberDiff line numberDiff line change
@@ -72,52 +72,55 @@ async def _run(self):
7272
self._sse_factory = AsyncSSEFactory(self._config, session=self._owned_session)
7373
log.info("Starting AsyncStreamingUpdateProcessor connecting to uri: " + self._uri)
7474
self._running = True
75-
self._sse = self._sse_factory.create(self._uri, self._config.initial_reconnect_delay)
76-
self._connection_attempt_start_time = time.time()
77-
async for action in self._sse.all:
78-
if isinstance(action, Start):
79-
# On reconnect after an error the timer was cleared; reset it here.
80-
# For the initial connect the pre-loop timestamp is already set.
81-
if self._connection_attempt_start_time is None:
82-
self._connection_attempt_start_time = time.time()
83-
elif isinstance(action, Event):
84-
message_ok = False
85-
try:
86-
message_ok = await self._process_message(action)
87-
except json.decoder.JSONDecodeError as e:
88-
log.info("Error while handling stream event; will restart stream: %s" % e)
89-
await self._sse.interrupt()
90-
91-
await self._handle_error(e)
92-
except Exception as e:
93-
log.warning("Error while handling stream event; will restart stream: %s" % e)
94-
await self._sse.interrupt()
95-
96-
if self._data_source_update_sink is not None:
97-
error_info = DataSourceErrorInfo(DataSourceErrorKind.UNKNOWN, 0, time.time(), str(e))
98-
99-
self._data_source_update_sink.update_status(DataSourceState.INTERRUPTED, error_info)
100-
101-
if message_ok:
102-
self._record_stream_init(False)
103-
self._connection_attempt_start_time = None
104-
105-
if self._data_source_update_sink is not None:
106-
self._data_source_update_sink.update_status(DataSourceState.VALID, None)
107-
108-
if not self._ready.is_set():
109-
log.info("AsyncStreamingUpdateProcessor initialized ok.")
110-
self._ready.set()
111-
elif isinstance(action, Fault):
112-
# If the SSE client detects the stream has closed, then it will emit a fault with no-error. We can
113-
# ignore this since we want the connection to continue.
114-
if action.error is None:
115-
continue
116-
117-
if not await self._handle_error(action.error):
118-
break
119-
await self._sse.close()
120-
await self._close_owned_session()
75+
try:
76+
self._sse = self._sse_factory.create(self._uri, self._config.initial_reconnect_delay)
77+
self._connection_attempt_start_time = time.time()
78+
async for action in self._sse.all:
79+
if isinstance(action, Start):
80+
# On reconnect after an error the timer was cleared; reset it here.
81+
# For the initial connect the pre-loop timestamp is already set.
82+
if self._connection_attempt_start_time is None:
83+
self._connection_attempt_start_time = time.time()
84+
elif isinstance(action, Event):
85+
message_ok = False
86+
try:
87+
message_ok = await self._process_message(action)
88+
except json.decoder.JSONDecodeError as e:
89+
log.info("Error while handling stream event; will restart stream: %s" % e)
90+
await self._sse.interrupt()
91+
92+
await self._handle_error(e)
93+
except Exception as e:
94+
log.warning("Error while handling stream event; will restart stream: %s" % e)
95+
await self._sse.interrupt()
96+
97+
if self._data_source_update_sink is not None:
98+
error_info = DataSourceErrorInfo(DataSourceErrorKind.UNKNOWN, 0, time.time(), str(e))
99+
100+
self._data_source_update_sink.update_status(DataSourceState.INTERRUPTED, error_info)
101+
102+
if message_ok:
103+
self._record_stream_init(False)
104+
self._connection_attempt_start_time = None
105+
106+
if self._data_source_update_sink is not None:
107+
self._data_source_update_sink.update_status(DataSourceState.VALID, None)
108+
109+
if not self._ready.is_set():
110+
log.info("AsyncStreamingUpdateProcessor initialized ok.")
111+
self._ready.set()
112+
elif isinstance(action, Fault):
113+
# If the SSE client detects the stream has closed, then it will emit a fault with no-error. We can
114+
# ignore this since we want the connection to continue.
115+
if action.error is None:
116+
continue
117+
118+
if not await self._handle_error(action.error):
119+
break
120+
finally:
121+
if self._sse:
122+
await self._sse.close()
123+
await self._close_owned_session()
121124

122125
async def _close_owned_session(self):
123126
"""Close the aiohttp session if the SDK created it. A caller-supplied

‎ldclient/testing/impl/datasource/test_async_streaming.py‎

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -440,6 +440,28 @@ async def test_default_construction_session_closed_on_stop():
440440
assert proc._owned_session is None
441441

442442

443+
@pytest.mark.asyncio
444+
async def test_default_construction_session_closed_when_run_fails():
445+
"""If _run fails after building the session (e.g. create() raises), the
446+
SDK-created session is closed rather than leaked."""
447+
config = _make_config()
448+
store = MockAsyncFeatureStore()
449+
ready = asyncio.Event()
450+
fake_session = _FakeSession()
451+
452+
with mock.patch.object(
453+
async_streaming, "make_client_session", return_value=fake_session
454+
), mock.patch.object(async_streaming, "AsyncSSEFactory") as factory_cls:
455+
factory_cls.return_value.create.side_effect = RuntimeError("boom")
456+
proc = AsyncStreamingUpdateProcessor(config, store, ready, None)
457+
proc.start()
458+
459+
# _run builds the session, then create() raises. The finally must still
460+
# close the SDK-created session instead of leaking it.
461+
await _wait_until(lambda: fake_session.closed)
462+
assert proc._owned_session is None
463+
464+
443465
@pytest.mark.asyncio
444466
async def test_injected_factory_leaves_session_unowned():
445467
"""When a factory is injected, no session is built and none is owned."""

0 commit comments

Comments
 (0)