From 2e561f61b5787b89748a6c8aec4b0f5a8bc6c1e9 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 18 Aug 2026 19:05:08 +0000 Subject: [PATCH] fix: flush pending assignment and exposure events on LocalEvaluationClient.stop() stop() previously stopped the flag poller and closed the connection pool but never drained the assignment/exposure Amplitude buffers, silently dropping up to flush_queue_size events per instance whenever the analytics SDK's atexit hook doesn't fire (recycled workers, forked job runners, SIGKILL'd containers). stop() now flushes both instances and waits for the returned futures with a bounded timeout (new parameter, default 10s, None waits forever), logging a warning if batches remain. The instances are deliberately NOT shut down: Amplitude.shutdown() sets opt_out=True on the caller-provided AssignmentConfig/ExposureConfig, which outlives the client and silently drops all events for any instance reusing that config. Flush-only alternative to #79. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01FQQHvcxnQ1Qvo6CaQuwKGn --- src/amplitude_experiment/local/client.py | 29 ++++++++- tests/local/stop_flush_test.py | 79 ++++++++++++++++++++++++ 2 files changed, 105 insertions(+), 3 deletions(-) create mode 100644 tests/local/stop_flush_test.py diff --git a/src/amplitude_experiment/local/client.py b/src/amplitude_experiment/local/client.py index 4e87eda..2a73954 100644 --- a/src/amplitude_experiment/local/client.py +++ b/src/amplitude_experiment/local/client.py @@ -1,5 +1,6 @@ +from concurrent.futures import wait from threading import Lock -from typing import Any, List, Dict, Set +from typing import Any, List, Dict, Set, Optional from amplitude import Amplitude @@ -158,12 +159,34 @@ def __setup_connection_pool(self): self._connection_pool = HTTPConnectionPool(host, max_size=1, idle_timeout=30, read_timeout=timeout, scheme=scheme) - def stop(self) -> None: + def stop(self, timeout: Optional[float] = 10.0) -> None: """ - Stop polling for flag configurations. Close resource like connection pool with client + Stop polling for flag configurations, flush pending assignment and exposure events, and close resources + like the connection pool. + + The assignment and exposure Amplitude instances are flushed but not shut down, so they remain usable + and their caller-provided configurations are not mutated. + + Parameters: + timeout (float | None): Maximum time, in seconds, to wait for pending assignment and exposure + events to finish sending before returning. Defaults to 10 seconds. Pass None to wait + indefinitely. """ self.deployment_runner.stop() self._connection_pool.close() + self.__flush_event_services(timeout) + + def __flush_event_services(self, timeout: Optional[float]) -> None: + instances = [service.amplitude for service in (self.assignment_service, self.exposure_service) + if service is not None] + futures = [] + for instance in instances: + futures.extend(f for f in (instance.flush() or []) if f is not None) + if futures: + _, not_done = wait(futures, timeout=timeout) + if not_done: + self.logger.warning(f"[Experiment] Stop timed out after {timeout}s waiting for " + f"{len(not_done)} pending event batch(es) to flush") def __enter__(self) -> 'LocalEvaluationClient': return self diff --git a/tests/local/stop_flush_test.py b/tests/local/stop_flush_test.py new file mode 100644 index 0000000..21cc1af --- /dev/null +++ b/tests/local/stop_flush_test.py @@ -0,0 +1,79 @@ +import time +import unittest +from concurrent.futures import Future +from unittest.mock import MagicMock + +from src.amplitude_experiment import LocalEvaluationClient, LocalEvaluationConfig +from src.amplitude_experiment.assignment import AssignmentConfig +from src.amplitude_experiment.exposure.exposure_config import ExposureConfig + +API_KEY = 'server-api-key' + + +def completed_future() -> Future: + future = Future() + future.set_result(None) + return future + + +class LocalEvaluationClientStopTestCase(unittest.TestCase): + + def _client_with_event_services(self) -> LocalEvaluationClient: + config = LocalEvaluationConfig( + assignment_config=AssignmentConfig(api_key='analytics-api-key'), + exposure_config=ExposureConfig(api_key='analytics-api-key'), + ) + return LocalEvaluationClient(API_KEY, config) + + def test_stop_flushes_assignment_and_exposure_without_shutdown(self): + client = self._client_with_event_services() + assignment_amplitude = MagicMock() + assignment_amplitude.flush.return_value = [completed_future()] + exposure_amplitude = MagicMock() + exposure_amplitude.flush.return_value = [None] + client.assignment_service.amplitude = assignment_amplitude + client.exposure_service.amplitude = exposure_amplitude + + client.stop() + + assignment_amplitude.flush.assert_called_once() + exposure_amplitude.flush.assert_called_once() + # The instances are not shut down: shutdown() would set opt_out=True on the + # caller-provided config, silently dropping events for any client reusing it. + assignment_amplitude.shutdown.assert_not_called() + exposure_amplitude.shutdown.assert_not_called() + + def test_stop_timeout_bounds_wait_on_pending_events(self): + client = self._client_with_event_services() + never_completes = Future() + assignment_amplitude = MagicMock() + assignment_amplitude.flush.return_value = [never_completes] + client.assignment_service.amplitude = assignment_amplitude + client.exposure_service.amplitude = MagicMock(flush=MagicMock(return_value=[])) + + start = time.monotonic() + client.stop(timeout=0.2) + elapsed = time.monotonic() - start + + self.assertLess(elapsed, 2) + + def test_stop_without_event_services(self): + client = LocalEvaluationClient(API_KEY, LocalEvaluationConfig()) + client.stop() + + def test_context_manager_exit_flushes(self): + client = self._client_with_event_services() + exposure_amplitude = MagicMock() + exposure_amplitude.flush.return_value = [completed_future()] + client.assignment_service.amplitude = MagicMock(flush=MagicMock(return_value=[])) + client.exposure_service.amplitude = exposure_amplitude + + with client: + pass + + exposure_amplitude.flush.assert_called_once() + exposure_amplitude.shutdown.assert_not_called() + + +if __name__ == '__main__': + unittest.main()