From 41c2df708f768fd0acc24e38d3a9a06ac0fcea19 Mon Sep 17 00:00:00 2001 From: cmyui Date: Tue, 4 Aug 2026 13:04:13 -0400 Subject: [PATCH] fix: flush pending assignment and exposure events on stop() LocalEvaluationClient.stop() stopped the flag config poller and closed the connection pool, but never flushed or shut down the Amplitude analytics instances backing the assignment and exposure services. Any events still in their buffers (up to flush_queue_size per instance, accumulating for up to flush_interval_millis) were silently dropped unless the interpreter happened to exit cleanly enough for the analytics SDK's atexit hook to fire - which it often doesn't in forked/reaped server workers. stop() now flushes both instances, waits up to a configurable timeout (new parameter, default 10s, None = wait indefinitely) for the pending batches to send, then shuts the instances down. Instances shut down after the flush so late-tracked events are dropped deliberately rather than accumulating in a stopped client. Co-Authored-By: Claude Fable 5 --- src/amplitude_experiment/local/client.py | 30 ++++++++- tests/local/stop_flush_test.py | 78 ++++++++++++++++++++++++ 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..b671549 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,35 @@ 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. + + 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.__shutdown_event_services(timeout) + + def __shutdown_event_services(self, timeout: Optional[float]) -> None: + instances = [service.amplitude for service in (self.assignment_service, self.exposure_service) + if service is not None] + if not instances: + return + 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") + for instance in instances: + instance.shutdown() 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..0e791a2 --- /dev/null +++ b/tests/local/stop_flush_test.py @@ -0,0 +1,78 @@ +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_then_shuts_down_assignment_and_exposure(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() + assignment_amplitude.shutdown.assert_called_once() + exposure_amplitude.shutdown.assert_called_once() + + 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) + assignment_amplitude.shutdown.assert_called_once() + + 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_called_once() + + +if __name__ == '__main__': + unittest.main()