diff --git a/README.md b/README.md index 92b24f49..465ed54b 100644 --- a/README.md +++ b/README.md @@ -279,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 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`. +**Processes.** With the SDK's default HTTP client, the sync client polls through an HTTP client of its own and closes it in `close()`; with an HTTP client you supplied, close the runtime client before closing that client. Except on macOS, a runtime client created before a fork (for example in an app gunicorn loads with `--preload`) can be started in each worker, and a sync client that was already polling resumes on new connections the first time it is used there. If the parent already sent requests through an HTTP client you supplied, or through the `AsyncWorkOSClient`, whether from the runtime client, any other SDK call, or another WorkOS client given the same HTTP client, the runtime client does not poll in the worker: `start()` and `refresh()` raise, it keeps serving its last configuration, and if it was polling it logs a warning the first time it is used. The SDK can't see requests your own code sends through an HTTP client you supplied, so don't use it before forking either. On macOS, once a process has created an HTTP client (every `WorkOSClient` has one), a child it forks can crash when it opens a network connection, so a runtime client created before the fork never polls in a worker there: `start()` logs a warning instead of raising, evaluation serves the bootstrap or last configuration (otherwise your defaults), `refresh()` raises, and so does `wait_until_ready()` without a configuration. Create the `WorkOSClient` and the runtime client after forking instead, for example in gunicorn's `post_fork` hook, or run without `--preload` locally. Background threads must be enabled, so run uWSGI with `--enable-threads`. ## Type Safety diff --git a/src/workos/_base_client.py b/src/workos/_base_client.py index ec98fc43..4f743d42 100644 --- a/src/workos/_base_client.py +++ b/src/workos/_base_client.py @@ -33,6 +33,7 @@ TransportConnectError, TransportError, TransportTimeout, + record_request, resolve_async_backend, resolve_sync_backend, ) @@ -533,6 +534,9 @@ def __init__( self._backend, self._owns_http_client = resolve_sync_backend( http_client, self._request_timeout ) + # Whether a request went out through the HTTP client, so a forked + # process can tell whether it would share the parent's connections. + self._sent_requests = False def close(self) -> None: """Release the HTTP client if the SDK created it. @@ -592,6 +596,8 @@ def request( timeout = self._resolve_timeout(request_options) max_retries = self._resolve_max_retries(request_options) last_error: Optional[Exception] = None + self._sent_requests = True + record_request(self._backend) for attempt in range(max_retries + 1): try: response = self._backend.request( @@ -777,6 +783,9 @@ def __init__( self._backend, self._owns_http_client = resolve_async_backend( http_client, self._request_timeout ) + # Whether a request went out through the HTTP client, so a forked + # process can tell whether it would share the parent's connections. + self._sent_requests = False async def close(self) -> None: """Release the HTTP client if the SDK created it. @@ -836,6 +845,8 @@ async def request( timeout = self._resolve_timeout(request_options) max_retries = self._resolve_max_retries(request_options) last_error: Optional[Exception] = None + self._sent_requests = True + record_request(self._backend) for attempt in range(max_retries + 1): try: response = await self._backend.request( diff --git a/src/workos/_http.py b/src/workos/_http.py index ca67746c..96cb71f1 100644 --- a/src/workos/_http.py +++ b/src/workos/_http.py @@ -12,6 +12,7 @@ import importlib import inspect import json +import weakref from dataclasses import dataclass from typing import ( Any, @@ -240,6 +241,42 @@ async def close(self) -> None: await self._client.aclose() +# --- connection tracking -------------------------------------------------------- + +# The HTTP clients requests have gone out through. Several WorkOS clients can be +# given the same HTTP client, so a forked process looks here to tell whether it +# would share the parent's connections. Keyed by id, since a caller's own +# backend need not be hashable; an entry goes away with its client. +_clients_with_requests: weakref.WeakValueDictionary[int, Any] = ( + weakref.WeakValueDictionary() +) + + +def _connection_owner(backend: object) -> object: + # WorkOS clients given the same httpx client each wrap it in a backend of + # their own; the httpx client holds the connections. + if isinstance(backend, (HttpxBackend, AsyncHttpxBackend)): + return backend._client + return backend + + +def record_request(backend: object) -> None: + """Note that a request is going out through ``backend``'s connections.""" + owner = _connection_owner(backend) + try: + _clients_with_requests[id(owner)] = owner + except TypeError: + # It can't be weakly referenced. Each WorkOS client still tracks its + # own requests. + pass + + +def sent_requests(backend: object) -> bool: + """Whether any WorkOS client sent a request through ``backend``'s connections.""" + owner = _connection_owner(backend) + return _clients_with_requests.get(id(owner)) is owner + + # --- resolution ----------------------------------------------------------------- diff --git a/src/workos/feature_flags/runtime/_async.py b/src/workos/feature_flags/runtime/_async.py index cbc80c33..90b023ae 100644 --- a/src/workos/feature_flags/runtime/_async.py +++ b/src/workos/feature_flags/runtime/_async.py @@ -7,15 +7,18 @@ import asyncio import inspect +import os import time +import weakref from types import TracebackType from typing import TYPE_CHECKING, Any, Awaitable, Callable, Optional, Type, Union from ..._errors import WorkOSTimeoutError -from . import _policy +from ..._http import sent_requests +from . import _core, _policy from ._configuration import FeatureFlagsConfiguration from ._context import LoggerLike -from ._core import Notification, RuntimeClientCore, utc_now +from ._core import Notification, RuntimeClientCore, fork_blocked_message, utc_now from ._types import FlagsChangedEvent, RuntimeStatus if TYPE_CHECKING: @@ -44,6 +47,9 @@ class AsyncFeatureFlagsRuntimeClient(RuntimeClientCore): The client polls through its ``AsyncWorkOSClient``, so use it on that client's event loop, and close it before closing the ``AsyncWorkOSClient``. + A client created before a fork can start in the child unless the + ``AsyncWorkOSClient`` already sent requests in the parent, or the platform + is macOS. """ def __init__( @@ -81,6 +87,11 @@ def __init__( # waiting for, if any. self._drainer: Optional[asyncio.Task[Any]] = None self._listener_task: Optional[asyncio.Future[Any]] = None + # Why polling cannot run in this forked child, if it cannot, and + # whether the platform is why. + self._fork_blocker: Optional[str] = None + self._fork_blocked_by_platform = False + self._fork_warned = False self._init_synchronization() if on_change is not None: self.add_change_listener(on_change) @@ -88,6 +99,7 @@ def __init__( self.add_error_listener(on_error) if on_status is not None: self.add_status_listener(on_status) + _clients.add(self) def _init_synchronization(self) -> None: # None of these binds to an event loop until first awaited. @@ -95,7 +107,7 @@ def _init_synchronization(self) -> None: 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: + if self._closed or self._status.ready or self._readiness_error is not None: self._ready.set() # --- Lifecycle ------------------------------------------------------------ @@ -105,15 +117,23 @@ async def start(self) -> None: 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. + polling permanently; the last good configuration stays usable. In a + forked child on macOS, a client created before the fork does not + poll: this logs a warning and returns. Raises: - RuntimeError: If the client is closed or no event loop is running. + RuntimeError: If the client is closed, no event loop is running, or + the client cannot poll in this forked process because the + parent used the ``AsyncWorkOSClient`` or its HTTP client. """ if self._closed: raise RuntimeError( "This feature flags runtime client is closed; create a new one" ) + if self._fork_blocked_by_platform: + self._skip_start_after_fork() + return + self._raise_if_fork_blocked() if self._started: return loop = asyncio.get_running_loop() @@ -138,10 +158,13 @@ async def wait_until_ready(self, timeout: Optional[float] = None) -> None: 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. + RuntimeError: If the client was closed before it became ready, or + has no configuration and cannot poll in this forked process. """ if self._status.ready: return + # Nothing can make the client ready in a forked child that cannot poll. + self._raise_if_fork_blocked() try: await asyncio.wait_for(self._ready.wait(), timeout) except asyncio.TimeoutError: @@ -159,12 +182,14 @@ async def refresh(self) -> None: Raises: WorkOSError: If the request fails or the payload cannot be read. - RuntimeError: If the client is closed. + RuntimeError: If the client is closed or cannot poll in this forked + process. """ if self._closed: raise RuntimeError( "This feature flags runtime client is closed; create a new one" ) + self._raise_if_fork_blocked() try: await self._refresh(background=False) except _Closed: @@ -286,7 +311,11 @@ async def _refresh(self, *, background: bool) -> None: await self._drain() async def _refresh_locked(self, *, background: bool) -> None: - if self._closed or (background and self._status.polling_stopped): + if ( + self._closed + or self._fork_blocker is not None + or (background and self._status.polling_stopped) + ): raise _Closed() attempted = utc_now() try: @@ -396,3 +425,96 @@ async def _await_listener(self, awaitable: Awaitable[Any]) -> None: self._logger.error("A feature flags runtime client listener was cancelled") else: listener_task.result() + + # --- Fork safety ---------------------------------------------------------- + + def _reinitialize_after_fork(self) -> None: + # Once the AsyncWorkOSClient's HTTP client has sent requests, its + # connections belong to the parent's event loop, so a copy of a client + # that polled, or whose HTTP client carried other calls (from this or + # another AsyncWorkOSClient), stops polling here (even if the child + # keeps running a copy of the loop) and keeps serving its last + # configuration. Otherwise it starts over on the child's own loop. + self._reinitialize_core_after_fork() + polling = ( + self._started and not self._closed and not self._status.polling_stopped + ) + # The polling task and anything waiting on these primitives belong to + # the parent's event loop, which may not run here: close() must not + # wait for that task. + self._task = None + self._drainer = None + self._listener_task = None + self._init_synchronization() + if self._closed: + self._fork_blocker = None + elif not _core.FORK_SAFE_PLATFORM: + self._fork_blocker = ( + f"{_core.MACOS_FORK_HAZARD}. Create the AsyncWorkOSClient and the " + "runtime client after forking instead, for example in gunicorn's " + "post_fork hook" + ) + elif ( + self._used_before_fork() + or self._client._sent_requests + or sent_requests(self._client._backend) + ): + self._fork_blocker = ( + "the AsyncWorkOSClient or its HTTP client was used in the parent, " + "on the parent's event loop and connections. Create the " + "AsyncWorkOSClient and the runtime client after forking instead" + ) + else: + self._fork_blocker = None + # On macOS every client created before the fork is blocked, however + # the parent used it, so start() warns instead of raising there. + self._fork_blocked_by_platform = ( + self._fork_blocker is not None and not _core.FORK_SAFE_PLATFORM + ) + # Only a client that was polling warns, on first use, that it no longer + # polls. On macOS any other warns once started; elsewhere start() + # raises instead. + self._fork_warned = not polling + + def _before_use(self) -> None: + self._warn_if_fork_blocked() + + def _warn_if_fork_blocked(self) -> None: + if self._fork_blocker is not None and not self._fork_warned: + self._fork_warned = True + self._logger.warning( + "%s. Evaluation keeps serving the last configuration.", + fork_blocked_message(self._fork_blocker), + ) + + def _raise_if_fork_blocked(self) -> None: + if self._fork_blocker is not None: + raise RuntimeError(fork_blocked_message(self._fork_blocker)) + + def _skip_start_after_fork(self) -> None: + # A client that was polling warns once, on first use. Starting one + # that was not makes it such a client, which has not warned yet. + if not self._started: + self._started = True + self._fork_warned = False + self._warn_if_fork_blocked() + + +_clients: weakref.WeakSet[AsyncFeatureFlagsRuntimeClient] = weakref.WeakSet() + + +def _reinitialize_clients_after_fork() -> None: + for client in list(_clients): + # One client that fails must not leave the others unprepared. + try: + client._reinitialize_after_fork() + except Exception: + client._logger.exception( + "Could not prepare a feature flags runtime client for forked " + "process %d", + os.getpid(), + ) + + +if hasattr(os, "register_at_fork"): + os.register_at_fork(after_in_child=_reinitialize_clients_after_fork) diff --git a/src/workos/feature_flags/runtime/_context.py b/src/workos/feature_flags/runtime/_context.py index 7920155a..a3bb333e 100644 --- a/src/workos/feature_flags/runtime/_context.py +++ b/src/workos/feature_flags/runtime/_context.py @@ -66,6 +66,9 @@ def warn(self, reason: str, key: str, message: str) -> None: self._seen.add((reason, key)) self._logger.warning(message) + def reinitialize_after_fork(self) -> None: + self._lock = threading.Lock() + def normalize_context(context: object, warnings: ContextWarnings) -> dict[str, str]: """Reduce an evaluation context to valid target type to ID pairs. diff --git a/src/workos/feature_flags/runtime/_core.py b/src/workos/feature_flags/runtime/_core.py index aab608e7..19497118 100644 --- a/src/workos/feature_flags/runtime/_core.py +++ b/src/workos/feature_flags/runtime/_core.py @@ -7,6 +7,8 @@ import copy import logging +import os +import sys import threading from collections import deque from dataclasses import replace @@ -33,6 +35,23 @@ Listener = Callable[[Any], object] Notification = Tuple[Listener, Any] +# On macOS, once a process has created an HTTP client, a child it forks without +# exec'ing can crash inside system libraries (DNS resolution, proxy settings, +# the system trust store) when it opens a network connection. Every +# WorkOSClient has an HTTP client, so a runtime client created before the fork +# never polls in the child there. +FORK_SAFE_PLATFORM = sys.platform != "darwin" +MACOS_FORK_HAZARD = ( + "on macOS, a forked process cannot safely open network connections once its " + "parent has created an HTTP client" +) + + +def fork_blocked_message(reason: str) -> str: + return ( + f"Feature flag polling does not run in forked process {os.getpid()}: {reason}" + ) + def utc_now() -> datetime: return datetime.now(timezone.utc) @@ -114,6 +133,7 @@ def is_enabled( "organization": "org_123", "workspace": "ws_1"}``. default: Returned when the flag has no answer. """ + self._before_use() configuration = self._status.configuration flag = configuration._compiled(slug) if configuration is not None else None if flag is None: @@ -130,6 +150,7 @@ def get_all_flags( omitted so that ``flags.get(slug, my_default)`` applies your default. An unready client returns an empty dict. """ + self._before_use() configuration = self._status.configuration if configuration is None: return {} @@ -143,27 +164,32 @@ def get_all_flags( def get_flag(self, slug: str) -> Optional[FlagDefinition]: """Return one flag's configuration from the current snapshot, for inspection.""" + self._before_use() configuration = self._status.configuration return configuration.get_flag(slug) if configuration is not None else None @property def status(self) -> RuntimeStatus: """Readiness, the current snapshot, refresh timestamps, counts and the last error.""" + self._before_use() return self._status @property def configuration(self) -> Optional[FeatureFlagsConfiguration]: """The current snapshot, suitable for persisting with ``to_json()``.""" + self._before_use() return self._status.configuration @property def is_ready(self) -> bool: """Whether a validated configuration is available for local evaluation.""" + self._before_use() return self._status.configuration is not None @property def is_closed(self) -> bool: """Whether ``close()`` has been called.""" + self._before_use() return self._closed # --- Listeners ------------------------------------------------------------ @@ -297,6 +323,21 @@ def _readiness_failure(self) -> Exception: error.__suppress_context__ = True return error + def _before_use(self) -> None: + """Hook run before public calls; the sync client restarts polling after a fork.""" + + def _used_before_fork(self) -> bool: + """Whether this client sent requests, or was polling, in the parent.""" + return self._started or self._status.refresh_count > 0 + + def _reinitialize_core_after_fork(self) -> None: + # The child has only the forking thread, so these locks may have been + # held by threads that no longer exist. + self._status_lock = threading.Lock() + self._listeners_lock = threading.Lock() + self._context_warnings.reinitialize_after_fork() + self._pending.clear() + def _log_listener_error(self) -> None: self._logger.exception("A feature flags runtime client listener raised") diff --git a/src/workos/feature_flags/runtime/_sync.py b/src/workos/feature_flags/runtime/_sync.py index a8f03d5c..cce0c365 100644 --- a/src/workos/feature_flags/runtime/_sync.py +++ b/src/workos/feature_flags/runtime/_sync.py @@ -6,15 +6,23 @@ from __future__ import annotations import inspect +import os import threading import time +import weakref from types import TracebackType from typing import TYPE_CHECKING, Any, Callable, Optional, Type -from . import _policy +from . import _core, _policy from ._configuration import FeatureFlagsConfiguration from ._context import LoggerLike -from ._core import Listener, Notification, RuntimeClientCore, utc_now +from ._core import ( + Listener, + Notification, + RuntimeClientCore, + fork_blocked_message, + utc_now, +) from ._transport import PollTransport from ._types import FlagsChangedEvent, RuntimeStatus @@ -38,9 +46,13 @@ class FeatureFlagsRuntimeClient(RuntimeClientCore): flags.wait_until_ready(timeout=5) flags.is_enabled("new-checkout", {"organization": "org_123"}) - Polling uses the ``WorkOSClient``'s API key, base URL, headers, retry - settings and HTTP client, so close this client before closing the - ``WorkOSClient`` or an HTTP client you supplied. + Polling uses the ``WorkOSClient``'s API key, base URL, headers and retry + settings. With the SDK's default HTTP client it polls through an HTTP + client of its own, which a forked child replaces; with an HTTP client you + supplied, close this client before closing yours. A client created before + a fork can start in the child, and one that was polling resumes there, + except through an HTTP client you supplied that the parent already sent + requests through, and on macOS. """ def __init__( @@ -95,6 +107,12 @@ def __init__( self._transport = PollTransport(client) self._thread: Optional[threading.Thread] = None self._delivering: Optional[int] = None + self._fork_restart_pending = False + # Why polling cannot run in this forked child, if it cannot, and + # whether the platform is why. + self._fork_blocker: Optional[str] = None + self._fork_blocked_by_platform = False + self._fork_warned = False self._init_synchronization() if on_change is not None: self.add_change_listener(on_change) @@ -102,6 +120,7 @@ def __init__( self.add_error_listener(on_error) if on_status is not None: self.add_status_listener(on_status) + _clients.add(self) def _init_synchronization(self) -> None: self._lifecycle_lock = threading.Lock() @@ -112,8 +131,11 @@ def _init_synchronization(self) -> None: self._stop = threading.Event() self._wake = threading.Event() self._ready = threading.Event() - if self._status.ready: + if self._closed or self._status.ready or self._readiness_error is not None: self._ready.set() + if self._closed: + self._stop.set() + self._wake.set() # --- Lifecycle ------------------------------------------------------------ @@ -122,16 +144,25 @@ def start(self) -> None: 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. + polling permanently; the last good configuration stays usable. In a + forked child on macOS, a client created before the fork does not + poll: this logs a warning and returns. Raises: - RuntimeError: If the client is closed. + RuntimeError: If the client is closed, or cannot poll in this forked + process because the parent sent requests through the HTTP + client you supplied. """ + self._before_use() with self._lifecycle_lock: if self._closed: raise RuntimeError( "This feature flags runtime client is closed; create a new one" ) + if self._fork_blocked_by_platform: + self._skip_start_after_fork() + return + self._raise_if_fork_blocked() if self._started: return self._started = True @@ -156,10 +187,14 @@ def wait_until_ready(self, timeout: Optional[float] = None) -> None: 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. + RuntimeError: If the client was closed before it became ready, or + has no configuration and cannot poll in this forked process. """ + self._before_use() if self._status.ready: return + # Nothing can make the client ready in a forked child that cannot poll. + self._raise_if_fork_blocked() if not self._ready.wait(timeout): raise self._readiness_timeout(timeout) if not self._status.ready: @@ -174,12 +209,15 @@ def refresh(self) -> None: Raises: WorkOSError: If the request fails or the payload cannot be read. - RuntimeError: If the client is closed. + RuntimeError: If the client is closed or cannot poll in this forked + process. """ + self._before_use() if self._closed: raise RuntimeError( "This feature flags runtime client is closed; create a new one" ) + self._raise_if_fork_blocked() try: self._refresh(background=False) except _Closed: @@ -194,7 +232,8 @@ def refresh(self) -> None: def close(self) -> None: """Stop polling and release the background thread. Idempotent. - Waits up to ``request_timeout`` for an in-flight poll to finish. A poll + Waits up to ``request_timeout`` for an in-flight poll to finish, then + closes the HTTP client polling used, if the runtime created it. A poll still running after that, such as one reading a stalled response, is abandoned: it ends on its own and changes nothing. Every caller waits, including one that calls while another close() is still waiting. @@ -214,6 +253,7 @@ def close(self) -> None: and thread.is_alive() ): thread.join(self._request_timeout + 1.0) + self._transport.close() def __enter__(self) -> FeatureFlagsRuntimeClient: self.start() @@ -412,8 +452,97 @@ def _deliver(self, notifications: list[Notification]) -> None: "async listeners" ) + # --- Fork safety ---------------------------------------------------------- + + def _before_use(self) -> None: + if self._fork_restart_pending: + self._restart_after_fork() + else: + self._warn_if_fork_blocked() + + def _warn_if_fork_blocked(self) -> None: + if self._fork_blocker is not None and not self._fork_warned: + self._fork_warned = True + self._logger.warning( + "%s. Evaluation keeps serving the last configuration.", + fork_blocked_message(self._fork_blocker), + ) + + def _raise_if_fork_blocked(self) -> None: + if self._fork_blocker is not None: + raise RuntimeError(fork_blocked_message(self._fork_blocker)) + + def _skip_start_after_fork(self) -> None: + # A client that was polling warns once, on first use. Starting one + # that was not makes it such a client, which has not warned yet. + if not self._started: + self._started = True + self._fork_warned = False + self._warn_if_fork_blocked() + + def _restart_after_fork(self) -> None: + with self._lifecycle_lock: + if not self._fork_restart_pending: + return + self._fork_restart_pending = False + if self._closed: + return + self._logger.info( + "Restarting feature flag polling in forked process %d", os.getpid() + ) + self._spawn_poller() + + def _reinitialize_after_fork(self) -> None: + # The child has only the forking thread. Locks may have been held by + # threads that no longer exist, the poller thread is gone, and the + # parent's HTTP connections must not be shared, so rebuild all of it + # and restart polling, with a new HTTP client, on the next use. + self._reinitialize_core_after_fork() + self._transport.reinitialize_after_fork() + self._delivering = None + self._thread = None + self._init_synchronization() + polling = ( + self._started and not self._closed and not self._status.polling_stopped + ) + self._fork_blocker = ( + None + if self._closed + else self._transport.fork_blocker(self._used_before_fork()) + ) + # On macOS every client created before the fork is blocked, however + # the parent used it, so start() warns instead of raising there. + self._fork_blocked_by_platform = ( + self._fork_blocker is not None and not _core.FORK_SAFE_PLATFORM + ) + self._fork_restart_pending = polling and self._fork_blocker is None + # Only a client that was polling warns, on first use, that it no longer + # polls. On macOS any other warns once started; elsewhere start() + # raises instead. + self._fork_warned = not polling + def _discard_awaitable(awaitable: object) -> None: close = getattr(awaitable, "close", None) if callable(close): close() + + +_clients: weakref.WeakSet[FeatureFlagsRuntimeClient] = weakref.WeakSet() + + +def _reinitialize_clients_after_fork() -> None: + for client in list(_clients): + # One client that fails must not leave the others unprepared. + try: + client._reinitialize_after_fork() + except Exception: + client._logger.exception( + "Could not prepare a feature flags runtime client for forked " + "process %d", + os.getpid(), + ) + + +if hasattr(os, "register_at_fork"): + os.register_at_fork(after_in_child=_reinitialize_clients_after_fork) diff --git a/src/workos/feature_flags/runtime/_transport.py b/src/workos/feature_flags/runtime/_transport.py index 50f91820..067c4502 100644 --- a/src/workos/feature_flags/runtime/_transport.py +++ b/src/workos/feature_flags/runtime/_transport.py @@ -1,14 +1,17 @@ # @oagen-ignore-file # This file is hand-maintained. The sync runtime client's poll request, sent -# with a total time budget. +# with a total time budget and without sharing connections across a fork. # pyright: strict, reportPrivateUsage=false from __future__ import annotations +import threading import time from typing import TYPE_CHECKING, Any, Dict, Iterator, Optional from urllib.parse import urlsplit +import httpx2 + from ..._errors import WorkOSConnectionError, WorkOSError, WorkOSTimeoutError from ..._http import ( HTTPResponse, @@ -17,8 +20,10 @@ TransportError, TransportTimeout, _HttpxFamily, + record_request, + sent_requests, ) -from . import _policy +from . import _core, _policy if TYPE_CHECKING: from _collections_abc import dict_items, dict_values @@ -77,13 +82,20 @@ class PollTransport: """Sends ``GET /sdk/feature-flags`` for the sync runtime client. The request carries the same URL, headers and User-Agent as any other - request from the ``WorkOSClient``, goes through its HTTP client, and maps - responses and failures to the same errors. Unlike other requests, a refresh - has a total budget. The body is streamed, every network wait is limited to - the time left, and the response is abandoned once the deadline passes, so a - refresh ends within ``request_timeout`` plus at most one read of a stalling - response body. A caller-supplied ``HTTPBackend`` that is not an httpx client - cannot stream; it gets the remaining budget as its per-request timeout. + request from the ``WorkOSClient`` and maps responses and failures to the + same errors. Two things differ: + + - A refresh has a total budget. The body is streamed, every network wait is + limited to the time left, and the response is abandoned once the deadline + passes, so a refresh ends within ``request_timeout`` plus at most one read + of a stalling response body. A caller-supplied ``HTTPBackend`` that is + not an httpx client cannot stream; it gets the remaining budget as its + per-request timeout. + - With the SDK's default HTTP client, polling uses its own httpx2 client + with the same settings, so a forked child can replace it instead of + sharing the parent's connections. A caller-supplied client cannot be + rebuilt that way, so a child does not poll through one the parent sent + requests through. On macOS, a child does not poll at all. """ def __init__(self, workos: WorkOSClient) -> None: @@ -91,9 +103,50 @@ def __init__(self, workos: WorkOSClient) -> None: backend = workos._backend self._shared_httpx: Any = None self._family = _HttpxFamily("httpx2") - if isinstance(backend, HttpxBackend): + if isinstance(backend, HttpxBackend) and not workos._owns_http_client: self._shared_httpx = backend._client self._family = backend._family + self._uses_default_client = workos._owns_http_client + self._owned: Optional[httpx2.Client] = None + self._closed = False + self._lock = threading.Lock() + + def fork_blocker(self, used_before_fork: bool) -> Optional[str]: + """Why a forked child cannot poll safely, or ``None`` if it can.""" + if not _core.FORK_SAFE_PLATFORM: + return ( + f"{_core.MACOS_FORK_HAZARD}. Create the WorkOSClient and the " + "runtime client after forking instead, for example in gunicorn's " + "post_fork hook" + ) + # The WorkOSClient's other requests count too, and so do those of any + # other WorkOSClient given the same HTTP client: they open the same + # connections. + if ( + used_before_fork + or self._workos._sent_requests + or sent_requests(self._workos._backend) + ) and not self._uses_default_client: + return ( + "the parent used its caller-supplied HTTP client, whose " + "connections this process would share. Create that HTTP client, " + "the WorkOSClient and the runtime client after forking instead" + ) + return None + + def reinitialize_after_fork(self) -> None: + # Abandon, never close, the inherited client: its sockets are the + # parent's connections, and closing it could wait on a pool lock that + # a thread missing from this process held at fork time. + self._owned = None + self._lock = threading.Lock() + + def close(self) -> None: + with self._lock: + self._closed = True + owned, self._owned = self._owned, None + if owned is not None: + owned.close() def get(self, deadline: float) -> Any: """Send one poll request and return the decoded body.""" @@ -115,7 +168,10 @@ def get(self, deadline: float) -> Any: def _send(self, url: str, headers: Dict[str, str], deadline: float) -> HTTPResponse: if deadline - time.monotonic() <= 0: raise TransportTimeout("no time was left in the request timeout") - client = self._shared_httpx + if not self._uses_default_client: + self._workos._sent_requests = True + record_request(self._workos._backend) + client = self._httpx_client() if client is None: return self._workos._backend.request( "GET", @@ -133,6 +189,23 @@ def _send(self, url: str, headers: Dict[str, str], deadline: float) -> HTTPRespo except self._family.http_error as error: raise TransportError(str(error)) from error + def _httpx_client(self) -> Any: + if self._shared_httpx is not None: + return self._shared_httpx + if not self._uses_default_client: + return None + with self._lock: + # A poll that close() stopped waiting for must not open a client + # that nothing would close. + if self._closed: + raise RuntimeError("The feature flags runtime client is closed") + if self._owned is None: + # The settings resolve_sync_backend gives the SDK's default client. + self._owned = httpx2.Client( + timeout=self._workos._request_timeout, follow_redirects=True + ) + return self._owned + def _stream( self, client: Any, url: str, headers: Dict[str, str], deadline: float ) -> HTTPResponse: diff --git a/tests/feature_flags_runtime/helpers.py b/tests/feature_flags_runtime/helpers.py index 24a5a91b..c106ff3f 100644 --- a/tests/feature_flags_runtime/helpers.py +++ b/tests/feature_flags_runtime/helpers.py @@ -236,7 +236,7 @@ class LocalHTTPServer: ``mode`` selects the next responses: ``"ok"`` serves ``payload``, ``"stall"`` waits ``stall`` seconds before answering, and ``"drip"`` sends the body one byte every ``drip_interval`` seconds. Every request is - recorded with its headers. + recorded with its headers and the client's TCP port. """ def __init__(self, payload: Dict[str, Any]) -> None: @@ -268,6 +268,14 @@ def do_GET(self) -> None: def base_url(self) -> str: return f"http://127.0.0.1:{self._server.server_address[1]}" + def ports(self, authorization: str) -> set: + with self._lock: + return { + request["port"] + for request in self.requests + if request["headers"].get("Authorization") == authorization + } + def close(self) -> None: self._server.shutdown() self._server.server_close() @@ -278,6 +286,7 @@ def _handle(self, handler: Any) -> None: { "path": handler.path, "headers": dict(handler.headers.items()), + "port": handler.client_address[1], } ) mode, stall, interval = self.mode, self.stall, self.drip_interval diff --git a/tests/feature_flags_runtime/test_fork.py b/tests/feature_flags_runtime/test_fork.py new file mode 100644 index 00000000..3ff000c1 --- /dev/null +++ b/tests/feature_flags_runtime/test_fork.py @@ -0,0 +1,761 @@ +# @oagen-ignore-file +"""What runtime clients do in a forked child process. + +Each scenario runs in a subprocess so a hang or a fork cannot wedge pytest. +""" + +from __future__ import annotations + +import asyncio +import logging +import os +import subprocess +import sys +import textwrap +import time +import weakref +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any, Callable, Dict, List, Optional + +import httpx2 +import pytest + +from workos import AsyncWorkOSClient, HTTPResponse, WorkOSClient +from workos.feature_flags.runtime import ( + AsyncFeatureFlagsRuntimeClient, + FeatureFlagsConfiguration, + FeatureFlagsRuntimeClient, + _async, + _core, + _sync, +) + +from tests.feature_flags_runtime.helpers import ( + API_KEY, + FlagServer, + LocalHTTPServer, + flag, + payload, +) + +ROOT = Path(__file__).resolve().parents[2] +FORK_BLOCKED = "does not run in forked process" +BOOTSTRAP = payload(f=flag("f", default_value=True)) + +Make = Callable[..., FeatureFlagsRuntimeClient] +MakeAsync = Callable[..., AsyncFeatureFlagsRuntimeClient] + +PRELUDE = """ +import asyncio +import os +import time +import warnings + +from workos import AsyncWorkOSClient, WorkOSClient +from workos.feature_flags.runtime import _policy +from tests.feature_flags_runtime.helpers import LocalHTTPServer, flag, payload + +_policy.next_delay = lambda interval, failures, max_retry_delay: 0.05 +server = LocalHTTPServer(payload(f=flag("f", default_value=True))) + + +def wait_for(predicate, timeout=5.0): + deadline = time.monotonic() + timeout + while not predicate(): + if time.monotonic() > deadline: + raise SystemExit("timed out") + time.sleep(0.01) + + +def fork(): + with warnings.catch_warnings(): + warnings.simplefilter("ignore", DeprecationWarning) + return os.fork() + + +def finish(pid): + _, status = os.waitpid(pid, 0) + raise SystemExit(os.waitstatus_to_exitcode(status)) +""" + +# The parent and child poll the parent's server with different API keys, so +# the server can tell whose connection each request came on. +DEFAULT_CLIENT = """ +workos = WorkOSClient(api_key="sk_parent", base_url=server.base_url) +runtime = workos.create_feature_flags_runtime_client() +runtime.start() +runtime.wait_until_ready(timeout=5) +wait_for(lambda: len(server.ports("Bearer sk_parent")) >= 1 and len(server.requests) >= 3) +pid = fork() +if pid == 0: + workos._api_key = "sk_child" + # Reading any property is a use of the client that resumes polling. + assert runtime.is_ready + time.sleep(1.0) + runtime.close() + os._exit(0) +_, status = os.waitpid(pid, 0) +assert os.waitstatus_to_exitcode(status) == 0 +time.sleep(0.2) +runtime.close() +child_requests = [r for r in server.requests if r["headers"].get("Authorization") == "Bearer sk_child"] +parent_ports = server.ports("Bearer sk_parent") +child_ports = server.ports("Bearer sk_child") +assert len(child_requests) >= 2, f"the child sent {len(child_requests)} polls" +assert not parent_ports & child_ports, f"shared connections: {parent_ports & child_ports}" +print(f"parent ports {sorted(parent_ports)}, child ports {sorted(child_ports)}") +""" + +CALLER_CLIENT = """ +import httpx2 +import logging + +logging.basicConfig(level=logging.WARNING) +workos = WorkOSClient(api_key="sk_parent", base_url=server.base_url, http_client=httpx2.Client()) +runtime = workos.create_feature_flags_runtime_client() +runtime.start() +runtime.wait_until_ready(timeout=5) +pid = fork() +if pid == 0: + workos._api_key = "sk_child" + assert runtime.is_enabled("f") is True + try: + runtime.refresh() + os._exit(2) + except RuntimeError as error: + assert "caller-supplied HTTP client" in str(error), error + time.sleep(0.5) + os._exit(0) +_, status = os.waitpid(pid, 0) +assert os.waitstatus_to_exitcode(status) == 0 +runtime.close() +assert server.ports("Bearer sk_child") == set(), "the child polled" +""" + +# asyncio refuses to run a loop inherited across fork, so the parent runs its +# loop on a thread (which the child does not inherit) and forks from the main +# thread, as a pre-fork server would. +ASYNC_CLIENT = """ +import signal +import threading + +loop = asyncio.new_event_loop() +threading.Thread(target=loop.run_forever, daemon=True).start() + + +def call(coroutine): + return asyncio.run_coroutine_threadsafe(coroutine, loop).result(10) + + +workos = AsyncWorkOSClient(api_key="sk_parent", base_url=server.base_url) +runtime = call(_make_runtime(workos)) +call(runtime.start()) +call(runtime.wait_until_ready(timeout=5)) +pid = fork() +if pid == 0: + workos._api_key = "sk_child" + assert runtime.is_enabled("f") is True + for method in ("refresh", "start"): + try: + asyncio.run(getattr(runtime, method)()) + os._exit(2) + except RuntimeError as error: + assert "was used in the parent" in str(error), error + # The parent's polling task is on a loop that does not run here, so + # close() must not wait for it. + signal.signal(signal.SIGALRM, lambda *args: os._exit(3)) + signal.alarm(5) + asyncio.run(runtime.close()) + signal.alarm(0) + time.sleep(0.5) + os._exit(0) +_, status = os.waitpid(pid, 0) +assert os.waitstatus_to_exitcode(status) == 0 +call(runtime.close()) +call(workos.close()) +assert server.ports("Bearer sk_child") == set(), "the child polled" +""" + +MAKE_ASYNC_RUNTIME = """ +async def _make_runtime(workos): + return workos.create_feature_flags_runtime_client() +""" + +MACOS_DEFAULT_CLIENT = """ +import logging + +logging.basicConfig(level=logging.WARNING) +workos = WorkOSClient(api_key="sk_parent", base_url=server.base_url) +runtime = workos.create_feature_flags_runtime_client() +runtime.start() +runtime.wait_until_ready(timeout=5) +pid = fork() +if pid == 0: + workos._api_key = "sk_child" + assert runtime.is_enabled("f") is True + runtime.start() + try: + runtime.refresh() + os._exit(2) + except RuntimeError as error: + assert "cannot safely open network connections" in str(error), error + time.sleep(0.5) + os._exit(0) +_, status = os.waitpid(pid, 0) +assert os.waitstatus_to_exitcode(status) == 0 +runtime.close() +assert server.ports("Bearer sk_child") == set(), "the child polled" +""" + +# The app is preloaded and each worker starts the runtime clients, as with +# gunicorn --preload. +MACOS_PRELOADED = """ +import logging +import signal + +from workos import FeatureFlagsConfiguration + +logging.basicConfig(level=logging.WARNING) +bootstrap = FeatureFlagsConfiguration.from_dict(payload(f=flag("f", default_value=True))) +workos = WorkOSClient(api_key="sk_parent", base_url=server.base_url) +bootstrapped = workos.create_feature_flags_runtime_client(bootstrap=bootstrap) +empty = workos.create_feature_flags_runtime_client() +async_workos = AsyncWorkOSClient(api_key="sk_parent", base_url=server.base_url) +async_bootstrapped = async_workos.create_feature_flags_runtime_client(bootstrap=bootstrap) +pid = fork() +if pid == 0: + workos._api_key = "sk_child" + async_workos._api_key = "sk_child" + signal.signal(signal.SIGALRM, lambda *args: os._exit(3)) + signal.alarm(10) + bootstrapped.start() + bootstrapped.wait_until_ready() + assert bootstrapped.is_enabled("f") is True + empty.start() + try: + empty.wait_until_ready() + os._exit(2) + except RuntimeError as error: + assert "does not run in forked process" in str(error), error + + async def serve(): + await async_bootstrapped.start() + await async_bootstrapped.wait_until_ready() + assert async_bootstrapped.is_enabled("f") is True + + asyncio.run(serve()) + signal.alarm(0) + time.sleep(0.5) + os._exit(0) +_, status = os.waitpid(pid, 0) +assert os.waitstatus_to_exitcode(status) == 0 +assert server.ports("Bearer sk_child") == set(), "the child polled" +""" + +# The parent never polls, as when an app is preloaded and each worker starts +# the runtime client. +CALLER_CLIENT_UNUSED = """ +import httpx2 + +workos = WorkOSClient(api_key="sk_parent", base_url=server.base_url, http_client=httpx2.Client()) +runtime = workos.create_feature_flags_runtime_client() +pid = fork() +if pid == 0: + workos._api_key = "sk_child" + runtime.start() + runtime.wait_until_ready(timeout=5) + runtime.close() + os._exit(0) +_, status = os.waitpid(pid, 0) +assert os.waitstatus_to_exitcode(status) == 0 +assert server.ports("Bearer sk_child"), "the child never polled" +""" + +ASYNC_CLIENT_UNUSED = """ +workos = AsyncWorkOSClient(api_key="sk_parent", base_url=server.base_url) +runtime = workos.create_feature_flags_runtime_client() +pid = fork() +if pid == 0: + workos._api_key = "sk_child" + + async def serve(): + await runtime.start() + await runtime.wait_until_ready(timeout=5) + await runtime.close() + await workos.close() + + asyncio.run(serve()) + os._exit(0) +_, status = os.waitpid(pid, 0) +assert os.waitstatus_to_exitcode(status) == 0 +assert server.ports("Bearer sk_child"), "the child never polled" +""" + +# The runtime client never polls in the parent, but another SDK call does, on +# the HTTP client the runtime client would poll through. +CALLER_CLIENT_SDK_USED = """ +import httpx2 + +workos = WorkOSClient(api_key="sk_parent", base_url=server.base_url, http_client=httpx2.Client()) +try: + workos.organizations.list_organizations() +except Exception: + pass +assert server.ports("Bearer sk_parent"), "the parent sent no request" +runtime = workos.create_feature_flags_runtime_client() +pid = fork() +if pid == 0: + workos._api_key = "sk_child" + try: + runtime.start() + os._exit(2) + except RuntimeError as error: + assert "caller-supplied HTTP client" in str(error), error + os._exit(0) +_, status = os.waitpid(pid, 0) +assert os.waitstatus_to_exitcode(status) == 0 +assert server.ports("Bearer sk_child") == set(), "the child polled" +""" + +ASYNC_CLIENT_SDK_USED = """ +workos = AsyncWorkOSClient(api_key="sk_parent", base_url=server.base_url) + + +async def call_api(): + try: + await workos.organizations.list_organizations() + except Exception: + pass + + +asyncio.run(call_api()) +assert server.ports("Bearer sk_parent"), "the parent sent no request" +runtime = workos.create_feature_flags_runtime_client() +pid = fork() +if pid == 0: + workos._api_key = "sk_child" + try: + asyncio.run(runtime.start()) + os._exit(2) + except RuntimeError as error: + assert "was used in the parent" in str(error), error + os._exit(0) +_, status = os.waitpid(pid, 0) +assert os.waitstatus_to_exitcode(status) == 0 +assert server.ports("Bearer sk_child") == set(), "the child polled" +""" + +# Another WorkOS client sends the request, through an HTTP client it shares +# with the one the runtime client would poll through. +CALLER_CLIENT_SHARED_SDK_USED = """ +import httpx2 + +shared = httpx2.Client() +other = WorkOSClient(api_key="sk_parent", base_url=server.base_url, http_client=shared) +try: + other.organizations.list_organizations() +except Exception: + pass +assert server.ports("Bearer sk_parent"), "the parent sent no request" +workos = WorkOSClient(api_key="sk_parent", base_url=server.base_url, http_client=shared) +runtime = workos.create_feature_flags_runtime_client() +pid = fork() +if pid == 0: + workos._api_key = "sk_child" + try: + runtime.start() + os._exit(2) + except RuntimeError as error: + assert "caller-supplied HTTP client" in str(error), error + os._exit(0) +_, status = os.waitpid(pid, 0) +assert os.waitstatus_to_exitcode(status) == 0 +assert server.ports("Bearer sk_child") == set(), "the child polled" +""" + +ASYNC_CLIENT_SHARED_SDK_USED = """ +import httpx2 + +shared = httpx2.AsyncClient() +other = AsyncWorkOSClient(api_key="sk_parent", base_url=server.base_url, http_client=shared) + + +async def call_api(): + try: + await other.organizations.list_organizations() + except Exception: + pass + + +asyncio.run(call_api()) +assert server.ports("Bearer sk_parent"), "the parent sent no request" +workos = AsyncWorkOSClient(api_key="sk_parent", base_url=server.base_url, http_client=shared) +runtime = workos.create_feature_flags_runtime_client() +pid = fork() +if pid == 0: + workos._api_key = "sk_child" + try: + asyncio.run(runtime.start()) + os._exit(2) + except RuntimeError as error: + assert "was used in the parent" in str(error), error + os._exit(0) +_, status = os.waitpid(pid, 0) +assert os.waitstatus_to_exitcode(status) == 0 +assert server.ports("Bearer sk_child") == set(), "the child polled" +""" + +# On macOS a forked child never polls, so to check what runs elsewhere, this +# pretends the platform is fork-safe. The scenarios reach the server by IP over +# plain HTTP, and NO_PROXY keeps httpx2 from reading macOS proxy settings, so +# the child opens connections without the system libraries that crash there. +FORCE_FORK_SAFE = """ +from workos.feature_flags.runtime import _core + +_core.FORK_SAFE_PLATFORM = True +""" + + +def run(*scenario: str) -> subprocess.CompletedProcess: + script = "".join(textwrap.dedent(part) for part in (PRELUDE, *scenario)) + return subprocess.run( + [sys.executable, "-c", script], + cwd=ROOT, + env={**os.environ, "NO_PROXY": "*"}, + capture_output=True, + text=True, + timeout=60, + ) + + +pytestmark = pytest.mark.skipif(not hasattr(os, "fork"), reason="requires os.fork") + + +def test_a_child_polls_on_its_own_connections_once_used() -> None: + result = run(FORCE_FORK_SAFE, DEFAULT_CLIENT) + assert result.returncode == 0, result.stdout + result.stderr + + +@pytest.mark.skipif(sys.platform != "darwin", reason="macOS behavior") +def test_on_macos_a_child_does_not_poll() -> None: + result = run(MACOS_DEFAULT_CLIENT) + assert result.returncode == 0, result.stdout + result.stderr + assert "cannot safely open network connections" in result.stderr + assert result.stderr.count(FORK_BLOCKED) == 1, result.stderr + + +@pytest.mark.skipif(sys.platform != "darwin", reason="macOS behavior") +def test_on_macos_a_preloaded_client_serves_its_bootstrap_in_a_child() -> None: + result = run(MACOS_PRELOADED) + assert result.returncode == 0, result.stdout + result.stderr + # One warning from each start(). + assert result.stderr.count(FORK_BLOCKED) == 3, result.stderr + + +def test_a_child_does_not_poll_through_a_caller_supplied_client() -> None: + result = run(FORCE_FORK_SAFE, CALLER_CLIENT) + assert result.returncode == 0, result.stdout + result.stderr + assert "does not run in forked process" in result.stderr + + +def test_an_async_client_stops_polling_in_a_child() -> None: + result = run(FORCE_FORK_SAFE, MAKE_ASYNC_RUNTIME, ASYNC_CLIENT) + assert result.returncode == 0, result.stdout + result.stderr + + +@pytest.mark.parametrize( + "scenario", + [CALLER_CLIENT_UNUSED, ASYNC_CLIENT_UNUSED], + ids=["sync-caller-supplied-client", "async"], +) +def test_a_client_that_never_polled_starts_in_a_child(scenario: str) -> None: + result = run(FORCE_FORK_SAFE, scenario) + assert result.returncode == 0, result.stdout + result.stderr + + +@pytest.mark.parametrize( + "scenario", + [ + CALLER_CLIENT_SDK_USED, + ASYNC_CLIENT_SDK_USED, + CALLER_CLIENT_SHARED_SDK_USED, + ASYNC_CLIENT_SHARED_SDK_USED, + ], + ids=[ + "sync-caller-supplied-client", + "async", + "sync-client-shared-with-another-workos-client", + "async-client-shared-with-another-workos-client", + ], +) +def test_a_child_does_not_poll_once_another_call_used_the_http_client( + scenario: str, +) -> None: + result = run(FORCE_FORK_SAFE, scenario) + assert result.returncode == 0, result.stdout + result.stderr + + +def fork_warnings(caplog: pytest.LogCaptureFixture) -> List[str]: + return [r.getMessage() for r in caplog.records if FORK_BLOCKED in r.getMessage()] + + +@pytest.mark.parametrize( + "was_polling", [False, True], ids=["never-polled", "was-polling"] +) +def test_on_macos_start_warns_once_and_does_not_poll( + was_polling: bool, + server: FlagServer, + make_runtime: Make, + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, +) -> None: + monkeypatch.setattr(_core, "FORK_SAFE_PLATFORM", False) + runtime = make_runtime() + if was_polling: + runtime.start() + runtime.wait_until_ready(timeout=5) + sent = len(server.requests) + + # What the at-fork hook runs in a child, without forking pytest. + runtime._reinitialize_after_fork() + with caplog.at_level(logging.WARNING): + if not was_polling: + runtime.is_enabled("f") + assert fork_warnings(caplog) == [] + runtime.start() + runtime.start() + runtime.is_enabled("f") + + [warning] = fork_warnings(caplog) + assert "Create the WorkOSClient and the runtime client after forking" in warning + assert "post_fork" in warning + assert runtime._thread is None + assert len(server.requests) == sent + + +def test_on_macos_a_child_serves_its_bootstrap_and_cannot_refresh( + server: FlagServer, make_runtime: Make, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(_core, "FORK_SAFE_PLATFORM", False) + bootstrapped = make_runtime( + bootstrap=FeatureFlagsConfiguration.from_dict(BOOTSTRAP) + ) + empty = make_runtime() + for runtime in (bootstrapped, empty): + runtime._reinitialize_after_fork() + runtime.start() + with pytest.raises(RuntimeError, match=FORK_BLOCKED): + runtime.refresh() + + bootstrapped.wait_until_ready(timeout=5) + assert bootstrapped.is_enabled("f") is True + started = time.monotonic() + with pytest.raises(RuntimeError, match=FORK_BLOCKED): + empty.wait_until_ready(timeout=5) + assert time.monotonic() - started < 1 + assert empty.is_enabled("f") is False + assert server.requests == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "was_polling", [False, True], ids=["never-polled", "was-polling"] +) +async def test_on_macos_async_start_warns_once_and_does_not_poll( + was_polling: bool, + server: FlagServer, + make_async_runtime: MakeAsync, + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, +) -> None: + monkeypatch.setattr(_core, "FORK_SAFE_PLATFORM", False) + runtime = make_async_runtime() + if was_polling: + await runtime.start() + await runtime.wait_until_ready(timeout=5) + parent_task = runtime._task + sent = len(server.requests) + + # What the at-fork hook runs in a child, without forking pytest. + runtime._reinitialize_after_fork() + with caplog.at_level(logging.WARNING): + if not was_polling: + runtime.is_enabled("f") + assert fork_warnings(caplog) == [] + await runtime.start() + await runtime.start() + runtime.is_enabled("f") + + [warning] = fork_warnings(caplog) + assert ( + "Create the AsyncWorkOSClient and the runtime client after forking" in warning + ) + assert "post_fork" in warning + assert runtime._task is None + assert len(server.requests) == sent + if parent_task is not None: + # The parent's polling task, which a forked child would not run. + parent_task.cancel() + await asyncio.wait({parent_task}) + + +@pytest.mark.asyncio +async def test_on_macos_an_async_child_serves_its_bootstrap_and_cannot_refresh( + server: FlagServer, make_async_runtime: MakeAsync, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(_core, "FORK_SAFE_PLATFORM", False) + bootstrapped = make_async_runtime( + bootstrap=FeatureFlagsConfiguration.from_dict(BOOTSTRAP) + ) + empty = make_async_runtime() + for runtime in (bootstrapped, empty): + runtime._reinitialize_after_fork() + await runtime.start() + with pytest.raises(RuntimeError, match=FORK_BLOCKED): + await runtime.refresh() + + await bootstrapped.wait_until_ready() + assert bootstrapped.is_enabled("f") is True + with pytest.raises(RuntimeError, match=FORK_BLOCKED): + await asyncio.wait_for(empty.wait_until_ready(), 1) + assert empty.is_enabled("f") is False + assert server.requests == [] + + +def test_a_child_raises_from_start_and_wait_once_the_parent_used_the_http_client( + make_runtime: Make, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(_core, "FORK_SAFE_PLATFORM", True) + # Another runtime client polls through the caller-supplied HTTP client. + make_runtime().refresh() + runtime = make_runtime() + + # What the at-fork hook runs in a child, without forking pytest. + runtime._reinitialize_after_fork() + with pytest.raises(RuntimeError, match="caller-supplied HTTP client"): + runtime.start() + started = time.monotonic() + with pytest.raises(RuntimeError, match="caller-supplied HTTP client"): + runtime.wait_until_ready(timeout=5) + assert time.monotonic() - started < 1 + + +@pytest.mark.asyncio +async def test_an_async_child_raises_from_start_and_wait_once_the_parent_used_it( + make_async_runtime: MakeAsync, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(_core, "FORK_SAFE_PLATFORM", True) + await make_async_runtime().refresh() + runtime = make_async_runtime() + + # What the at-fork hook runs in a child, without forking pytest. + runtime._reinitialize_after_fork() + with pytest.raises(RuntimeError, match="was used in the parent"): + await runtime.start() + with pytest.raises(RuntimeError, match="was used in the parent"): + await asyncio.wait_for(runtime.wait_until_ready(), 1) + + +@pytest.mark.parametrize( + "read", ["is_ready", "configuration", "is_closed", "status", "get_all_flags"] +) +def test_every_read_resumes_polling_after_a_fork( + read: str, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(_core, "FORK_SAFE_PLATFORM", True) + server = LocalHTTPServer(payload(f=flag("f", default_value=True))) + workos = WorkOSClient(api_key="sk_test", base_url=server.base_url) + runtime = workos.create_feature_flags_runtime_client() + runtime.start() + runtime.wait_until_ready(timeout=5) + first_poller = runtime._thread + + # What the at-fork hook runs in a child, without forking pytest. + runtime._reinitialize_after_fork() + assert runtime._fork_restart_pending + value = getattr(runtime, read) + if callable(value): + value() + + assert runtime._thread is not None and runtime._thread is not first_poller + assert runtime._thread.is_alive() + runtime.close() + workos.close() + server.close() + + +@dataclass +class UnhashableBackend: + """An HTTPBackend that, like any dataclass with eq=True, can't be hashed.""" + + requests: List[str] = field(default_factory=list) + + def request( + self, + method: str, + url: str, + *, + headers: Dict[str, str], + content: Optional[bytes], + timeout: float, + ) -> HTTPResponse: + self.requests.append(url) + return HTTPResponse( + status_code=200, + headers=httpx2.Headers(), + content=b'{"version": 2, "flags": {}}', + request_method=method, + request_url=url, + ) + + def close(self) -> None: + pass + + +def test_an_unhashable_custom_backend_survives_the_fork_hook( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(_core, "FORK_SAFE_PLATFORM", True) + backend = UnhashableBackend() + workos = WorkOSClient(api_key=API_KEY, http_client=backend) + unused = workos.create_feature_flags_runtime_client() + + # What the at-fork hook runs in a child, without forking pytest. + unused._reinitialize_after_fork() + unused.refresh() + assert backend.requests + + # Another WorkOSClient on the same backend sees that request too. + sharing = WorkOSClient(api_key=API_KEY, http_client=backend) + blocked = sharing.create_feature_flags_runtime_client() + blocked._reinitialize_after_fork() + with pytest.raises(RuntimeError, match="caller-supplied HTTP client"): + blocked.start() + unused.close() + blocked.close() + + +@pytest.mark.parametrize("module", [_sync, _async], ids=["sync", "async"]) +def test_one_client_failing_to_reinitialize_does_not_stop_the_others( + module: Any, monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture +) -> None: + workos = WorkOSClient if module is _sync else AsyncWorkOSClient + runtimes = [ + workos(api_key=API_KEY).create_feature_flags_runtime_client() for _ in range(2) + ] + attempts: List[int] = [] + + def fail() -> None: + attempts.append(1) + raise RuntimeError("could not rebuild") + + for runtime in runtimes: + monkeypatch.setattr(runtime, "_reinitialize_after_fork", fail) + monkeypatch.setattr(module, "_clients", weakref.WeakSet(runtimes)) + + with caplog.at_level(logging.ERROR): + module._reinitialize_clients_after_fork() + + assert len(attempts) == 2 + assert len([r for r in caplog.records if r.levelno == logging.ERROR]) == 2 diff --git a/tests/feature_flags_runtime/test_transport.py b/tests/feature_flags_runtime/test_transport.py index 1b141097..25cfd7fa 100644 --- a/tests/feature_flags_runtime/test_transport.py +++ b/tests/feature_flags_runtime/test_transport.py @@ -32,7 +32,9 @@ def local() -> Iterator[LocalHTTPServer]: server.close() -def test_sends_the_sdk_url_and_headers(local: LocalHTTPServer) -> None: +def test_sends_the_sdk_url_and_headers_and_closes_its_client( + local: LocalHTTPServer, +) -> None: workos = WorkOSClient(api_key=API_KEY, base_url=local.base_url) runtime = workos.create_feature_flags_runtime_client() runtime.refresh() @@ -44,6 +46,20 @@ def test_sends_the_sdk_url_and_headers(local: LocalHTTPServer) -> None: assert runtime.is_enabled("checkout") is True runtime.close() + assert runtime._transport._owned is None + workos.close() + + +def test_a_poll_after_close_opens_no_client(local: LocalHTTPServer) -> None: + workos = WorkOSClient(api_key=API_KEY, base_url=local.base_url) + runtime = workos.create_feature_flags_runtime_client() + runtime.close() + + # As a poll does when close() stopped waiting for it. + with pytest.raises(RuntimeError, match="closed"): + runtime._transport.get(time.monotonic() + 5) + assert runtime._transport._owned is None + assert local.requests == [] workos.close()