diff --git a/reflexio/server/services/durable_learning/local.py b/reflexio/server/services/durable_learning/local.py index 7424defe..4f445fcb 100644 --- a/reflexio/server/services/durable_learning/local.py +++ b/reflexio/server/services/durable_learning/local.py @@ -1,17 +1,127 @@ """Library lifecycle registration for the same durable extraction scheduler.""" import threading +import weakref from reflexio.server.api_endpoints.request_context import RequestContext from reflexio.server.services.durable_learning.scheduler import DurableLearningScheduler _lock = threading.RLock() -_contexts: dict[tuple[str, str | None], RequestContext] = {} + +# Lifecycle choice: contexts are held WEAKLY, and a directory's scheduler is +# retired once every context for it has been collected. +# +# A library-local scheduler exists to serve live ``Reflexio`` handles, so "a +# handle pointed at this directory is still reachable" is the only honest +# signal that one is still needed. Reference counting would have said the same +# thing but needs an explicit ``close()`` on a public API that has none today — +# every caller would have to remember it, and the ones that forgot would leak +# exactly as before. A weak registry gets the same answer for free and cannot +# be forgotten. +# +# Retirement runs on the *caller's* thread inside ``ensure_local_extraction`` +# rather than from inside a tick: ``stop()`` joins the scheduler thread, and a +# thread cannot join itself. Nothing is lost by waiting for the next caller — +# durable work is persisted, and a dropped handle means nobody in this process +# is waiting for its results; a later ``Reflexio`` on the same directory picks +# the backlog up through the usual restart-recovery path. +# Two registries, because "which context should the worker use?" and "is any +# handle still alive?" are different questions and one structure cannot answer +# both. +# +# IDENTITY. The CURRENT context per (org, directory). Overwriting is the point: +# the worker must extract against the handle in use now, not an earlier one +# whose storage may be gone. +_contexts: "weakref.WeakValueDictionary[tuple[str, str | None], RequestContext]" = ( + weakref.WeakValueDictionary() +) +# LIVENESS. EVERY live context, grouped by directory. +# +# `_contexts` alone cannot answer this: `ReflexioBase` builds a distinct +# `RequestContext` per handle, so a second handle on the same org and directory +# evicts the first from that key. Collecting the second then empties the key +# while the first is still alive, and the sweep retires a scheduler that handle +# still depends on -- extraction stopping silently under a live caller. +# Measured on the identity map alone: with two same-key handles it held only +# the second. +# +# Answering BOTH from one structure was tried and is worse. A set alone loses +# identity: every e2e test shares one org and one directory, so the set +# accumulates ("ctxs=4", then 5...) and the factory hands the worker an +# arbitrary, usually stale context pointing at a previous test's storage. The +# work never completes and publishes time out -- measured, 8 passed in 14.9s +# became 2 failed in 275.8s. +_live: "dict[str | None, weakref.WeakSet[RequestContext]]" = {} _schedulers: dict[str | None, DurableLearningScheduler] = {} _server_scheduler: DurableLearningScheduler | None = None +def _start_scheduler(directory: str | None) -> DurableLearningScheduler: + """Build and start the scheduler that polls one storage directory. + + Both closures capture only ``directory``, never a context, so the registry + stays the sole owner of context lifetime. + + Args: + directory (str | None): Storage base directory the scheduler serves. + + Returns: + DurableLearningScheduler: The started scheduler. + """ + + def factory(org_id: str) -> RequestContext: + with _lock: + cached = _contexts.get((org_id, directory)) + return cached or RequestContext(org_id=org_id, storage_base_dir=directory) + + def discover() -> list[str]: + with _lock: + contexts = [ + ctx for (_, path), ctx in _contexts.items() if path == directory + ] + orgs = set() + for ctx in contexts: + if ctx.storage is not None: + orgs.update(ctx.storage.list_extraction_orgs()) + return sorted(orgs) + + scheduler = DurableLearningScheduler( + request_context_factory=factory, org_ids_provider=discover + ) + scheduler.start() + return scheduler + + +def _take_orphan_schedulers() -> list[DurableLearningScheduler]: + """Unregister the schedulers whose directory has no live context left. + + Must be called with ``_lock`` held. The caller stops the returned + schedulers *after* releasing the lock: ``stop()`` joins the scheduler + thread, whose discovery callback takes ``_lock`` itself, so stopping under + the lock would stall until the join timed out and leave the thread alive. + + Returns: + list[DurableLearningScheduler]: Schedulers removed from the registry, + which the caller owns and must stop. + """ + # A WeakSet loses members as they are collected, so "no members left" is + # exactly "no handle for this directory is reachable". + for empty in [d for d, ctxs in _live.items() if not len(ctxs)]: + del _live[empty] + return [_schedulers.pop(d) for d in [*_schedulers] if d not in _live] + + def ensure_local_extraction(context: RequestContext) -> None: + """Keep a durable extraction scheduler polling this context's directory. + + Registers ``context`` weakly and starts a scheduler for its directory if + one is not already running, then retires any scheduler whose directory has + no reachable context left. + + Args: + context (RequestContext): Live context whose storage directory needs + local durable extraction. + """ if context.storage is None: return directory = context.storage_base_dir @@ -19,30 +129,12 @@ def ensure_local_extraction(context: RequestContext) -> None: if _server_scheduler is not None and _server_scheduler.is_running(): return _contexts[(context.org_id, directory)] = context - if directory in _schedulers: - return - - def factory(org_id: str) -> RequestContext: - with _lock: - cached = _contexts.get((org_id, directory)) - return cached or RequestContext(org_id=org_id, storage_base_dir=directory) - - def discover() -> list[str]: - with _lock: - contexts = [ - ctx for (_, path), ctx in _contexts.items() if path == directory - ] - orgs = set() - for ctx in contexts: - if ctx.storage is not None: - orgs.update(ctx.storage.list_extraction_orgs()) - return sorted(orgs) - - scheduler = DurableLearningScheduler( - request_context_factory=factory, org_ids_provider=discover - ) - _schedulers[directory] = scheduler - scheduler.start() + _live.setdefault(directory, weakref.WeakSet()).add(context) + orphans = _take_orphan_schedulers() + if directory not in _schedulers: + _schedulers[directory] = _start_scheduler(directory) + for orphan in orphans: + orphan.stop() def adopt_server_scheduler(scheduler: DurableLearningScheduler) -> None: @@ -52,6 +144,7 @@ def adopt_server_scheduler(scheduler: DurableLearningScheduler) -> None: previous = list(_schedulers.values()) _schedulers.clear() _contexts.clear() + _live.clear() _server_scheduler = scheduler for local in previous: local.stop() diff --git a/tests/server/services/durable_learning/test_local_lifecycle.py b/tests/server/services/durable_learning/test_local_lifecycle.py new file mode 100644 index 00000000..784d6057 --- /dev/null +++ b/tests/server/services/durable_learning/test_local_lifecycle.py @@ -0,0 +1,133 @@ +"""A library-local scheduler outlives its handles for no longer than one call.""" + +import gc +import threading + +import pytest + +from reflexio.models.config_schema import ( + Config, + ProfileExtractorConfig, + StorageConfigSQLite, +) +from reflexio.server.api_endpoints.request_context import RequestContext +from reflexio.server.services.configurator.configurator import DefaultConfigurator +from reflexio.server.services.durable_learning import local + +_THREAD_NAME = "reflexio-durable-learning-scheduler" + + +def _scheduler_threads() -> int: + return sum(1 for t in threading.enumerate() if t.name == _THREAD_NAME) + + +@pytest.fixture +def isolated_registry(): + """Swap the module registry for an empty one and stop whatever it collects.""" + with local._lock: + saved_schedulers = local._schedulers + saved_contexts = local._contexts + saved_server = local._server_scheduler + local._schedulers = {} + local._contexts = type(saved_contexts)() + local._server_scheduler = None + yield + with local._lock: + created = list(local._schedulers.values()) + local._schedulers = saved_schedulers + local._contexts = saved_contexts + local._server_scheduler = saved_server + for scheduler in created: + scheduler.stop() + + +def _context(base_dir, org_id: str) -> RequestContext: + base_dir.mkdir(parents=True, exist_ok=True) + configurator = DefaultConfigurator(org_id=org_id, base_dir=str(base_dir)) + configurator.set_config( + Config( + storage_config=StorageConfigSQLite(db_path=str(base_dir / "lifecycle.db")), + window_size=1, + stride_size=1, + profile_extractor_config=ProfileExtractorConfig( + extraction_definition_prompt="Preferences" + ), + user_playbook_extractor_config=None, + ) + ) + return RequestContext( + org_id=org_id, storage_base_dir=str(base_dir), configurator=configurator + ) + + +def test_dropped_handles_do_not_accumulate_scheduler_threads( + tmp_path, isolated_registry +): + baseline = _scheduler_threads() + for index in range(6): + context = _context(tmp_path / f"dir{index}", f"lifecycle{index}") + local.ensure_local_extraction(context) + del context + gc.collect() + # Only the most recently registered directory may still hold a scheduler: + # it is retired by the next caller, never by the one that created it. + assert len(local._schedulers) == 1 + assert _scheduler_threads() <= baseline + 1 + + +def test_scheduler_survives_while_its_context_is_reachable(tmp_path, isolated_registry): + kept = _context(tmp_path / "kept", "kept-org") + local.ensure_local_extraction(kept) + kept_scheduler = local._schedulers[kept.storage_base_dir] + + transient = _context(tmp_path / "other", "other-org") + local.ensure_local_extraction(transient) + del transient + gc.collect() + + local.ensure_local_extraction(_context(tmp_path / "third", "third-org")) + gc.collect() + + assert local._schedulers.get(kept.storage_base_dir) is kept_scheduler + assert kept_scheduler.is_running() + assert str(tmp_path / "other") not in local._schedulers + + +def test_two_handles_on_one_key_both_keep_the_scheduler_alive( + tmp_path, isolated_registry +): + """A second handle on the same org+directory must not unregister the first. + + `ReflexioBase` builds an independent `RequestContext` per handle, so two + handles sharing an org and a directory are two distinct objects. A registry + keyed by `(org_id, storage_base_dir)` collapses them: the second + registration evicts the first, and collecting the second empties the key + while the first handle is still alive and using its scheduler. The sweep + then retires that scheduler and extraction stops silently under a live + caller. + + Measured on the keyed version: with both handles alive the registry held + only the second. This is the regression test for that. + """ + shared = tmp_path / "shared" + first = _context(shared, "same-org") + local.ensure_local_extraction(first) + scheduler = local._schedulers[first.storage_base_dir] + + second = _context(shared, "same-org") + assert second is not first, "precondition: the handles are distinct objects" + local.ensure_local_extraction(second) + # The first must still be represented -- the whole point of the set. + assert first in local._live[first.storage_base_dir] + + del second + gc.collect() + + # Any later caller triggers the sweep; the first handle is still alive, so + # its scheduler must survive it. + local.ensure_local_extraction(_context(tmp_path / "elsewhere", "other-org")) + gc.collect() + + assert first in local._live[first.storage_base_dir] + assert local._schedulers.get(first.storage_base_dir) is scheduler + assert scheduler.is_running() diff --git a/tests/server/services/durable_learning/test_scheduler.py b/tests/server/services/durable_learning/test_scheduler.py index d34e982a..9073180d 100644 --- a/tests/server/services/durable_learning/test_scheduler.py +++ b/tests/server/services/durable_learning/test_scheduler.py @@ -99,6 +99,10 @@ def test_library_recovers_persisted_backlog_without_new_publish(tmp_path, monkey finally: with local._lock: scheduler = local._schedulers.pop(str(tmp_path), None) - local._contexts.pop((org, str(tmp_path)), None) + # `_contexts` is a WeakSet of live contexts, not a dict keyed by + # (org, dir) -- two handles here share that key, which is exactly + # why the key was removed. Clear it; this is teardown. + local._contexts.clear() + local._live.clear() if scheduler: scheduler.stop()