diff --git a/ldclient/impl/aio/concurrency.py b/ldclient/impl/aio/concurrency.py index 42cf80b1..ed3f7962 100644 --- a/ldclient/impl/aio/concurrency.py +++ b/ldclient/impl/aio/concurrency.py @@ -242,8 +242,6 @@ async def _run(self): result = self.__action() if inspect.isawaitable(result): await result - except asyncio.CancelledError: - raise except Exception as e: log.exception("Unexpected exception on worker task: %s" % e) delay = next_time - time.time() diff --git a/ldclient/impl/datasource/async_status.py b/ldclient/impl/datasource/async_status.py index 09f9880d..1bebc281 100644 --- a/ldclient/impl/datasource/async_status.py +++ b/ldclient/impl/datasource/async_status.py @@ -3,7 +3,6 @@ from ldclient.impl.dependency_tracker import DependencyTracker, KindAndKey from ldclient.impl.listeners import Listeners -from ldclient.impl.rwlock import ReadWriteLock from ldclient.interfaces import ( AsyncDataSourceUpdateSink, AsyncFeatureStore, @@ -23,13 +22,11 @@ def __init__(self, store: AsyncFeatureStore, status_listeners: Listeners, flag_c self.__flag_change_listeners = flag_change_listeners self.__tracker = DependencyTracker() - self.__lock = ReadWriteLock() self.__status = DataSourceStatus(DataSourceState.INITIALIZING, time.time(), None) @property def status(self) -> DataSourceStatus: - with self.__lock.read(): - return self.__status + return self.__status async def init(self, all_data: Mapping[VersionedDataKind, Mapping[str, dict]]) -> None: old_data: Optional[Dict[VersionedDataKind, Mapping[str, dict]]] = None @@ -73,22 +70,21 @@ async def delete(self, kind: VersionedDataKind, key: str, version: int) -> None: def update_status(self, new_state: DataSourceState, new_error: Optional[DataSourceErrorInfo]) -> None: status_to_broadcast = None - with self.__lock.write(): - old_status = self.__status + old_status = self.__status - if new_state == DataSourceState.INTERRUPTED and old_status.state == DataSourceState.INITIALIZING: - new_state = DataSourceState.INITIALIZING + if new_state == DataSourceState.INTERRUPTED and old_status.state == DataSourceState.INITIALIZING: + new_state = DataSourceState.INITIALIZING - if new_state == old_status.state and new_error is None: - return + if new_state == old_status.state and new_error is None: + return - self.__status = DataSourceStatus( - new_state, - self.__status.since if new_state == self.__status.state else time.time(), - self.__status.error if new_error is None else new_error, - ) + self.__status = DataSourceStatus( + new_state, + self.__status.since if new_state == self.__status.state else time.time(), + self.__status.error if new_error is None else new_error, + ) - status_to_broadcast = self.__status + status_to_broadcast = self.__status if status_to_broadcast is not None: self.__status_listeners.notify(status_to_broadcast) diff --git a/ldclient/impl/datasourcev2/async_polling.py b/ldclient/impl/datasourcev2/async_polling.py new file mode 100644 index 00000000..80c3052c --- /dev/null +++ b/ldclient/impl/datasourcev2/async_polling.py @@ -0,0 +1,434 @@ +""" +This module contains the implementations of a polling synchronizer and +initializer, along with any required supporting classes and protocols. +""" + +import json +from abc import abstractmethod +from collections import namedtuple +from typing import AsyncGenerator, Mapping, Optional, Protocol, Tuple +from urllib import parse + +from ldclient.config import ( + DataSourceBuilder, + DataSourceBuilderConfig, + HTTPConfig +) +from ldclient.impl.aio.concurrency import AsyncEvent +from ldclient.impl.aio.transport import AsyncHTTPTransport +from ldclient.impl.datasource.async_feature_requester import ( + FDV1_POLLING_ENDPOINT +) +from ldclient.impl.datasourcev2.polling_common import ( + PollAction, + fdv1_polling_payload_to_changeset, + map_polling_result, + polling_payload_to_changeset, + polling_result_to_basis +) +from ldclient.impl.util import ( + UnsuccessfulResponseException, + _Fail, + _headers, + _Result, + _Success, + log +) +from ldclient.interfaces import ( + AsyncInitializer, + AsyncSynchronizer, + BasisResult, + ChangeSet, + ChangeSetBuilder, + Selector, + SelectorStore, + Update +) + +FDV2_POLLING_ENDPOINT = "/sdk/poll" + + +PollingResult = _Result[Tuple[ChangeSet, Mapping], str] + + +class AsyncRequester(Protocol): # pylint: disable=too-few-public-methods + """ + AsyncRequester allows AsyncPollingDataSource to delegate fetching data to + another component. + + This is useful for testing the AsyncPollingDataSource without needing to set up + a test HTTP server. + """ + + @abstractmethod + async def fetch(self, selector: Optional[Selector]) -> PollingResult: + """ + Fetches the data for the given selector. + Returns a Result containing a tuple of ChangeSet and any request headers, + or an error if the data could not be retrieved. + """ + raise NotImplementedError + + @abstractmethod + async def close(self) -> None: + """ + Releases any resources (such as an HTTP transport) owned by the + requester. + """ + raise NotImplementedError + + +CacheEntry = namedtuple("CacheEntry", ["data", "etag"]) + + +class AsyncPollingDataSource(AsyncInitializer, AsyncSynchronizer): + """ + AsyncPollingDataSource is a data source that can retrieve information from + LaunchDarkly either as an initializer or as a synchronizer. + """ + + def __init__( + self, + poll_interval: float, + requester: AsyncRequester, + ): + self._requester = requester + self._poll_interval = poll_interval + self._interrupt_event = AsyncEvent() + self._stop = AsyncEvent() + + @property + def name(self) -> str: + """Returns the name of the initializer.""" + return "PollingDataSourceV2" + + async def fetch(self, ss: SelectorStore) -> BasisResult: + """ + Fetch returns a Basis, or an error if the Basis could not be retrieved. + """ + return await self._poll(ss) + + async def sync(self, ss: SelectorStore) -> AsyncGenerator[Update, None]: + """ + sync begins the synchronization process for the data source, yielding + Update objects until the connection is closed or an unrecoverable error + occurs. + """ + log.info("Starting PollingDataSourceV2 synchronizer") + self._interrupt_event.clear() + self._stop.clear() + try: + while self._stop.is_set() is False: + result = await self._requester.fetch(ss.selector()) + decision = map_polling_result(result) + yield decision.update + + if decision.control is PollAction.BREAK: + break + if decision.control is PollAction.WAIT_CONTINUE: + await self._interrupt_event.wait(self._poll_interval) + continue + if await self._interrupt_event.wait(self._poll_interval): + break + finally: + await self._requester.close() + + async def stop(self): + """Signals the synchronizer to stop.""" + log.info("Stopping PollingDataSourceV2 synchronizer") + self._interrupt_event.set() + self._stop.set() + + async def _poll(self, ss: SelectorStore) -> BasisResult: + try: + result = await self._requester.fetch(ss.selector()) + return polling_result_to_basis(result) + except Exception as e: # pylint: disable=broad-except + msg = f"Error: Exception encountered when updating flags. {e}" + log.exception(msg) + + return _Fail(error=msg, exception=e) + + +# pylint: disable=too-few-public-methods +class AiohttpPollingRequester(AsyncRequester): + """ + A requester implementation that issues HTTP requests through the SDK's + HTTP transport. + """ + + def __init__( + self, + config: DataSourceBuilderConfig, + base_uri: str, + http_options: HTTPConfig, + session=None, + ): + self._etag: Optional[str] = None + self._http = AsyncHTTPTransport( + config, + client=session, + http_options=http_options, + ) + self._http_options = http_options + self._config = config + self._poll_uri = base_uri + FDV2_POLLING_ENDPOINT + + async def fetch(self, selector: Optional[Selector]) -> PollingResult: + """ + Fetches the data for the given selector. + Returns a Result containing a tuple of ChangeSet and any request headers, + or an error if the data could not be retrieved. + """ + query_params = {} + if self._config.payload_filter_key is not None: + query_params["filter"] = self._config.payload_filter_key + + if selector is not None and selector.is_defined(): + query_params["selector"] = selector.state + + uri = self._poll_uri + if len(query_params) > 0: + filter_query = parse.urlencode(query_params) + uri += f"?{filter_query}" + + hdrs = _headers(self._config) + hdrs["Accept-Encoding"] = "gzip" + + if self._etag is not None: + hdrs["If-None-Match"] = self._etag + + response = await self._http.request( + "GET", + uri, + headers=hdrs, + ) + headers = response.headers + + if response.status >= 400: + return _Fail( + f"HTTP error {response.status}", UnsuccessfulResponseException(response.status), + headers=headers, + ) + + if response.status == 304: + return _Success(value=(ChangeSetBuilder.no_changes(), headers)) + + data = json.loads(response.body) + etag = headers.get("ETag") + + if etag is not None: + self._etag = etag + + log.debug( + "%s response status:[%d] ETag:[%s]", + uri, + response.status, + etag, + ) + + changeset_result = polling_payload_to_changeset(data) + if isinstance(changeset_result, _Success): + return _Success(value=(changeset_result.value, headers)) + + return _Fail( + error=changeset_result.error, + exception=changeset_result.exception, + headers=headers, # type: ignore + ) + + async def close(self) -> None: + """Closes the requester's HTTP transport.""" + await self._http.close() + + +class AsyncPollingDataSourceBuilder(DataSourceBuilder): + """ + Builder for a AsyncPollingDataSource. + """ + + def __init__(self): + self.__base_uri: Optional[str] = None + self.__poll_interval: Optional[float] = None + self.__http_options: Optional[HTTPConfig] = None + self.__requester: Optional[AsyncRequester] = None + self.__session = None + + def base_uri(self, uri: str) -> 'AsyncPollingDataSourceBuilder': + """Sets the base URI for the streaming data source.""" + self.__base_uri = uri.rstrip('/') + return self + + def poll_interval(self, poll_interval: float) -> 'AsyncPollingDataSourceBuilder': + """Sets the polling interval for the AsyncPollingDataSource.""" + self.__poll_interval = poll_interval + return self + + def http_options(self, http_options: HTTPConfig) -> 'AsyncPollingDataSourceBuilder': + """Sets the HTTP options for the streaming data source.""" + self.__http_options = http_options + return self + + def requester(self, requester: AsyncRequester) -> 'AsyncPollingDataSourceBuilder': + """Sets a custom AsyncRequester for the AsyncPollingDataSource.""" + self.__requester = requester + return self + + def session(self, session) -> 'AsyncPollingDataSourceBuilder': + """Sets the aiohttp session used for HTTP requests.""" + self.__session = session + return self + + def build(self, config: DataSourceBuilderConfig) -> AsyncPollingDataSource: + """Builds the AsyncPollingDataSource with the configured parameters.""" + requester = ( + self.__requester + if self.__requester is not None + else AiohttpPollingRequester( + config, + self.__base_uri or config.base_uri, + self.__http_options or config.http, + session=self.__session, + ) + ) + + return AsyncPollingDataSource( + poll_interval=self.__poll_interval or config.poll_interval, + requester=requester + ) + + +class AsyncFallbackToFDv1PollingDataSourceBuilder(DataSourceBuilder): + """ + Builder for a AsyncPollingDataSource that falls back to Flag Delivery v1. + """ + + def __init__(self): + self.__base_uri: Optional[str] = None + self.__poll_interval: Optional[float] = None + self.__http_options: Optional[HTTPConfig] = None + self.__session = None + + def base_uri(self, uri: str) -> 'AsyncFallbackToFDv1PollingDataSourceBuilder': + """Sets the base URI for the data source.""" + self.__base_uri = uri.rstrip('/') + return self + + def poll_interval(self, poll_interval: float) -> 'AsyncFallbackToFDv1PollingDataSourceBuilder': + """Sets the polling interval for the data source.""" + self.__poll_interval = poll_interval + return self + + def http_options(self, http_options: HTTPConfig) -> 'AsyncFallbackToFDv1PollingDataSourceBuilder': + """Sets the HTTP options for the data source.""" + self.__http_options = http_options + return self + + def session(self, session) -> 'AsyncFallbackToFDv1PollingDataSourceBuilder': + """Sets the aiohttp session used for HTTP requests.""" + self.__session = session + return self + + def build(self, config: DataSourceBuilderConfig) -> AsyncPollingDataSource: + """Builds the AsyncPollingDataSource with the configured parameters.""" + builder = AsyncPollingDataSourceBuilder() + builder.requester( + AiohttpFDv1PollingRequester( + config, + self.__base_uri or config.base_uri, + self.__http_options or config.http, + session=self.__session, + ) + ) + builder.poll_interval(self.__poll_interval or config.poll_interval) + + return builder.build(config) + + +# pylint: disable=too-few-public-methods +class AiohttpFDv1PollingRequester(AsyncRequester): + """ + A requester implementation for the Flag Delivery v1 polling endpoint that + issues HTTP requests through the SDK's HTTP transport. + """ + + def __init__( + self, + config: DataSourceBuilderConfig, + base_uri: str, + http_options: HTTPConfig, + session=None, + ): + self._etag: Optional[str] = None + self._http = AsyncHTTPTransport( + config, + client=session, + http_options=http_options, + ) + self._http_options = http_options + self._config = config + self._poll_uri = base_uri + FDV1_POLLING_ENDPOINT + + async def fetch(self, selector: Optional[Selector]) -> PollingResult: + """ + Fetches the data for the given selector. + Returns a Result containing a tuple of ChangeSet and any request headers, + or an error if the data could not be retrieved. + """ + query_params = {} + if self._config.payload_filter_key is not None: + query_params["filter"] = self._config.payload_filter_key + + uri = self._poll_uri + if len(query_params) > 0: + filter_query = parse.urlencode(query_params) + uri += f"?{filter_query}" + + hdrs = _headers(self._config) + hdrs["Accept-Encoding"] = "gzip" + + if self._etag is not None: + hdrs["If-None-Match"] = self._etag + + response = await self._http.request( + "GET", + uri, + headers=hdrs, + ) + + headers = response.headers + if response.status >= 400: + return _Fail( + f"HTTP error {response.status}", UnsuccessfulResponseException(response.status), + headers=headers + ) + + if response.status == 304: + return _Success(value=(ChangeSetBuilder.no_changes(), headers)) + + data = json.loads(response.body) + etag = headers.get("ETag") + + if etag is not None: + self._etag = etag + + log.debug( + "%s response status:[%d] ETag:[%s]", + uri, + response.status, + etag, + ) + + changeset_result = fdv1_polling_payload_to_changeset(data) + if isinstance(changeset_result, _Success): + return _Success(value=(changeset_result.value, headers)) + + return _Fail( + error=changeset_result.error, + exception=changeset_result.exception, + headers=headers, + ) + + async def close(self) -> None: + """Closes the requester's HTTP transport.""" + await self._http.close() diff --git a/ldclient/impl/datasourcev2/async_streaming.py b/ldclient/impl/datasourcev2/async_streaming.py new file mode 100644 index 00000000..42013ea6 --- /dev/null +++ b/ldclient/impl/datasourcev2/async_streaming.py @@ -0,0 +1,326 @@ +""" +This module contains the implementations of a streaming synchronizer, along +with any required supporting classes and protocols. +""" + +import json +from time import time +from typing import AsyncGenerator, Callable, Optional, Tuple +from urllib import parse + +import aiohttp +from ld_eventsource import AsyncSSEClient +from ld_eventsource.actions import Event, Fault, Start + +from ldclient.config import ( + DataSourceBuilder, + DataSourceBuilderConfig, + HTTPConfig +) +from ldclient.impl.aio.transport import AsyncSSEFactory, make_client_session +from ldclient.impl.datasourcev2.streaming_common import ( + classify_stream_error, + process_message, + with_fallback_signal +) +from ldclient.impl.datasystem import DiagnosticAccumulator, DiagnosticSource +from ldclient.impl.util import _LD_ENVID_HEADER, _LD_FD_FALLBACK_HEADER, log +from ldclient.interfaces import ( + AsyncSynchronizer, + ChangeSetBuilder, + DataSourceErrorInfo, + DataSourceErrorKind, + DataSourceState, + SelectorStore, + Update +) + +STREAMING_ENDPOINT = "/sdk/stream" + +SseClientBuilder = Callable[ + [str, HTTPConfig, float, DataSourceBuilderConfig, SelectorStore], + Tuple[AsyncSSEClient, Optional[aiohttp.ClientSession]], +] + + +def create_sse_client( + base_uri: str, + http_options: HTTPConfig, + initial_reconnect_delay: float, + config: DataSourceBuilderConfig, + ss: SelectorStore, + session=None, +) -> Tuple[AsyncSSEClient, Optional[aiohttp.ClientSession]]: + """ + create_sse_client creates an SSE client configured to connect to the + LaunchDarkly streaming endpoint, along with the aiohttp session backing it + when the SDK created that session itself. + + When no ``session`` is supplied, one is built from the SDK's HTTP options + via ``make_client_session`` (CA certs, client cert, SSL verification, proxy + trust, connector limits) and returned as the second element so the caller + can close it on shutdown -- the SSE client treats the supplied session as + externally owned and never closes it. When a ``session`` is supplied, the + caller owns it and ``None`` is returned in its place. + """ + uri = base_uri + STREAMING_ENDPOINT + if config.payload_filter_key is not None: + uri += "?%s" % parse.urlencode({"filter": config.payload_filter_key}) + + def query_params() -> dict: + selector = ss.selector() + return {"basis": selector.state} if selector.is_defined() else {} + + if session is None: + session = make_client_session(config, http_options) + owned_session: Optional[aiohttp.ClientSession] = session + else: + owned_session = None + + factory = AsyncSSEFactory( + config, + session=session, + http_options=http_options, + ) + sse_client = factory.create(uri, initial_reconnect_delay, query_params=query_params) + return sse_client, owned_session + + +class AsyncStreamingDataSource(AsyncSynchronizer, DiagnosticSource): + """ + AsyncStreamingDataSource is a specific type of synchronizer that handles + streaming data sources. + + It should implement the sync method to yield updates as they are received + from the streaming data source. + """ + + def __init__( + self, + uri: str, + http_options: HTTPConfig, + initial_reconnect_delay: float, + config: DataSourceBuilderConfig, + session=None, + ): + self.__uri = uri + self.__http_options = http_options + self.__initial_reconnect_delay = initial_reconnect_delay + + self._sse_client_builder: SseClientBuilder = create_sse_client + self._config = config + self._session = session + # Build the default SSE builder here, not in __init__, so the session + # passed to the constructor reaches the SSE client. + self._sse_client_builder = lambda *args: create_sse_client(*args, session=self._session) # type: ignore[misc] + self._sse: Optional[AsyncSSEClient] = None + self._owned_session: Optional[aiohttp.ClientSession] = None + self._running = False + self._diagnostic_accumulator: Optional[DiagnosticAccumulator] = None + self._connection_attempt_start_time: Optional[float] = None + + def set_diagnostic_accumulator(self, diagnostic_accumulator: DiagnosticAccumulator): + self._diagnostic_accumulator = diagnostic_accumulator + + @property + def name(self) -> str: + """ + Returns the name of the synchronizer, which is used for logging and debugging. + """ + return "streaming" + + async def sync(self, ss: SelectorStore) -> AsyncGenerator[Update, None]: + """ + sync should begin the synchronization process for the data source, yielding + Update objects until the connection is closed or an unrecoverable error + occurs. + """ + self._sse, self._owned_session = self._sse_client_builder( + self.__uri, + self.__http_options, + self.__initial_reconnect_delay, + self._config, + ss + ) + + if self._sse is None: + log.error("Failed to create SSE client for streaming updates.") + await self._close_owned_session() + return + + change_set_builder = ChangeSetBuilder() + self._running = True + self._connection_attempt_start_time = time() + + envid = None + # A Start action with the X-LD-FD-Fallback: true header sets fallback_requested. + # We apply the current payload before halting, so consumers can serve that data + # while FDv1 takes over. Once set, the flag stays set: every later payload event + # or error re-emits the signal and halts the stream, even when that event has no + # directive header of its own. + fallback_requested = False + + try: + async for action in self._sse.all: + if isinstance(action, Fault): + # If the SSE client detects the stream has closed, then it will + # emit a fault with no-error. We can ignore this since we want + # the connection to continue. + if action.error is None: + continue + + if action.headers is not None: + envid = action.headers.get(_LD_ENVID_HEADER, envid) + + (update, should_continue) = await self._handle_error(action.error, envid) + if update is not None: + yield with_fallback_signal(update, fallback_requested) + + # The FDv1 Fallback Directive is one-way and terminal: if it + # was latched on a prior Start, we must not keep retrying the + # FDv2 endpoint even when the failure itself looks recoverable. + if fallback_requested or not should_continue: + break + continue + + if isinstance(action, Start) and action.headers is not None: + envid = action.headers.get(_LD_ENVID_HEADER, envid) + if action.headers.get(_LD_FD_FALLBACK_HEADER) == 'true': + fallback_requested = True + + if not isinstance(action, Event): + continue + + try: + update = process_message(action, change_set_builder, envid) + if update is not None: + self._record_stream_init(False) + self._connection_attempt_start_time = None + if fallback_requested: + # The completed update is the natural moment to honor + # the latched directive: yield once with the signal, + # then halt — the consumer will switch to FDv1. + yield with_fallback_signal(update, fallback_requested) + break + yield update + except json.decoder.JSONDecodeError as e: + log.info( + "Error while handling stream event; will restart stream: %s", e + ) + await self._sse.interrupt() + + (update, should_continue) = await self._handle_error(e, envid) + if update is not None: + yield with_fallback_signal(update, fallback_requested) + if fallback_requested or not should_continue: + break + except Exception as e: # pylint: disable=broad-except + log.info( + "Error while handling stream event; will restart stream: %s", e + ) + await self._sse.interrupt() + + yield with_fallback_signal(Update( + state=DataSourceState.INTERRUPTED, + error=DataSourceErrorInfo( + DataSourceErrorKind.UNKNOWN, 0, time(), str(e) + ), + fallback_to_fdv1=False, + environment_id=envid, + ), fallback_requested) + if fallback_requested: + break + finally: + await self._sse.close() + await self._close_owned_session() + + async def stop(self): + """ + Stops the streaming synchronizer, closing any open connections. + """ + log.info("Stopping StreamingUpdateProcessor") + self._running = False + if self._sse: + await self._sse.close() + await self._close_owned_session() + + async def _close_owned_session(self): + """Close the aiohttp session if the SDK created it. A caller-supplied + session is owned by the caller and is never closed here. Closing sets + the reference back to ``None`` so it isn't closed twice.""" + if self._owned_session is not None: + await self._owned_session.close() + self._owned_session = None + + def _record_stream_init(self, failed: bool): + if self._diagnostic_accumulator and self._connection_attempt_start_time: + current_time = int(time() * 1000) + elapsed = current_time - int(self._connection_attempt_start_time * 1000) + self._diagnostic_accumulator.record_stream_init(current_time, elapsed if elapsed >= 0 else 0, failed) + + async def _handle_error(self, error: Exception, envid: Optional[str]) -> Tuple[Optional[Update], bool]: + """ + This method handles errors that occur during the streaming process. + + It may return an update indicating the error state, and a boolean + indicating whether the synchronizer should continue retrying the connection. + + If an update is provided, it should be forward upstream, regardless of + whether or not we are going to retry this failure. + + The return should be thought of (update, should_continue) + """ + if not self._running: + return (None, False) # don't retry if we've been deliberately stopped + + decision = classify_stream_error( + error, self._sse.next_retry_delay, envid # type: ignore + ) + self._record_stream_init(True) + self._connection_attempt_start_time = decision.next_attempt_start + if decision.should_stop: + await self.stop() + + return (decision.update, decision.should_continue) + + +class AsyncStreamingDataSourceBuilder(DataSourceBuilder): + """ + Builder for a AsyncStreamingDataSource. + """ + + def __init__(self): + self.__base_uri: Optional[str] = None + self.__initial_reconnect_delay: Optional[float] = None + self.__http_options: Optional[HTTPConfig] = None + self.__session = None + + def base_uri(self, uri: str) -> 'AsyncStreamingDataSourceBuilder': + """Sets the base URI for the streaming data source.""" + self.__base_uri = uri.rstrip('/') + return self + + def initial_reconnect_delay(self, delay: float) -> 'AsyncStreamingDataSourceBuilder': + """Sets the initial reconnect delay for the streaming data source.""" + self.__initial_reconnect_delay = delay + return self + + def http_options(self, http_options: HTTPConfig) -> 'AsyncStreamingDataSourceBuilder': + """Sets the HTTP options for the streaming data source.""" + self.__http_options = http_options + return self + + def session(self, session) -> 'AsyncStreamingDataSourceBuilder': + """Sets the aiohttp session for the streaming data source.""" + self.__session = session + return self + + def build(self, config: DataSourceBuilderConfig) -> AsyncStreamingDataSource: + """Builds a AsyncStreamingDataSource instance with the configured parameters.""" + return AsyncStreamingDataSource( + self.__base_uri or config.stream_base_uri, + self.__http_options or config.http, + self.__initial_reconnect_delay or config.initial_reconnect_delay, + config, + session=self.__session, + ) diff --git a/ldclient/impl/datasourcev2/polling.py b/ldclient/impl/datasourcev2/polling.py index b0333113..54877d41 100644 --- a/ldclient/impl/datasourcev2/polling.py +++ b/ldclient/impl/datasourcev2/polling.py @@ -7,7 +7,6 @@ from abc import abstractmethod from collections import namedtuple from threading import Event -from time import time from typing import Generator, Mapping, Optional, Protocol, Tuple from urllib import parse @@ -19,38 +18,29 @@ HTTPConfig ) from ldclient.impl.datasource.feature_requester import FDV1_POLLING_ENDPOINT -from ldclient.impl.datasystem.protocolv2 import ( - DeleteObject, - EventName, - PutObject +from ldclient.impl.datasourcev2.polling_common import ( + PollAction, + fdv1_polling_payload_to_changeset, + map_polling_result, + polling_payload_to_changeset, + polling_result_to_basis ) from ldclient.impl.http import HTTPFactory, _base_headers from ldclient.impl.util import ( - _LD_ENVID_HEADER, - _LD_FD_FALLBACK_HEADER, UnsuccessfulResponseException, _Fail, _headers, _Result, _Success, - http_error_message, - is_http_error_recoverable, log ) from ldclient.interfaces import ( - Basis, BasisResult, ChangeSet, ChangeSetBuilder, - DataSourceErrorInfo, - DataSourceErrorKind, - DataSourceState, Initializer, - IntentCode, - ObjectKind, Selector, SelectorStore, - ServerIntent, Synchronizer, Update ) @@ -121,82 +111,14 @@ def sync(self, ss: SelectorStore) -> Generator[Update, None, None]: self._stop.clear() while self._stop.is_set() is False: result = self._requester.fetch(ss.selector()) - if isinstance(result, _Fail): - fallback = None - envid = None - - if result.headers is not None: - fallback = result.headers.get(_LD_FD_FALLBACK_HEADER) == 'true' - envid = result.headers.get(_LD_ENVID_HEADER) - - if isinstance(result.exception, UnsuccessfulResponseException): - error_info = DataSourceErrorInfo( - kind=DataSourceErrorKind.ERROR_RESPONSE, - status_code=result.exception.status, - time=time(), - message=http_error_message( - result.exception.status, "polling request" - ), - ) - - if fallback: - yield Update( - state=DataSourceState.OFF, - error=error_info, - fallback_to_fdv1=True, - environment_id=envid, - ) - break - - status_code = result.exception.status - if is_http_error_recoverable(status_code): - yield Update( - state=DataSourceState.INTERRUPTED, - error=error_info, - environment_id=envid, - ) - self._interrupt_event.wait(self._poll_interval) - continue - - yield Update( - state=DataSourceState.OFF, - error=error_info, - environment_id=envid, - ) - break - - error_info = DataSourceErrorInfo( - kind=DataSourceErrorKind.NETWORK_ERROR, - time=time(), - status_code=0, - message=result.error, - ) - - # Even a non-HTTP error (e.g. malformed JSON) can carry the fallback - # header. If so, halt rather than retrying the FDv2 endpoint. - if fallback: - yield Update( - state=DataSourceState.OFF, - error=error_info, - fallback_to_fdv1=True, - environment_id=envid, - ) - break - - yield Update( - state=DataSourceState.INTERRUPTED, - error=error_info, - environment_id=envid, - ) - else: - (change_set, headers) = result.value - yield Update( - state=DataSourceState.VALID, - change_set=change_set, - environment_id=headers.get(_LD_ENVID_HEADER), - fallback_to_fdv1=headers.get(_LD_FD_FALLBACK_HEADER) == 'true' - ) + decision = map_polling_result(result) + yield decision.update + if decision.control is PollAction.BREAK: + break + if decision.control is PollAction.WAIT_CONTINUE: + self._interrupt_event.wait(self._poll_interval) + continue if self._interrupt_event.wait(self._poll_interval): break @@ -209,44 +131,7 @@ def stop(self): def _poll(self, ss: SelectorStore) -> BasisResult: try: result = self._requester.fetch(ss.selector()) - - if isinstance(result, _Fail): - if isinstance(result.exception, UnsuccessfulResponseException): - status_code = result.exception.status - http_error_message_result = http_error_message( - status_code, "polling request" - ) - if is_http_error_recoverable(status_code): - log.warning(http_error_message_result) - - # Forward any response headers so callers (e.g. FDv2 datasystem) - # can read the X-LD-FD-Fallback directive even on error. - return _Fail( - error=http_error_message_result, - exception=result.exception, - headers=result.headers, - ) - - return _Fail( - error=result.error or "Failed to request payload", - exception=result.exception, - headers=result.headers, - ) - - (change_set, headers) = result.value - - env_id = headers.get(_LD_ENVID_HEADER) - if not isinstance(env_id, str): - env_id = None - - basis = Basis( - change_set=change_set, - persist=change_set.selector.is_defined(), - environment_id=env_id, - fallback_to_fdv1=headers.get(_LD_FD_FALLBACK_HEADER) == 'true', - ) - - return _Success(value=basis) + return polling_result_to_basis(result) except Exception as e: # pylint: disable=broad-except msg = f"Error: Exception encountered when updating flags. {e}" log.exception(msg) @@ -338,63 +223,6 @@ def fetch(self, selector: Optional[Selector]) -> PollingResult: ) -# pylint: disable=too-many-branches,too-many-return-statements -def polling_payload_to_changeset(data: dict) -> _Result[ChangeSet, str]: - """ - Converts a polling payload into a ChangeSet. - """ - if "events" not in data or not isinstance(data["events"], list): - return _Fail(error="Invalid payload: 'events' key is missing or not a list") - - builder = ChangeSetBuilder() - - for event in data["events"]: - if not isinstance(event, dict): - return _Fail(error="Invalid payload: 'events' must be a list of objects") - - if "event" not in event: - continue - - if event["event"] == EventName.SERVER_INTENT: - try: - server_intent = ServerIntent.from_dict(event["data"]) - except ValueError as err: - return _Fail(error="Invalid JSON in server intent", exception=err) - - if server_intent.payload.code == IntentCode.TRANSFER_NONE: - return _Success(ChangeSetBuilder.no_changes()) - - builder.start(server_intent.payload.code) - elif event["event"] == EventName.PUT_OBJECT: - try: - put = PutObject.from_dict(event["data"]) - except ValueError as err: - return _Fail(error="Invalid JSON in put object", exception=err) - - builder.add_put(put.kind, put.key, put.version, put.object) - elif event["event"] == EventName.DELETE_OBJECT: - try: - delete_object = DeleteObject.from_dict(event["data"]) - except ValueError as err: - return _Fail(error="Invalid JSON in delete object", exception=err) - - builder.add_delete( - delete_object.kind, delete_object.key, delete_object.version - ) - elif event["event"] == EventName.PAYLOAD_TRANSFERRED: - try: - selector = Selector.from_dict(event["data"]) - changeset = builder.finish(selector) - - return _Success(value=changeset) - except ValueError as err: - return _Fail( - error="Invalid JSON in payload transferred object", exception=err - ) - - return _Fail(error="didn't receive any known protocol events in polling payload") - - class PollingDataSourceBuilder(DataSourceBuilder): """ Builder for a PollingDataSource. @@ -564,40 +392,3 @@ def fetch(self, selector: Optional[Selector]) -> PollingResult: exception=changeset_result.exception, headers=headers, ) - - -# pylint: disable=too-many-branches,too-many-return-statements -def fdv1_polling_payload_to_changeset(data: dict) -> _Result[ChangeSet, str]: - """ - Converts a fdv1 polling payload into a ChangeSet. - """ - builder = ChangeSetBuilder() - builder.start(IntentCode.TRANSFER_FULL) - selector = Selector.no_selector() - - # FDv1 uses "flags" instead of "features", so we need to map accordingly - # Map FDv1 JSON keys to ObjectKind enum values - kind_mappings = [ - (ObjectKind.FLAG, "flags"), - (ObjectKind.SEGMENT, "segments") - ] - - for kind, fdv1_key in kind_mappings: - kind_data = data.get(fdv1_key) - if kind_data is None: - continue - if not isinstance(kind_data, dict): - return _Fail(error=f"Invalid format: {fdv1_key} is not a dictionary") - - for key in kind_data: - flag_or_segment = kind_data.get(key) - if flag_or_segment is None or not isinstance(flag_or_segment, dict): - return _Fail(error=f"Invalid format: {key} is not a dictionary") - - version = flag_or_segment.get('version') - if version is None: - return _Fail(error=f"Invalid format: {key} does not have a version set") - - builder.add_put(kind, key, version, flag_or_segment) - - return _Success(builder.finish(selector)) diff --git a/ldclient/impl/datasourcev2/polling_common.py b/ldclient/impl/datasourcev2/polling_common.py new file mode 100644 index 00000000..532efd9b --- /dev/null +++ b/ldclient/impl/datasourcev2/polling_common.py @@ -0,0 +1,301 @@ +""" +Shared, transport-agnostic parsers that convert polling payloads into +ChangeSets. Used by both the sync and async polling requesters. +""" + +from dataclasses import dataclass +from enum import Enum +from time import time +from typing import Mapping, Tuple + +from ldclient.impl.datasystem.protocolv2 import ( + DeleteObject, + EventName, + PutObject +) +from ldclient.impl.util import ( + _LD_ENVID_HEADER, + _LD_FD_FALLBACK_HEADER, + UnsuccessfulResponseException, + _Fail, + _Result, + _Success, + http_error_message, + is_http_error_recoverable, + log +) +from ldclient.interfaces import ( + Basis, + BasisResult, + ChangeSet, + ChangeSetBuilder, + DataSourceErrorInfo, + DataSourceErrorKind, + DataSourceState, + IntentCode, + ObjectKind, + Selector, + ServerIntent, + Update +) + + +# pylint: disable=too-many-branches,too-many-return-statements +def polling_payload_to_changeset(data: dict) -> _Result[ChangeSet, str]: + """ + Converts a polling payload into a ChangeSet. + """ + if "events" not in data or not isinstance(data["events"], list): + return _Fail(error="Invalid payload: 'events' key is missing or not a list") + + builder = ChangeSetBuilder() + + for event in data["events"]: + if not isinstance(event, dict): + return _Fail(error="Invalid payload: 'events' must be a list of objects") + + if "event" not in event: + continue + + if event["event"] == EventName.SERVER_INTENT: + try: + server_intent = ServerIntent.from_dict(event["data"]) + except ValueError as err: + return _Fail(error="Invalid JSON in server intent", exception=err) + + if server_intent.payload.code == IntentCode.TRANSFER_NONE: + return _Success(ChangeSetBuilder.no_changes()) + + builder.start(server_intent.payload.code) + elif event["event"] == EventName.PUT_OBJECT: + try: + put = PutObject.from_dict(event["data"]) + except ValueError as err: + return _Fail(error="Invalid JSON in put object", exception=err) + + builder.add_put(put.kind, put.key, put.version, put.object) + elif event["event"] == EventName.DELETE_OBJECT: + try: + delete_object = DeleteObject.from_dict(event["data"]) + except ValueError as err: + return _Fail(error="Invalid JSON in delete object", exception=err) + + builder.add_delete( + delete_object.kind, delete_object.key, delete_object.version + ) + elif event["event"] == EventName.PAYLOAD_TRANSFERRED: + try: + selector = Selector.from_dict(event["data"]) + changeset = builder.finish(selector) + + return _Success(value=changeset) + except ValueError as err: + return _Fail( + error="Invalid JSON in payload transferred object", exception=err + ) + + return _Fail(error="didn't receive any known protocol events in polling payload") + + +# pylint: disable=too-many-branches,too-many-return-statements +def fdv1_polling_payload_to_changeset(data: dict) -> _Result[ChangeSet, str]: + """ + Converts a fdv1 polling payload into a ChangeSet. + """ + builder = ChangeSetBuilder() + builder.start(IntentCode.TRANSFER_FULL) + selector = Selector.no_selector() + + # FDv1 uses "flags" instead of "features", so we need to map accordingly + # Map FDv1 JSON keys to ObjectKind enum values + kind_mappings = [ + (ObjectKind.FLAG, "flags"), + (ObjectKind.SEGMENT, "segments") + ] + + for kind, fdv1_key in kind_mappings: + kind_data = data.get(fdv1_key) + if kind_data is None: + continue + if not isinstance(kind_data, dict): + return _Fail(error=f"Invalid format: {fdv1_key} is not a dictionary") + + for key in kind_data: + flag_or_segment = kind_data.get(key) + if flag_or_segment is None or not isinstance(flag_or_segment, dict): + return _Fail(error=f"Invalid format: {key} is not a dictionary") + + version = flag_or_segment.get('version') + if version is None: + return _Fail(error=f"Invalid format: {key} does not have a version set") + + builder.add_put(kind, key, version, flag_or_segment) + + return _Success(builder.finish(selector)) + + +def polling_result_to_basis( + result: _Result[Tuple[ChangeSet, Mapping], str] +) -> BasisResult: + """ + Convert a requester fetch result into a Basis, or a failure. Used by both + the sync and async polling initializers after they fetch a payload. + """ + if isinstance(result, _Fail): + if isinstance(result.exception, UnsuccessfulResponseException): + status_code = result.exception.status + http_error_message_result = http_error_message( + status_code, "polling request" + ) + if is_http_error_recoverable(status_code): + log.warning(http_error_message_result) + + # Forward any response headers so callers (e.g. FDv2 datasystem) + # can read the X-LD-FD-Fallback directive even on error. + return _Fail( + error=http_error_message_result, + exception=result.exception, + headers=result.headers, + ) + + return _Fail( + error=result.error or "Failed to request payload", + exception=result.exception, + headers=result.headers, + ) + + (change_set, headers) = result.value + + env_id = headers.get(_LD_ENVID_HEADER) + if not isinstance(env_id, str): + env_id = None + + basis = Basis( + change_set=change_set, + persist=change_set.selector.is_defined(), + environment_id=env_id, + fallback_to_fdv1=headers.get(_LD_FD_FALLBACK_HEADER) == 'true', + ) + + return _Success(value=basis) + + +class PollAction(Enum): + """ + How a polling synchronizer's loop should proceed after a fetch. + + ``BREAK`` stops the loop. ``WAIT_CONTINUE`` waits the poll interval and then + starts the next iteration. ``WAIT_BREAK_IF_SET`` waits the poll interval and + stops only when the interrupt event fired during the wait. + """ + BREAK = "break" + WAIT_CONTINUE = "wait_continue" + WAIT_BREAK_IF_SET = "wait_break_if_set" + + +@dataclass +class PollDecision: + """ + The update a polling synchronizer emits after a fetch, and how its loop + should proceed. + """ + update: Update + control: PollAction + + +def map_polling_result( + result: _Result[Tuple[ChangeSet, Mapping], str] +) -> PollDecision: + """ + Convert a requester fetch result into the update a polling synchronizer + emits and the control signal for its loop. Used by both the sync and async + polling synchronizers. + """ + if isinstance(result, _Fail): + fallback = None + envid = None + + if result.headers is not None: + fallback = result.headers.get(_LD_FD_FALLBACK_HEADER) == 'true' + envid = result.headers.get(_LD_ENVID_HEADER) + + if isinstance(result.exception, UnsuccessfulResponseException): + error_info = DataSourceErrorInfo( + kind=DataSourceErrorKind.ERROR_RESPONSE, + status_code=result.exception.status, + time=time(), + message=http_error_message( + result.exception.status, "polling request" + ), + ) + + if fallback: + return PollDecision( + Update( + state=DataSourceState.OFF, + error=error_info, + fallback_to_fdv1=True, + environment_id=envid, + ), + PollAction.BREAK, + ) + + status_code = result.exception.status + if is_http_error_recoverable(status_code): + return PollDecision( + Update( + state=DataSourceState.INTERRUPTED, + error=error_info, + environment_id=envid, + ), + PollAction.WAIT_CONTINUE, + ) + + return PollDecision( + Update( + state=DataSourceState.OFF, + error=error_info, + environment_id=envid, + ), + PollAction.BREAK, + ) + + error_info = DataSourceErrorInfo( + kind=DataSourceErrorKind.NETWORK_ERROR, + time=time(), + status_code=0, + message=result.error, + ) + + # Even a non-HTTP error (e.g. malformed JSON) can carry the fallback + # header. If so, halt rather than retrying the FDv2 endpoint. + if fallback: + return PollDecision( + Update( + state=DataSourceState.OFF, + error=error_info, + fallback_to_fdv1=True, + environment_id=envid, + ), + PollAction.BREAK, + ) + + return PollDecision( + Update( + state=DataSourceState.INTERRUPTED, + error=error_info, + environment_id=envid, + ), + PollAction.WAIT_BREAK_IF_SET, + ) + + (change_set, headers) = result.value + return PollDecision( + Update( + state=DataSourceState.VALID, + change_set=change_set, + environment_id=headers.get(_LD_ENVID_HEADER), + fallback_to_fdv1=headers.get(_LD_FD_FALLBACK_HEADER) == 'true', + ), + PollAction.WAIT_BREAK_IF_SET, + ) diff --git a/ldclient/impl/datasourcev2/streaming.py b/ldclient/impl/datasourcev2/streaming.py index 40217468..16e533c1 100644 --- a/ldclient/impl/datasourcev2/streaming.py +++ b/ldclient/impl/datasourcev2/streaming.py @@ -15,38 +15,26 @@ ErrorStrategy, RetryDelayStrategy ) -from ld_eventsource.errors import HTTPStatusError from ldclient.config import ( DataSourceBuilder, DataSourceBuilderConfig, HTTPConfig ) -from ldclient.impl.datasystem import DiagnosticAccumulator, DiagnosticSource -from ldclient.impl.datasystem.protocolv2 import ( - DeleteObject, - Error, - EventName, - Goodbye, - PutObject +from ldclient.impl.datasourcev2.streaming_common import ( + classify_stream_error, + process_message, + with_fallback_signal ) +from ldclient.impl.datasystem import DiagnosticAccumulator, DiagnosticSource from ldclient.impl.http import HTTPFactory, _base_headers -from ldclient.impl.util import ( - _LD_ENVID_HEADER, - _LD_FD_FALLBACK_HEADER, - http_error_message, - is_http_error_recoverable, - log -) +from ldclient.impl.util import _LD_ENVID_HEADER, _LD_FD_FALLBACK_HEADER, log from ldclient.interfaces import ( ChangeSetBuilder, DataSourceErrorInfo, DataSourceErrorKind, DataSourceState, - IntentCode, - Selector, SelectorStore, - ServerIntent, Synchronizer, Update ) @@ -182,19 +170,6 @@ def sync(self, ss: SelectorStore) -> Generator[Update, None, None]: # itself doesn't see the directive header. fallback_requested = False - def _with_fallback_signal(update: Update) -> Update: - """Return ``update`` decorated with ``fallback_to_fdv1=True`` when - the directive has been latched. Idempotent if already set.""" - if not fallback_requested or update.fallback_to_fdv1: - return update - return Update( - state=update.state, - change_set=update.change_set, - error=update.error, - fallback_to_fdv1=True, - environment_id=update.environment_id, - ) - for action in self._sse.all: if isinstance(action, Fault): # If the SSE client detects the stream has closed, then it will @@ -208,7 +183,7 @@ def _with_fallback_signal(update: Update) -> Update: (update, should_continue) = self._handle_error(action.error, envid) if update is not None: - yield _with_fallback_signal(update) + yield with_fallback_signal(update, fallback_requested) # The FDv1 Fallback Directive is one-way and terminal: if it # was latched on a prior Start, we must not keep retrying the @@ -226,7 +201,7 @@ def _with_fallback_signal(update: Update) -> Update: continue try: - update = self._process_message(action, change_set_builder, envid) + update = process_message(action, change_set_builder, envid) if update is not None: self._record_stream_init(False) self._connection_attempt_start_time = None @@ -234,7 +209,7 @@ def _with_fallback_signal(update: Update) -> Update: # The completed update is the natural moment to honor # the latched directive: yield once with the signal, # then halt — the consumer will switch to FDv1. - yield _with_fallback_signal(update) + yield with_fallback_signal(update, fallback_requested) break yield update except json.decoder.JSONDecodeError as e: @@ -245,7 +220,7 @@ def _with_fallback_signal(update: Update) -> Update: (update, should_continue) = self._handle_error(e, envid) if update is not None: - yield _with_fallback_signal(update) + yield with_fallback_signal(update, fallback_requested) if fallback_requested or not should_continue: break except Exception as e: # pylint: disable=broad-except @@ -254,14 +229,14 @@ def _with_fallback_signal(update: Update) -> Update: ) self._sse.interrupt() - yield _with_fallback_signal(Update( + yield with_fallback_signal(Update( state=DataSourceState.INTERRUPTED, error=DataSourceErrorInfo( DataSourceErrorKind.UNKNOWN, 0, time(), str(e) ), fallback_to_fdv1=False, environment_id=envid, - )) + ), fallback_requested) if fallback_requested: break @@ -282,76 +257,6 @@ def _record_stream_init(self, failed: bool): elapsed = current_time - int(self._connection_attempt_start_time * 1000) self._diagnostic_accumulator.record_stream_init(current_time, elapsed if elapsed >= 0 else 0, failed) - # pylint: disable=too-many-return-statements - def _process_message( - self, msg: Event, change_set_builder: ChangeSetBuilder, envid: Optional[str] - ) -> Optional[Update]: - """ - Processes a single message from the SSE stream and returns an Update - object if applicable. - - This method may raise exceptions if the message is malformed or if an - error occurs while processing the message. The caller should handle these - exceptions appropriately. - """ - if msg.event == EventName.HEARTBEAT: - return None - - if msg.event == EventName.SERVER_INTENT: - server_intent = ServerIntent.from_dict(json.loads(msg.data)) - change_set_builder.start(server_intent.payload.code) - - if server_intent.payload.code == IntentCode.TRANSFER_NONE: - change_set_builder.expect_changes() - return Update( - state=DataSourceState.VALID, - environment_id=envid, - ) - return None - - if msg.event == EventName.PUT_OBJECT: - put = PutObject.from_dict(json.loads(msg.data)) - change_set_builder.add_put(put.kind, put.key, put.version, put.object) - return None - - if msg.event == EventName.DELETE_OBJECT: - delete = DeleteObject.from_dict(json.loads(msg.data)) - change_set_builder.add_delete(delete.kind, delete.key, delete.version) - return None - - if msg.event == EventName.GOODBYE: - goodbye = Goodbye.from_dict(json.loads(msg.data)) - log.info("SSE server sent goodbye: %s", goodbye.reason) - - return None - - if msg.event == EventName.ERROR: - error = Error.from_dict(json.loads(msg.data)) - log.error("Error on %s: %s", error.payload_id, error.reason) - - # The protocol should "reset" any previous change events it has - # received, but should continue to operate under the assumption the - # last server intent was in effect. - # - # The server may choose to send a new server-intent, at which point - # we will set that as well. - change_set_builder.reset() - - return None - - if msg.event == EventName.PAYLOAD_TRANSFERRED: - selector = Selector.from_dict(json.loads(msg.data)) - change_set = change_set_builder.finish(selector) - - return Update( - state=DataSourceState.VALID, - change_set=change_set, - environment_id=envid, - ) - - log.info("Unexpected event found in stream: %s", msg.event) - return None - def _handle_error(self, error: Exception, envid: Optional[str]) -> Tuple[Optional[Update], bool]: """ This method handles errors that occur during the streaming process. @@ -367,89 +272,15 @@ def _handle_error(self, error: Exception, envid: Optional[str]) -> Tuple[Optiona if not self._running: return (None, False) # don't retry if we've been deliberately stopped - update: Optional[Update] = None - - if isinstance(error, json.decoder.JSONDecodeError): - log.error("Unexpected error on stream connection: %s, will retry", error) - self._record_stream_init(True) - self._connection_attempt_start_time = time() + \ - self._sse.next_retry_delay # type: ignore - - update = Update( - state=DataSourceState.INTERRUPTED, - error=DataSourceErrorInfo( - DataSourceErrorKind.INVALID_DATA, 0, time(), str(error) - ), - fallback_to_fdv1=False, - environment_id=envid, - ) - return (update, True) - - if isinstance(error, HTTPStatusError): - self._record_stream_init(True) - self._connection_attempt_start_time = time() + \ - self._sse.next_retry_delay # type: ignore - - error_info = DataSourceErrorInfo( - DataSourceErrorKind.ERROR_RESPONSE, - error.status, - time(), - str(error), - ) - - if envid is None and error.headers is not None: - envid = error.headers.get(_LD_ENVID_HEADER) - - if error.headers is not None and error.headers.get(_LD_FD_FALLBACK_HEADER) == 'true': - update = Update( - state=DataSourceState.OFF, - error=error_info, - fallback_to_fdv1=True, - environment_id=envid, - ) - self.stop() - return (update, False) - - http_error_message_result = http_error_message( - error.status, "stream connection" - ) - is_recoverable = is_http_error_recoverable(error.status) - update = Update( - state=( - DataSourceState.INTERRUPTED - if is_recoverable - else DataSourceState.OFF - ), - error=error_info, - fallback_to_fdv1=False, - environment_id=envid, - ) - - if not is_recoverable: - self._connection_attempt_start_time = None - log.error(http_error_message_result) - self.stop() - return (update, False) - - log.warning(http_error_message_result) - return (update, True) - - log.warning("Unexpected error on stream connection: %s, will retry", error) - self._record_stream_init(True) - self._connection_attempt_start_time = time() + self._sse.next_retry_delay # type: ignore - - update = Update( - state=DataSourceState.INTERRUPTED, - error=DataSourceErrorInfo( - DataSourceErrorKind.UNKNOWN, 0, time(), str(error) - ), - fallback_to_fdv1=False, - environment_id=envid, + decision = classify_stream_error( + error, self._sse.next_retry_delay, envid # type: ignore ) - # no stacktrace here because, for a typical connection error, it'll - # just be a lengthy tour of urllib3 internals + self._record_stream_init(True) + self._connection_attempt_start_time = decision.next_attempt_start + if decision.should_stop: + self.stop() - return (update, True) + return (decision.update, decision.should_continue) class StreamingDataSourceBuilder(DataSourceBuilder): diff --git a/ldclient/impl/datasourcev2/streaming_common.py b/ldclient/impl/datasourcev2/streaming_common.py new file mode 100644 index 00000000..cb0428d8 --- /dev/null +++ b/ldclient/impl/datasourcev2/streaming_common.py @@ -0,0 +1,248 @@ +""" +Shared, transport-agnostic parser for FDv2 streaming messages. Used by both +the sync and async streaming synchronizers. +""" + +import json +from dataclasses import dataclass +from time import time +from typing import Optional + +from ld_eventsource.actions import Event +from ld_eventsource.errors import HTTPStatusError + +from ldclient.impl.datasystem.protocolv2 import ( + DeleteObject, + Error, + EventName, + Goodbye, + PutObject +) +from ldclient.impl.util import ( + _LD_ENVID_HEADER, + _LD_FD_FALLBACK_HEADER, + http_error_message, + is_http_error_recoverable, + log +) +from ldclient.interfaces import ( + ChangeSetBuilder, + DataSourceErrorInfo, + DataSourceErrorKind, + DataSourceState, + IntentCode, + Selector, + ServerIntent, + Update +) + + +# pylint: disable=too-many-return-statements +def process_message( + msg: Event, change_set_builder: ChangeSetBuilder, envid: Optional[str] +) -> Optional[Update]: + """ + Processes a single message from the SSE stream and returns an Update + object if applicable. + + This function may raise exceptions if the message is malformed or if an + error occurs while processing the message. The caller should handle these + exceptions appropriately. + """ + if msg.event == EventName.HEARTBEAT: + return None + + if msg.event == EventName.SERVER_INTENT: + server_intent = ServerIntent.from_dict(json.loads(msg.data)) + change_set_builder.start(server_intent.payload.code) + + if server_intent.payload.code == IntentCode.TRANSFER_NONE: + change_set_builder.expect_changes() + return Update( + state=DataSourceState.VALID, + environment_id=envid, + ) + return None + + if msg.event == EventName.PUT_OBJECT: + put = PutObject.from_dict(json.loads(msg.data)) + change_set_builder.add_put(put.kind, put.key, put.version, put.object) + return None + + if msg.event == EventName.DELETE_OBJECT: + delete = DeleteObject.from_dict(json.loads(msg.data)) + change_set_builder.add_delete(delete.kind, delete.key, delete.version) + return None + + if msg.event == EventName.GOODBYE: + goodbye = Goodbye.from_dict(json.loads(msg.data)) + log.info("SSE server sent goodbye: %s", goodbye.reason) + + return None + + if msg.event == EventName.ERROR: + error = Error.from_dict(json.loads(msg.data)) + log.error("Error on %s: %s", error.payload_id, error.reason) + + # The protocol should "reset" any previous change events it has + # received, but should continue to operate under the assumption the + # last server intent was in effect. + # + # The server may choose to send a new server-intent, at which point + # we will set that as well. + change_set_builder.reset() + + return None + + if msg.event == EventName.PAYLOAD_TRANSFERRED: + selector = Selector.from_dict(json.loads(msg.data)) + change_set = change_set_builder.finish(selector) + + return Update( + state=DataSourceState.VALID, + change_set=change_set, + environment_id=envid, + ) + + log.info("Unexpected event found in stream: %s", msg.event) + return None + + +def with_fallback_signal(update: Update, fallback_requested: bool) -> Update: + """ + Return ``update`` marked with ``fallback_to_fdv1=True`` when the FDv1 + fallback directive has been latched. Returns the update unchanged when the + directive is not set or the update already carries it. + """ + if not fallback_requested or update.fallback_to_fdv1: + return update + return Update( + state=update.state, + change_set=update.change_set, + error=update.error, + fallback_to_fdv1=True, + environment_id=update.environment_id, + ) + + +@dataclass +class StreamErrorDecision: + """ + Describes how a streaming synchronizer should react to a connection error. + + The values are computed without any I/O so both the sync and async + synchronizers share one decision, then apply the side effects themselves: + they record a failed stream init, set the connection attempt start time to + ``next_attempt_start``, stop the synchronizer when ``should_stop`` is set, + and retry the connection when ``should_continue`` is set. + """ + update: Update + should_continue: bool + should_stop: bool + next_attempt_start: Optional[float] + + +def classify_stream_error( + error: Exception, next_retry_delay: float, envid: Optional[str] +) -> StreamErrorDecision: + """ + Decide how to react to a streaming connection error. + + The caller must first check that the synchronizer is still running; this + function assumes it is. It returns a decision the caller applies. The caller + always records a failed stream init before it sets the connection attempt + start time to the returned ``next_attempt_start``. + """ + if isinstance(error, json.decoder.JSONDecodeError): + log.error("Unexpected error on stream connection: %s, will retry", error) + update = Update( + state=DataSourceState.INTERRUPTED, + error=DataSourceErrorInfo( + DataSourceErrorKind.INVALID_DATA, 0, time(), str(error) + ), + fallback_to_fdv1=False, + environment_id=envid, + ) + return StreamErrorDecision( + update=update, + should_continue=True, + should_stop=False, + next_attempt_start=time() + next_retry_delay, + ) + + if isinstance(error, HTTPStatusError): + next_attempt_start: Optional[float] = time() + next_retry_delay + + error_info = DataSourceErrorInfo( + DataSourceErrorKind.ERROR_RESPONSE, + error.status, + time(), + str(error), + ) + + if envid is None and error.headers is not None: + envid = error.headers.get(_LD_ENVID_HEADER) + + if error.headers is not None and error.headers.get(_LD_FD_FALLBACK_HEADER) == 'true': + update = Update( + state=DataSourceState.OFF, + error=error_info, + fallback_to_fdv1=True, + environment_id=envid, + ) + return StreamErrorDecision( + update=update, + should_continue=False, + should_stop=True, + next_attempt_start=next_attempt_start, + ) + + http_error_message_result = http_error_message( + error.status, "stream connection" + ) + is_recoverable = is_http_error_recoverable(error.status) + update = Update( + state=( + DataSourceState.INTERRUPTED + if is_recoverable + else DataSourceState.OFF + ), + error=error_info, + fallback_to_fdv1=False, + environment_id=envid, + ) + + if not is_recoverable: + log.error(http_error_message_result) + return StreamErrorDecision( + update=update, + should_continue=False, + should_stop=True, + next_attempt_start=None, + ) + + log.warning(http_error_message_result) + return StreamErrorDecision( + update=update, + should_continue=True, + should_stop=False, + next_attempt_start=next_attempt_start, + ) + + log.warning("Unexpected error on stream connection: %s, will retry", error) + update = Update( + state=DataSourceState.INTERRUPTED, + error=DataSourceErrorInfo( + DataSourceErrorKind.UNKNOWN, 0, time(), str(error) + ), + fallback_to_fdv1=False, + environment_id=envid, + ) + # no stacktrace here because, for a typical connection error, it'll + # just be a lengthy tour of the HTTP client internals + return StreamErrorDecision( + update=update, + should_continue=True, + should_stop=False, + next_attempt_start=time() + next_retry_delay, + ) diff --git a/ldclient/impl/integrations/test_datav2/async_test_data_sourcev2.py b/ldclient/impl/integrations/test_datav2/async_test_data_sourcev2.py new file mode 100644 index 00000000..569c82dd --- /dev/null +++ b/ldclient/impl/integrations/test_datav2/async_test_data_sourcev2.py @@ -0,0 +1,163 @@ +import asyncio +from typing import AsyncGenerator, Optional + +from ldclient.impl.util import _Fail, _Success, current_time_millis +from ldclient.interfaces import ( + Basis, + BasisResult, + ChangeSetBuilder, + DataSourceErrorInfo, + DataSourceErrorKind, + DataSourceState, + IntentCode, + ObjectKind, + Selector, + SelectorStore, + Update +) + + +class _AsyncTestDataSourceV2: + """ + Async implementation of both the Initializer and Synchronizer protocols for TestDataV2. + + The async twin of :class:`_TestDataSourceV2`: it shares the same TestDataV2 flag + management but exposes ``async`` ``fetch``/``sync``/``stop`` and delivers updates + through an :class:`asyncio.Queue` so it can drive the async data system. + """ + + def __init__(self, test_data): + self._test_data = test_data + self._closed = False + self._update_queue: asyncio.Queue = asyncio.Queue() + # The event loop that sync() runs on. Captured when sync() starts so + # cross-thread updates can be scheduled onto it safely. + self._loop: Optional[asyncio.AbstractEventLoop] = None + + # Register for change notifications; upsert_flag is invoked on updates. + self._test_data._add_instance(self) + + @property + def name(self) -> str: + """Return the name of this data source.""" + return "TestDataV2" + + async def fetch(self, ss: SelectorStore) -> BasisResult: + """Implementation of the AsyncInitializer.fetch method.""" + return self._make_basis() + + async def sync(self, ss: SelectorStore) -> AsyncGenerator[Update, None]: + """Implementation of the AsyncSynchronizer.sync method: yields the initial + data, then each update as it is queued, until the source is stopped.""" + self._loop = asyncio.get_running_loop() + + initial_result = self._make_basis() + if isinstance(initial_result, _Fail): + yield Update( + state=DataSourceState.OFF, + error=DataSourceErrorInfo( + kind=DataSourceErrorKind.STORE_ERROR, + status_code=0, + time=current_time_millis(), + message=initial_result.error, + ), + ) + return + + yield Update( + state=DataSourceState.VALID, change_set=initial_result.value.change_set + ) + + while not self._closed: + update = await self._update_queue.get() + if update is None: # Sentinel value for shutdown + break + yield update + + async def stop(self): + """Stop the data source and clean up resources.""" + if self._closed: + return + self._closed = True + self._test_data._closed_instance(self) + # Wake the sync generator so it can exit. + self._enqueue(None) + + def _enqueue(self, item) -> None: + """Put an item on the update queue safely from any thread. + + Flag updates can arrive on any thread (TestDataV2.update may be called + from user code), but an asyncio.Queue must only be touched from its own + event loop. So schedule the put on the loop captured when sync() started. + Before sync() runs there is no consumer and no loop to cross into, so put + directly. If the loop has since closed, drop the item.""" + loop = self._loop + if loop is None: + self._update_queue.put_nowait(item) + return + try: + loop.call_soon_threadsafe(self._update_queue.put_nowait, item) + except RuntimeError: + # The event loop has been closed; there is nothing left to deliver to. + pass + + def upsert_flag(self, flag_data: dict): + """Called by TestDataV2 when a flag is updated; queues the change for + delivery through the sync() generator.""" + if self._closed: + return + try: + version = self._test_data._get_version() + + builder = ChangeSetBuilder() + builder.start(IntentCode.TRANSFER_CHANGES) + builder.add_put( + ObjectKind.FLAG, + flag_data["key"], + flag_data.get("version", 1), + flag_data, + ) + + selector = Selector.new_selector(str(version), version) + change_set = builder.finish(selector) + + self._enqueue( + Update(state=DataSourceState.VALID, change_set=change_set) + ) + except Exception as e: + self._enqueue( + Update( + state=DataSourceState.OFF, + error=DataSourceErrorInfo( + kind=DataSourceErrorKind.STORE_ERROR, + status_code=0, + time=current_time_millis(), + message=f"Error processing flag update: {str(e)}", + ), + ) + ) + + def _make_basis(self) -> BasisResult: + """Builds a full-transfer Basis from the current test data. Shared by + fetch() and the initial yield of sync().""" + try: + if self._closed: + return _Fail("TestDataV2 source has been closed") + + init_data = self._test_data._make_init_data() + version = self._test_data._get_version() + + builder = ChangeSetBuilder() + builder.start(IntentCode.TRANSFER_FULL) + for key, flag_data in init_data.items(): + builder.add_put( + ObjectKind.FLAG, key, flag_data.get("version", 1), flag_data + ) + + selector = Selector.new_selector(str(version), version) + change_set = builder.finish(selector) + basis = Basis(change_set=change_set, persist=False, environment_id=None) + + return _Success(basis) + except Exception as e: + return _Fail(f"Error fetching test data: {str(e)}") diff --git a/ldclient/integrations/test_datav2.py b/ldclient/integrations/test_datav2.py index 954f2aeb..5c13dc45 100644 --- a/ldclient/integrations/test_datav2.py +++ b/ldclient/integrations/test_datav2.py @@ -5,6 +5,9 @@ from ldclient.config import Config, DataSourceBuilder, DataSourceBuilderConfig from ldclient.context import Context +from ldclient.impl.integrations.test_datav2.async_test_data_sourcev2 import ( + _AsyncTestDataSourceV2 +) from ldclient.impl.integrations.test_datav2.test_data_sourcev2 import ( _TestDataSourceV2 ) @@ -682,6 +685,20 @@ def builder(self) -> DataSourceBuilder: """ return TestDataSourceBuilder(self) + @property + def async_builder(self) -> DataSourceBuilder: + """ + Creates a builder that can be used with the async FDv2 data system. + + .. caution:: + This feature is experimental and should NOT be considered ready for production + use. It may change or be removed without notice and is not subject to backwards + compatibility guarantees. + + :return: an async test data data source builder + """ + return AsyncTestDataSourceBuilder(self) + class TestDataSourceBuilder(DataSourceBuilder[_TestDataSourceV2]): # pylint: disable=too-few-public-methods """Builder for TestDataV2 data sources that implements the DataSourceBuilder protocol.""" @@ -692,3 +709,21 @@ def __init__(self, test_data: TestDataV2): def build(self, config: DataSourceBuilderConfig) -> _TestDataSourceV2: # pylint: disable=unused-argument """Builds the TestDataSourceV2 instance.""" return _TestDataSourceV2(self._test_data) + + +class AsyncTestDataSourceBuilder(DataSourceBuilder[_AsyncTestDataSourceV2]): # pylint: disable=too-few-public-methods + """ + Builder for async TestDataV2 data sources that implements the DataSourceBuilder protocol. + + .. caution:: + This feature is experimental and should NOT be considered ready for production + use. It may change or be removed without notice and is not subject to backwards + compatibility guarantees. + """ + + def __init__(self, test_data: TestDataV2): + self._test_data = test_data + + def build(self, config: DataSourceBuilderConfig) -> _AsyncTestDataSourceV2: # pylint: disable=unused-argument + """Builds the async TestDataSourceV2 instance.""" + return _AsyncTestDataSourceV2(self._test_data) diff --git a/ldclient/testing/impl/datasourcev2/test_async_polling.py b/ldclient/testing/impl/datasourcev2/test_async_polling.py new file mode 100644 index 00000000..181d397c --- /dev/null +++ b/ldclient/testing/impl/datasourcev2/test_async_polling.py @@ -0,0 +1,316 @@ +# pylint: disable=missing-docstring + +import asyncio +import time +from typing import List, Optional + +import pytest + +from ldclient.impl.datasourcev2.async_polling import ( + AsyncPollingDataSource, + PollingResult +) +from ldclient.impl.util import ( + _LD_ENVID_HEADER, + _LD_FD_FALLBACK_HEADER, + UnsuccessfulResponseException, + _Fail, + _Success +) +from ldclient.interfaces import ( + ChangeSetBuilder, + DataSourceErrorKind, + DataSourceState, + IntentCode, + ObjectKind, + Selector +) +from ldclient.testing.mock_components import MockSelectorStore + + +class MockPollingRequester: # pylint: disable=too-few-public-methods + def __init__(self, results: List[PollingResult]): + self._results = list(results) + self._index = 0 + self.call_times: List[float] = [] + + async def fetch(self, selector: Optional[Selector]) -> PollingResult: + self.call_times.append(time.monotonic()) + result = self._results[self._index % len(self._results)] + self._index += 1 + return result + + async def close(self) -> None: + pass + + +class MockExceptionThrowingRequester: # pylint: disable=too-few-public-methods + async def fetch(self, selector: Optional[Selector]) -> PollingResult: + raise RuntimeError("requester blew up") + + async def close(self) -> None: + pass + + +def _valid_change_set(): + builder = ChangeSetBuilder() + builder.start(IntentCode.TRANSFER_FULL) + builder.add_put(ObjectKind.FLAG, "flag-key", 100, {"key": "flag-key"}) + return builder.finish(Selector(state="p:SOMETHING:300", version=300)) + + +def _make_source(results: List[PollingResult], poll_interval: float = 0.01) -> AsyncPollingDataSource: + return AsyncPollingDataSource( + poll_interval=poll_interval, + requester=MockPollingRequester(results), + ) + + +def _ss() -> MockSelectorStore: + return MockSelectorStore(Selector.no_selector()) + + +def test_name(): + src = _make_source([_Fail(error="failure message")]) + assert src.name == "PollingDataSourceV2" + + +@pytest.mark.asyncio +async def test_fetch_success(): + change_set = _valid_change_set() + src = _make_source([_Success(value=(change_set, {_LD_ENVID_HEADER: "env1"}))]) + + result = await src.fetch(_ss()) + assert isinstance(result, _Success) + basis = result.value + assert basis.change_set is change_set + assert basis.persist is True + assert basis.environment_id == "env1" + assert basis.fallback_to_fdv1 is False + + +@pytest.mark.asyncio +async def test_fetch_failure_passes_through(): + src = _make_source([_Fail(error="failure message")]) + + result = await src.fetch(_ss()) + assert isinstance(result, _Fail) + assert result.error == "failure message" + + +@pytest.mark.asyncio +async def test_fetch_recoverable_error(): + src = _make_source([_Fail(error="500", exception=UnsuccessfulResponseException(500))]) + + result = await src.fetch(_ss()) + assert isinstance(result, _Fail) + assert result.error.startswith("Received HTTP error 500") + + +@pytest.mark.asyncio +async def test_fetch_unrecoverable_error(): + src = _make_source([_Fail(error="401", exception=UnsuccessfulResponseException(401))]) + + result = await src.fetch(_ss()) + assert isinstance(result, _Fail) + assert result.error.startswith("Received HTTP error 401") + + +@pytest.mark.asyncio +async def test_fetch_no_changes(): + src = _make_source([_Success(value=(ChangeSetBuilder.no_changes(), {}))]) + + result = await src.fetch(_ss()) + assert isinstance(result, _Success) + basis = result.value + assert basis.persist is False + assert not basis.change_set.selector.is_defined() + + +@pytest.mark.asyncio +async def test_fetch_requester_exception_is_caught(): + src = AsyncPollingDataSource(poll_interval=0.01, requester=MockExceptionThrowingRequester()) + + result = await src.fetch(_ss()) + assert isinstance(result, _Fail) + assert "Exception encountered when updating flags" in result.error + + +@pytest.mark.asyncio +async def test_sync_yields_valid_update(): + change_set = _valid_change_set() + src = _make_source([_Success(value=(change_set, {_LD_ENVID_HEADER: "env1"}))], poll_interval=60) + + gen = src.sync(_ss()) + update = await gen.__anext__() + await gen.aclose() + + assert update.state == DataSourceState.VALID + assert update.change_set is change_set + assert update.environment_id == "env1" + assert update.fallback_to_fdv1 is False + + +@pytest.mark.asyncio +async def test_sync_yields_interrupted_on_recoverable_error(): + src = _make_source([_Fail(error="500", exception=UnsuccessfulResponseException(500))], poll_interval=60) + + gen = src.sync(_ss()) + update = await gen.__anext__() + await gen.aclose() + + assert update.state == DataSourceState.INTERRUPTED + assert update.error is not None + assert update.error.kind == DataSourceErrorKind.ERROR_RESPONSE + assert update.error.status_code == 500 + + +@pytest.mark.asyncio +async def test_sync_yields_off_on_unrecoverable_error(): + src = _make_source([_Fail(error="401", exception=UnsuccessfulResponseException(401))], poll_interval=60) + + updates = [update async for update in src.sync(_ss())] + + assert len(updates) == 1 + assert updates[0].state == DataSourceState.OFF + + +@pytest.mark.asyncio +async def test_sync_network_error_yields_interrupted(): + call_count = 0 + + class _StoppingRequester: + async def fetch(self, selector): + nonlocal call_count + call_count += 1 + if call_count >= 2: + await src.stop() + return _Fail(error="connection refused") + + async def close(self) -> None: + pass + + src = AsyncPollingDataSource(poll_interval=0.01, requester=_StoppingRequester()) + + updates = [update async for update in src.sync(_ss())] + + assert len(updates) >= 1 + assert updates[0].state == DataSourceState.INTERRUPTED + assert updates[0].error.kind == DataSourceErrorKind.NETWORK_ERROR + + +@pytest.mark.asyncio +async def test_sync_fallback_to_fdv1_on_error(): + headers = {_LD_FD_FALLBACK_HEADER: 'true', _LD_ENVID_HEADER: 'env1'} + src = _make_source( + [_Fail(error="403", exception=UnsuccessfulResponseException(403), headers=headers)], + poll_interval=60, + ) + + updates = [update async for update in src.sync(_ss())] + + assert len(updates) == 1 + assert updates[0].state == DataSourceState.OFF + assert updates[0].fallback_to_fdv1 is True + assert updates[0].environment_id == 'env1' + + +@pytest.mark.asyncio +async def test_sync_fallback_to_fdv1_on_success(): + change_set = _valid_change_set() + headers = {_LD_FD_FALLBACK_HEADER: 'true'} + src = _make_source([_Success(value=(change_set, headers))], poll_interval=60) + + gen = src.sync(_ss()) + update = await gen.__anext__() + await gen.aclose() + + assert update.state == DataSourceState.VALID + assert update.fallback_to_fdv1 is True + + +@pytest.mark.asyncio +async def test_sync_interval_is_respected(): + requester = MockPollingRequester([_Success(value=(_valid_change_set(), {}))]) + src = AsyncPollingDataSource(poll_interval=0.1, requester=requester) + + updates = [] + async for update in src.sync(_ss()): + updates.append(update) + if len(updates) >= 2: + await src.stop() + + assert len(requester.call_times) >= 2 + elapsed = requester.call_times[1] - requester.call_times[0] + assert elapsed >= 0.05, f"Poll interval too short: {elapsed:.3f}s" + + +@pytest.mark.asyncio +async def test_stop_halts_sync(): + src = _make_source([_Success(value=(_valid_change_set(), {}))], poll_interval=60) + + updates = [] + first_update_received = asyncio.Event() + + async def consume(): + async for update in src.sync(_ss()): + updates.append(update) + first_update_received.set() + + task = asyncio.create_task(consume()) + await asyncio.wait_for(first_update_received.wait(), timeout=2) + await src.stop() + await asyncio.wait_for(task, timeout=2) + + assert len(updates) == 1 + assert updates[0].state == DataSourceState.VALID + + +@pytest.mark.asyncio +async def test_sync_recoverable_error_continues_polling(): + change_set = _valid_change_set() + src = _make_source( + [ + _Fail(error="500", exception=UnsuccessfulResponseException(500)), + _Success(value=(change_set, {})), + ], + poll_interval=0.01, + ) + + gen = src.sync(_ss()) + first = await asyncio.wait_for(gen.__anext__(), timeout=2) + second = await asyncio.wait_for(gen.__anext__(), timeout=2) + await gen.aclose() + + assert first.state == DataSourceState.INTERRUPTED + assert first.error is not None + assert first.error.status_code == 500 + assert second.state == DataSourceState.VALID + assert second.change_set is change_set + + +@pytest.mark.asyncio +async def test_requester_is_closed_after_stop(): + class CloseTrackingRequester(MockPollingRequester): + def __init__(self, results): + super().__init__(results) + self.closed = False + + async def close(self) -> None: + self.closed = True + + requester = CloseTrackingRequester([_Success(value=(_valid_change_set(), {}))]) + src = AsyncPollingDataSource(poll_interval=60, requester=requester) + + first_update_received = asyncio.Event() + + async def consume(): + async for _ in src.sync(_ss()): + first_update_received.set() + + task = asyncio.create_task(consume()) + await asyncio.wait_for(first_update_received.wait(), timeout=2) + await src.stop() + await asyncio.wait_for(task, timeout=2) + + assert requester.closed is True diff --git a/ldclient/testing/impl/datasourcev2/test_async_streaming.py b/ldclient/testing/impl/datasourcev2/test_async_streaming.py new file mode 100644 index 00000000..b3c55f02 --- /dev/null +++ b/ldclient/testing/impl/datasourcev2/test_async_streaming.py @@ -0,0 +1,551 @@ +# pylint: disable=missing-docstring, too-few-public-methods + +import json +from typing import AsyncIterable, List, Optional +from unittest import mock + +import pytest +from ld_eventsource.actions import Start +from ld_eventsource.http import HTTPStatusError +from ld_eventsource.sse_client import Event, Fault + +from ldclient.config import Config +from ldclient.impl.datasourcev2 import async_streaming +from ldclient.impl.datasourcev2.async_streaming import ( + STREAMING_ENDPOINT, + AsyncStreamingDataSource, + create_sse_client +) +from ldclient.impl.datasystem.protocolv2 import ( + DeleteObject, + Error, + EventName, + Goodbye, + PutObject +) +from ldclient.impl.util import _LD_ENVID_HEADER, _LD_FD_FALLBACK_HEADER +from ldclient.interfaces import ( + ChangeType, + DataSourceErrorKind, + DataSourceState, + IntentCode, + ObjectKind, + Payload, + Selector, + ServerIntent +) +from ldclient.testing.mock_components import MockSelectorStore + + +class MockAsyncSSEClient: + """An async SSE client backed by a static list of actions.""" + + def __init__(self, actions: List): + self._actions = actions + self.interrupted = False + self.closed = False + self.next_retry_delay = 0.1 + + @property + def all(self) -> AsyncIterable: + return self._async_gen() + + async def _async_gen(self): + for action in self._actions: + if self.interrupted or self.closed: + return + yield action + + async def interrupt(self): + self.interrupted = True + + async def close(self): + self.closed = True + + +def list_sse_client(actions: List): + """Returns an SseClientBuilder producing a MockAsyncSSEClient.""" + + def builder(base_uri, http_options, initial_reconnect_delay, config, ss, session=None): # pylint: disable=unused-argument + return MockAsyncSSEClient(actions), None + + return builder + + +def make_streaming_data_source() -> AsyncStreamingDataSource: + config = Config("key") + return AsyncStreamingDataSource( + config.stream_base_uri + STREAMING_ENDPOINT, + config.http, + config.initial_reconnect_delay, + config, + ) + + +def server_intent_event(code: IntentCode) -> Event: + si = ServerIntent(payload=Payload(id="p1", target=1, code=code, reason="test")) + return Event(event=EventName.SERVER_INTENT, data=json.dumps(si.to_dict())) + + +def put_object_event(key: str = "my-flag", version: int = 1) -> Event: + put = PutObject(version=version, kind=ObjectKind.FLAG, key=key, object={"key": key, "version": version}) + return Event(event=EventName.PUT_OBJECT, data=json.dumps(put.to_dict())) + + +def payload_transferred_event(version: int = 1) -> Event: + sel = Selector(state=f"p:test:{version}", version=version) + return Event(event=EventName.PAYLOAD_TRANSFERRED, data=json.dumps(sel.to_dict())) + + +def delete_object_event(key: str = "my-flag", version: int = 2) -> Event: + d = DeleteObject(version=version, kind=ObjectKind.FLAG, key=key) + return Event(event=EventName.DELETE_OBJECT, data=json.dumps(d.to_dict())) + + +async def collect_updates(src: AsyncStreamingDataSource, actions: List, ss=None): + """Drive sync() with the given mock actions, collecting yielded updates.""" + if ss is None: + ss = MockSelectorStore(Selector.no_selector()) + + src._sse_client_builder = list_sse_client(actions) + + return [update async for update in src.sync(ss)] + + +@pytest.mark.asyncio +async def test_full_transfer(): + src = make_streaming_data_source() + + actions = [ + server_intent_event(IntentCode.TRANSFER_FULL), + put_object_event("flag-1"), + payload_transferred_event(), + ] + + updates = await collect_updates(src, actions) + assert len(updates) == 1 + update = updates[0] + assert update.state == DataSourceState.VALID + assert update.change_set is not None + changes = [c for c in update.change_set.changes if c.action == ChangeType.PUT] + assert any(c.key == "flag-1" for c in changes) + + +@pytest.mark.asyncio +async def test_transfer_none(): + src = make_streaming_data_source() + + si = ServerIntent(payload=Payload(id="p1", target=1, code=IntentCode.TRANSFER_NONE, reason="up-to-date")) + actions = [ + Event(event=EventName.SERVER_INTENT, data=json.dumps(si.to_dict())), + ] + + updates = await collect_updates(src, actions) + assert len(updates) == 1 + assert updates[0].state == DataSourceState.VALID + + +@pytest.mark.asyncio +async def test_heartbeat_is_ignored(): + src = make_streaming_data_source() + + actions = [ + Event(event=EventName.HEARTBEAT), + server_intent_event(IntentCode.TRANSFER_FULL), + put_object_event("flag-1"), + payload_transferred_event(), + ] + + updates = await collect_updates(src, actions) + assert len(updates) == 1 + assert updates[0].state == DataSourceState.VALID + + +@pytest.mark.asyncio +async def test_delete_object(): + src = make_streaming_data_source() + + actions = [ + server_intent_event(IntentCode.TRANSFER_FULL), + put_object_event("flag-1", 1), + delete_object_event("flag-1", 2), + payload_transferred_event(), + ] + + updates = await collect_updates(src, actions) + assert len(updates) == 1 + changes = updates[0].change_set.changes + assert any(c.key == "flag-1" and c.action == ChangeType.PUT for c in changes) + assert any(c.key == "flag-1" and c.action == ChangeType.DELETE for c in changes) + + +@pytest.mark.asyncio +async def test_goodbye_is_ignored(): + src = make_streaming_data_source() + + goodbye = Goodbye(reason="test reason") + actions = [ + Event(event=EventName.GOODBYE, data=json.dumps(goodbye.to_dict())), + server_intent_event(IntentCode.TRANSFER_FULL), + put_object_event("flag-1"), + payload_transferred_event(), + ] + + updates = await collect_updates(src, actions) + assert len(updates) == 1 + assert updates[0].state == DataSourceState.VALID + + +@pytest.mark.asyncio +async def test_error_resets_changeset(): + src = make_streaming_data_source() + + err = Error(payload_id="p1", reason="test error") + actions = [ + server_intent_event(IntentCode.TRANSFER_FULL), + put_object_event("flag-1"), + # Error mid-transfer — builder is reset + Event(event=EventName.ERROR, data=json.dumps(err.to_dict())), + # Re-transmit + server_intent_event(IntentCode.TRANSFER_FULL), + put_object_event("flag-2"), + payload_transferred_event(), + ] + + updates = await collect_updates(src, actions) + assert len(updates) == 1 + changes = [c for c in updates[0].change_set.changes if c.action == ChangeType.PUT] + keys = {c.key for c in changes} + assert "flag-2" in keys + + +@pytest.mark.asyncio +async def test_errorless_fault_is_ignored(): + src = make_streaming_data_source() + + actions = [ + Fault(error=None), + server_intent_event(IntentCode.TRANSFER_FULL), + put_object_event("flag-1"), + payload_transferred_event(), + ] + + updates = await collect_updates(src, actions) + assert len(updates) == 1 + assert updates[0].state == DataSourceState.VALID + + +@pytest.mark.asyncio +async def test_recoverable_fault_yields_interrupted_and_continues(): + src = make_streaming_data_source() + + actions = [ + Fault(error=HTTPStatusError(503)), + server_intent_event(IntentCode.TRANSFER_FULL), + put_object_event("flag-1"), + payload_transferred_event(), + ] + + updates = await collect_updates(src, actions) + assert len(updates) == 2 + assert updates[0].state == DataSourceState.INTERRUPTED + assert updates[0].error.kind == DataSourceErrorKind.ERROR_RESPONSE + assert updates[0].error.status_code == 503 + assert updates[1].state == DataSourceState.VALID + + +@pytest.mark.asyncio +async def test_unrecoverable_fault_yields_off_and_halts(): + src = make_streaming_data_source() + + actions = [ + Fault(error=HTTPStatusError(401)), + server_intent_event(IntentCode.TRANSFER_FULL), + put_object_event("flag-1"), + payload_transferred_event(), + ] + + updates = await collect_updates(src, actions) + assert len(updates) == 1 + assert updates[0].state == DataSourceState.OFF + assert src._running is False + + +@pytest.mark.asyncio +async def test_fault_with_fallback_header_halts_with_signal(): + src = make_streaming_data_source() + + headers = {_LD_FD_FALLBACK_HEADER: 'true', _LD_ENVID_HEADER: 'env1'} + actions = [ + Fault(error=HTTPStatusError(503, headers=headers)), + server_intent_event(IntentCode.TRANSFER_FULL), + ] + + updates = await collect_updates(src, actions) + assert len(updates) == 1 + assert updates[0].state == DataSourceState.OFF + assert updates[0].fallback_to_fdv1 is True + assert updates[0].environment_id == 'env1' + + +@pytest.mark.asyncio +async def test_fallback_to_fdv1_on_start_header(): + """When Start has X-LD-FD-Fallback: true the next completed update signals fallback.""" + src = make_streaming_data_source() + + fallback_headers = {_LD_FD_FALLBACK_HEADER: 'true', _LD_ENVID_HEADER: 'env1'} + + actions = [ + Start(headers=fallback_headers), + server_intent_event(IntentCode.TRANSFER_FULL), + put_object_event("flag-1"), + payload_transferred_event(), + # Should never be reached — the latched directive halts the stream. + put_object_event("flag-2"), + ] + + updates = await collect_updates(src, actions) + assert len(updates) == 1 + assert updates[0].fallback_to_fdv1 is True + assert updates[0].state == DataSourceState.VALID + + +@pytest.mark.asyncio +async def test_env_id_propagated(): + """Environment ID from Start headers should be included in updates.""" + src = make_streaming_data_source() + + start_headers = {_LD_ENVID_HEADER: 'my-env'} + actions = [ + Start(headers=start_headers), + server_intent_event(IntentCode.TRANSFER_FULL), + put_object_event("flag-1"), + payload_transferred_event(), + ] + + updates = await collect_updates(src, actions) + assert len(updates) == 1 + assert updates[0].environment_id == 'my-env' + + +@pytest.mark.asyncio +async def test_invalid_json_interrupts_stream(): + src = make_streaming_data_source() + + actions = [ + Event(event=EventName.SERVER_INTENT, data="this is not json"), + ] + + updates = await collect_updates(src, actions) + assert len(updates) == 1 + assert updates[0].state == DataSourceState.INTERRUPTED + assert updates[0].error.kind == DataSourceErrorKind.INVALID_DATA + + +@pytest.mark.asyncio +async def test_stop_sets_running_false(): + src = make_streaming_data_source() + src._running = True + await src.stop() + assert src._running is False + + +@pytest.mark.asyncio +async def test_stop_closes_sse_client(): + src = make_streaming_data_source() + sse = MockAsyncSSEClient([]) + src._sse = sse + await src.stop() + assert sse.closed is True + + +class FakeSession: + """Minimal stand-in for an aiohttp.ClientSession that records closure.""" + + def __init__(self): + self.closed = False + + async def close(self): + self.closed = True + + +def test_create_sse_client_builds_session_when_none_supplied(): + config = Config("key") + ss = MockSelectorStore(Selector.no_selector()) + fake_session = FakeSession() + + with mock.patch.object( + async_streaming, "make_client_session", return_value=fake_session + ) as make_session, mock.patch.object( + async_streaming, "AsyncSSEFactory" + ) as factory_cls: + factory_cls.return_value.create.return_value = MockAsyncSSEClient([]) + + sse, owned = create_sse_client( + config.stream_base_uri, + config.http, + config.initial_reconnect_delay, + config, + ss, + ) + + make_session.assert_called_once_with(config, config.http) + # The configured session is handed to the SSE factory... + assert factory_cls.call_args.kwargs["session"] is fake_session + # ...and returned as the owned session the SDK must close. + assert owned is fake_session + + +def test_create_sse_client_does_not_build_session_when_supplied(): + config = Config("key") + ss = MockSelectorStore(Selector.no_selector()) + supplied = FakeSession() + + with mock.patch.object( + async_streaming, "make_client_session" + ) as make_session, mock.patch.object( + async_streaming, "AsyncSSEFactory" + ) as factory_cls: + factory_cls.return_value.create.return_value = MockAsyncSSEClient([]) + + sse, owned = create_sse_client( + config.stream_base_uri, + config.http, + config.initial_reconnect_delay, + config, + ss, + session=supplied, + ) + + make_session.assert_not_called() + assert factory_cls.call_args.kwargs["session"] is supplied + # A caller-supplied session is not owned by the SDK. + assert owned is None + + +@pytest.mark.asyncio +async def test_owned_session_closed_on_sync_completion(): + """When the SDK creates the session, it is closed once sync() finishes.""" + src = make_streaming_data_source() + owned = FakeSession() + + def builder(base_uri, http_options, initial_reconnect_delay, config, ss, session=None): # pylint: disable=unused-argument + return MockAsyncSSEClient([ + server_intent_event(IntentCode.TRANSFER_FULL), + put_object_event("flag-1"), + payload_transferred_event(), + ]), owned + + src._sse_client_builder = builder + ss = MockSelectorStore(Selector.no_selector()) + + updates = [u async for u in src.sync(ss)] + assert len(updates) == 1 + assert owned.closed is True + assert src._owned_session is None + + +@pytest.mark.asyncio +async def test_owned_session_closed_on_stop(): + """When the SDK creates the session, stop() closes it.""" + src = make_streaming_data_source() + owned = FakeSession() + src._sse = MockAsyncSSEClient([]) + src._owned_session = owned + + await src.stop() + assert owned.closed is True + assert src._owned_session is None + + +@pytest.mark.asyncio +async def test_supplied_session_not_closed_on_sync_completion(): + """A caller-supplied session (owned=None) is never closed by the SDK.""" + src = make_streaming_data_source() + supplied = FakeSession() + + def builder(base_uri, http_options, initial_reconnect_delay, config, ss, session=None): # pylint: disable=unused-argument + # Mirrors create_sse_client when a session is supplied: owned is None. + return MockAsyncSSEClient([ + server_intent_event(IntentCode.TRANSFER_FULL), + put_object_event("flag-1"), + payload_transferred_event(), + ]), None + + src._sse_client_builder = builder + ss = MockSelectorStore(Selector.no_selector()) + + [u async for u in src.sync(ss)] + assert supplied.closed is False + + +@pytest.mark.asyncio +async def test_supplied_session_flows_through_builder_as_unowned(): + """A session supplied via the ctor reaches create_sse_client and is not owned.""" + config = Config("key") + supplied = FakeSession() + src = AsyncStreamingDataSource( + config.stream_base_uri + STREAMING_ENDPOINT, + config.http, + config.initial_reconnect_delay, + config, + session=supplied, + ) + + with mock.patch.object( + async_streaming, "make_client_session" + ) as make_session, mock.patch.object( + async_streaming, "AsyncSSEFactory" + ) as factory_cls: + factory_cls.return_value.create.return_value = MockAsyncSSEClient([]) + ss = MockSelectorStore(Selector.no_selector()) + + [u async for u in src.sync(ss)] + + make_session.assert_not_called() + assert factory_cls.call_args.kwargs["session"] is supplied + assert supplied.closed is False + assert src._owned_session is None + + +@pytest.mark.asyncio +async def test_fallback_latched_on_start_carries_through_recoverable_fault(): + """A directive latched on Start must halt the stream even on a recoverable fault.""" + src = make_streaming_data_source() + + actions = [ + Start(headers={_LD_FD_FALLBACK_HEADER: 'true'}), + Fault(error=HTTPStatusError(503)), + # Must never be reached — the latched directive halts the stream. + server_intent_event(IntentCode.TRANSFER_FULL), + put_object_event("flag-1"), + payload_transferred_event(), + ] + + updates = await collect_updates(src, actions) + assert len(updates) == 1 + assert updates[0].state == DataSourceState.INTERRUPTED + assert updates[0].error.status_code == 503 + assert updates[0].fallback_to_fdv1 is True + + +@pytest.mark.asyncio +async def test_envid_preserved_across_headerless_reconnect_start(): + """A reconnect Start without X-LD-EnvID must not clear the environment ID.""" + src = make_streaming_data_source() + + actions = [ + Start(headers={_LD_ENVID_HEADER: 'my-env'}), + server_intent_event(IntentCode.TRANSFER_FULL), + put_object_event("flag-1"), + payload_transferred_event(1), + Start(headers={}), + server_intent_event(IntentCode.TRANSFER_FULL), + put_object_event("flag-2"), + payload_transferred_event(2), + ] + + updates = await collect_updates(src, actions) + assert len(updates) == 2 + assert updates[0].environment_id == 'my-env' + assert updates[1].environment_id == 'my-env' diff --git a/ldclient/testing/impl/datasourcev2/test_polling_payload_parsing.py b/ldclient/testing/impl/datasourcev2/test_polling_payload_parsing.py index 11fd2d68..8e4c6369 100644 --- a/ldclient/testing/impl/datasourcev2/test_polling_payload_parsing.py +++ b/ldclient/testing/impl/datasourcev2/test_polling_payload_parsing.py @@ -1,12 +1,11 @@ import json from ldclient.impl.datasourcev2.polling import ( - IntentCode, fdv1_polling_payload_to_changeset, polling_payload_to_changeset ) from ldclient.impl.util import _Fail, _Success -from ldclient.interfaces import ChangeType, ObjectKind +from ldclient.interfaces import ChangeType, IntentCode, ObjectKind def test_payload_is_missing_events_key():