From dd34842e74c0560e59d0c82d7a80b2abda4084f7 Mon Sep 17 00:00:00 2001 From: Tristan Simas Date: Mon, 12 Oct 2026 00:25:54 -0400 Subject: [PATCH] Keep atomic and deferral block state per context; name the navigation payload atomic, atomic_success and defer_live_invalidations kept their nesting in class attributes shared by the process, while config_context keeps its scope in context variables. The open block is now an AtomicBlock held in a ContextVar and the deferral nesting is a ContextVar count, so a thread outside any block no longer sees another thread's block (#30). TimeTravelChangeSet.legacy_entries is the only completion-callback payload, so it is renamed states_needing_navigation (#34 item 3). Co-Authored-By: Claude Opus 5.5 --- src/objectstate/object_state_registry.py | 179 +++++++++++------- tests/test_atomic_block_context.py | 139 ++++++++++++++ tests/test_callable_baseline.py | 3 - tests/test_object_state_metadata_contracts.py | 3 - tests/test_object_state_update_instance.py | 3 - tests/test_undo_redo_cascade.py | 7 +- 6 files changed, 248 insertions(+), 86 deletions(-) create mode 100644 tests/test_atomic_block_context.py diff --git a/src/objectstate/object_state_registry.py b/src/objectstate/object_state_registry.py index 3c765e0..8620fc5 100644 --- a/src/objectstate/object_state_registry.py +++ b/src/objectstate/object_state_registry.py @@ -4,6 +4,7 @@ """ from contextlib import contextmanager +from contextvars import ContextVar, Token from dataclasses import dataclass, is_dataclass, replace import logging from typing import ( @@ -136,8 +137,10 @@ class TimeTravelChangeSet: triggering_scope: Optional[str] scope_changes: Dict[str, TimeTravelScopeChange] - def legacy_entries(self, states: Dict[str, 'ObjectState']) -> List[TimeTravelStateEntry]: - """Project to the existing callback payload shape.""" + def states_needing_navigation( + self, states: dict[str, 'ObjectState'] + ) -> list[TimeTravelStateEntry]: + """Return the registered states this restore changed, as completion callbacks take them.""" return [ (scope_key, states[scope_key]) for scope_key, change in self.scope_changes.items() @@ -156,6 +159,35 @@ class DeferredFieldInvalidation: include_origin: bool = True +@dataclass +class AtomicBlock: + """One open ``atomic`` / ``atomic_success`` block: the undo step being coalesced. + + The block belongs to the context that opened it, exactly as a + ``config_context`` layer does: a new thread starts outside every block, and + an asyncio task started inside a block inherits it. + """ + + label: str + triggering_scope: str | None + entry_signature: RegistryMutationSignature + depth: int = 1 + snapshot_requested: bool = False + + +# Which block, and how many deferrals, the running context is inside. These +# answer "whose block is this", so they live where ``config_context`` keeps its +# scope (``contextvars``, see ``context_manager``). What the blocks act on, the +# registered states, the history and the queue of owed invalidations, is one +# per process and stays on ``ObjectStateRegistry``. +_atomic_block: ContextVar[AtomicBlock | None] = ContextVar( + "objectstate_atomic_block", default=None +) +_live_invalidation_deferrals: ContextVar[int] = ContextVar( + "objectstate_live_invalidation_deferrals", default=0 +) + + class ObjectStateRegistry: """Singleton registry of all ObjectState instances. @@ -167,7 +199,16 @@ class ObjectStateRegistry: - ImageBrowser: registers when opened, unregisters when closed - Config window: registers PipelineConfig/GlobalPipelineConfig - Thread safety: Not thread-safe (all operations expected on main thread). + One registry per process, by design: the registered states, the snapshot + history, the callbacks and the queue of owed live-cache invalidations are + class attributes, so two independent registries cannot exist. + + Per context (thread or asyncio task), like ``config_context``: the open + ``atomic`` / ``atomic_success`` block and the ``defer_live_invalidations`` + nesting. A thread outside any block never sees another thread's block. + + Thread safety: the process-wide stores above are not synchronized; callers + that mutate the registry from several threads must serialize those calls. """ _states: Dict[str, 'ObjectState'] = {} # Keys are always strings (None normalized to "") @@ -175,7 +216,6 @@ class ObjectStateRegistry: # Callbacks receive (scope_id: str, object_state: ObjectState) _on_register_callbacks: List[Callable[[str, 'ObjectState'], None]] = [] _on_unregister_callbacks: List[Callable[[str, 'ObjectState'], None]] = [] - _deferred_invalidation_depth: int = 0 _deferred_invalidations: Set[DeferredFieldInvalidation] = set() _snapshot_dirty_scopes: Set[str] = set() @@ -290,11 +330,11 @@ def _fire_history_changed_callbacks(cls) -> None: @contextmanager def defer_live_invalidations(cls) -> Generator[None, None, None]: """Queue descendant live-cache invalidations for an explicit later flush.""" - cls._deferred_invalidation_depth += 1 + token = _live_invalidation_deferrals.set(_live_invalidation_deferrals.get() + 1) try: yield finally: - cls._deferred_invalidation_depth -= 1 + _live_invalidation_deferrals.reset(token) @classmethod def has_deferred_invalidations(cls) -> bool: @@ -309,8 +349,7 @@ def flush_deferred_invalidations(cls) -> None: requests = tuple(cls._deferred_invalidations) cls._deferred_invalidations.clear() - previous_depth = cls._deferred_invalidation_depth - cls._deferred_invalidation_depth = 0 + token = _live_invalidation_deferrals.set(0) try: for request in requests: cls._invalidate_by_type_and_scope_now( @@ -321,7 +360,7 @@ def flush_deferred_invalidations(cls) -> None: include_origin=request.include_origin, ) finally: - cls._deferred_invalidation_depth = previous_depth + _live_invalidation_deferrals.reset(token) @classmethod def _fire_callbacks(cls, callbacks: List[Callable], role: str, *args: Any) -> None: @@ -498,17 +537,13 @@ def clear(cls) -> None: cls._time_travel_limbo.clear() cls._graveyard.clear() cls._deferred_invalidations.clear() - cls._deferred_invalidation_depth = 0 + _live_invalidation_deferrals.set(0) cls._snapshot_dirty_scopes.clear() cls._snapshots.clear() cls._timelines.clear() cls._current_head = None cls._current_timeline = cls.MAIN_BRANCH - cls._atomic_depth = 0 - cls._atomic_label = None - cls._atomic_triggering_scope = None - cls._atomic_snapshot_requested = False - cls._atomic_entry_signature = None + _atomic_block.set(None) logger.debug("Cleared ObjectState registry, history, limbo, and graveyard") # ========== TOKEN MANAGEMENT AND CHANGE NOTIFICATION ========== @@ -735,7 +770,7 @@ def invalidate_by_type_and_scope( well as its descendants. """ changed_scope = cls._normalize_scope_id(scope_id) - if cls._deferred_invalidation_depth > 0: + if _live_invalidation_deferrals.get() > 0: cls._deferred_invalidations.add( DeferredFieldInvalidation( scope_id=changed_scope, @@ -877,13 +912,6 @@ def _invalidate_field_in_matching_states( _max_history_size: int = 1000 # Max snapshots in DAG (increased since we don't truncate branches) _in_time_travel: bool = False # Flag for PFM to know to refresh widget values - # Atomic operation state - when >0, snapshots are deferred until operation completes - _atomic_depth: int = 0 - _atomic_label: Optional[str] = None # Label for the coalesced snapshot - _atomic_triggering_scope: Optional[str] = None - _atomic_snapshot_requested: bool = False - _atomic_entry_signature: Optional[RegistryMutationSignature] = None - # Limbo: ObjectStates temporarily removed during time-travel # When traveling to a snapshot, ObjectStates not in that snapshot are moved here. # When traveling forward or to head, they're restored from here. @@ -926,28 +954,12 @@ def atomic(cls, label: str, scope_id: Optional[str] = None) -> Generator[None, N the first scoped snapshot request inside the atomic block becomes the owner. """ - cls._atomic_depth += 1 - if cls._atomic_depth == 1: - cls._atomic_label = label - cls._atomic_triggering_scope = scope_id - cls._atomic_snapshot_requested = False - cls._atomic_entry_signature = cls._registry_mutation_signature() - + block, token = cls._enter_atomic_block(label, scope_id) try: yield finally: - cls._atomic_depth -= 1 - if cls._atomic_depth == 0: - # Outermost atomic block - record the coalesced snapshot - final_label = cls._atomic_label or label - triggering_scope = cls._atomic_triggering_scope - should_record = cls._atomic_should_record() - cls._atomic_label = None - cls._atomic_triggering_scope = None - cls._atomic_snapshot_requested = False - cls._atomic_entry_signature = None - if should_record: - cls.record_snapshot(final_label, triggering_scope) + if cls._leave_atomic_block(block, token) and cls._atomic_should_record(block): + cls.record_snapshot(block.label, block.triggering_scope) @classmethod @contextmanager @@ -963,7 +975,7 @@ def atomic_success( raises, the outermost success-only operation unwinds without recording a partial snapshot. """ - outermost = cls._atomic_depth == 0 + outermost = cls.current_atomic_block() is None snapshots_before = set(cls._snapshots) if outermost else None timelines_before = copy.deepcopy(cls._timelines) if outermost else None current_head_before = cls._current_head if outermost else None @@ -973,30 +985,16 @@ def atomic_success( set(cls._snapshot_dirty_scopes) if outermost else None ) - cls._atomic_depth += 1 - if cls._atomic_depth == 1: - cls._atomic_label = label - cls._atomic_triggering_scope = scope_id - cls._atomic_snapshot_requested = False - cls._atomic_entry_signature = cls._registry_mutation_signature() - + block, token = cls._enter_atomic_block(label, scope_id) succeeded = False try: yield succeeded = True finally: - cls._atomic_depth -= 1 - if cls._atomic_depth == 0: - final_label = cls._atomic_label or label - triggering_scope = cls._atomic_triggering_scope - should_record = cls._atomic_should_record() - cls._atomic_label = None - cls._atomic_triggering_scope = None - cls._atomic_snapshot_requested = False - cls._atomic_entry_signature = None + if cls._leave_atomic_block(block, token): if succeeded: - if should_record: - cls.record_snapshot(final_label, triggering_scope) + if cls._atomic_should_record(block): + cls.record_snapshot(block.label, block.triggering_scope) elif snapshots_before is not None and timelines_before is not None: for snapshot_id in set(cls._snapshots) - snapshots_before: del cls._snapshots[snapshot_id] @@ -1007,6 +1005,44 @@ def atomic_success( if snapshot_dirty_scopes_before is not None: cls._snapshot_dirty_scopes = snapshot_dirty_scopes_before + @classmethod + def current_atomic_block(cls) -> AtomicBlock | None: + """Return the block the running context is inside, or None outside any block.""" + block = _atomic_block.get() + if block is None or block.depth == 0: + return None + return block + + @classmethod + def _enter_atomic_block( + cls, label: str, scope_id: str | None + ) -> tuple[AtomicBlock, Token[AtomicBlock | None] | None]: + """Join the running context's block, opening one when it is outside any.""" + block = cls.current_atomic_block() + if block is not None: + block.depth += 1 + return block, None + block = AtomicBlock( + label=label, + triggering_scope=scope_id, + entry_signature=cls._registry_mutation_signature(), + ) + return block, _atomic_block.set(block) + + @classmethod + def _leave_atomic_block( + cls, block: AtomicBlock, token: Token[AtomicBlock | None] | None + ) -> bool: + """Leave ``block``; return whether that closed it, so its step is now due. + + A block dropped by ``clear()`` while open is not due: its states are gone. + """ + block.depth -= 1 + still_current = _atomic_block.get() is block + if token is not None: + _atomic_block.reset(token) + return block.depth == 0 and still_current + @classmethod def _registry_mutation_signature( cls, @@ -1027,12 +1063,12 @@ def _registry_mutation_signature( return tuple(entries) @classmethod - def _atomic_should_record(cls) -> bool: - """Return whether the outer atomic scope produced a user-visible state change.""" + def _atomic_should_record(cls, block: AtomicBlock) -> bool: + """Return whether the closed block produced a user-visible state change.""" - if cls._atomic_snapshot_requested: + if block.snapshot_requested: return True - return cls._registry_mutation_signature() != cls._atomic_entry_signature + return cls._registry_mutation_signature() != block.entry_signature @classmethod def ensure_baseline_snapshot(cls, label: str = "init") -> None: @@ -1120,11 +1156,12 @@ def record_snapshot(cls, label: str = "", scope_id: Optional[str] = None) -> Non # ATOMIC OPERATIONS: Defer snapshot until atomic block exits # The atomic() context manager will call record_snapshot() when it completes - if cls._atomic_depth > 0: - cls._atomic_snapshot_requested = True - if scope_id is not None and cls._atomic_triggering_scope is None: - cls._atomic_triggering_scope = scope_id - logger.debug(f"⏱️ ATOMIC: Deferring snapshot '{label}' (depth={cls._atomic_depth})") + block = cls.current_atomic_block() + if block is not None: + block.snapshot_requested = True + if scope_id is not None and block.triggering_scope is None: + block.triggering_scope = scope_id + logger.debug(f"⏱️ ATOMIC: Deferring snapshot '{label}' (depth={block.depth})") return import time @@ -1886,7 +1923,7 @@ def _fire_time_travel_complete_callbacks( if not cls._on_time_travel_complete_callbacks: return - states_with_changes = change_set.legacy_entries(cls._states) + states_with_changes = change_set.states_needing_navigation(cls._states) logger.debug( "⏱️ TIME_TRAVEL: Firing %d callback(s) with %d changed state(s)", len(cls._on_time_travel_complete_callbacks), diff --git a/tests/test_atomic_block_context.py b/tests/test_atomic_block_context.py new file mode 100644 index 0000000..9777811 --- /dev/null +++ b/tests/test_atomic_block_context.py @@ -0,0 +1,139 @@ +"""Block state belongs to the context that opened the block (issue #30).""" + +import asyncio +import threading +from dataclasses import dataclass + +import pytest + +from objectstate import ObjectState +from objectstate.object_state_registry import ObjectStateRegistry + + +@dataclass +class Dummy: + x: int = 1 + + +@pytest.fixture(autouse=True) +def empty_registry(): + ObjectStateRegistry.clear() + yield + ObjectStateRegistry.clear() + + +class HeldOpen: + """Run a body on another thread and hold it inside until released.""" + + def __init__(self, body) -> None: + self._inside = threading.Event() + self._release = threading.Event() + self._errors: list[BaseException] = [] + self._thread = threading.Thread(target=self._run, args=(body,)) + + def _run(self, body) -> None: + try: + with body(): + self._inside.set() + self._release.wait(timeout=10) + except BaseException as error: # reported to the test thread + self._errors.append(error) + self._inside.set() + + def __enter__(self) -> "HeldOpen": + self._thread.start() + assert self._inside.wait(timeout=10) + return self + + def __exit__(self, *exc_info) -> None: + self._release.set() + self._thread.join(timeout=10) + assert not self._errors, self._errors + + +def _block_facts() -> tuple[int, str | None]: + block = ObjectStateRegistry.current_atomic_block() + return (0, None) if block is None else (block.depth, block.label) + + +def test_thread_outside_any_block_does_not_see_another_threads_block() -> None: + """The reproducer of issue #30.""" + with HeldOpen(lambda: ObjectStateRegistry.atomic("worker edit")): + assert _block_facts() == (0, None) + with ObjectStateRegistry.atomic("main edit"): + assert _block_facts() == (1, "main edit") + assert _block_facts() == (0, None) + + +def test_snapshot_is_not_absorbed_into_another_threads_block() -> None: + state = ObjectState(Dummy(), scope_id="/plate::step_0") + ObjectStateRegistry.register(state, _skip_snapshot=True) + + with HeldOpen(lambda: ObjectStateRegistry.atomic("worker edit")): + ObjectStateRegistry.record_snapshot("main snapshot") + assert [s.label for s in ObjectStateRegistry.get_branch_history()] == ["main snapshot"] + + with ObjectStateRegistry.atomic("main edit", scope_id="/plate::step_0"): + state.update_parameter("x", 2) + labels = [s.label for s in ObjectStateRegistry.get_branch_history()] + assert "main edit [/plate::step_0]" in labels + + +def test_atomic_success_in_another_thread_is_outermost_for_its_own_context() -> None: + state = ObjectState(Dummy(), scope_id="/plate::step_0") + ObjectStateRegistry.register(state, _skip_snapshot=True) + + with HeldOpen(lambda: ObjectStateRegistry.atomic_success("worker edit")): + with ObjectStateRegistry.atomic_success("main edit", scope_id="/plate::step_0"): + assert _block_facts() == (1, "main edit") + state.update_parameter("x", 2) + labels = [s.label for s in ObjectStateRegistry.get_branch_history()] + assert "main edit [/plate::step_0]" in labels + + +def test_nested_blocks_in_one_context_share_the_outer_block() -> None: + with ObjectStateRegistry.atomic("outer"): + with ObjectStateRegistry.atomic_success("inner"): + assert _block_facts() == (2, "outer") + assert _block_facts() == (1, "outer") + assert _block_facts() == (0, None) + + +def test_block_follows_the_same_owner_as_config_context_across_tasks() -> None: + """A task started inside a block inherits it; a sibling task does not.""" + + async def facts() -> tuple[int, str | None]: + return _block_facts() + + async def scenario() -> tuple[tuple[int, str | None], tuple[int, str | None]]: + sibling = asyncio.ensure_future(facts()) + with ObjectStateRegistry.atomic("task edit"): + child = asyncio.ensure_future(facts()) + inherited = await child + return inherited, await sibling + + inherited, sibling = asyncio.run(scenario()) + assert inherited == (1, "task edit") + assert sibling == (0, None) + + +def test_deferral_in_one_thread_does_not_queue_another_threads_invalidation() -> None: + state = ObjectState(Dummy(), scope_id="/plate::step_0") + ObjectStateRegistry.register(state, _skip_snapshot=True) + + def invalidate() -> None: + ObjectStateRegistry.invalidate_by_type_and_scope( + scope_id="/plate::step_0", + changed_type=Dummy, + field_name="x", + ) + + with HeldOpen(ObjectStateRegistry.defer_live_invalidations): + invalidate() + assert not ObjectStateRegistry.has_deferred_invalidations() + + with ObjectStateRegistry.defer_live_invalidations(): + invalidate() + assert ObjectStateRegistry.has_deferred_invalidations() + ObjectStateRegistry.flush_deferred_invalidations() + assert not ObjectStateRegistry.has_deferred_invalidations() diff --git a/tests/test_callable_baseline.py b/tests/test_callable_baseline.py index 67b930e..206e164 100644 --- a/tests/test_callable_baseline.py +++ b/tests/test_callable_baseline.py @@ -21,9 +21,6 @@ def _reset_registry_and_history() -> None: ObjectStateRegistry._current_timeline = "main" ObjectStateRegistry._current_head = None ObjectStateRegistry._in_time_travel = False - ObjectStateRegistry._atomic_depth = 0 - ObjectStateRegistry._atomic_label = None - ObjectStateRegistry._atomic_triggering_scope = None class RebuiltCallable: diff --git a/tests/test_object_state_metadata_contracts.py b/tests/test_object_state_metadata_contracts.py index f983377..375ed1e 100644 --- a/tests/test_object_state_metadata_contracts.py +++ b/tests/test_object_state_metadata_contracts.py @@ -21,9 +21,6 @@ def _reset_registry_and_history() -> None: ObjectStateRegistry._current_timeline = "main" ObjectStateRegistry._current_head = None ObjectStateRegistry._in_time_travel = False - ObjectStateRegistry._atomic_depth = 0 - ObjectStateRegistry._atomic_label = None - ObjectStateRegistry._atomic_triggering_scope = None @dataclass diff --git a/tests/test_object_state_update_instance.py b/tests/test_object_state_update_instance.py index 6060d25..0eab449 100644 --- a/tests/test_object_state_update_instance.py +++ b/tests/test_object_state_update_instance.py @@ -28,9 +28,6 @@ def _reset_registry() -> None: ObjectStateRegistry._current_timeline = "main" ObjectStateRegistry._current_head = None ObjectStateRegistry._in_time_travel = False - ObjectStateRegistry._atomic_depth = 0 - ObjectStateRegistry._atomic_label = None - ObjectStateRegistry._atomic_triggering_scope = None ObjectStateRegistry._token = 0 ObjectStateRegistry._resolved_changed_callbacks.clear() diff --git a/tests/test_undo_redo_cascade.py b/tests/test_undo_redo_cascade.py index f980a60..6d052fb 100644 --- a/tests/test_undo_redo_cascade.py +++ b/tests/test_undo_redo_cascade.py @@ -25,11 +25,6 @@ def _reset_registry_and_history() -> None: ObjectStateRegistry._current_timeline = "main" ObjectStateRegistry._current_head = None ObjectStateRegistry._in_time_travel = False - ObjectStateRegistry._atomic_depth = 0 - ObjectStateRegistry._atomic_label = None - ObjectStateRegistry._atomic_triggering_scope = None - ObjectStateRegistry._atomic_snapshot_requested = False - ObjectStateRegistry._atomic_entry_signature = None @dataclass @@ -236,7 +231,7 @@ def test_time_travel_callback_entries_include_nested_resolved_path_changes() -> }, ) - assert change_set.legacy_entries(ObjectStateRegistry._states) == [ + assert change_set.states_needing_navigation(ObjectStateRegistry._states) == [ (step_scope, step_state) ]