diff --git a/instructor/v2/core/registry.py b/instructor/v2/core/registry.py index 66419ebd8..8501bda36 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,12 @@ def __init__(self) -> None: """Initialize empty registry.""" self._handlers: dict[tuple[Provider, Mode], ModeHandlers] = {} self._lazy_loaders: dict[tuple[Provider, Mode], Callable[[], ModeHandlers]] = {} + # Guards the lazy-load resolution (check -> pop -> import -> set) in + # get_handlers(). Without this, concurrent first-callers for the same + # mode_key race: one pops the loader, the others see neither dict + # populated yet and raise KeyError. Held for the whole resolution + # (not per-key) since lazy-loading only ever runs once per key. + self._lazy_load_lock = threading.Lock() def register( self, @@ -181,12 +188,22 @@ def get_handlers(self, provider: Provider, mode: Mode) -> ModeHandlers: 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 + # Try lazy loading. Locked because the pop -> import -> set sequence + # below is not atomic: without the lock, a thread that loses the race + # to pop self._lazy_loaders[mode_key] would find it already gone and + # self._handlers[mode_key] not yet set, and raise KeyError even + # though the mode genuinely is registered (just still resolving). + with self._lazy_load_lock: + # Re-check: another thread may have finished loading this key + # while we were waiting for the lock. + 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. " diff --git a/tests/v2/test_registry.py b/tests/v2/test_registry.py index 5ffc2142b..d58bb3270 100644 --- a/tests/v2/test_registry.py +++ b/tests/v2/test_registry.py @@ -116,3 +116,69 @@ def test_registry_invalid_handler_type(provider: Provider, mode: Mode): """Test error for invalid handler type.""" with pytest.raises(ValueError, match="Invalid handler_type"): mode_registry.get_handler(provider, mode, "invalid_type") + + +def test_get_handlers_concurrent_first_access_does_not_race(): + """Regression test for #2422. + + Concurrent first callers for the same lazily-registered (provider, mode) + key must all get the same handlers back, never a KeyError, even when the + loader is slow (simulating a module import in flight). + + The loader sleeps briefly so every other thread has a chance to reach + get_handlers() while the first thread's load is still in progress, this + is the exact window in which the unlocked version raced: losing threads + would see the lazy loader already popped and self._handlers not yet set, + and raise KeyError for a mode that genuinely is registered. + """ + import threading + import time + + from instructor.v2.core.registry import ModeHandlers, ModeRegistry + + registry = ModeRegistry() + n_threads = 8 + started = threading.Barrier(n_threads, timeout=5) + load_count = 0 + load_count_lock = threading.Lock() + + def slow_loader() -> ModeHandlers: + nonlocal load_count + with load_count_lock: + load_count += 1 + # Simulate a slow module import: gives every other thread time to + # reach get_handlers() and start contending before this resolves. + time.sleep(0.2) + return ModeHandlers( + request_handler=lambda *_a, **_k: None, + reask_handler=lambda *_a, **_k: None, + response_parser=lambda *_a, **_k: None, + ) + + registry.register_lazy(Provider.DEEPSEEK, Mode.TOOLS, slow_loader) + + results: list[object] = [None] * n_threads + errors: list[BaseException] = [] + errors_lock = threading.Lock() + + def worker(idx: int) -> None: + # Line up all threads so they call get_handlers() as close to + # simultaneously as possible, maximizing contention. + started.wait(timeout=5) + try: + results[idx] = registry.get_handlers(Provider.DEEPSEEK, Mode.TOOLS) + except BaseException as exc: # noqa: BLE001 + with errors_lock: + errors.append(exc) + + threads = [threading.Thread(target=worker, args=(i,)) for i in range(n_threads)] + for t in threads: + t.start() + for t in threads: + t.join(timeout=10) + + assert errors == [], f"get_handlers raised for concurrent callers: {errors}" + assert all(r is not None for r in results) + assert len({id(r) for r in results}) == 1, "all callers should get the same ModeHandlers instance" + # The loader must run exactly once, not once per racing thread. + assert load_count == 1