diff --git a/backend/fluksio/flow/executor.py b/backend/fluksio/flow/executor.py index 89f15d0..a77baac 100644 --- a/backend/fluksio/flow/executor.py +++ b/backend/fluksio/flow/executor.py @@ -259,10 +259,19 @@ class ExecutionService: """Stop claiming, and wait for what is already claimed to finish. Called around a rebuild: items claimed against the old pipeline should - finish there rather than half-run against the new one. + finish there rather than half-run against the new one. That includes + a claim still under way: the consumer holds its lock from the claim to + the dispatch, so once that lock is free here the claim is counted, and + every later pass finds the intake closed. """ self._intake.clear() + # So a consumer waiting for capacity sees the intake closed. + self._capacity_changed() deadline = time.monotonic() + DRAIN_TIMEOUT_S + if not self._consumer_lock.acquire(timeout=DRAIN_TIMEOUT_S): + logger.warning("Rebuild did not wait out a claim") + return + self._consumer_lock.release() with self._inflight_lock: while self._running: remaining = deadline - time.monotonic() @@ -284,8 +293,8 @@ class ExecutionService: """Stop claiming, promoting, touching and reclaiming, once every claimed cascade has finished. For a rename, which rewrites the queue. - Unlike `pause_intake` this also waits out a claim already blocking - when intake clears, and a timer pass. It waits without a limit and + Unlike `pause_intake` this also waits out a timer pass, and keeps the + consumer lock until `release`. It waits without a limit and tells ``report`` the nodes whose cascades are still running each time they change. Setting ``abort`` gives up: claiming resumes and it returns False. True means held until `release`. diff --git a/backend/fluksio/flow/state.py b/backend/fluksio/flow/state.py index cf2a70f..d6fc341 100644 --- a/backend/fluksio/flow/state.py +++ b/backend/fluksio/flow/state.py @@ -834,35 +834,31 @@ class RedisState(StateBackend): watch_keys = [self._key(k) for k in expected.keys()] - try: - # Start watching the keys for changes - self._client.watch(*watch_keys) + # WATCH has to go through the pipeline: it pins one connection, and + # the client-level call is a no-op in redis-py 7. Leaving the block + # resets the pipeline, which unwatches. + with self._client.pipeline() as pipe: + try: + pipe.watch(*watch_keys) # type: ignore[no-untyped-call] + # A watching pipeline runs commands immediately. + current = cast(list[bytes | None], pipe.mget(watch_keys)) + for raw, exp_value in zip(current, expected.values(), strict=True): + if self._deserialize(raw) != exp_value: + return False - # Check current values - current = self.get_multi(list(expected.keys())) - - for key, exp_value in expected.items(): - if current.get(key) != exp_value: - self._client.unwatch() - return False - - # All values match - execute the transaction - pipe = self._client.pipeline(True) # True = use MULTI/EXEC - - for key, value in updates.items(): - data = self._serialize(value) - if self._ttl: - pipe.setex(self._key(key), self._ttl, data) - else: - pipe.set(self._key(key), data) - self._note(pipe, updates) - - pipe.execute() - return True - - except redis.WatchError: - # Another client modified one of the watched keys - return False + pipe.multi() + for key, value in updates.items(): + data = self._serialize(value) + if self._ttl: + pipe.setex(self._key(key), self._ttl, data) + else: + pipe.set(self._key(key), data) + self._note(pipe, updates) + pipe.execute() + return True + except redis.WatchError: + # Another client modified one of the watched keys + return False def close(self) -> None: """Give the connection pool back.""" diff --git a/backend/tests/flow/test_engine_hold.py b/backend/tests/flow/test_engine_hold.py index c796d4c..a815df4 100644 --- a/backend/tests/flow/test_engine_hold.py +++ b/backend/tests/flow/test_engine_hold.py @@ -125,8 +125,31 @@ class SlowClaim(MemoryWorkQueue): return items +def test_pause_intake_waits_out_a_claim_already_under_way(monkeypatch): + seen: list = [] + queue = SlowClaim() + _, service, source = _engine(seen.append, queue) + wake = service._capacity_changed + + def closing() -> None: + wake() + # The claim lands only once the intake has closed. + queue.go.set() + + monkeypatch.setattr(service, "_capacity_changed", closing) + service.start() + try: + source.inject({"reading": 1.0}) + assert queue.claiming.wait(5) + service.pause_intake() + assert seen == [1.0] + assert service.inflight == 0 + finally: + queue.go.set() + service.stop() + + def test_a_hold_waits_out_a_claim_already_under_way(): - """`pause_intake` returns here and the item runs after it.""" seen: list = [] queue = SlowClaim() _, service, source = _engine(seen.append, queue) diff --git a/backend/tests/flow/test_redis_state.py b/backend/tests/flow/test_redis_state.py index 6868c4b..ae84d53 100644 --- a/backend/tests/flow/test_redis_state.py +++ b/backend/tests/flow/test_redis_state.py @@ -38,3 +38,17 @@ def test_a_compare_and_swap_names_what_it_wrote(state): state.set("heat.a", 1) assert state.compare_and_swap_multi({"heat.a": 1}, {"heat.b": 2}) assert "heat.b" in state.message_names() + + +def test_a_compare_and_swap_loses_to_a_write_after_its_read(state, monkeypatch): + state.set("heat.a", 1) + note = RedisState._note + + def racing_note(self, pipe, keys): + # Another writer lands between the compare and EXEC. + self._client.set(self._key("heat.a"), b"2") + note(self, pipe, keys) + + monkeypatch.setattr(RedisState, "_note", racing_note) + assert not state.compare_and_swap_multi({"heat.a": 1}, {"heat.a": 3}) + assert state.get("heat.a") == 2