fix(feature-flags): Poll in uWSGI workers forked without Python's fork hooks

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. 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`.
+**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`, or that uWSGI loads before forking its workers) 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. That includes a fork that skips Python's fork hooks, as uWSGI's does unless you pass `--py-call-uwsgi-fork-hooks`, since the runtime client checks for a new process each time it is used. 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

@@ -534,9 +534,10 @@ 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
+        # The process that sent the first request through the HTTP client, so
+        # a forked process can tell whether it would share an earlier
+        # process's connections.
+        self._first_request_pid: Optional[int] = None
 
     def close(self) -> None:
         """Release the HTTP client if the SDK created it.
@@ -596,7 +597,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
+        if self._first_request_pid is None:
+            self._first_request_pid = os.getpid()
         record_request(self._backend)
         for attempt in range(max_retries + 1):
             try:
@@ -783,9 +785,10 @@ 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
+        # The process that sent the first request through the HTTP client, so
+        # a forked process can tell whether it would share an earlier
+        # process's connections.
+        self._first_request_pid: Optional[int] = None
 
     async def close(self) -> None:
         """Release the HTTP client if the SDK created it.
@@ -845,7 +848,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
+        if self._first_request_pid is None:
+            self._first_request_pid = os.getpid()
         record_request(self._backend)
         for attempt in range(max_retries + 1):
             try:

src/workos/_http.py

@@ -12,6 +12,7 @@
 import importlib
 import inspect
 import json
+import os
 import weakref
 from dataclasses import dataclass
 from typing import (
@@ -243,13 +244,20 @@ async def close(self) -> None:
 
 # --- connection tracking --------------------------------------------------------
 
-# The HTTP clients requests have gone out through. Several WorkOS clients can be
+# The HTTP clients requests have gone out through, each held weakly with the
+# ID of the process that sent its first request. 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()
-)
+# would share an earlier process's connections. Keyed by id, since a caller's
+# own backend need not be hashable; an entry goes away with its client.
+_first_requests: Dict[int, Tuple[weakref.ref[Any], int]] = {}
+
+
+def _forget(key: int, ref: weakref.ref[Any]) -> None:
+    entry = _first_requests.get(key)
+    # Another thread's first request through the same client may have
+    # replaced this entry; that one goes when its own reference dies.
+    if entry is not None and entry[0] is ref:
+        _first_requests.pop(key, None)
 
 
 def _connection_owner(backend: object) -> object:
@@ -263,18 +271,28 @@ 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)
+    key = id(owner)
+    entry = _first_requests.get(key)
+    if entry is not None and entry[0]() is owner:
+        return
     try:
-        _clients_with_requests[id(owner)] = owner
+        ref = weakref.ref(owner, lambda ref: _forget(key, ref))
     except TypeError:
         # It can't be weakly referenced. Each WorkOS client still tracks its
         # own requests.
-        pass
+        return
+    _first_requests[key] = (ref, os.getpid())
 
 
-def sent_requests(backend: object) -> bool:
-    """Whether any WorkOS client sent a request through ``backend``'s connections."""
+def sent_requests_from_another_process(backend: object) -> bool:
+    """Whether a WorkOS client in another process sent a request through ``backend``'s connections.
+
+    In a forked process, those are requests the parent (or an earlier
+    ancestor) sent before the fork. Requests this process sent don't count.
+    """
     owner = _connection_owner(backend)
-    return _clients_with_requests.get(id(owner)) is owner
+    entry = _first_requests.get(id(owner))
+    return entry is not None and entry[0]() is owner and entry[1] != os.getpid()
 
 
 # --- resolution -----------------------------------------------------------------

src/workos/feature_flags/runtime/_async.py

@@ -14,7 +14,6 @@
 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
@@ -47,9 +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.
+    A client created before a fork, including one that skips Python's fork
+    hooks, can start in the child unless the ``AsyncWorkOSClient`` already
+    sent requests in the parent, or the platform is macOS.
     """
 
     def __init__(
@@ -126,6 +125,7 @@ async def start(self) -> None:
                 the client cannot poll in this forked process because the
                 parent used the ``AsyncWorkOSClient`` or its HTTP client.
         """
+        self._reinitialize_if_forked()
         if self._closed:
             raise RuntimeError(
                 "This feature flags runtime client is closed; create a new one"
@@ -161,6 +161,7 @@ async def wait_until_ready(self, timeout: Optional[float] = None) -> None:
             RuntimeError: If the client was closed before it became ready, or
                 has no configuration and cannot poll in this forked process.
         """
+        self._reinitialize_if_forked()
         if self._status.ready:
             return
         # Nothing can make the client ready in a forked child that cannot poll.
@@ -185,6 +186,7 @@ async def refresh(self) -> None:
             RuntimeError: If the client is closed or cannot poll in this forked
                 process.
         """
+        self._reinitialize_if_forked()
         if self._closed:
             raise RuntimeError(
                 "This feature flags runtime client is closed; create a new one"
@@ -208,6 +210,7 @@ async def close(self) -> None:
         still waiting. Evaluation keeps serving the last configuration after
         closing.
         """
+        self._reinitialize_if_forked()
         first = not self._closed
         self._closed = True
         self._ready.set()
@@ -454,11 +457,7 @@ def _reinitialize_after_fork(self) -> None:
                 "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)
-        ):
+        elif self._used_before_fork() or _core.used_by_another_process(self._client):
             self._fork_blocker = (
                 "the AsyncWorkOSClient or its HTTP client was used in the parent, "
                 "on the parent's event loop and connections. Create the "
@@ -477,6 +476,9 @@ def _reinitialize_after_fork(self) -> None:
         self._fork_warned = not polling
 
     def _before_use(self) -> None:
+        # Evaluation runs this on every call, so the common case is one compare.
+        if self._pid != os.getpid():
+            self._reinitialize_after_unhooked_fork()
         self._warn_if_fork_blocked()
 
     def _warn_if_fork_blocked(self) -> None:
@@ -507,7 +509,7 @@ 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()
+            client._after_fork_in_child()
         except Exception:
             client._logger.exception(
                 "Could not prepare a feature flags runtime client for forked "

src/workos/feature_flags/runtime/_core.py

@@ -13,9 +13,10 @@
 from collections import deque
 from dataclasses import replace
 from datetime import datetime, timezone
-from typing import TYPE_CHECKING, Any, Callable, Deque, Optional, Tuple
+from typing import TYPE_CHECKING, Any, Callable, Deque, Optional, Tuple, Union
 
 from ..._errors import APIError
+from ..._http import sent_requests_from_another_process
 from . import _policy
 from ._configuration import (
     FeatureFlagsConfiguration,
@@ -28,7 +29,7 @@
 from ._types import FlagsChangedEvent, RuntimeStatus
 
 if TYPE_CHECKING:
-    from ..._base_client import _BaseWorkOSClient
+    from ..._base_client import AsyncWorkOSClient, WorkOSClient, _BaseWorkOSClient
 
 LOGGER_NAME = "workos.feature_flags.runtime"
 
@@ -57,6 +58,18 @@ def utc_now() -> datetime:
     return datetime.now(timezone.utc)
 
 
+def used_by_another_process(client: Union[WorkOSClient, AsyncWorkOSClient]) -> bool:
+    """Whether another process sent requests through ``client`` or its HTTP client.
+
+    In a forked process, those are the parent's requests from before the
+    fork, so the answer doesn't change when the child sends its own.
+    """
+    first_request_pid = client._first_request_pid
+    return (
+        first_request_pid is not None and first_request_pid != os.getpid()
+    ) or sent_requests_from_another_process(client._backend)
+
+
 class RuntimeClientCore:
     """Readiness, status, listeners and local evaluation for a runtime client.
 
@@ -109,6 +122,12 @@ def __init__(
         self._pending: Deque[list[Notification]] = deque()
         # The 401 that stopped polling before any configuration was available.
         self._readiness_error: Optional[Exception] = None
+        # The process this client was created in or last prepared for. A
+        # different one means a fork that skipped Python's fork hooks.
+        self._pid = os.getpid()
+        # The threads preparing this client for a process, by process ID; see
+        # _reinitialize_after_unhooked_fork.
+        self._fork_claims: dict[int, threading.Event] = {}
 
     # --- Evaluation -----------------------------------------------------------
 
@@ -197,10 +216,12 @@ def is_closed(self) -> bool:
     def _add_listener(self, kind: str, listener: Listener) -> None:
         if not callable(listener):
             raise TypeError(f"{kind} listener must be callable")
+        self._reinitialize_if_forked()
         with self._listeners_lock:
             self._listeners[kind] = (*self._listeners[kind], listener)
 
     def _remove_listener(self, kind: str, listener: Listener) -> None:
+        self._reinitialize_if_forked()
         with self._listeners_lock:
             self._listeners[kind] = tuple(
                 registered
@@ -330,9 +351,58 @@ 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
 
+    # --- Fork safety ----------------------------------------------------------
+
+    def _after_fork_in_child(self) -> None:
+        """What the at-fork hook runs in a forked child, which has only one thread."""
+        # Recorded first, so a client that fails to rebuild isn't retried on use.
+        self._pid = os.getpid()
+        self._fork_claims = {}
+        self._reinitialize_after_fork()
+
+    def _reinitialize_if_forked(self) -> None:
+        if self._pid != os.getpid():
+            self._reinitialize_after_unhooked_fork()
+
+    def _reinitialize_after_unhooked_fork(self) -> None:
+        """Run the at-fork rebuild in a process forked without Python's fork hooks.
+
+        uWSGI, for one, forks its workers that way unless told otherwise.
+        Public calls check the process ID first and come here once it has
+        changed, so a client is rebuilt on its first use in the new process,
+        with the rules the at-fork hook applies. Any of this process's threads
+        may get here, and possibly at once: the first claims the process and
+        rebuilds the client, and the others wait for it.
+        """
+        pid = os.getpid()
+        claim = threading.Event()
+        claimed = self._fork_claims.setdefault(pid, claim)
+        if claimed is not claim:
+            claimed.wait()
+            return
+        try:
+            # A thread that read the old process ID can claim the process after
+            # another thread finished rebuilding the client.
+            if self._pid != pid:
+                self._reinitialize_after_fork()
+        except Exception:
+            self._logger.exception(
+                "Could not prepare a feature flags runtime client for forked "
+                "process %d",
+                pid,
+            )
+        finally:
+            self._pid = pid
+            self._fork_claims.pop(pid, None)
+            claim.set()
+
+    def _reinitialize_after_fork(self) -> None:
+        """Rebuild what a forked child can't use from the parent."""
+        raise NotImplementedError
+
     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.
+        # The parent's other threads don't exist in the child, and any of them
+        # may have held these locks.
         self._status_lock = threading.Lock()
         self._listeners_lock = threading.Lock()
         self._context_warnings.reinitialize_after_fork()