diff --git a/README.md b/README.md index 465ed54b..8c76f54c 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. 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 diff --git a/src/workos/_base_client.py b/src/workos/_base_client.py index 4f743d42..d29b6c46 100644 --- a/src/workos/_base_client.py +++ b/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: diff --git a/src/workos/_http.py b/src/workos/_http.py index 96cb71f1..1d2061da 100644 --- a/src/workos/_http.py +++ b/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 ----------------------------------------------------------------- diff --git a/src/workos/feature_flags/runtime/_async.py b/src/workos/feature_flags/runtime/_async.py index 90b023ae..c82e1bac 100644 --- a/src/workos/feature_flags/runtime/_async.py +++ b/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 " diff --git a/src/workos/feature_flags/runtime/_core.py b/src/workos/feature_flags/runtime/_core.py index 19497118..623ccf7c 100644 --- a/src/workos/feature_flags/runtime/_core.py +++ b/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() diff --git a/src/workos/feature_flags/runtime/_sync.py b/src/workos/feature_flags/runtime/_sync.py index cce0c365..912c35bd 100644 --- a/src/workos/feature_flags/runtime/_sync.py +++ b/src/workos/feature_flags/runtime/_sync.py @@ -52,7 +52,9 @@ class FeatureFlagsRuntimeClient(RuntimeClientCore): 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. + requests through, and on macOS. That includes a fork that skips Python's + fork hooks, as uWSGI's does by default: the client checks for a new + process each time it is used. """ def __init__( @@ -239,6 +241,7 @@ def close(self) -> None: including one that calls while another close() is still waiting. Evaluation keeps serving the last configuration after closing. """ + self._reinitialize_if_forked() with self._lifecycle_lock: if not self._closed: self._closed = True @@ -455,6 +458,9 @@ def _deliver(self, notifications: list[Notification]) -> None: # --- Fork safety ---------------------------------------------------------- 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() if self._fork_restart_pending: self._restart_after_fork() else: @@ -493,10 +499,10 @@ def _restart_after_fork(self) -> None: 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. + # The parent's other threads don't exist in the child. Locks may have + # been held by them, 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 @@ -535,7 +541,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 " diff --git a/src/workos/feature_flags/runtime/_transport.py b/src/workos/feature_flags/runtime/_transport.py index 067c4502..78999822 100644 --- a/src/workos/feature_flags/runtime/_transport.py +++ b/src/workos/feature_flags/runtime/_transport.py @@ -5,6 +5,7 @@ from __future__ import annotations +import os import threading import time from typing import TYPE_CHECKING, Any, Dict, Iterator, Optional @@ -21,7 +22,6 @@ TransportTimeout, _HttpxFamily, record_request, - sent_requests, ) from . import _core, _policy @@ -123,9 +123,7 @@ def fork_blocker(self, used_before_fork: bool) -> Optional[str]: # other WorkOSClient given the same HTTP client: they open the same # connections. if ( - used_before_fork - or self._workos._sent_requests - or sent_requests(self._workos._backend) + used_before_fork or _core.used_by_another_process(self._workos) ) and not self._uses_default_client: return ( "the parent used its caller-supplied HTTP client, whose " @@ -169,7 +167,8 @@ def _send(self, url: str, headers: Dict[str, str], deadline: float) -> HTTPRespo if deadline - time.monotonic() <= 0: raise TransportTimeout("no time was left in the request timeout") if not self._uses_default_client: - self._workos._sent_requests = True + if self._workos._first_request_pid is None: + self._workos._first_request_pid = os.getpid() record_request(self._workos._backend) client = self._httpx_client() if client is None: diff --git a/tests/feature_flags_runtime/helpers.py b/tests/feature_flags_runtime/helpers.py index c106ff3f..78de9317 100644 --- a/tests/feature_flags_runtime/helpers.py +++ b/tests/feature_flags_runtime/helpers.py @@ -4,6 +4,7 @@ from __future__ import annotations import json +import socket import threading import time from dataclasses import dataclass @@ -261,7 +262,12 @@ def do_GET(self) -> None: self._server = http.server.ThreadingHTTPServer(("127.0.0.1", 0), Handler) self._server.daemon_threads = True - self._thread = threading.Thread(target=self._server.serve_forever, daemon=True) + # Without a poll interval the server thread sleeps until a connection + # arrives, so it can't be waiting for the GIL when a test forks without + # Python's fork hooks. + self._thread = threading.Thread( + target=self._server.serve_forever, args=(None,), daemon=True + ) self._thread.start() @property @@ -277,7 +283,16 @@ def ports(self, authorization: str) -> set: } def close(self) -> None: - self._server.shutdown() + # The server thread only sees the shutdown once a connection wakes it. + stopper = threading.Thread(target=self._server.shutdown) + stopper.start() + while stopper.is_alive(): + try: + address = ("127.0.0.1", self._server.server_address[1]) + socket.create_connection(address, timeout=1).close() + except OSError: + pass + stopper.join(0.05) self._server.server_close() def _handle(self, handler: Any) -> None: diff --git a/tests/feature_flags_runtime/test_fork.py b/tests/feature_flags_runtime/test_fork.py index 3ff000c1..3b1fa8ec 100644 --- a/tests/feature_flags_runtime/test_fork.py +++ b/tests/feature_flags_runtime/test_fork.py @@ -2,16 +2,21 @@ """What runtime clients do in a forked child process. Each scenario runs in a subprocess so a hang or a fork cannot wedge pytest. +Most run with both kinds of fork: ``os.fork()``, which runs Python's fork +hooks, and libc's ``fork()``, which skips them, as uWSGI forks its workers by +default. """ from __future__ import annotations import asyncio +import itertools import logging import os import subprocess import sys import textwrap +import threading import time import weakref from dataclasses import dataclass, field @@ -416,6 +421,138 @@ async def call_api(): _core.FORK_SAFE_PLATFORM = True """ +# Replaces fork() with libc's fork(), which skips the hooks Python runs around +# os.fork(), as uWSGI does when it forks workers without +# --py-call-uwsgi-fork-hooks. PyDLL keeps the GIL across the call. Without the +# hooks, CPython doesn't rebuild the GIL's own locks in the child, so a parent +# thread waiting for the GIL at the fork can leave the child unable to switch +# threads; the parent's pollers are idled first. A child that fails exits +# without finalizing the interpreter, which would wait for threads that exist +# only in the parent. +WITHOUT_FORK_HOOKS = """ +import ctypes +import sys +import traceback + +_libc_fork = ctypes.PyDLL(None).fork + + +def _exit_on_error(*exc_info): + traceback.print_exception(*exc_info) + sys.stderr.flush() + os._exit(1) + + +def fork(): + next_delay = _policy.next_delay + _policy.next_delay = lambda interval, failures, max_retry_delay: 60.0 + time.sleep(0.3) + pid = _libc_fork() + _policy.next_delay = next_delay + if pid == 0: + sys.excepthook = _exit_on_error + return pid +""" + +FORKS = pytest.mark.parametrize( + "fork", ["", WITHOUT_FORK_HOOKS], ids=["os-fork", "fork-without-hooks"] +) + +# The parent never uses the HTTP client, and the child sends another SDK call +# through it before starting the runtime client. +CALLER_CLIENT_USED_IN_CHILD = """ +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" + try: + workos.organizations.list_organizations() + except Exception: + pass + 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 +paths = [r["path"] for r in server.requests if r["headers"].get("Authorization") == "Bearer sk_child"] +assert paths[0].startswith("/organizations"), paths +assert any(path.startswith("/sdk/feature-flags") for path in paths), paths +""" + +ASYNC_CLIENT_USED_IN_CHILD = """ +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(): + try: + await workos.organizations.list_organizations() + except Exception: + pass + 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 +paths = [r["path"] for r in server.requests if r["headers"].get("Authorization") == "Bearer sk_child"] +assert paths[0].startswith("/organizations"), paths +assert any(path.startswith("/sdk/feature-flags") for path in paths), paths +""" + +# Counts how often each client is prepared for the child, however many times +# the child uses it. +PREPARED_ONCE = """ +from workos.feature_flags.runtime import AsyncFeatureFlagsRuntimeClient, FeatureFlagsRuntimeClient + +prepared = [] + + +def count_preparations(cls): + reinitialize = cls._reinitialize_after_fork + + def counted(self): + prepared.append(cls.__name__) + reinitialize(self) + + cls._reinitialize_after_fork = counted + + +count_preparations(FeatureFlagsRuntimeClient) +count_preparations(AsyncFeatureFlagsRuntimeClient) +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) +async_workos = AsyncWorkOSClient(api_key="sk_parent", base_url=server.base_url) +async_runtime = async_workos.create_feature_flags_runtime_client() +pid = fork() +if pid == 0: + workos._api_key = "sk_child" + for _ in range(3): + assert runtime.is_enabled("f") is True + runtime.start() + runtime.refresh() + async_runtime.is_enabled("f") + asyncio.run(async_runtime.close()) + runtime.close() + assert sorted(prepared) == ["AsyncFeatureFlagsRuntimeClient", "FeatureFlagsRuntimeClient"], prepared + os._exit(0) +_, status = os.waitpid(pid, 0) +assert os.waitstatus_to_exitcode(status) == 0 +runtime.close() +assert server.ports("Bearer sk_child"), "the child never polled" +""" + def run(*scenario: str) -> subprocess.CompletedProcess: script = "".join(textwrap.dedent(part) for part in (PRELUDE, *scenario)) @@ -432,48 +569,70 @@ def run(*scenario: str) -> subprocess.CompletedProcess: 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) +@FORKS +def test_a_child_polls_on_its_own_connections_once_used(fork: str) -> None: + result = run(FORCE_FORK_SAFE, fork, 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) +@FORKS +def test_on_macos_a_child_does_not_poll(fork: str) -> None: + result = run(fork, 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) +@FORKS +def test_on_macos_a_preloaded_client_serves_its_bootstrap_in_a_child( + fork: str, +) -> None: + result = run(fork, MACOS_PRELOADED) assert result.returncode == 0, result.stdout + result.stderr # One warning from each start(). assert result.stderr.count(FORK_BLOCKED) == 3, result.stderr -def test_a_child_does_not_poll_through_a_caller_supplied_client() -> None: - result = run(FORCE_FORK_SAFE, CALLER_CLIENT) +@FORKS +def test_a_child_does_not_poll_through_a_caller_supplied_client(fork: str) -> None: + result = run(FORCE_FORK_SAFE, fork, 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) +@FORKS +def test_an_async_client_stops_polling_in_a_child(fork: str) -> None: + result = run(FORCE_FORK_SAFE, fork, MAKE_ASYNC_RUNTIME, ASYNC_CLIENT) assert result.returncode == 0, result.stdout + result.stderr +@FORKS @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) +def test_a_client_that_never_polled_starts_in_a_child(fork: str, scenario: str) -> None: + result = run(FORCE_FORK_SAFE, fork, scenario) + assert result.returncode == 0, result.stdout + result.stderr + + +@FORKS +@pytest.mark.parametrize( + "scenario", + [CALLER_CLIENT_USED_IN_CHILD, ASYNC_CLIENT_USED_IN_CHILD], + ids=["sync-caller-supplied-client", "async"], +) +def test_requests_the_child_sent_itself_do_not_stop_it_polling( + fork: str, scenario: str +) -> None: + result = run(FORCE_FORK_SAFE, fork, scenario) assert result.returncode == 0, result.stdout + result.stderr +@FORKS @pytest.mark.parametrize( "scenario", [ @@ -490,9 +649,15 @@ def test_a_client_that_never_polled_starts_in_a_child(scenario: str) -> None: ], ) def test_a_child_does_not_poll_once_another_call_used_the_http_client( - scenario: str, + fork: str, scenario: str ) -> None: - result = run(FORCE_FORK_SAFE, scenario) + result = run(FORCE_FORK_SAFE, fork, scenario) + assert result.returncode == 0, result.stdout + result.stderr + + +@FORKS +def test_a_child_prepares_each_client_once(fork: str) -> None: + result = run(FORCE_FORK_SAFE, fork, PREPARED_ONCE) assert result.returncode == 0, result.stdout + result.stderr @@ -500,11 +665,49 @@ def fork_warnings(caplog: pytest.LogCaptureFixture) -> List[str]: return [r.getMessage() for r in caplog.records if FORK_BLOCKED in r.getMessage()] +ForkInPlace = Callable[..., None] + +HOOKS = pytest.mark.parametrize( + "hooks", [True, False], ids=["at-fork-hook", "without-fork-hooks"] +) + + +@pytest.fixture +def fork_in_place(monkeypatch: pytest.MonkeyPatch) -> ForkInPlace: + """Continue the test as a forked child, without forking pytest. + + Each call gives the test a new process ID, as a fork does. With + ``hooks=True`` the at-fork hook then runs for ``runtimes``, as after + ``os.fork()``; with ``hooks=False`` nothing runs, as when uWSGI forks a + worker. Request this fixture before the ones that create runtime clients, + so those are closed while the test is still the child. + """ + pids = itertools.count(os.getpid() + 1) + + def fork(*runtimes: Any, hooks: bool) -> None: + pid = next(pids) + monkeypatch.setattr(os, "getpid", lambda: pid) + if not hooks: + return + for module, kind in ( + (_sync, FeatureFlagsRuntimeClient), + (_async, AsyncFeatureFlagsRuntimeClient), + ): + mine = [runtime for runtime in runtimes if isinstance(runtime, kind)] + monkeypatch.setattr(module, "_clients", weakref.WeakSet(mine)) + module._reinitialize_clients_after_fork() + + return fork + + +@HOOKS @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, + hooks: bool, + fork_in_place: ForkInPlace, server: FlagServer, make_runtime: Make, monkeypatch: pytest.MonkeyPatch, @@ -517,8 +720,7 @@ def test_on_macos_start_warns_once_and_does_not_poll( 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() + fork_in_place(runtime, hooks=hooks) with caplog.at_level(logging.WARNING): if not was_polling: runtime.is_enabled("f") @@ -534,16 +736,21 @@ def test_on_macos_start_warns_once_and_does_not_poll( assert len(server.requests) == sent +@HOOKS def test_on_macos_a_child_serves_its_bootstrap_and_cannot_refresh( - server: FlagServer, make_runtime: Make, monkeypatch: pytest.MonkeyPatch + hooks: bool, + fork_in_place: ForkInPlace, + 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() + fork_in_place(bootstrapped, empty, hooks=hooks) for runtime in (bootstrapped, empty): - runtime._reinitialize_after_fork() runtime.start() with pytest.raises(RuntimeError, match=FORK_BLOCKED): runtime.refresh() @@ -559,11 +766,14 @@ def test_on_macos_a_child_serves_its_bootstrap_and_cannot_refresh( @pytest.mark.asyncio +@HOOKS @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, + hooks: bool, + fork_in_place: ForkInPlace, server: FlagServer, make_async_runtime: MakeAsync, monkeypatch: pytest.MonkeyPatch, @@ -577,8 +787,7 @@ async def test_on_macos_async_start_warns_once_and_does_not_poll( parent_task = runtime._task sent = len(server.requests) - # What the at-fork hook runs in a child, without forking pytest. - runtime._reinitialize_after_fork() + fork_in_place(runtime, hooks=hooks) with caplog.at_level(logging.WARNING): if not was_polling: runtime.is_enabled("f") @@ -601,16 +810,21 @@ async def test_on_macos_async_start_warns_once_and_does_not_poll( @pytest.mark.asyncio +@HOOKS async def test_on_macos_an_async_child_serves_its_bootstrap_and_cannot_refresh( - server: FlagServer, make_async_runtime: MakeAsync, monkeypatch: pytest.MonkeyPatch + hooks: bool, + fork_in_place: ForkInPlace, + 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() + fork_in_place(bootstrapped, empty, hooks=hooks) for runtime in (bootstrapped, empty): - runtime._reinitialize_after_fork() await runtime.start() with pytest.raises(RuntimeError, match=FORK_BLOCKED): await runtime.refresh() @@ -623,16 +837,19 @@ async def test_on_macos_an_async_child_serves_its_bootstrap_and_cannot_refresh( assert server.requests == [] +@HOOKS def test_a_child_raises_from_start_and_wait_once_the_parent_used_the_http_client( - make_runtime: Make, monkeypatch: pytest.MonkeyPatch + hooks: bool, + fork_in_place: ForkInPlace, + 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() + fork_in_place(runtime, hooks=hooks) with pytest.raises(RuntimeError, match="caller-supplied HTTP client"): runtime.start() started = time.monotonic() @@ -642,26 +859,67 @@ def test_a_child_raises_from_start_and_wait_once_the_parent_used_the_http_client @pytest.mark.asyncio +@HOOKS async def test_an_async_child_raises_from_start_and_wait_once_the_parent_used_it( - make_async_runtime: MakeAsync, monkeypatch: pytest.MonkeyPatch + hooks: bool, + fork_in_place: ForkInPlace, + 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() + fork_in_place(runtime, hooks=hooks) 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) +@dataclass +class ParentPoller: + """What a client polling before an in-place fork leaves in this process. + + A forked child wouldn't have the parent's poller thread or HTTP client; + a test that forks in place stops them once the child is done. + """ + + thread: threading.Thread + stop_event: threading.Event + wake: threading.Event + http_client: Optional[httpx2.Client] + + @classmethod + def of(cls, runtime: FeatureFlagsRuntimeClient) -> ParentPoller: + assert runtime._thread is not None + return cls( + runtime._thread, runtime._stop, runtime._wake, runtime._transport._owned + ) + + def stop(self) -> None: + self.stop_event.set() + self.wake.set() + self.thread.join(5) + if self.http_client is not None: + self.http_client.close() + + +@HOOKS @pytest.mark.parametrize( - "read", ["is_ready", "configuration", "is_closed", "status", "get_all_flags"] + "read", + [ + "is_ready", + "configuration", + "is_closed", + "status", + "get_all_flags", + "is_enabled", + "get_flag", + ], ) def test_every_read_resumes_polling_after_a_fork( - read: str, monkeypatch: pytest.MonkeyPatch + read: str, hooks: bool, fork_in_place: ForkInPlace, monkeypatch: pytest.MonkeyPatch ) -> None: monkeypatch.setattr(_core, "FORK_SAFE_PLATFORM", True) server = LocalHTTPServer(payload(f=flag("f", default_value=True))) @@ -669,22 +927,122 @@ def test_every_read_resumes_polling_after_a_fork( runtime = workos.create_feature_flags_runtime_client() runtime.start() runtime.wait_until_ready(timeout=5) - first_poller = runtime._thread + parent = ParentPoller.of(runtime) - # What the at-fork hook runs in a child, without forking pytest. - runtime._reinitialize_after_fork() - assert runtime._fork_restart_pending + fork_in_place(runtime, hooks=hooks) value = getattr(runtime, read) - if callable(value): + if read in ("is_enabled", "get_flag"): + value("f") + elif callable(value): value() - assert runtime._thread is not None and runtime._thread is not first_poller + assert runtime._thread is not None and runtime._thread is not parent.thread assert runtime._thread.is_alive() runtime.close() + parent.stop() workos.close() server.close() +@HOOKS +def test_closing_in_a_child_leaves_the_parents_poller_and_http_client_alone( + hooks: bool, fork_in_place: ForkInPlace, 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) + parent = ParentPoller.of(runtime) + assert parent.http_client is not None + + fork_in_place(runtime, hooks=hooks) + runtime.close() + + assert not parent.stop_event.is_set() + assert not parent.http_client.is_closed + parent.stop() + workos.close() + server.close() + + +def test_threads_racing_to_first_use_after_a_fork_without_hooks_prepare_it_once( + fork_in_place: ForkInPlace, 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) + parent = ParentPoller.of(runtime) + prepared: List[int] = [] + spawned: List[int] = [] + reinitialize = runtime._reinitialize_after_fork + spawn_poller = runtime._spawn_poller + + def slow_reinitialize() -> None: + prepared.append(1) + # Long enough for the other threads to arrive while it runs. + time.sleep(0.2) + reinitialize() + + def counted_spawn_poller() -> None: + spawned.append(1) + spawn_poller() + + monkeypatch.setattr(runtime, "_reinitialize_after_fork", slow_reinitialize) + monkeypatch.setattr(runtime, "_spawn_poller", counted_spawn_poller) + + fork_in_place(runtime, hooks=False) + barrier = threading.Barrier(8) + answers: List[bool] = [] + + def first_use() -> None: + barrier.wait() + answers.append(runtime.is_enabled("f")) + + threads = [threading.Thread(target=first_use) for _ in range(8)] + for thread in threads: + thread.start() + for thread in threads: + thread.join(5) + + assert answers == [True] * 8 + assert prepared == [1] + assert spawned == [1] + runtime.close() + parent.stop() + workos.close() + server.close() + + +def test_a_client_that_cannot_be_prepared_after_a_fork_without_hooks_keeps_evaluating( + fork_in_place: ForkInPlace, + make_runtime: Make, + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, +) -> None: + runtime = make_runtime(bootstrap=FeatureFlagsConfiguration.from_dict(BOOTSTRAP)) + attempts: List[int] = [] + + def fail() -> None: + attempts.append(1) + raise RuntimeError("could not rebuild") + + monkeypatch.setattr(runtime, "_reinitialize_after_fork", fail) + + fork_in_place(runtime, hooks=False) + with caplog.at_level(logging.ERROR): + assert runtime.is_enabled("f") is True + assert runtime.is_enabled("f") is True + + assert attempts == [1] + [error] = [r for r in caplog.records if r.levelno == logging.ERROR] + assert "Could not prepare a feature flags runtime client" in error.getMessage() + + @dataclass class UnhashableBackend: """An HTTPBackend that, like any dataclass with eq=True, can't be hashed.""" @@ -713,23 +1071,24 @@ def close(self) -> None: pass +@HOOKS def test_an_unhashable_custom_backend_survives_the_fork_hook( - monkeypatch: pytest.MonkeyPatch, + hooks: bool, fork_in_place: ForkInPlace, 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() + fork_in_place(unused, hooks=hooks) unused.refresh() assert backend.requests - # Another WorkOSClient on the same backend sees that request too. + # Another WorkOSClient on the same backend sees that request too, in a + # process forked after it. sharing = WorkOSClient(api_key=API_KEY, http_client=backend) blocked = sharing.create_feature_flags_runtime_client() - blocked._reinitialize_after_fork() + fork_in_place(blocked, hooks=hooks) with pytest.raises(RuntimeError, match="caller-supplied HTTP client"): blocked.start() unused.close() diff --git a/tests/test_http_backends.py b/tests/test_http_backends.py index 3b169fd0..afd30f76 100644 --- a/tests/test_http_backends.py +++ b/tests/test_http_backends.py @@ -4,6 +4,7 @@ from __future__ import annotations +import gc import json import subprocess import sys @@ -26,7 +27,8 @@ from workos import _base_client as base_client_module from workos._base_client import _BaseWorkOSClient from workos._errors import WorkOSConnectionError, WorkOSTimeoutError -from workos._http import AsyncHttpxBackend, HttpxBackend +from workos import _http +from workos._http import AsyncHttpxBackend, HttpxBackend, record_request from tests.generated_helpers import load_fixture API_KEY = "sk_test_123" @@ -414,6 +416,20 @@ def handler(request: Any) -> Any: ) +class TestConnectionTracking: + def test_forgets_an_http_client_once_it_is_collected(self) -> None: + http_client = httpx2.Client() + record_request(HttpxBackend(http_client)) + key = id(http_client) + assert key in _http._first_requests + + http_client.close() + del http_client + gc.collect() + + assert key not in _http._first_requests + + class TestResolution: @pytest.mark.parametrize( "httpx_installed", [False, True], ids=["without-httpx", "with-httpx"]