From c64c9dc5c2457724007ddd845a7e7c31c4ade5d8 Mon Sep 17 00:00:00 2001 From: Stanley Phu Date: Tue, 3 Nov 2026 23:27:29 -0800 Subject: [PATCH 1/4] chore(feature-flags): Keep the runtime clients safe across fork() A runtime client created before a fork (gunicorn --preload) can be used in each worker. With the SDK's default HTTP client, the sync client polls through an HTTP client of its own, so a child replaces it and resumes polling on new connections the first time it is used. Where the child would share the parent's connections or event loop, or on macOS, start() and refresh() raise and the client keeps serving its last configuration. Co-Authored-By: Claude Opus 5.5 --- README.md | 2 +- src/workos/_base_client.py | 8 + src/workos/feature_flags/runtime/_async.py | 91 +++- src/workos/feature_flags/runtime/_context.py | 3 + src/workos/feature_flags/runtime/_core.py | 41 ++ src/workos/feature_flags/runtime/_sync.py | 108 ++++- .../feature_flags/runtime/_transport.py | 88 +++- tests/feature_flags_runtime/helpers.py | 11 +- tests/feature_flags_runtime/test_fork.py | 387 ++++++++++++++++++ tests/feature_flags_runtime/test_transport.py | 18 +- 10 files changed, 729 insertions(+), 28 deletions(-) create mode 100644 tests/feature_flags_runtime/test_fork.py diff --git a/README.md b/README.md index 92b24f49..ba072189 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. 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 or any other SDK call, 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; create the `WorkOSClient` and the runtime client after forking instead. 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..1f77c2e4 100644 --- a/src/workos/_base_client.py +++ b/src/workos/_base_client.py @@ -533,6 +533,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 +595,7 @@ 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 for attempt in range(max_retries + 1): try: response = self._backend.request( @@ -777,6 +781,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 +843,7 @@ 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 for attempt in range(max_retries + 1): try: response = await self._backend.request( diff --git a/src/workos/feature_flags/runtime/_async.py b/src/workos/feature_flags/runtime/_async.py index cbc80c33..35bf8eab 100644 --- a/src/workos/feature_flags/runtime/_async.py +++ b/src/workos/feature_flags/runtime/_async.py @@ -7,15 +7,17 @@ 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 . 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 +46,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 +86,9 @@ 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. + self._fork_blocker: Optional[str] = None + self._fork_warned = False self._init_synchronization() if on_change is not None: self.add_change_listener(on_change) @@ -88,6 +96,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 +104,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 ------------------------------------------------------------ @@ -108,12 +117,14 @@ async def start(self) -> None: polling permanently; the last good configuration stays usable. 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. """ if self._closed: raise RuntimeError( "This feature flags runtime client is closed; create a new one" ) + self._raise_if_fork_blocked() if self._started: return loop = asyncio.get_running_loop() @@ -159,12 +170,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 +299,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 +413,65 @@ 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 has sent requests, its connections belong + # to the parent's event loop, so a copy of a client that polled, or of + # one whose AsyncWorkOSClient made other calls, 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" + ) + elif self._used_before_fork() or self._client._sent_requests: + self._fork_blocker = ( + "the AsyncWorkOSClient 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 + # Only a client that was polling warns, on first use, that it no longer + # polls; any other raises from start() and refresh() instead. + self._fork_warned = not polling + + def _before_use(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)) + + +_clients: weakref.WeakSet[AsyncFeatureFlagsRuntimeClient] = weakref.WeakSet() + + +def _reinitialize_clients_after_fork() -> None: + for client in list(_clients): + client._reinitialize_after_fork() + + +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..78df641e 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 ._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,10 @@ 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. + self._fork_blocker: Optional[str] = None + self._fork_warned = False self._init_synchronization() if on_change is not None: self.add_change_listener(on_change) @@ -102,6 +118,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 +129,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 ------------------------------------------------------------ @@ -125,13 +145,16 @@ def start(self) -> None: polling permanently; the last good configuration stays usable. Raises: - RuntimeError: If the client is closed. + RuntimeError: If the client is closed or cannot poll in this forked + process. """ + self._before_use() with self._lifecycle_lock: if self._closed: raise RuntimeError( "This feature flags runtime client is closed; create a new one" ) + self._raise_if_fork_blocked() if self._started: return self._started = True @@ -158,6 +181,7 @@ def wait_until_ready(self, timeout: Optional[float] = None) -> None: configuration was available. RuntimeError: If the client was closed before it became ready. """ + self._before_use() if self._status.ready: return if not self._ready.wait(timeout): @@ -174,12 +198,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 +221,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 +242,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 +441,71 @@ 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() + elif 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 _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()) + ) + 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; any other raises from start() and refresh() 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): + client._reinitialize_after_fork() + + +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..5cdb00b0 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, @@ -18,7 +21,7 @@ TransportTimeout, _HttpxFamily, ) -from . import _policy +from . import _core, _policy if TYPE_CHECKING: from _collections_abc import dict_items, dict_values @@ -77,13 +80,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 +101,46 @@ 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" + ) + # The WorkOSClient's other requests count too: they open the same + # HTTP client's connections. + if ( + used_before_fork or self._workos._sent_requests + ) 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 +162,9 @@ 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 + client = self._httpx_client() if client is None: return self._workos._backend.request( "GET", @@ -133,6 +182,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..f92e364a --- /dev/null +++ b/tests/feature_flags_runtime/test_fork.py @@ -0,0 +1,387 @@ +# @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 logging +import os +import subprocess +import sys +import textwrap +from pathlib import Path + +import pytest + +from workos import WorkOSClient +from workos.feature_flags.runtime import _core + +from tests.feature_flags_runtime.helpers import LocalHTTPServer, flag, payload + +ROOT = Path(__file__).resolve().parents[2] + +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 + 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 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" +""" + +# 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 + + +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], + ids=["sync-caller-supplied-client", "async"], +) +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 test_on_macos_a_client_that_never_polled_raises_from_start_without_warning( + monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture +) -> None: + monkeypatch.setattr(_core, "FORK_SAFE_PLATFORM", False) + workos = WorkOSClient(api_key="sk_test") + runtime = workos.create_feature_flags_runtime_client() + + runtime._reinitialize_after_fork() + with caplog.at_level(logging.WARNING): + assert runtime.is_enabled("f", default=True) is True + assert caplog.records == [] + with pytest.raises(RuntimeError, match="Create the WorkOSClient and the runtime"): + runtime.start() + runtime.close() + workos.close() + + +@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() 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() From f98387e0f61556b40aec1e182cbc0941c30b57e8 Mon Sep 17 00:00:00 2001 From: Stanley Phu Date: Wed, 4 Nov 2026 00:02:07 -0800 Subject: [PATCH 2/4] fix(feature-flags): Track parent requests per HTTP client, across WorkOS clients Two WorkOS clients can share one HTTP client you supplied. A request sent through one marked only that WorkOS client, so an unused runtime client on the other still polled in a forked child over the parent's connections, or on the parent's event loop for async. Requests are now recorded against the HTTP client itself, and the fork check reads that. Co-Authored-By: Claude Opus 5.5 --- README.md | 2 +- src/workos/_base_client.py | 3 + src/workos/_http.py | 32 ++++++++ src/workos/feature_flags/runtime/_async.py | 24 +++--- .../feature_flags/runtime/_transport.py | 12 ++- tests/feature_flags_runtime/test_fork.py | 74 ++++++++++++++++++- 6 files changed, 132 insertions(+), 15 deletions(-) diff --git a/README.md b/README.md index ba072189..886e6d2c 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.** 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. 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 or any other SDK call, 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; create the `WorkOSClient` and the runtime client after forking instead. 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. 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; create the `WorkOSClient` and the runtime client after forking instead. 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 1f77c2e4..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, ) @@ -596,6 +597,7 @@ def request( 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( @@ -844,6 +846,7 @@ async def request( 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..27d8fc20 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,37 @@ 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. +_clients_with_requests: weakref.WeakSet[object] = weakref.WeakSet() + + +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.""" + try: + _clients_with_requests.add(_connection_owner(backend)) + 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.""" + return _connection_owner(backend) in _clients_with_requests + + # --- resolution ----------------------------------------------------------------- diff --git a/src/workos/feature_flags/runtime/_async.py b/src/workos/feature_flags/runtime/_async.py index 35bf8eab..93960b83 100644 --- a/src/workos/feature_flags/runtime/_async.py +++ b/src/workos/feature_flags/runtime/_async.py @@ -14,6 +14,7 @@ from typing import TYPE_CHECKING, Any, Awaitable, Callable, Optional, Type, Union from ..._errors import WorkOSTimeoutError +from ..._http import sent_requests from . import _core, _policy from ._configuration import FeatureFlagsConfiguration from ._context import LoggerLike @@ -417,11 +418,12 @@ async def _await_listener(self, awaitable: Awaitable[Any]) -> None: # --- Fork safety ---------------------------------------------------------- def _reinitialize_after_fork(self) -> None: - # Once the AsyncWorkOSClient has sent requests, its connections belong - # to the parent's event loop, so a copy of a client that polled, or of - # one whose AsyncWorkOSClient made other calls, 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. + # 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 @@ -440,11 +442,15 @@ def _reinitialize_after_fork(self) -> None: f"{_core.MACOS_FORK_HAZARD}. Create the AsyncWorkOSClient and the " "runtime client after forking instead" ) - elif self._used_before_fork() or self._client._sent_requests: + elif ( + self._used_before_fork() + or self._client._sent_requests + or sent_requests(self._client._backend) + ): self._fork_blocker = ( - "the AsyncWorkOSClient was used in the parent, on the parent's " - "event loop and connections. Create the AsyncWorkOSClient and the " - "runtime client after forking instead" + "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 diff --git a/src/workos/feature_flags/runtime/_transport.py b/src/workos/feature_flags/runtime/_transport.py index 5cdb00b0..634b20b5 100644 --- a/src/workos/feature_flags/runtime/_transport.py +++ b/src/workos/feature_flags/runtime/_transport.py @@ -20,6 +20,8 @@ TransportError, TransportTimeout, _HttpxFamily, + record_request, + sent_requests, ) from . import _core, _policy @@ -116,10 +118,13 @@ def fork_blocker(self, used_before_fork: bool) -> Optional[str]: f"{_core.MACOS_FORK_HAZARD}. Create the WorkOSClient and the " "runtime client after forking instead" ) - # The WorkOSClient's other requests count too: they open the same - # HTTP client's connections. + # 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 + 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 " @@ -164,6 +169,7 @@ def _send(self, url: str, headers: Dict[str, str], deadline: float) -> HTTPRespo raise TransportTimeout("no time was left in the request timeout") 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( diff --git a/tests/feature_flags_runtime/test_fork.py b/tests/feature_flags_runtime/test_fork.py index f92e364a..9eb2aaaf 100644 --- a/tests/feature_flags_runtime/test_fork.py +++ b/tests/feature_flags_runtime/test_fork.py @@ -271,6 +271,66 @@ async def call_api(): 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 @@ -332,8 +392,18 @@ def test_a_client_that_never_polled_starts_in_a_child(scenario: str) -> None: @pytest.mark.parametrize( "scenario", - [CALLER_CLIENT_SDK_USED, ASYNC_CLIENT_SDK_USED], - ids=["sync-caller-supplied-client", "async"], + [ + 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, From c1be99a9f587eab36c80c8549ca27d38631b0267 Mon Sep 17 00:00:00 2001 From: Stanley Phu Date: Wed, 4 Nov 2026 00:44:21 -0800 Subject: [PATCH 3/4] fix(feature-flags): Prepare forked runtime clients with unhashable backends The record of HTTP clients that sent requests was a WeakSet, which hashes its members, so a caller's own backend that can't be hashed (any dataclass with eq=True) raised TypeError in the fork hook and stopped it from rebuilding the remaining runtime clients. The record is now keyed by id, and a client that fails to rebuild is logged without stopping the others. Co-Authored-By: Claude Opus 5.5 --- src/workos/_http.py | 13 +++- src/workos/feature_flags/runtime/_async.py | 10 ++- src/workos/feature_flags/runtime/_sync.py | 10 ++- tests/feature_flags_runtime/test_fork.py | 91 +++++++++++++++++++++- 4 files changed, 115 insertions(+), 9 deletions(-) diff --git a/src/workos/_http.py b/src/workos/_http.py index 27d8fc20..96cb71f1 100644 --- a/src/workos/_http.py +++ b/src/workos/_http.py @@ -245,8 +245,11 @@ async def close(self) -> None: # 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. -_clients_with_requests: weakref.WeakSet[object] = weakref.WeakSet() +# 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: @@ -259,8 +262,9 @@ def _connection_owner(backend: object) -> object: 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.add(_connection_owner(backend)) + _clients_with_requests[id(owner)] = owner except TypeError: # It can't be weakly referenced. Each WorkOS client still tracks its # own requests. @@ -269,7 +273,8 @@ def record_request(backend: object) -> None: def sent_requests(backend: object) -> bool: """Whether any WorkOS client sent a request through ``backend``'s connections.""" - return _connection_owner(backend) in _clients_with_requests + 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 93960b83..ccd6c180 100644 --- a/src/workos/feature_flags/runtime/_async.py +++ b/src/workos/feature_flags/runtime/_async.py @@ -476,7 +476,15 @@ def _raise_if_fork_blocked(self) -> None: def _reinitialize_clients_after_fork() -> None: for client in list(_clients): - client._reinitialize_after_fork() + # 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"): diff --git a/src/workos/feature_flags/runtime/_sync.py b/src/workos/feature_flags/runtime/_sync.py index 78df641e..1c53593c 100644 --- a/src/workos/feature_flags/runtime/_sync.py +++ b/src/workos/feature_flags/runtime/_sync.py @@ -504,7 +504,15 @@ def _discard_awaitable(awaitable: object) -> None: def _reinitialize_clients_after_fork() -> None: for client in list(_clients): - client._reinitialize_after_fork() + # 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"): diff --git a/tests/feature_flags_runtime/test_fork.py b/tests/feature_flags_runtime/test_fork.py index 9eb2aaaf..a92dff4b 100644 --- a/tests/feature_flags_runtime/test_fork.py +++ b/tests/feature_flags_runtime/test_fork.py @@ -11,14 +11,23 @@ import subprocess import sys import textwrap +import weakref +from dataclasses import dataclass, field from pathlib import Path +from typing import Any, Dict, List, Optional +import httpx2 import pytest -from workos import WorkOSClient -from workos.feature_flags.runtime import _core +from workos import AsyncWorkOSClient, HTTPResponse, WorkOSClient +from workos.feature_flags.runtime import _async, _core, _sync -from tests.feature_flags_runtime.helpers import LocalHTTPServer, flag, payload +from tests.feature_flags_runtime.helpers import ( + API_KEY, + LocalHTTPServer, + flag, + payload, +) ROOT = Path(__file__).resolve().parents[2] @@ -455,3 +464,79 @@ def test_every_read_resumes_polling_after_a_fork( 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 From d60a6207a9529f4aa7441f34258e8029fcfe5173 Mon Sep 17 00:00:00 2001 From: Stanley Phu Date: Wed, 4 Nov 2026 11:15:18 -0800 Subject: [PATCH 4/4] fix(feature-flags): Warn and skip polling on start() in a macOS child On macOS, start() raised in every forked child of a process that had created a client, so gunicorn --preload on a Mac failed at worker boot with no opt-out. Every client created before the fork is blocked there, however the parent used it, so start() now logs the fork warning once and returns without polling. The client serves its bootstrap or last configuration. refresh() still raises, since it would open a connection, and wait_until_ready() raises at once for a blocked child without a configuration instead of waiting for a poll that never comes. The Linux blockers, where the parent used a caller-supplied HTTP client or the AsyncWorkOSClient, still raise from start(). The warning and README now point to gunicorn's post_fork hook. Co-Authored-By: Claude Opus 5.5 --- README.md | 2 +- src/workos/feature_flags/runtime/_async.py | 41 ++- src/workos/feature_flags/runtime/_sync.py | 45 +++- .../feature_flags/runtime/_transport.py | 3 +- tests/feature_flags_runtime/test_fork.py | 241 +++++++++++++++++- 5 files changed, 305 insertions(+), 27 deletions(-) diff --git a/README.md b/README.md index 886e6d2c..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.** 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. 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; create the `WorkOSClient` and the runtime client after forking instead. 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/feature_flags/runtime/_async.py b/src/workos/feature_flags/runtime/_async.py index ccd6c180..90b023ae 100644 --- a/src/workos/feature_flags/runtime/_async.py +++ b/src/workos/feature_flags/runtime/_async.py @@ -87,8 +87,10 @@ 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. + # 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: @@ -115,16 +117,22 @@ 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, no event loop is running, or - the client cannot poll in this forked process. + 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 @@ -150,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: @@ -440,7 +451,8 @@ def _reinitialize_after_fork(self) -> 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" + "runtime client after forking instead, for example in gunicorn's " + "post_fork hook" ) elif ( self._used_before_fork() @@ -454,11 +466,20 @@ def _reinitialize_after_fork(self) -> None: ) 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; any other raises from start() and refresh() instead. + # 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( @@ -470,6 +491,14 @@ 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() diff --git a/src/workos/feature_flags/runtime/_sync.py b/src/workos/feature_flags/runtime/_sync.py index 1c53593c..cce0c365 100644 --- a/src/workos/feature_flags/runtime/_sync.py +++ b/src/workos/feature_flags/runtime/_sync.py @@ -13,7 +13,7 @@ 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 ( @@ -108,8 +108,10 @@ def __init__( 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. + # 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: @@ -142,11 +144,14 @@ 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 cannot poll in this forked - process. + 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: @@ -154,6 +159,9 @@ def start(self) -> None: 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 @@ -179,11 +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: @@ -446,7 +457,11 @@ def _deliver(self, notifications: list[Notification]) -> None: def _before_use(self) -> None: if self._fork_restart_pending: self._restart_after_fork() - elif self._fork_blocker is not None and not self._fork_warned: + 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.", @@ -457,6 +472,14 @@ 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: @@ -487,9 +510,15 @@ def _reinitialize_after_fork(self) -> 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; any other raises from start() and refresh() instead. + # polls. On macOS any other warns once started; elsewhere start() + # raises instead. self._fork_warned = not polling diff --git a/src/workos/feature_flags/runtime/_transport.py b/src/workos/feature_flags/runtime/_transport.py index 634b20b5..067c4502 100644 --- a/src/workos/feature_flags/runtime/_transport.py +++ b/src/workos/feature_flags/runtime/_transport.py @@ -116,7 +116,8 @@ def fork_blocker(self, used_before_fork: bool) -> Optional[str]: if not _core.FORK_SAFE_PLATFORM: return ( f"{_core.MACOS_FORK_HAZARD}. Create the WorkOSClient and the " - "runtime client after forking instead" + "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 diff --git a/tests/feature_flags_runtime/test_fork.py b/tests/feature_flags_runtime/test_fork.py index a92dff4b..3ff000c1 100644 --- a/tests/feature_flags_runtime/test_fork.py +++ b/tests/feature_flags_runtime/test_fork.py @@ -6,30 +6,45 @@ 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, Dict, List, Optional +from typing import Any, Callable, Dict, List, Optional import httpx2 import pytest from workos import AsyncWorkOSClient, HTTPResponse, WorkOSClient -from workos.feature_flags.runtime import _async, _core, _sync +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 @@ -179,6 +194,12 @@ async def _make_runtime(workos): 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) @@ -187,6 +208,51 @@ async def _make_runtime(workos): 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 = """ @@ -376,6 +442,15 @@ 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: @@ -421,21 +496,165 @@ def test_a_child_does_not_poll_once_another_call_used_the_http_client( assert result.returncode == 0, result.stdout + result.stderr -def test_on_macos_a_client_that_never_polled_raises_from_start_without_warning( - monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture +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) - workos = WorkOSClient(api_key="sk_test") - runtime = workos.create_feature_flags_runtime_client() + 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): - assert runtime.is_enabled("f", default=True) is True - assert caplog.records == [] - with pytest.raises(RuntimeError, match="Create the WorkOSClient and the runtime"): + if not was_polling: + runtime.is_enabled("f") + assert fork_warnings(caplog) == [] runtime.start() - runtime.close() - workos.close() + 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(