diff --git a/CHANGELOG.md b/CHANGELOG.md index 55bb558c0..c815ff66f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,13 @@ Versioning: [Semantic Versioning](https://semver.org/spec/v2.0.0.html) --- +## [Unreleased] + +### Fixed +- **v2 registry**: Fix thread-safety race condition in `ModeRegistry.get_handlers()` where concurrent first-access to a lazy-loaded mode caused `KeyError`/`RegistryError` for all threads except the winner (49/50 failures in testing). Added double-checked locking with `threading.Lock` around the lazy-load path. ([#2422](https://github.com/567-labs/instructor/issues/2422)) + +--- + ## [1.15.5] - 2026-06-28 ### Fixed diff --git a/instructor/v2/core/registry.py b/instructor/v2/core/registry.py index 66419ebd8..6b2e66e1f 100644 --- a/instructor/v2/core/registry.py +++ b/instructor/v2/core/registry.py @@ -6,6 +6,7 @@ from __future__ import annotations +import threading from dataclasses import dataclass from typing import Callable @@ -80,6 +81,7 @@ def __init__(self) -> None: """Initialize empty registry.""" self._handlers: dict[tuple[Provider, Mode], ModeHandlers] = {} self._lazy_loaders: dict[tuple[Provider, Mode], Callable[[], ModeHandlers]] = {} + self._lock = threading.Lock() def register( self, @@ -110,18 +112,19 @@ def register( ConfigurationError: If mode is already registered with different handlers """ mode_key = (provider, mode) - if mode_key in self._lazy_loaders: - self._lazy_loaders.pop(mode_key, None) - - self._handlers[mode_key] = ModeHandlers( - request_handler=request_handler, - reask_handler=reask_handler, - response_parser=response_parser, - stream_extractor=stream_extractor, - stream_extractor_async=stream_extractor_async, - message_converter=message_converter, - template_handler=template_handler, - ) + with self._lock: + if mode_key in self._lazy_loaders: + self._lazy_loaders.pop(mode_key, None) + + self._handlers[mode_key] = ModeHandlers( + request_handler=request_handler, + reask_handler=reask_handler, + response_parser=response_parser, + stream_extractor=stream_extractor, + stream_extractor_async=stream_extractor_async, + message_converter=message_converter, + template_handler=template_handler, + ) def register_lazy( self, @@ -144,10 +147,11 @@ def register_lazy( from instructor.v2.core.errors import ConfigurationError mode_key = (provider, mode) - if mode_key in self._handlers or mode_key in self._lazy_loaders: - raise ConfigurationError(f"Mode {mode_key} is already registered") + with self._lock: + if mode_key in self._handlers or mode_key in self._lazy_loaders: + raise ConfigurationError(f"Mode {mode_key} is already registered") - self._lazy_loaders[mode_key] = loader + self._lazy_loaders[mode_key] = loader def get_handlers(self, provider: Provider, mode: Mode) -> ModeHandlers: """Get all handlers for a mode. @@ -177,21 +181,31 @@ def get_handlers(self, provider: Provider, mode: Mode) -> ModeHandlers: normalized_mode = normalize_mode(provider, mode) mode_key = (provider, normalized_mode) - # Check if already loaded + # Fast path: check if already loaded (no lock needed for dict read + # in CPython — the GIL makes single dict.__contains__ atomic, and + # once a key is in _handlers it's never removed). if mode_key in self._handlers: return self._handlers[mode_key] - # Try lazy loading - if mode_key in self._lazy_loaders: - loader = self._lazy_loaders.pop(mode_key) - handlers = loader() - self._handlers[mode_key] = handlers - return handlers - - raise KeyError( - f"Mode {mode_key} is not registered. " - f"Available modes: {list(self._handlers.keys())}" - ) + # Slow path: lazy loading requires a lock to prevent a race where + # multiple threads concurrently pop the lazy loader entry, leaving + # losing threads with no entry in either dict and a permanent KeyError. + with self._lock: + # Double-checked locking: re-check _handlers inside the lock + # in case another thread already resolved it while we waited. + if mode_key in self._handlers: + return self._handlers[mode_key] + + if mode_key in self._lazy_loaders: + loader = self._lazy_loaders.pop(mode_key) + handlers = loader() + self._handlers[mode_key] = handlers + return handlers + + raise KeyError( + f"Mode {mode_key} is not registered. " + f"Available modes: {list(self._handlers.keys())}" + ) def get_handler( self, diff --git a/tests/core/test_registry_thread_safety.py b/tests/core/test_registry_thread_safety.py new file mode 100644 index 000000000..cc1881641 --- /dev/null +++ b/tests/core/test_registry_thread_safety.py @@ -0,0 +1,174 @@ +"""Tests for ModeRegistry thread safety (issue #2422). + +Verifies that concurrent first-access to a lazy-loaded mode handler does not +cause permanent KeyError/RegistryError for losing threads. +""" + +from __future__ import annotations + +import threading +from concurrent.futures import ThreadPoolExecutor + +import pytest + +from instructor.v2.core.mode import Mode +from instructor.v2.core.providers import Provider +from instructor.v2.core.registry import ModeHandlers, ModeRegistry + + +def _make_dummy_handlers() -> ModeHandlers: + """Create minimal ModeHandlers for testing.""" + + def dummy_request(_response_model=None, kwargs=None): # noqa: ARG001 + return kwargs or {} + + def dummy_reask(kwargs=None, _response=None, _exception=None): # noqa: ARG001 + return kwargs or {} + + def dummy_response( + response, + _response_model=None, # noqa: ARG001 + _validation_context=None, # noqa: ARG001 + _mode=None, # noqa: ARG001 + _stream=None, # noqa: ARG001 + **_kw, # noqa: ARG001 + ): + return response + + return ModeHandlers( + request_handler=dummy_request, + reask_handler=dummy_reask, + response_parser=dummy_response, + ) + + +class TestRegistryThreadSafety: + """Tests that ModeRegistry.get_handlers() is thread-safe.""" + + def test_concurrent_first_access_no_failures(self): + """All threads should succeed when concurrently resolving the same + lazy-loaded mode for the first time. + + Before the fix, 49/50 threads would fail with KeyError because the + first thread pops the lazy loader entry before others can use it. + """ + registry = ModeRegistry() + call_count = 0 + count_lock = threading.Lock() + + def slow_loader(): + """Simulate a slow module import to widen the race window.""" + nonlocal call_count + with count_lock: + call_count += 1 + # Small delay to ensure threads overlap + threading.Event().wait(0.01) + return _make_dummy_handlers() + + registry.register_lazy(Provider.OPENAI, Mode.TOOLS, slow_loader) + + def check(i: int) -> tuple[int, str]: + try: + registry.get_handlers(Provider.OPENAI, Mode.TOOLS) + return (i, "OK") + except Exception as e: + return (i, f"{type(e).__name__}: {str(e)[:80]}") + + with ThreadPoolExecutor(max_workers=10) as pool: + results = list(pool.map(check, range(50))) + + failures = [r for r in results if r[1] != "OK"] + assert not failures, ( + f"{len(failures)}/50 threads failed under concurrent access:\n" + + "\n".join(f" thread {i}: {err}" for i, err in failures[:5]) + ) + + def test_loader_called_exactly_once(self): + """The lazy loader should be called exactly once even under heavy + concurrency, because the lock serializes the lazy-load path.""" + registry = ModeRegistry() + call_count = 0 + count_lock = threading.Lock() + + def loader(): + nonlocal call_count + with count_lock: + call_count += 1 + threading.Event().wait(0.01) + return _make_dummy_handlers() + + registry.register_lazy(Provider.OPENAI, Mode.TOOLS, loader) + + with ThreadPoolExecutor(max_workers=10) as pool: + list( + pool.map( + lambda _: registry.get_handlers(Provider.OPENAI, Mode.TOOLS), + range(30), + ) + ) + + assert call_count == 1, ( + f"Loader was called {call_count} times, expected exactly 1. " + "The lock should serialize lazy loading." + ) + + def test_double_checked_locking_resolves_quickly(self): + """After the first thread resolves the handler, subsequent threads + should get it from _handlers (fast path), not re-enter the lock.""" + registry = ModeRegistry() + + def loader(): + return _make_dummy_handlers() + + registry.register_lazy(Provider.OPENAI, Mode.TOOLS, loader) + + # First call triggers lazy load + registry.get_handlers(Provider.OPENAI, Mode.TOOLS) + # The loader entry should be consumed + assert (Provider.OPENAI, Mode.TOOLS) not in registry._lazy_loaders + # The handler should be cached + assert (Provider.OPENAI, Mode.TOOLS) in registry._handlers + + # Subsequent calls should hit the fast path + handlers2 = registry.get_handlers(Provider.OPENAI, Mode.TOOLS) + assert handlers2 is registry._handlers[(Provider.OPENAI, Mode.TOOLS)] + + def test_unregistered_mode_still_raises(self): + """An unregistered mode should still raise KeyError, not silently + succeed or hang.""" + registry = ModeRegistry() + + with pytest.raises(KeyError, match="is not registered"): + registry.get_handlers(Provider.OPENAI, Mode.TOOLS) + + def test_concurrent_different_modes(self): + """Concurrent access to different mode keys should not interfere.""" + registry = ModeRegistry() + + def make_loader(_p, _m): # noqa: ARG001 + return lambda: _make_dummy_handlers() + + modes = [ + (Provider.OPENAI, Mode.TOOLS), + (Provider.ANTHROPIC, Mode.TOOLS), + (Provider.GEMINI, Mode.TOOLS), + ] + + for provider, mode in modes: + registry.register_lazy(provider, mode, make_loader(provider, mode)) + + def check(key): + provider, mode = key + try: + registry.get_handlers(provider, mode) + return "OK" + except Exception as e: + return f"{type(e).__name__}: {e}" + + # Each key accessed by multiple threads simultaneously + keys = modes * 10 # 30 tasks, 10 per key + with ThreadPoolExecutor(max_workers=10) as pool: + results = list(pool.map(check, keys)) + + failures = [r for r in results if r != "OK"] + assert not failures, f"Concurrent different-mode access failed: {failures}"