chore(feature-flags): Keep the runtime clients safe across fork()

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
 

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(

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 -----------------------------------------------------------------
 
 

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,21 +87,27 @@ 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)
         if on_error is not None:
             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.
         self._refresh_lock = asyncio.Lock()
         self._ready = asyncio.Event()
         # Set when polling stops after a 401, to end the polling task's sleep.
         self._wake = asyncio.Event()
-        if self._status.ready:
+        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)

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.