From aab60fb90aa068cc462018cbb98c274c07664415 Mon Sep 17 00:00:00 2001 From: Aikins Laryea Date: Sat, 31 Oct 2026 16:10:22 +0000 Subject: [PATCH] fix(sdk-python): recover prompt fetches and the cache lock in forked children A fork in a signal handler during a prompt fetch wait now starts a new fetch in the child instead of waiting for the timeout (#107). The prompt cache has one lock per process, keyed by os.getpid(), so a fork that runs no hooks cannot leave the child waiting for a lock held by a parent thread (#106). The prompt tests count only requests under a random path of the capture server URL (#114). --- sdks/python/src/telemetry_dev/_prompts.py | 65 +++++---- sdks/python/tests/test_prompts.py | 156 +++++++++++++++++----- 2 files changed, 163 insertions(+), 58 deletions(-) diff --git a/sdks/python/src/telemetry_dev/_prompts.py b/sdks/python/src/telemetry_dev/_prompts.py index 98f2af5d..63739f96 100644 --- a/sdks/python/src/telemetry_dev/_prompts.py +++ b/sdks/python/src/telemetry_dev/_prompts.py @@ -41,6 +41,8 @@ from ._serialize import AttributeValue, fit_object_values, json_dumps, serialize_content _FETCH_TIMEOUT_S = 10.0 +# How often a call that waits for a fetch checks that it still runs in the process of the fetch. +_FORK_CHECK_S = 0.1 # How long a failed fetch with no copy answers calls, so that they do not each fetch again. _ERROR_TTL_S = 10.0 # The longest Retry-After that get_prompt honors, so that a bad header cannot stop fetches. @@ -617,37 +619,47 @@ class _Entry: _cache: OrderedDict[tuple[str, str, str, str], _Entry] = OrderedDict() -# Reentrant, so that a fork in a thread that holds the lock, for example in a signal handler, does -# not wait for itself in _lock_for_fork. _reset_in_child can replace it, so read it from the module -# each time. -_cache_lock = threading.RLock() +# The cache lock of each process, by process ID. A forked process gets a copy of the lock of its +# parent, and when a thread of the parent held it, no thread in the child can release the copy. A +# fork that runs no hooks, as a fork in C code, does not reset it, but the child has a new ID. +_cache_locks: dict[int, threading.RLock] = {} # Whether _lock_for_fork got the lock in this thread. os.fork() runs the after hooks also when a # before hook raises, for example in a signal handler. _forking = threading.local() +def _cache_lock() -> threading.RLock: + """Reentrant, so that a fork in a thread that holds the lock, for example in a signal handler, + does not wait for itself in _lock_for_fork.""" + pid = os.getpid() + lock = _cache_locks.get(pid) + if lock is None: + # The first call in this process. The locks of other processes are of no use here. + # list() copies the keys at once: other threads can change the dict. + for other in list(_cache_locks): + if other != pid: + _cache_locks.pop(other, None) + # setdefault, so that threads that get here at the same time share one lock. + lock = _cache_locks.setdefault(pid, threading.RLock()) + return lock + + def _lock_for_fork() -> None: - _cache_lock.acquire() + _cache_lock().acquire() _forking.locked = True def _unlock_after_fork() -> None: if getattr(_forking, "locked", False): _forking.locked = False - _cache_lock.release() + _cache_lock().release() def _reset_in_child() -> None: - """Unlock the cache, and drop the fetches of the parent: the child has none of its threads.""" - global _cache_lock - if getattr(_forking, "locked", False): - _forking.locked = False - _cache_lock.release() - else: - # The before hook did not get the lock, because a signal stopped its wait, or did not run, - # as in uWSGI with py-call-osafterfork. A thread of the parent can hold the lock, and only - # this thread exists here. - _cache_lock = threading.RLock() + """Drop the locks and the fetches of the parent: the child has none of its threads. The locks + go too, because the child can get the ID of an ended process whose lock _cache_locks keeps.""" + _forking.locked = False + _cache_locks.clear() for entry in _cache.values(): entry.pending = None @@ -673,7 +685,7 @@ def _load( older than the TTL. Only a call with no copy waits for a fetch, until the fetch's deadline. A failed fetch with no copy fails the calls of the next ``_ERROR_TTL_S`` (or the TTL, if it is shorter) without a fetch.""" - with _cache_lock: + with _cache_lock(): entry = _cache.get(key) if entry is None: if len(_cache) >= _MAX_CACHE_ENTRIES: @@ -682,9 +694,7 @@ def _load( now = monotonic() ttl = ttl_seconds if entry.response is not None else min(ttl_seconds, _ERROR_TTL_S) if entry.pending is not None and entry.pending_pid != os.getpid(): - # A fork in C code that runs no hooks copied the fetch but not its thread. This check - # does not replace the lock. If a thread of the parent held the lock at that fork, a - # call in the child waits for the lock with no limit. + # A fork in C code that runs no hooks copied the fetch but not its thread. entry.pending = None # A local, because a fork in this thread, as in a signal handler, can clear entry.pending. pending = entry.pending @@ -698,9 +708,16 @@ def _load( if entry.response is not None: return entry.response deadline = entry.deadline - if pending is not None and not pending.wait(deadline - monotonic()): - raise _timeout() - with _cache_lock: + pending_pid = entry.pending_pid + if pending is not None: + while not pending.wait(min(_FORK_CHECK_S, deadline - monotonic())): + if os.getpid() != pending_pid: + # A fork in a signal handler during the wait. The child has the wait, but not the + # thread of the fetch. + return _load(key, ttl_seconds, fetch, on_error) + if monotonic() >= deadline: + raise _timeout() + with _cache_lock(): response, error = entry.response, entry.error if response is not None: return response @@ -723,7 +740,7 @@ def _refresh( response = fetch() except Exception as exc: error = exc - with _cache_lock: + with _cache_lock(): stale = entry.response delay = ( error._retry_after # pyright: ignore[reportPrivateUsage] diff --git a/sdks/python/tests/test_prompts.py b/sdks/python/tests/test_prompts.py index af449c9c..a37814d0 100644 --- a/sdks/python/tests/test_prompts.py +++ b/sdks/python/tests/test_prompts.py @@ -11,6 +11,7 @@ import os import pickle import re +import secrets import signal import socket import sys @@ -21,6 +22,7 @@ import urllib.request import weakref from collections.abc import Callable, Generator +from contextlib import contextmanager from datetime import datetime, timedelta, timezone from email.utils import format_datetime from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer @@ -48,7 +50,7 @@ from telemetry_dev._serialize import utf16_length from .conftest import MakeClient -from .otlp_capture import CaptureServer, otlp_capture_server +from .otlp_capture import CapturedRequest, CaptureServer, otlp_capture_server # Python 3.12+ gives a warning for fork() in a process with threads, and the capture server has one. pytestmark = pytest.mark.filterwarnings( @@ -83,9 +85,28 @@ def api_error(status: int, code: str, message: str) -> tuple[int, object]: return status, {"error": {"code": code, "message": message}} +@contextmanager +def prompt_server() -> Generator[CaptureServer]: + """A capture server whose URL has a random path. Another local process can send requests to its + port; ``own_requests`` leaves them out.""" + with otlp_capture_server() as capture: + capture.url += f"/{secrets.token_hex(8)}" + yield capture + + +def own_requests(server: CaptureServer) -> list[CapturedRequest]: + """The requests to the URL of ``server``, with the paths that the SDK asked for.""" + prefix = urlsplit(server.url).path + return [ + dataclasses.replace(request, path=request.path.removeprefix(prefix)) + for request in server.requests + if request.path.startswith(f"{prefix}/") + ] + + @pytest.fixture def server() -> Generator[CaptureServer]: - with otlp_capture_server() as capture: + with prompt_server() as capture: capture.replies = [(200, reply())] yield capture @@ -207,7 +228,7 @@ def test_get_prompt_sends_query_with_key_and_sdk_header( ) -> None: get_prompt("support/reply", api_key="td_live_a", base_url=f"{server.url}/", **options) - [request] = server.requests + [request] = own_requests(server) assert urlsplit(request.path).path == "/v1/prompts" assert query_of(request.path) == {"name": "support/reply", **query} assert request.headers["authorization"] == "Bearer td_live_a" @@ -220,14 +241,14 @@ def test_get_prompt_reads_the_initialized_client_then_the_environment( ) -> None: make(api_key="td_live_client", base_url=server.url) assert get_prompt("support/reply").version == 12 - assert server.requests[0].headers["authorization"] == "Bearer td_live_client" + assert own_requests(server)[0].headers["authorization"] == "Bearer td_live_client" telemetry_dev.shutdown() telemetry_dev.init(api_key="td_live_disabled", enabled=False, disable_atexit=True) monkeypatch.setenv("TELEMETRY_DEV_API_KEY", "td_live_env") monkeypatch.setenv("TELEMETRY_DEV_BASE_URL", server.url) get_prompt("support/reply", label="staging") - assert server.requests[1].headers["authorization"] == "Bearer td_live_env" + assert own_requests(server)[1].headers["authorization"] == "Bearer td_live_env" def test_get_prompt_reads_the_host_that_telemetry_goes_to(server: CaptureServer) -> None: @@ -240,7 +261,7 @@ def test_get_prompt_reads_the_host_that_telemetry_goes_to(server: CaptureServer) get_prompt("support/reply") - assert [urlsplit(request.path).path for request in server.requests] == [ + assert [urlsplit(request.path).path for request in own_requests(server)] == [ "/v1/traces", "/v1/prompts", ] @@ -281,19 +302,19 @@ def version() -> int | None: assert version() == 12 clock[0] += 59 assert version() == 12 - assert len(server.requests) == 1 + assert len(own_requests(server)) == 1 # After the TTL, calls get the copy at once while one refresh runs. server.delay = 0.5 clock[0] += 1 assert [version() for _ in range(3)] == [12, 12, 12] wait_for(lambda: version() == 13) - assert len(server.requests) == 2 + assert len(own_requests(server)) == 2 # A different key does not share the copy. server.delay = 0 get_prompt("support/reply", api_key="td_live_b", base_url=server.url) - assert len(server.requests) == 3 + assert len(own_requests(server)) == 3 def test_concurrent_calls_share_one_fetch(server: CaptureServer) -> None: @@ -311,7 +332,7 @@ def fetch() -> None: for thread in threads: thread.join() - assert len(server.requests) == 1 + assert len(own_requests(server)) == 1 assert versions == [12] * 5 @@ -335,7 +356,7 @@ def test_a_failed_refresh_returns_the_last_copy_and_reports_the_error( # The stale copy counts as fresh for another TTL, so calls do not each start a fetch. get_prompt("support/reply") - assert len(server.requests) == 2 + assert len(own_requests(server)) == 2 clock[0] += 60 assert get_prompt("support/reply").version == 12 @@ -370,15 +391,15 @@ def fallback_prompt(cache_ttl_seconds: float = 60) -> ChatPrompt: clock[0] += 9 assert fallback_prompt().is_fallback - assert len(server.requests) == 1 + assert len(own_requests(server)) == 1 clock[0] += 1 assert fallback_prompt().is_fallback - assert len(server.requests) == 2 + assert len(own_requests(server)) == 2 # A TTL under 10 seconds bounds the wait too. fallback_prompt(cache_ttl_seconds=0) - assert len(server.requests) == 3 + assert len(own_requests(server)) == 3 @pytest.mark.parametrize( @@ -418,11 +439,11 @@ def fallback_prompt() -> ChatPrompt: # The HTTP date counts from when the server answered, a moment before this. clock[0] += wait - 2 assert fallback_prompt().is_fallback - assert len(server.requests) == 1 + assert len(own_requests(server)) == 1 clock[0] += 2 fallback_prompt() - assert len(server.requests) == 2 + assert len(own_requests(server)) == 2 def test_a_refresh_that_gets_a_retry_after_keeps_the_copy_until_then( @@ -447,7 +468,7 @@ def version() -> int | None: for entry in prompts_module._cache.values() # pyright: ignore[reportPrivateUsage] ) ) - assert len(server.requests) == 2 + assert len(own_requests(server)) == 2 # The TTL ends after 60 seconds, but the server asked for 300. clock[0] += 299 @@ -458,12 +479,12 @@ def version() -> int | None: for entry in prompts_module._cache.values() # pyright: ignore[reportPrivateUsage] ) ) - assert len(server.requests) == 2 + assert len(own_requests(server)) == 2 clock[0] += 1 version() wait_for(lambda: version() == 13) - assert len(server.requests) == 3 + assert len(own_requests(server)) == 3 def test_with_no_copy_and_no_fallback_get_prompt_raises_a_prompt_error( @@ -512,7 +533,7 @@ def test_a_reply_without_a_label_is_an_invalid_response(server: CaptureServer) - def test_a_redirect_fails_and_does_not_send_the_key(server: CaptureServer) -> None: - with otlp_capture_server() as other: + with prompt_server() as other: other.replies = [(200, reply())] server.replies = [(302, {})] server.location = f"{other.url}/v1/prompts?name=support%2Freply&label=production" @@ -521,7 +542,7 @@ def test_a_redirect_fails_and_does_not_send_the_key(server: CaptureServer) -> No get_prompt("support/reply", api_key="td_live_a", base_url=server.url) assert (raised.value.code, raised.value.status) == ("http_302", 302) - assert other.requests == [] + assert own_requests(other) == [] @pytest.mark.parametrize("slow", ["reply", "address lookup"]) @@ -637,7 +658,7 @@ def test_a_forked_process_fetches_the_prompt_that_the_parent_was_fetching( versions: list[int | None] = [] parent = threading.Thread(target=lambda: versions.append(version_of(server))) parent.start() - wait_for(lambda: len(server.requests) == 1) + wait_for(lambda: len(own_requests(server)) == 1) server.delay = 0 parent_pid = os.getpid() @@ -651,7 +672,7 @@ def child() -> str: parent.join() assert versions == [12] # The fetch of the parent, and the fetch of the child. - assert len(server.requests) == 2 + assert len(own_requests(server)) == 2 def test_a_call_does_not_wait_for_a_fetch_that_another_process_started( @@ -660,7 +681,7 @@ def test_a_call_does_not_wait_for_a_fetch_that_another_process_started( server.delay = 1 first = threading.Thread(target=version_of, args=(server,)) first.start() - wait_for(lambda: len(server.requests) == 1) + wait_for(lambda: len(own_requests(server)) == 1) server.delay = 0 # A fork that runs no fork hooks, as a fork in C code, copies the fetch of the parent but not # the thread that runs it. Here this process is the child: the fetch gets the id of its parent. @@ -668,10 +689,38 @@ def test_a_call_does_not_wait_for_a_fetch_that_another_process_started( entry.pending_pid = os.getppid() assert version_of(server) == 12 - assert len(server.requests) == 2 + assert len(own_requests(server)) == 2 first.join() +def test_a_call_does_not_wait_for_a_cache_lock_that_another_process_held( + server: CaptureServer, monkeypatch: pytest.MonkeyPatch +) -> None: + held, done = threading.Event(), threading.Event() + released_at: list[float] = [] + + def hold() -> None: + with prompts_module._cache_lock(): # pyright: ignore[reportPrivateUsage] + held.set() + done.wait(5) + released_at.append(time.monotonic()) + + holder = threading.Thread(target=hold) + holder.start() + held.wait() + # A fork that runs no fork hooks, as a fork in C code, copies the lock while a thread of the + # parent holds it. Here this process is the child: it gets a new process ID. + pid = os.getpid() + monkeypatch.setattr(os, "getpid", lambda: pid + 1) + try: + assert version_of(server) == 12 + returned_at = time.monotonic() + finally: + done.set() + holder.join() + assert returned_at < released_at[0] + + @needs_fork def test_a_forked_process_refreshes_the_copy_that_the_parent_was_refreshing( server: CaptureServer, clock: list[float] @@ -681,7 +730,7 @@ def test_a_forked_process_refreshes_the_copy_that_the_parent_was_refreshing( server.replies = [(200, reply(version=13))] server.delay = 1 assert version_of(server) == 12 - wait_for(lambda: len(server.requests) == 2) + wait_for(lambda: len(own_requests(server)) == 2) server.delay = 0 def child() -> str: @@ -703,7 +752,7 @@ def test_a_fork_waits_for_a_thread_that_holds_the_cache_lock(server: CaptureServ def hold() -> None: try: - with prompts_module._cache_lock: # pyright: ignore[reportPrivateUsage] + with prompts_module._cache_lock(): # pyright: ignore[reportPrivateUsage] held.set() time.sleep(0.2) released_at.append(time.monotonic()) @@ -734,7 +783,7 @@ def test_a_fork_in_the_thread_that_holds_the_cache_lock_returns( monkeypatch.setattr(prompts_module, "_FETCH_TIMEOUT_S", 3) def fork_in_the_lock() -> int: - with prompts_module._cache_lock: # pyright: ignore[reportPrivateUsage] + with prompts_module._cache_lock(): # pyright: ignore[reportPrivateUsage] return os.fork() def fork_and_fetch() -> str: @@ -782,6 +831,47 @@ def fork() -> int: assert results == ["12"] +@needs_fork +@pytest.mark.skipif(not hasattr(signal, "setitimer"), reason="needs signal.setitimer") +def test_a_fork_in_a_signal_handler_during_the_wait_for_a_fetch_fetches_again_in_the_child( + server: CaptureServer, monkeypatch: pytest.MonkeyPatch +) -> None: + # A child that waits for the fetch of the parent gets no answer in 3 s. + monkeypatch.setattr(prompts_module, "_FETCH_TIMEOUT_S", 3) + server.delay = 1 + pids: list[int] = [] + results: list[str] = [] + forked_in: list[str | None] = [] + + def fork_in_the_wait(signum: int, frame: FrameType | None) -> None: + server.delay = 0 + forked_in.append(frame and frame.f_code.co_name) + pids.append(os.fork()) + + def fork() -> int: + previous = signal.signal(signal.SIGALRM, fork_in_the_wait) + # Between two 0.1 s steps of the wait. + signal.setitimer(signal.ITIMER_REAL, 0.35) + started = time.monotonic() + try: + version = version_of(server) + results.append(f"{version} in under 2 s: {time.monotonic() - started < 2}") + except Exception as exc: + # The child must return the pid to in_fork, and not raise into pytest. + results.append(repr(exc)) + finally: + signal.setitimer(signal.ITIMER_REAL, 0) + signal.signal(signal.SIGALRM, previous) + return pids[0] + + assert in_fork(lambda: results[0], fork) == "12 in under 2 s: True" + assert results == ["12 in under 2 s: True"] + # The signal came in the wait of threading.Event, not before it. + assert forked_in == ["wait"] + # The fetch of the parent, and the fetch of the child. + assert len(own_requests(server)) == 2 + + @needs_fork def test_a_fork_after_a_reload_of_the_prompt_module_returns( server: CaptureServer, monkeypatch: pytest.MonkeyPatch @@ -803,14 +893,12 @@ def test_a_fork_after_a_reload_of_the_prompt_module_returns( @pytest.mark.skipif(not hasattr(signal, "setitimer"), reason="needs signal.setitimer") @pytest.mark.parametrize("before_hook", ["stopped by a signal", "not run"]) def test_the_fork_hooks_handle_a_lock_that_the_before_hook_did_not_get( - server: CaptureServer, monkeypatch: pytest.MonkeyPatch, before_hook: str + server: CaptureServer, before_hook: str ) -> None: class Interrupted(Exception): pass - lock = prompts_module._cache_lock # pyright: ignore[reportPrivateUsage] - # The child hook replaces the lock. monkeypatch puts it back. - monkeypatch.setattr(prompts_module, "_cache_lock", lock) + lock = prompts_module._cache_lock() # pyright: ignore[reportPrivateUsage] held, done = threading.Event(), threading.Event() errors: list[BaseException] = [] @@ -856,7 +944,7 @@ def record(unraisable: sys.UnraisableHookArgs) -> None: else: # uWSGI with py-call-osafterfork forks in C and runs only the child hook. prompts_module._reset_in_child() # pyright: ignore[reportPrivateUsage] - child_lock = prompts_module._cache_lock # pyright: ignore[reportPrivateUsage] + child_lock = prompts_module._cache_lock() # pyright: ignore[reportPrivateUsage] assert child_lock.acquire(blocking=False) child_lock.release() finally: @@ -1985,7 +2073,7 @@ async def tick() -> None: ticker.cancel() assert prompt.version == 12 - assert query_of(server.requests[0].path)["label"] == "production" + assert query_of(own_requests(server)[0].path)["label"] == "production" # The event loop ran while the fetch waited for the slow reply. A fetch on the loop's thread # would block it, and the ticker would not run until the fetch ended. assert ticks >= 5