diff --git a/README.md b/README.md index f5d307a3..92b24f49 100644 --- a/README.md +++ b/README.md @@ -258,11 +258,13 @@ flags.get_all_flags({"organization": "org_123"}) # {"new-checkout": True, ...} flags.close() # on shutdown; or use `with workos.create_feature_flags_runtime_client() as flags:` ``` +With `AsyncWorkOSClient`, polling runs in an asyncio task: `await flags.start()`, `await flags.wait_until_ready(timeout=5)` and `await flags.close()` (or `async with`). Evaluation stays synchronous. + **Evaluation.** The context maps a target type to the exact, case-sensitive ID being evaluated. `user` and `organization` are target types like any [custom target](https://workos.com/docs/feature-flags/custom-targets) type, and an ID of `None` is skipped. A disabled flag serves its off value. For an enabled flag, the first rule whose conditions all match wins, so a rule serving `False` can exclude a user before a later rule includes their organization; otherwise the flag's default value applies. A missing flag, an unready client, or a flag without a boolean answer returns `default` (`False` unless you pass one), and `get_all_flags` omits flags without an answer. Evaluation never raises: invalid context entries are ignored with a logged warning and the rest is evaluated. -**Failures.** A failed poll keeps the last good configuration and retries with exponential backoff. Transient failures (429, 500, 502, 503, 504, timeouts, connection errors) are also retried within one refresh, up to the client's `max_retries` and within `request_timeout`. A 429's `Retry-After` is honored up to 60 seconds, both by those retries and before the next poll. The client streams the response and gives every network wait only the time left, so a refresh takes at most `request_timeout` plus one read of a response body that stalls; a custom `HTTPBackend` that is not an httpx client receives the remaining time as its per-request timeout instead. A 401, from a poll or from `refresh()`, stops polling for good: create a new client with a valid API key. Failures are logged to the `workos.feature_flags.runtime` logger (pass `logger=` to use your own) and reported to error listeners. `flags.status` shows readiness, the last error, refresh counts and timestamps, `polling_stopped`, `flag_count` and `cache_age()`. +**Failures.** A failed poll keeps the last good configuration and retries with exponential backoff. Transient failures (429, 500, 502, 503, 504, timeouts, connection errors) are also retried within one refresh, up to the client's `max_retries` and within `request_timeout`. A 429's `Retry-After` is honored up to 60 seconds, both by those retries and before the next poll. The async client enforces `request_timeout` exactly. The sync client streams the response and gives every network wait only the time left, so a refresh takes at most `request_timeout` plus one read of a response body that stalls; a custom `HTTPBackend` that is not an httpx client receives the remaining time as its per-request timeout instead. A 401, from a poll or from `refresh()`, stops polling for good: create a new client with a valid API key. Failures are logged to the `workos.feature_flags.runtime` logger (pass `logger=` to use your own) and reported to error listeners. `flags.status` shows readiness, the last error, refresh counts and timestamps, `polling_stopped`, `flag_count` and `cache_age()`. -**Notifications.** `on_change` (or `add_change_listener`) fires after a refresh installs a different configuration, including the first load, with the previous and current snapshots and the changed slugs. `on_error` fires for failed background polls and `on_status` after every refresh. Listeners run one at a time; their exceptions are logged and never affect polling. +**Notifications.** `on_change` (or `add_change_listener`) fires after a refresh installs a different configuration, including the first load, with the previous and current snapshots and the changed slugs. `on_error` fires for failed background polls and `on_status` after every refresh. Listeners run one at a time; their exceptions are logged and never affect polling. Async clients also accept coroutine listeners. **Bootstrap.** Persist the configuration and pass it back to make a new client ready before its first poll: @@ -277,7 +279,7 @@ flags = workos.create_feature_flags_runtime_client( **Inspection and manual refresh.** `flags.get_flag(slug)` returns a flag's configuration as received, with its rules in order. `flags.refresh()` fetches the configuration immediately, raising on failure, and works without `start()`, e.g. in scripts or serverless functions. -**Processes.** The runtime client polls through the `WorkOSClient`'s HTTP client, so close it before closing the `WorkOSClient` or an HTTP client you supplied. Create the `WorkOSClient` and the runtime client in each worker process, after forking, rather than in an app gunicorn loads with `--preload`. Background threads must be enabled, so run uWSGI with `--enable-threads`. +**Processes.** The runtime client polls through its `WorkOSClient` or `AsyncWorkOSClient`, so close it before closing that client or an HTTP client you supplied. Create the client and the runtime client in each worker process, after forking, rather than in an app gunicorn loads with `--preload`. Background threads must be enabled, so run uWSGI with `--enable-threads`. ## Type Safety diff --git a/src/workos/__init__.py b/src/workos/__init__.py index f29ac8fe..fc234f95 100644 --- a/src/workos/__init__.py +++ b/src/workos/__init__.py @@ -24,6 +24,7 @@ ) from ._pagination import AsyncPage, ListMetadata, SyncPage from .feature_flags.runtime import ( + AsyncFeatureFlagsRuntimeClient, FeatureFlagsConfiguration, FeatureFlagsRuntimeClient, ) @@ -56,5 +57,6 @@ "TransportTimeout", "TransportConnectError", "FeatureFlagsRuntimeClient", + "AsyncFeatureFlagsRuntimeClient", "FeatureFlagsConfiguration", ] diff --git a/src/workos/_client.py b/src/workos/_client.py index 592c6480..8425e8f8 100644 --- a/src/workos/_client.py +++ b/src/workos/_client.py @@ -344,4 +344,10 @@ def pkce(self) -> PKCE: """PKCE (Proof Key for Code Exchange) utilities.""" return PKCE() + # A plain function bound as a method, so its signature and docs live next + # to the runtime client's constructor. + from .feature_flags.runtime._factory import ( + create_async_feature_flags_runtime_client as create_feature_flags_runtime_client, + ) + # @oagen-ignore-end diff --git a/src/workos/feature_flags/runtime/__init__.py b/src/workos/feature_flags/runtime/__init__.py index 1b488922..e13e2b12 100644 --- a/src/workos/feature_flags/runtime/__init__.py +++ b/src/workos/feature_flags/runtime/__init__.py @@ -4,9 +4,11 @@ """Local evaluation of WorkOS feature flags from a polled configuration. -Use :class:`FeatureFlagsRuntimeClient` with :class:`workos.WorkOSClient`. +Use :class:`FeatureFlagsRuntimeClient` with :class:`workos.WorkOSClient` and +:class:`AsyncFeatureFlagsRuntimeClient` with :class:`workos.AsyncWorkOSClient`. """ +from ._async import AsyncFeatureFlagsRuntimeClient from ._configuration import ( ConditionDefinition, FeatureFlagsConfiguration, @@ -19,6 +21,7 @@ from ._types import FlagsChangedEvent, RuntimeStatus __all__ = [ + "AsyncFeatureFlagsRuntimeClient", "ConditionDefinition", "EvaluationContext", "FeatureFlagsConfiguration", diff --git a/src/workos/feature_flags/runtime/_async.py b/src/workos/feature_flags/runtime/_async.py new file mode 100644 index 00000000..cbc80c33 --- /dev/null +++ b/src/workos/feature_flags/runtime/_async.py @@ -0,0 +1,398 @@ +# @oagen-ignore-file +# This file is hand-maintained. The asynchronous feature flags runtime client +# polls in an asyncio task and evaluates flags locally. +# pyright: strict, reportPrivateUsage=false + +from __future__ import annotations + +import asyncio +import inspect +import time +from types import TracebackType +from typing import TYPE_CHECKING, Any, Awaitable, Callable, Optional, Type, Union + +from ..._errors import WorkOSTimeoutError +from . import _policy +from ._configuration import FeatureFlagsConfiguration +from ._context import LoggerLike +from ._core import Notification, RuntimeClientCore, utc_now +from ._types import FlagsChangedEvent, RuntimeStatus + +if TYPE_CHECKING: + from ..._base_client import AsyncWorkOSClient + +ChangeListener = Callable[[FlagsChangedEvent], Union[None, Awaitable[None]]] +ErrorListener = Callable[[Exception], Union[None, Awaitable[None]]] +StatusListener = Callable[[RuntimeStatus], Union[None, Awaitable[None]]] + + +class _Closed(Exception): + """The client was closed while a refresh was in progress.""" + + +class AsyncFeatureFlagsRuntimeClient(RuntimeClientCore): + """Polls an environment's feature flag configuration in an asyncio task and evaluates flags locally. + + Create one long-lived instance per event loop and share it. Evaluation + (``is_enabled``, ``get_all_flags``, ``get_flag``) is synchronous, never + makes a network request and never raises:: + + flags = workos.create_feature_flags_runtime_client() + await flags.start() + await flags.wait_until_ready(timeout=5) + flags.is_enabled("new-checkout", {"organization": "org_123"}) + + The client polls through its ``AsyncWorkOSClient``, so use it on that + client's event loop, and close it before closing the ``AsyncWorkOSClient``. + """ + + def __init__( + self, + client: AsyncWorkOSClient, + *, + polling_interval: float = _policy.DEFAULT_POLLING_INTERVAL, + request_timeout: float = _policy.DEFAULT_REQUEST_TIMEOUT, + max_retry_delay: float = _policy.DEFAULT_MAX_RETRY_DELAY, + bootstrap: Optional[FeatureFlagsConfiguration] = None, + logger: Optional[LoggerLike] = None, + on_change: Optional[ChangeListener] = None, + on_error: Optional[ErrorListener] = None, + on_status: Optional[StatusListener] = None, + ) -> None: + """Create a runtime client. No request is made until :meth:`start` or :meth:`refresh`. + + Takes the same options as :class:`FeatureFlagsRuntimeClient`. Listeners + may be plain functions or coroutine functions. + + Raises: + ConfigurationError: If ``client`` has no API key. + """ + super().__init__( + client, + polling_interval=polling_interval, + request_timeout=request_timeout, + max_retry_delay=max_retry_delay, + bootstrap=bootstrap, + logger=logger, + ) + self._client = client + self._task: Optional[asyncio.Task[None]] = None + # The task delivering notifications, and the coroutine listener it is + # waiting for, if any. + self._drainer: Optional[asyncio.Task[Any]] = None + self._listener_task: Optional[asyncio.Future[Any]] = None + self._init_synchronization() + if on_change is not None: + self.add_change_listener(on_change) + if on_error is not None: + self.add_error_listener(on_error) + if on_status is not None: + self.add_status_listener(on_status) + + def _init_synchronization(self) -> None: + # None of these binds to an event loop until first awaited. + self._refresh_lock = asyncio.Lock() + self._ready = asyncio.Event() + # Set when polling stops after a 401, to end the polling task's sleep. + self._wake = asyncio.Event() + if self._status.ready: + self._ready.set() + + # --- Lifecycle ------------------------------------------------------------ + + async def start(self) -> None: + """Start background polling in a task on the running event loop. Idempotent. + + The first poll is sent right away, then every ``polling_interval`` + seconds. Failed polls retry with exponential backoff. A 401 stops + polling permanently; the last good configuration stays usable. + + Raises: + RuntimeError: If the client is closed or no event loop is running. + """ + if self._closed: + raise RuntimeError( + "This feature flags runtime client is closed; create a new one" + ) + if self._started: + return + loop = asyncio.get_running_loop() + self._started = True + if self._status.polling_stopped: + self._logger.warning( + "Not starting feature flag polling: it stopped after a 401. " + "Create a new runtime client with a valid API key." + ) + return + self._task = loop.create_task(self._poll_forever(), name="workos-feature-flags") + + async def wait_until_ready(self, timeout: Optional[float] = None) -> None: + """Wait until a bootstrap or a successful refresh makes flags available. + + Does not start polling: call :meth:`start` (or :meth:`refresh`). + + Args: + timeout: Seconds to wait. ``None`` waits indefinitely. + + Raises: + TimeoutError: If flags are not available within ``timeout``. + AuthenticationError: If polling stopped after a 401 before any + configuration was available. + RuntimeError: If the client was closed before it became ready. + """ + if self._status.ready: + return + try: + await asyncio.wait_for(self._ready.wait(), timeout) + except asyncio.TimeoutError: + # The readiness timeout names the last refresh error as its cause. + raise self._readiness_timeout(timeout) + if not self._status.ready: + raise self._readiness_failure() + + async def refresh(self) -> None: + """Fetch the configuration now. + + Serialized with background polling. Errors are raised to the caller and + the last good configuration is kept; ``on_error`` is not called. A 401 + also stops background polling, as it would from a background poll. + + Raises: + WorkOSError: If the request fails or the payload cannot be read. + RuntimeError: If the client is closed. + """ + if self._closed: + raise RuntimeError( + "This feature flags runtime client is closed; create a new one" + ) + try: + await self._refresh(background=False) + except _Closed: + raise RuntimeError( + "The feature flags runtime client was closed during the refresh" + ) from None + except Exception as error: + if _policy.is_unauthorized(error): + await self._stop_polling(error) + raise + + async def close(self) -> None: + """Stop polling and wait for the polling task to finish. Idempotent. + + Every caller waits, including one that calls while another close() is + still waiting. Evaluation keeps serving the last configuration after + closing. + """ + first = not self._closed + self._closed = True + self._ready.set() + task = self._task + current = asyncio.current_task() + if ( + task is None + or task.done() + or task is current + or (current is self._listener_task and self._drainer is task) + ): + # Called from a listener the polling task is delivering to: the + # task exits on its own once the listener returns. + return + if first: + task.cancel() + # `wait` never raises the task's own cancellation, but still + # propagates a cancellation of the caller. + await asyncio.wait({task}) + + async def __aenter__(self) -> AsyncFeatureFlagsRuntimeClient: + await self.start() + return self + + async def __aexit__( + self, + exc_type: Optional[Type[BaseException]], + exc_val: Optional[BaseException], + exc_tb: Optional[TracebackType], + ) -> None: + await self.close() + + # --- Listeners ------------------------------------------------------------ + + def add_change_listener(self, listener: ChangeListener) -> None: + """Call ``listener`` after a refresh installs a different configuration. + + Fires on the first load and on the first refresh that differs from a + bootstrap; never for identical refreshes or failures. Coroutine + listeners are awaited. Listeners run one at a time in refresh order; + their exceptions are logged and never affect polling. + """ + self._add_listener("change", listener) + + def remove_change_listener(self, listener: ChangeListener) -> None: + """Stop calling a change listener.""" + self._remove_listener("change", listener) + + def add_error_listener(self, listener: ErrorListener) -> None: + """Call ``listener`` with the error of every failed background poll.""" + self._add_listener("error", listener) + + def remove_error_listener(self, listener: ErrorListener) -> None: + """Stop calling an error listener.""" + self._remove_listener("error", listener) + + def add_status_listener(self, listener: StatusListener) -> None: + """Call ``listener`` with the new status after every refresh and when polling stops.""" + self._add_listener("status", listener) + + def remove_status_listener(self, listener: StatusListener) -> None: + """Stop calling a status listener.""" + self._remove_listener("status", listener) + + # --- Polling -------------------------------------------------------------- + + async def _poll_forever(self) -> None: + failures = 0 + while not self._closed and not self._status.polling_stopped: + retry_after: Optional[float] = None + try: + await self._refresh(background=True) + failures = 0 + except _Closed: + return + except Exception as error: + if self._closed: + return + if _policy.is_unauthorized(error): + await self._stop_polling(error) + return + failures = min(failures + 1, _policy.MAX_COUNTED_FAILURES) + retry_after = _policy.retry_after(error) + if self._closed: + return + delay = _policy.next_delay( + self._polling_interval, failures, self._max_retry_delay + ) + try: + await asyncio.wait_for( + self._wake.wait(), timeout=max(delay, retry_after or 0.0) + ) + except asyncio.TimeoutError: + pass + + async def _refresh(self, *, background: bool) -> None: + try: + async with self._refresh_lock: + await self._refresh_locked(background=background) + finally: + await self._drain() + + async def _refresh_locked(self, *, background: bool) -> None: + if self._closed or (background and self._status.polling_stopped): + raise _Closed() + attempted = utc_now() + try: + payload = await asyncio.wait_for(self._fetch(), self._request_timeout) + configuration = self._parse(payload) + except asyncio.TimeoutError as error: + if self._closed: + raise _Closed() from error + timeout = WorkOSTimeoutError( + f"Feature flag refresh exceeded its request timeout of {self._request_timeout} seconds" + ) + self._pending.append( + self._record_failure(timeout, attempted, background=background) + ) + raise timeout from error + except Exception as error: + if self._closed: + raise _Closed() from error + self._pending.append( + self._record_failure(error, attempted, background=background) + ) + raise + if self._closed: + raise _Closed() + notifications = self._record_success( + configuration, attempted, background=background + ) + self._ready.set() + self._pending.append(notifications) + + async def _fetch(self) -> Any: + deadline = time.monotonic() + self._request_timeout + attempt = 0 + while True: + remaining = max(deadline - time.monotonic(), 0.001) + try: + return await self._client.request( + "get", + _policy.POLL_PATH, + params=dict(_policy.POLL_PARAMS), + request_options={"timeout": remaining, "max_retries": 0}, + ) + except Exception as error: + delay = _policy.attempt_retry_delay( + error, attempt, deadline - time.monotonic(), self._max_retries + ) + if delay is None: + raise + await asyncio.sleep(delay) + attempt += 1 + + async def _stop_polling(self, error: Exception) -> None: + async with self._refresh_lock: + if self._status.polling_stopped: + return + notifications = self._record_polling_stopped(error) + if not self._status.ready: + self._readiness_error = error + self._ready.set() + self._wake.set() + self._pending.append(notifications) + await self._drain() + + async def _drain(self) -> None: + """Deliver pending notifications in refresh order, one batch at a time. + + A refresh that finishes while another task is delivering, including a + listener calling refresh(), leaves its batch to that task, so batches + are never interleaved. + """ + if self._drainer is not None: + return + self._drainer = asyncio.current_task() + try: + while self._pending: + batch = self._pending.popleft() + if not self._closed: + await self._deliver(batch) + finally: + self._drainer = None + + async def _deliver(self, notifications: list[Notification]) -> None: + for listener, argument in notifications: + try: + result = listener(argument) + if inspect.isawaitable(result): + await self._await_listener(result) + except Exception: + self._log_listener_error() + + async def _await_listener(self, awaitable: Awaitable[Any]) -> None: + # The listener runs in a task of its own, so a CancelledError raised + # inside it, such as from awaiting work something else cancelled, ends + # only that listener. Cancelling the delivering task, as close() does, + # cancels the listener with it. + listener_task = asyncio.ensure_future(awaitable) + self._listener_task = listener_task + try: + await asyncio.wait({listener_task}) + except asyncio.CancelledError: + listener_task.cancel() + await asyncio.wait({listener_task}) + raise + finally: + self._listener_task = None + if listener_task.cancelled(): + self._logger.error("A feature flags runtime client listener was cancelled") + else: + listener_task.result() diff --git a/src/workos/feature_flags/runtime/_core.py b/src/workos/feature_flags/runtime/_core.py index cfd9fb6d..aab608e7 100644 --- a/src/workos/feature_flags/runtime/_core.py +++ b/src/workos/feature_flags/runtime/_core.py @@ -1,6 +1,6 @@ # @oagen-ignore-file -# This file is hand-maintained. State and local evaluation for the feature -# flags runtime client. +# This file is hand-maintained. State and local evaluation shared by the sync +# and async feature flags runtime clients. # pyright: strict, reportPrivateUsage=false from __future__ import annotations diff --git a/src/workos/feature_flags/runtime/_factory.py b/src/workos/feature_flags/runtime/_factory.py index f9966a80..393b3e02 100644 --- a/src/workos/feature_flags/runtime/_factory.py +++ b/src/workos/feature_flags/runtime/_factory.py @@ -1,6 +1,6 @@ # @oagen-ignore-file -# This file is hand-maintained. The factory method that WorkOSClient exposes -# as `create_feature_flags_runtime_client`. +# This file is hand-maintained. Factory methods that WorkOSClient and +# AsyncWorkOSClient expose as `create_feature_flags_runtime_client`. # pyright: strict from __future__ import annotations @@ -8,13 +8,19 @@ from typing import TYPE_CHECKING, Callable, Optional from . import _policy +from ._async import ( + AsyncFeatureFlagsRuntimeClient, + ChangeListener, + ErrorListener, + StatusListener, +) from ._configuration import FeatureFlagsConfiguration from ._context import LoggerLike from ._sync import FeatureFlagsRuntimeClient from ._types import FlagsChangedEvent, RuntimeStatus if TYPE_CHECKING: - from ..._base_client import WorkOSClient + from ..._base_client import AsyncWorkOSClient, WorkOSClient def create_feature_flags_runtime_client( @@ -47,3 +53,35 @@ def create_feature_flags_runtime_client( on_error=on_error, on_status=on_status, ) + + +def create_async_feature_flags_runtime_client( + self: AsyncWorkOSClient, + *, + polling_interval: float = _policy.DEFAULT_POLLING_INTERVAL, + request_timeout: float = _policy.DEFAULT_REQUEST_TIMEOUT, + max_retry_delay: float = _policy.DEFAULT_MAX_RETRY_DELAY, + bootstrap: Optional[FeatureFlagsConfiguration] = None, + logger: Optional[LoggerLike] = None, + on_change: Optional[ChangeListener] = None, + on_error: Optional[ErrorListener] = None, + on_status: Optional[StatusListener] = None, +) -> AsyncFeatureFlagsRuntimeClient: + """Create an async feature flags runtime client that polls through this client. + + The runtime client evaluates flags locally from a configuration it polls in + an asyncio task. Call ``await start()`` to begin polling. See + :class:`~workos.feature_flags.runtime.AsyncFeatureFlagsRuntimeClient` for + the options. + """ + return AsyncFeatureFlagsRuntimeClient( + self, + polling_interval=polling_interval, + request_timeout=request_timeout, + max_retry_delay=max_retry_delay, + bootstrap=bootstrap, + logger=logger, + on_change=on_change, + on_error=on_error, + on_status=on_status, + ) diff --git a/src/workos/feature_flags/runtime/_policy.py b/src/workos/feature_flags/runtime/_policy.py index 395473bf..6b5a2374 100644 --- a/src/workos/feature_flags/runtime/_policy.py +++ b/src/workos/feature_flags/runtime/_policy.py @@ -1,6 +1,6 @@ # @oagen-ignore-file -# This file is hand-maintained. It holds the polling and retry policy of the -# feature flags runtime client. +# This file is hand-maintained. It holds the polling and retry policy shared +# by the sync and async feature flags runtime clients. # pyright: strict from __future__ import annotations diff --git a/src/workos/feature_flags/runtime/_sync.py b/src/workos/feature_flags/runtime/_sync.py index e614d46c..a8f03d5c 100644 --- a/src/workos/feature_flags/runtime/_sync.py +++ b/src/workos/feature_flags/runtime/_sync.py @@ -234,7 +234,8 @@ def _add_listener(self, kind: str, listener: Listener) -> None: # never run. if inspect.iscoroutinefunction(listener): raise TypeError( - f"{kind} listener must be a plain function, not an async function" + f"{kind} listener must be a plain function, not an async function; " + "use AsyncFeatureFlagsRuntimeClient for async listeners" ) super()._add_listener(kind, listener) @@ -407,7 +408,8 @@ def _deliver(self, notifications: list[Notification]) -> None: _discard_awaitable(result) self._logger.warning( "A FeatureFlagsRuntimeClient listener returned an awaitable, " - "which was not awaited; listeners must be plain functions" + "which was not awaited; use AsyncFeatureFlagsRuntimeClient for " + "async listeners" ) diff --git a/tests/feature_flags_runtime/conftest.py b/tests/feature_flags_runtime/conftest.py index 76e5f8d6..8106ef8b 100644 --- a/tests/feature_flags_runtime/conftest.py +++ b/tests/feature_flags_runtime/conftest.py @@ -3,16 +3,22 @@ from __future__ import annotations -from typing import Any, Callable, Iterator, List +from typing import Any, AsyncIterator, Callable, Iterator, List import pytest +import pytest_asyncio -from workos import WorkOSClient -from workos.feature_flags.runtime import FeatureFlagsRuntimeClient, _policy +from workos import AsyncWorkOSClient, WorkOSClient +from workos.feature_flags.runtime import ( + AsyncFeatureFlagsRuntimeClient, + FeatureFlagsRuntimeClient, + _policy, +) from tests.feature_flags_runtime.helpers import ( API_KEY, BASE_URL, + AsyncBackend, FlagServer, SyncBackend, ) @@ -32,6 +38,13 @@ def client(server: FlagServer) -> Iterator[WorkOSClient]: workos.close() +@pytest.fixture +def async_client(server: FlagServer) -> AsyncWorkOSClient: + return AsyncWorkOSClient( + api_key=API_KEY, base_url=BASE_URL, http_client=AsyncBackend(server) + ) + + @pytest.fixture def fast_polling(monkeypatch: pytest.MonkeyPatch) -> List[int]: """Poll every 10 ms and retry inside a refresh after 5 ms. @@ -65,3 +78,20 @@ def make(**options: Any) -> FeatureFlagsRuntimeClient: yield make for runtime in created: runtime.close() + + +@pytest_asyncio.fixture +async def make_async_runtime( + async_client: AsyncWorkOSClient, +) -> AsyncIterator[Callable[..., AsyncFeatureFlagsRuntimeClient]]: + """Create async runtime clients that are always closed at teardown.""" + created: List[AsyncFeatureFlagsRuntimeClient] = [] + + def make(**options: Any) -> AsyncFeatureFlagsRuntimeClient: + runtime = async_client.create_feature_flags_runtime_client(**options) + created.append(runtime) + return runtime + + yield make + for runtime in created: + await runtime.close() diff --git a/tests/feature_flags_runtime/helpers.py b/tests/feature_flags_runtime/helpers.py index 8019e5aa..24a5a91b 100644 --- a/tests/feature_flags_runtime/helpers.py +++ b/tests/feature_flags_runtime/helpers.py @@ -172,6 +172,35 @@ def close(self) -> None: pass +class AsyncBackend: + def __init__(self, server: FlagServer) -> None: + self.server = server + + async def request( + self, + method: str, + url: str, + *, + headers: Dict[str, str], + content: Optional[bytes], + timeout: float, + ) -> HTTPResponse: + import asyncio + + reply = self.server._next(method, url, headers, timeout) + try: + if reply is FlagServer.BLOCK: + while not self.server._released.is_set(): + await asyncio.sleep(0.005) + reply = (200, payload()) + return self.server._response(reply, method, url) + finally: + self.server._finish() + + async def close(self) -> None: + pass + + class Collector: """Thread-safe list of listener calls with a blocking wait.""" diff --git a/tests/feature_flags_runtime/test_async_client.py b/tests/feature_flags_runtime/test_async_client.py new file mode 100644 index 00000000..290e2935 --- /dev/null +++ b/tests/feature_flags_runtime/test_async_client.py @@ -0,0 +1,417 @@ +# @oagen-ignore-file +"""AsyncFeatureFlagsRuntimeClient: polling, readiness, notifications and shutdown.""" + +from __future__ import annotations + +import asyncio +import traceback +from typing import Any, Callable, List + +import pytest + +from workos import AsyncWorkOSClient +from workos._errors import AuthenticationError, NotFoundError, WorkOSTimeoutError +from workos.feature_flags.runtime import ( + AsyncFeatureFlagsRuntimeClient, + FeatureFlagsConfiguration, + FeatureFlagsConfigurationError, + FlagsChangedEvent, +) + +from tests.feature_flags_runtime.helpers import ( + API_KEY, + BASE_URL, + FlagServer, + flag, + payload, + rule, +) + +pytestmark = pytest.mark.asyncio + +Make = Callable[..., AsyncFeatureFlagsRuntimeClient] + +ORDERED = payload( + ordered=flag( + "ordered", + default_value=True, + rules=[ + rule("flag_rule_1", False, "organization", "org_1"), + rule("flag_rule_2", True, "user", "user_1"), + ], + ) +) +CHANGED = payload( + ordered=flag("ordered", default_value=False), + added=flag("added", default_value=True), +) + + +async def until(predicate: Callable[[], bool], timeout: float = 5.0) -> None: + async def poll() -> None: + while not predicate(): + await asyncio.sleep(0.002) + + await asyncio.wait_for(poll(), timeout) + + +async def test_polls_the_v2_endpoint_and_becomes_ready( + server: FlagServer, make_async_runtime: Make +) -> None: + server.respond(200, ORDERED) + runtime = make_async_runtime(request_timeout=7) + assert server.requests == [] + await runtime.start() + await runtime.wait_until_ready(timeout=5) + + request = server.requests[0] + assert request.url == f"{BASE_URL}/sdk/feature-flags?payload_version=2" + assert request.headers["Authorization"] == f"Bearer {API_KEY}" + assert 0 < request.timeout <= 7 + assert ( + runtime.is_enabled("ordered", {"user": "user_1", "organization": "org_1"}) + is False + ) + assert runtime.is_enabled("ordered", {"user": "user_1"}) is True + assert runtime.get_all_flags() == {"ordered": True} + assert runtime.get_flag("ordered") is not None + assert runtime.status.refresh_count == 1 + + +async def test_bootstrap_is_ready_without_a_request( + server: FlagServer, make_async_runtime: Make +) -> None: + runtime = make_async_runtime(bootstrap=FeatureFlagsConfiguration.from_dict(ORDERED)) + await runtime.wait_until_ready(timeout=0) + assert runtime.is_enabled("ordered") is True + assert server.requests == [] + + +async def test_wait_times_out_without_a_configuration(make_async_runtime: Make) -> None: + with pytest.raises(TimeoutError): + await make_async_runtime().wait_until_ready(timeout=0.01) + + +async def test_a_wait_timeout_names_the_last_refresh_error( + server: FlagServer, make_async_runtime: Make +) -> None: + server.respond(404, {"message": "Not Found"}) + runtime = make_async_runtime() + with pytest.raises(NotFoundError): + await runtime.refresh() + with pytest.raises( + TimeoutError, match="the last refresh failed: NotFoundError" + ) as raised: + await runtime.wait_until_ready(timeout=0.01) + assert raised.value.__cause__ is runtime.status.last_error + + +async def test_each_wait_after_a_401_raises_a_fresh_error( + server: FlagServer, make_async_runtime: Make +) -> None: + server.respond(401, {"message": "Unauthorized"}) + runtime = make_async_runtime() + with pytest.raises(AuthenticationError): + await runtime.refresh() + errors: List[BaseException] = [] + for _ in range(4): + try: + await runtime.wait_until_ready(timeout=1) + except AuthenticationError as error: + errors.append(error) + assert len({id(error) for error in errors}) == 4 + stored = errors[0].__cause__ + assert stored is not None and all(error.__cause__ is stored for error in errors) + depths = {len(traceback.extract_tb(error.__traceback__)) for error in errors} + assert len(depths) == 1 + + +async def test_awaits_coroutine_and_plain_listeners( + server: FlagServer, make_async_runtime: Make, fast_polling: List[int] +) -> None: + server.respond(200, ORDERED) + seen: List[Any] = [] + + async def on_change(event: FlagsChangedEvent) -> None: + await asyncio.sleep(0) + seen.append(("change", event.changed_flags)) + + runtime = make_async_runtime(on_change=on_change) + runtime.add_status_listener( + lambda status: seen.append(("status", status.refresh_count)) + ) + await runtime.start() + await until(lambda: len(seen) >= 2) + assert seen[:2] == [("status", 1), ("change", frozenset({"ordered"}))] + + server.respond(200, CHANGED) + await until( + lambda: any( + item == ("change", frozenset({"ordered", "added"})) for item in seen + ) + ) + + +async def test_a_listener_cancelled_from_within_does_not_stop_polling( + server: FlagServer, make_async_runtime: Make, fast_polling: List[int] +) -> None: + server.respond(200, ORDERED) + changes: List[FlagsChangedEvent] = [] + + async def awaits_cancelled_work(event: FlagsChangedEvent) -> None: + if event.previous is None: + # Work someone else cancelled raises CancelledError in the listener. + cancelled = asyncio.get_running_loop().create_future() + cancelled.cancel() + await cancelled + + runtime = make_async_runtime(on_change=awaits_cancelled_work) + runtime.add_change_listener(changes.append) + await runtime.start() + await until(lambda: len(changes) == 1) + + server.respond(200, CHANGED) + await until(lambda: runtime.is_enabled("added")) + await until(lambda: len(changes) == 2) + assert runtime._task is not None and not runtime._task.done() + + +async def test_close_cancels_a_running_listener( + server: FlagServer, make_async_runtime: Make +) -> None: + server.respond(200, ORDERED) + started = asyncio.Event() + cancelled = asyncio.Event() + + async def slow_listener(event: FlagsChangedEvent) -> None: + started.set() + try: + await asyncio.sleep(10) + except asyncio.CancelledError: + cancelled.set() + raise + + runtime = make_async_runtime(on_change=slow_listener) + await runtime.start() + await asyncio.wait_for(started.wait(), 5) + await asyncio.wait_for(runtime.close(), 1) + + assert cancelled.is_set() + assert runtime._task is not None and runtime._task.done() + + +async def test_listeners_can_await_refresh_and_events_stay_in_order( + server: FlagServer, make_async_runtime: Make +) -> None: + server.respond(200, ORDERED) + log: List[str] = [] + holder: List[AsyncFeatureFlagsRuntimeClient] = [] + + async def on_change(event: FlagsChangedEvent) -> None: + label = "first-load" if event.previous is None else "refresh" + log.append(f"start {label}") + if event.previous is None: + server.respond(200, CHANGED) + await holder[0].refresh() + log.append(f"end {label}") + + runtime = make_async_runtime(on_change=on_change) + holder.append(runtime) + await asyncio.wait_for(runtime.refresh(), 2) + + assert log == ["start first-load", "end first-load", "start refresh", "end refresh"] + assert runtime.is_enabled("added") is True + + +@pytest.mark.parametrize( + "status, body, error_type", + [ + (404, {"message": "Not Found"}, NotFoundError), + (200, {"version": 3, "flags": {}}, FeatureFlagsConfigurationError), + ], +) +async def test_failures_keep_the_last_good_flags_and_recover( + server: FlagServer, + make_async_runtime: Make, + fast_polling: List[int], + status: int, + body: Any, + error_type: type, +) -> None: + server.respond(200, ORDERED) + errors: List[Exception] = [] + runtime = make_async_runtime(on_error=errors.append) + await runtime.start() + await runtime.wait_until_ready(timeout=5) + good = runtime.configuration + + server.respond(status, body) + await until(lambda: len(errors) >= 1) + assert isinstance(errors[0], error_type) + assert runtime.configuration is good + + server.respond(200, CHANGED) + await until(lambda: runtime.is_enabled("added")) + assert runtime.status.last_error is None + + +async def test_a_401_stops_polling( + server: FlagServer, make_async_runtime: Make, fast_polling: List[int] +) -> None: + server.respond(401, {"message": "Unauthorized"}) + runtime = make_async_runtime() + await runtime.start() + with pytest.raises(AuthenticationError): + await runtime.wait_until_ready(timeout=5) + await asyncio.sleep(0.05) + assert len(server.requests) == 1 + assert runtime.status.polling_stopped + + +async def test_a_401_from_refresh_stops_polling( + server: FlagServer, make_async_runtime: Make, fast_polling: List[int] +) -> None: + server.respond(200, ORDERED) + runtime = make_async_runtime() + await runtime.start() + await until(lambda: len(server.requests) >= 2) + + server.respond(401, {"message": "Unauthorized"}) + with pytest.raises(AuthenticationError): + await runtime.refresh() + assert runtime.status.polling_stopped + await until(lambda: runtime._task is not None and runtime._task.done()) + sent = len(server.requests) + await asyncio.sleep(0.05) + assert len(server.requests) == sent + assert runtime.is_enabled("ordered", {"user": "user_1"}) is True + + +async def test_a_429_retry_after_delays_the_next_poll( + server: FlagServer, make_async_runtime: Make, fast_polling: List[int] +) -> None: + server.enqueue((429, {"message": "slow down"}, {"Retry-After": "30"})) + server.respond(200, ORDERED) + runtime = make_async_runtime(request_timeout=1) + await runtime.start() + await until(lambda: len(server.requests) == 1) + await asyncio.sleep(0.3) + assert len(server.requests) == 1 + assert runtime.status.failed_refresh_count == 1 + + +async def test_a_403_keeps_polling( + server: FlagServer, make_async_runtime: Make, fast_polling: List[int] +) -> None: + server.respond(403, {"message": "Forbidden"}) + runtime = make_async_runtime() + await runtime.start() + await until(lambda: len(server.requests) >= 3) + assert runtime.status.polling_stopped is False + + +async def test_a_refresh_is_bounded_by_its_request_timeout( + server: FlagServer, make_async_runtime: Make +) -> None: + server.enqueue(FlagServer.BLOCK) + runtime = make_async_runtime(request_timeout=0.05) + with pytest.raises(WorkOSTimeoutError, match="request timeout"): + await runtime.refresh() + server.release() + assert isinstance(runtime.status.last_error, WorkOSTimeoutError) + assert runtime.status.failed_refresh_count == 1 + + +async def test_manual_refreshes_are_serialized( + server: FlagServer, make_async_runtime: Make +) -> None: + server.respond(200, ORDERED) + runtime = make_async_runtime() + await asyncio.gather(*(runtime.refresh() for _ in range(5))) + assert server.max_in_flight == 1 + assert runtime.status.refresh_count == 5 + + +async def test_close_stops_polling_and_keeps_the_last_flags( + server: FlagServer, make_async_runtime: Make, fast_polling: List[int] +) -> None: + server.respond(200, ORDERED) + runtime = make_async_runtime() + await runtime.start() + await until(lambda: len(server.requests) >= 2) + await runtime.close() + sent = len(server.requests) + await asyncio.sleep(0.05) + + assert len(server.requests) == sent + assert runtime.is_enabled("ordered", {"user": "user_1"}) is True + with pytest.raises(RuntimeError): + await runtime.start() + with pytest.raises(RuntimeError): + await runtime.refresh() + await runtime.close() + + +async def test_close_cancels_a_refresh_in_flight( + server: FlagServer, make_async_runtime: Make +) -> None: + server.enqueue(FlagServer.BLOCK) + runtime = make_async_runtime() + await runtime.start() + await until(lambda: len(server.requests) == 1) + await asyncio.wait_for(runtime.close(), 1) + server.release() + assert runtime.is_ready is False + + +async def test_every_close_waits_for_the_polling_task_to_finish( + server: FlagServer, make_async_runtime: Make +) -> None: + server.respond(200, ORDERED) + started = asyncio.Event() + cleaned_up = asyncio.Event() + + async def slow_listener(event: FlagsChangedEvent) -> None: + started.set() + try: + await asyncio.sleep(10) + finally: + await asyncio.sleep(0.05) + cleaned_up.set() + + runtime = make_async_runtime(on_change=slow_listener) + await runtime.start() + await asyncio.wait_for(started.wait(), 5) + first = asyncio.ensure_future(runtime.close()) + await asyncio.sleep(0) + await asyncio.wait_for(runtime.close(), 1) + + assert cleaned_up.is_set() + await first + + +async def test_close_can_be_called_from_a_listener( + server: FlagServer, make_async_runtime: Make +) -> None: + server.respond(200, ORDERED) + holder: List[AsyncFeatureFlagsRuntimeClient] = [] + closed = asyncio.Event() + + async def close_on_change(event: FlagsChangedEvent) -> None: + await holder[0].close() + closed.set() + + runtime = make_async_runtime(on_change=close_on_change) + holder.append(runtime) + await runtime.start() + await asyncio.wait_for(closed.wait(), 5) + await until(lambda: runtime._task is not None and runtime._task.done()) + + +async def test_context_manager_starts_and_closes( + server: FlagServer, async_client: AsyncWorkOSClient +) -> None: + server.respond(200, ORDERED) + async with async_client.create_feature_flags_runtime_client() as runtime: + await runtime.wait_until_ready(timeout=5) + assert runtime.is_closed diff --git a/tests/feature_flags_runtime/test_policy.py b/tests/feature_flags_runtime/test_policy.py index 56c6cd53..b45589d9 100644 --- a/tests/feature_flags_runtime/test_policy.py +++ b/tests/feature_flags_runtime/test_policy.py @@ -1,5 +1,5 @@ # @oagen-ignore-file -"""Polling delays and retry decisions of the runtime client.""" +"""Polling delays and retry decisions shared by both runtime clients.""" from __future__ import annotations diff --git a/tests/smoke_test.py b/tests/smoke_test.py index 03eee258..3ce23670 100644 --- a/tests/smoke_test.py +++ b/tests/smoke_test.py @@ -199,8 +199,10 @@ def test_pagination_importable() -> None: def test_feature_flags_runtime_client_importable() -> None: - """Verify the runtime client constructs without making a request.""" + """Verify both runtime clients construct without making a request.""" from workos import ( + AsyncFeatureFlagsRuntimeClient, + AsyncWorkOSClient, FeatureFlagsConfiguration, FeatureFlagsRuntimeClient, WorkOSClient, @@ -210,10 +212,14 @@ def test_feature_flags_runtime_client_importable() -> None: sync_runtime = WorkOSClient( api_key="sk_test_smoke" ).create_feature_flags_runtime_client(bootstrap=configuration) + async_runtime = AsyncWorkOSClient( + api_key="sk_test_smoke" + ).create_feature_flags_runtime_client(bootstrap=configuration) assert isinstance(sync_runtime, FeatureFlagsRuntimeClient) + assert isinstance(async_runtime, AsyncFeatureFlagsRuntimeClient) assert sync_runtime.is_enabled("missing", default=True) is True sync_runtime.close() - print("✓ Feature flags runtime client importable") + print("✓ Feature flags runtime clients importable") def test_dependencies_available() -> None: