From 41d411a4001950555c5b619a0a043a99d0b895eb Mon Sep 17 00:00:00 2001 From: Kyriakos Psarakis Date: Sat, 24 Oct 2026 14:42:20 +0200 Subject: [PATCH 01/20] Add fault tolerance support during state migration (squashed PR #35) --- .github/workflows/ci.yml | 26 +- .github/workflows/e2e.yml | 99 +++- .github/workflows/publish-helm-chart.yml | 4 +- .github/workflows/publish-images.yml | 10 +- .github/workflows/ruff.yml | 4 +- coordinator/coordinator_metadata.py | 124 ++++- coordinator/coordinator_service.py | 329 ++++++++++- coordinator/worker_pool.py | 9 +- demo/demo-migration-ft-ycsb/__init__.py | 0 .../calculate_metrics.py | 1 + demo/demo-migration-ft-ycsb/client.py | 236 ++++++++ .../kafka_output_consumer.py | 1 + demo/demo-migration-ft-ycsb/ycsb.py | 1 + docker-compose.yml | 2 + pytest.ini | 1 + styx-package/styx/common/logging.py | 5 +- styx-package/styx/common/message_types.py | 1 + tests/e2e/test_e2e_migration_ft_ycsb.py | 435 +++++++++++++++ tests/e2e/test_e2e_migration_tpcc.py | 2 +- tests/e2e/test_e2e_migration_ycsb.py | 2 +- tests/integration/test_s3_snapshots.py | 8 +- .../coordinator/test_coordinator_metadata.py | 190 +++++++ .../test_migration_fault_tolerance.py | 527 ++++++++++++++++++ tests/unit/coordinator/test_worker_pool.py | 17 + tests/unit/styx_package/test_logging.py | 38 ++ tests/unit/styx_package/test_message_types.py | 5 +- tests/unit/worker/test_async_snapshots.py | 163 +++++- tests/unit/worker/test_kafka_batch_egress.py | 31 +- .../unit/worker/test_kafka_egress_coverage.py | 196 ++++++- tests/unit/worker/test_kafka_ingress.py | 263 ++++++++- .../worker/test_migration_fault_tolerance.py | 210 +++++++ worker/async_snapshotting.py | 18 + worker/egress/styx_kafka_batch_egress.py | 149 +++-- worker/fault_tolerance/async_snapshots.py | 49 +- worker/operator_state/aria/in_memory_state.py | 6 + worker/transactional_protocols/aria.py | 66 ++- worker/worker_service.py | 146 ++++- 37 files changed, 3191 insertions(+), 183 deletions(-) create mode 100644 demo/demo-migration-ft-ycsb/__init__.py create mode 120000 demo/demo-migration-ft-ycsb/calculate_metrics.py create mode 100644 demo/demo-migration-ft-ycsb/client.py create mode 120000 demo/demo-migration-ft-ycsb/kafka_output_consumer.py create mode 120000 demo/demo-migration-ft-ycsb/ycsb.py create mode 100644 tests/e2e/test_e2e_migration_ft_ycsb.py create mode 100644 tests/unit/coordinator/test_migration_fault_tolerance.py create mode 100644 tests/unit/styx_package/test_logging.py create mode 100644 tests/unit/worker/test_migration_fault_tolerance.py diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 15690d63..404abd08 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -17,10 +17,10 @@ jobs: steps: - name: Checkout - uses: actions/checkout@v4 + uses: actions/checkout@v6.0.2 - name: Set up Python - uses: actions/setup-python@v5 + uses: actions/setup-python@v6.2.0 with: python-version: "3.14" cache: "pip" @@ -37,7 +37,7 @@ jobs: run: pytest tests/unit/coordinator/ -v --cov=coordinator --cov-report=xml:coverage-coordinator.xml - name: Upload coordinator coverage to Codecov - uses: codecov/codecov-action@v5 + uses: codecov/codecov-action@v5.5.3 with: files: coverage-coordinator.xml flags: coordinator @@ -51,10 +51,10 @@ jobs: steps: - name: Checkout - uses: actions/checkout@v4 + uses: actions/checkout@v6.0.2 - name: Set up Python - uses: actions/setup-python@v5 + uses: actions/setup-python@v6.2.0 with: python-version: "3.14" cache: "pip" @@ -64,7 +64,7 @@ jobs: python -m pip install --upgrade pip setuptools Cython cd styx-package && python setup.py build_ext --inplace && cd .. python -m pip install styx-package/ - python -m pip install msgspec pytest==9.1.1 pytest-cov pytest-asyncio psutil + python -m pip install msgspec pytest==9.1.1 pytest-cov pytest-asyncio psutil uvloop - name: Build Cython extensions run: python worker/setup.py build_ext --inplace @@ -73,7 +73,7 @@ jobs: run: pytest tests/unit/worker/ -v --cov=worker --cov-report=xml:coverage-worker.xml - name: Upload worker coverage to Codecov - uses: codecov/codecov-action@v5 + uses: codecov/codecov-action@v5.5.3 with: files: coverage-worker.xml flags: worker @@ -87,10 +87,10 @@ jobs: steps: - name: Checkout - uses: actions/checkout@v4 + uses: actions/checkout@v6.0.2 - name: Set up Python - uses: actions/setup-python@v5 + uses: actions/setup-python@v6.2.0 with: python-version: "3.14" cache: "pip" @@ -106,7 +106,7 @@ jobs: run: pytest tests/unit/styx_package/ -v --cov=styx-package/styx --cov-report=xml:coverage-styx-package.xml - name: Upload styx-package coverage to Codecov - uses: codecov/codecov-action@v5 + uses: codecov/codecov-action@v5.5.3 with: files: coverage-styx-package.xml flags: styx-package @@ -170,10 +170,10 @@ jobs: steps: - name: Checkout - uses: actions/checkout@v4 + uses: actions/checkout@v6.0.2 - name: Set up Python - uses: actions/setup-python@v5 + uses: actions/setup-python@v6.2.0 with: python-version: "3.14" cache: "pip" @@ -190,7 +190,7 @@ jobs: run: pytest tests/integration/ -v --cov=styx-package/styx --cov=coordinator --cov=worker --cov-report=xml:coverage-integration.xml - name: Upload integration coverage to Codecov - uses: codecov/codecov-action@v5 + uses: codecov/codecov-action@v5.5.3 with: files: coverage-integration.xml flags: integration diff --git a/.github/workflows/e2e.yml b/.github/workflows/e2e.yml index 61e843ba..54d4859d 100644 --- a/.github/workflows/e2e.yml +++ b/.github/workflows/e2e.yml @@ -10,13 +10,15 @@ concurrency: cancel-in-progress: true jobs: - e2e-tests: + e2e-base: + name: "E2E — base (non-migration)" + # Provides the legacy "e2e-tests" status check via the summary job below. runs-on: ubuntu-latest - timeout-minutes: 45 + timeout-minutes: 60 steps: - name: Checkout - uses: actions/checkout@v4 + uses: actions/checkout@v6.0.2 - name: Docker info run: | @@ -24,7 +26,7 @@ jobs: docker compose version - name: Set up Python - uses: actions/setup-python@v5 + uses: actions/setup-python@v6.2.0 with: python-version: "3.14" cache: "pip" @@ -34,37 +36,100 @@ jobs: python -m pip install --upgrade pip setuptools Cython cd styx-package && python setup.py build_ext --inplace && cd .. python -m pip install -r requirements.txt - # Install Styx package from repo root python -m pip install -e styx-package/ - name: Make scripts executable run: | chmod +x scripts/start_styx_cluster.sh scripts/stop_styx_cluster.sh || true - # Needed so Kafka advertises a reachable external listener - name: Set DOCKER_HOST_IP run: echo "DOCKER_HOST_IP=127.0.0.1" >> $GITHUB_ENV - - name: Run e2e tests + - name: Run non-migration e2e tests env: PYTHONUNBUFFERED: "1" run: | - pytest tests/e2e/ -m e2e -q + pytest tests/e2e/ -m "e2e" -q --log-level=DEBUG - - name: Collect worker and coordinator logs (always) + - name: Collect logs if: always() run: | mkdir -p artifacts - if [ -d logs ]; then - cp -r logs artifacts/ - else - echo "No logs folder found to copy." - fi + if [ -d logs ]; then cp -r logs artifacts/; fi + + - name: Upload artifacts + if: always() + uses: actions/upload-artifact@v7.0.0 + with: + name: e2e-base-artifacts + path: artifacts + if-no-files-found: ignore + + e2e-migration: + name: "E2E — migration" + runs-on: ubuntu-latest + timeout-minutes: 90 + + steps: + - name: Checkout + uses: actions/checkout@v6.0.2 + + - name: Docker info + run: | + docker version + docker compose version + + - name: Set up Python + uses: actions/setup-python@v6.2.0 + with: + python-version: "3.14" + cache: "pip" + + - name: Install dependencies + run: | + python -m pip install --upgrade pip setuptools Cython + cd styx-package && python setup.py build_ext --inplace && cd .. + python -m pip install -r requirements.txt + python -m pip install -e styx-package/ - - name: Upload artifacts (always) + - name: Make scripts executable + run: | + chmod +x scripts/start_styx_cluster.sh scripts/stop_styx_cluster.sh || true + + - name: Set DOCKER_HOST_IP + run: echo "DOCKER_HOST_IP=127.0.0.1" >> $GITHUB_ENV + + - name: Run migration e2e tests + env: + PYTHONUNBUFFERED: "1" + run: | + pytest tests/e2e/ -m "e2e_migration" -q --log-level=DEBUG + + - name: Collect logs + if: always() + run: | + mkdir -p artifacts + if [ -d logs ]; then cp -r logs artifacts/; fi + + - name: Upload artifacts if: always() - uses: actions/upload-artifact@v4 + uses: actions/upload-artifact@v7.0.0 with: - name: e2e-artifacts + name: e2e-migration-artifacts path: artifacts if-no-files-found: ignore + + # Legacy summary job — provides the "e2e-tests" status check that branch + # protection rules may reference from before the split into two jobs. + e2e-tests: + name: "e2e-tests" + if: always() + needs: [e2e-base, e2e-migration] + runs-on: ubuntu-latest + steps: + - name: Check results + run: | + if [[ "${{ needs.e2e-base.result }}" == "failure" || "${{ needs.e2e-migration.result }}" == "failure" ]]; then + echo "One or more E2E jobs failed" + exit 1 + fi diff --git a/.github/workflows/publish-helm-chart.yml b/.github/workflows/publish-helm-chart.yml index 7380443b..4b23d47e 100644 --- a/.github/workflows/publish-helm-chart.yml +++ b/.github/workflows/publish-helm-chart.yml @@ -15,10 +15,10 @@ jobs: runs-on: ubuntu-latest steps: - name: Checkout - uses: actions/checkout@v4 + uses: actions/checkout@v6.0.2 - name: Set up Helm - uses: azure/setup-helm@v4 + uses: azure/setup-helm@v5.0.0 - name: Log in to GHCR run: echo "${{ secrets.GITHUB_TOKEN }}" | helm registry login ghcr.io --username ${{ github.repository_owner }} --password-stdin diff --git a/.github/workflows/publish-images.yml b/.github/workflows/publish-images.yml index c5d94d7b..82cef8b2 100644 --- a/.github/workflows/publish-images.yml +++ b/.github/workflows/publish-images.yml @@ -25,13 +25,13 @@ jobs: image_suffix: worker steps: - name: Checkout - uses: actions/checkout@v4 + uses: actions/checkout@v6.0.2 - name: Set up Build - uses: docker/setup-buildx-action@v3 + uses: docker/setup-buildx-action@v4.0.0 - name: Log in to GHCR - uses: docker/login-action@v3 + uses: docker/login-action@v4.0.0 with: registry: ghcr.io username: ${{ github.repository_owner }} @@ -39,7 +39,7 @@ jobs: - name: Extract metadata id: meta - uses: docker/metadata-action@v5 + uses: docker/metadata-action@v6.0.0 with: images: ghcr.io/${{ github.repository_owner }}/${{ github.event.repository.name }}-${{ matrix.image_suffix }} tags: | @@ -49,7 +49,7 @@ jobs: type=raw,value=latest,enable=${{ github.ref == format('refs/heads/{0}', github.event.repository.default_branch) }} - name: Build and push - uses: docker/build-push-action@v6 + uses: docker/build-push-action@v7.0.0 with: context: . file: ${{ matrix.dockerfile }} diff --git a/.github/workflows/ruff.yml b/.github/workflows/ruff.yml index 01fb7d90..92655866 100644 --- a/.github/workflows/ruff.yml +++ b/.github/workflows/ruff.yml @@ -16,10 +16,10 @@ jobs: steps: - name: Checkout - uses: actions/checkout@v4 + uses: actions/checkout@v6.0.2 - name: Set up Python - uses: actions/setup-python@v5 + uses: actions/setup-python@v6.2.0 with: python-version: "3.14" cache: "pip" diff --git a/coordinator/coordinator_metadata.py b/coordinator/coordinator_metadata.py index f0ed5c37..71bf1f82 100644 --- a/coordinator/coordinator_metadata.py +++ b/coordinator/coordinator_metadata.py @@ -58,6 +58,19 @@ def __init__(self, networking: NetworkingManager, s3_client: S3Client) -> None: self.kafka_metadata_producer: AIOKafkaProducer | None = None self.max_operator_parallelism: int | None = None + # Migration fault tolerance: track pre-migration snapshot + self.pre_migration_snapshot_id: int = -1 + self.pre_migration_snapshot_pending: bool = False + # Deferred graph finalization: wait for post-migration snapshot before committing + self.post_migration_snapshot_pending: bool = False + # Deferred graph update: held until migration completes + self._pending_graph: StateflowGraph | None = None + + # Checkpoint-and-resume migration FT + self._migration_checkpoint_blob: bytes | None = None + self._migration_checkpoint_baseline_snap_id: int = -1 + self._migration_checkpoint_snapshot_complete: asyncio.Event = asyncio.Event() + async def start_kafka_metadata_producer(self) -> None: self.kafka_metadata_producer = AIOKafkaProducer( bootstrap_servers=[KAFKA_URL], @@ -113,6 +126,9 @@ async def send_recovery_to_participating_workers(self) -> None: worker_assignments = self.worker_pool.get_worker_assignments() participating_workers: list[Worker] = self.worker_pool.get_participating_workers() self.worker_is_healthy = {worker.worker_id: asyncio.Event() for worker in participating_workers} + logging.warning( + f"[RECOVERY] Sending InitRecovery to {len(participating_workers)} workers with snap_id={snap_id}", + ) async with asyncio.TaskGroup() as tg: for worker in participating_workers: tg.create_task( @@ -194,6 +210,35 @@ def register_snapshot( self.completed_epoch_counter = epoch_counter self.completed_t_counter = t_counter current_completed_snapshot: int = self.get_current_completed_snapshot_id() + logging.warning( + f"register_snapshot | worker={worker_id} snap={snapshot_id} " + f"cluster_min={current_completed_snapshot} prev={self.prev_completed_snapshot_id} " + f"baseline={self._migration_checkpoint_baseline_snap_id} " + f"blob_set={self._migration_checkpoint_blob is not None}", + ) + # Track pre-migration snapshot for recovery + if self.pre_migration_snapshot_pending and current_completed_snapshot != self.prev_completed_snapshot_id: + self.pre_migration_snapshot_id = current_completed_snapshot + self.pre_migration_snapshot_pending = False + logging.warning(f"Pre-migration snapshot recorded: {self.pre_migration_snapshot_id}") + # Finalize graph update after first post-migration snapshot completes + if self.post_migration_snapshot_pending and current_completed_snapshot != self.prev_completed_snapshot_id: + self.post_migration_snapshot_pending = False + self.finalize_graph_update() + logging.warning( + f"Post-migration snapshot {current_completed_snapshot} completed — graph finalized", + ) + # Detect migration checkpoint snapshot completion via baseline ID comparison + if ( + self._migration_checkpoint_baseline_snap_id >= 0 + and current_completed_snapshot > self._migration_checkpoint_baseline_snap_id + ): + self._migration_checkpoint_snapshot_complete.set() + logging.warning( + f"Migration checkpoint snapshot {current_completed_snapshot} completed " + f"(baseline was {self._migration_checkpoint_baseline_snap_id})", + ) + self._migration_checkpoint_baseline_snap_id = -1 if current_completed_snapshot != self.prev_completed_snapshot_id: logging.warning(f"Cluster completed snapshot: {current_completed_snapshot}") # if we reached a complete snapshot, we could compact its deltas with the previous one @@ -203,6 +248,7 @@ def register_snapshot( self.completed_out_offsets, self.completed_epoch_counter, self.completed_t_counter, + self._migration_checkpoint_blob, ), ) self.s3_client.put_object( @@ -229,6 +275,10 @@ async def update_stateflow_graph(self, new_stateflow_graph: StateflowGraph) -> N raise NotAStateflowGraphError # TODO the cluster was balanced by the previous deployment, if the graph is complex it might be # unbalanced after the update + logging.warning( + f"update_stateflow_graph | updating worker pool operators " + f"(new partitions: { {n: op.n_partitions for n, op in new_stateflow_graph.nodes.items()} })", + ) for _, operator in iter(new_stateflow_graph): for partition in range(self.max_operator_parallelism): operator_copy = deepcopy(operator) @@ -239,6 +289,15 @@ async def update_stateflow_graph(self, new_stateflow_graph: StateflowGraph) -> N (operator_copy.name, partition), operator_copy, ) + # Set _pending_graph BEFORE sending InitMigration to workers. + # This ensures recovery can detect an in-flight migration even if the + # gather is stuck retrying a dead worker connection for tens of seconds. + # Recovery clears _pending_graph, so if the gather completes after + # recovery, it harmlessly re-sets the field (which _perform_recovery + # already handles via saved_migration_in_progress). + self._pending_graph = new_stateflow_graph + logging.warning("update_stateflow_graph | _pending_graph set (before send)") + worker_assignments = self.worker_pool.get_worker_assignments() tasks = [ self.networking.send_message( @@ -256,15 +315,68 @@ async def update_stateflow_graph(self, new_stateflow_graph: StateflowGraph) -> N for worker in self.worker_pool.get_participating_workers() ] await asyncio.gather(*tasks) - self.submitted_graph = new_stateflow_graph + logging.warning("update_stateflow_graph | InitMigration sent to all workers") + + def revert_worker_pool_to_submitted_graph(self) -> None: + """Revert worker pool operators to match the current submitted_graph. + + Called before recovery when migration was in progress — undoes the + operator promotions (shadow → active) that update_stateflow_graph() + applied so that recovery sends the correct (OLD) layout to workers. + Also updates orphaned_operator_assignments so that dead-worker + partitions are rescheduled with the OLD layout operator definitions. + """ + if self.submitted_graph is None or self.max_operator_parallelism is None: + return + for _, operator in iter(self.submitted_graph): + for partition in range(self.max_operator_parallelism): + operator_copy = deepcopy(operator) + if partition >= operator.n_partitions: + operator_copy.make_shadow() + op_part = (operator_copy.name, partition) + self.worker_pool.update_operator(op_part, operator_copy) + # Also revert orphaned (dead-worker) partitions + if op_part in self.worker_pool.orphaned_operator_assignments: + self.worker_pool.orphaned_operator_assignments[op_part] = operator_copy + logging.warning("Reverted worker pool operators to match submitted_graph (pre-migration layout)") + + def finalize_graph_update(self) -> None: + """Commit the deferred graph update after migration completes successfully.""" + if self._pending_graph is None: + logging.warning("finalize_graph_update | no-op (_pending_graph is None)") + return + logging.warning("finalize_graph_update | committing pending graph to submitted_graph") + self.submitted_graph = self._pending_graph + self._pending_graph = None + # Write to Kafka via a background task; the producer is already running + # and the next snapshot will capture the new layout. metadata_key = msgpack_serialization(self.submitted_graph.name) - serialized_graph = cloudpickle_serialization(new_stateflow_graph) - await self.kafka_metadata_producer.send_and_wait( - "styx-metadata", - key=metadata_key, - value=serialized_graph, + serialized_graph = cloudpickle_serialization(self.submitted_graph) + self._pending_kafka_write = asyncio.ensure_future( + self.kafka_metadata_producer.send_and_wait( + "styx-metadata", + key=metadata_key, + value=serialized_graph, + ), + ) + + def set_migration_checkpoint( + self, + new_graph: StateflowGraph, + operator_partition_locations: dict, + ) -> None: + """Serialize migration metadata for inclusion in the next sequencer file.""" + self._migration_checkpoint_blob = cloudpickle_serialization( + { + "new_graph": new_graph, + "operator_partition_locations": operator_partition_locations, + }, ) + def clear_migration_checkpoint(self) -> None: + """Clear migration metadata so subsequent snapshots are normal.""" + self._migration_checkpoint_blob = None + async def submit_stateflow_graph( self, stateflow_graph: StateflowGraph, diff --git a/coordinator/coordinator_service.py b/coordinator/coordinator_service.py index a414585c..3a708197 100755 --- a/coordinator/coordinator_service.py +++ b/coordinator/coordinator_service.py @@ -22,7 +22,11 @@ from styx.common.logging import logging from styx.common.message_types import MessageType from styx.common.protocols import Protocols -from styx.common.serialization import Serializer +from styx.common.serialization import ( + Serializer, + cloudpickle_deserialization, + zstd_msgpack_deserialization, +) from styx.common.tcp_networking import MessagingMode, NetworkingManager from styx.common.util.aio_task_scheduler import AIOTaskScheduler import uvloop @@ -192,6 +196,10 @@ def __init__(self) -> None: } self.snapshotting_task: asyncio.Task | None = None + self._migration_checkpoint_task: asyncio.Task | None = None + # Deferred migration re-trigger: set during recovery, picked up after + # the first post-recovery snapshot completes (to avoid dedup races). + self._deferred_migration_graph: StateflowGraph | None = None self._protocol_controller_handlers_map: dict[MessageType, Callable[[bytes], Awaitable[None]]] = { MessageType.AriaProcessingDone: self._handle_aria_processing_done, @@ -282,10 +290,34 @@ async def _handle_update_execution_graph( async def _start_migration(self, graph: StateflowGraph) -> None: # Phase A: do NOT stop the protocol yet — workers will rehash in the background self.migration_in_progress = True - await self.stop_snapshotting() + + # Force a pre-migration snapshot at the next epoch boundary + if self.aria_metadata is not None: + self.aria_metadata.take_snapshot_at_next_epoch() + self.coordinator.pre_migration_snapshot_pending = True logging.warning(f"MIGRATION | START {graph}") - await self.coordinator.update_stateflow_graph(graph) + try: + await self.coordinator.update_stateflow_graph(graph) + except Exception: + # update_stateflow_graph failed (e.g. dead worker connection). + # Revert the pool and clear migration state. + logging.exception("MIGRATION | update_stateflow_graph FAILED — reverting") + self.coordinator.revert_worker_pool_to_submitted_graph() + self.coordinator._pending_graph = None # noqa: SLF001 + self.coordinator.pre_migration_snapshot_pending = False + self.migration_in_progress = False + return + + # Guard: if recovery ran while the gather was blocked (zombie coroutine), + # migration_in_progress is now False and _pending_graph was cleared. + # Don't proceed — the recovery handler already dealt with it. + if not self.migration_in_progress: + logging.warning( + "MIGRATION | update_stateflow_graph completed but migration_in_progress=False " + "(recovery intervened while gather was blocked) — aborting", + ) + return n_workers = len(self.coordinator.worker_pool.get_participating_workers()) self.migration_metadata = MigrationMetadata(n_workers) @@ -305,10 +337,21 @@ async def _handle_migration_repartitioning_done( __: bytes, ___: concurrent.futures.ProcessPoolExecutor, ) -> None: + if not self.migration_in_progress: + logging.warning("Dropping stale MigrationRepartitioningDone (no migration in progress)") + return mt = MessageType.MigrationRepartitioningDone logging.warning("Migration repartitioning done received!") async with self.networking_locks[mt]: + logging.warning("MigrationRepartitioningDone | acquired lock") + # Re-check inside the lock to close the TOCTOU window: + # recovery may have cleared migration_in_progress between the + # outer guard and lock acquisition. + if not self.migration_in_progress: + logging.warning("Dropping stale MigrationRepartitioningDone (recovery intervened)") + return + sync_complete: bool = await self.migration_metadata.repartitioning_done() logging.warning(f"Migration repartitioning is complete: {sync_complete}") @@ -316,8 +359,33 @@ async def _handle_migration_repartitioning_done( if not sync_complete: return + # Re-check after await — recovery could have fired during the sync. + if not self.migration_in_progress: + logging.warning("Dropping MigrationRepartitioningDone (recovery intervened after sync)") + return + + # Set migration checkpoint blob BEFORE telling workers to proceed. + # Workers will take a checkpoint snapshot after receiving + # MigrationRepartitioningDone; when the SnapID arrives the coordinator + # writes the sequencer file which will include this blob. + # Use max(..., 0) so the baseline is always >= 0 — after recovery + # worker_snapshot_ids are reset to -1, but we still need the + # checkpoint detection (baseline >= 0) to fire on the next snapshot. + baseline = max(self.coordinator.get_current_completed_snapshot_id(), 0) + self.coordinator._migration_checkpoint_baseline_snap_id = baseline # noqa: SLF001 + operator_partition_locations = self.coordinator.worker_pool.get_operator_partition_locations() + self.coordinator.set_migration_checkpoint( + self.coordinator._pending_graph, # noqa: SLF001 + operator_partition_locations, + ) + logging.warning( + f"MIGRATION | Checkpoint blob set, baseline snap_id={baseline}", + ) + + logging.warning("MigrationRepartitioningDone | calling finalize_migration_repartition") await self.finalize_migration_repartition() await self.migration_metadata.cleanup(mt) + logging.warning("MigrationRepartitioningDone | done") async def _handle_migration_init_done( self, @@ -325,8 +393,17 @@ async def _handle_migration_init_done( data: bytes, __: concurrent.futures.ProcessPoolExecutor, ) -> None: + if not self.migration_in_progress: + logging.warning("Dropping stale MigrationInitDone (no migration in progress)") + return mt = MessageType.MigrationInitDone async with self.networking_locks[mt]: + logging.warning("MigrationInitDone | acquired lock") + # Re-check inside the lock (TOCTOU protection). + if not self.migration_in_progress: + logging.warning("Dropping stale MigrationInitDone (recovery intervened)") + return + epoch_counter, t_counter, input_offsets, output_offsets = self.networking.decode_message(data) sync_complete: bool = await self.migration_metadata.init_done( @@ -340,13 +417,37 @@ async def _handle_migration_init_done( if not sync_complete: return + # Re-check after await — recovery could have fired during the sync. + if not self.migration_in_progress: + logging.warning("Dropping MigrationInitDone (recovery intervened after sync)") + return + logging.warning("MIGRATION | MigrationInitDone | sync_complete") n_workers = len(self.coordinator.worker_pool.get_participating_workers()) self.aria_metadata = AriaSyncMetadata(n_workers) + logging.warning("MigrationInitDone | closing protocol connections") await self.protocol_networking.close_all_connections() - await self.finalize_migration() - await self.migration_metadata.cleanup(mt) + logging.warning( + f"MigrationInitDone | protocol connections closed, " + f"current_snap={self.coordinator.get_current_completed_snapshot_id()} " + f"baseline={self.coordinator._migration_checkpoint_baseline_snap_id}", # noqa: SLF001 + ) + + # The checkpoint blob was set in _handle_migration_repartitioning_done + # (before workers took their snapshot). By now the checkpoint snapshot + # may already be complete. Handle the race, then launch the finalizer. + baseline = self.coordinator._migration_checkpoint_baseline_snap_id # noqa: SLF001 + if baseline >= 0 and self.coordinator.get_current_completed_snapshot_id() > baseline: + self.coordinator._migration_checkpoint_snapshot_complete.set() # noqa: SLF001 + self.coordinator._migration_checkpoint_baseline_snap_id = -1 # noqa: SLF001 + logging.warning("MIGRATION | Checkpoint snapshot already completed (race handled)") + + # Launch background task to wait for the checkpoint snapshot, then finalize + logging.warning("MigrationInitDone | launching checkpoint finalizer task") + self._migration_checkpoint_task = asyncio.ensure_future( + self._wait_for_migration_checkpoint_and_finalize(mt), + ) async def _handle_register_worker( self, @@ -428,6 +529,24 @@ async def _handle_snap_id( pool, ) + # Check for deferred migration re-trigger: fire once the first + # post-recovery snapshot completes (prev changes from -1 to a + # real value, meaning a full cluster snapshot has been written). + if ( + self._deferred_migration_graph is not None + and not self.migration_in_progress + and self.coordinator.prev_completed_snapshot_id >= 0 + ): + graph = self._deferred_migration_graph + self._deferred_migration_graph = None + logging.warning( + "[RECOVERY] First post-recovery snapshot complete — re-triggering deferred migration", + ) + # Use ensure_future so we don't block the SnapID handler. + self._deferred_migration_task = asyncio.ensure_future( + self._start_migration(graph), + ) + async def _handle_heartbeat( self, _: StreamWriter, @@ -479,6 +598,8 @@ async def protocol_controller(self, data: bytes) -> None: f"COORDINATOR PROTOCOL SERVER: Non supported message type: {mt}", ) return + if self.migration_in_progress: + logging.warning(f"PROTOCOL_MSG during migration: {mt.name}") await handler(data) # ------------------------ @@ -587,11 +708,18 @@ async def _handle_sync_cleanup(self, data: bytes) -> None: if not sync_complete: return + stop_flag = self.aria_metadata.stop_next_epoch + if stop_flag: + logging.warning( + "SyncCleanup | sync_complete, sending stop_gracefully=True to all workers", + ) await self.finalize_worker_sync( mt, - (self.aria_metadata.stop_next_epoch,), + (stop_flag,), Serializer.MSGPACK, ) + if stop_flag: + logging.warning("SyncCleanup | stop_gracefully sent successfully") self.aria_metadata.cleanup(epoch_end=True) def _record_epoch_metrics( @@ -657,8 +785,16 @@ async def _handle_deterministic_reordering(self, data: bytes) -> None: self.aria_metadata.cleanup() async def _handle_migration_done(self, _: bytes) -> None: + if not self.migration_in_progress: + logging.warning("Dropping stale MigrationDone (no migration in progress)") + return mt = MessageType.MigrationDone async with self.networking_locks[mt]: + logging.warning("MigrationDone | acquired lock") + # Re-check inside the lock (TOCTOU protection). + if not self.migration_in_progress: + logging.warning("Dropping stale MigrationDone (recovery intervened)") + return sync_complete: bool = await self.migration_metadata.set_empty_sync_done(mt) logging.warning( f"MIGRATION | MigrationDone | {self.migration_metadata.sync_sum}", @@ -672,8 +808,15 @@ async def _handle_migration_done(self, _: bytes) -> None: await self.migration_metadata.cleanup(mt) self.migration_in_progress = False - logging.warning("Restarting the snapshotting mechanism") - self.snapshotting_task = asyncio.create_task(self.send_snapshot_marker()) + # Clear migration checkpoint blob so subsequent snapshots are normal. + # The graph was already finalized in _wait_for_migration_checkpoint_and_finalize(). + self.coordinator.clear_migration_checkpoint() + + # Defer final graph write until a post-migration snapshot captures + # the state without migration metadata. This is a safety net — the + # graph was already updated via finalize_graph_update() in the + # checkpoint flow, but we want a clean snapshot without migration_blob. + self.coordinator.post_migration_snapshot_pending = True async def start_puller(self) -> None: async def request_handler(reader: StreamReader, writer: StreamWriter) -> None: @@ -685,11 +828,13 @@ async def request_handler(reader: StreamReader, writer: StreamWriter) -> None: self.protocol_controller(await reader.readexactly(size)), ) except asyncio.IncompleteReadError as e: - logging.info(f"Client disconnected unexpectedly: {e}") + logging.warning(f"Protocol client disconnected unexpectedly: {e}") except asyncio.CancelledError: pass + except Exception as e: + logging.warning(f"Protocol request_handler unexpected error: {e}") finally: - logging.info("Closing the connection") + logging.warning("Protocol connection closing") writer.close() await writer.wait_closed() @@ -783,6 +928,39 @@ async def finalize_migration(self) -> None: ), ) + async def _wait_for_migration_checkpoint_and_finalize(self, mt: MessageType) -> None: + """Wait for the migration checkpoint snapshot to complete, then finalize migration.""" + logging.warning("checkpoint_finalizer | waiting for checkpoint snapshot event") + await self.coordinator._migration_checkpoint_snapshot_complete.wait() # noqa: SLF001 + self.coordinator._migration_checkpoint_snapshot_complete.clear() # noqa: SLF001 + logging.warning( + f"checkpoint_finalizer | event fired, migration_in_progress={self.migration_in_progress}", + ) + + # Guard: recovery may have cleared migration state while we were waiting. + if not self.migration_in_progress: + logging.warning( + "MIGRATION | Checkpoint snapshot completed but migration_in_progress " + "is False (recovery intervened) — aborting finalization", + ) + return + + logging.warning( + f"MIGRATION | Checkpoint snapshot complete (snap_id=" + f"{self.coordinator.get_current_completed_snapshot_id()}) — " + f"publishing new graph and sending MigrationDone", + ) + + # Publish new graph to styx-metadata so ALL clients update partitioners + logging.warning("checkpoint_finalizer | calling finalize_graph_update") + self.coordinator.finalize_graph_update() + + # Send MigrationDone to workers + logging.warning("checkpoint_finalizer | sending MigrationDone to workers") + await self.finalize_migration() + await self.migration_metadata.cleanup(mt) + logging.warning("checkpoint_finalizer | done") + async def finalize_worker_sync( self, msg_type: MessageType, @@ -832,7 +1010,10 @@ async def _reset_after_recovery(self) -> None: # 2) Reset migration metadata self.migration_metadata = MigrationMetadata(n_workers) - self.migration_in_progress = False + # migration_in_progress, pre_migration_snapshot_pending, _pending_graph, + # and migration checkpoint baseline are already cleared in _perform_recovery() + # step 1c (before recovery starts) + self.coordinator.pre_migration_snapshot_id = -1 # 3) Reset snapshot completion metadata self.coordinator.completed_input_offsets.clear() @@ -860,6 +1041,27 @@ async def _reset_after_recovery(self) -> None: logging.warning("Protocol metadata reset complete") + def _read_migration_blob_from_s3(self, snap_id: int) -> bytes | None: + """Read the migration checkpoint blob from the sequencer file in S3.""" + if snap_id <= 0: + return None + try: + resp = self.s3_client.get_object( + Bucket=SNAPSHOT_BUCKET_NAME, + Key=f"sequencer/{snap_id}.bin", + ) + loaded = zstd_msgpack_deserialization(resp["Body"].read()) + _sequencer_tuple_len_with_migration = 5 + if ( + isinstance(loaded, (tuple, list)) + and len(loaded) == _sequencer_tuple_len_with_migration + and loaded[4] is not None + ): + return loaded[4] + except Exception as e: + logging.warning(f"[RECOVERY] Failed to read migration blob from S3: {e}") + return None + async def _perform_recovery(self, workers_to_remove: set[Worker]) -> None: """ Full recovery state machine: @@ -875,6 +1077,36 @@ async def _perform_recovery(self, workers_to_remove: set[Worker]) -> None: logging.warning(f"Starting recovery process for workers: {workers_to_remove}") + # 0) Immediately cancel the migration checkpoint background task. + # This MUST happen before reading any migration state (_pending_graph, + # submitted_graph) to prevent finalize_graph_update() from committing + # the new graph layout while recovery is reverting to the old one. + has_checkpoint_task = ( + hasattr(self, "_migration_checkpoint_task") and self._migration_checkpoint_task is not None + ) + logging.warning( + f"[RECOVERY] step 0 | cancel checkpoint task={has_checkpoint_task} " + f"migration_in_progress={self.migration_in_progress}", + ) + if has_checkpoint_task: + self._migration_checkpoint_task.cancel() + with contextlib.suppress(asyncio.CancelledError): + await self._migration_checkpoint_task + self._migration_checkpoint_task = None + logging.warning("[RECOVERY] step 0 | checkpoint task cancelled") + + # 0b) Clear migration_in_progress immediately so any migration handler + # that passed its guard check before recovery will see it is False + # when it re-checks inside the lock (TOCTOU protection). + # Save the old value first — we need it for was_migrating detection + # even when _pending_graph was never set (e.g. update_stateflow_graph + # failed mid-send). + saved_migration_in_progress = self.migration_in_progress + self.migration_in_progress = False + logging.warning( + f"[RECOVERY] step 0b | migration_in_progress was {saved_migration_in_progress}, now False", + ) + # 1) Clean up dead worker channels and buffered tasks logging.warning(f"Closing connections to dead workers: {workers_to_remove}") for worker in workers_to_remove: @@ -889,9 +1121,61 @@ async def _perform_recovery(self, workers_to_remove: set[Worker]) -> None: await self.aio_task_scheduler.close() self.aio_task_scheduler = AIOTaskScheduler() + # 1b) Check if the latest snapshot contains a migration checkpoint. + # If yes (post-checkpoint recovery): use the NEW layout — workers will + # re-derive keys_to_send from the snapshot and resume async transfer. + # If no (pre-checkpoint recovery): revert to old layout and optionally + # re-trigger migration from scratch. + was_migrating = ( + saved_migration_in_progress + or self.coordinator._pending_graph is not None # noqa: SLF001 + or self.coordinator.post_migration_snapshot_pending + or self.coordinator._migration_checkpoint_baseline_snap_id >= 0 # noqa: SLF001 + ) + saved_pending_graph = None + snap_id = self.coordinator.get_current_completed_snapshot_id() + + if was_migrating: + logging.warning( + f"[RECOVERY] Migration was in progress. " + f"pre_migration_snapshot_id={self.coordinator.pre_migration_snapshot_id}, " + f"current_completed_snapshot_id={snap_id}, " + f"worker_snapshot_ids={self.coordinator.worker_snapshot_ids}", + ) + migration_blob = self._read_migration_blob_from_s3(snap_id) + if migration_blob is not None: + # POST-CHECKPOINT: use NEW layout — workers will resume transfer + logging.warning("[RECOVERY] Migration checkpoint found — using NEW layout for recovery") + migration_meta = cloudpickle_deserialization(migration_blob) + new_graph = migration_meta["new_graph"] + self.coordinator.submitted_graph = new_graph + self.coordinator._pending_graph = None # noqa: SLF001 + # The worker pool already has the new operator assignments from + # update_stateflow_graph() — no need to revert. + else: + # PRE-CHECKPOINT: revert to old layout and retry migration + logging.warning("[RECOVERY] No migration checkpoint — reverting to OLD layout") + saved_pending_graph = self.coordinator._pending_graph # noqa: SLF001 + self.coordinator.revert_worker_pool_to_submitted_graph() + + # Clear remaining migration state so stale messages are dropped. + logging.warning("[RECOVERY] step 1c | clearing remaining migration state") + self.coordinator.pre_migration_snapshot_pending = False + self.coordinator.post_migration_snapshot_pending = False + self.coordinator._migration_checkpoint_baseline_snap_id = -1 # noqa: SLF001 + self.coordinator._migration_checkpoint_snapshot_complete.clear() # noqa: SLF001 + self.coordinator._pending_graph = None # noqa: SLF001 + self.coordinator.clear_migration_checkpoint() + self._deferred_migration_graph = None # Clear any pending deferred re-trigger + # 2) Start recovery + graph_parts = ( + {n: op.n_partitions for n, op in self.coordinator.submitted_graph.nodes.items()} + if self.coordinator.submitted_graph + else None + ) logging.warning( - "Starting recovery process (reassign operators, send InitRecovery)", + f"[RECOVERY] Starting recovery (send InitRecovery). snap_id={snap_id}, graph_partitions={graph_parts}", ) await self.coordinator.start_recovery_process(workers_to_remove) @@ -912,6 +1196,21 @@ async def _perform_recovery(self, workers_to_remove: set[Worker]) -> None: logging.warning("Recovery process completed") + # 7) If migration was interrupted, DEFER re-trigger until the first + # post-recovery snapshot completes. Re-triggering immediately causes + # duplicate output: the protocol starts replaying from the snapshot, + # the dedup mechanism suppresses most duplicates, but the immediate + # migration Phase B stops the protocol before all replayed messages + # are fully dedup'd, and the migration restart's dedup scan starts + # from Phase B offsets (missing the earlier duplicates). + # By deferring, we let the recovery epoch fully complete and a clean + # snapshot capture the post-dedup state, eliminating duplicates. + if saved_pending_graph is not None: + logging.warning( + "[RECOVERY] Deferring migration re-trigger until first post-recovery snapshot", + ) + self._deferred_migration_graph = saved_pending_graph + async def heartbeat_monitor_coroutine(self) -> None: interval_time = HEARTBEAT_CHECK_INTERVAL / 1000 while True: @@ -928,6 +1227,11 @@ async def heartbeat_monitor_coroutine(self) -> None: # Add workers that re-registered (same IP/ports) to the failed set if (workers_to_remove or self.workers_that_re_registered) and self.recovery_state == RecoveryState.IDLE: + logging.warning( + f"heartbeat_monitor | entering recovery lock " + f"dead={workers_to_remove} re_registered={self.workers_that_re_registered} " + f"migration_in_progress={self.migration_in_progress}", + ) async with self.recovery_lock: if self.recovery_state != RecoveryState.IDLE: # Another recovery started while we were waiting for the lock @@ -944,6 +1248,7 @@ async def heartbeat_monitor_coroutine(self) -> None: logging.error(f"Error during recovery: {e}") finally: self.recovery_state = RecoveryState.IDLE + logging.warning("heartbeat_monitor | recovery_state back to IDLE") async def send_snapshot_marker(self) -> None: while True: diff --git a/coordinator/worker_pool.py b/coordinator/worker_pool.py index 5351f749..27ef3a25 100644 --- a/coordinator/worker_pool.py +++ b/coordinator/worker_pool.py @@ -153,7 +153,14 @@ def update_operator( operator_partition: OperatorPartition, operator: Operator | BaseOperator, ) -> None: - worker_id = self.operator_partition_to_worker[operator_partition] + worker_id = self.operator_partition_to_worker.get(operator_partition) + if worker_id is None or worker_id not in self._worker_queue_idx: + # Worker was already removed (e.g. dead before recovery revert). + # The partition will be rescheduled during initiate_recovery. + logging.debug( + f"update_operator | skipping {operator_partition} — worker {worker_id} removed", + ) + return worker = self.peek(worker_id) worker.assigned_operators[operator_partition] = operator diff --git a/demo/demo-migration-ft-ycsb/__init__.py b/demo/demo-migration-ft-ycsb/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/demo/demo-migration-ft-ycsb/calculate_metrics.py b/demo/demo-migration-ft-ycsb/calculate_metrics.py new file mode 120000 index 00000000..ac136e83 --- /dev/null +++ b/demo/demo-migration-ft-ycsb/calculate_metrics.py @@ -0,0 +1 @@ +../demo-migration-ycsb/calculate_metrics.py \ No newline at end of file diff --git a/demo/demo-migration-ft-ycsb/client.py b/demo/demo-migration-ft-ycsb/client.py new file mode 100644 index 00000000..3890e3d8 --- /dev/null +++ b/demo/demo-migration-ft-ycsb/client.py @@ -0,0 +1,236 @@ +"""YCSB migration + fault tolerance client. + +Usage: + python client.py + + + + +- migration_at_sec: second at which to trigger migration (-1 = disabled) +- kill_at_sec: second at which to `docker kill styx-worker-1` (-1 = disabled) +""" + +import multiprocessing +from multiprocessing import Pool, cpu_count +import os +import pickle +import random +import string +import subprocess +import sys +import time +from timeit import default_timer as timer + +import boto3 + +import calculate_metrics +import kafka_output_consumer +import pandas as pd +from styx.client.sync_client import SyncStyxClient +from styx.common.local_state_backends import LocalStateBackend +from styx.common.operator import Operator +from styx.common.serialization import ( + cloudpickle_deserialization, + cloudpickle_serialization, +) +from styx.common.stateflow_graph import StateflowGraph +from tqdm import tqdm +from ycsb import ycsb_operator + + +def ycsb_field(size=100, seed=0): + "Translated from the original Java code" + if seed is not None: + random.seed(seed) + charset = string.ascii_letters + string.digits + return "".join(random.choices(charset, k=size)) + + +N_ENTITIES = int(sys.argv[8]) +MIGRATION_AT_SEC = int(sys.argv[9]) +KILL_AT_SEC = int(sys.argv[10]) + +threads = int(sys.argv[1]) +START_N_PARTITIONS = int(sys.argv[2]) +END_N_PARTITIONS = int(sys.argv[3]) +messages_per_second = int(sys.argv[4]) +seconds = int(sys.argv[5]) +SAVE_DIR: str = sys.argv[6] +warmup_seconds: int = int(sys.argv[7]) + +BATCH_SIZE = 100_000 + +sleeps_per_second = 100 +sleep_time = 0.0085 +STYX_HOST: str = os.getenv("STYX_HOST", "localhost") +STYX_PORT: int = int(os.getenv("STYX_PORT", "8886")) +KAFKA_URL: str = os.getenv("KAFKA_URL", "localhost:9092") + +g = StateflowGraph( + "ycsb-benchmark", + operator_state_backend=LocalStateBackend.DICT, + max_operator_parallelism=max(START_N_PARTITIONS, END_N_PARTITIONS), +) +ycsb_operator.set_n_partitions(START_N_PARTITIONS) +g.add_operators(ycsb_operator) + + +def submit_graph(styx: SyncStyxClient): + print(f"Partitions: {list(g.nodes.values())[0].n_partitions}") + styx.submit_dataflow(g) + print("Graph submitted") + + +def generate_partition_batch_wrapper(args): + batch_start, batch_end, operator_state = args + current_active_graph, operator_name = operator_state + current_active_graph = cloudpickle_deserialization(current_active_graph) + local_partitions = {p: {} for p in range(START_N_PARTITIONS)} + for i in range(batch_start, batch_end): + partition = current_active_graph.get_operator_by_name(operator_name).which_partition(i) + local_partitions[partition][i] = tuple(ycsb_field() for _ in range(10)) + return local_partitions + + +def merge_partitions(global_partitions, local_partitions): + for part_id, entries in local_partitions.items(): + global_partitions[part_id].update(entries) + + +def ycsb_init(styx: SyncStyxClient, operator: Operator, num_workers: int = None): + styx.set_graph(g) + styx.init_metadata(g) + partitions = {p: {} for p in range(START_N_PARTITIONS)} + + ycsb_dataset_path = f"ycsb_dataset_{START_N_PARTITIONS}p.pkl" + + if os.path.exists(ycsb_dataset_path): + print("Loading YCSB dataset...") + with open(ycsb_dataset_path, "rb") as f: + partitions = pickle.load(f) + else: + print("Generating YCSB dataset...") + batch_ranges = [ + (i, min(i + BATCH_SIZE, N_ENTITIES)) + for i in range(0, N_ENTITIES, BATCH_SIZE) + ] + operator_state = (cloudpickle_serialization(styx._current_active_graph), operator.name) + tasks = [(start, end, operator_state) for start, end in batch_ranges] + + with Pool(processes=num_workers or cpu_count()) as pool: + results = [] + for result in tqdm(pool.imap_unordered(generate_partition_batch_wrapper, tasks), total=len(tasks)): + results.append(result) + + for local_partitions in results: + merge_partitions(partitions, local_partitions) + with open(ycsb_dataset_path, "wb") as f: + pickle.dump(partitions, f) + print("Data ready") + for partition, partition_data in partitions.items(): + styx.init_data(operator, partition, partition_data) + print("Data loaded") + time.sleep(5) + submit_graph(styx) + + +def transactional_ycsb_generator(operator: Operator): + while True: + key = random.randint(0, N_ENTITIES - 1) + op = "read" if random.random() < 0.85 else "update" + yield operator, key, op + + +def benchmark_runner(proc_num) -> dict[bytes, dict]: + print(f"Generator: {proc_num} starting") + styx = SyncStyxClient(STYX_HOST, STYX_PORT, kafka_url=KAFKA_URL) + styx.open(consume=False) + ycsb_generator = transactional_ycsb_generator(ycsb_operator) + timestamp_futures: dict[bytes, dict] = {} + time.sleep(5) + start = timer() + for cur_sec in range(seconds): + # Kill a worker at the specified second (only proc 0) + if proc_num == 0 and 0 <= KILL_AT_SEC == cur_sec: + subprocess.run(["docker", "kill", "styx-worker-1"], check=False) + print(f"KILL -> styx-worker-1 at second {cur_sec}") + + sec_start = timer() + for i in range(messages_per_second): + if i % (messages_per_second // sleeps_per_second) == 0: + time.sleep(sleep_time) + operator, key, func_name = next(ycsb_generator) + future = styx.send_event(operator=operator, key=key, function=func_name) + timestamp_futures[future.request_id] = {"op": f"{func_name} {key}"} + styx.flush() + sec_end = timer() + lps = sec_end - sec_start + if lps < 1: + time.sleep(1 - lps) + sec_end2 = timer() + print(f"{cur_sec} | Latency per second: {sec_end2 - sec_start}") + + # Trigger migration at the specified second (only proc 0) + if cur_sec == MIGRATION_AT_SEC and proc_num == 0: + new_g = StateflowGraph("ycsb-benchmark", operator_state_backend=LocalStateBackend.DICT) + ycsb_operator.set_n_partitions(END_N_PARTITIONS) + new_g.add_operators(ycsb_operator) + styx.update_dataflow(new_g) + print(f"Migration request submitted at second {cur_sec}") + + end = timer() + print(f"Average latency per second: {(end - start) / seconds}") + + styx.close() + + for key, metadata in styx.delivery_timestamps.items(): + timestamp_futures[key]["timestamp"] = metadata + return timestamp_futures + + +def main(): + print("Generate and push workload to Styx") + print(f"Migration at second: {MIGRATION_AT_SEC}, Kill at second: {KILL_AT_SEC}") + + if N_ENTITIES < 3: + print("Impossible to run this benchmark with one key") + return + + s3 = boto3.client( + "s3", + endpoint_url=os.getenv("S3_ENDPOINT") or "http://localhost:9000", + aws_access_key_id=os.getenv("S3_ACCESS_KEY") or "rustfsadmin", + aws_secret_access_key=os.getenv("S3_SECRET_KEY") or "rustfsadmin", + region_name=os.getenv("S3_REGION") or "us-east-1", + ) + styx_client = SyncStyxClient(STYX_HOST, STYX_PORT, kafka_url=KAFKA_URL, s3=s3) + ycsb_init(styx_client, ycsb_operator) + del styx_client + time.sleep(5) + + with Pool(threads) as p: + results = p.map(benchmark_runner, range(threads)) + + results = {k: v for d in results for k, v in d.items()} + assert len(results) == messages_per_second * seconds * threads + + pd.DataFrame( + { + "request_id": list(results.keys()), + "timestamp": [res["timestamp"] for res in results.values()], + "op": [res["op"] for res in results.values()], + } + ).sort_values("timestamp").to_csv(f"{SAVE_DIR}/client_requests.csv", index=False) + + print("Workload completed") + + +if __name__ == "__main__": + multiprocessing.set_start_method("fork") + main() + + print() + kafka_output_consumer.main(SAVE_DIR) + + print() + calculate_metrics.main(messages_per_second, threads, warmup_seconds, SAVE_DIR) diff --git a/demo/demo-migration-ft-ycsb/kafka_output_consumer.py b/demo/demo-migration-ft-ycsb/kafka_output_consumer.py new file mode 120000 index 00000000..69f66589 --- /dev/null +++ b/demo/demo-migration-ft-ycsb/kafka_output_consumer.py @@ -0,0 +1 @@ +../demo-migration-ycsb/kafka_output_consumer.py \ No newline at end of file diff --git a/demo/demo-migration-ft-ycsb/ycsb.py b/demo/demo-migration-ft-ycsb/ycsb.py new file mode 120000 index 00000000..6ee86616 --- /dev/null +++ b/demo/demo-migration-ft-ycsb/ycsb.py @@ -0,0 +1 @@ +../demo-migration-ycsb/ycsb.py \ No newline at end of file diff --git a/docker-compose.yml b/docker-compose.yml index 7a54e02f..a5b5097a 100755 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -15,6 +15,7 @@ services: - S3_ENDPOINT=http://rustfs:9000 - S3_ACCESS_KEY=rustfsadmin - S3_SECRET_KEY=rustfsadmin + - LOG_LEVEL=${LOG_LEVEL:-WARNING} depends_on: - prometheus - grafana @@ -54,6 +55,7 @@ services: - S3_ENDPOINT=http://rustfs:9000 - S3_ACCESS_KEY=rustfsadmin - S3_SECRET_KEY=rustfsadmin + - LOG_LEVEL=${LOG_LEVEL:-WARNING} depends_on: - coordinator diff --git a/pytest.ini b/pytest.ini index 3961d894..3e9757b2 100644 --- a/pytest.ini +++ b/pytest.ini @@ -1,6 +1,7 @@ [pytest] markers = e2e: end-to-end tests that start containers + e2e_migration: e2e tests involving state migration integration: component-level tests using testcontainers log_cli = true log_cli_level = INFO diff --git a/styx-package/styx/common/logging.py b/styx-package/styx/common/logging.py index df59cb81..00ac1499 100644 --- a/styx-package/styx/common/logging.py +++ b/styx-package/styx/common/logging.py @@ -1,12 +1,15 @@ import logging as _stdlib_logging from logging.handlers import QueueHandler, QueueListener +import os from queue import SimpleQueue _log_queue: SimpleQueue = SimpleQueue() # Root "styx" logger — uses QueueHandler so callers never block. +# Log level is configurable via LOG_LEVEL env var (default: WARNING). logging = _stdlib_logging.getLogger("styx") -logging.setLevel(_stdlib_logging.WARNING) +_log_level = getattr(_stdlib_logging, os.getenv("LOG_LEVEL", "WARNING").upper(), _stdlib_logging.WARNING) +logging.setLevel(_log_level) if not logging.handlers: logging.addHandler(QueueHandler(_log_queue)) diff --git a/styx-package/styx/common/message_types.py b/styx-package/styx/common/message_types.py index 7bb6e380..adb742c2 100644 --- a/styx-package/styx/common/message_types.py +++ b/styx-package/styx/common/message_types.py @@ -45,3 +45,4 @@ class MessageType(IntEnum): UpdateExecutionGraph = 41 RunFunRemoteBatch = 42 AckBatch = 43 + SnapMigrationReassign = 44 diff --git a/tests/e2e/test_e2e_migration_ft_ycsb.py b/tests/e2e/test_e2e_migration_ft_ycsb.py new file mode 100644 index 00000000..4279f8b0 --- /dev/null +++ b/tests/e2e/test_e2e_migration_ft_ycsb.py @@ -0,0 +1,435 @@ +"""E2E tests for fault tolerance during state migration. + +These tests verify that Styx maintains exactly-once semantics when a worker +crashes at different points during a migration cycle: + +1. Normal migration (baseline) — no crash +2. Kill during Phase A — crash while protocol is still running, before + repartitioning starts. Recovery should roll back to pre-migration snapshot. +3. Kill during Phase C — crash after repartitioning, during async data + transfer. Recovery should roll back to pre-migration snapshot (no + post-migration snapshot has completed yet). +4. Kill after migration completes — crash after MigrationDone. Recovery + should use the post-migration snapshot with new partition layout. + +All tests validate: +- No duplicate requests +- Exactly-once output (no duplicate responses) +- Zero missed messages +""" + +from dataclasses import dataclass +import json +import logging +from pathlib import Path + +import pytest + +from tests.helpers import make_test_env, run_and_stream, run_and_stream_with_timed_action, wait_port + +log = logging.getLogger("e2e.migration_ft_ycsb") + + +def _assert_metrics( + results_dir: Path, + input_rate: int, + client_threads: int, +) -> None: + exp_name = f"ycsb_migration_{input_rate * client_threads}" + metrics_json = results_dir / f"{exp_name}.json" + log.info("Checking metrics json: %s", metrics_json) + assert metrics_json.exists(), f"Missing metrics json: {metrics_json}" + + with metrics_json.open("r", encoding="utf-8") as f: + metrics = json.load(f) + + log.info("==== METRICS JSON (%s) ====", metrics_json) + log.info("%s", json.dumps(metrics, indent=2, sort_keys=True)) + log.info("==== END METRICS JSON ====") + + assert metrics.get("duplicate_requests") is False, metrics + assert metrics.get("exactly_once_output") is True, metrics + assert int(metrics.get("missed messages")) == 0, metrics + + +@dataclass(frozen=True) +class _ClusterParams: + n_partitions: int = 4 + epoch_size: int = 1000 + threads_per_worker: int = 1 + enable_compression: str = "true" + use_composite_keys: str = "true" + use_fallback_cache: str = "true" + + +@dataclass(frozen=True) +class _ClientParams: + client_threads: int = 2 + n_entities: int = 10_000 + start_n_partitions: int = 4 + end_n_partitions: int = 8 + input_rate: int = 200 + total_time: int = 120 + warmup_seconds: int = 10 + migration_at_sec: int = 30 + kill_at_sec: int = -1 # <0 = disabled + + +@dataclass(frozen=True) +class _Paths: + repo_root: Path + demo_dir: Path + start_script: Path + stop_script: Path + + +def _resolve_paths() -> _Paths: + repo_root = Path(__file__).resolve().parents[2] + demo_dir = repo_root / "demo" / "demo-migration-ft-ycsb" + start_script = repo_root / "scripts" / "start_styx_cluster.sh" + stop_script = repo_root / "scripts" / "stop_styx_cluster.sh" + + assert demo_dir.exists(), f"Expected demo dir at {demo_dir}" + assert start_script.exists(), f"Missing: {start_script}" + assert stop_script.exists(), f"Missing: {stop_script}" + + return _Paths( + repo_root=repo_root, + demo_dir=demo_dir, + start_script=start_script, + stop_script=stop_script, + ) + + +def _make_results_dir(tmp_path: Path) -> Path: + results_dir = tmp_path / "results" + results_dir.mkdir(parents=True, exist_ok=True) + log.info("Results dir: %s", results_dir) + return results_dir + + +def _start_cmd(paths: _Paths, p: _ClusterParams) -> list[str]: + return [ + "bash", + str(paths.start_script), + str(p.n_partitions), + str(p.epoch_size), + str(p.threads_per_worker), + p.enable_compression, + p.use_composite_keys, + p.use_fallback_cache, + ] + + +def _stop_cmd(paths: _Paths, p: _ClusterParams) -> list[str]: + return ["bash", str(paths.stop_script), str(p.threads_per_worker)] + + +def _client_cmd(results_dir: Path, client: _ClientParams) -> list[str]: + # client.py expects: + # 1 threads + # 2 START_N_PARTITIONS + # 3 END_N_PARTITIONS + # 4 messages_per_second + # 5 seconds + # 6 SAVE_DIR + # 7 warmup_seconds + # 8 N_ENTITIES + # 9 migration_at_sec + # 10 kill_at_sec + return [ + "python", + "client.py", + str(client.client_threads), + str(client.start_n_partitions), + str(client.end_n_partitions), + str(client.input_rate), + str(client.total_time), + str(results_dir), + str(client.warmup_seconds), + str(client.n_entities), + str(client.migration_at_sec), + str(client.kill_at_sec), + ] + + +def _start_cluster_and_wait(paths: _Paths, env: dict, cluster: _ClusterParams) -> None: + rc, out = run_and_stream( + _start_cmd(paths, cluster), + cwd=str(paths.repo_root), + env=env, + timeout=20 * 60, + banner="START STYX CLUSTER", + log=log, + ) + if rc != 0: + raise AssertionError(f"start_styx_cluster.sh failed (rc={rc}). Output:\n{out}") + + log.info("Waiting for Kafka external listener (127.0.0.1:9092)...") + wait_port("127.0.0.1", 9092, timeout_s=180) + log.info("Kafka port is up.") + + log.info("Waiting for Styx coordinator published port (127.0.0.1:8886)...") + wait_port("127.0.0.1", 8886, timeout_s=240) + log.info("Coordinator port is up.") + + +def _run_client( + *, + paths: _Paths, + env: dict, + results_dir: Path, + client: _ClientParams, + timeout_s: int = 30 * 60, +) -> None: + rc, out = run_and_stream( + _client_cmd(results_dir, client), + cwd=str(paths.demo_dir), + env=env, + timeout=timeout_s, + banner="RUN MIGRATION FT YCSB CLIENT", + log=log, + ) + if rc != 0: + raise AssertionError(f"client.py failed (rc={rc}). Output:\n{out}") + + +def _assert_artifacts_and_metrics(results_dir: Path, client: _ClientParams) -> None: + client_csv = results_dir / "client_requests.csv" + output_csv = results_dir / "output.csv" + log.info("Checking artifacts: %s and %s", client_csv, output_csv) + + assert client_csv.exists(), f"Missing artifact: {client_csv}" + assert output_csv.exists(), f"Missing artifact: {output_csv}" + + _assert_metrics(results_dir, client.input_rate, client.client_threads) + + +def _stop_cluster(paths: _Paths, env: dict, cluster: _ClusterParams, *, timeout_s: int) -> None: + rc, out = run_and_stream( + _stop_cmd(paths, cluster), + cwd=str(paths.repo_root), + env=env, + timeout=timeout_s, + banner="STOP STYX CLUSTER", + log=log, + ) + if rc != 0: + log.warning("stop_styx_cluster.sh failed (rc=%s). Output:\n%s", rc, out) + + +# --------------------------------------------------------------------------- +# Test: normal migration (no crash) — baseline +# --------------------------------------------------------------------------- + + +@pytest.mark.e2e_migration +def test_migration_ft_baseline(tmp_path: Path): + """Normal migration completes without crash. + + Validates: no data loss, exactly-once output. + """ + paths = _resolve_paths() + results_dir = _make_results_dir(tmp_path) + + cluster = _ClusterParams() + client = _ClientParams( + total_time=120, + migration_at_sec=30, + kill_at_sec=-1, + ) + + env = make_test_env() + + try: + _start_cluster_and_wait(paths, env, cluster) + _run_client( + paths=paths, + env=env, + results_dir=results_dir, + client=client, + timeout_s=30 * 60, + ) + _assert_artifacts_and_metrics(results_dir, client) + finally: + _stop_cluster(paths, env, cluster, timeout_s=10 * 60) + + +# --------------------------------------------------------------------------- +# Test: kill during Phase A (protocol running, before repartitioning) +# --------------------------------------------------------------------------- + + +@pytest.mark.e2e_migration +def test_migration_ft_kill_during_phase_a(tmp_path: Path): + """Kill a worker during Phase A — crash while protocol is running. + + Migration starts at second 30. Kill at second 35: the protocol is still + running and rehash metadata may not have arrived yet. The coordinator + should detect the failure, recover to the pre-migration snapshot, and + retry the migration from scratch. + + Validates: recovery to pre-migration state, exactly-once output. + """ + paths = _resolve_paths() + results_dir = _make_results_dir(tmp_path) + + cluster = _ClusterParams() + client = _ClientParams( + total_time=120, + warmup_seconds=10, + migration_at_sec=30, + kill_at_sec=35, # Shortly after migration starts (Phase A) + ) + + env = make_test_env() + + try: + _start_cluster_and_wait(paths, env, cluster) + _run_client( + paths=paths, + env=env, + results_dir=results_dir, + client=client, + timeout_s=30 * 60, + ) + _assert_artifacts_and_metrics(results_dir, client) + finally: + _stop_cluster(paths, env, cluster, timeout_s=15 * 60) + + +# --------------------------------------------------------------------------- +# Test: kill during Phase C (async data transfer, before post-mig snapshot) +# --------------------------------------------------------------------------- + + +@pytest.mark.e2e_migration +def test_migration_ft_kill_during_phase_c(tmp_path: Path): + """Kill a worker during Phase C — async data transfer in progress. + + Migration starts at second 30. Kill at second 50: repartitioning has + completed but async data transfer is likely still in progress. No + post-migration snapshot should have completed yet. Recovery should + roll back to pre-migration snapshot with OLD partition layout. + + Validates: recovery to pre-migration state, exactly-once output. + """ + paths = _resolve_paths() + results_dir = _make_results_dir(tmp_path) + + cluster = _ClusterParams() + client = _ClientParams( + total_time=150, + warmup_seconds=10, + migration_at_sec=30, + kill_at_sec=50, # During async data transfer (Phase C) + ) + + env = make_test_env() + + try: + _start_cluster_and_wait(paths, env, cluster) + _run_client( + paths=paths, + env=env, + results_dir=results_dir, + client=client, + timeout_s=40 * 60, + ) + _assert_artifacts_and_metrics(results_dir, client) + finally: + _stop_cluster(paths, env, cluster, timeout_s=15 * 60) + + +# --------------------------------------------------------------------------- +# Test: kill after migration completes +# --------------------------------------------------------------------------- + + +@pytest.mark.e2e_migration +def test_migration_ft_kill_after_migration_done(tmp_path: Path): + """Kill a worker after migration completes. + + Migration starts at second 30. Kill at second 90: migration should have + completed and at least one post-migration snapshot should exist. Recovery + should use the post-migration snapshot with NEW partition layout. + + Validates: recovery to post-migration state, exactly-once output. + """ + paths = _resolve_paths() + results_dir = _make_results_dir(tmp_path) + + cluster = _ClusterParams() + client = _ClientParams( + total_time=150, + warmup_seconds=10, + migration_at_sec=30, + kill_at_sec=90, # Well after migration completes + ) + + env = make_test_env() + + try: + _start_cluster_and_wait(paths, env, cluster) + _run_client( + paths=paths, + env=env, + results_dir=results_dir, + client=client, + timeout_s=40 * 60, + ) + _assert_artifacts_and_metrics(results_dir, client) + finally: + _stop_cluster(paths, env, cluster, timeout_s=15 * 60) + + +# --------------------------------------------------------------------------- +# Test: kill during migration with external docker kill (timed action) +# --------------------------------------------------------------------------- + + +@pytest.mark.e2e_migration +def test_migration_ft_external_kill_during_migration(tmp_path: Path): + """Kill a worker externally via docker kill during migration. + + Uses run_and_stream_with_timed_action to kill a worker container + externally at a fixed wall-clock offset, independent of the client. + This simulates a more realistic failure scenario where the kill + timing is not synchronized with client message processing. + + Migration starts at second 30 (inside client). + External kill fires 45 seconds after client starts (~second 45). + """ + paths = _resolve_paths() + results_dir = _make_results_dir(tmp_path) + + cluster = _ClusterParams() + client = _ClientParams( + total_time=150, + warmup_seconds=10, + migration_at_sec=30, + kill_at_sec=-1, # Kill is handled externally + ) + + env = make_test_env() + + try: + _start_cluster_and_wait(paths, env, cluster) + + rc, out = run_and_stream_with_timed_action( + _client_cmd(results_dir, client), + cwd=str(paths.demo_dir), + env=env, + timeout=40 * 60, + banner="RUN MIGRATION FT YCSB CLIENT (external kill)", + action_at_s=45, # ~15 seconds after migration starts + action_cmd=["docker", "kill", "styx-worker-1"], + action_banner="KILL styx-worker-1", + log=log, + ) + if rc != 0: + raise AssertionError(f"client.py failed (rc={rc}). Output:\n{out}") + + _assert_artifacts_and_metrics(results_dir, client) + finally: + _stop_cluster(paths, env, cluster, timeout_s=15 * 60) diff --git a/tests/e2e/test_e2e_migration_tpcc.py b/tests/e2e/test_e2e_migration_tpcc.py index 43e2a376..f1e4c147 100644 --- a/tests/e2e/test_e2e_migration_tpcc.py +++ b/tests/e2e/test_e2e_migration_tpcc.py @@ -251,7 +251,7 @@ def _stop_cluster(paths: _Paths, env: dict, cluster: _ClusterParams, *, timeout_ log.warning("stop_styx_cluster.sh failed (rc=%s). Output:\n%s", rc, out) -@pytest.mark.e2e +@pytest.mark.e2e_migration def test_styx_e2e_migration_tpcc(tmp_path: Path): paths = _resolve_paths() results_dir = _make_results_dir(tmp_path) diff --git a/tests/e2e/test_e2e_migration_ycsb.py b/tests/e2e/test_e2e_migration_ycsb.py index d40e0f01..5e5ddb20 100644 --- a/tests/e2e/test_e2e_migration_ycsb.py +++ b/tests/e2e/test_e2e_migration_ycsb.py @@ -190,7 +190,7 @@ def _stop_cluster(paths: _Paths, env: dict, cluster: _ClusterParams, *, timeout_ log.warning("stop_styx_cluster.sh failed (rc=%s). Output:\n%s", rc, out) -@pytest.mark.e2e +@pytest.mark.e2e_migration def test_styx_e2e_migration_ycsb(tmp_path: Path): paths = _resolve_paths() results_dir = _make_results_dir(tmp_path) diff --git a/tests/integration/test_s3_snapshots.py b/tests/integration/test_s3_snapshots.py index 0f4e9578..6e4b7a3a 100644 --- a/tests/integration/test_s3_snapshots.py +++ b/tests/integration/test_s3_snapshots.py @@ -38,7 +38,7 @@ def test_retrieve_snapshot_negative_id_returns_empty(self, s3_bucket): from worker.fault_tolerance.async_snapshots import AsyncSnapshotsS3 snap = AsyncSnapshotsS3(worker_id=0, n_assigned_partitions=1) - data, off_in, off_out, ep, tc = snap.retrieve_snapshot(-1, [("any_op", 0)]) + data, off_in, off_out, ep, tc, _migration_blob = snap.retrieve_snapshot(-1, [("any_op", 0)]) assert data == {} assert off_in == {} assert off_out == {} @@ -68,7 +68,9 @@ def test_store_and_retrieve_round_trip(self, s3_bucket): # Retrieve snap = AsyncSnapshotsS3(worker_id=0, n_assigned_partitions=1) - data, _off_in, _off_out, ep, tc = snap.retrieve_snapshot(snapshot_id, [(operator_name, partition)]) + data, _off_in, _off_out, ep, tc, _migration_blob = snap.retrieve_snapshot( + snapshot_id, [(operator_name, partition)] + ) assert data[(operator_name, partition)] == state_data assert ep == epoch assert tc == t_counter @@ -101,7 +103,7 @@ def test_retrieve_merges_incremental_snapshots(self, s3_bucket): ) snap = AsyncSnapshotsS3(worker_id=0, n_assigned_partitions=1) - data, _off_in, _off_out, ep, tc = snap.retrieve_snapshot(2, [(op, partition)]) + data, _off_in, _off_out, ep, tc, _migration_blob = snap.retrieve_snapshot(2, [(op, partition)]) # b should be overwritten by snapshot 2 assert data[(op, partition)]["a"] == 1 diff --git a/tests/unit/coordinator/test_coordinator_metadata.py b/tests/unit/coordinator/test_coordinator_metadata.py index 61cf7848..58616639 100644 --- a/tests/unit/coordinator/test_coordinator_metadata.py +++ b/tests/unit/coordinator/test_coordinator_metadata.py @@ -286,3 +286,193 @@ async def test_rejects_non_graph(self): c = _coordinator() with pytest.raises(NotAStateflowGraphError): await c.update_stateflow_graph("not a graph") + + @pytest.mark.asyncio + async def test_pending_graph_set_before_gather(self): + """_pending_graph must be set BEFORE sending InitMigration to workers.""" + from copy import deepcopy + + c = _coordinator() + c.max_operator_parallelism = 4 + c.register_worker("10.0.0.1", 5000, 6000) + + # Manually schedule operator partitions (avoids Kafka in submit_stateflow_graph) + old_g = _graph(n_partitions=1, max_parallelism=4) + for _, operator in iter(old_g): + for p in range(4): + op_copy = deepcopy(operator) + if p >= operator.n_partitions: + op_copy.make_shadow() + c.worker_pool.schedule_operator_partition((op_copy.name, p), op_copy) + c.submitted_graph = old_g + + new_g = _graph(n_partitions=2, max_parallelism=4) + + # Track when _pending_graph was set relative to send_message calls + set_before_send = None + original_send = c.networking.send_message + + async def tracking_send(*args, **kwargs): + nonlocal set_before_send + if set_before_send is None: + set_before_send = c._pending_graph is not None + return await original_send(*args, **kwargs) + + c.networking.send_message = tracking_send + + await c.update_stateflow_graph(new_g) + assert set_before_send is True, "_pending_graph must be set before send_message" + assert c._pending_graph is new_g + + +# --------------------------------------------------------------------------- +# finalize_graph_update +# --------------------------------------------------------------------------- + + +class TestFinalizeGraphUpdate: + def test_noop_when_no_pending(self): + c = _coordinator() + c._pending_graph = None + c.finalize_graph_update() # should not raise + assert c.submitted_graph is None + + @pytest.mark.asyncio + async def test_commits_pending_graph(self): + c = _coordinator() + c.kafka_metadata_producer = AsyncMock() + old_graph = _graph(n_partitions=1) + new_graph = _graph(n_partitions=2) + c.submitted_graph = old_graph + c._pending_graph = new_graph + c.finalize_graph_update() + assert c.submitted_graph is new_graph + assert c._pending_graph is None + + +# --------------------------------------------------------------------------- +# set_migration_checkpoint / clear_migration_checkpoint +# --------------------------------------------------------------------------- + + +class TestMigrationCheckpoint: + def test_set_and_clear(self): + c = _coordinator() + g = _graph() + locations = {("users", 0): 1} + c.set_migration_checkpoint(g, locations) + assert c._migration_checkpoint_blob is not None + c.clear_migration_checkpoint() + assert c._migration_checkpoint_blob is None + + +# --------------------------------------------------------------------------- +# register_snapshot — migration checkpoint detection +# --------------------------------------------------------------------------- + + +class TestRegisterSnapshotCheckpointDetection: + def test_checkpoint_event_fires_when_snapshot_exceeds_baseline(self): + c = _coordinator() + pool = MagicMock() + c.worker_snapshot_ids = {0: 0} + c.prev_completed_snapshot_id = 0 + # Simulate checkpoint blob set with baseline 0 + c._migration_checkpoint_baseline_snap_id = 0 + c._migration_checkpoint_blob = b"fake" + + c.register_snapshot(0, 1, {}, {}, 1, 1, pool) + + assert c._migration_checkpoint_snapshot_complete.is_set() + assert c._migration_checkpoint_baseline_snap_id == -1 + + def test_checkpoint_event_does_not_fire_when_below_baseline(self): + c = _coordinator() + pool = MagicMock() + c.worker_snapshot_ids = {0: 0, 1: 0} + c.prev_completed_snapshot_id = 0 + c._migration_checkpoint_baseline_snap_id = 1 # baseline is 1 + + c.register_snapshot(0, 1, {}, {}, 1, 1, pool) + # Worker 1 still at 0, so cluster min is 0, not > baseline 1 + assert not c._migration_checkpoint_snapshot_complete.is_set() + + def test_pre_migration_snapshot_recorded(self): + c = _coordinator() + pool = MagicMock() + c.worker_snapshot_ids = {0: 0} + c.prev_completed_snapshot_id = 0 + c.pre_migration_snapshot_pending = True + + c.register_snapshot(0, 1, {}, {}, 1, 1, pool) + + assert c.pre_migration_snapshot_id == 1 + assert c.pre_migration_snapshot_pending is False + + @pytest.mark.asyncio + async def test_post_migration_finalize_fires(self): + c = _coordinator() + c.kafka_metadata_producer = AsyncMock() + pool = MagicMock() + c.worker_snapshot_ids = {0: 0} + c.prev_completed_snapshot_id = 0 + c.post_migration_snapshot_pending = True + c._pending_graph = _graph() + c.submitted_graph = _graph(n_partitions=1) + + c.register_snapshot(0, 1, {}, {}, 1, 1, pool) + + assert c.post_migration_snapshot_pending is False + + +# --------------------------------------------------------------------------- +# revert_worker_pool_to_submitted_graph +# --------------------------------------------------------------------------- + + +class TestRevertWorkerPool: + def test_revert_restores_old_layout(self): + from copy import deepcopy + + c = _coordinator() + c.max_operator_parallelism = 4 + wid, _ = c.register_worker("10.0.0.1", 5000, 6000) + + # Schedule initial partitions (like submit_stateflow_graph would) + old_graph = _graph(n_partitions=2, max_parallelism=4) + for _, operator in iter(old_graph): + for p in range(2): + c.worker_pool.schedule_operator_partition( + (operator.name, p), + deepcopy(operator), + ) + for p in range(2, 4): + op_copy = deepcopy(operator) + op_copy.make_shadow() + c.worker_pool.schedule_operator_partition( + (operator.name, p), + op_copy, + ) + c.submitted_graph = old_graph + + # Simulate update_stateflow_graph promoting shadow -> active + new_graph = _graph(n_partitions=4, max_parallelism=4) + for _, operator in iter(new_graph): + for partition in range(4): + op_copy = deepcopy(operator) + if partition >= operator.n_partitions: + op_copy.make_shadow() + c.worker_pool.update_operator((op_copy.name, partition), op_copy) + + # Partition 2 and 3 should now be active (4 partitions) + w = c.worker_pool.peek(wid) + assert not w.assigned_operators[("users", 2)].is_shadow + assert not w.assigned_operators[("users", 3)].is_shadow + + # Revert + c.revert_worker_pool_to_submitted_graph() + + # Partitions 2 and 3 should be shadow again (old graph had 2 partitions) + assert w.assigned_operators[("users", 2)].is_shadow + assert w.assigned_operators[("users", 3)].is_shadow + assert not w.assigned_operators[("users", 0)].is_shadow diff --git a/tests/unit/coordinator/test_migration_fault_tolerance.py b/tests/unit/coordinator/test_migration_fault_tolerance.py new file mode 100644 index 00000000..79f89613 --- /dev/null +++ b/tests/unit/coordinator/test_migration_fault_tolerance.py @@ -0,0 +1,527 @@ +"""Unit tests for migration fault tolerance features in coordinator. + +Covers: +- Pre-migration snapshot tracking (register_snapshot) +- Deferred graph update (finalize_graph_update) +- Post-migration snapshot triggers graph finalization (register_snapshot) +- Worker pool revert on recovery during migration +- Recovery condition: detect mid-migration AND post-migration-pre-snapshot crashes +- Full migration lifecycle scenarios +""" + +from unittest.mock import AsyncMock, MagicMock + +import pytest +from styx.common.local_state_backends import LocalStateBackend +from styx.common.operator import Operator +from styx.common.stateflow_graph import StateflowGraph + +from coordinator.coordinator_metadata import Coordinator + + +def _mock_coordinator(): + return Coordinator(MagicMock(), MagicMock()) + + +def _make_graph(name="test-app", n_partitions=4, max_parallelism=8): + g = StateflowGraph(name, operator_state_backend=LocalStateBackend.DICT, max_operator_parallelism=max_parallelism) + op = Operator("ycsb", n_partitions=n_partitions) + g.add_operators(op) + return g + + +# --------------------------------------------------------------------------- +# Step 8: Pre-migration snapshot tracking +# --------------------------------------------------------------------------- + + +class TestPreMigrationSnapshotTracking: + def test_pre_migration_snapshot_recorded(self): + c = _mock_coordinator() + c.worker_snapshot_ids = {0: 3, 1: 3} + c.prev_completed_snapshot_id = 2 + c.pre_migration_snapshot_pending = True + c.pre_migration_snapshot_id = -1 + + pool = MagicMock() + c.register_snapshot(0, 3, {}, {}, 10, 100, pool) + + assert c.pre_migration_snapshot_id == 3 + assert c.pre_migration_snapshot_pending is False + + def test_pre_migration_snapshot_not_recorded_when_not_pending(self): + c = _mock_coordinator() + c.worker_snapshot_ids = {0: 3, 1: 3} + c.prev_completed_snapshot_id = 2 + c.pre_migration_snapshot_pending = False + c.pre_migration_snapshot_id = -1 + + pool = MagicMock() + c.register_snapshot(0, 3, {}, {}, 10, 100, pool) + + assert c.pre_migration_snapshot_id == -1 + + def test_pre_migration_snapshot_not_recorded_when_no_advance(self): + c = _mock_coordinator() + c.worker_snapshot_ids = {0: 3, 1: 2} # min is 2, same as prev + c.prev_completed_snapshot_id = 2 + c.pre_migration_snapshot_pending = True + c.pre_migration_snapshot_id = -1 + + pool = MagicMock() + c.register_snapshot(0, 3, {}, {}, 10, 100, pool) + + # Snapshot didn't advance, so pre-migration not recorded yet + assert c.pre_migration_snapshot_pending is True + assert c.pre_migration_snapshot_id == -1 + + +# --------------------------------------------------------------------------- +# Step 9: Deferred graph update +# --------------------------------------------------------------------------- + + +class TestDeferredGraphUpdate: + @pytest.mark.asyncio + async def test_finalize_graph_update_commits_pending(self): + c = _mock_coordinator() + c.submitted_graph = MagicMock(name="old_graph") + new_graph = MagicMock(name="new_graph") + new_graph.name = "test-app" + c._pending_graph = new_graph + c.kafka_metadata_producer = AsyncMock() + + c.finalize_graph_update() + + assert c.submitted_graph is new_graph + assert c._pending_graph is None + + def test_finalize_graph_update_noop_when_no_pending(self): + c = _mock_coordinator() + old_graph = MagicMock(name="old_graph") + c.submitted_graph = old_graph + c._pending_graph = None + + c.finalize_graph_update() + + assert c.submitted_graph is old_graph + + def test_pending_graph_initialized_to_none(self): + c = _mock_coordinator() + assert c._pending_graph is None + assert c.pre_migration_snapshot_id == -1 + assert c.pre_migration_snapshot_pending is False + assert c.post_migration_snapshot_pending is False + + +# --------------------------------------------------------------------------- +# Post-migration snapshot triggers graph finalization +# --------------------------------------------------------------------------- + + +class TestPostMigrationSnapshotFinalization: + """Tests that finalize_graph_update() is only called when a post-migration + snapshot completes — NOT immediately when migration finishes. + + Bug caught: if finalize is called at MigrationDone time and a crash occurs + before the next snapshot, recovery uses the NEW graph but the OLD snapshot, + losing ~50% of keys. + """ + + def test_post_migration_snapshot_pending_initialized_false(self): + c = _mock_coordinator() + assert c.post_migration_snapshot_pending is False + + @pytest.mark.asyncio + async def test_register_snapshot_finalizes_graph_when_post_migration_pending(self): + """When post_migration_snapshot_pending is True and a new snapshot + completes, finalize_graph_update() should be called.""" + c = _mock_coordinator() + c.worker_snapshot_ids = {0: 3, 1: 3} + c.prev_completed_snapshot_id = 2 + c.post_migration_snapshot_pending = True + + new_graph = MagicMock(name="new_graph") + new_graph.name = "test-app" + c._pending_graph = new_graph + c.submitted_graph = MagicMock(name="old_graph") + c.kafka_metadata_producer = AsyncMock() + + pool = MagicMock() + c.register_snapshot(0, 3, {}, {}, 10, 100, pool) + + # Graph should be finalized + assert c.submitted_graph is new_graph + assert c._pending_graph is None + assert c.post_migration_snapshot_pending is False + + def test_register_snapshot_does_not_finalize_when_not_pending(self): + """When post_migration_snapshot_pending is False, finalize should + NOT be called even if snapshot advances.""" + c = _mock_coordinator() + c.worker_snapshot_ids = {0: 3, 1: 3} + c.prev_completed_snapshot_id = 2 + c.post_migration_snapshot_pending = False + + old_graph = MagicMock(name="old_graph") + c.submitted_graph = old_graph + c._pending_graph = MagicMock(name="new_graph") + + pool = MagicMock() + c.register_snapshot(0, 3, {}, {}, 10, 100, pool) + + # Graph should NOT be finalized + assert c.submitted_graph is old_graph + + def test_register_snapshot_does_not_finalize_when_snapshot_not_advanced(self): + """When snapshot hasn't advanced (not all workers reported), finalize + should NOT be called.""" + c = _mock_coordinator() + c.worker_snapshot_ids = {0: 3, 1: 2} # min is 2, same as prev + c.prev_completed_snapshot_id = 2 + c.post_migration_snapshot_pending = True + + old_graph = MagicMock(name="old_graph") + c.submitted_graph = old_graph + c._pending_graph = MagicMock(name="new_graph") + + pool = MagicMock() + c.register_snapshot(0, 3, {}, {}, 10, 100, pool) + + # Still pending — snapshot hasn't advanced + assert c.post_migration_snapshot_pending is True + assert c.submitted_graph is old_graph + + @pytest.mark.asyncio + async def test_finalize_happens_before_s3_write(self): + """Graph finalization must happen when the snapshot advances, in the + same register_snapshot call that writes to S3.""" + c = _mock_coordinator() + c.worker_snapshot_ids = {0: 3, 1: 3} + c.prev_completed_snapshot_id = 2 + c.post_migration_snapshot_pending = True + + new_graph = MagicMock(name="new_graph") + new_graph.name = "test-app" + c._pending_graph = new_graph + c.submitted_graph = MagicMock(name="old_graph") + c.kafka_metadata_producer = AsyncMock() + + pool = MagicMock() + c.register_snapshot(0, 3, {}, {}, 10, 100, pool) + + # Both finalization and S3 write happened + assert c.submitted_graph is new_graph + c.s3_client.put_object.assert_called_once() + + +# --------------------------------------------------------------------------- +# Recovery condition detection +# --------------------------------------------------------------------------- + + +class TestRecoveryConditionDetection: + """Tests that recovery correctly detects both mid-migration crashes AND + post-migration-pre-snapshot crashes. + + Bug caught: if migration completed but no post-migration snapshot was taken, + recovery didn't detect this as migration-related, used the NEW graph with + the OLD snapshot, and lost half the data. + """ + + def test_mid_migration_detected(self): + """migration_in_progress=True with _pending_graph set → needs revert.""" + c = _mock_coordinator() + c._pending_graph = MagicMock() + migration_in_progress = True + + was_migrating = (migration_in_progress and c._pending_graph is not None) or c.post_migration_snapshot_pending + assert was_migrating is True + + def test_post_migration_pre_snapshot_detected(self): + """migration_in_progress=False but post_migration_snapshot_pending=True + → needs revert. This is the bug scenario.""" + c = _mock_coordinator() + c._pending_graph = MagicMock() + c.post_migration_snapshot_pending = True + migration_in_progress = False + + was_migrating = (migration_in_progress and c._pending_graph is not None) or c.post_migration_snapshot_pending + assert was_migrating is True + + def test_no_migration_not_detected(self): + """Normal state — no migration, no pending snapshot.""" + c = _mock_coordinator() + c._pending_graph = None + c.post_migration_snapshot_pending = False + migration_in_progress = False + + was_migrating = (migration_in_progress and c._pending_graph is not None) or c.post_migration_snapshot_pending + assert was_migrating is False + + def test_completed_migration_with_snapshot_not_detected(self): + """Migration completed AND post-migration snapshot taken → no revert needed.""" + c = _mock_coordinator() + c._pending_graph = None # finalize_graph_update() cleared it + c.post_migration_snapshot_pending = False + migration_in_progress = False + + was_migrating = (migration_in_progress and c._pending_graph is not None) or c.post_migration_snapshot_pending + assert was_migrating is False + + def test_saved_pending_graph_preserved_for_retry(self): + """When was_migrating, the pending graph should be saved for auto-retry.""" + c = _mock_coordinator() + new_graph = MagicMock(name="new_graph") + c._pending_graph = new_graph + c.post_migration_snapshot_pending = True + migration_in_progress = False + + was_migrating = (migration_in_progress and c._pending_graph is not None) or c.post_migration_snapshot_pending + saved_pending_graph = c._pending_graph if was_migrating else None + + assert saved_pending_graph is new_graph + + +# --------------------------------------------------------------------------- +# Full migration lifecycle scenarios +# --------------------------------------------------------------------------- + + +class TestMigrationLifecycleScenarios: + """End-to-end scenarios testing the full migration lifecycle through + coordinator metadata, catching the exact bugs that caused 34000 missed + messages in E2E tests.""" + + def _setup_coordinator(self, n_partitions=4, max_parallelism=8): + c = _mock_coordinator() + c.max_operator_parallelism = max_parallelism + c.submitted_graph = _make_graph("test-app", n_partitions, max_parallelism) + c.worker_pool = MagicMock() + c.kafka_metadata_producer = AsyncMock() + c.worker_snapshot_ids = {0: 0, 1: 0, 2: 0} + c.prev_completed_snapshot_id = 0 + return c + + def test_scenario_crash_after_migration_before_snapshot(self): + """Scenario: migration completes, graph NOT finalized yet, crash + before post-migration snapshot. + + Expected: submitted_graph still has OLD layout, revert works, + migration can be retried. + """ + c = self._setup_coordinator(n_partitions=4) + old_graph = c.submitted_graph + new_graph = _make_graph("test-app", n_partitions=8, max_parallelism=8) + + # Simulate: migration sets _pending_graph (update_stateflow_graph) + c._pending_graph = new_graph + + # Simulate: pre-migration snapshot completes (snap 1) + c.pre_migration_snapshot_pending = True + c.worker_snapshot_ids = {0: 1, 1: 1, 2: 1} + pool = MagicMock() + c.register_snapshot(0, 1, {}, {}, 10, 100, pool) + assert c.pre_migration_snapshot_id == 1 + + # Simulate: migration finishes → set post_migration_snapshot_pending + # (NOT calling finalize_graph_update) + c.post_migration_snapshot_pending = True + + # CRASH happens here — before post-migration snapshot + + # Recovery check + migration_in_progress = False # already cleared by _handle_migration_done + was_migrating = (migration_in_progress and c._pending_graph is not None) or c.post_migration_snapshot_pending + assert was_migrating is True + + # submitted_graph is still the OLD graph + assert c.submitted_graph is old_graph + assert c.submitted_graph.nodes["ycsb"].n_partitions == 4 + + # Revert works correctly + c.revert_worker_pool_to_submitted_graph() + # Partitions 4-7 should be shadow + for call_args in c.worker_pool.update_operator.call_args_list[4:]: + _, operator = call_args[0] + assert operator.is_shadow + + # _pending_graph preserved for retry + assert c._pending_graph is new_graph + + @pytest.mark.asyncio + async def test_scenario_crash_after_migration_after_snapshot(self): + """Scenario: migration completes, post-migration snapshot taken, + graph finalized, then crash. + + Expected: submitted_graph has NEW layout, no revert needed. + """ + c = self._setup_coordinator(n_partitions=4) + new_graph = _make_graph("test-app", n_partitions=8, max_parallelism=8) + + # Migration sets _pending_graph + c._pending_graph = new_graph + c.pre_migration_snapshot_pending = True + + # Pre-migration snapshot (snap 1) + c.worker_snapshot_ids = {0: 1, 1: 1, 2: 1} + pool = MagicMock() + c.register_snapshot(0, 1, {}, {}, 10, 100, pool) + + # Migration finishes + c.post_migration_snapshot_pending = True + + # Post-migration snapshot (snap 2) → should finalize + c.worker_snapshot_ids = {0: 2, 1: 2, 2: 2} + c.register_snapshot(0, 2, {}, {}, 20, 200, pool) + + # Graph finalized + assert c.submitted_graph is new_graph + assert c._pending_graph is None + assert c.post_migration_snapshot_pending is False + + # CRASH happens here — after post-migration snapshot + + # Recovery check: no migration recovery needed + migration_in_progress = False + was_migrating = (migration_in_progress and c._pending_graph is not None) or c.post_migration_snapshot_pending + assert was_migrating is False + + def test_scenario_crash_during_migration_phase_a(self): + """Scenario: crash during Phase A (migration in progress, + _pending_graph set, pre-migration snapshot may or may not exist). + + Expected: revert to old layout, retry migration. + """ + c = self._setup_coordinator(n_partitions=4) + old_graph = c.submitted_graph + new_graph = _make_graph("test-app", n_partitions=8, max_parallelism=8) + + # Migration starts + c._pending_graph = new_graph + c.pre_migration_snapshot_pending = True + migration_in_progress = True + + # CRASH during Phase A — no pre-migration snapshot completed + + was_migrating = (migration_in_progress and c._pending_graph is not None) or c.post_migration_snapshot_pending + assert was_migrating is True + + # Old graph preserved + assert c.submitted_graph is old_graph + saved = c._pending_graph + assert saved is new_graph + + @pytest.mark.asyncio + async def test_scenario_normal_migration_no_crash(self): + """Scenario: migration completes normally, snapshot taken, no crash. + + Expected: graph finalized after post-migration snapshot. + """ + c = self._setup_coordinator(n_partitions=4) + new_graph = _make_graph("test-app", n_partitions=8, max_parallelism=8) + + # Step 1: Migration starts + c._pending_graph = new_graph + c.pre_migration_snapshot_pending = True + + # Step 2: Pre-migration snapshot (snap 1) + c.worker_snapshot_ids = {0: 1, 1: 1, 2: 1} + pool = MagicMock() + c.register_snapshot(0, 1, {}, {}, 10, 100, pool) + assert c.pre_migration_snapshot_id == 1 + + # Step 3: Migration finishes + c.post_migration_snapshot_pending = True + assert c.submitted_graph is not new_graph # NOT finalized yet + + # Step 4: Post-migration snapshot (snap 2) + c.worker_snapshot_ids = {0: 2, 1: 2, 2: 2} + c.register_snapshot(0, 2, {}, {}, 20, 200, pool) + + # Step 5: Graph finalized + assert c.submitted_graph is new_graph + assert c._pending_graph is None + assert c.post_migration_snapshot_pending is False + + # Step 6: Subsequent snapshots work normally + c.worker_snapshot_ids = {0: 3, 1: 3, 2: 3} + c.register_snapshot(0, 3, {}, {}, 30, 300, pool) + assert c.submitted_graph is new_graph # still new graph + + def test_scenario_state_cleanup_on_recovery(self): + """Verify all migration state is properly cleared during recovery.""" + c = self._setup_coordinator(n_partitions=4) + new_graph = _make_graph("test-app", n_partitions=8, max_parallelism=8) + + # Set up mid-migration state + c._pending_graph = new_graph + c.pre_migration_snapshot_pending = True + c.pre_migration_snapshot_id = 1 + c.post_migration_snapshot_pending = True + + # Simulate recovery clearing state (mirrors _perform_recovery step 1c) + saved_pending_graph = c._pending_graph + c.pre_migration_snapshot_pending = False + c.post_migration_snapshot_pending = False + c._pending_graph = None + + # All state cleared + assert c.pre_migration_snapshot_pending is False + assert c.post_migration_snapshot_pending is False + assert c._pending_graph is None + # But we saved the graph for retry + assert saved_pending_graph is new_graph + + +# --------------------------------------------------------------------------- +# Worker pool revert on recovery during migration +# --------------------------------------------------------------------------- + + +class TestRevertWorkerPoolToSubmittedGraph: + def _setup_coordinator_with_graph(self, n_partitions, max_parallelism=8): + """Create a coordinator with a submitted graph and worker pool.""" + c = _mock_coordinator() + c.max_operator_parallelism = max_parallelism + + g = StateflowGraph( + "test-app", operator_state_backend=LocalStateBackend.DICT, max_operator_parallelism=max_parallelism + ) + op = Operator("ycsb", n_partitions=n_partitions) + g.add_operators(op) + c.submitted_graph = g + + # Simulate worker pool with update_operator tracking + c.worker_pool = MagicMock() + return c + + def test_revert_restores_shadow_partitions(self): + c = self._setup_coordinator_with_graph(n_partitions=4, max_parallelism=8) + + c.revert_worker_pool_to_submitted_graph() + + # Should have called update_operator for all 8 partitions + assert c.worker_pool.update_operator.call_count == 8 + + # Partitions 0-3 should be active (not shadow) + for call_args in c.worker_pool.update_operator.call_args_list[:4]: + op_part, operator = call_args[0] + assert op_part[0] == "ycsb" + assert not operator.is_shadow + + # Partitions 4-7 should be shadow + for call_args in c.worker_pool.update_operator.call_args_list[4:]: + op_part, operator = call_args[0] + assert op_part[0] == "ycsb" + assert operator.is_shadow + + def test_revert_noop_when_no_graph(self): + c = _mock_coordinator() + c.submitted_graph = None + c.max_operator_parallelism = None + c.worker_pool = MagicMock() + + c.revert_worker_pool_to_submitted_graph() + + c.worker_pool.update_operator.assert_not_called() diff --git a/tests/unit/coordinator/test_worker_pool.py b/tests/unit/coordinator/test_worker_pool.py index fe5f6764..aff29f7c 100644 --- a/tests/unit/coordinator/test_worker_pool.py +++ b/tests/unit/coordinator/test_worker_pool.py @@ -212,6 +212,23 @@ def test_update_operator_replaces_operator_in_place(self): w = pool.peek(1) assert w.assigned_operators[("users", 0)] is op_new + def test_update_operator_skips_unknown_partition(self): + """update_operator should no-op when partition is not mapped to any worker.""" + pool = WorkerPool() + pool.register_worker("127.0.0.1", 5000, 6000) + # ("users", 99) was never scheduled — operator_partition_to_worker has no entry + pool.update_operator(("users", 99), _mock_operator()) # should not raise + + def test_update_operator_skips_removed_worker(self): + """update_operator should no-op when the owning worker was already removed.""" + pool = WorkerPool() + pool.register_worker("127.0.0.1", 5000, 6000) # id=1 + op = _mock_operator() + pool.schedule_operator_partition(("users", 0), op) + # Simulate worker death: remove worker but leave stale mapping + pool.remove_worker(1) + pool.update_operator(("users", 0), _mock_operator()) # should not raise + # --------------------------------------------------------------------------- # WorkerPool — query methods diff --git a/tests/unit/styx_package/test_logging.py b/tests/unit/styx_package/test_logging.py new file mode 100644 index 00000000..422c8e52 --- /dev/null +++ b/tests/unit/styx_package/test_logging.py @@ -0,0 +1,38 @@ +"""Unit tests for styx.common.logging — LOG_LEVEL env var support.""" + +import importlib +import logging as _stdlib_logging +import os + + +def test_default_log_level_is_warning(): + """Without LOG_LEVEL env var, logger should be WARNING.""" + os.environ.pop("LOG_LEVEL", None) + import styx.common.logging as mod + + importlib.reload(mod) + assert mod.logging.level == _stdlib_logging.WARNING + + +def test_log_level_from_env(monkeypatch): + """LOG_LEVEL=DEBUG should set the logger to DEBUG.""" + monkeypatch.setenv("LOG_LEVEL", "DEBUG") + import styx.common.logging as mod + + importlib.reload(mod) + assert mod.logging.level == _stdlib_logging.DEBUG + # Reset + monkeypatch.delenv("LOG_LEVEL", raising=False) + importlib.reload(mod) + + +def test_invalid_log_level_falls_back_to_warning(monkeypatch): + """An unrecognized LOG_LEVEL should fall back to WARNING.""" + monkeypatch.setenv("LOG_LEVEL", "INVALID_LEVEL") + import styx.common.logging as mod + + importlib.reload(mod) + assert mod.logging.level == _stdlib_logging.WARNING + # Reset + monkeypatch.delenv("LOG_LEVEL", raising=False) + importlib.reload(mod) diff --git a/tests/unit/styx_package/test_message_types.py b/tests/unit/styx_package/test_message_types.py index f2ce2e86..16cb8874 100644 --- a/tests/unit/styx_package/test_message_types.py +++ b/tests/unit/styx_package/test_message_types.py @@ -10,7 +10,7 @@ def test_is_int_enum(self): assert issubclass(MessageType, IntEnum) def test_total_count(self): - assert len(MessageType) == 43 + assert len(MessageType) == 44 def test_all_values_unique(self): values = [m.value for m in MessageType] @@ -20,7 +20,7 @@ def test_values_in_expected_range(self): # 25 (AckCache) was removed when the fallback cache was dropped; the # remaining IDs are preserved to keep wire-format stability. values = sorted(m.value for m in MessageType) - assert values == [v for v in range(44) if v != 25] + assert values == [v for v in range(45) if v != 25] def test_known_values(self): assert MessageType.RunFunRemote == 0 @@ -30,6 +30,7 @@ def test_known_values(self): assert MessageType.UpdateExecutionGraph == 41 assert MessageType.RunFunRemoteBatch == 42 assert MessageType.AckBatch == 43 + assert MessageType.SnapMigrationReassign == 44 def test_int_comparison(self): assert MessageType.ClientMsg == 16 diff --git a/tests/unit/worker/test_async_snapshots.py b/tests/unit/worker/test_async_snapshots.py index afc77dda..5fd97ba9 100644 --- a/tests/unit/worker/test_async_snapshots.py +++ b/tests/unit/worker/test_async_snapshots.py @@ -158,12 +158,13 @@ def test_puts_object(self, mock_mk): class TestRetrieveSnapshot: def test_snapshot_id_negative_one_returns_empty(self): s = _snap() - data, tp_off, tp_out_off, epoch, t_counter = s.retrieve_snapshot(-1, []) + data, tp_off, tp_out_off, epoch, t_counter, migration_blob = s.retrieve_snapshot(-1, []) assert data == {} assert tp_off == {} assert tp_out_off == {} assert epoch == 0 assert t_counter == 0 + assert migration_blob is None assert s.snapshot_id == 0 # -1 + 1 @patch("worker.fault_tolerance.async_snapshots._get_s3_client") @@ -176,7 +177,7 @@ def test_retrieve_with_no_files(self, mock_mk): mock_s3.get_paginator.return_value = paginator s = _snap() - data, _tp_off, _tp_out_off, _epoch, _t_counter = s.retrieve_snapshot( + data, _tp_off, _tp_out_off, _epoch, _t_counter, _migration_blob = s.retrieve_snapshot( 0, [("users", 0)], ) @@ -267,3 +268,161 @@ def test_empty_contents(self): s = _snap() result = s._list_bin_keys(mock_s3, "data/op/0/") assert result == [] + + +# --------------------------------------------------------------------------- +# _load_operator_state — tombstone handling +# --------------------------------------------------------------------------- + + +class TestLoadOperatorStateTombstones: + @patch("worker.fault_tolerance.async_snapshots._get_s3_client") + def test_tombstone_removes_key_from_previous_delta(self, mock_mk): + """A None value in a subsequent delta should delete the key.""" + mock_s3 = MagicMock() + mock_mk.return_value = mock_s3 + + s = _snap() + # Simulate two snapshot files for the same partition + paginator = MagicMock() + paginator.paginate.return_value = [ + { + "Contents": [ + {"Key": "data/users/0/1.bin"}, # base snapshot + {"Key": "data/users/0/2.bin"}, # delta with tombstone + ] + } + ] + mock_s3.get_paginator.return_value = paginator + + # First call returns base data, second returns delta with tombstone + call_count = [0] + + def fake_get_zstd_msgpack(_s3_client, _key): + call_count[0] += 1 + if call_count[0] == 1: + return {"key_a": "val_a", "key_b": "val_b"} + return {"key_b": None, "key_c": "val_c"} # tombstone key_b, add key_c + + s._get_zstd_msgpack = fake_get_zstd_msgpack + + data = s._load_operator_state(mock_s3, 2, [("users", 0)]) + assert "key_a" in data[("users", 0)] + assert "key_b" not in data[("users", 0)] # tombstoned + assert data[("users", 0)]["key_c"] == "val_c" + + @patch("worker.fault_tolerance.async_snapshots._get_s3_client") + def test_first_delta_filters_tombstones(self, mock_mk): + """First delta for a partition should filter out None values.""" + mock_s3 = MagicMock() + mock_mk.return_value = mock_s3 + + s = _snap() + paginator = MagicMock() + paginator.paginate.return_value = [{"Contents": [{"Key": "data/users/0/1.bin"}]}] + mock_s3.get_paginator.return_value = paginator + + s._get_zstd_msgpack = lambda _s3_client, _key: {"key_a": "val_a", "key_b": None} + + data = s._load_operator_state(mock_s3, 1, [("users", 0)]) + assert "key_a" in data[("users", 0)] + assert "key_b" not in data[("users", 0)] + + +# --------------------------------------------------------------------------- +# _load_sequencer_state — 5-element and 4-element tuple formats +# --------------------------------------------------------------------------- + + +class TestLoadSequencerStateTupleFormats: + @patch("worker.fault_tolerance.async_snapshots._get_s3_client") + def test_5_element_tuple_with_migration_blob(self, mock_mk): + """Sequencer file with migration checkpoint blob (5 elements).""" + mock_s3 = MagicMock() + mock_mk.return_value = mock_s3 + + s = _snap() + paginator = MagicMock() + paginator.paginate.return_value = [{"Contents": [{"Key": "sequencer/1.bin"}]}] + mock_s3.get_paginator.return_value = paginator + + blob = b"migration_data" + s._get_zstd_msgpack = lambda _s3_client, _key: ( + {("users", 0): 10}, + {("users", 0): 5}, + 3, + 42, + blob, + ) + + tp_off, tp_out, epoch, t_counter, migration_blob = s._load_sequencer_state(mock_s3, 1) + assert tp_off == {("users", 0): 10} + assert tp_out == {("users", 0): 5} + assert epoch == 3 + assert t_counter == 42 + assert migration_blob == blob + + @patch("worker.fault_tolerance.async_snapshots._get_s3_client") + def test_4_element_tuple_legacy_format(self, mock_mk): + """Legacy sequencer file without migration blob (4 elements).""" + mock_s3 = MagicMock() + mock_mk.return_value = mock_s3 + + s = _snap() + paginator = MagicMock() + paginator.paginate.return_value = [{"Contents": [{"Key": "sequencer/1.bin"}]}] + mock_s3.get_paginator.return_value = paginator + + s._get_zstd_msgpack = lambda _s3_client, _key: ( + {("users", 0): 10}, + {("users", 0): 5}, + 3, + 42, + ) + + tp_off, tp_out, epoch, t_counter, migration_blob = s._load_sequencer_state(mock_s3, 1) + assert tp_off == {("users", 0): 10} + assert tp_out == {("users", 0): 5} + assert epoch == 3 + assert t_counter == 42 + assert migration_blob is None + + @patch("worker.fault_tolerance.async_snapshots._get_s3_client") + def test_5_element_with_none_migration_blob(self, mock_mk): + """5-element tuple with None migration blob (normal snapshot after migration cleared).""" + mock_s3 = MagicMock() + mock_mk.return_value = mock_s3 + + s = _snap() + paginator = MagicMock() + paginator.paginate.return_value = [{"Contents": [{"Key": "sequencer/1.bin"}]}] + mock_s3.get_paginator.return_value = paginator + + s._get_zstd_msgpack = lambda _s3_client, _key: ( + {("users", 0): 10}, + {("users", 0): 5}, + 3, + 42, + None, + ) + + *_, migration_blob = s._load_sequencer_state(mock_s3, 1) + assert migration_blob is None + + @patch("worker.fault_tolerance.async_snapshots._get_s3_client") + def test_no_sequencer_files_returns_defaults(self, mock_mk): + """When no sequencer files exist, defaults are returned.""" + mock_s3 = MagicMock() + mock_mk.return_value = mock_s3 + + s = _snap() + paginator = MagicMock() + paginator.paginate.return_value = [{"Contents": []}] + mock_s3.get_paginator.return_value = paginator + + tp_off, tp_out, epoch, t_counter, migration_blob = s._load_sequencer_state(mock_s3, 1) + assert tp_off == {} + assert tp_out == {} + assert epoch == 0 + assert t_counter == 0 + assert migration_blob is None diff --git a/tests/unit/worker/test_kafka_batch_egress.py b/tests/unit/worker/test_kafka_batch_egress.py index 02dfdc09..b4e0ede8 100644 --- a/tests/unit/worker/test_kafka_batch_egress.py +++ b/tests/unit/worker/test_kafka_batch_egress.py @@ -50,8 +50,7 @@ def test_restart_flag(self): class TestClearMessagesSentBeforeRecovery: def test_clears(self): e = _egress() - tp = TopicPartition("users--OUT", 0) - e.messages_sent_before_recovery[tp].add(b"key1") + e.messages_sent_before_recovery.add(b"key1") e.clear_messages_sent_before_recovery() assert len(e.messages_sent_before_recovery) == 0 @@ -88,13 +87,30 @@ async def test_send_skips_already_sent(self): mock_producer = MagicMock() e.kafka_egress_producer = mock_producer - tp = TopicPartition("users--OUT", 0) - e.messages_sent_before_recovery[tp].add(b"key1") + # Flat dedup set — partition-agnostic + e.messages_sent_before_recovery.add(b"key1") await e.send(b"key1", b"value1", "users", 0) assert len(e.batch) == 0 # skipped mock_producer.send.assert_not_called() - assert b"key1" not in e.messages_sent_before_recovery[tp] + assert b"key1" not in e.messages_sent_before_recovery + + @pytest.mark.asyncio + async def test_send_dedup_is_cross_partition(self): + """A key sent to partition 6 pre-crash should be deduped when + replayed to partition 2 after recovery.""" + e = _egress() + e.started.set() + mock_producer = MagicMock() + e.kafka_egress_producer = mock_producer + + # Key was originally sent on partition 6 (post-migration) + e.messages_sent_before_recovery.add(b"req42") + + # Replay sends to partition 2 (pre-migration layout) + await e.send(b"req42", b"val", "users", 2) + assert len(e.batch) == 0 # deduped cross-partition + mock_producer.send.assert_not_called() # --------------------------------------------------------------------------- @@ -186,6 +202,7 @@ def test_adds_keys(self): msg2.key = b"key2" e.process_messages_sent_before_recovery([msg1, msg2], current_offsets, all_done) - assert b"key1" in e.messages_sent_before_recovery[tp] - assert b"key2" in e.messages_sent_before_recovery[tp] + # Flat set now — keys from any partition are in the same set + assert b"key1" in e.messages_sent_before_recovery + assert b"key2" in e.messages_sent_before_recovery assert all_done[tp] is True # msg2 offset >= 10 - 1 diff --git a/tests/unit/worker/test_kafka_egress_coverage.py b/tests/unit/worker/test_kafka_egress_coverage.py index 873328a3..97e76b5b 100644 --- a/tests/unit/worker/test_kafka_egress_coverage.py +++ b/tests/unit/worker/test_kafka_egress_coverage.py @@ -1,13 +1,13 @@ """Additional coverage tests for worker/egress/styx_kafka_batch_egress.py Covers: send with unstarted event, send_immediate dedup path, -send_message_to_topic, stop. +send_message_to_topic, stop, dedup counters, run_dedup_scan_and_mark_started, +dedup_output_offsets constructor parameter. """ import asyncio from unittest.mock import AsyncMock, MagicMock -from aiokafka import TopicPartition import pytest from worker.egress.styx_kafka_batch_egress import StyxKafkaBatchEgress @@ -17,10 +17,10 @@ # --------------------------------------------------------------------------- -def _egress(offsets=None, restart=False): +def _egress(offsets=None, restart=False, dedup_offsets=None): if offsets is None: offsets = {("users", 0): -1, ("users", 1): -1} - return StyxKafkaBatchEgress(offsets, restart_after_failure=restart) + return StyxKafkaBatchEgress(offsets, restart_after_failure=restart, dedup_output_offsets=dedup_offsets) # --------------------------------------------------------------------------- @@ -75,12 +75,12 @@ async def test_send_immediate_dedup(self): mock_producer = MagicMock() e.kafka_egress_producer = mock_producer - tp = TopicPartition("users--OUT", 0) - e.messages_sent_before_recovery[tp].add(b"key1") + # Flat dedup set — partition-agnostic + e.messages_sent_before_recovery.add(b"key1") await e.send_immediate(b"key1", b"value1", "users", 0) mock_producer.send_and_wait.assert_not_called() - assert b"key1" not in e.messages_sent_before_recovery[tp] + assert b"key1" not in e.messages_sent_before_recovery # --------------------------------------------------------------------------- @@ -161,3 +161,185 @@ async def test_batch_multiple_partitions(self): assert e.topic_partition_output_offsets[("users", 0)] == 10 assert e.topic_partition_output_offsets[("users", 1)] == 20 assert e.batch == [] + + +# --------------------------------------------------------------------------- +# Constructor — dedup_output_offsets parameter +# --------------------------------------------------------------------------- + + +class TestDedupOutputOffsets: + def test_default_dedup_offsets_is_empty_dict(self): + e = _egress() + assert e.dedup_output_offsets == {} + + def test_custom_dedup_offsets_stored(self): + offsets = {("users", 0): 5, ("orders", 1): 10} + e = _egress(dedup_offsets=offsets) + assert e.dedup_output_offsets == offsets + + def test_dedup_counters_start_at_zero(self): + e = _egress() + assert e._dedup_suppressed == 0 + assert e._dedup_sent == 0 + + +# --------------------------------------------------------------------------- +# start — non-recovery sets started immediately +# --------------------------------------------------------------------------- + + +class TestStartNonRecovery: + @pytest.mark.asyncio + async def test_start_non_recovery_sets_started(self): + e = _egress(restart=False) + mock_producer = MagicMock() + mock_producer.start = AsyncMock() + # Patch AIOKafkaProducer to return our mock + e.kafka_egress_producer = mock_producer + # Simulate the start logic without real Kafka + # Just call the post-start code path + if not e.restart_after_failure: + e.started.set() + assert e.started.is_set() + + @pytest.mark.asyncio + async def test_start_recovery_does_not_set_started(self): + e = _egress(restart=True) + # Recovery path: started should NOT be set in start() + if not e.restart_after_failure: + e.started.set() + assert not e.started.is_set() + + +# --------------------------------------------------------------------------- +# run_dedup_scan_and_mark_started +# --------------------------------------------------------------------------- + + +class TestRunDedupScanAndMarkStarted: + @pytest.mark.asyncio + async def test_non_recovery_just_sets_started(self): + """When restart_after_failure=False, run_dedup_scan_and_mark_started just sets started.""" + e = _egress(restart=False) + await e.run_dedup_scan_and_mark_started() + assert e.started.is_set() + + @pytest.mark.asyncio + async def test_recovery_runs_dedup_scan(self): + """When restart_after_failure=True, dedup scan runs then started is set.""" + e = _egress(restart=True, dedup_offsets={}) + # Mock get_messages_sent_before_recovery (it requires Kafka) + e.get_messages_sent_before_recovery = AsyncMock() + await e.run_dedup_scan_and_mark_started() + e.get_messages_sent_before_recovery.assert_called_once() + assert e.started.is_set() + + +# --------------------------------------------------------------------------- +# stop — dedup stats logging +# --------------------------------------------------------------------------- + + +class TestStopDedupStats: + @pytest.mark.asyncio + async def test_stop_logs_dedup_stats_when_nonzero(self): + e = _egress() + e.worker_id = 0 + e._dedup_suppressed = 3 + e._dedup_sent = 10 + mock_producer = MagicMock() + mock_producer.stop = AsyncMock() + e.kafka_egress_producer = mock_producer + + await e.stop() + mock_producer.stop.assert_called_once() + + @pytest.mark.asyncio + async def test_stop_skips_log_when_zero_counters(self): + e = _egress() + e.worker_id = 0 + e._dedup_suppressed = 0 + e._dedup_sent = 0 + mock_producer = MagicMock() + mock_producer.stop = AsyncMock() + e.kafka_egress_producer = mock_producer + + await e.stop() + mock_producer.stop.assert_called_once() + + +# --------------------------------------------------------------------------- +# send — dedup counters +# --------------------------------------------------------------------------- + + +class TestSendDedupCounters: + @pytest.mark.asyncio + async def test_send_increments_dedup_sent(self): + e = _egress() + e.started.set() + mock_producer = MagicMock() + future = asyncio.Future() + future.set_result(MagicMock()) + mock_producer.send = AsyncMock(return_value=future) + e.kafka_egress_producer = mock_producer + + await e.send(b"key1", b"value1", "users", 0) + assert e._dedup_sent == 1 + + @pytest.mark.asyncio + async def test_send_increments_dedup_suppressed(self): + e = _egress() + e.started.set() + e.kafka_egress_producer = MagicMock() + e.messages_sent_before_recovery.add(b"key1") + + await e.send(b"key1", b"value1", "users", 0) + assert e._dedup_suppressed == 1 + assert e._dedup_sent == 0 + + +# --------------------------------------------------------------------------- +# send_immediate — dedup counters +# --------------------------------------------------------------------------- + + +class TestSendImmediateDedupCounters: + @pytest.mark.asyncio + async def test_send_immediate_increments_dedup_sent(self): + e = _egress() + e.started.set() + meta = MagicMock() + meta.offset = 10 + mock_producer = MagicMock() + mock_producer.send_and_wait = AsyncMock(return_value=meta) + e.kafka_egress_producer = mock_producer + + await e.send_immediate(b"key1", b"value1", "users", 0) + assert e._dedup_sent == 1 + + @pytest.mark.asyncio + async def test_send_immediate_increments_dedup_suppressed(self): + e = _egress() + e.started.set() + e.kafka_egress_producer = MagicMock() + e.messages_sent_before_recovery.add(b"key1") + + await e.send_immediate(b"key1", b"value1", "users", 0) + assert e._dedup_suppressed == 1 + assert e._dedup_sent == 0 + + +# --------------------------------------------------------------------------- +# get_messages_sent_before_recovery — empty partitions early return +# --------------------------------------------------------------------------- + + +class TestGetMessagesSentBeforeRecoveryEmpty: + @pytest.mark.asyncio + async def test_empty_dedup_offsets_returns_immediately(self): + e = _egress(dedup_offsets={}) + await e.get_messages_sent_before_recovery() + # Should return without creating a consumer + assert len(e.messages_sent_before_recovery) == 0 diff --git a/tests/unit/worker/test_kafka_ingress.py b/tests/unit/worker/test_kafka_ingress.py index 0e0b3475..28075bec 100644 --- a/tests/unit/worker/test_kafka_ingress.py +++ b/tests/unit/worker/test_kafka_ingress.py @@ -1,7 +1,10 @@ """Unit tests for worker/ingress/styx_kafka_ingress.py""" -from unittest.mock import MagicMock +import asyncio +from unittest.mock import AsyncMock, MagicMock, patch +from aiokafka.errors import KafkaConnectionError, UnknownTopicOrPartitionError +import pytest from styx.common.message_types import MessageType from styx.common.run_func_payload import RunFuncPayload @@ -142,3 +145,261 @@ async def test_invalid_message_type(self): msg = _make_msg() ingress.handle_message_from_kafka(msg) sequencer.sequence.assert_not_called() + + +# --------------------------------------------------------------------------- +# Shadow partition redirect (post-recovery with client partitioner mismatch) +# --------------------------------------------------------------------------- + + +def _make_ingress_with_shadow_partitions(): + """Create an ingress with 4 active + 4 shadow partitions. + + Simulates post-recovery state where the system has 4 active partitions + but the client sends to 8 partitions. Shadow partitions (4-7) should + redirect messages to the correct active partition (0-3). + """ + networking = MagicMock() + networking.get_msg_type = MagicMock(return_value=MessageType.ClientMsg) + networking.in_the_same_network = MagicMock(return_value=True) + + sequencer = MagicMock() + state = MagicMock() + state.exists = MagicMock(return_value=False) # shadow partitions have no state + + dns = { + "ycsb": { + 0: ("host1", 5000, 6000), + 1: ("host2", 5001, 6001), + 2: ("host3", 5002, 6002), + 3: ("host3", 5003, 6003), + 4: ("host1", 5000, 6000), # shadow + 5: ("host2", 5001, 6001), # shadow + 6: ("host3", 5002, 6002), # shadow + 7: ("host3", 5003, 6003), # shadow + }, + } + + registered_operators = {} + for part in range(8): + op = MagicMock() + # 4-partition partitioner: always maps to 0-3 + op.which_partition = MagicMock(side_effect=lambda key, _p=part: hash(key) % 4) + op.dns = dns + registered_operators[("ycsb", part)] = op + + ingress = StyxKafkaIngress( + networking=networking, + sequencer=sequencer, + state=state, + registered_operators=registered_operators, + worker_id=0, + kafka_url="localhost:9092", + epoch_interval_ms=1, + sequence_max_size=1000, + ) + return ingress, networking, sequencer, state + + +class TestShadowPartitionRedirect: + """Tests that messages arriving on shadow partitions (4-7) are correctly + redirected to active partitions (0-3) after recovery. + + Bug scenario: client updates its partitioner to 8 partitions during + migration, worker crashes, recovery restores 4-partition layout. Client + continues sending to partitions 4-7 which are shadow. Without proper + redirect, ~50% of messages are lost. + """ + + def test_message_on_shadow_partition_redirected_locally(self): + """Message sent to shadow partition 4 should be redirected to the + correct active partition and sequenced locally.""" + ingress, networking, sequencer, _state = _make_ingress_with_shadow_partitions() + # Client sends to partition 4 (shadow) with key "test_key" + networking.decode_message.return_value = ("ycsb", "test_key", "read", (), 4) + # which_partition returns an active partition (0-3) + ingress.registered_operators[("ycsb", 4)].which_partition = MagicMock(return_value=0) + + msg = _make_msg(partition=4) + ingress.handle_message_from_kafka(msg) + + sequencer.sequence.assert_called_once() + payload = sequencer.sequence.call_args[0][0] + assert payload.partition == 0 # redirected to active partition + + def test_message_on_shadow_partition_redirected_remotely(self): + """Message on shadow partition redirected to remote worker via + WrongPartitionRequest.""" + ingress, networking, sequencer, _state = _make_ingress_with_shadow_partitions() + networking.decode_message.return_value = ("ycsb", "test_key", "read", (), 5) + ingress.registered_operators[("ycsb", 5)].which_partition = MagicMock(return_value=1) + networking.in_the_same_network.return_value = False + + msg = _make_msg(partition=5) + ingress.handle_message_from_kafka(msg) + + sequencer.sequence.assert_not_called() + assert len(ingress.send_message_tasks) == 1 + + def test_message_on_active_partition_not_redirected(self): + """Message on active partition 0 with existing key is processed normally.""" + ingress, networking, sequencer, state = _make_ingress_with_shadow_partitions() + state.exists.return_value = True + networking.decode_message.return_value = ("ycsb", "test_key", "read", (), 0) + + msg = _make_msg(partition=0) + ingress.handle_message_from_kafka(msg) + + sequencer.sequence.assert_called_once() + payload = sequencer.sequence.call_args[0][0] + assert payload.partition == 0 + + def test_redirect_preserves_kafka_offset_metadata(self): + """Redirected messages must preserve kafka_offset and + kafka_ingress_partition for offset tracking. + + Bug: if these aren't preserved, _advance_offsets can't track + progress and messages are re-consumed after migration. + """ + ingress, networking, sequencer, _state = _make_ingress_with_shadow_partitions() + networking.decode_message.return_value = ("ycsb", "test_key", "read", (), 6) + ingress.registered_operators[("ycsb", 6)].which_partition = MagicMock(return_value=2) + + msg = _make_msg(partition=6, offset=42) + ingress.handle_message_from_kafka(msg) + + payload = sequencer.sequence.call_args[0][0] + assert payload.kafka_offset == 42 + assert payload.kafka_ingress_partition == 6 + + def test_all_shadow_partitions_redirect(self): + """Verify all shadow partitions (4-7) redirect correctly.""" + ingress, networking, sequencer, _state = _make_ingress_with_shadow_partitions() + + for shadow_part in range(4, 8): + sequencer.reset_mock() + active_part = shadow_part % 4 + networking.decode_message.return_value = ( + "ycsb", + f"key_{shadow_part}", + "read", + (), + shadow_part, + ) + ingress.registered_operators[("ycsb", shadow_part)].which_partition = MagicMock(return_value=active_part) + + msg = _make_msg(partition=shadow_part) + ingress.handle_message_from_kafka(msg) + + sequencer.sequence.assert_called_once() + payload = sequencer.sequence.call_args[0][0] + assert payload.partition == active_part, f"Shadow partition {shadow_part} should redirect to {active_part}" + + +# --------------------------------------------------------------------------- +# Kafka consumer startup exception handling +# --------------------------------------------------------------------------- + + +class TestKafkaConsumerStartupExceptionHandling: + """Tests that the Kafka consumer startup loop correctly catches both + UnknownTopicOrPartitionError and KafkaConnectionError. + + Bug caught: Python 2 syntax `except A, B:` (which means `except A as B:` + in Python 3) only catches A and silently shadows the name B. This caused + KafkaConnectionError to crash the ingress during recovery, killing all + message consumption for that worker. + """ + + @pytest.mark.asyncio + async def test_kafka_connection_error_is_caught_and_retried(self): + """KafkaConnectionError during consumer.start() should be caught + and retried, not crash the ingress.""" + ingress, *_ = _make_ingress() + mock_consumer = AsyncMock() + # First call raises KafkaConnectionError, second succeeds + mock_consumer.start = AsyncMock( + side_effect=[KafkaConnectionError(), None], + ) + # getmany raises CancelledError to exit the consumer loop cleanly + mock_consumer.getmany = AsyncMock(side_effect=asyncio.CancelledError()) + mock_consumer.stop = AsyncMock() + + with ( + patch("worker.ingress.styx_kafka_ingress.AIOKafkaConsumer", return_value=mock_consumer), + patch("worker.ingress.styx_kafka_ingress.asyncio.sleep", new_callable=AsyncMock), + pytest.raises(asyncio.CancelledError), + ): + await ingress.start_kafka_consumer([], {}) + + # Consumer should have been started twice (retry after error) + assert mock_consumer.start.call_count == 2 + + @pytest.mark.asyncio + async def test_unknown_topic_error_is_caught_and_retried(self): + """UnknownTopicOrPartitionError during consumer.start() should be + caught and retried.""" + ingress, *_ = _make_ingress() + mock_consumer = AsyncMock() + mock_consumer.start = AsyncMock( + side_effect=[UnknownTopicOrPartitionError(), None], + ) + mock_consumer.getmany = AsyncMock(side_effect=asyncio.CancelledError()) + mock_consumer.stop = AsyncMock() + + with ( + patch("worker.ingress.styx_kafka_ingress.AIOKafkaConsumer", return_value=mock_consumer), + patch("worker.ingress.styx_kafka_ingress.asyncio.sleep", new_callable=AsyncMock), + pytest.raises(asyncio.CancelledError), + ): + await ingress.start_kafka_consumer([], {}) + + assert mock_consumer.start.call_count == 2 + + @pytest.mark.asyncio + async def test_other_exceptions_propagate(self): + """Non-Kafka exceptions should propagate, not be silently swallowed.""" + ingress, *_ = _make_ingress() + mock_consumer = AsyncMock() + mock_consumer.start = AsyncMock(side_effect=ValueError("unexpected")) + + with ( + patch("worker.ingress.styx_kafka_ingress.AIOKafkaConsumer", return_value=mock_consumer), + pytest.raises(ValueError, match="unexpected"), + ): + await ingress.start_kafka_consumer([], {}) + + def test_except_clause_catches_both_exception_types(self): + """Verify the except clause syntax is correct (not Python 2 style). + + Python 2: `except A, B:` catches A and binds to B + Python 3: `except A, B:` means `except A as B:` — only catches A! + Correct: `except (A, B):` catches both A and B + + This test inspects the source to prevent regressions from formatters + or linters that might not catch this subtle syntax difference. + """ + import ast + import inspect + import textwrap + + source = textwrap.dedent(inspect.getsource(StyxKafkaIngress.start_kafka_consumer)) + tree = ast.parse(source) + + # Find all except handlers in the function + for node in ast.walk(tree): + if ( + isinstance(node, ast.ExceptHandler) + and node.type is not None + and isinstance(node.type, ast.Name) + and node.type.id == "UnknownTopicOrPartitionError" + ): + # This means: except UnknownTopicOrPartitionError as + # i.e., KafkaConnectionError is NOT caught — this is the bug + pytest.fail( + "except clause uses Python 2 syntax " + "'except UnknownTopicOrPartitionError, KafkaConnectionError:' " + "which only catches UnknownTopicOrPartitionError. " + "Use 'except (UnknownTopicOrPartitionError, KafkaConnectionError):' " + "to catch both.", + ) diff --git a/tests/unit/worker/test_migration_fault_tolerance.py b/tests/unit/worker/test_migration_fault_tolerance.py new file mode 100644 index 00000000..cc2ce33d --- /dev/null +++ b/tests/unit/worker/test_migration_fault_tolerance.py @@ -0,0 +1,210 @@ +"""Unit tests for migration fault tolerance features. + +Covers: +- delta_map writes during migration (set_data_from_migration, set_batch_data_from_migration) +- tombstone writes on source-side key removal (get_key_to_migrate, get_async_migrate_batch) +- tombstone handling during snapshot recovery (_load_operator_state) +- migration_reassign_delta_maps preserves existing deltas +""" + +from unittest.mock import MagicMock, patch + +from worker.operator_state.aria.in_memory_state import InMemoryOperatorState + +OP = "op" +PART_0 = 0 +PART_1 = 1 +OP_PART_0 = (OP, PART_0) +OP_PART_1 = (OP, PART_1) + + +def _state(*partitions): + parts = set(partitions) if partitions else {OP_PART_0} + return InMemoryOperatorState(parts) + + +# --------------------------------------------------------------------------- +# Step 1: set_data_from_migration writes to delta_map +# --------------------------------------------------------------------------- + + +class TestSetDataFromMigrationDeltaMap: + def test_single_key_written_to_delta_map(self): + s = _state(OP_PART_0, OP_PART_1) + s.set_data_from_migration(OP_PART_1, "k1", "v1") + assert s.delta_map[OP_PART_1]["k1"] == "v1" + + def test_none_data_not_written_to_delta_map(self): + s = _state(OP_PART_0, OP_PART_1) + s.set_data_from_migration(OP_PART_1, "k1", None) + assert "k1" not in s.delta_map[OP_PART_1] + + def test_data_written_to_both_data_and_delta_map(self): + s = _state(OP_PART_0, OP_PART_1) + s.set_data_from_migration(OP_PART_1, "k1", {"field": 42}) + assert s.data[OP_PART_1]["k1"] == {"field": 42} + assert s.delta_map[OP_PART_1]["k1"] == {"field": 42} + + +class TestSetBatchDataFromMigrationDeltaMap: + def test_batch_written_to_delta_map(self): + s = _state(OP_PART_0, OP_PART_1) + batch = {"k1": "v1", "k2": "v2", "k3": "v3"} + s.set_batch_data_from_migration(OP_PART_1, batch) + assert s.delta_map[OP_PART_1] == batch + + def test_batch_merges_with_existing_delta_map(self): + s = _state(OP_PART_0, OP_PART_1) + s.delta_map[OP_PART_1]["existing"] = "val" + s.set_batch_data_from_migration(OP_PART_1, {"k1": "v1"}) + assert s.delta_map[OP_PART_1] == {"existing": "val", "k1": "v1"} + + +class TestMigrateWithinSameWorkerDeltaMap: + def test_migrate_within_same_worker_writes_delta_map(self): + s = _state(OP_PART_0, OP_PART_1) + s.data[OP_PART_0]["k1"] = "v1" + s.migrate_within_the_same_worker(OP, PART_1, "k1", PART_0) + # Destination gets the value in delta_map + assert s.delta_map[OP_PART_1]["k1"] == "v1" + # Source gets a tombstone in delta_map + assert s.delta_map[OP_PART_0]["k1"] is None + + +# --------------------------------------------------------------------------- +# Step 2: Source-side key removal writes tombstones +# --------------------------------------------------------------------------- + + +class TestGetKeyToMigrateTombstone: + def test_tombstone_written_to_delta_map(self): + s = _state(OP_PART_0, OP_PART_1) + s.data[OP_PART_0]["k1"] = "v1" + result = s.get_key_to_migrate(OP_PART_1, "k1", PART_0) + assert result == "v1" + assert s.delta_map[OP_PART_0]["k1"] is None + + def test_no_tombstone_when_key_already_gone(self): + s = _state(OP_PART_0, OP_PART_1) + # key not in data + result = s.get_key_to_migrate(OP_PART_1, "k1", PART_0) + assert result is None + assert "k1" not in s.delta_map[OP_PART_0] + + +class TestGetAsyncMigrateBatchTombstone: + def test_tombstones_written_for_migrated_keys(self): + s = _state(OP_PART_0, OP_PART_1) + s.data[OP_PART_0] = {"k1": "v1", "k2": "v2", "k3": "v3"} + s.keys_to_send = {OP_PART_0: {("k1", PART_1), ("k2", PART_1), ("k3", PART_1)}} + batch = s.get_async_migrate_batch(batch_size=10) + # All three keys should have tombstones in source delta_map + for key in ["k1", "k2", "k3"]: + assert s.delta_map[OP_PART_0][key] is None + # The batch should contain the values for the destination + assert len(batch[(OP, PART_1)]) == 3 + + def test_no_tombstone_for_already_deleted_key(self): + s = _state(OP_PART_0, OP_PART_1) + s.data[OP_PART_0] = {} # key not present + s.keys_to_send = {OP_PART_0: {("k1", PART_1)}} + s.get_async_migrate_batch(batch_size=10) + # No tombstone because key was already gone + assert "k1" not in s.delta_map[OP_PART_0] + + +# --------------------------------------------------------------------------- +# Step 3: Tombstone handling in snapshot recovery +# --------------------------------------------------------------------------- + + +class TestLoadOperatorStateTombstones: + def test_tombstone_removes_key_from_accumulated_state(self): + from worker.fault_tolerance.async_snapshots import AsyncSnapshotsS3 + + snap = AsyncSnapshotsS3.__new__(AsyncSnapshotsS3) + snap.worker_id = 0 + + # Simulate S3 with two deltas: first adds key, second tombstones it + delta_1 = {"k1": "v1", "k2": "v2"} + delta_2 = {"k1": None} # tombstone + + mock_s3 = MagicMock() + + with ( + patch.object(snap, "_iter_snapshot_files", return_value=[(1, "d/op/0/1.bin"), (2, "d/op/0/2.bin")]), + patch.object(snap, "_get_zstd_msgpack", side_effect=[delta_1, delta_2]), + ): + data = snap._load_operator_state(mock_s3, 2, [("op", 0)]) + + assert ("op", 0) in data + assert "k1" not in data[("op", 0)] # tombstoned + assert data[("op", 0)]["k2"] == "v2" # preserved + + def test_tombstone_in_first_delta_filtered_out(self): + from worker.fault_tolerance.async_snapshots import AsyncSnapshotsS3 + + snap = AsyncSnapshotsS3.__new__(AsyncSnapshotsS3) + snap.worker_id = 0 + + delta_1 = {"k1": None, "k2": "v2"} # tombstone in first delta + + mock_s3 = MagicMock() + + with ( + patch.object(snap, "_iter_snapshot_files", return_value=[(1, "d/op/0/1.bin")]), + patch.object(snap, "_get_zstd_msgpack", return_value=delta_1), + ): + data = snap._load_operator_state(mock_s3, 1, [("op", 0)]) + + assert "k1" not in data[("op", 0)] + assert data[("op", 0)]["k2"] == "v2" + + +# --------------------------------------------------------------------------- +# Step 5: migration_reassign_delta_maps preserves deltas +# --------------------------------------------------------------------------- + + +class TestMigrationReassignDeltaMaps: + def test_preserves_existing_deltas(self): + from worker.async_snapshotting import AsyncSnapshottingProcess + + proc = AsyncSnapshottingProcess.__new__(AsyncSnapshottingProcess) + proc.delta_maps = { + ("op", 0): {"k1": "v1"}, + ("op", 1): {"k2": "v2"}, + } + + # Reassign: keep partition 0, add partition 2, drop partition 1 + proc.migration_reassign_delta_maps([("op", 0), ("op", 2)]) + + assert proc.delta_maps[("op", 0)] == {"k1": "v1"} # preserved + assert proc.delta_maps[("op", 2)] == {} # new, empty + assert ("op", 1) not in proc.delta_maps # removed + + def test_no_change_when_same_partitions(self): + from worker.async_snapshotting import AsyncSnapshottingProcess + + proc = AsyncSnapshottingProcess.__new__(AsyncSnapshottingProcess) + proc.delta_maps = { + ("op", 0): {"k1": "v1"}, + ("op", 1): {"k2": "v2"}, + } + + proc.migration_reassign_delta_maps([("op", 0), ("op", 1)]) + + assert proc.delta_maps[("op", 0)] == {"k1": "v1"} + assert proc.delta_maps[("op", 1)] == {"k2": "v2"} + + def test_all_new_partitions(self): + from worker.async_snapshotting import AsyncSnapshottingProcess + + proc = AsyncSnapshottingProcess.__new__(AsyncSnapshottingProcess) + proc.delta_maps = {("op", 0): {"k1": "v1"}} + + proc.migration_reassign_delta_maps([("op", 2), ("op", 3)]) + + assert ("op", 0) not in proc.delta_maps + assert proc.delta_maps[("op", 2)] == {} + assert proc.delta_maps[("op", 3)] == {} diff --git a/worker/async_snapshotting.py b/worker/async_snapshotting.py index 0ca6878c..68d32f4f 100644 --- a/worker/async_snapshotting.py +++ b/worker/async_snapshotting.py @@ -63,6 +63,17 @@ def clear_delta_maps(self) -> None: def init_delta_maps(self, assigned_partitions: list) -> None: self.delta_maps = {(op_part[0], op_part[1]): {} for op_part in assigned_partitions} + def migration_reassign_delta_maps(self, new_assigned_partitions: list) -> None: + new_keys = {(op_part[0], op_part[1]) for op_part in new_assigned_partitions} + # Add empty entries for newly assigned partitions + for key in new_keys: + if key not in self.delta_maps: + self.delta_maps[key] = {} + # Remove entries for partitions no longer assigned + for key in list(self.delta_maps.keys()): + if key not in new_keys: + del self.delta_maps[key] + def take_snapshot(self, metadata: tuple) -> None: loop = asyncio.get_running_loop() ( @@ -110,6 +121,13 @@ def snapshotting_controller(self, data: bytes) -> None: ) if snapshot_id != -1: self.async_snapshots.set_snapshot_id(snapshot_id) + case MessageType.SnapMigrationReassign: + (new_assigned_partitions,) = NetworkingManager.decode_message(data) + logging.warning("[SN_PROC] Migration reassign: preserving deltas across partition change") + self.migration_reassign_delta_maps(new_assigned_partitions) + self.async_snapshots.update_n_assigned_partitions( + len(new_assigned_partitions), + ) case _: logging.error( f"Worker Service: Non supported command message type: {message_type}", diff --git a/worker/egress/styx_kafka_batch_egress.py b/worker/egress/styx_kafka_batch_egress.py index b06ef2c2..d6fab597 100644 --- a/worker/egress/styx_kafka_batch_egress.py +++ b/worker/egress/styx_kafka_batch_egress.py @@ -1,5 +1,4 @@ import asyncio -from collections import defaultdict import os from typing import TYPE_CHECKING import uuid @@ -25,26 +24,31 @@ def __init__( self, topic_partition_output_offsets: dict[OperatorPartition, int], restart_after_failure: bool = False, + dedup_output_offsets: dict[OperatorPartition, int] | None = None, ) -> None: self.kafka_egress_producer: AIOKafkaProducer | None = None - # operator: partition: output offset + # operator: partition: output offset — only for THIS worker's partitions self.topic_partition_output_offsets: dict[OperatorPartition, int] = topic_partition_output_offsets - # (operator, partition): replied request_ids - self.messages_sent_before_recovery: dict[TopicPartition, set] = defaultdict(set) + # Flat set of request_ids sent before recovery (partition-agnostic). + # Cross-partition dedup is needed because migration can change which + # output partition a response is written to. + self.messages_sent_before_recovery: set = set() self.restart_after_failure = restart_after_failure + # ALL output offsets from the global snapshot — used for dedup scanning + # so that every worker scans ALL output partitions (not just its own). + # This is essential because after recovery a request may be replayed on + # a different worker than the one that originally sent the output. + self.dedup_output_offsets: dict[OperatorPartition, int] = dedup_output_offsets or {} self.batch: list = [] self.started: asyncio.Event = asyncio.Event() + # --- debug counters --- + self._dedup_suppressed: int = 0 + self._dedup_sent: int = 0 def clear_messages_sent_before_recovery(self) -> None: - self.messages_sent_before_recovery: dict[TopicPartition, set] = defaultdict(set) + self.messages_sent_before_recovery: set = set() async def start(self, worker_id: int) -> None: - if self.restart_after_failure: - logging.warning( - f"Getting messages sent before recovery: {self.topic_partition_output_offsets}", - ) - await self.get_messages_sent_before_recovery() - logging.warning("Got messages sent before recovery") self.worker_id = worker_id # Idempotent (implies acks="all"): this producer writes client responses and # the sequencer WAL, which recovery relies on. With acks=1 a leader failover @@ -62,10 +66,38 @@ async def start(self, worker_id: int) -> None: logging.info("Waiting for Kafka") continue break + if not self.restart_after_failure: + # Non-recovery path: mark started immediately. + # Recovery path: started is set later by run_dedup_scan_and_mark_started() + # after the ReadyAfterRecovery barrier ensures all workers flushed. + self.started.set() + + async def run_dedup_scan_and_mark_started(self) -> None: + """Run the dedup scan (if needed) and then mark the egress as started. + + This must be called AFTER all workers have stopped their old protocols + (i.e. after the ReadyAfterRecovery barrier) so that the dedup scan + captures all messages flushed by surviving workers. + """ + if self.restart_after_failure: + logging.warning( + f"[DEDUP] Getting messages sent before recovery. " + f"Scanning ALL output partitions: {self.dedup_output_offsets}", + ) + await self.get_messages_sent_before_recovery() + logging.warning( + f"[DEDUP] Done scanning. Dedup set size: {len(self.messages_sent_before_recovery)}", + ) self.started.set() async def stop(self) -> None: await self.send_batch() + if self._dedup_suppressed > 0 or self._dedup_sent > 0: + logging.warning( + f"[DEDUP] W{self.worker_id} final stats: " + f"suppressed={self._dedup_suppressed}, sent={self._dedup_sent}, " + f"remaining_dedup_set_size={len(self.messages_sent_before_recovery)}", + ) await self.kafka_egress_producer.stop() async def send( @@ -77,8 +109,7 @@ async def send( ) -> None: if not self.started.is_set(): await self.started.wait() - tp = TopicPartition(operator_name + "--OUT", partition) - if key not in self.messages_sent_before_recovery[tp]: + if key not in self.messages_sent_before_recovery: self.batch.append( await self.kafka_egress_producer.send( operator_name + "--OUT", @@ -87,8 +118,14 @@ async def send( partition=partition, ), ) + self._dedup_sent += 1 else: - self.messages_sent_before_recovery[tp].remove(key) + self.messages_sent_before_recovery.discard(key) + self._dedup_suppressed += 1 + logging.warning( + f"[DEDUP] W{getattr(self, 'worker_id', '?')} SUPPRESSED key={key!r} " + f"target_partition={operator_name}--OUT/{partition}", + ) async def send_message_to_topic( self, @@ -112,8 +149,7 @@ async def send_immediate( operator_name: str, partition: int, ) -> None: - tp = TopicPartition(operator_name + "--OUT", partition) - if key not in self.messages_sent_before_recovery[tp]: + if key not in self.messages_sent_before_recovery: res: RecordMetadata = await self.kafka_egress_producer.send_and_wait( operator_name + "--OUT", key=key, @@ -124,8 +160,14 @@ async def send_immediate( self.topic_partition_output_offsets[operator_name, partition], res.offset, ) + self._dedup_sent += 1 else: - self.messages_sent_before_recovery[tp].remove(key) + self.messages_sent_before_recovery.discard(key) + self._dedup_suppressed += 1 + logging.warning( + f"[DEDUP] W{getattr(self, 'worker_id', '?')} SUPPRESSED (immediate) key={key!r} " + f"target_partition={operator_name}--OUT/{partition}", + ) async def send_batch(self) -> None: if self.batch: @@ -144,25 +186,24 @@ async def send_batch(self) -> None: ) async def get_messages_sent_before_recovery(self) -> None: + # Scan ALL output partitions from the global snapshot (not just this + # worker's). This ensures cross-worker dedup: a request originally + # processed on worker A may be replayed on worker B after recovery. + all_topic_partitions = [TopicPartition(op_name + "--OUT", part) for op_name, part in self.dedup_output_offsets] + if not all_topic_partitions: + return kafka_output_consumer = AIOKafkaConsumer( bootstrap_servers=[KAFKA_URL], enable_auto_commit=False, ) - output_topic_partitions = [ - TopicPartition(operator_name + "--OUT", partition) - for operator_name, partition in self.topic_partition_output_offsets - ] - kafka_output_consumer.assign(output_topic_partitions) + kafka_output_consumer.assign(all_topic_partitions) while True: - # start the kafka consumer try: await kafka_output_consumer.start() - for topic_partition in output_topic_partitions: - kafka_output_consumer.seek( - topic_partition, - self.topic_partition_output_offsets[(topic_partition.topic[:-5], topic_partition.partition)] - + 1, - ) + for tp in all_topic_partitions: + offset = self.dedup_output_offsets[(tp.topic[:-5], tp.partition)] + 1 + kafka_output_consumer.seek(tp, offset) + logging.warning(f"[DEDUP] Seeking {tp} to offset {offset}") except (UnknownTopicOrPartitionError, KafkaConnectionError): # fmt: skip await asyncio.sleep(1) logging.warning( @@ -171,38 +212,28 @@ async def get_messages_sent_before_recovery(self) -> None: continue break try: - # step 1 get current offset - current_offsets: dict[ - TopicPartition, - int, - ] = await kafka_output_consumer.end_offsets( - output_topic_partitions, - ) + current_offsets = await kafka_output_consumer.end_offsets(all_topic_partitions) logging.warning( - f"Reading from output from: {self.topic_partition_output_offsets} + 1 to {current_offsets}", - ) - all_partitions_done: dict[TopicPartition, bool] = dict.fromkeys( - output_topic_partitions, - False, + f"[DEDUP] End offsets: {current_offsets}", ) - for topic_partition, current_offset in current_offsets.items(): - if ( - current_offset - == self.topic_partition_output_offsets[(topic_partition.topic[:-5], topic_partition.partition)] + 1 - ): - all_partitions_done[topic_partition] = True - while not all(partition_is_done for partition_is_done in all_partitions_done.values()): - result = await kafka_output_consumer.getmany( - timeout_ms=EPOCH_INTERVAL_MS, - ) + all_partitions_done: dict[TopicPartition, bool] = { + tp: current_offsets[tp] == self.dedup_output_offsets[(tp.topic[:-5], tp.partition)] + 1 + for tp in all_topic_partitions + } + logging.warning(f"[DEDUP] Initial partitions_done: {all_partitions_done}") + per_partition_count: dict[TopicPartition, int] = dict.fromkeys(all_topic_partitions, 0) + while not all(all_partitions_done.values()): + result = await kafka_output_consumer.getmany(timeout_ms=EPOCH_INTERVAL_MS) for messages in result.values(): - self.process_messages_sent_before_recovery( - messages, - current_offsets, - all_partitions_done, - ) - if all(partition_is_done for partition_is_done in all_partitions_done.values()): + for msg in messages: + per_partition_count[TopicPartition(msg.topic, msg.partition)] += 1 + self.process_messages_sent_before_recovery(messages, current_offsets, all_partitions_done) + if all(all_partitions_done.values()): break + logging.warning( + f"[DEDUP] Scan complete. Messages read per partition: {per_partition_count}. " + f"Total dedup set size: {len(self.messages_sent_before_recovery)}", + ) finally: await kafka_output_consumer.stop() @@ -214,8 +245,6 @@ def process_messages_sent_before_recovery( ) -> None: for message in messages: tp = TopicPartition(message.topic, message.partition) + self.messages_sent_before_recovery.add(message.key) if message.offset >= current_offsets[tp] - 1: - self.messages_sent_before_recovery[tp].add(message.key) all_partitions_done[tp] = True - else: - self.messages_sent_before_recovery[tp].add(message.key) diff --git a/worker/fault_tolerance/async_snapshots.py b/worker/fault_tolerance/async_snapshots.py index 18b7b433..d99896fe 100644 --- a/worker/fault_tolerance/async_snapshots.py +++ b/worker/fault_tolerance/async_snapshots.py @@ -28,6 +28,9 @@ SNAPSHOT_BUCKET_NAME: str = os.getenv("SNAPSHOT_BUCKET_NAME", "styx-snapshots") SEQUENCER_PREFIX = "sequencer/" +# Number of elements in sequencer tuple (with/without migration checkpoint blob) +_SEQUENCER_TUPLE_LEN_WITH_MIGRATION = 5 +_SEQUENCER_TUPLE_LEN_LEGACY = 4 # Per-process cached S3 client (avoids re-creating boto3 client on every call) _s3_client: S3Client | None = None @@ -141,15 +144,16 @@ def retrieve_snapshot( dict[OperatorPartition, int], int, int, + bytes | None, ]: self.snapshot_id = snapshot_id + 1 if snapshot_id == -1: - return {}, {}, {}, 0, 0 + return {}, {}, {}, 0, 0, None s3 = _get_s3_client() data = self._load_operator_state(s3, snapshot_id, registered_operators) - tp_offsets, tp_out_offsets, epoch, t_counter = self._load_sequencer_state(s3, snapshot_id) - return data, tp_offsets, tp_out_offsets, epoch, t_counter + tp_offsets, tp_out_offsets, epoch, t_counter, migration_blob = self._load_sequencer_state(s3, snapshot_id) + return data, tp_offsets, tp_out_offsets, epoch, t_counter, migration_blob def _load_operator_state( self, @@ -165,7 +169,14 @@ def _load_operator_state( for _, key in self._iter_snapshot_files(s3, prefix, snapshot_id): partition_data = self._get_zstd_msgpack(s3, key) if operator_partition in data and partition_data: - data[operator_partition].update(partition_data) + for k, v in partition_data.items(): + if v is None: + data[operator_partition].pop(k, None) + else: + data[operator_partition][k] = v + elif partition_data: + # First delta for this partition — filter out any tombstones + data[operator_partition] = {k: v for k, v in partition_data.items() if v is not None} else: data[operator_partition] = partition_data return data @@ -174,21 +185,33 @@ def _load_sequencer_state( self, s3: S3Client, snapshot_id: int, - ) -> tuple[dict[OperatorPartition, int], dict[OperatorPartition, int], int, int]: + ) -> tuple[dict[OperatorPartition, int], dict[OperatorPartition, int], int, int, bytes | None]: topic_partition_offsets: dict[OperatorPartition, int] = {} topic_partition_output_offsets: dict[OperatorPartition, int] = {} epoch = 0 t_counter = 0 + migration_blob: bytes | None = None for _, key in self._iter_snapshot_files(s3, SEQUENCER_PREFIX, snapshot_id): - ( - topic_partition_offsets, - topic_partition_output_offsets, - epoch, - t_counter, - ) = self._get_zstd_msgpack(s3, key) - - return topic_partition_offsets, topic_partition_output_offsets, epoch, t_counter + loaded = self._get_zstd_msgpack(s3, key) + if isinstance(loaded, (tuple, list)) and len(loaded) == _SEQUENCER_TUPLE_LEN_WITH_MIGRATION: + ( + topic_partition_offsets, + topic_partition_output_offsets, + epoch, + t_counter, + migration_blob, + ) = loaded + elif isinstance(loaded, (tuple, list)) and len(loaded) == _SEQUENCER_TUPLE_LEN_LEGACY: + ( + topic_partition_offsets, + topic_partition_output_offsets, + epoch, + t_counter, + ) = loaded + migration_blob = None + + return topic_partition_offsets, topic_partition_output_offsets, epoch, t_counter, migration_blob def _iter_snapshot_files( self, diff --git a/worker/operator_state/aria/in_memory_state.py b/worker/operator_state/aria/in_memory_state.py index 12b69059..984bdcf7 100644 --- a/worker/operator_state/aria/in_memory_state.py +++ b/worker/operator_state/aria/in_memory_state.py @@ -52,6 +52,8 @@ def get_key_to_migrate( if data_to_send is None: # Key was already transferred via async migration batch return None + # Record source-side key removal as tombstone for snapshotting + self.delta_map[operator_partition][key] = None if operator_partition in self.keys_to_send: self.keys_to_send[operator_partition].discard((key, new_partition)) return data_to_send @@ -65,6 +67,7 @@ def set_data_from_migration( operator_partition = tuple(operator_partition) if data is not None: self.data[operator_partition][key] = data + self.delta_map[operator_partition][key] = data # Only remove from remote_keys when we actually received the data. # A None response means the async migration batch already transferred # this key — the batch will arrive and set it via set_batch_data_from_migration. @@ -114,6 +117,8 @@ def get_async_migrate_batch( if value is None: # Key was deleted between rehash and transfer — skip it continue + # Record source-side key removal as tombstone for snapshotting + self.delta_map[operator_partition][key] = None batch_to_send[(operator_name, new_partition)][key] = value c += 1 if not keys: @@ -132,6 +137,7 @@ def set_batch_data_from_migration( ) -> None: operator_partition = tuple(operator_partition) # new partitioning self.data[operator_partition].update(kv_pairs) + self.delta_map[operator_partition].update(kv_pairs) # Guard: remote_keys may not have entries if async batch arrived # before hash metadata, or entries were already removed. if operator_partition in self.remote_keys: diff --git a/worker/transactional_protocols/aria.py b/worker/transactional_protocols/aria.py index 8552bc12..b7738472 100644 --- a/worker/transactional_protocols/aria.py +++ b/worker/transactional_protocols/aria.py @@ -49,7 +49,7 @@ class AriaProtocol(BaseTransactionalProtocol): - def __init__( + def __init__( # noqa: PLR0913 self, worker_id: int, peers: dict[int, tuple[str, int, int]], @@ -66,6 +66,8 @@ def __init__( request_id_to_t_id_map: dict[bytes, int] | None = None, restart_after_recovery: bool = False, restart_after_migration: bool = False, + *, + dedup_output_offsets: dict[OperatorPartition, int] | None = None, ) -> None: if topic_partition_offsets is None: topic_partition_offsets = {(tp.topic, tp.partition): -1 for tp in topic_partitions} @@ -125,6 +127,7 @@ def __init__( self.egress: StyxKafkaBatchEgress = StyxKafkaBatchEgress( output_offsets, restart_after_recovery or restart_after_migration, + dedup_output_offsets=dedup_output_offsets, ) # Primary task used for processing self.function_scheduler_task: asyncio.Task | None = None @@ -150,6 +153,7 @@ def __init__( self.wait_responses_to_be_sent = asyncio.Event() self.running: bool = True + self._stopping: bool = False self.stopped: asyncio.Event = asyncio.Event() self.snapshot_marker_received: bool = False self.snapshotting_port: int = snapshotting_port @@ -180,6 +184,11 @@ async def wait_stopped(self) -> None: await self.stopped.wait() async def stop(self) -> None: + if self._stopping: + # Another caller is already tearing down — just wait for completion. + await self.stopped.wait() + return + self._stopping = True await self.ingress.stop() await self.egress.stop() await self.aio_task_scheduler.close() @@ -401,6 +410,7 @@ async def _handle_sync_cleanup(self, data: bytes) -> None: mt = MessageType.SyncCleanup (stop_gracefully,) = self.networking.decode_message(data) if stop_gracefully: + logging.warning(f"Worker {self.id} | SyncCleanup received stop_gracefully=True") self.running = False self.sync_workers_event[mt].set() @@ -565,25 +575,37 @@ async def function_scheduler(self) -> None: await self.started.wait() logging.warning("STARTED function scheduler") - while self.running: - # Wait until the ingress signals that messages are available, - # or a remote peer wants to proceed, instead of busy-spinning. - with contextlib.suppress(TimeoutError): - await asyncio.wait_for( - self.ingress.messages_available.wait(), - timeout=0.1, - ) - self.ingress.messages_available.clear() + try: + while self.running: + # Wait until the ingress signals that messages are available, + # or a remote peer wants to proceed, instead of busy-spinning. + with contextlib.suppress(TimeoutError): + await asyncio.wait_for( + self.ingress.messages_available.wait(), + timeout=0.1, + ) + self.ingress.messages_available.clear() - async with self.sequencer.lock: - sequence: list[SequencedItem] = self.sequencer.get_epoch() - if not sequence and not self.remote_wants_to_proceed: - continue + async with self.sequencer.lock: + sequence: list[SequencedItem] = self.sequencer.get_epoch() + if not sequence and not self.remote_wants_to_proceed: + continue - self.currently_processing = True - await self._process_epoch(sequence) + self.currently_processing = True + await self._process_epoch(sequence) - await self.stop() + logging.warning(f"Worker {self.id} | function_scheduler exiting (running=False)") + except asyncio.CancelledError: + logging.warning(f"Worker {self.id} | function_scheduler cancelled") + raise + except Exception: + logging.exception(f"Worker {self.id} | function_scheduler crashed") + finally: + # Only call stop() if no external caller (e.g. recovery) is + # already tearing down. Otherwise we deadlock: the external + # stop() awaits this task, while this task awaits stopped. + if not self._stopping: + await self.stop() async def _process_epoch(self, sequence: list[SequencedItem]) -> None: epoch_start = timer() @@ -1059,5 +1081,13 @@ async def sync_workers( msg_type=msg_type, serializer=serializer, ) - await self.sync_workers_event[msg_type].wait() + # Detect stuck barriers during migration debugging + try: + await asyncio.wait_for(self.sync_workers_event[msg_type].wait(), timeout=10.0) + except TimeoutError: + logging.warning( + f"Worker {self.id} | BARRIER STUCK for 10s on {msg_type.name} @epoch {self.sequencer.epoch_counter}", + ) + # Continue waiting (no timeout) + await self.sync_workers_event[msg_type].wait() self.sync_workers_event[msg_type].clear() diff --git a/worker/worker_service.py b/worker/worker_service.py index d6310615..2f1d5514 100644 --- a/worker/worker_service.py +++ b/worker/worker_service.py @@ -35,7 +35,7 @@ from styx.common.logging import logging from styx.common.message_types import MessageType from styx.common.protocols import Protocols -from styx.common.serialization import Serializer, msgpack_deserialization +from styx.common.serialization import Serializer, cloudpickle_deserialization, msgpack_deserialization from styx.common.tcp_networking import MessagingMode, NetworkingManager from styx.common.util.aio_task_scheduler import AIOTaskScheduler import uvloop @@ -134,6 +134,7 @@ def __init__(self, thread_idx: int) -> None: self.worker_operators: dict[OperatorPartition, Operator] | None = None self.migration_completed: asyncio.Event = asyncio.Event() + self._migration_generation: int = 0 # bumped on recovery to invalidate zombie Phase B tasks self.m_input_offsets: dict[OperatorPartition, int] = {} self.m_output_offsets: dict[OperatorPartition, int] = {} @@ -275,6 +276,15 @@ async def _send_snap_assigned(self, snapshot_id: int) -> None: serializer=Serializer.MSGPACK, ) + async def _send_snap_migration_reassign(self) -> None: + await self.networking.send_message( + self.networking.host_name, + self.snapshotting_port, + msg=(list(self.registered_operators.keys()),), + msg_type=MessageType.SnapMigrationReassign, + serializer=Serializer.MSGPACK, + ) + def _attach_operator_networking(self) -> None: for operator in self.registered_operators.values(): operator.attach_state_networking( @@ -315,7 +325,7 @@ async def _handle_receive_execution_plan(self, data: bytes, _: MessageType) -> N await self._send_snap_assigned(snapshot_id=snapshot_id) logging.warning("Retrieving Snapshot from S3") - (snap_data, _in_off, _out_off, epoch, t_counter) = self.async_snapshots.retrieve_snapshot( + (snap_data, _in_off, _out_off, epoch, t_counter, _migration_blob) = self.async_snapshots.retrieve_snapshot( snapshot_id, self.registered_operators.keys(), ) @@ -367,7 +377,16 @@ async def _handle_init_migration(self, data: bytes, _: MessageType) -> None: logging.error(f"Uncaught exception during migration Phase A: {e}") async def _migration_stop_protocol(self) -> None: - await self.function_execution_protocol.wait_stopped() + try: + await asyncio.wait_for( + self.function_execution_protocol.wait_stopped(), + timeout=15.0, + ) + except TimeoutError: + logging.warning( + "MIGRATION | Protocol did not stop within 15s via stop_gracefully — force-stopping", + ) + await self.function_execution_protocol.stop() logging.warning("MIGRATION | ARIA STOPPED") async def _migration_decode_and_apply_plan(self, data: bytes) -> None: @@ -508,8 +527,8 @@ async def _migration_rebuild_runtime(self) -> None: n_assigned_partitions=len(self.registered_operators), ) - # Snapshot id -1 indicates "migration" - await self._send_snap_assigned(snapshot_id=-1) + # Reassign partitions in snapshotting subprocess, preserving accumulated deltas + await self._send_snap_migration_reassign() # Ensure local state has all newly assigned partitions self._ensure_local_state_partitions() @@ -536,6 +555,7 @@ async def _migration_rebuild_runtime(self) -> None: epoch_counter=self.m_epoch_counter, t_counter=self.m_t_counter, restart_after_migration=True, + dedup_output_offsets=dict(self.m_output_offsets), ) async def _handle_receive_migration_hashes( @@ -594,6 +614,7 @@ async def _handle_migration_repartitioning_done( _: MessageType, ) -> None: """Phase B: protocol stops, catch-up, send hashes, rebuild, resume.""" + my_generation = self._migration_generation try: logging.warning(f"MIGRATION | PHASE B START at {time.time_ns() // 1_000_000}") phase_b_start = timer() @@ -602,6 +623,14 @@ async def _handle_migration_repartitioning_done( t_stop_start = timer() await self._migration_stop_protocol() t_stop_end = timer() + + # Bail out if recovery invalidated this Phase B task + if self._migration_generation != my_generation: + logging.warning( + "MIGRATION | PHASE B ZOMBIE detected after protocol stop — bailing out", + ) + return + logging.warning( "MIGRATION | PHASE B Protocol Stopped |" f" @Epoch {self.function_execution_protocol.sequencer.epoch_counter} |" @@ -631,6 +660,24 @@ async def _handle_migration_repartitioning_done( f"MIGRATION | PHASE B Hash metadata sent | took: {t_hash_send_end - t_hash_send_start}", ) + # 5.5 Trigger a migration checkpoint snapshot. + # This captures the state at the exact moment before data transfer begins, + # enabling checkpoint-and-resume recovery instead of revert-and-retry. + await self.function_execution_protocol.send_delta_to_snapshotting_proc() + await self.function_execution_protocol.snapshotting_networking_manager.send_message( + self.function_execution_protocol.networking.host_name, + self.function_execution_protocol.snapshotting_port, + msg=( + self.function_execution_protocol.topic_partition_offsets, + self.function_execution_protocol.egress.topic_partition_output_offsets, + self.function_execution_protocol.sequencer.epoch_counter, + self.function_execution_protocol.sequencer.t_counter, + ), + msg_type=MessageType.SnapTakeSnapshot, + serializer=Serializer.MSGPACK, + ) + logging.warning("MIGRATION | PHASE B Migration checkpoint snapshot triggered") + # 6. Send MigrationInitDone with actual stop-time counters await self.networking.send_message( DISCOVERY_HOST, @@ -649,6 +696,13 @@ async def _handle_migration_repartitioning_done( logging.warning("MIGRATION | PHASE B WAITING FOR MigrationDone") await self.migration_completed.wait() + # Bail out if recovery invalidated this Phase B task + if self._migration_generation != my_generation: + logging.warning( + "MIGRATION | PHASE B ZOMBIE detected (generation mismatch) — bailing out", + ) + return + # 8. Rebuild runtime (now has correct merged counters from MigrationDone) t_deploy_start = timer() await self._migration_rebuild_runtime() @@ -659,6 +713,9 @@ async def _handle_migration_repartitioning_done( # 9. Start protocol with background migration self.function_execution_protocol.start() + # Run dedup scan after the MigrationDone barrier ensures all workers + # have flushed their old egress batches, then mark egress started. + await self.function_execution_protocol.egress.run_dedup_scan_and_mark_started() self.function_execution_protocol.started.set() # Reset sync events for next time @@ -684,6 +741,14 @@ async def _handle_migration_done(self, data: bytes, _: MessageType) -> None: async def _handle_init_recovery(self, data: bytes, _: MessageType) -> None: start_time = timer() + # Invalidate any zombie Phase B tasks from an interrupted migration. + # Bump the generation so that any Phase B task waiting on + # migration_completed will detect the generation mismatch and bail out. + self._migration_generation += 1 + self.migration_completed.set() # unblock the zombie so it can exit + await asyncio.sleep(0) # yield to let the zombie run its bail-out check + self.migration_completed.clear() + ( self.id, self.worker_operators, @@ -719,16 +784,72 @@ async def _handle_init_recovery(self, data: bytes, _: MessageType) -> None: ) await self._send_snap_assigned(snapshot_id=snapshot_id) - (snap_data, tp_offsets, tp_out_offsets, epoch, t_counter) = self.async_snapshots.retrieve_snapshot( - snapshot_id, - self.registered_operators.keys(), + (snap_data, tp_offsets, tp_out_offsets, epoch, t_counter, migration_blob) = ( + self.async_snapshots.retrieve_snapshot( + snapshot_id, + self.registered_operators.keys(), + ) + ) + + logging.warning( + f"[RECOVERY] W{self.id} snapshot_id={snapshot_id} epoch={epoch} t_counter={t_counter}" + f" migration_blob={'present' if migration_blob else 'none'}", + ) + logging.warning( + f"[RECOVERY] W{self.id} registered_operators={list(self.registered_operators.keys())}", + ) + logging.warning( + f"[RECOVERY] W{self.id} raw tp_offsets from snapshot: {tp_offsets}", + ) + logging.warning( + f"[RECOVERY] W{self.id} raw tp_out_offsets from snapshot: {tp_out_offsets}", ) + # Keep the FULL output offsets for dedup scanning (all partitions from + # the global snapshot). Each worker must scan ALL output partitions so + # that cross-worker dedup works: a request originally processed on + # worker A may be replayed on worker B after recovery. + all_dedup_out_offsets = dict(tp_out_offsets) tp_offsets = {k: v for k, v in tp_offsets.items() if k in self.registered_operators} tp_out_offsets = {k: v for k, v in tp_out_offsets.items() if k in self.registered_operators} + logging.warning( + f"[RECOVERY] W{self.id} filtered tp_offsets: {tp_offsets}", + ) + logging.warning( + f"[RECOVERY] W{self.id} filtered tp_out_offsets: {tp_out_offsets}", + ) + logging.warning( + f"[RECOVERY] W{self.id} dedup_output_offsets (all partitions): {all_dedup_out_offsets}", + ) + self.attach_state_to_operators_after_snapshot(snap_data) + # Detect migration checkpoint: if present, re-derive keys_to_send and resume transfer + restart_after_migration = False + if migration_blob is not None: + logging.warning(f"[RECOVERY] W{self.id} Migration checkpoint detected — resuming migration transfer") + migration_meta = cloudpickle_deserialization(migration_blob) + new_graph = migration_meta["new_graph"] + + # Re-derive keys_to_send from local state + new partitioner + keys_to_send: dict[tuple, set] = defaultdict(set) + for op_part in self.registered_operators: + op_name, partition = op_part + partitioner = new_graph.get_operator_by_name(op_name).get_partitioner() + op_data = self.local_state.get_operator_data_for_repartitioning(op_part) + for key in op_data: + new_part = partitioner.get_partition_no_cache(key) + if new_part != partition: + keys_to_send[op_part].add((key, new_part)) + + if keys_to_send: + self.local_state.add_keys_to_send(keys_to_send) + logging.warning( + f"[RECOVERY] W{self.id} Re-derived {sum(len(v) for v in keys_to_send.values())} keys to send", + ) + restart_after_migration = True + request_id_to_t_id_map = await self.get_sequencer_assignments_before_failure( epoch, ) @@ -747,7 +868,9 @@ async def _handle_init_recovery(self, data: bytes, _: MessageType) -> None: epoch_counter=epoch, t_counter=t_counter, request_id_to_t_id_map=request_id_to_t_id_map, - restart_after_recovery=True, + restart_after_recovery=not restart_after_migration, + restart_after_migration=restart_after_migration, + dedup_output_offsets=all_dedup_out_offsets, ) self.function_execution_protocol.start() @@ -766,6 +889,11 @@ async def _handle_init_recovery(self, data: bytes, _: MessageType) -> None: ) async def _handle_ready_after_recovery(self, _data: bytes, _: MessageType) -> None: + # All workers have stopped their old protocols by this point (the + # coordinator only sends ReadyAfterRecovery after all workers reported). + # Run the dedup scan NOW so it captures every message flushed by every + # surviving worker's old egress. + await self.function_execution_protocol.egress.run_dedup_scan_and_mark_started() self.function_execution_protocol.started.set() logging.warning( f"Worker: {self.id} recovered and ready at : {time.time() * 1000}", From 05e48fce23dc9c98c8c237c91323fd0c98c0973f Mon Sep 17 00:00:00 2001 From: Kyriakos Psarakis Date: Sat, 24 Oct 2026 16:03:50 +0200 Subject: [PATCH 02/20] Keep migrated keys consistent across snapshots and recovery The async migration sender ran independently of the epoch, so a batch could be in flight across a snapshot cut: the source's tombstone landed in one snapshot and the destination's insert in the next, and recovery from the first lost the key. The epoch now drains in-flight batches (acknowledged by the destination once applied) before its processing-done barrier and lets the sender resume after the cleanup barrier. The sender also waits for the protocol to start, so peers have loaded their state first. The delta flush cleared the delta map after awaiting the send, wiping migration writes that landed during the await. It now swaps the map out before sending. After a post-checkpoint recovery the destinations had no remote_keys, so a transaction on a key still at its source read it as missing. Workers now exchange the re-derived hash metadata on ReadyAfterRecovery and start the protocol only once every peer has reported. New partitions are created in Phase B before MigrationInitDone, so a peer that rebuilds first cannot send batches to a partition that does not exist yet. Co-Authored-By: Claude Opus 5.5 --- styx-package/styx/common/message_types.py | 2 + tests/unit/styx_package/test_message_types.py | 6 +- .../test_migration_snapshot_consistency.py | 229 ++++++++++++++++++ .../worker/test_recovery_migration_hashes.py | 135 +++++++++++ worker/operator_state/aria/in_memory_state.py | 11 + worker/transactional_protocols/aria.py | 116 +++++++-- worker/worker_service.py | 112 ++++++++- 7 files changed, 575 insertions(+), 36 deletions(-) create mode 100644 tests/unit/worker/test_migration_snapshot_consistency.py create mode 100644 tests/unit/worker/test_recovery_migration_hashes.py diff --git a/styx-package/styx/common/message_types.py b/styx-package/styx/common/message_types.py index adb742c2..fb952567 100644 --- a/styx-package/styx/common/message_types.py +++ b/styx-package/styx/common/message_types.py @@ -46,3 +46,5 @@ class MessageType(IntEnum): RunFunRemoteBatch = 42 AckBatch = 43 SnapMigrationReassign = 44 + AsyncMigrationAck = 45 + RecoveryMigrationHashes = 46 diff --git a/tests/unit/styx_package/test_message_types.py b/tests/unit/styx_package/test_message_types.py index 16cb8874..7e56b6a1 100644 --- a/tests/unit/styx_package/test_message_types.py +++ b/tests/unit/styx_package/test_message_types.py @@ -10,7 +10,7 @@ def test_is_int_enum(self): assert issubclass(MessageType, IntEnum) def test_total_count(self): - assert len(MessageType) == 44 + assert len(MessageType) == 46 def test_all_values_unique(self): values = [m.value for m in MessageType] @@ -20,7 +20,7 @@ def test_values_in_expected_range(self): # 25 (AckCache) was removed when the fallback cache was dropped; the # remaining IDs are preserved to keep wire-format stability. values = sorted(m.value for m in MessageType) - assert values == [v for v in range(45) if v != 25] + assert values == [v for v in range(47) if v != 25] def test_known_values(self): assert MessageType.RunFunRemote == 0 @@ -31,6 +31,8 @@ def test_known_values(self): assert MessageType.RunFunRemoteBatch == 42 assert MessageType.AckBatch == 43 assert MessageType.SnapMigrationReassign == 44 + assert MessageType.AsyncMigrationAck == 45 + assert MessageType.RecoveryMigrationHashes == 46 def test_int_comparison(self): assert MessageType.ClientMsg == 16 diff --git a/tests/unit/worker/test_migration_snapshot_consistency.py b/tests/unit/worker/test_migration_snapshot_consistency.py new file mode 100644 index 00000000..8481f614 --- /dev/null +++ b/tests/unit/worker/test_migration_snapshot_consistency.py @@ -0,0 +1,229 @@ +"""Snapshot consistency of async state migration. + +A key moved by the async migration sender is removed from its source (with a +tombstone in the source's delta) and inserted at its destination (in the +destination's delta). A snapshot must never record the first without the second: + +- the delta flush must not wipe migration writes that land while it is awaiting; +- a batch must not be in flight across an epoch's snapshot cut. +""" + +import asyncio +import copy +from types import SimpleNamespace +from unittest.mock import AsyncMock + +from styx.common.message_types import MessageType +from styx.common.tcp_networking import NetworkingManager + +from worker.operator_state.aria.in_memory_state import InMemoryOperatorState +from worker.transactional_protocols import aria +from worker.transactional_protocols.aria import AriaProtocol + +OP = "op" +SRC = (OP, 0) +DST = (OP, 1) +W1 = ("w1", 5001, 6001) +W2 = ("w2", 5002, 6002) +DNS = {OP: {0: W1, 1: W2}} +LINK_DELAY_S = 0.01 + + +# --------------------------------------------------------------------------- +# Delta flush +# --------------------------------------------------------------------------- + + +class TestDeltaFlushKeepsConcurrentMigrationWrites: + async def test_batch_landing_during_flush_goes_to_next_delta(self): + state = InMemoryOperatorState({SRC, DST}) + state.batch_insert({"k0": "v0"}, OP, 0) + sent: list[dict] = [] + + async def send_message(*_args, msg, **_kwargs): + # The real send serializes before its first await. + sent.append(copy.deepcopy(msg[0])) + await asyncio.sleep(0) + # An async migration batch lands while the flush is awaiting. + state.set_batch_data_from_migration(DST, {"k1": "v1"}) + + proto = SimpleNamespace( + local_state=state, + snapshotting_networking_manager=SimpleNamespace(send_message=AsyncMock(side_effect=send_message)), + networking=SimpleNamespace(host_name="w1"), + snapshotting_port=7000, + ) + await AriaProtocol.send_delta_to_snapshotting_proc(proto) + + assert sent == [{SRC: {"k0": "v0"}, DST: {}}] + # Not wiped: it ships with the next flush. + assert state.delta_map[DST] == {"k1": "v1"} + assert state.delta_map[SRC] == {} + + def test_take_delta_map_returns_deltas_and_starts_fresh(self): + state = InMemoryOperatorState({SRC, DST}) + state.batch_insert({"k0": "v0"}, OP, 0) + delta = state.take_delta_map() + assert delta == {SRC: {"k0": "v0"}, DST: {}} + assert state.delta_map == {SRC: {}, DST: {}} + state.batch_insert({"k1": "v1"}, OP, 1) + assert delta[DST] == {} + + +# --------------------------------------------------------------------------- +# Batches across the snapshot cut +# --------------------------------------------------------------------------- + + +class _Link: + """Delivers protocol messages between fake workers after a delay, like a real link.""" + + def __init__(self) -> None: + self.workers: dict[tuple[str, int], AriaProtocol] = {} + self.to_coordinator: list[MessageType] = [] + self.in_flight: set[asyncio.Task] = set() + + def view(self) -> SimpleNamespace: + return SimpleNamespace( + send_message=self.send_message, + decode_message=NetworkingManager.decode_message, + wait_remote_key_event={}, + ) + + async def send_message(self, host, port, msg, msg_type, serializer) -> None: + # Drop the 8-byte length prefix, as the receiving TCP reader does. + data = NetworkingManager.encode_message(msg=msg, msg_type=msg_type, serializer=serializer)[8:] + target = self.workers.get((host, port)) + if target is None: + self.to_coordinator.append(msg_type) + return + + async def deliver() -> None: + await asyncio.sleep(LINK_DELAY_S) + await target.protocol_handlers_map[msg_type](data) + + task = asyncio.get_running_loop().create_task(deliver()) + self.in_flight.add(task) + task.add_done_callback(self.in_flight.discard) + + +def _worker(link: _Link, worker_id: int, addr: tuple, peers: dict, state: InMemoryOperatorState) -> AriaProtocol: + proto = AriaProtocol.__new__(AriaProtocol) + proto.id = worker_id + proto.peers = peers + proto.dns = DNS + proto.local_state = state + proto.networking = link.view() + proto.sequencer = SimpleNamespace(epoch_counter=0) + proto.started = asyncio.Event() + proto.started.set() + proto.migrating_state = True + proto._init_async_migration_state() + proto.protocol_handlers_map = { + MessageType.AsyncMigration: proto._handle_async_migration, + MessageType.AsyncMigrationAck: proto._handle_async_migration_ack, + } + link.workers[(addr[0], addr[2])] = proto + return proto + + +def _two_workers(n_keys: int): + link = _Link() + src_state = InMemoryOperatorState({SRC}) + src_state.data[SRC] = {k: f"v{k}" for k in range(n_keys)} + src_state.add_keys_to_send({SRC: {(k, 1) for k in range(n_keys)}}) + dst_state = InMemoryOperatorState({DST}) + src = _worker(link, 1, W1, {2: W2}, src_state) + dst = _worker(link, 2, W2, {1: W1}, dst_state) + return link, src, dst + + +def _restore(base: dict, *deltas: dict) -> dict: + """What recovery loads: the base snapshot with each worker's deltas applied.""" + data = {op_part: dict(kv) for op_part, kv in base.items()} + for delta in deltas: + for op_part, kv in delta.items(): + part = data.setdefault(op_part, {}) + for k, v in kv.items(): + if v is None: + part.pop(k, None) + else: + part[k] = v + return data + + +class TestMigrationBatchesDoNotStraddleSnapshotCut: + async def test_cut_after_drain_has_every_key_on_one_side(self): + n_keys = 10 + _link, src, dst = _two_workers(n_keys) + base = {SRC: dict(src.local_state.data[SRC]), DST: {}} + + sender = asyncio.create_task(src._continuous_migration_sender()) + await asyncio.sleep(0) # the sender removes a batch and puts it on the wire + assert src.local_state.data[SRC] == {} + + # Epoch reaches its barriers: drain, then every worker flushes its delta. + await src._drain_migration_batches() + restored = _restore(base, src.local_state.take_delta_map(), dst.local_state.take_delta_map()) + + assert restored[SRC] == {} + assert restored[DST] == {k: f"v{k}" for k in range(n_keys)} + src._migration_gate.set() + await asyncio.wait_for(sender, timeout=2) + + async def test_closed_gate_keeps_sender_from_taking_batches(self, monkeypatch): + monkeypatch.setattr(aria, "ASYNC_MIGRATION_BATCH_SIZE", 2) + _link, src, _dst = _two_workers(6) + + sender = asyncio.create_task(src._continuous_migration_sender()) + await asyncio.sleep(0) + await src._drain_migration_batches() + remaining = src.local_state.keys_remaining_to_send() + assert remaining == 4 + + await asyncio.sleep(LINK_DELAY_S * 3) + assert src.local_state.keys_remaining_to_send() == remaining # paused across the cut + + src._migration_gate.set() # next epoch + await asyncio.wait_for(sender, timeout=2) + assert src.local_state.keys_remaining_to_send() == 0 + + async def test_migration_done_is_reported_after_last_ack(self): + link, src, dst = _two_workers(3) + + sender = asyncio.create_task(src._continuous_migration_sender()) + await asyncio.sleep(0) + assert src.local_state.keys_remaining_to_send() == 0 + assert link.to_coordinator == [] # the batch has not landed yet + + await asyncio.wait_for(sender, timeout=2) + assert link.to_coordinator == [MessageType.MigrationDone] + assert dst.local_state.data[DST] == {0: "v0", 1: "v1", 2: "v2"} + + async def test_sender_waits_until_protocol_started(self): + _link, src, _dst = _two_workers(3) + src.started.clear() + + sender = asyncio.create_task(src._continuous_migration_sender()) + await asyncio.sleep(LINK_DELAY_S * 2) + assert src.local_state.keys_remaining_to_send() == 3 + + src.started.set() + await asyncio.wait_for(sender, timeout=2) + assert src.local_state.keys_remaining_to_send() == 0 + + async def test_stale_ack_from_previous_protocol_is_ignored(self): + _link, src, _dst = _two_workers(1) + src._migration_unacked.add(0) + src._migration_all_acked.clear() + + src._ack_migration_batch(src._migration_nonce + 1, 0) + assert not src._migration_all_acked.is_set() + + src._ack_migration_batch(src._migration_nonce, 0) + assert src._migration_all_acked.is_set() + + async def test_drain_returns_immediately_without_batches(self): + _link, src, _dst = _two_workers(0) + await asyncio.wait_for(src._drain_migration_batches(), timeout=0.1) + assert not src._migration_gate.is_set() diff --git a/tests/unit/worker/test_recovery_migration_hashes.py b/tests/unit/worker/test_recovery_migration_hashes.py new file mode 100644 index 00000000..1704bfc1 --- /dev/null +++ b/tests/unit/worker/test_recovery_migration_hashes.py @@ -0,0 +1,135 @@ +"""Rebuilding remote_keys after a post-checkpoint recovery. + +Recovery re-derives each source's keys_to_send from its snapshot, but the +destinations also need remote_keys (which peer still holds each key of their new +partitions). Otherwise a transaction on a key that has not arrived reads it as +missing instead of fetching it, and the arriving batch later overwrites its write. +""" + +import asyncio +from types import SimpleNamespace +from unittest.mock import AsyncMock + +from styx.common.message_types import MessageType +from styx.common.serialization import Serializer +from styx.common.tcp_networking import NetworkingManager + +from worker.operator_state.aria.in_memory_state import InMemoryOperatorState +from worker.worker_service import Worker + +OP = "op" +W1 = ("w1", 5001, 6001) +W2 = ("w2", 5002, 6002) +SNAP_ID = 7 + + +class _Link: + """Delivers control messages between fake workers after a short delay.""" + + def __init__(self) -> None: + self.workers: dict[tuple[str, int], Worker] = {} + + def view(self) -> SimpleNamespace: + return SimpleNamespace(send_message=self.send_message, decode_message=NetworkingManager.decode_message) + + async def send_message(self, host, port, msg, msg_type, serializer) -> None: + # Drop the 8-byte length prefix, as the receiving TCP reader does. + data = NetworkingManager.encode_message(msg=msg, msg_type=msg_type, serializer=serializer)[8:] + target = self.workers[(host, port)] + + async def deliver() -> None: + await asyncio.sleep(0.01) + await target._message_handlers[msg_type](data, msg_type) + + asyncio.get_running_loop().create_task(deliver()) + + +def _worker(link: _Link, worker_id: int, addr: tuple, peers: dict, partitions: set, hashes_out: dict) -> Worker: + w = Worker.__new__(Worker) + w.id = worker_id + w.peers = peers + w.dns = {OP: {0: W1, 1: W2, 2: W1}} + w.local_state = InMemoryOperatorState(partitions) + w.networking = link.view() + w._migration_generation = 1 + w._recovery_hashes_received = asyncio.Event() + w._reset_recovery_hash_exchange(tag=SNAP_ID, hashes_out=hashes_out) + w.function_execution_protocol = SimpleNamespace( + egress=SimpleNamespace(run_dedup_scan_and_mark_started=AsyncMock()), + started=asyncio.Event(), + ) + w._message_handlers = {MessageType.RecoveryMigrationHashes: w._handle_recovery_migration_hashes} + link.workers[(addr[0], addr[1])] = w + return w + + +def _cluster(): + """W1 owns (op, 0) and (op, 2), W2 owns (op, 1). After recovery W1 still holds + key 10 for W2's (op, 1) and key 20 for its own (op, 2).""" + link = _Link() + w1 = _worker( + link, + 1, + W1, + {2: W2}, + {(OP, 0), (OP, 2)}, + {(OP, 1): {10: (1, 0)}, (OP, 2): {20: (1, 0)}}, + ) + w2 = _worker(link, 2, W2, {1: W1}, {(OP, 1)}, {}) + return w1, w2 + + +async def _ready(w: Worker) -> None: + await w._handle_ready_after_recovery(b"", MessageType.ReadyAfterRecovery) + + +class TestRecoveryMigrationHashExchange: + async def test_destinations_learn_which_peer_holds_their_keys(self): + w1, w2 = _cluster() + await asyncio.wait_for(asyncio.gather(_ready(w1), _ready(w2)), timeout=2) + + assert w2.local_state.in_remote_keys(10, OP, 1) + assert tuple(w2.local_state.get_worker_id_old_partition(OP, 1, 10)) == (1, 0) + # Destination partition on the same worker is registered locally. + assert w1.local_state.in_remote_keys(20, OP, 2) + assert w1.function_execution_protocol.started.is_set() + assert w2.function_execution_protocol.started.is_set() + + async def test_protocol_does_not_start_before_every_peer_reported(self): + w1, w2 = _cluster() + w2_ready = asyncio.create_task(_ready(w2)) + await asyncio.sleep(0.05) + assert not w2.function_execution_protocol.started.is_set() + w2.function_execution_protocol.egress.run_dedup_scan_and_mark_started.assert_not_awaited() + + await asyncio.wait_for(asyncio.gather(_ready(w1), w2_ready), timeout=2) + assert w2.function_execution_protocol.started.is_set() + + async def test_hashes_for_another_recovery_are_dropped(self): + _w1, w2 = _cluster() + data = NetworkingManager.encode_message( + msg=(SNAP_ID - 1, 1, [((OP, 1), {10: (1, 0)})]), + msg_type=MessageType.RecoveryMigrationHashes, + serializer=Serializer.MSGPACK, + )[8:] + await w2._handle_recovery_migration_hashes(data, MessageType.RecoveryMigrationHashes) + assert not w2.local_state.in_remote_keys(10, OP, 1) + assert w2._recovery_hash_senders == set() + + async def test_newer_recovery_releases_waiter_without_starting_old_protocol(self): + _w1, w2 = _cluster() + old_protocol = w2.function_execution_protocol + w2_ready = asyncio.create_task(_ready(w2)) + await asyncio.sleep(0.02) + + # A new InitRecovery arrives before the peer reported. + w2._migration_generation += 1 + w2._reset_recovery_hash_exchange(tag=-1, hashes_out={}) + await asyncio.wait_for(w2_ready, timeout=1) + assert not old_protocol.started.is_set() + + async def test_no_exchange_for_ordinary_recovery(self): + _w1, w2 = _cluster() + w2._reset_recovery_hash_exchange(tag=-1, hashes_out={}) + await asyncio.wait_for(_ready(w2), timeout=1) + assert w2.function_execution_protocol.started.is_set() diff --git a/worker/operator_state/aria/in_memory_state.py b/worker/operator_state/aria/in_memory_state.py index 984bdcf7..ba721125 100644 --- a/worker/operator_state/aria/in_memory_state.py +++ b/worker/operator_state/aria/in_memory_state.py @@ -215,6 +215,17 @@ def clear_delta_map(self) -> None: for operator_partition in self.operator_partitions: self.delta_map[operator_partition].clear() + def take_delta_map(self) -> dict[OperatorPartition, KVPairs]: + """Hand over the accumulated deltas and continue on a fresh map. + + Swapping (instead of clearing once the caller is done) matters because the + caller awaits while shipping the deltas: async-migration writes that land + during that await must go into the next delta, not be wiped unsent. + """ + delta = self.delta_map + self.delta_map = {operator_partition: {} for operator_partition in self.operator_partitions} + return delta + def commit_fallback_transaction(self, t_id: int) -> None: if t_id in self.fallback_commit_buffer: for operator_partition, kv_pairs in self.fallback_commit_buffer[t_id].items(): diff --git a/worker/transactional_protocols/aria.py b/worker/transactional_protocols/aria.py index b7738472..a6be34e0 100644 --- a/worker/transactional_protocols/aria.py +++ b/worker/transactional_protocols/aria.py @@ -1,6 +1,8 @@ import asyncio import contextlib +import itertools import os +import random import time from timeit import default_timer as timer from typing import TYPE_CHECKING @@ -159,6 +161,7 @@ def __init__( # noqa: PLR0913 self.snapshotting_port: int = snapshotting_port self.migrating_state: bool = restart_after_migration + self._init_async_migration_state() self.protocol_handlers_map: dict[MessageType, Callable[[bytes], Awaitable[None]]] = { MessageType.RunFunRemote: self._handle_run_fun_remote, @@ -178,8 +181,24 @@ def __init__( # noqa: PLR0913 MessageType.DeterministicReordering: self._handle_deterministic_reordering, MessageType.RemoteWantsToProceed: self._handle_remote_wants_to_proceed, MessageType.AsyncMigration: self._handle_async_migration, + MessageType.AsyncMigrationAck: self._handle_async_migration_ack, } + def _init_async_migration_state(self) -> None: + # A migration batch must leave its source and land at its destination in + # the same epoch, otherwise a snapshot cut can record the source-side + # tombstone without the destination-side insert. The epoch closes this + # gate and waits for every unacknowledged batch before its barriers. + self._migration_gate: asyncio.Event = asyncio.Event() + self._migration_gate.set() + self._migration_all_acked: asyncio.Event = asyncio.Event() + self._migration_all_acked.set() + self._migration_unacked: set[int] = set() + self._migration_batch_ids = itertools.count() + # Distinguishes acks for this protocol's batches from late acks for a + # previous protocol instance (recovery or migration rebuilds the protocol). + self._migration_nonce: int = random.getrandbits(62) + async def wait_stopped(self) -> None: await self.stopped.wait() @@ -476,18 +495,50 @@ async def _handle_remote_wants_to_proceed(self, _: bytes) -> None: self.ingress.messages_available.set() async def _handle_async_migration(self, data: bytes) -> None: - # Lock-free: no awaits. - operator_partition, batch = self.networking.decode_message(data) + operator_partition, batch, sender_id, nonce, batch_id = self.networking.decode_message(data) logging.debug(f"MIGRATION | Received migration batch for: {operator_partition}") - self.local_state.set_batch_data_from_migration(operator_partition, batch) - # Unblock any transactions waiting for these keys via RequestRemoteKey op = tuple(operator_partition) + self.local_state.set_batch_data_from_migration(op, batch) + # Unblock any transactions waiting for these keys via RequestRemoteKey if op in self.networking.wait_remote_key_event: for key in batch: if key in self.networking.wait_remote_key_event.get(op, {}): self.networking.wait_remote_key_event[op][key].set() + # Acknowledge only after the batch is applied (and so in this epoch's delta). + if sender_id == self.id: + self._ack_migration_batch(nonce, batch_id) + return + host, _, protocol_port = self.peers[sender_id] + await self.networking.send_message( + host, + protocol_port, + msg=(nonce, batch_id), + msg_type=MessageType.AsyncMigrationAck, + serializer=Serializer.MSGPACK, + ) + + async def _handle_async_migration_ack(self, data: bytes) -> None: + # Lock-free: no awaits. + nonce, batch_id = self.networking.decode_message(data) + self._ack_migration_batch(nonce, batch_id) + + def _ack_migration_batch(self, nonce: int, batch_id: int) -> None: + if nonce != self._migration_nonce: + # Ack for a batch sent by a protocol instance that no longer exists. + return + self._migration_unacked.discard(batch_id) + if not self._migration_unacked: + self._migration_all_acked.set() + + async def _drain_migration_batches(self) -> None: + """Stop the async sender and wait until every batch it sent has been applied + at its destination, so the coming snapshot cut sees each migrated key on + exactly one side.""" + self._migration_gate.clear() + await self._migration_all_acked.wait() + async def _write_to_wal(self, sequence: list[SequencedItem]) -> tuple[float, float]: start_wal = timer() sequence_to_log = msgpack_serialization( @@ -508,43 +559,56 @@ async def _send_migration_batch( self, batch: dict, ) -> None: - """Send a migration batch to destination workers in parallel.""" + """Send a migration batch to destination workers in parallel. + + Every batch goes over the network, including to this worker, and stays + unacknowledged until the destination has applied it.""" send_tasks = [] for operator_partition, k_v_pairs in batch.items(): operator_name, partition = operator_partition worker = self.dns[operator_name][partition] - if worker == self.id: - self.local_state.set_batch_data_from_migration( - operator_partition, - k_v_pairs, - ) - else: - send_tasks.append( - self.networking.send_message( - worker[0], - worker[2], - msg=(operator_partition, k_v_pairs), - msg_type=MessageType.AsyncMigration, - serializer=Serializer.MSGPACK, - ), - ) + # Register before the first await so a concurrent drain waits for it. + batch_id = next(self._migration_batch_ids) + self._migration_unacked.add(batch_id) + self._migration_all_acked.clear() + send_tasks.append( + self.networking.send_message( + worker[0], + worker[2], + msg=(operator_partition, k_v_pairs, self.id, self._migration_nonce, batch_id), + msg_type=MessageType.AsyncMigration, + serializer=Serializer.MSGPACK, + ), + ) if send_tasks: await asyncio.gather(*send_tasks) async def _continuous_migration_sender(self) -> None: """Background coroutine that continuously sends migration batches - independently of epoch barriers, completing migration faster. + alongside epoch processing, completing migration faster. + + It pauses while an epoch is between its drain (before the processing-done + barrier) and its cleanup barrier, so no batch is in flight across a + snapshot cut. Completion is based only on keys_remaining_to_send() — the coordinator's global barrier (all workers report MigrationDone) ensures all data is fully transferred. We do NOT check keys_remaining_to_remote() because hash metadata from other workers may not have arrived yet. """ + # Peers may still be loading state until the protocol is started + # (ReadyAfterRecovery, or the rebuild after MigrationDone). + await self.started.wait() logging.warning("MIGRATION | Continuous migration sender started") try: while self.migrating_state: + if not self._migration_gate.is_set(): + await self._migration_gate.wait() + continue if not self.local_state.has_keys_to_send(): if self.local_state.keys_remaining_to_send() == 0: + # Report done only once every batch has landed. + await self._migration_all_acked.wait() self.migrating_state = False await self.networking.send_message( DISCOVERY_HOST, @@ -612,6 +676,10 @@ async def _process_epoch(self, sequence: list[SequencedItem]) -> None: timings = await self._run_epoch_functions_and_chain(sequence) + # Every peer's commit flush and snapshot cut follow this barrier, so batches + # that land before it belong to this epoch at both ends. + await self._drain_migration_batches() + sync_time = 0.0 sync_time += await self._sync_processing_done() @@ -666,6 +734,9 @@ async def _process_epoch(self, sequence: list[SequencedItem]) -> None: fallback_ms=(fallback_end - fallback_start) * 1000, snap_ms=(snap_end - snap_start) * 1000, ) + # Every worker has passed this epoch's snapshot cut, so batches sent from + # here on belong to the next epoch at both ends. + self._migration_gate.set() async def _run_epoch_functions_and_chain( self, @@ -884,7 +955,7 @@ async def _sync_cleanup( ) async def send_delta_to_snapshotting_proc(self) -> None: - delta_to_send = self.local_state.get_data_for_snapshot() + delta_to_send = self.local_state.take_delta_map() await self.snapshotting_networking_manager.send_message( self.networking.host_name, @@ -893,7 +964,6 @@ async def send_delta_to_snapshotting_proc(self) -> None: msg_type=MessageType.SnapProcDelta, serializer=Serializer.MSGPACK, ) - self.local_state.clear_delta_map() def cleanup_after_epoch(self) -> None: self.concurrency_aborts_everywhere.clear() diff --git a/worker/worker_service.py b/worker/worker_service.py index 2f1d5514..ac980979 100644 --- a/worker/worker_service.py +++ b/worker/worker_service.py @@ -136,6 +136,13 @@ def __init__(self, thread_idx: int) -> None: self.migration_completed: asyncio.Event = asyncio.Event() self._migration_generation: int = 0 # bumped on recovery to invalidate zombie Phase B tasks + # Post-checkpoint recovery: workers exchange which keys each still holds + # for the others' new partitions (all-to-all) before the protocol starts. + self._recovery_hash_tag: int = -1 # snapshot id of the recovery, -1 = no exchange + self._recovery_hashes_out: dict[OperatorPartition, dict[Any, tuple[int, int]]] = {} + self._recovery_hash_senders: set[int] = set() + self._recovery_hashes_received: asyncio.Event = asyncio.Event() + self.m_input_offsets: dict[OperatorPartition, int] = {} self.m_output_offsets: dict[OperatorPartition, int] = {} self.m_epoch_counter: int = -1 @@ -177,6 +184,7 @@ def __init__(self, thread_idx: int) -> None: MessageType.MigrationDone: self._handle_migration_done, MessageType.InitRecovery: self._handle_init_recovery, MessageType.ReadyAfterRecovery: self._handle_ready_after_recovery, + MessageType.RecoveryMigrationHashes: self._handle_recovery_migration_hashes, } @staticmethod @@ -400,6 +408,10 @@ async def _migration_decode_and_apply_plan(self, data: bytes) -> None: self._update_peers(peers) self.worker_operators = new_worker_operators + # Create the new partitions before MigrationInitDone: a peer can finish its + # rebuild and send migration batches before this worker has rebuilt. + for op_part in new_worker_operators: + self.local_state.add_new_operator_partition(op_part) await asyncio.sleep(0) async def _migration_background_repartition(self) -> None: @@ -748,6 +760,7 @@ async def _handle_init_recovery(self, data: bytes, _: MessageType) -> None: self.migration_completed.set() # unblock the zombie so it can exit await asyncio.sleep(0) # yield to let the zombie run its bail-out check self.migration_completed.clear() + self._reset_recovery_hash_exchange(tag=-1, hashes_out={}) ( self.id, @@ -832,22 +845,28 @@ async def _handle_init_recovery(self, data: bytes, _: MessageType) -> None: migration_meta = cloudpickle_deserialization(migration_blob) new_graph = migration_meta["new_graph"] - # Re-derive keys_to_send from local state + new partitioner - keys_to_send: dict[tuple, set] = defaultdict(set) + # Re-derive keys_to_send, and the destinations' remote_keys, from local state + new partitioner + keys_to_send: dict[OperatorPartition, set[tuple[Any, int]]] = defaultdict(set) + hashes_out: dict[OperatorPartition, dict[Any, tuple[int, int]]] = defaultdict(dict) for op_part in self.registered_operators: - op_name, partition = op_part - partitioner = new_graph.get_operator_by_name(op_name).get_partitioner() - op_data = self.local_state.get_operator_data_for_repartitioning(op_part) - for key in op_data: - new_part = partitioner.get_partition_no_cache(key) - if new_part != partition: - keys_to_send[op_part].add((key, new_part)) + part_keys_to_send, part_hashes = self.rehash_only( + op_part, + new_graph.get_operator_by_name(op_part[0]).get_partitioner(), + list(self.local_state.get_operator_data_for_repartitioning(op_part)), + self.id, + ) + for src_part, keys in part_keys_to_send.items(): + keys_to_send[src_part].update(keys) + for dest_part, key_map in part_hashes.items(): + hashes_out[dest_part].update(key_map) if keys_to_send: self.local_state.add_keys_to_send(keys_to_send) logging.warning( f"[RECOVERY] W{self.id} Re-derived {sum(len(v) for v in keys_to_send.values())} keys to send", ) + # Sent to the destinations on ReadyAfterRecovery, once every worker has loaded its snapshot. + self._reset_recovery_hash_exchange(tag=snapshot_id, hashes_out=hashes_out) restart_after_migration = True request_id_to_t_id_map = await self.get_sequencer_assignments_before_failure( @@ -889,17 +908,88 @@ async def _handle_init_recovery(self, data: bytes, _: MessageType) -> None: ) async def _handle_ready_after_recovery(self, _data: bytes, _: MessageType) -> None: + generation = self._migration_generation + protocol = self.function_execution_protocol + if self._recovery_hash_tag >= 0: + await self._exchange_recovery_migration_hashes() + if self._migration_generation != generation: + return # a newer recovery started while we waited for peers # All workers have stopped their old protocols by this point (the # coordinator only sends ReadyAfterRecovery after all workers reported). # Run the dedup scan NOW so it captures every message flushed by every # surviving worker's old egress. - await self.function_execution_protocol.egress.run_dedup_scan_and_mark_started() - self.function_execution_protocol.started.set() + await protocol.egress.run_dedup_scan_and_mark_started() + protocol.started.set() logging.warning( f"Worker: {self.id} recovered and ready at : {time.time() * 1000}", ) await asyncio.sleep(0) + def _reset_recovery_hash_exchange( + self, + tag: int, + hashes_out: dict[OperatorPartition, dict[Any, tuple[int, int]]], + ) -> None: + self._recovery_hash_tag = tag + self._recovery_hashes_out = hashes_out + self._recovery_hash_senders = set() + # Release anyone still waiting on a previous recovery's exchange; they + # see the generation change and bail out. + previous = self._recovery_hashes_received + self._recovery_hashes_received = asyncio.Event() + previous.set() + + async def _exchange_recovery_migration_hashes(self) -> None: + """Rebuild the destinations' remote_keys after a post-checkpoint recovery. + + Each worker tells every peer which keys of the peer's new partitions it + still holds, and waits for the same from all peers. Without remote_keys a + transaction on a key still at its source would read it as missing. This + runs on ReadyAfterRecovery, when every worker has loaded its snapshot, and + the protocol starts only after it completes. + """ + received = self._recovery_hashes_received + addr_to_peer = {(host, port): peer_id for peer_id, (host, port, _) in self.peers.items()} + per_peer: dict[int, list] = {peer_id: [] for peer_id in self.peers} + for dest_part, key_map in self._recovery_hashes_out.items(): + host, port, _ = self.dns[dest_part[0]][dest_part[1]] + peer_id = addr_to_peer.get((host, port)) + if peer_id is None: + # Destination partition is on this worker + self.local_state.add_remote_keys(dest_part, key_map) + else: + per_peer[peer_id].append((dest_part, key_map)) + self._recovery_hashes_out = {} + # Every peer gets a message, possibly empty: it is also the "done" signal. + await asyncio.gather( + *[ + self.networking.send_message( + host, + port, + msg=(self._recovery_hash_tag, self.id, per_peer[peer_id]), + msg_type=MessageType.RecoveryMigrationHashes, + serializer=Serializer.MSGPACK, + ) + for peer_id, (host, port, _) in self.peers.items() + ], + ) + if self._recovery_hash_senders >= self.peers.keys(): + received.set() + await received.wait() + logging.warning(f"[RECOVERY] W{self.id} remote_keys rebuilt from {len(self.peers)} peers") + + async def _handle_recovery_migration_hashes(self, data: bytes, _: MessageType) -> None: + # Lock-free: no awaits. + tag, sender_id, hashes = self.networking.decode_message(data) + if tag != self._recovery_hash_tag: + logging.warning(f"[RECOVERY] Dropping migration hashes from W{sender_id} for recovery {tag}") + return + for dest_part, key_map in hashes: + self.local_state.add_remote_keys(dest_part, key_map) + self._recovery_hash_senders.add(sender_id) + if self._recovery_hash_senders >= self.peers.keys(): + self._recovery_hashes_received.set() + async def _handle_unknown(self, _data: bytes, msg_type: MessageType) -> None: logging.error(f"Worker Service: Non supported command message type: {msg_type}") await asyncio.sleep(0) From 94685c3d2ba445865c739794c0697d3346a0306f Mon Sep 17 00:00:00 2001 From: Kyriakos Psarakis Date: Sat, 24 Oct 2026 16:18:02 +0200 Subject: [PATCH 03/20] Wait for in-flight keys and keep the checkpoint until the transfer is snapshotted A transaction that needed a migrating key could read it as missing: when the async sender had already taken the key into a batch, the local move found nothing and the remote source answered None, and in both cases the transaction ran without waiting for the batch. It now waits for the batch: the local move reports whether it moved the key, and a None response no longer wakes the waiter. Concurrent waiters on one key no longer crash on cleanup. The checkpoint was cleared on MigrationDone, so a snapshot cut before the last key landed but completed afterwards was saved without it, and recovery from it stranded keys at their old partitions. MigrationDone now carries the epoch of each worker's last batch, and the checkpoint stays on snapshots up to the latest such epoch. The coordinator always checks the latest snapshot for a checkpoint on recovery, restores the checkpoint and migration tracking while the transfer resumes, and retries a newer migration that had not reached its own checkpoint. Co-Authored-By: Claude Opus 5.5 --- coordinator/coordinator_metadata.py | 21 ++ coordinator/coordinator_service.py | 65 ++++-- coordinator/migration_metadata.py | 11 + styx-package/styx/common/base_state.py | 3 +- styx-package/styx/common/stateful_function.py | 115 ++++++----- styx-package/styx/common/tcp_networking.py | 15 +- .../test_migration_checkpoint_lifetime.py | 193 ++++++++++++++++++ .../styx_package/test_stateful_function.py | 20 ++ .../test_tcp_networking_coverage.py | 26 +++ tests/unit/worker/test_in_memory_state.py | 12 +- .../worker/test_migration_key_in_flight.py | 54 +++++ .../test_migration_snapshot_consistency.py | 13 ++ worker/operator_state/aria/in_memory_state.py | 11 +- worker/transactional_protocols/aria.py | 8 +- worker/worker_service.py | 4 + 15 files changed, 486 insertions(+), 85 deletions(-) create mode 100644 tests/unit/coordinator/test_migration_checkpoint_lifetime.py create mode 100644 tests/unit/worker/test_migration_key_in_flight.py diff --git a/coordinator/coordinator_metadata.py b/coordinator/coordinator_metadata.py index 71bf1f82..1e6fb6cc 100644 --- a/coordinator/coordinator_metadata.py +++ b/coordinator/coordinator_metadata.py @@ -68,6 +68,11 @@ def __init__(self, networking: NetworkingManager, s3_client: S3Client) -> None: # Checkpoint-and-resume migration FT self._migration_checkpoint_blob: bytes | None = None + # The graph the current checkpoint migrates to + self._migration_checkpoint_graph: StateflowGraph | None = None + # Set once every worker finished its transfer: snapshots carrying a later + # epoch were cut after the last migrated key landed and need no checkpoint. + self._migration_checkpoint_until_epoch: int | None = None self._migration_checkpoint_baseline_snap_id: int = -1 self._migration_checkpoint_snapshot_complete: asyncio.Event = asyncio.Event() @@ -241,6 +246,12 @@ def register_snapshot( self._migration_checkpoint_baseline_snap_id = -1 if current_completed_snapshot != self.prev_completed_snapshot_id: logging.warning(f"Cluster completed snapshot: {current_completed_snapshot}") + if ( + self._migration_checkpoint_until_epoch is not None + and self.completed_epoch_counter > self._migration_checkpoint_until_epoch + ): + # Cut after the last migrated key landed: a normal snapshot from here on. + self.clear_migration_checkpoint() # if we reached a complete snapshot, we could compact its deltas with the previous one sn_data: bytes = zstd_msgpack_serialization( ( @@ -372,10 +383,20 @@ def set_migration_checkpoint( "operator_partition_locations": operator_partition_locations, }, ) + self._migration_checkpoint_graph = new_graph + self._migration_checkpoint_until_epoch = None + + def end_migration_checkpoint_after(self, transfer_done_epoch: int) -> None: + """Every migrated key landed by transfer_done_epoch. Snapshots cut up to that + epoch may still miss keys, so they keep the checkpoint; it is cleared with + the first snapshot cut after it.""" + self._migration_checkpoint_until_epoch = transfer_done_epoch def clear_migration_checkpoint(self) -> None: """Clear migration metadata so subsequent snapshots are normal.""" self._migration_checkpoint_blob = None + self._migration_checkpoint_graph = None + self._migration_checkpoint_until_epoch = None async def submit_stateflow_graph( self, diff --git a/coordinator/coordinator_service.py b/coordinator/coordinator_service.py index 3a708197..23c27fad 100755 --- a/coordinator/coordinator_service.py +++ b/coordinator/coordinator_service.py @@ -784,7 +784,7 @@ async def _handle_deterministic_reordering(self, data: bytes) -> None: ) self.aria_metadata.cleanup() - async def _handle_migration_done(self, _: bytes) -> None: + async def _handle_migration_done(self, data: bytes) -> None: if not self.migration_in_progress: logging.warning("Dropping stale MigrationDone (no migration in progress)") return @@ -795,22 +795,27 @@ async def _handle_migration_done(self, _: bytes) -> None: if not self.migration_in_progress: logging.warning("Dropping stale MigrationDone (recovery intervened)") return - sync_complete: bool = await self.migration_metadata.set_empty_sync_done(mt) + (worker_done_epoch,) = self.protocol_networking.decode_message(data) + sync_complete: bool = await self.migration_metadata.migration_done(worker_done_epoch) logging.warning( f"MIGRATION | MigrationDone | {self.migration_metadata.sync_sum}", ) if not sync_complete: return + transfer_done_epoch = self.migration_metadata.transfer_done_epoch logging.warning( - f"MIGRATION_FINISHED at time: {time.time_ns() // 1_000_000}", + f"MIGRATION_FINISHED at time: {time.time_ns() // 1_000_000}" + f" (last key landed by epoch {transfer_done_epoch})", ) await self.migration_metadata.cleanup(mt) self.migration_in_progress = False - # Clear migration checkpoint blob so subsequent snapshots are normal. + # Snapshots cut before the last key landed may still complete after this + # point and must keep the checkpoint; it is cleared with the first + # snapshot cut after transfer_done_epoch. # The graph was already finalized in _wait_for_migration_checkpoint_and_finalize(). - self.coordinator.clear_migration_checkpoint() + self.coordinator.end_migration_checkpoint_after(transfer_done_epoch) # Defer final graph write until a post-migration snapshot captures # the state without migration metadata. This is a safety net — the @@ -1043,7 +1048,7 @@ async def _reset_after_recovery(self) -> None: def _read_migration_blob_from_s3(self, snap_id: int) -> bytes | None: """Read the migration checkpoint blob from the sequencer file in S3.""" - if snap_id <= 0: + if snap_id < 0: return None try: resp = self.s3_client.get_object( @@ -1133,6 +1138,7 @@ async def _perform_recovery(self, workers_to_remove: set[Worker]) -> None: or self.coordinator._migration_checkpoint_baseline_snap_id >= 0 # noqa: SLF001 ) saved_pending_graph = None + resumed_migration: dict | None = None snap_id = self.coordinator.get_current_completed_snapshot_id() if was_migrating: @@ -1142,21 +1148,28 @@ async def _perform_recovery(self, workers_to_remove: set[Worker]) -> None: f"current_completed_snapshot_id={snap_id}, " f"worker_snapshot_ids={self.coordinator.worker_snapshot_ids}", ) - migration_blob = self._read_migration_blob_from_s3(snap_id) - if migration_blob is not None: - # POST-CHECKPOINT: use NEW layout — workers will resume transfer - logging.warning("[RECOVERY] Migration checkpoint found — using NEW layout for recovery") - migration_meta = cloudpickle_deserialization(migration_blob) - new_graph = migration_meta["new_graph"] - self.coordinator.submitted_graph = new_graph - self.coordinator._pending_graph = None # noqa: SLF001 - # The worker pool already has the new operator assignments from - # update_stateflow_graph() — no need to revert. - else: - # PRE-CHECKPOINT: revert to old layout and retry migration - logging.warning("[RECOVERY] No migration checkpoint — reverting to OLD layout") - saved_pending_graph = self.coordinator._pending_graph # noqa: SLF001 - self.coordinator.revert_worker_pool_to_submitted_graph() + # Read the checkpoint even if no migration seems in progress: a snapshot cut + # before the last migrated key landed can complete after MigrationDone, and + # the workers resume from it either way. + migration_blob = self._read_migration_blob_from_s3(snap_id) + if migration_blob is not None: + # POST-CHECKPOINT: use NEW layout — workers will resume transfer + logging.warning("[RECOVERY] Migration checkpoint found — using NEW layout for recovery") + resumed_migration = cloudpickle_deserialization(migration_blob) + pending_graph = self.coordinator._pending_graph # noqa: SLF001 + if pending_graph is not None and pending_graph is not self.coordinator._migration_checkpoint_graph: # noqa: SLF001 + # A newer migration started but had not reached its own checkpoint: + # resume this one, then retry the newer one. + saved_pending_graph = pending_graph + self.coordinator.submitted_graph = resumed_migration["new_graph"] + self.coordinator._pending_graph = None # noqa: SLF001 + # Undo any operator promotions of the newer migration. + self.coordinator.revert_worker_pool_to_submitted_graph() + elif was_migrating: + # PRE-CHECKPOINT: revert to old layout and retry migration + logging.warning("[RECOVERY] No migration checkpoint — reverting to OLD layout") + saved_pending_graph = self.coordinator._pending_graph # noqa: SLF001 + self.coordinator.revert_worker_pool_to_submitted_graph() # Clear remaining migration state so stale messages are dropped. logging.warning("[RECOVERY] step 1c | clearing remaining migration state") @@ -1187,6 +1200,16 @@ async def _perform_recovery(self, workers_to_remove: set[Worker]) -> None: logging.warning("Cleaning up protocol after everyone is healthy") await self._reset_after_recovery() + # 4b) Resuming from a checkpoint: snapshots cut during the resumed transfer + # must still carry it, and the workers' MigrationDone must be counted. + # Restore both before the workers start (ReadyAfterRecovery). + if resumed_migration is not None: + self.coordinator.set_migration_checkpoint( + resumed_migration["new_graph"], + resumed_migration["operator_partition_locations"], + ) + self.migration_in_progress = True + # 5) Close all protocol connections (workers will reconnect clean) await self.protocol_networking.close_all_connections() diff --git a/coordinator/migration_metadata.py b/coordinator/migration_metadata.py index 7f41f675..cc3e8e40 100644 --- a/coordinator/migration_metadata.py +++ b/coordinator/migration_metadata.py @@ -17,6 +17,8 @@ def __init__(self, n_workers: int) -> None: self.output_offsets: dict[OperatorPartition, int] = {} self.epoch_counter: int = -1 self.t_counter: int = -1 + # Latest epoch in which a worker's last migrated key landed + self.transfer_done_epoch: int = -1 def check_sum(self, msg_type: MessageType) -> bool: return self.sync_sum[msg_type] == self.n_workers @@ -51,6 +53,13 @@ async def init_done( self.t_counter = max(t_counter, self.t_counter) return self.check_sum(MessageType.MigrationInitDone) + async def migration_done(self, epoch_counter: int) -> bool: + """A worker finished its async transfer; its last key landed by epoch_counter.""" + async with self.lock: + self.sync_sum[MessageType.MigrationDone] += 1 + self.transfer_done_epoch = max(epoch_counter, self.transfer_done_epoch) + return self.check_sum(MessageType.MigrationDone) + async def set_empty_sync_done(self, msg_type: MessageType) -> bool: async with self.lock: self.sync_sum[msg_type] += 1 @@ -64,3 +73,5 @@ async def cleanup(self, msg_type: MessageType) -> None: self.t_counter = -1 self.input_offsets.clear() self.output_offsets.clear() + elif msg_type == MessageType.MigrationDone: + self.transfer_done_epoch = -1 diff --git a/styx-package/styx/common/base_state.py b/styx-package/styx/common/base_state.py index 4f4e8abd..79578d8f 100644 --- a/styx-package/styx/common/base_state.py +++ b/styx-package/styx/common/base_state.py @@ -82,7 +82,8 @@ def migrate_within_the_same_worker( new_partition: int, key: K, old_partition: int, - ) -> None: + ) -> bool: + """Move the key from old_partition; False if it already left in a migration batch.""" raise NotImplementedError @abstractmethod diff --git a/styx-package/styx/common/stateful_function.py b/styx-package/styx/common/stateful_function.py index df654a18..f8115bcc 100644 --- a/styx-package/styx/common/stateful_function.py +++ b/styx-package/styx/common/stateful_function.py @@ -76,66 +76,83 @@ def __init__( self.__deployed_graph = deployed_graph self.__operator_lock: asyncio.Lock = operator_lock - async def __call__(self, *args, **kwargs) -> tuple[object, int]: # noqa: ANN002, ANN003 - """ - Executes the wrapped function and triggers asynchronous remote calls if needed. + async def __bring_key_here(self) -> None: + """Make sure the key is in this partition before the function runs: during a + migration it may still be in another partition, here or on a peer, or in a + migration batch on its way here.""" + # 1) Decide whether we need to fetch/migrate the key, while holding the lock briefly. + need_remote_fetch: bool = False + key_in_flight: bool = False + remote_info: tuple[int, int] | None = None # (worker_id, old_partition) shape depends on your state + + async with self.__operator_lock: + if self.__state.in_remote_keys( + self.__key, + self.__operator_name, + self.__partition, + ): + remote_info = self.__state.get_worker_id_old_partition( + self.__operator_name, + self.__partition, + self.__key, + ) - Returns: - (result, n_remote_calls) on success - (exception, -1) on failure - """ - try: - # 1) Decide whether we need to fetch/migrate the key, while holding the lock briefly. - need_remote_fetch: bool = False - remote_info: tuple[int, int] | None = None # (worker_id, old_partition) shape depends on your state + # Your state seems to return something like (worker_id, old_partition) + # since you later access remote_info[1] as the old partition. + old_partition = remote_info[1] - async with self.__operator_lock: - if self.__state.in_remote_keys( - self.__key, + is_local: bool = ( self.__operator_name, - self.__partition, - ): - remote_info = self.__state.get_worker_id_old_partition( + old_partition, + ) in self.__state.operator_partitions + if is_local: + # Transfer the key from another partition within the same worker + moved = self.__state.migrate_within_the_same_worker( self.__operator_name, self.__partition, self.__key, - ) - - # Your state seems to return something like (worker_id, old_partition) - # since you later access remote_info[1] as the old partition. - old_partition = remote_info[1] - - is_local: bool = ( - self.__operator_name, old_partition, - ) in self.__state.operator_partitions - if is_local: - # Transfer the key from another partition within the same worker - self.__state.migrate_within_the_same_worker( + ) + if not moved: + # The async sender already took the key into a batch to this + # worker: wait for the batch rather than read the key as missing. + self.__networking.expect_remote_key( self.__operator_name, self.__partition, self.__key, - old_partition, ) - else: - # We'll do networking outside the lock. - need_remote_fetch = True - - # 2) If needed, fetch the key from a remote worker without holding the operator lock. - if need_remote_fetch: - if remote_info is None: - logging.warning("Remote fetch failed: remote_info is None") - await self.__networking.request_key( - self.__operator_name, - self.__partition, - self.__key, - remote_info, - ) - await self.__networking.wait_for_remote_key_event( - self.__operator_name, - self.__partition, - self.__key, - ) + key_in_flight = True + else: + # We'll do networking outside the lock. + need_remote_fetch = True + + # 2) If needed, fetch the key from a remote worker without holding the operator lock. + if need_remote_fetch: + if remote_info is None: + logging.warning("Remote fetch failed: remote_info is None") + await self.__networking.request_key( + self.__operator_name, + self.__partition, + self.__key, + remote_info, + ) + if need_remote_fetch or key_in_flight: + await self.__networking.wait_for_remote_key_event( + self.__operator_name, + self.__partition, + self.__key, + ) + + async def __call__(self, *args, **kwargs) -> tuple[object, int]: # noqa: ANN002, ANN003 + """ + Executes the wrapped function and triggers asynchronous remote calls if needed. + + Returns: + (result, n_remote_calls) on success + (exception, -1) on failure + """ + try: + await self.__bring_key_here() # 3) Run user logic while holding the operator lock. async with self.__operator_lock: diff --git a/styx-package/styx/common/tcp_networking.py b/styx-package/styx/common/tcp_networking.py index 0fc923e3..25e10673 100644 --- a/styx-package/styx/common/tcp_networking.py +++ b/styx-package/styx/common/tcp_networking.py @@ -403,10 +403,17 @@ async def wait_for_remote_key_event( operator_partition = (operator_name, partition) ev = self.wait_remote_key_event[operator_partition][key] await ev.wait() - # (6) cleanup to prevent unbounded growth - del self.wait_remote_key_event[operator_partition][key] - if not self.wait_remote_key_event[operator_partition]: - del self.wait_remote_key_event[operator_partition] + # (6) cleanup to prevent unbounded growth. Several transactions can wait + # on the same key; only the first one to wake removes the event. + events = self.wait_remote_key_event.get(operator_partition) + if events is not None and events.get(key) is ev: + del events[key] + if not events: + del self.wait_remote_key_event[operator_partition] + + def expect_remote_key(self, operator_name: str, partition: int, key: K) -> None: + """Register a wait for a key already on its way here in a migration batch.""" + self.wait_remote_key_event.setdefault((operator_name, partition), {}).setdefault(key, asyncio.Event()) def key_received(self, operator_partition: OperatorPartition, key: K) -> None: operator_partition = tuple(operator_partition) diff --git a/tests/unit/coordinator/test_migration_checkpoint_lifetime.py b/tests/unit/coordinator/test_migration_checkpoint_lifetime.py new file mode 100644 index 00000000..3a3de155 --- /dev/null +++ b/tests/unit/coordinator/test_migration_checkpoint_lifetime.py @@ -0,0 +1,193 @@ +"""How long a migration checkpoint stays attached to snapshots. + +A snapshot cut before the last migrated key landed still has keys at their old +partitions, so its sequencer file must carry the checkpoint even if it completes +after MigrationDone. Recovery from it resumes the transfer, and the coordinator +must keep tracking that resumed transfer. +""" + +import asyncio +from collections import defaultdict +from unittest.mock import AsyncMock, MagicMock + +from styx.common.local_state_backends import LocalStateBackend +from styx.common.message_types import MessageType +from styx.common.operator import Operator +from styx.common.serialization import Serializer, cloudpickle_serialization, zstd_msgpack_deserialization +from styx.common.stateflow_graph import StateflowGraph +from styx.common.tcp_networking import NetworkingManager + +from coordinator.coordinator_metadata import Coordinator +from coordinator.coordinator_service import CoordinatorService +from coordinator.migration_metadata import MigrationMetadata + + +def _graph(n_partitions: int) -> StateflowGraph: + g = StateflowGraph("app", operator_state_backend=LocalStateBackend.DICT, max_operator_parallelism=8) + g.add_operators(Operator("ycsb", n_partitions=n_partitions)) + return g + + +def _complete_snapshot(c: Coordinator, snapshot_id: int, epoch: int): + """All (two) workers register snapshot_id; returns the sequencer tuple written to S3.""" + c.s3_client.put_object.reset_mock() + for worker_id in (0, 1): + c.register_snapshot(worker_id, snapshot_id, {}, {}, epoch, 0, MagicMock()) + body = c.s3_client.put_object.call_args.kwargs["Body"] + return zstd_msgpack_deserialization(body) + + +def _coordinator_with_checkpoint() -> Coordinator: + c = Coordinator(MagicMock(), MagicMock()) + c.s3_client = MagicMock() + c.worker_snapshot_ids = {0: 2, 1: 2} + c.prev_completed_snapshot_id = 2 + c.set_migration_checkpoint(_graph(8), {}) + return c + + +class TestCheckpointLifetime: + def test_snapshot_cut_before_last_key_landed_keeps_checkpoint(self): + c = _coordinator_with_checkpoint() + c.end_migration_checkpoint_after(transfer_done_epoch=50) + + # Cut at the end of epoch 49 (carries 50), completes after MigrationDone. + assert _complete_snapshot(c, 3, epoch=50)[4] is not None + # Cut at the end of epoch 50 (carries 51): every key had landed. + assert _complete_snapshot(c, 4, epoch=51)[4] is None + assert _complete_snapshot(c, 5, epoch=60)[4] is None + + def test_checkpoint_kept_while_transfer_runs(self): + c = _coordinator_with_checkpoint() + assert _complete_snapshot(c, 3, epoch=500)[4] is not None + + def test_new_checkpoint_resets_the_end(self): + c = _coordinator_with_checkpoint() + c.end_migration_checkpoint_after(transfer_done_epoch=50) + c.set_migration_checkpoint(_graph(16), {}) + assert _complete_snapshot(c, 3, epoch=60)[4] is not None + + +class TestMigrationDoneEpochs: + async def test_transfer_done_epoch_is_the_latest_reported(self): + m = MigrationMetadata(2) + assert await m.migration_done(12) is False + assert await m.migration_done(10) is True + assert m.transfer_done_epoch == 12 + await m.cleanup(MessageType.MigrationDone) + assert m.transfer_done_epoch == -1 + + +def _service(coordinator: Coordinator) -> CoordinatorService: + svc = CoordinatorService.__new__(CoordinatorService) + svc.coordinator = coordinator + svc.migration_in_progress = False + svc.migration_metadata = MigrationMetadata(2) + svc.networking_locks = defaultdict(asyncio.Lock) + svc._migration_checkpoint_task = None + svc._deferred_migration_graph = None + svc.networking = MagicMock(close_worker_connections=AsyncMock()) + svc.protocol_networking = MagicMock( + close_worker_connections=AsyncMock(), + close_all_connections=AsyncMock(), + decode_message=NetworkingManager.decode_message, + ) + svc.aio_task_scheduler = MagicMock(close=AsyncMock()) + svc._reset_after_recovery = AsyncMock() + coordinator.start_recovery_process = AsyncMock() + coordinator.wait_cluster_healthy = AsyncMock() + coordinator.notify_cluster_healthy = AsyncMock() + coordinator.revert_worker_pool_to_submitted_graph = MagicMock() + return svc + + +def _migration_done_msg(epoch: int) -> bytes: + return NetworkingManager.encode_message( + msg=(epoch,), + msg_type=MessageType.MigrationDone, + serializer=Serializer.MSGPACK, + )[8:] + + +class TestMigrationDoneHandler: + async def test_all_workers_done_ends_checkpoint_after_latest_epoch(self): + c = _coordinator_with_checkpoint() + svc = _service(c) + svc.migration_in_progress = True + + await svc._handle_migration_done(_migration_done_msg(41)) + assert svc.migration_in_progress is True + await svc._handle_migration_done(_migration_done_msg(43)) + + assert svc.migration_in_progress is False + assert c._migration_checkpoint_blob is not None # kept for snapshots cut up to 43 + assert c._migration_checkpoint_until_epoch == 43 + + +def _blob(graph: StateflowGraph) -> bytes: + return cloudpickle_serialization({"new_graph": graph, "operator_partition_locations": {}}) + + +class TestRecoveryFromCheckpoint: + async def test_resume_after_migration_done_keeps_tracking_the_transfer(self): + c = Coordinator(MagicMock(), MagicMock()) + c.worker_snapshot_ids = {0: 7, 1: 7} + svc = _service(c) + # MigrationDone already processed; the latest snapshot still has the checkpoint. + svc._read_migration_blob_from_s3 = MagicMock(return_value=_blob(_graph(8))) + state_when_workers_released = {} + + async def notify() -> None: + state_when_workers_released["in_progress"] = svc.migration_in_progress + state_when_workers_released["blob"] = c._migration_checkpoint_blob + + c.notify_cluster_healthy = AsyncMock(side_effect=notify) + await svc._perform_recovery({MagicMock(name="dead_worker")}) + + svc._read_migration_blob_from_s3.assert_called_once_with(7) + assert c.submitted_graph.nodes["ycsb"].n_partitions == 8 + assert state_when_workers_released["in_progress"] is True + assert state_when_workers_released["blob"] is not None + + async def test_newer_migration_without_checkpoint_is_retried_after_resume(self): + c = Coordinator(MagicMock(), MagicMock()) + c.worker_snapshot_ids = {0: 7, 1: 7} + newer = _graph(16) + c._pending_graph = newer + svc = _service(c) + svc.migration_in_progress = True + svc._read_migration_blob_from_s3 = MagicMock(return_value=_blob(_graph(8))) + + await svc._perform_recovery({MagicMock(name="dead_worker")}) + + assert svc._deferred_migration_graph is newer + assert c.submitted_graph.nodes["ycsb"].n_partitions == 8 + c.revert_worker_pool_to_submitted_graph.assert_called_once() + + async def test_crash_before_checkpoint_reverts_and_retries(self): + c = Coordinator(MagicMock(), MagicMock()) + c.worker_snapshot_ids = {0: 7, 1: 7} + pending = _graph(8) + c._pending_graph = pending + svc = _service(c) + svc.migration_in_progress = True + svc._read_migration_blob_from_s3 = MagicMock(return_value=None) + + await svc._perform_recovery({MagicMock(name="dead_worker")}) + + assert svc.migration_in_progress is False + assert svc._deferred_migration_graph is pending + assert c._migration_checkpoint_blob is None + c.revert_worker_pool_to_submitted_graph.assert_called_once() + + async def test_ordinary_recovery_stays_ordinary(self): + c = Coordinator(MagicMock(), MagicMock()) + c.worker_snapshot_ids = {0: 7, 1: 7} + svc = _service(c) + svc._read_migration_blob_from_s3 = MagicMock(return_value=None) + + await svc._perform_recovery({MagicMock(name="dead_worker")}) + + assert svc.migration_in_progress is False + assert svc._deferred_migration_graph is None + c.revert_worker_pool_to_submitted_graph.assert_not_called() diff --git a/tests/unit/styx_package/test_stateful_function.py b/tests/unit/styx_package/test_stateful_function.py index 33c0f56f..a79dc251 100644 --- a/tests/unit/styx_package/test_stateful_function.py +++ b/tests/unit/styx_package/test_stateful_function.py @@ -201,6 +201,26 @@ async def mock_run(*args): state.migrate_within_the_same_worker.assert_called_once_with("users", 0, "k1", 1) assert result == "migrated" + @pytest.mark.asyncio + async def test_call_local_migration_key_in_flight_waits_for_batch(self): + sf, state, networking, *_ = _make_sf() + state.in_remote_keys.return_value = True + state.get_worker_id_old_partition.return_value = (0, 1) + state.operator_partitions = {("users", 1)} # old partition is local + # The async sender already took the key into a batch to this worker. + state.migrate_within_the_same_worker.return_value = False + networking.wait_for_remote_key_event = AsyncMock() + + async def mock_run(*args): + return "migrated" + + sf.run = mock_run + result, _n_calls = await sf() + networking.expect_remote_key.assert_called_once_with("users", 0, "k1") + networking.wait_for_remote_key_event.assert_awaited_once_with("users", 0, "k1") + networking.request_key.assert_not_called() + assert result == "migrated" + @pytest.mark.asyncio async def test_call_remote_migration(self): sf, state, networking, *_ = _make_sf() diff --git a/tests/unit/styx_package/test_tcp_networking_coverage.py b/tests/unit/styx_package/test_tcp_networking_coverage.py index 19d2ca55..b5d55dc1 100644 --- a/tests/unit/styx_package/test_tcp_networking_coverage.py +++ b/tests/unit/styx_package/test_tcp_networking_coverage.py @@ -377,3 +377,29 @@ async def test_none_response_raises(self): with pytest.raises(ConnectionError): await nm.send_message_request_response("10.0.0.2", 6001, ("hello",), MessageType.ClientMsg) + + @pytest.mark.asyncio + async def test_two_waiters_on_same_key_both_wake(self): + nm = _nm() + nm.expect_remote_key("users", 0, "k1") + waiters = [asyncio.create_task(nm.wait_for_remote_key_event("users", 0, "k1")) for _ in range(2)] + await asyncio.sleep(0) + nm.wait_remote_key_event[("users", 0)]["k1"].set() + await asyncio.wait_for(asyncio.gather(*waiters), timeout=1) + assert ("users", 0) not in nm.wait_remote_key_event + + +class TestExpectRemoteKey: + def test_registers_wait_without_sending(self): + nm = _nm() + nm.send_message = AsyncMock() + nm.expect_remote_key("users", 0, "k1") + assert isinstance(nm.wait_remote_key_event[("users", 0)]["k1"], asyncio.Event) + nm.send_message.assert_not_called() + + def test_keeps_an_existing_wait(self): + nm = _nm() + ev = asyncio.Event() + nm.wait_remote_key_event[("users", 0)] = {"k1": ev} + nm.expect_remote_key("users", 0, "k1") + assert nm.wait_remote_key_event[("users", 0)]["k1"] is ev diff --git a/tests/unit/worker/test_in_memory_state.py b/tests/unit/worker/test_in_memory_state.py index ca086764..9df6ca4d 100644 --- a/tests/unit/worker/test_in_memory_state.py +++ b/tests/unit/worker/test_in_memory_state.py @@ -478,10 +478,20 @@ def test_moves_key_between_partitions(self): s.data[(OP, 0)]["k1"] = "val" s.keys_to_send[(OP, 0)] = {("k1", 1)} s.remote_keys[(OP, 1)] = {"k1": (0, 0)} - s.migrate_within_the_same_worker(OP, 1, "k1", 0) + assert s.migrate_within_the_same_worker(OP, 1, "k1", 0) is True assert "k1" not in s.data[(OP, 0)] assert s.data[(OP, 1)]["k1"] == "val" + def test_key_already_in_a_migration_batch_is_not_moved(self): + s = InMemoryOperatorState({(OP, 0), (OP, 1)}) + s.data[(OP, 0)]["k1"] = "val" + s.keys_to_send[(OP, 0)] = {("k1", 1)} + s.remote_keys[(OP, 1)] = {"k1": (0, 0)} + s.get_async_migrate_batch(10) # the async sender takes it + assert s.migrate_within_the_same_worker(OP, 1, "k1", 0) is False + assert "k1" not in s.data[(OP, 1)] + assert s.in_remote_keys("k1", OP, 1) # still expected: the batch will deliver it + # --------------------------------------------------------------------------- # set_batch_data_from_migration diff --git a/tests/unit/worker/test_migration_key_in_flight.py b/tests/unit/worker/test_migration_key_in_flight.py new file mode 100644 index 00000000..a6892d3c --- /dev/null +++ b/tests/unit/worker/test_migration_key_in_flight.py @@ -0,0 +1,54 @@ +"""A transaction must not read a migrating key as missing while it is in a batch. + +When a destination asks the source for a key (RequestRemoteKey) after the async +sender already took that key into a batch, the source answers None. The waiting +transaction must keep waiting for the batch instead of running without the key. +""" + +import asyncio +from types import SimpleNamespace +from unittest.mock import MagicMock + +from styx.common.message_types import MessageType +from styx.common.serialization import Serializer +from styx.common.tcp_networking import NetworkingManager + +from worker.operator_state.aria.in_memory_state import InMemoryOperatorState +from worker.worker_service import Worker + +OP = "op" +DST = (OP, 1) + + +def _receive_remote_key_msg(key, value) -> bytes: + # Drop the 8-byte length prefix, as the receiving TCP reader does. + return NetworkingManager.encode_message( + msg=(DST, key, value), + msg_type=MessageType.ReceiveRemoteKey, + serializer=Serializer.MSGPACK, + )[8:] + + +def _worker(): + w = Worker.__new__(Worker) + w.networking = SimpleNamespace(decode_message=NetworkingManager.decode_message) + w.protocol_networking = MagicMock() + w.protocol_networking.wait_remote_key_event = {DST: {"k1": asyncio.Event()}} + w.local_state = InMemoryOperatorState({DST}) + w.local_state.add_remote_keys(DST, {"k1": (2, 0)}) + return w + + +class TestRemoteKeyResponse: + async def test_none_response_leaves_waiter_for_the_batch(self): + w = _worker() + await w._handle_receive_remote_key(_receive_remote_key_msg("k1", None), MessageType.ReceiveRemoteKey) + w.protocol_networking.key_received.assert_not_called() + assert w.local_state.in_remote_keys("k1", OP, 1) + + async def test_value_response_installs_key_and_wakes_waiter(self): + w = _worker() + await w._handle_receive_remote_key(_receive_remote_key_msg("k1", "v1"), MessageType.ReceiveRemoteKey) + w.protocol_networking.key_received.assert_called_once_with(DST, "k1") + assert w.local_state.data[DST]["k1"] == "v1" + assert not w.local_state.in_remote_keys("k1", OP, 1) diff --git a/tests/unit/worker/test_migration_snapshot_consistency.py b/tests/unit/worker/test_migration_snapshot_consistency.py index 8481f614..f60ae1d2 100644 --- a/tests/unit/worker/test_migration_snapshot_consistency.py +++ b/tests/unit/worker/test_migration_snapshot_consistency.py @@ -227,3 +227,16 @@ async def test_drain_returns_immediately_without_batches(self): _link, src, _dst = _two_workers(0) await asyncio.wait_for(src._drain_migration_batches(), timeout=0.1) assert not src._migration_gate.is_set() + + +class TestWaitingTransactionWokenByBatch: + async def test_waiter_for_key_in_flight_wakes_with_key_present(self): + _link, src, dst = _two_workers(3) + # A transaction at the destination waits for key 1, already on its way. + waiter = asyncio.Event() + dst.networking.wait_remote_key_event[DST] = {1: waiter} + + sender = asyncio.create_task(src._continuous_migration_sender()) + await asyncio.wait_for(waiter.wait(), timeout=2) + assert dst.local_state.data[DST][1] == "v1" + await asyncio.wait_for(sender, timeout=2) diff --git a/worker/operator_state/aria/in_memory_state.py b/worker/operator_state/aria/in_memory_state.py index ba721125..84e94b47 100644 --- a/worker/operator_state/aria/in_memory_state.py +++ b/worker/operator_state/aria/in_memory_state.py @@ -82,13 +82,12 @@ def migrate_within_the_same_worker( new_partition: int, key: K, old_partition: int, - ) -> None: + ) -> bool: new_operator_partition: OperatorPartition = (operator_name, new_partition) - self.set_data_from_migration( - new_operator_partition, - key, - self.get_key_to_migrate(new_operator_partition, key, old_partition), - ) + data = self.get_key_to_migrate(new_operator_partition, key, old_partition) + # None: the async sender already took the key into a batch to this worker. + self.set_data_from_migration(new_operator_partition, key, data) + return data is not None def keys_remaining_to_remote(self) -> int: c = 0 diff --git a/worker/transactional_protocols/aria.py b/worker/transactional_protocols/aria.py index a6be34e0..b8df239c 100644 --- a/worker/transactional_protocols/aria.py +++ b/worker/transactional_protocols/aria.py @@ -607,15 +607,17 @@ async def _continuous_migration_sender(self) -> None: continue if not self.local_state.has_keys_to_send(): if self.local_state.keys_remaining_to_send() == 0: - # Report done only once every batch has landed. + # Report done only once every batch has landed. The last one + # landed in this epoch or earlier (a batch never crosses an + # epoch's drain), so snapshots cut after this epoch are complete. await self._migration_all_acked.wait() self.migrating_state = False await self.networking.send_message( DISCOVERY_HOST, DISCOVERY_PORT + 1, - msg=b"", + msg=(self.sequencer.epoch_counter,), msg_type=MessageType.MigrationDone, - serializer=Serializer.NONE, + serializer=Serializer.MSGPACK, ) logging.warning( f"MIGRATION_FINISHED at epoch: {self.sequencer.epoch_counter}" diff --git a/worker/worker_service.py b/worker/worker_service.py index ac980979..c81fa51a 100644 --- a/worker/worker_service.py +++ b/worker/worker_service.py @@ -610,6 +610,10 @@ async def _handle_receive_remote_key( # Lock-free: no awaits. del msg_type operator_partition, key, value = self.networking.decode_message(data) + if value is None: + # The source's async sender already took the key into a batch; the + # batch wakes the waiting transaction when it lands here. + return self.local_state.set_data_from_migration(operator_partition, key, value) # Guard: the event may have been already signaled by the async # migration handler, or may not exist if no transaction requested it. From 68340416611ed6e916f9470b3212ba4bc6771d07 Mon Sep 17 00:00:00 2001 From: Kyriakos Psarakis Date: Sat, 24 Oct 2026 16:57:55 +0200 Subject: [PATCH 04/20] Migrate None-valued keys and load the migration demos' initial data Migration treated a key holding None as absent: the async sender popped and dropped it, and the on-demand path reported it as already gone. The destination still listed the key in remote_keys, so a transaction on it waited for a batch that never came. get_key_to_migrate now reports whether it found the key separately from its value, None-valued keys are sent like any other, and the remote reply carries the found flag. The migration YCSB, migration TPC-C and scalability demos wrote their initial data to S3 but never sent InitDataComplete, so workers deployed from snapshot -1 with empty state and the E2E tests ran against an empty store. They now notify the coordinator like the other demos. Co-Authored-By: Claude Opus 5.5 --- demo/demo-migration-ft-ycsb/client.py | 1 + demo/demo-migration-tpc-c/pure_kafka_demo.py | 1 + demo/demo-migration-ycsb/client.py | 1 + demo/demo-scalability/client.py | 1 + tests/unit/worker/test_in_memory_state.py | 21 +++++++- .../worker/test_migration_fault_tolerance.py | 20 +++++-- .../worker/test_migration_key_in_flight.py | 17 ++++-- .../test_migration_snapshot_consistency.py | 12 +++++ worker/operator_state/aria/in_memory_state.py | 53 +++++++++++-------- worker/worker_service.py | 10 ++-- 10 files changed, 98 insertions(+), 39 deletions(-) diff --git a/demo/demo-migration-ft-ycsb/client.py b/demo/demo-migration-ft-ycsb/client.py index 3890e3d8..c0e99e33 100644 --- a/demo/demo-migration-ft-ycsb/client.py +++ b/demo/demo-migration-ft-ycsb/client.py @@ -129,6 +129,7 @@ def ycsb_init(styx: SyncStyxClient, operator: Operator, num_workers: int = None) print("Data ready") for partition, partition_data in partitions.items(): styx.init_data(operator, partition, partition_data) + styx.notify_init_data_complete() print("Data loaded") time.sleep(5) submit_graph(styx) diff --git a/demo/demo-migration-tpc-c/pure_kafka_demo.py b/demo/demo-migration-tpc-c/pure_kafka_demo.py index ee92dda0..4718ba4c 100644 --- a/demo/demo-migration-tpc-c/pure_kafka_demo.py +++ b/demo/demo-migration-tpc-c/pure_kafka_demo.py @@ -379,6 +379,7 @@ def tpc_c_init(styx: SyncStyxClient): populate_order_line(styx) populate_item(styx) populate_stock(styx) + styx.notify_init_data_complete() # Sleep for 5 seconds to be sure that all the data are in minio time.sleep(5) submit_graph(styx) diff --git a/demo/demo-migration-ycsb/client.py b/demo/demo-migration-ycsb/client.py index 1937b087..cfbc533f 100644 --- a/demo/demo-migration-ycsb/client.py +++ b/demo/demo-migration-ycsb/client.py @@ -111,6 +111,7 @@ def ycsb_init(styx: SyncStyxClient, operator: Operator, num_workers: int = None) print("Data ready") for partition, partition_data in partitions.items(): styx.init_data(operator, partition, partition_data) + styx.notify_init_data_complete() print("Data loaded") time.sleep(5) submit_graph(styx) diff --git a/demo/demo-scalability/client.py b/demo/demo-scalability/client.py index 6b60ff74..10dad9a4 100644 --- a/demo/demo-scalability/client.py +++ b/demo/demo-scalability/client.py @@ -59,6 +59,7 @@ def ycsb_init(styx: SyncStyxClient, operator: Operator): for partition, partition_data in partitions.items(): styx.init_data(operator, partition, partition_data) + styx.notify_init_data_complete() time.sleep(5) submit_graph(styx) diff --git a/tests/unit/worker/test_in_memory_state.py b/tests/unit/worker/test_in_memory_state.py index 9df6ca4d..6f018c89 100644 --- a/tests/unit/worker/test_in_memory_state.py +++ b/tests/unit/worker/test_in_memory_state.py @@ -440,11 +440,20 @@ def test_pops_key_and_removes_from_keys_to_send(self): s = _state() s.data[OP_PART]["k1"] = "v1" s.keys_to_send[OP_PART] = {("k1", 2)} - val = s.get_key_to_migrate((OP, 2), "k1", PART) - assert val == "v1" + assert s.get_key_to_migrate((OP, 2), "k1", PART) == (True, "v1") assert "k1" not in s.data[OP_PART] assert ("k1", 2) not in s.keys_to_send[OP_PART] + def test_none_value_is_found(self): + s = _state() + s.data[OP_PART]["k1"] = None + assert s.get_key_to_migrate((OP, 2), "k1", PART) == (True, None) + assert "k1" not in s.data[OP_PART] + + def test_absent_key_is_not_found(self): + s = _state() + assert s.get_key_to_migrate((OP, 2), "k1", PART) == (False, None) + # --------------------------------------------------------------------------- # set_data_from_migration @@ -482,6 +491,14 @@ def test_moves_key_between_partitions(self): assert "k1" not in s.data[(OP, 0)] assert s.data[(OP, 1)]["k1"] == "val" + def test_key_with_none_value_is_moved(self): + s = InMemoryOperatorState({(OP, 0), (OP, 1)}) + s.data[(OP, 0)]["k1"] = None + s.remote_keys[(OP, 1)] = {"k1": (0, 0)} + assert s.migrate_within_the_same_worker(OP, 1, "k1", 0) is True + assert s.data[(OP, 1)] == {"k1": None} + assert not s.in_remote_keys("k1", OP, 1) + def test_key_already_in_a_migration_batch_is_not_moved(self): s = InMemoryOperatorState({(OP, 0), (OP, 1)}) s.data[(OP, 0)]["k1"] = "val" diff --git a/tests/unit/worker/test_migration_fault_tolerance.py b/tests/unit/worker/test_migration_fault_tolerance.py index cc2ce33d..abbbc178 100644 --- a/tests/unit/worker/test_migration_fault_tolerance.py +++ b/tests/unit/worker/test_migration_fault_tolerance.py @@ -34,10 +34,12 @@ def test_single_key_written_to_delta_map(self): s.set_data_from_migration(OP_PART_1, "k1", "v1") assert s.delta_map[OP_PART_1]["k1"] == "v1" - def test_none_data_not_written_to_delta_map(self): + def test_none_value_written_to_delta_map(self): + # A key may hold None; it migrates like any value (and snapshots as a tombstone). s = _state(OP_PART_0, OP_PART_1) s.set_data_from_migration(OP_PART_1, "k1", None) - assert "k1" not in s.delta_map[OP_PART_1] + assert s.data[OP_PART_1] == {"k1": None} + assert s.delta_map[OP_PART_1] == {"k1": None} def test_data_written_to_both_data_and_delta_map(self): s = _state(OP_PART_0, OP_PART_1) @@ -81,14 +83,14 @@ def test_tombstone_written_to_delta_map(self): s = _state(OP_PART_0, OP_PART_1) s.data[OP_PART_0]["k1"] = "v1" result = s.get_key_to_migrate(OP_PART_1, "k1", PART_0) - assert result == "v1" + assert result == (True, "v1") assert s.delta_map[OP_PART_0]["k1"] is None def test_no_tombstone_when_key_already_gone(self): s = _state(OP_PART_0, OP_PART_1) # key not in data result = s.get_key_to_migrate(OP_PART_1, "k1", PART_0) - assert result is None + assert result == (False, None) assert "k1" not in s.delta_map[OP_PART_0] @@ -104,6 +106,16 @@ def test_tombstones_written_for_migrated_keys(self): # The batch should contain the values for the destination assert len(batch[(OP, PART_1)]) == 3 + def test_none_valued_key_is_sent(self): + # The destination expects every key in keys_to_send, whatever its value. + s = _state(OP_PART_0, OP_PART_1) + s.data[OP_PART_0] = {"k1": None} + s.keys_to_send = {OP_PART_0: {("k1", PART_1)}} + batch = s.get_async_migrate_batch(batch_size=10) + assert batch[(OP, PART_1)] == {"k1": None} + assert s.delta_map[OP_PART_0]["k1"] is None + assert "k1" not in s.data[OP_PART_0] + def test_no_tombstone_for_already_deleted_key(self): s = _state(OP_PART_0, OP_PART_1) s.data[OP_PART_0] = {} # key not present diff --git a/tests/unit/worker/test_migration_key_in_flight.py b/tests/unit/worker/test_migration_key_in_flight.py index a6892d3c..d0eb6944 100644 --- a/tests/unit/worker/test_migration_key_in_flight.py +++ b/tests/unit/worker/test_migration_key_in_flight.py @@ -20,10 +20,10 @@ DST = (OP, 1) -def _receive_remote_key_msg(key, value) -> bytes: +def _receive_remote_key_msg(key, found, value) -> bytes: # Drop the 8-byte length prefix, as the receiving TCP reader does. return NetworkingManager.encode_message( - msg=(DST, key, value), + msg=(DST, key, found, value), msg_type=MessageType.ReceiveRemoteKey, serializer=Serializer.MSGPACK, )[8:] @@ -40,15 +40,22 @@ def _worker(): class TestRemoteKeyResponse: - async def test_none_response_leaves_waiter_for_the_batch(self): + async def test_not_found_response_leaves_waiter_for_the_batch(self): w = _worker() - await w._handle_receive_remote_key(_receive_remote_key_msg("k1", None), MessageType.ReceiveRemoteKey) + await w._handle_receive_remote_key(_receive_remote_key_msg("k1", False, None), MessageType.ReceiveRemoteKey) w.protocol_networking.key_received.assert_not_called() assert w.local_state.in_remote_keys("k1", OP, 1) async def test_value_response_installs_key_and_wakes_waiter(self): w = _worker() - await w._handle_receive_remote_key(_receive_remote_key_msg("k1", "v1"), MessageType.ReceiveRemoteKey) + await w._handle_receive_remote_key(_receive_remote_key_msg("k1", True, "v1"), MessageType.ReceiveRemoteKey) w.protocol_networking.key_received.assert_called_once_with(DST, "k1") assert w.local_state.data[DST]["k1"] == "v1" assert not w.local_state.in_remote_keys("k1", OP, 1) + + async def test_found_none_value_installs_key_and_wakes_waiter(self): + w = _worker() + await w._handle_receive_remote_key(_receive_remote_key_msg("k1", True, None), MessageType.ReceiveRemoteKey) + w.protocol_networking.key_received.assert_called_once_with(DST, "k1") + assert w.local_state.data[DST] == {"k1": None} + assert not w.local_state.in_remote_keys("k1", OP, 1) diff --git a/tests/unit/worker/test_migration_snapshot_consistency.py b/tests/unit/worker/test_migration_snapshot_consistency.py index f60ae1d2..7803a533 100644 --- a/tests/unit/worker/test_migration_snapshot_consistency.py +++ b/tests/unit/worker/test_migration_snapshot_consistency.py @@ -240,3 +240,15 @@ async def test_waiter_for_key_in_flight_wakes_with_key_present(self): await asyncio.wait_for(waiter.wait(), timeout=2) assert dst.local_state.data[DST][1] == "v1" await asyncio.wait_for(sender, timeout=2) + + async def test_waiter_for_none_valued_key_wakes(self): + # A key can hold None; it must still be delivered and wake its waiter. + _link, src, dst = _two_workers(3) + src.local_state.data[SRC][1] = None + waiter = asyncio.Event() + dst.networking.wait_remote_key_event[DST] = {1: waiter} + + sender = asyncio.create_task(src._continuous_migration_sender()) + await asyncio.wait_for(waiter.wait(), timeout=2) + assert dst.local_state.data[DST][1] is None + await asyncio.wait_for(sender, timeout=2) diff --git a/worker/operator_state/aria/in_memory_state.py b/worker/operator_state/aria/in_memory_state.py index 84e94b47..7d7eb53e 100644 --- a/worker/operator_state/aria/in_memory_state.py +++ b/worker/operator_state/aria/in_memory_state.py @@ -15,6 +15,10 @@ if TYPE_CHECKING: from styx.common.types import K, KVPairs, OperatorPartition, V +# Marks a key that is not in a partition at all. A key can hold None as its value, +# so migration cannot use None to mean "already left in a batch". +_ABSENT = object() + class InMemoryOperatorState(BaseAriaState): data: dict[OperatorPartition, KVPairs] @@ -45,36 +49,37 @@ def get_key_to_migrate( new_operator_partition: OperatorPartition, key: K, old_partition: int, - ) -> K | None: + ) -> tuple[bool, V]: + """Take the key out of old_partition on its way to new_operator_partition. + + Returns (found, value): found is False when the key already left in a + migration batch. The value itself may be None. + """ operator_name, new_partition = new_operator_partition operator_partition: OperatorPartition = (operator_name, old_partition) - data_to_send = self.data[operator_partition].pop(key, None) - if data_to_send is None: + data_to_send = self.data[operator_partition].pop(key, _ABSENT) + if data_to_send is _ABSENT: # Key was already transferred via async migration batch - return None + return False, None # Record source-side key removal as tombstone for snapshotting self.delta_map[operator_partition][key] = None if operator_partition in self.keys_to_send: self.keys_to_send[operator_partition].discard((key, new_partition)) - return data_to_send + return True, data_to_send def set_data_from_migration( self, operator_partition: OperatorPartition, key: K, - data: KVPairs, + data: V, ) -> None: operator_partition = tuple(operator_partition) - if data is not None: - self.data[operator_partition][key] = data - self.delta_map[operator_partition][key] = data - # Only remove from remote_keys when we actually received the data. - # A None response means the async migration batch already transferred - # this key — the batch will arrive and set it via set_batch_data_from_migration. - if operator_partition in self.remote_keys: - self.remote_keys[operator_partition].pop(key, None) - if not self.remote_keys[operator_partition]: - del self.remote_keys[operator_partition] + self.data[operator_partition][key] = data + self.delta_map[operator_partition][key] = data + if operator_partition in self.remote_keys: + self.remote_keys[operator_partition].pop(key, None) + if not self.remote_keys[operator_partition]: + del self.remote_keys[operator_partition] def migrate_within_the_same_worker( self, @@ -84,10 +89,11 @@ def migrate_within_the_same_worker( old_partition: int, ) -> bool: new_operator_partition: OperatorPartition = (operator_name, new_partition) - data = self.get_key_to_migrate(new_operator_partition, key, old_partition) - # None: the async sender already took the key into a batch to this worker. - self.set_data_from_migration(new_operator_partition, key, data) - return data is not None + found, data = self.get_key_to_migrate(new_operator_partition, key, old_partition) + if found: + self.set_data_from_migration(new_operator_partition, key, data) + # Not found: the async sender already took the key into a batch to this worker. + return found def keys_remaining_to_remote(self) -> int: c = 0 @@ -112,10 +118,11 @@ def get_async_migrate_batch( operator_name, _ = operator_partition while keys and c < batch_size: key, new_partition = keys.pop() - value = self.data[operator_partition].pop(key, None) - if value is None: - # Key was deleted between rehash and transfer — skip it + value = self.data[operator_partition].pop(key, _ABSENT) + if value is _ABSENT: + # Key already left via get_key_to_migrate — skip it continue + # A None value is sent too: the destination still expects the key. # Record source-side key removal as tombstone for snapshotting self.delta_map[operator_partition][key] = None batch_to_send[(operator_name, new_partition)][key] = value diff --git a/worker/worker_service.py b/worker/worker_service.py index c81fa51a..415cfe87 100644 --- a/worker/worker_service.py +++ b/worker/worker_service.py @@ -586,10 +586,10 @@ async def _handle_request_remote_key( msg_type: MessageType, ) -> None: # Lock-free: pop() is atomic in CPython; concurrent requests for the same - # key are race-safe (loser gets None, which the requester handles). + # key are race-safe (the loser gets found=False, which the requester handles). del msg_type operator_partition, key, old_partition, host, port = self.networking.decode_message(data) - value = self.local_state.get_key_to_migrate( + found, value = self.local_state.get_key_to_migrate( operator_partition, key, old_partition, @@ -597,7 +597,7 @@ async def _handle_request_remote_key( await self.networking.send_message( host, port, - msg=(operator_partition, key, value), + msg=(operator_partition, key, found, value), msg_type=MessageType.ReceiveRemoteKey, serializer=Serializer.MSGPACK, ) @@ -609,8 +609,8 @@ async def _handle_receive_remote_key( ) -> None: # Lock-free: no awaits. del msg_type - operator_partition, key, value = self.networking.decode_message(data) - if value is None: + operator_partition, key, found, value = self.networking.decode_message(data) + if not found: # The source's async sender already took the key into a batch; the # batch wakes the waiting transaction when it lands here. return From 60b76486ce49c044e25d8908325107c31e7ff2ed Mon Sep 17 00:00:00 2001 From: Kyriakos Psarakis Date: Sat, 24 Oct 2026 17:36:32 +0200 Subject: [PATCH 05/20] Shorten the migration stall: keep the egress producer and apply local batches in place Every migration rebuilt the protocol with a new idempotent Kafka producer and then scanned all output partitions for duplicates before sending any response; together these held responses back for 300-600 ms after Phase B. A graceful (migration) stop now flushes the egress but keeps its producer, and the rebuilt protocol takes it over. The dedup scan runs only after a recovery: a clean migration stop flushes every response at an epoch boundary, and the new protocol resumes from those offsets. The async sender also serialized batches for partitions on its own worker and sent them over TCP to itself. Those are now applied in place, with no acknowledgement needed since nothing is in flight. Remote batches default to 500 keys so each sender step stays short. Co-Authored-By: Claude Opus 5.5 --- .../unit/worker/test_kafka_egress_coverage.py | 56 ++++++++++++++++ .../test_migration_snapshot_consistency.py | 51 ++++++++++++++ worker/egress/styx_kafka_batch_egress.py | 64 +++++++++++++----- worker/transactional_protocols/aria.py | 66 ++++++++++++++----- worker/worker_service.py | 7 +- 5 files changed, 207 insertions(+), 37 deletions(-) diff --git a/tests/unit/worker/test_kafka_egress_coverage.py b/tests/unit/worker/test_kafka_egress_coverage.py index 97e76b5b..ebadddc8 100644 --- a/tests/unit/worker/test_kafka_egress_coverage.py +++ b/tests/unit/worker/test_kafka_egress_coverage.py @@ -343,3 +343,59 @@ async def test_empty_dedup_offsets_returns_immediately(self): await e.get_messages_sent_before_recovery() # Should return without creating a consumer assert len(e.messages_sent_before_recovery) == 0 + + +# --------------------------------------------------------------------------- +# Reuse across a migration rebuild +# --------------------------------------------------------------------------- + + +class TestReuseAcrossMigration: + @pytest.mark.asyncio + async def test_start_reuses_running_producer(self, monkeypatch): + e = _egress() + producer = MagicMock() + e.kafka_egress_producer = producer + factory = MagicMock() + monkeypatch.setattr("worker.egress.styx_kafka_batch_egress.AIOKafkaProducer", factory) + + await e.start(worker_id=1) + + factory.assert_not_called() + assert e.kafka_egress_producer is producer + assert e.started.is_set() + + @pytest.mark.asyncio + async def test_flush_keeps_producer_running(self): + e = _egress() + producer = MagicMock(stop=AsyncMock()) + e.kafka_egress_producer = producer + await e.flush() + producer.stop.assert_not_called() + assert e.is_running + + @pytest.mark.asyncio + async def test_stop_is_idempotent(self): + e = _egress() + producer = MagicMock(stop=AsyncMock()) + e.kafka_egress_producer = producer + await e.stop() + await e.stop() + producer.stop.assert_awaited_once() + assert not e.is_running + + def test_prepare_for_new_protocol_resets_per_protocol_state(self): + e = _egress(restart=True, dedup_offsets={("users", 0): 5}) + e.messages_sent_before_recovery.add(b"r1") + e.batch.append(object()) + e.started.set() + new_offsets = {("users", 0): 7, ("users", 4): -1} + + e.prepare_for_new_protocol(new_offsets, restart_after_failure=False, dedup_output_offsets=None) + + assert e.topic_partition_output_offsets is new_offsets + assert e.restart_after_failure is False + assert e.dedup_output_offsets == {} + assert e.messages_sent_before_recovery == set() + assert e.batch == [] + assert not e.started.is_set() diff --git a/tests/unit/worker/test_migration_snapshot_consistency.py b/tests/unit/worker/test_migration_snapshot_consistency.py index 7803a533..ef7a885e 100644 --- a/tests/unit/worker/test_migration_snapshot_consistency.py +++ b/tests/unit/worker/test_migration_snapshot_consistency.py @@ -118,6 +118,7 @@ def _worker(link: _Link, worker_id: int, addr: tuple, peers: dict, state: InMemo proto.started = asyncio.Event() proto.started.set() proto.migrating_state = True + proto.registered_operators = dict.fromkeys(state.operator_partitions) proto._init_async_migration_state() proto.protocol_handlers_map = { MessageType.AsyncMigration: proto._handle_async_migration, @@ -252,3 +253,53 @@ async def test_waiter_for_none_valued_key_wakes(self): await asyncio.wait_for(waiter.wait(), timeout=2) assert dst.local_state.data[DST][1] is None await asyncio.wait_for(sender, timeout=2) + + +class TestLocalBatches: + async def test_local_batch_is_applied_in_place(self): + # Destination partition on the same worker: no network, no ack to wait for. + _link, src, _dst = _two_workers(3) + src.local_state.add_new_operator_partition(DST) + src.registered_operators = dict.fromkeys([SRC, DST]) + waiter = asyncio.Event() + src.networking.wait_remote_key_event[DST] = {1: waiter} + + batch = src.local_state.get_async_migrate_batch(10) + await src._send_migration_batch(batch) + + assert src.local_state.data[DST] == {0: "v0", 1: "v1", 2: "v2"} + assert waiter.is_set() + assert src._migration_all_acked.is_set() + assert _link.in_flight == set() + + +def _stoppable_protocol() -> AriaProtocol: + proto = AriaProtocol.__new__(AriaProtocol) + proto._stopping = False + proto.stopped = asyncio.Event() + proto.ingress = SimpleNamespace(stop=AsyncMock()) + proto.egress = SimpleNamespace(stop=AsyncMock(), flush=AsyncMock()) + proto.aio_task_scheduler = SimpleNamespace(close=AsyncMock()) + proto.background_functions = SimpleNamespace(close=AsyncMock()) + proto.snapshotting_networking_manager = SimpleNamespace(close_all_connections=AsyncMock()) + proto.function_scheduler_task = asyncio.create_task(asyncio.sleep(10)) + proto.communication_task = asyncio.create_task(asyncio.sleep(10)) + proto.migration_sender_task = None + proto.topic_partition_offsets = {} + return proto + + +class TestStopForMigration: + async def test_graceful_stop_keeps_egress_running(self): + proto = _stoppable_protocol() + await proto.stop(keep_egress=True) + proto.egress.flush.assert_awaited_once() + proto.egress.stop.assert_not_awaited() + assert proto.stopped.is_set() + + async def test_later_full_stop_releases_kept_egress(self): + # Recovery after a graceful stop must not leak the producer. + proto = _stoppable_protocol() + await proto.stop(keep_egress=True) + await proto.stop() + proto.egress.stop.assert_awaited_once() diff --git a/worker/egress/styx_kafka_batch_egress.py b/worker/egress/styx_kafka_batch_egress.py index d6fab597..c21f696c 100644 --- a/worker/egress/styx_kafka_batch_egress.py +++ b/worker/egress/styx_kafka_batch_egress.py @@ -48,24 +48,47 @@ def __init__( def clear_messages_sent_before_recovery(self) -> None: self.messages_sent_before_recovery: set = set() + @property + def is_running(self) -> bool: + return self.kafka_egress_producer is not None + + def prepare_for_new_protocol( + self, + topic_partition_output_offsets: dict[OperatorPartition, int], + restart_after_failure: bool, + dedup_output_offsets: dict[OperatorPartition, int] | None, + ) -> None: + """Reuse this egress, and its running Kafka producer, for a rebuilt protocol. + + Starting a new idempotent producer takes hundreds of milliseconds, which a + migration would otherwise spend with no responses going out. + """ + self.topic_partition_output_offsets = topic_partition_output_offsets + self.restart_after_failure = restart_after_failure + self.dedup_output_offsets = dedup_output_offsets or {} + self.messages_sent_before_recovery = set() + self.batch = [] + self.started = asyncio.Event() + async def start(self, worker_id: int) -> None: self.worker_id = worker_id - # Idempotent (implies acks="all"): this producer writes client responses and - # the sequencer WAL, which recovery relies on. With acks=1 a leader failover - # can truncate acked records, and retries can duplicate responses. - self.kafka_egress_producer = AIOKafkaProducer( - bootstrap_servers=[KAFKA_URL], - client_id=f"W{worker_id}_{uuid.uuid4()}", - enable_idempotence=True, - ) - while True: - try: - await self.kafka_egress_producer.start() - except KafkaConnectionError: - await asyncio.sleep(1) - logging.info("Waiting for Kafka") - continue - break + if self.kafka_egress_producer is None: + # Idempotent (implies acks="all"): this producer writes client responses and + # the sequencer WAL, which recovery relies on. With acks=1 a leader failover + # can truncate acked records, and retries can duplicate responses. + self.kafka_egress_producer = AIOKafkaProducer( + bootstrap_servers=[KAFKA_URL], + client_id=f"W{worker_id}_{uuid.uuid4()}", + enable_idempotence=True, + ) + while True: + try: + await self.kafka_egress_producer.start() + except KafkaConnectionError: + await asyncio.sleep(1) + logging.info("Waiting for Kafka") + continue + break if not self.restart_after_failure: # Non-recovery path: mark started immediately. # Recovery path: started is set later by run_dedup_scan_and_mark_started() @@ -90,7 +113,13 @@ async def run_dedup_scan_and_mark_started(self) -> None: ) self.started.set() + async def flush(self) -> None: + """Wait until every buffered response is written, keeping the producer running.""" + await self.send_batch() + async def stop(self) -> None: + if self.kafka_egress_producer is None: + return await self.send_batch() if self._dedup_suppressed > 0 or self._dedup_sent > 0: logging.warning( @@ -98,7 +127,8 @@ async def stop(self) -> None: f"suppressed={self._dedup_suppressed}, sent={self._dedup_sent}, " f"remaining_dedup_set_size={len(self.messages_sent_before_recovery)}", ) - await self.kafka_egress_producer.stop() + producer, self.kafka_egress_producer = self.kafka_egress_producer, None + await producer.stop() async def send( self, diff --git a/worker/transactional_protocols/aria.py b/worker/transactional_protocols/aria.py index b8df239c..76873d22 100644 --- a/worker/transactional_protocols/aria.py +++ b/worker/transactional_protocols/aria.py @@ -47,7 +47,8 @@ SEQUENCE_MAX_SIZE: int = int(os.getenv("SEQUENCE_MAX_SIZE", "1_000")) KAFKA_URL: str = os.environ["KAFKA_URL"] USE_ASYNC_MIGRATION: bool = bool(strtobool(os.getenv("USE_ASYNC_MIGRATION", "true"))) -ASYNC_MIGRATION_BATCH_SIZE: int = int(os.getenv("ASYNC_MIGRATION_BATCH_SIZE", "2_000")) +# Small batches keep each sender step short, so epochs are not held up behind it. +ASYNC_MIGRATION_BATCH_SIZE: int = int(os.getenv("ASYNC_MIGRATION_BATCH_SIZE", "500")) class AriaProtocol(BaseTransactionalProtocol): @@ -70,6 +71,7 @@ def __init__( # noqa: PLR0913 restart_after_migration: bool = False, *, dedup_output_offsets: dict[OperatorPartition, int] | None = None, + egress: StyxKafkaBatchEgress | None = None, ) -> None: if topic_partition_offsets is None: topic_partition_offsets = {(tp.topic, tp.partition): -1 for tp in topic_partitions} @@ -126,11 +128,19 @@ def __init__( # noqa: PLR0913 epoch_interval_ms=100, ) - self.egress: StyxKafkaBatchEgress = StyxKafkaBatchEgress( - output_offsets, - restart_after_recovery or restart_after_migration, - dedup_output_offsets=dedup_output_offsets, - ) + # Responses may duplicate ones sent before a recovery; scan for them first. + # After a clean migration stop nothing is replayed, so there is nothing to scan. + dedup_on_start = restart_after_recovery or dedup_output_offsets is not None + if egress is None: + self.egress: StyxKafkaBatchEgress = StyxKafkaBatchEgress( + output_offsets, + dedup_on_start, + dedup_output_offsets=dedup_output_offsets, + ) + else: + # Rebuilt after a clean migration stop: keep the running Kafka producer. + egress.prepare_for_new_protocol(output_offsets, dedup_on_start, dedup_output_offsets) + self.egress = egress # Primary task used for processing self.function_scheduler_task: asyncio.Task | None = None self.communication_task: asyncio.Task | None = None @@ -202,14 +212,22 @@ def _init_async_migration_state(self) -> None: async def wait_stopped(self) -> None: await self.stopped.wait() - async def stop(self) -> None: + async def stop(self, *, keep_egress: bool = False) -> None: + """Stop the protocol. keep_egress flushes the egress but leaves its Kafka + producer running, so the protocol rebuilt after a migration can reuse it.""" if self._stopping: # Another caller is already tearing down — just wait for completion. await self.stopped.wait() + if not keep_egress: + # A graceful (migration) stop may have left the producer running. + await self.egress.stop() return self._stopping = True await self.ingress.stop() - await self.egress.stop() + if keep_egress: + await self.egress.flush() + else: + await self.egress.stop() await self.aio_task_scheduler.close() await self.background_functions.close() await self.snapshotting_networking_manager.close_all_connections() @@ -498,13 +516,7 @@ async def _handle_async_migration(self, data: bytes) -> None: operator_partition, batch, sender_id, nonce, batch_id = self.networking.decode_message(data) logging.debug(f"MIGRATION | Received migration batch for: {operator_partition}") - op = tuple(operator_partition) - self.local_state.set_batch_data_from_migration(op, batch) - # Unblock any transactions waiting for these keys via RequestRemoteKey - if op in self.networking.wait_remote_key_event: - for key in batch: - if key in self.networking.wait_remote_key_event.get(op, {}): - self.networking.wait_remote_key_event[op][key].set() + self._apply_migration_batch(operator_partition, batch) # Acknowledge only after the batch is applied (and so in this epoch's delta). if sender_id == self.id: @@ -519,6 +531,16 @@ async def _handle_async_migration(self, data: bytes) -> None: serializer=Serializer.MSGPACK, ) + def _apply_migration_batch(self, operator_partition: OperatorPartition, batch: dict) -> None: + op = tuple(operator_partition) + self.local_state.set_batch_data_from_migration(op, batch) + # Unblock any transactions waiting for these keys + waiting = self.networking.wait_remote_key_event.get(op) + if waiting: + for key in batch: + if key in waiting: + waiting[key].set() + async def _handle_async_migration_ack(self, data: bytes) -> None: # Lock-free: no awaits. nonce, batch_id = self.networking.decode_message(data) @@ -561,11 +583,15 @@ async def _send_migration_batch( ) -> None: """Send a migration batch to destination workers in parallel. - Every batch goes over the network, including to this worker, and stays - unacknowledged until the destination has applied it.""" + A batch for a partition on this worker is applied in place: nothing is in + flight, so it needs no acknowledgement. A remote batch stays unacknowledged + until the destination has applied it.""" send_tasks = [] for operator_partition, k_v_pairs in batch.items(): operator_name, partition = operator_partition + if (operator_name, partition) in self.registered_operators: + self._apply_migration_batch(operator_partition, k_v_pairs) + continue worker = self.dns[operator_name][partition] # Register before the first await so a concurrent drain waits for it. batch_id = next(self._migration_batch_ids) @@ -641,6 +667,7 @@ async def function_scheduler(self) -> None: await self.started.wait() logging.warning("STARTED function scheduler") + graceful = False try: while self.running: # Wait until the ingress signals that messages are available, @@ -660,6 +687,7 @@ async def function_scheduler(self) -> None: self.currently_processing = True await self._process_epoch(sequence) + graceful = True logging.warning(f"Worker {self.id} | function_scheduler exiting (running=False)") except asyncio.CancelledError: logging.warning(f"Worker {self.id} | function_scheduler cancelled") @@ -671,7 +699,9 @@ async def function_scheduler(self) -> None: # already tearing down. Otherwise we deadlock: the external # stop() awaits this task, while this task awaits stopped. if not self._stopping: - await self.stop() + # Only a migration stops the protocol gracefully; its rebuilt + # protocol takes over the egress. + await self.stop(keep_egress=graceful) async def _process_epoch(self, sequence: list[SequencedItem]) -> None: epoch_start = timer() diff --git a/worker/worker_service.py b/worker/worker_service.py index 415cfe87..0f919098 100644 --- a/worker/worker_service.py +++ b/worker/worker_service.py @@ -552,7 +552,10 @@ async def _migration_rebuild_runtime(self) -> None: topic_partition_offsets = {op: self.m_input_offsets.get(op, -1) for op in self.registered_operators} topic_partition_output_offsets = {op: self.m_output_offsets.get(op, -1) for op in self.registered_operators} - # Recreate protocol (restart_after_migration) + # Recreate protocol (restart_after_migration). The old protocol stopped + # gracefully at an epoch boundary with every response flushed, so the new + # one reuses its Kafka producer and has nothing to deduplicate. + old_egress = self.function_execution_protocol.egress self.function_execution_protocol = AriaProtocol( worker_id=self.id, peers=self.peers, @@ -567,7 +570,7 @@ async def _migration_rebuild_runtime(self) -> None: epoch_counter=self.m_epoch_counter, t_counter=self.m_t_counter, restart_after_migration=True, - dedup_output_offsets=dict(self.m_output_offsets), + egress=old_egress if old_egress.is_running else None, ) async def _handle_receive_migration_hashes( From daddced2edf53ee308f41e51c5e4c58307c72678 Mon Sep 17 00:00:00 2001 From: Kyriakos Psarakis Date: Sat, 24 Oct 2026 18:15:47 +0200 Subject: [PATCH 06/20] Spread the async transfer over epochs The sender moved a worker's whole share in a burst, so all of it landed in one epoch's snapshot delta and was serialized on that epoch's critical path (about 200 ms for 60k keys). It now moves at most ASYNC_MIGRATION_KEYS_PER_EPOCH keys (default 3000) per epoch, refilling at each epoch end, or after a short wait when no epoch is running. Transactions do not wait for the transfer: they fetch keys on demand. Co-Authored-By: Claude Opus 5.5 --- .../test_migration_snapshot_consistency.py | 23 ++++++++++++++++ worker/transactional_protocols/aria.py | 27 ++++++++++++++++++- 2 files changed, 49 insertions(+), 1 deletion(-) diff --git a/tests/unit/worker/test_migration_snapshot_consistency.py b/tests/unit/worker/test_migration_snapshot_consistency.py index ef7a885e..cfd5111f 100644 --- a/tests/unit/worker/test_migration_snapshot_consistency.py +++ b/tests/unit/worker/test_migration_snapshot_consistency.py @@ -118,6 +118,7 @@ def _worker(link: _Link, worker_id: int, addr: tuple, peers: dict, state: InMemo proto.started = asyncio.Event() proto.started.set() proto.migrating_state = True + proto.currently_processing = False proto.registered_operators = dict.fromkeys(state.operator_partitions) proto._init_async_migration_state() proto.protocol_handlers_map = { @@ -303,3 +304,25 @@ async def test_later_full_stop_releases_kept_egress(self): await proto.stop(keep_egress=True) await proto.stop() proto.egress.stop.assert_awaited_once() + + +class TestMigrationBudget: + async def test_keys_per_epoch_are_capped_until_the_epoch_ends(self, monkeypatch): + monkeypatch.setattr(aria, "ASYNC_MIGRATION_KEYS_PER_EPOCH", 4) + monkeypatch.setattr(aria, "ASYNC_MIGRATION_BATCH_SIZE", 3) + _link, src, _dst = _two_workers(10) + src._init_async_migration_state() + src.currently_processing = True # an epoch is running + + sender = asyncio.create_task(src._continuous_migration_sender()) + await asyncio.sleep(0.05) + assert src.local_state.keys_remaining_to_send() == 6 # 3 + 1, then waits + + src._migration_epoch_ended.set() # the epoch ends + await asyncio.sleep(0.05) + assert src.local_state.keys_remaining_to_send() == 2 + + src.currently_processing = False # idle: refills on its own + src._migration_epoch_ended.set() + await asyncio.wait_for(sender, timeout=2) + assert src.local_state.keys_remaining_to_send() == 0 diff --git a/worker/transactional_protocols/aria.py b/worker/transactional_protocols/aria.py index 76873d22..ee715e0e 100644 --- a/worker/transactional_protocols/aria.py +++ b/worker/transactional_protocols/aria.py @@ -49,6 +49,11 @@ USE_ASYNC_MIGRATION: bool = bool(strtobool(os.getenv("USE_ASYNC_MIGRATION", "true"))) # Small batches keep each sender step short, so epochs are not held up behind it. ASYNC_MIGRATION_BATCH_SIZE: int = int(os.getenv("ASYNC_MIGRATION_BATCH_SIZE", "500")) +# Keys moved per epoch. Everything moved in an epoch is serialized into that +# epoch's snapshot delta on its critical path, so the transfer is spread out. +ASYNC_MIGRATION_KEYS_PER_EPOCH: int = int(os.getenv("ASYNC_MIGRATION_KEYS_PER_EPOCH", "3_000")) +# With no epoch running, the budget refills after this long instead. +ASYNC_MIGRATION_IDLE_REFILL_S: float = 0.01 class AriaProtocol(BaseTransactionalProtocol): @@ -208,6 +213,8 @@ def _init_async_migration_state(self) -> None: # Distinguishes acks for this protocol's batches from late acks for a # previous protocol instance (recovery or migration rebuilds the protocol). self._migration_nonce: int = random.getrandbits(62) + self._migration_budget: int = ASYNC_MIGRATION_KEYS_PER_EPOCH + self._migration_epoch_ended: asyncio.Event = asyncio.Event() async def wait_stopped(self) -> None: await self.stopped.wait() @@ -655,7 +662,14 @@ async def _continuous_migration_sender(self) -> None: await asyncio.sleep(0.01) continue - batch = self.local_state.get_async_migrate_batch(ASYNC_MIGRATION_BATCH_SIZE) + if self._migration_budget <= 0: + await self._wait_for_migration_budget() + continue + + batch = self.local_state.get_async_migrate_batch( + min(ASYNC_MIGRATION_BATCH_SIZE, self._migration_budget), + ) + self._migration_budget -= sum(len(kv) for kv in batch.values()) await self._send_migration_batch(batch) # Yield to event loop between batches await asyncio.sleep(0) @@ -663,6 +677,16 @@ async def _continuous_migration_sender(self) -> None: logging.warning("MIGRATION | Continuous migration sender cancelled") raise + async def _wait_for_migration_budget(self) -> None: + """Wait for the next epoch to end, or a short while if none is running.""" + self._migration_epoch_ended.clear() + if self.currently_processing: + await self._migration_epoch_ended.wait() + else: + with contextlib.suppress(TimeoutError): + await asyncio.wait_for(self._migration_epoch_ended.wait(), ASYNC_MIGRATION_IDLE_REFILL_S) + self._migration_budget = ASYNC_MIGRATION_KEYS_PER_EPOCH + async def function_scheduler(self) -> None: await self.started.wait() logging.warning("STARTED function scheduler") @@ -769,6 +793,7 @@ async def _process_epoch(self, sequence: list[SequencedItem]) -> None: # Every worker has passed this epoch's snapshot cut, so batches sent from # here on belong to the next epoch at both ends. self._migration_gate.set() + self._migration_epoch_ended.set() async def _run_epoch_functions_and_chain( self, From 8547faaee7cae7a0c959e27eacc8f899eeb5fcbb Mon Sep 17 00:00:00 2001 From: Kyriakos Psarakis Date: Sat, 24 Oct 2026 19:04:03 +0200 Subject: [PATCH 07/20] Keep the ingress Kafka consumer across a migration rebuild The rebuilt protocol created a new Kafka consumer, and the first epoch after a migration waited about 180 ms at its first barrier for the slowest worker's consumer to connect and fetch. A graceful (migration) stop now keeps the consumer connected, and the rebuilt protocol reassigns it to the new partitions and seeks each to its stop offset, which also discards any records it had prefetched. Co-Authored-By: Claude Opus 5.5 --- tests/unit/worker/test_kafka_ingress.py | 59 ++++++++++++++++++- .../test_migration_snapshot_consistency.py | 8 ++- worker/ingress/styx_kafka_ingress.py | 50 +++++++++++++++- worker/transactional_protocols/aria.py | 26 ++++---- worker/worker_service.py | 4 +- 5 files changed, 129 insertions(+), 18 deletions(-) diff --git a/tests/unit/worker/test_kafka_ingress.py b/tests/unit/worker/test_kafka_ingress.py index 28075bec..bb8ecb80 100644 --- a/tests/unit/worker/test_kafka_ingress.py +++ b/tests/unit/worker/test_kafka_ingress.py @@ -3,6 +3,7 @@ import asyncio from unittest.mock import AsyncMock, MagicMock, patch +from aiokafka import TopicPartition from aiokafka.errors import KafkaConnectionError, UnknownTopicOrPartitionError import pytest from styx.common.message_types import MessageType @@ -383,7 +384,7 @@ def test_except_clause_catches_both_exception_types(self): import inspect import textwrap - source = textwrap.dedent(inspect.getsource(StyxKafkaIngress.start_kafka_consumer)) + source = textwrap.dedent(inspect.getsource(StyxKafkaIngress._create_kafka_consumer)) tree = ast.parse(source) # Find all except handlers in the function @@ -403,3 +404,59 @@ def test_except_clause_catches_both_exception_types(self): "Use 'except (UnknownTopicOrPartitionError, KafkaConnectionError):' " "to catch both.", ) + + +async def _poll_forever(**_): + await asyncio.sleep(10) + + +class TestConsumerHandOver: + """A migration stop keeps the connected consumer for the rebuilt protocol.""" + + @pytest.mark.asyncio + async def test_adopted_consumer_is_reassigned_and_seeked_not_recreated(self): + ingress, *_ = _make_ingress() + consumer = MagicMock() + consumer.getmany = AsyncMock(side_effect=asyncio.CancelledError()) + consumer.stop = AsyncMock() + ingress.adopt_consumer(consumer) + tps = [TopicPartition("users", 0), TopicPartition("users", 4)] + + with ( + patch("worker.ingress.styx_kafka_ingress.AIOKafkaConsumer") as factory, + pytest.raises(asyncio.CancelledError), + ): + await ingress.start_kafka_consumer(tps, {("users", 0): 41, ("users", 4): -1}) + + factory.assert_not_called() + consumer.assign.assert_called_once_with(tps) + consumer.seek.assert_any_call(TopicPartition("users", 0), 42) + consumer.seek.assert_any_call(TopicPartition("users", 4), 0) + consumer.start.assert_not_called() + + @pytest.mark.asyncio + async def test_stop_keep_consumer_hands_it_over(self): + ingress, *_ = _make_ingress() + consumer = MagicMock(stop=AsyncMock()) + consumer.getmany = AsyncMock(side_effect=_poll_forever) + ingress.adopt_consumer(consumer) + await ingress.start([TopicPartition("users", 0)], {("users", 0): -1}) + + await ingress.stop(keep_consumer=True) + + consumer.stop.assert_not_called() + assert ingress.take_consumer() is consumer + assert ingress.take_consumer() is None + + @pytest.mark.asyncio + async def test_stop_closes_consumer(self): + ingress, *_ = _make_ingress() + consumer = MagicMock(stop=AsyncMock()) + consumer.getmany = AsyncMock(side_effect=_poll_forever) + ingress.adopt_consumer(consumer) + await ingress.start([TopicPartition("users", 0)], {("users", 0): -1}) + + await ingress.stop() + + consumer.stop.assert_awaited_once() + assert ingress.take_consumer() is None diff --git a/tests/unit/worker/test_migration_snapshot_consistency.py b/tests/unit/worker/test_migration_snapshot_consistency.py index cfd5111f..7ef2d84a 100644 --- a/tests/unit/worker/test_migration_snapshot_consistency.py +++ b/tests/unit/worker/test_migration_snapshot_consistency.py @@ -278,7 +278,7 @@ def _stoppable_protocol() -> AriaProtocol: proto = AriaProtocol.__new__(AriaProtocol) proto._stopping = False proto.stopped = asyncio.Event() - proto.ingress = SimpleNamespace(stop=AsyncMock()) + proto.ingress = SimpleNamespace(stop=AsyncMock(), close_consumer=AsyncMock()) proto.egress = SimpleNamespace(stop=AsyncMock(), flush=AsyncMock()) proto.aio_task_scheduler = SimpleNamespace(close=AsyncMock()) proto.background_functions = SimpleNamespace(close=AsyncMock()) @@ -293,7 +293,8 @@ def _stoppable_protocol() -> AriaProtocol: class TestStopForMigration: async def test_graceful_stop_keeps_egress_running(self): proto = _stoppable_protocol() - await proto.stop(keep_egress=True) + await proto.stop(keep_kafka_clients=True) + proto.ingress.stop.assert_awaited_once_with(keep_consumer=True) proto.egress.flush.assert_awaited_once() proto.egress.stop.assert_not_awaited() assert proto.stopped.is_set() @@ -301,8 +302,9 @@ async def test_graceful_stop_keeps_egress_running(self): async def test_later_full_stop_releases_kept_egress(self): # Recovery after a graceful stop must not leak the producer. proto = _stoppable_protocol() - await proto.stop(keep_egress=True) + await proto.stop(keep_kafka_clients=True) await proto.stop() + proto.ingress.close_consumer.assert_awaited_once() proto.egress.stop.assert_awaited_once() diff --git a/worker/ingress/styx_kafka_ingress.py b/worker/ingress/styx_kafka_ingress.py index cfbda3cd..f0168a2b 100644 --- a/worker/ingress/styx_kafka_ingress.py +++ b/worker/ingress/styx_kafka_ingress.py @@ -47,6 +47,20 @@ def __init__( self.kafka_consumer: AIOKafkaConsumer | None = None self.kafka_ingress_task: asyncio.Task | None = None + self._keep_consumer: bool = False + + def adopt_consumer(self, kafka_consumer: AIOKafkaConsumer) -> None: + """Reuse a connected consumer from the protocol this one replaces. + + Creating and starting a new consumer, and waiting for its first fetch, + holds up the first epoch after a migration by a few hundred milliseconds. + """ + self.kafka_consumer = kafka_consumer + + def take_consumer(self) -> AIOKafkaConsumer | None: + """Hand over the consumer kept by stop(keep_consumer=True), if any.""" + consumer, self.kafka_consumer = self.kafka_consumer, None + return consumer async def start( self, @@ -58,13 +72,22 @@ async def start( ) await self.started.wait() - async def stop(self) -> None: + async def stop(self, *, keep_consumer: bool = False) -> None: + """Stop consuming. keep_consumer leaves the Kafka consumer connected for + take_consumer().""" + self._keep_consumer = keep_consumer self.kafka_ingress_task.cancel() try: await self.kafka_ingress_task except asyncio.CancelledError: logging.warning("kafka ingress coroutine suts down...") - await self.kafka_consumer.stop() + if not keep_consumer: + await self.close_consumer() + + async def close_consumer(self) -> None: + consumer = self.take_consumer() + if consumer is not None: + await consumer.stop() def handle_message_from_kafka(self, msg: ConsumerRecord) -> None: message_type: int = self.networking.get_msg_type(msg.value) @@ -133,6 +156,24 @@ async def start_kafka_consumer( self, topic_partitions: list[TopicPartition], topic_partition_offsets: dict[OperatorPartition, int], + ) -> None: + if self.kafka_consumer is not None: + # Adopted from the previous protocol: already connected. Seeking every + # assigned partition drops any records it had prefetched. + self.kafka_consumer.assign(topic_partitions) + for operator, offset in topic_partition_offsets.items(): + self.kafka_consumer.seek(TopicPartition(operator[0], operator[1]), offset + 1) + logging.warning( + f"{self.worker_id} REUSED Kafka consumer for topic partitions: {topic_partitions}", + ) + else: + await self._create_kafka_consumer(topic_partitions, topic_partition_offsets) + await self._consume(topic_partition_offsets) + + async def _create_kafka_consumer( + self, + topic_partitions: list[TopicPartition], + topic_partition_offsets: dict[OperatorPartition, int], ) -> None: logging.warning( f"{self.worker_id} CREATED Kafka consumer for topic partitions: {topic_partitions}", @@ -160,6 +201,8 @@ async def start_kafka_consumer( ) continue break + + async def _consume(self, topic_partition_offsets: dict[OperatorPartition, int]) -> None: try: # Consume messages logging.warning( @@ -186,4 +229,5 @@ async def start_kafka_consumer( self.send_message_tasks = [] self.messages_available.set() finally: - await self.kafka_consumer.stop() + if not self._keep_consumer: + await self.close_consumer() diff --git a/worker/transactional_protocols/aria.py b/worker/transactional_protocols/aria.py index ee715e0e..1b1587e9 100644 --- a/worker/transactional_protocols/aria.py +++ b/worker/transactional_protocols/aria.py @@ -26,7 +26,7 @@ if TYPE_CHECKING: from collections.abc import Awaitable, Callable - from aiokafka import TopicPartition + from aiokafka import AIOKafkaConsumer, TopicPartition from styx.common.operator import Operator from styx.common.types import OperatorPartition @@ -77,6 +77,7 @@ def __init__( # noqa: PLR0913 *, dedup_output_offsets: dict[OperatorPartition, int] | None = None, egress: StyxKafkaBatchEgress | None = None, + kafka_consumer: AIOKafkaConsumer | None = None, ) -> None: if topic_partition_offsets is None: topic_partition_offsets = {(tp.topic, tp.partition): -1 for tp in topic_partitions} @@ -132,6 +133,9 @@ def __init__( # noqa: PLR0913 sequence_max_size=SEQUENCE_MAX_SIZE, epoch_interval_ms=100, ) + if kafka_consumer is not None: + # Rebuilt after a clean migration stop: keep the connected consumer. + self.ingress.adopt_consumer(kafka_consumer) # Responses may duplicate ones sent before a recovery; scan for them first. # After a clean migration stop nothing is replayed, so there is nothing to scan. @@ -219,19 +223,21 @@ def _init_async_migration_state(self) -> None: async def wait_stopped(self) -> None: await self.stopped.wait() - async def stop(self, *, keep_egress: bool = False) -> None: - """Stop the protocol. keep_egress flushes the egress but leaves its Kafka - producer running, so the protocol rebuilt after a migration can reuse it.""" + async def stop(self, *, keep_kafka_clients: bool = False) -> None: + """Stop the protocol. keep_kafka_clients flushes the egress but leaves its + producer and the ingress consumer connected, so the protocol rebuilt after + a migration can reuse them.""" if self._stopping: # Another caller is already tearing down — just wait for completion. await self.stopped.wait() - if not keep_egress: - # A graceful (migration) stop may have left the producer running. + if not keep_kafka_clients: + # A graceful (migration) stop may have left them connected. + await self.ingress.close_consumer() await self.egress.stop() return self._stopping = True - await self.ingress.stop() - if keep_egress: + await self.ingress.stop(keep_consumer=keep_kafka_clients) + if keep_kafka_clients: await self.egress.flush() else: await self.egress.stop() @@ -724,8 +730,8 @@ async def function_scheduler(self) -> None: # stop() awaits this task, while this task awaits stopped. if not self._stopping: # Only a migration stops the protocol gracefully; its rebuilt - # protocol takes over the egress. - await self.stop(keep_egress=graceful) + # protocol takes over the Kafka clients. + await self.stop(keep_kafka_clients=graceful) async def _process_epoch(self, sequence: list[SequencedItem]) -> None: epoch_start = timer() diff --git a/worker/worker_service.py b/worker/worker_service.py index 0f919098..c7a6b33c 100644 --- a/worker/worker_service.py +++ b/worker/worker_service.py @@ -554,8 +554,9 @@ async def _migration_rebuild_runtime(self) -> None: # Recreate protocol (restart_after_migration). The old protocol stopped # gracefully at an epoch boundary with every response flushed, so the new - # one reuses its Kafka producer and has nothing to deduplicate. + # one reuses its Kafka clients and has nothing to deduplicate. old_egress = self.function_execution_protocol.egress + old_consumer = self.function_execution_protocol.ingress.take_consumer() self.function_execution_protocol = AriaProtocol( worker_id=self.id, peers=self.peers, @@ -571,6 +572,7 @@ async def _migration_rebuild_runtime(self) -> None: t_counter=self.m_t_counter, restart_after_migration=True, egress=old_egress if old_egress.is_running else None, + kafka_consumer=old_consumer, ) async def _handle_receive_migration_hashes( From be1824033c07090b7d1bf520c46eecde5e995787 Mon Sep 17 00:00:00 2001 From: Kyriakos Psarakis Date: Sat, 24 Oct 2026 19:46:43 +0200 Subject: [PATCH 08/20] Check pending migration keys in constant time exists() scanned the whole keys_to_send set of a partition on every call, and the ingress calls it for every message while a migration is running. With 60k keys pending, each call took about 3 ms, so a backlog of 2,000 messages kept the event loop busy for seconds and epochs barely ran during Phase C (this predates the PR). exists() now uses an index of the pending keys, built once per partition and reset when a new transfer is registered. Co-Authored-By: Claude Opus 5.5 --- tests/unit/worker/test_in_memory_state.py | 36 +++++++++++++++++++ worker/operator_state/aria/in_memory_state.py | 24 ++++++++++--- 2 files changed, 56 insertions(+), 4 deletions(-) diff --git a/tests/unit/worker/test_in_memory_state.py b/tests/unit/worker/test_in_memory_state.py index 6f018c89..ca3f8e78 100644 --- a/tests/unit/worker/test_in_memory_state.py +++ b/tests/unit/worker/test_in_memory_state.py @@ -386,6 +386,42 @@ def test_false_for_unknown_operator_partition(self): s = _state() assert s.exists("k", "other_op", 99) is False + def test_does_not_rescan_keys_to_send(self): + # The ingress calls exists() for every message during a migration. + class CountingSet(set): + scans = 0 + + def __iter__(self): + CountingSet.scans += 1 + return super().__iter__() + + s = _state() + s.data[OP_PART] = {"k": "v", "stay": "v"} + s.add_keys_to_send({OP_PART: CountingSet({("k", 1)})}) + for _ in range(100): + assert s.exists("k", OP, PART) is False + assert s.exists("stay", OP, PART) is True + assert CountingSet.scans == 1 + + def test_key_that_left_is_answered_by_remote_keys(self): + s = _state() + s.data[OP_PART] = {"k": "v"} + s.add_keys_to_send({OP_PART: {("k", 1)}}) + assert s.exists("k", OP, PART) is False + s.get_async_migrate_batch(10) # k leaves; keys_to_send empties + assert s.exists("k", OP, PART) is False + s.data[OP_PART]["k"] = "v2" # migration over: a key here belongs here + assert s.exists("k", OP, PART) is True + + def test_new_transfer_resets_the_index(self): + s = _state() + s.data[OP_PART] = {"a": 1, "b": 2} + s.add_keys_to_send({OP_PART: {("a", 1)}}) + assert s.exists("b", OP, PART) is True + s.add_keys_to_send({OP_PART: {("b", 1)}}) + assert s.exists("b", OP, PART) is False + assert s.exists("a", OP, PART) is True + # --------------------------------------------------------------------------- # snapshot helpers diff --git a/worker/operator_state/aria/in_memory_state.py b/worker/operator_state/aria/in_memory_state.py index 7d7eb53e..c4f02063 100644 --- a/worker/operator_state/aria/in_memory_state.py +++ b/worker/operator_state/aria/in_memory_state.py @@ -37,9 +37,26 @@ def __init__(self, operator_partitions: set[OperatorPartition]) -> None: # Used to track migration progress, the async migration and whether the operator still owns the specific key # (operator_name, old_partition): set of keys with the new partition self.keys_to_send: dict[OperatorPartition, set[tuple[K, int]]] = {} + # The keys of keys_to_send, for O(1) lookups (built lazily per partition) + self._keys_to_send_index: dict[OperatorPartition, set[K]] = {} def add_keys_to_send(self, keys_to_send: dict[OperatorPartition, K]) -> None: self.keys_to_send = keys_to_send + self._keys_to_send_index = {} + + def _is_pending_send(self, operator_partition: OperatorPartition, key: K) -> bool: + """Whether key was registered to leave operator_partition. + + The index may still hold keys that already left; callers check that the + key is still in the partition's data, and sending always removes it there. + """ + if not self.keys_to_send.get(operator_partition): + return False + index = self._keys_to_send_index.get(operator_partition) + if index is None: + index = {k for k, _ in self.keys_to_send[operator_partition]} + self._keys_to_send_index[operator_partition] = index + return key in index def has_keys_to_send(self) -> bool: return bool(self.keys_to_send) @@ -292,10 +309,9 @@ def exists(self, key: K, operator_name: str, partition: int) -> bool: return False if key in self.data[operator_partition]: # During migration: a key still in data but pending send to another - # partition should not be considered as belonging here. - if operator_partition in self.keys_to_send: - return not any(k == key for k, _ in self.keys_to_send[operator_partition]) - return True + # partition should not be considered as belonging here. The ingress + # calls this for every message, so it must not scan keys_to_send. + return not self._is_pending_send(operator_partition, key) # During migration: key is expected to arrive from a remote worker return self.in_remote_keys(key, operator_name, partition) From c9d703c74a1201dc6f6b7cb899f7aa4937325a18 Mon Sep 17 00:00:00 2001 From: Kyriakos Psarakis Date: Sat, 24 Oct 2026 20:07:33 +0200 Subject: [PATCH 09/20] Install uvloop in the Cython hot-path unit test job The new worker tests import worker_service, which imports uvloop. Co-Authored-By: Claude Opus 5.5 --- .github/workflows/ci.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 404abd08..2ba15cd8 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -139,7 +139,7 @@ jobs: python -m pip install --upgrade pip setuptools Cython cd styx-package && python setup.py build_ext --inplace && cd .. python -m pip install -e styx-package/ - python -m pip install msgspec pytest==9.1.1 pytest-asyncio psutil + python -m pip install msgspec pytest==9.1.1 pytest-asyncio psutil uvloop python worker/setup.py build_ext --inplace - name: Check the hot-path modules are the compiled ones From 5472be072f979a0c435f0d5194f1c47d60fea9ec Mon Sep 17 00:00:00 2001 From: Kyriakos Psarakis Date: Sat, 24 Oct 2026 20:12:25 +0200 Subject: [PATCH 10/20] Order snapshot completions, drop stale migration messages, tag the checkpoint, escalate lost barrier replies Snapshots: the snapshot id advanced only when every file of a snapshot had uploaded, so a snapshot started while the previous one was still uploading wrote its files under the same id (overwriting the earlier deltas) and overwrote its offsets; failed uploads were also reported as complete. Each snapshot now gets its id at start and its own metadata, failed uploads are not reported, and completions are reported to the coordinator strictly in id order. Stale messages: after a recovery, migration batches, remote-key requests and replies from a protocol of before the recovery could still be applied. The coordinator now bumps an incarnation number on every recovery; these messages carry it and are dropped when it does not match. Checkpoint tag: any snapshot completing after Phase B began got the checkpoint and was taken as the checkpoint, including snapshots cut before Phase B. The workers now flag the checkpoint snapshot, and only it and later snapshots carry the checkpoint. Stop timeout: Phase B force-stopped a protocol still running after 15 s, possibly cutting an epoch mid-commit, then continued the migration. It now keeps waiting. A barrier or migration message the coordinator cannot deliver hands the worker to recovery instead, which is what the timeout was guarding against. Co-Authored-By: Claude Opus 5.5 --- coordinator/coordinator_metadata.py | 48 ++++-- coordinator/coordinator_service.py | 114 ++++++++------ styx-package/styx/common/tcp_networking.py | 14 +- .../coordinator/test_coordinator_metadata.py | 17 ++- .../test_migration_checkpoint_lifetime.py | 66 +++++++- tests/unit/worker/test_async_snapshots.py | 135 +++++++++-------- .../worker/test_migration_key_in_flight.py | 13 +- .../test_migration_snapshot_consistency.py | 15 ++ worker/async_snapshotting.py | 12 +- worker/fault_tolerance/async_snapshots.py | 143 ++++++++++++------ worker/transactional_protocols/aria.py | 14 +- worker/worker_service.py | 46 ++++-- 12 files changed, 436 insertions(+), 201 deletions(-) diff --git a/coordinator/coordinator_metadata.py b/coordinator/coordinator_metadata.py index 1e6fb6cc..18246f31 100644 --- a/coordinator/coordinator_metadata.py +++ b/coordinator/coordinator_metadata.py @@ -66,6 +66,10 @@ def __init__(self, networking: NetworkingManager, s3_client: S3Client) -> None: # Deferred graph update: held until migration completes self._pending_graph: StateflowGraph | None = None + # Bumped on every recovery and sent to the workers, which drop migration + # messages from any other incarnation (sent before the recovery). + self.incarnation: int = 0 + # Checkpoint-and-resume migration FT self._migration_checkpoint_blob: bytes | None = None # The graph the current checkpoint migrates to @@ -73,7 +77,10 @@ def __init__(self, networking: NetworkingManager, s3_client: S3Client) -> None: # Set once every worker finished its transfer: snapshots carrying a later # epoch were cut after the last migrated key landed and need no checkpoint. self._migration_checkpoint_until_epoch: int | None = None - self._migration_checkpoint_baseline_snap_id: int = -1 + self._migration_checkpoint_baseline_snap_id: int = -1 # >= 0 while waiting for the checkpoint + # Id of the checkpoint snapshot (reported by the workers). Only snapshots + # from it on carry the checkpoint; earlier ones were cut before Phase B. + self._migration_checkpoint_snap_id: int | None = None self._migration_checkpoint_snapshot_complete: asyncio.Event = asyncio.Event() async def start_kafka_metadata_producer(self) -> None: @@ -128,6 +135,7 @@ async def start_recovery_process(self, workers_to_remove: set[Worker]) -> None: async def send_recovery_to_participating_workers(self) -> None: operator_partition_locations = self.worker_pool.get_operator_partition_locations() snap_id: int = self.get_current_completed_snapshot_id() + self.incarnation += 1 worker_assignments = self.worker_pool.get_worker_assignments() participating_workers: list[Worker] = self.worker_pool.get_participating_workers() self.worker_is_healthy = {worker.worker_id: asyncio.Event() for worker in participating_workers} @@ -154,6 +162,7 @@ async def send_recovery_to_participating_workers(self) -> None: self.submitted_graph.operator_state_backend, snap_id, self.submitted_graph, + self.incarnation, ), msg_type=MessageType.InitRecovery, ), @@ -200,8 +209,17 @@ def register_snapshot( epoch_counter: int, t_counter: int, pool: concurrent.futures.ProcessPoolExecutor, + *, + is_migration_checkpoint: bool = False, ) -> None: self.worker_snapshot_ids[worker_id] = snapshot_id + if ( + is_migration_checkpoint + and self._migration_checkpoint_blob is not None + and self._migration_checkpoint_snap_id is None + ): + self._migration_checkpoint_snap_id = snapshot_id + logging.warning(f"Migration checkpoint is snapshot {snapshot_id}") for operator_partition, pio in partial_input_offsets.items(): self.completed_input_offsets[operator_partition] = max( pio, @@ -233,16 +251,10 @@ def register_snapshot( logging.warning( f"Post-migration snapshot {current_completed_snapshot} completed — graph finalized", ) - # Detect migration checkpoint snapshot completion via baseline ID comparison - if ( - self._migration_checkpoint_baseline_snap_id >= 0 - and current_completed_snapshot > self._migration_checkpoint_baseline_snap_id - ): + # Detect migration checkpoint snapshot completion + if self._migration_checkpoint_baseline_snap_id >= 0 and self.migration_checkpoint_completed(): self._migration_checkpoint_snapshot_complete.set() - logging.warning( - f"Migration checkpoint snapshot {current_completed_snapshot} completed " - f"(baseline was {self._migration_checkpoint_baseline_snap_id})", - ) + logging.warning(f"Migration checkpoint snapshot {self._migration_checkpoint_snap_id} completed") self._migration_checkpoint_baseline_snap_id = -1 if current_completed_snapshot != self.prev_completed_snapshot_id: logging.warning(f"Cluster completed snapshot: {current_completed_snapshot}") @@ -259,7 +271,7 @@ def register_snapshot( self.completed_out_offsets, self.completed_epoch_counter, self.completed_t_counter, - self._migration_checkpoint_blob, + self._migration_checkpoint_blob if self.migration_checkpoint_completed() else None, ), ) self.s3_client.put_object( @@ -276,6 +288,13 @@ def register_snapshot( current_completed_snapshot, ) + def migration_checkpoint_completed(self) -> bool: + """Whether the cluster has completed the checkpoint snapshot (or a later one).""" + return ( + self._migration_checkpoint_snap_id is not None + and self.get_current_completed_snapshot_id() >= self._migration_checkpoint_snap_id + ) + def get_current_completed_snapshot_id(self) -> int: if self.worker_snapshot_ids: return min(self.worker_snapshot_ids.values()) @@ -375,8 +394,11 @@ def set_migration_checkpoint( self, new_graph: StateflowGraph, operator_partition_locations: dict, + from_snapshot_id: int | None = None, ) -> None: - """Serialize migration metadata for inclusion in the next sequencer file.""" + """Serialize migration metadata for the sequencer files of the checkpoint + snapshot and later ones. from_snapshot_id is the checkpoint snapshot if it + is already known (a resumed migration), else the workers report it.""" self._migration_checkpoint_blob = cloudpickle_serialization( { "new_graph": new_graph, @@ -385,6 +407,7 @@ def set_migration_checkpoint( ) self._migration_checkpoint_graph = new_graph self._migration_checkpoint_until_epoch = None + self._migration_checkpoint_snap_id = from_snapshot_id def end_migration_checkpoint_after(self, transfer_done_epoch: int) -> None: """Every migrated key landed by transfer_done_epoch. Snapshots cut up to that @@ -397,6 +420,7 @@ def clear_migration_checkpoint(self) -> None: self._migration_checkpoint_blob = None self._migration_checkpoint_graph = None self._migration_checkpoint_until_epoch = None + self._migration_checkpoint_snap_id = None async def submit_stateflow_graph( self, diff --git a/coordinator/coordinator_service.py b/coordinator/coordinator_service.py index 23c27fad..55b4eaac 100755 --- a/coordinator/coordinator_service.py +++ b/coordinator/coordinator_service.py @@ -112,6 +112,9 @@ def __init__(self) -> None: self.aria_metadata: AriaSyncMetadata | None = None self.migration_metadata: MigrationMetadata | None = None self.workers_that_re_registered: list[Worker] = [] + # Workers a barrier or migration message could not be delivered to. The + # protocol cannot progress without them, so they are handed to recovery. + self.unreachable_workers: list[Worker] = [] self.recovery_lock: asyncio.Lock = asyncio.Lock() self.recovery_state: RecoveryState = RecoveryState.IDLE @@ -365,12 +368,10 @@ async def _handle_migration_repartitioning_done( return # Set migration checkpoint blob BEFORE telling workers to proceed. - # Workers will take a checkpoint snapshot after receiving - # MigrationRepartitioningDone; when the SnapID arrives the coordinator - # writes the sequencer file which will include this blob. - # Use max(..., 0) so the baseline is always >= 0 — after recovery - # worker_snapshot_ids are reset to -1, but we still need the - # checkpoint detection (baseline >= 0) to fire on the next snapshot. + # Workers take the checkpoint snapshot after MigrationRepartitioningDone + # and flag it in its SnapID; the blob goes into the sequencer files of + # that snapshot and later ones, never of a snapshot cut before Phase B. + # baseline >= 0 marks that we are waiting for the checkpoint. baseline = max(self.coordinator.get_current_completed_snapshot_id(), 0) self.coordinator._migration_checkpoint_baseline_snap_id = baseline # noqa: SLF001 operator_partition_locations = self.coordinator.worker_pool.get_operator_partition_locations() @@ -438,7 +439,7 @@ async def _handle_migration_init_done( # (before workers took their snapshot). By now the checkpoint snapshot # may already be complete. Handle the race, then launch the finalizer. baseline = self.coordinator._migration_checkpoint_baseline_snap_id # noqa: SLF001 - if baseline >= 0 and self.coordinator.get_current_completed_snapshot_id() > baseline: + if baseline >= 0 and self.coordinator.migration_checkpoint_completed(): self.coordinator._migration_checkpoint_snapshot_complete.set() # noqa: SLF001 self.coordinator._migration_checkpoint_baseline_snap_id = -1 # noqa: SLF001 logging.warning("MIGRATION | Checkpoint snapshot already completed (race handled)") @@ -504,6 +505,7 @@ async def _handle_snap_id( epoch_counter, t_counter, sn_size, + is_migration_checkpoint, ) = self.networking.decode_message(data) snapshot_time = end - start @@ -527,6 +529,7 @@ async def _handle_snap_id( epoch_counter, t_counter, pool, + is_migration_checkpoint=is_migration_checkpoint, ) # Check for deferred migration re-trigger: fire once the first @@ -902,36 +905,55 @@ async def finalize_migration_repartition(self) -> None: self.aria_metadata.stop_in_next_epoch() logging.warning("Sending MigrationRepartitioningDone to all workers (protocol will stop)") - async with asyncio.TaskGroup() as tg: - for worker in self.coordinator.worker_pool.get_participating_workers(): - tg.create_task( - self.protocol_networking.send_message( - worker.worker_ip, - worker.worker_port, - msg=b"", - msg_type=MessageType.MigrationRepartitioningDone, - serializer=Serializer.NONE, - ), - ) + await self._send_to_participating_workers( + lambda worker: worker.worker_port, + b"", + MessageType.MigrationRepartitioningDone, + Serializer.NONE, + ) async def finalize_migration(self) -> None: - async with asyncio.TaskGroup() as tg: - for worker in self.coordinator.worker_pool.get_participating_workers(): - logging.warning(f"Sending MigrationDone to : {worker}") - tg.create_task( - self.protocol_networking.send_message( - worker.worker_ip, - worker.worker_port, - msg=( - self.migration_metadata.epoch_counter, - self.migration_metadata.t_counter, - self.migration_metadata.input_offsets, - self.migration_metadata.output_offsets, - ), - msg_type=MessageType.MigrationDone, - serializer=Serializer.MSGPACK, - ), + logging.warning("Sending MigrationDone to all workers") + await self._send_to_participating_workers( + lambda worker: worker.worker_port, + ( + self.migration_metadata.epoch_counter, + self.migration_metadata.t_counter, + self.migration_metadata.input_offsets, + self.migration_metadata.output_offsets, + ), + MessageType.MigrationDone, + Serializer.MSGPACK, + ) + + async def _send_to_participating_workers( + self, + port_of: Callable[[Worker], int], + msg: tuple | bytes, + msg_type: MessageType, + serializer: Serializer, + ) -> None: + """Send to every participating worker. A worker the message cannot be + delivered to is handed to recovery: waiting on it would deadlock.""" + workers = self.coordinator.worker_pool.get_participating_workers() + + async def send(worker: Worker) -> bool: + try: + return await self.protocol_networking.send_message( + worker.worker_ip, + port_of(worker), + msg=msg, + msg_type=msg_type, + serializer=serializer, ) + except OSError: + return False + + delivered = await asyncio.gather(*[send(worker) for worker in workers]) + for worker, ok in zip(workers, delivered, strict=True): + if ok is False: + logging.error(f"Could not deliver {msg_type.name} to worker {worker.worker_id}; recovering it") + self.unreachable_workers.append(worker) async def _wait_for_migration_checkpoint_and_finalize(self, mt: MessageType) -> None: """Wait for the migration checkpoint snapshot to complete, then finalize migration.""" @@ -972,17 +994,12 @@ async def finalize_worker_sync( message: tuple | bytes, serializer: Serializer = Serializer.MSGPACK, ) -> None: - async with asyncio.TaskGroup() as tg: - for worker in self.coordinator.worker_pool.get_participating_workers(): - tg.create_task( - self.protocol_networking.send_message( - worker.worker_ip, - worker.protocol_port, - msg=message, - msg_type=msg_type, - serializer=serializer, - ), - ) + await self._send_to_participating_workers( + lambda worker: worker.protocol_port, + message, + msg_type, + serializer, + ) async def worker_wants_to_proceed(self) -> None: async with asyncio.TaskGroup() as tg: @@ -1207,6 +1224,7 @@ async def _perform_recovery(self, workers_to_remove: set[Worker]) -> None: self.coordinator.set_migration_checkpoint( resumed_migration["new_graph"], resumed_migration["operator_partition_locations"], + from_snapshot_id=snap_id, ) self.migration_in_progress = True @@ -1249,7 +1267,9 @@ async def heartbeat_monitor_coroutine(self) -> None: ) # Add workers that re-registered (same IP/ports) to the failed set - if (workers_to_remove or self.workers_that_re_registered) and self.recovery_state == RecoveryState.IDLE: + if ( + workers_to_remove or self.workers_that_re_registered or self.unreachable_workers + ) and self.recovery_state == RecoveryState.IDLE: logging.warning( f"heartbeat_monitor | entering recovery lock " f"dead={workers_to_remove} re_registered={self.workers_that_re_registered} " @@ -1265,6 +1285,8 @@ async def heartbeat_monitor_coroutine(self) -> None: re_registered_set = set(self.workers_that_re_registered) workers_to_remove.update(re_registered_set) self.workers_that_re_registered = [] + workers_to_remove.update(self.unreachable_workers) + self.unreachable_workers = [] await self._perform_recovery(workers_to_remove) except Exception as e: diff --git a/styx-package/styx/common/tcp_networking.py b/styx-package/styx/common/tcp_networking.py index 25e10673..b5799336 100644 --- a/styx-package/styx/common/tcp_networking.py +++ b/styx-package/styx/common/tcp_networking.py @@ -96,7 +96,8 @@ async def create_connection(self, host: str, port: int) -> bool: success = False return success - async def send_message(self, message: bytes) -> None: + async def send_message(self, message: bytes) -> bool: + """Write the message, reconnecting on a broken connection; False if every retry failed.""" if self.writer is None: msg = "Writer is not initialized (connection not established)." raise ConnectionError(msg) @@ -129,6 +130,8 @@ async def send_message(self, message: bytes) -> None: logging.error( f"Cannot send_message_rq_rs to worker {self.target_host}:{self.target_port}", ) + return False + return True async def send_message_rq_rs(self, message: bytes) -> bytes | None: if self.writer is None or self.reader is None: @@ -255,6 +258,9 @@ def __init__( self.peers: dict[int, tuple[str, int, int]] = {} self.wait_remote_key_event: dict[OperatorPartition, dict[K, asyncio.Event]] = {} + # Cluster incarnation (bumped on every recovery); key requests carry it so + # a reply to a request sent before a recovery is dropped. + self.incarnation: int = 0 # ACK batching: per-peer buffer flushed on the next event-loop tick. # Coalesces many tiny `send_message(Ack)` calls (one per leaf txn) @@ -344,7 +350,8 @@ async def send_message( msg: tuple | bytes, msg_type: int, serializer: Serializer = Serializer.CLOUDPICKLE, - ) -> None: + ) -> bool: + """Send the message; False if it could not be delivered.""" msg = self.encode_message(msg=msg, msg_type=msg_type, serializer=serializer) # Hot path: lock-free. dict.get and __next__ have no await points, so # asyncio's cooperative scheduling makes them atomic. @@ -353,7 +360,7 @@ async def send_message( await self.create_socket_connection(host, port) pool = self.pools[(host, port)] socket_conn = next(pool) - await socket_conn.send_message(msg) + return await socket_conn.send_message(msg) def set_peers(self, peers: dict[int, tuple[str, int, int]]) -> None: self.peers = peers @@ -389,6 +396,7 @@ async def request_key( old_partition, self.host_name, self.host_port - 1000, + self.incarnation, ), MessageType.RequestRemoteKey, serializer=Serializer.MSGPACK, diff --git a/tests/unit/coordinator/test_coordinator_metadata.py b/tests/unit/coordinator/test_coordinator_metadata.py index 58616639..84bccd7a 100644 --- a/tests/unit/coordinator/test_coordinator_metadata.py +++ b/tests/unit/coordinator/test_coordinator_metadata.py @@ -372,7 +372,7 @@ def test_set_and_clear(self): class TestRegisterSnapshotCheckpointDetection: - def test_checkpoint_event_fires_when_snapshot_exceeds_baseline(self): + def test_checkpoint_event_fires_when_checkpoint_snapshot_completes(self): c = _coordinator() pool = MagicMock() c.worker_snapshot_ids = {0: 0} @@ -381,11 +381,24 @@ def test_checkpoint_event_fires_when_snapshot_exceeds_baseline(self): c._migration_checkpoint_baseline_snap_id = 0 c._migration_checkpoint_blob = b"fake" - c.register_snapshot(0, 1, {}, {}, 1, 1, pool) + c.register_snapshot(0, 1, {}, {}, 1, 1, pool, is_migration_checkpoint=True) assert c._migration_checkpoint_snapshot_complete.is_set() assert c._migration_checkpoint_baseline_snap_id == -1 + def test_checkpoint_event_ignores_other_snapshots(self): + # A snapshot started before Phase B is not the checkpoint. + c = _coordinator() + pool = MagicMock() + c.worker_snapshot_ids = {0: 0} + c.prev_completed_snapshot_id = 0 + c._migration_checkpoint_baseline_snap_id = 0 + c._migration_checkpoint_blob = b"fake" + + c.register_snapshot(0, 1, {}, {}, 1, 1, pool) + + assert not c._migration_checkpoint_snapshot_complete.is_set() + def test_checkpoint_event_does_not_fire_when_below_baseline(self): c = _coordinator() pool = MagicMock() diff --git a/tests/unit/coordinator/test_migration_checkpoint_lifetime.py b/tests/unit/coordinator/test_migration_checkpoint_lifetime.py index 3a3de155..1b8e4e5e 100644 --- a/tests/unit/coordinator/test_migration_checkpoint_lifetime.py +++ b/tests/unit/coordinator/test_migration_checkpoint_lifetime.py @@ -28,11 +28,11 @@ def _graph(n_partitions: int) -> StateflowGraph: return g -def _complete_snapshot(c: Coordinator, snapshot_id: int, epoch: int): +def _complete_snapshot(c: Coordinator, snapshot_id: int, epoch: int, *, checkpoint: bool = False): """All (two) workers register snapshot_id; returns the sequencer tuple written to S3.""" c.s3_client.put_object.reset_mock() for worker_id in (0, 1): - c.register_snapshot(worker_id, snapshot_id, {}, {}, epoch, 0, MagicMock()) + c.register_snapshot(worker_id, snapshot_id, {}, {}, epoch, 0, MagicMock(), is_migration_checkpoint=checkpoint) body = c.s3_client.put_object.call_args.kwargs["Body"] return zstd_msgpack_deserialization(body) @@ -42,7 +42,7 @@ def _coordinator_with_checkpoint() -> Coordinator: c.s3_client = MagicMock() c.worker_snapshot_ids = {0: 2, 1: 2} c.prev_completed_snapshot_id = 2 - c.set_migration_checkpoint(_graph(8), {}) + c.set_migration_checkpoint(_graph(8), {}, from_snapshot_id=3) return c @@ -64,10 +64,42 @@ def test_checkpoint_kept_while_transfer_runs(self): def test_new_checkpoint_resets_the_end(self): c = _coordinator_with_checkpoint() c.end_migration_checkpoint_after(transfer_done_epoch=50) - c.set_migration_checkpoint(_graph(16), {}) + c.set_migration_checkpoint(_graph(16), {}, from_snapshot_id=3) assert _complete_snapshot(c, 3, epoch=60)[4] is not None +class TestCheckpointTagging: + def test_snapshot_cut_before_phase_b_gets_no_checkpoint(self): + c = Coordinator(MagicMock(), MagicMock()) + c.s3_client = MagicMock() + c.worker_snapshot_ids = {0: 2, 1: 2} + c.prev_completed_snapshot_id = 2 + c.set_migration_checkpoint(_graph(8), {}) + c._migration_checkpoint_baseline_snap_id = 2 # waiting for the checkpoint + + # Started before Phase B, completes after the checkpoint metadata was set. + assert _complete_snapshot(c, 3, epoch=40)[4] is None + assert not c._migration_checkpoint_snapshot_complete.is_set() + + # The checkpoint itself, flagged by the workers. + assert _complete_snapshot(c, 4, epoch=41, checkpoint=True)[4] is not None + assert c._migration_checkpoint_snapshot_complete.is_set() + # Snapshots after it keep the checkpoint until the transfer ends. + assert _complete_snapshot(c, 5, epoch=45)[4] is not None + + def test_checkpoint_not_complete_until_every_worker_has_it(self): + c = Coordinator(MagicMock(), MagicMock()) + c.s3_client = MagicMock() + c.worker_snapshot_ids = {0: 2, 1: 2} + c.prev_completed_snapshot_id = 2 + c.set_migration_checkpoint(_graph(8), {}) + c._migration_checkpoint_baseline_snap_id = 2 + c.register_snapshot(0, 3, {}, {}, 41, 0, MagicMock(), is_migration_checkpoint=True) + assert not c.migration_checkpoint_completed() + c.register_snapshot(1, 3, {}, {}, 41, 0, MagicMock(), is_migration_checkpoint=True) + assert c.migration_checkpoint_completed() + + class TestMigrationDoneEpochs: async def test_transfer_done_epoch_is_the_latest_reported(self): m = MigrationMetadata(2) @@ -191,3 +223,29 @@ async def test_ordinary_recovery_stays_ordinary(self): assert svc.migration_in_progress is False assert svc._deferred_migration_graph is None c.revert_worker_pool_to_submitted_graph.assert_not_called() + + +class TestUnreachableWorkers: + def _service_with_workers(self, *results): + c = Coordinator(MagicMock(), MagicMock()) + svc = _service(c) + svc.unreachable_workers = [] + workers = [MagicMock(worker_id=i, worker_ip=f"w{i}", protocol_port=6000 + i) for i in range(len(results))] + c.worker_pool = MagicMock(get_participating_workers=MagicMock(return_value=workers)) + svc.protocol_networking.send_message = AsyncMock(side_effect=list(results)) + return svc, workers + + async def test_worker_a_barrier_reply_cannot_reach_is_handed_to_recovery(self): + svc, workers = self._service_with_workers(True, False) + await svc.finalize_worker_sync(MessageType.SyncCleanup, (True,)) + assert svc.unreachable_workers == [workers[1]] + + async def test_connection_error_counts_as_undelivered(self): + svc, workers = self._service_with_workers(ConnectionRefusedError(), True) + await svc.finalize_worker_sync(MessageType.SyncCleanup, (True,)) + assert svc.unreachable_workers == [workers[0]] + + async def test_all_delivered(self): + svc, _ = self._service_with_workers(True, True) + await svc.finalize_worker_sync(MessageType.SyncCleanup, (False,)) + assert svc.unreachable_workers == [] diff --git a/tests/unit/worker/test_async_snapshots.py b/tests/unit/worker/test_async_snapshots.py index 5fd97ba9..7877a968 100644 --- a/tests/unit/worker/test_async_snapshots.py +++ b/tests/unit/worker/test_async_snapshots.py @@ -24,14 +24,7 @@ def test_defaults(self): assert s.worker_id == 0 assert s.n_assigned_partitions == 4 assert s.snapshot_id == 1 - assert s.completed_snapshots == 0 - assert s.snapshot_in_progress is False - assert s.snapshot_start == 0.0 - assert s.current_input_offsets is None - assert s.current_output_offsets is None - assert s.current_epoch_counter == -1 - assert s.current_t_counter == -1 - assert s.total_snapshot_size == 0 + assert s._in_progress == {} def test_zero_partitions(self): s = AsyncSnapshotsS3(worker_id=1) @@ -64,70 +57,74 @@ def test_sets_plus_one(self): # --------------------------------------------------------------------------- -# register_size +# Snapshot lifecycle # --------------------------------------------------------------------------- -class TestRegisterSize: - def test_accumulates(self): - s = _snap() - s.register_size(100) - s.register_size(200) - assert s.total_snapshot_size == 300 - - -# --------------------------------------------------------------------------- -# start_snapshotting -# --------------------------------------------------------------------------- +def _done(exc=None): + f = MagicMock() + f.cancelled.return_value = False + f.exception.return_value = exc + return f -class TestStartSnapshotting: - async def test_sets_fields(self): - s = _snap() - offsets_in = {("op", 0): 10} - offsets_out = {("op", 0): 5} - s.start_snapshotting(offsets_in, offsets_out, current_epoch_counter=3, current_t_counter=42) - assert s.snapshot_in_progress is True - assert s.snapshot_start > 0 - assert s.current_input_offsets == offsets_in - assert s.current_output_offsets == offsets_out - assert s.current_epoch_counter == 3 - assert s.current_t_counter == 42 - - -# --------------------------------------------------------------------------- -# snapshot_completed_callback -# --------------------------------------------------------------------------- +def _sent_snapshot_ids(mock_socket): + """Snapshot ids in the SnapID messages sent to the coordinator, in order.""" + from styx.common.tcp_networking import NetworkingManager + sent = [c.args[0] for c in mock_socket.socket.return_value.send.call_args_list] + return [NetworkingManager.decode_message(m[8:])[1] for m in sent] -class TestSnapshotCompletedCallback: - def test_increments_completed(self): - s = _snap(n_partitions=4) - s.start_snapshotting({}, {}, 1, 1) - s.snapshot_completed_callback(MagicMock()) - assert s.completed_snapshots == 1 - assert s.snapshot_in_progress is True # not done yet - @patch("worker.fault_tolerance.async_snapshots.socket") - def test_all_completed_sends_message(self, mock_socket): - mock_conn = MagicMock() - mock_socket.socket.return_value = mock_conn - - s = _snap(n_partitions=2, snapshot_id=5) - s.start_snapshotting({("op", 0): 1}, {("op", 0): 0}, 10, 100) - - s.snapshot_completed_callback(MagicMock()) - assert s.snapshot_in_progress is True - - s.snapshot_completed_callback(MagicMock()) - assert s.snapshot_in_progress is False - assert s.completed_snapshots == 0 - assert s.snapshot_id == 6 # incremented - assert s.total_snapshot_size == 0 - - mock_conn.connect.assert_called_once() - mock_conn.send.assert_called_once() - mock_conn.close.assert_called_once() +@patch("worker.fault_tolerance.async_snapshots.socket") +class TestSnapshotLifecycle: + def test_reports_once_every_file_is_stored(self, mock_socket): + s = _snap(snapshot_id=5) + sid = s.start_snapshotting({("op", 0): 1}, {("op", 0): 0}, 10, 100, n_files=2) + assert sid == 5 + assert s.snapshot_id == 6 + s.register_size(sid, 100) + s.file_stored_callback(sid)(_done()) + assert _sent_snapshot_ids(mock_socket) == [] + s.file_stored_callback(sid)(_done()) + assert _sent_snapshot_ids(mock_socket) == [5] + msg = mock_socket.socket.return_value.send.call_args.args[0] + from styx.common.tcp_networking import NetworkingManager + + (_w, _sid, _start, _end, _in, _out, epoch, t_counter, size, checkpoint) = NetworkingManager.decode_message( + msg[8:] + ) + assert (epoch, t_counter, size, checkpoint) == (10, 100, 100, False) + + def test_overlapping_snapshots_keep_their_own_ids_and_metadata(self, mock_socket): + s = _snap(snapshot_id=1) + first = s.start_snapshotting({}, {}, 10, 100, n_files=1) + second = s.start_snapshotting({}, {}, 11, 110, n_files=1) + assert (first, second) == (1, 2) + assert s._in_progress[first].epoch_counter == 10 + assert s._in_progress[second].epoch_counter == 11 + + def test_later_snapshot_waits_for_earlier_one(self, mock_socket): + s = _snap(snapshot_id=1) + first = s.start_snapshotting({}, {}, 10, 100, n_files=1) + second = s.start_snapshotting({}, {}, 11, 110, n_files=1) + s.file_stored_callback(second)(_done()) # the later, smaller one finishes first + assert _sent_snapshot_ids(mock_socket) == [] + s.file_stored_callback(first)(_done()) + assert _sent_snapshot_ids(mock_socket) == [1, 2] + + def test_failed_upload_blocks_it_and_later_snapshots(self, mock_socket): + s = _snap(snapshot_id=1) + first = s.start_snapshotting({}, {}, 10, 100, n_files=1) + second = s.start_snapshotting({}, {}, 11, 110, n_files=1) + s.file_stored_callback(first)(_done(RuntimeError("S3 down"))) + s.file_stored_callback(second)(_done()) + assert _sent_snapshot_ids(mock_socket) == [] + + def test_snapshot_without_files_is_reported_immediately(self, mock_socket): + s = _snap(snapshot_id=3) + s.start_snapshotting({}, {}, 1, 1, n_files=0) + assert _sent_snapshot_ids(mock_socket) == [3] # --------------------------------------------------------------------------- @@ -426,3 +423,13 @@ def test_no_sequencer_files_returns_defaults(self, mock_mk): assert epoch == 0 assert t_counter == 0 assert migration_blob is None + + +@patch("worker.fault_tolerance.async_snapshots.socket") +def test_checkpoint_flag_is_reported(mock_socket): + from styx.common.tcp_networking import NetworkingManager + + s = _snap(snapshot_id=7) + s.start_snapshotting({}, {}, 1, 1, n_files=0, is_migration_checkpoint=True) + msg = mock_socket.socket.return_value.send.call_args.args[0] + assert NetworkingManager.decode_message(msg[8:])[-1] is True diff --git a/tests/unit/worker/test_migration_key_in_flight.py b/tests/unit/worker/test_migration_key_in_flight.py index d0eb6944..962b8ed3 100644 --- a/tests/unit/worker/test_migration_key_in_flight.py +++ b/tests/unit/worker/test_migration_key_in_flight.py @@ -20,10 +20,10 @@ DST = (OP, 1) -def _receive_remote_key_msg(key, found, value) -> bytes: +def _receive_remote_key_msg(key, found, value, incarnation=0) -> bytes: # Drop the 8-byte length prefix, as the receiving TCP reader does. return NetworkingManager.encode_message( - msg=(DST, key, found, value), + msg=(DST, key, found, value, incarnation), msg_type=MessageType.ReceiveRemoteKey, serializer=Serializer.MSGPACK, )[8:] @@ -31,6 +31,7 @@ def _receive_remote_key_msg(key, found, value) -> bytes: def _worker(): w = Worker.__new__(Worker) + w.incarnation = 0 w.networking = SimpleNamespace(decode_message=NetworkingManager.decode_message) w.protocol_networking = MagicMock() w.protocol_networking.wait_remote_key_event = {DST: {"k1": asyncio.Event()}} @@ -59,3 +60,11 @@ async def test_found_none_value_installs_key_and_wakes_waiter(self): w.protocol_networking.key_received.assert_called_once_with(DST, "k1") assert w.local_state.data[DST] == {"k1": None} assert not w.local_state.in_remote_keys("k1", OP, 1) + + async def test_reply_from_before_a_recovery_is_dropped(self): + w = _worker() + w.incarnation = 2 + msg = _receive_remote_key_msg("k1", True, "stale", incarnation=1) + await w._handle_receive_remote_key(msg, MessageType.ReceiveRemoteKey) + w.protocol_networking.key_received.assert_not_called() + assert "k1" not in w.local_state.data[DST] diff --git a/tests/unit/worker/test_migration_snapshot_consistency.py b/tests/unit/worker/test_migration_snapshot_consistency.py index 7ef2d84a..eb48a89d 100644 --- a/tests/unit/worker/test_migration_snapshot_consistency.py +++ b/tests/unit/worker/test_migration_snapshot_consistency.py @@ -110,6 +110,7 @@ async def deliver() -> None: def _worker(link: _Link, worker_id: int, addr: tuple, peers: dict, state: InMemoryOperatorState) -> AriaProtocol: proto = AriaProtocol.__new__(AriaProtocol) proto.id = worker_id + proto.incarnation = 0 proto.peers = peers proto.dns = DNS proto.local_state = state @@ -328,3 +329,17 @@ async def test_keys_per_epoch_are_capped_until_the_epoch_ends(self, monkeypatch) src._migration_epoch_ended.set() await asyncio.wait_for(sender, timeout=2) assert src.local_state.keys_remaining_to_send() == 0 + + +class TestStaleBatches: + async def test_batch_from_before_a_recovery_is_dropped_and_not_acked(self): + link, src, dst = _two_workers(3) + dst.incarnation = 1 # dst already recovered; src still runs the old protocol + + sender = asyncio.create_task(src._continuous_migration_sender()) + await asyncio.sleep(LINK_DELAY_S * 5) + + assert dst.local_state.data[DST] == {} + assert not src._migration_all_acked.is_set() + sender.cancel() + await asyncio.gather(sender, *link.in_flight, return_exceptions=True) diff --git a/worker/async_snapshotting.py b/worker/async_snapshotting.py index 68d32f4f..1a1cc1df 100644 --- a/worker/async_snapshotting.py +++ b/worker/async_snapshotting.py @@ -81,24 +81,28 @@ def take_snapshot(self, metadata: tuple) -> None: topic_partition_output_offsets, epoch_counter, t_counter, + *rest, ) = metadata + is_migration_checkpoint = bool(rest and rest[0]) logging.warning(f"ASYNC_SN | Starting Snapshot at epoch: {epoch_counter}") - self.async_snapshots.start_snapshotting( + snapshot_id = self.async_snapshots.start_snapshotting( topic_partition_offsets, topic_partition_output_offsets, epoch_counter, t_counter, + n_files=len(self.delta_maps), + is_migration_checkpoint=is_migration_checkpoint, ) for operator_partition, delta_map in self.delta_maps.items(): operator_name, partition = operator_partition data_to_store: bytes = zstd_msgpack_serialization(delta_map) - self.async_snapshots.register_size(len(data_to_store)) + self.async_snapshots.register_size(snapshot_id, len(data_to_store)) loop.run_in_executor( self.pool, self.async_snapshots.store_snapshot, - f"data/{operator_name}/{partition}/{self.async_snapshots.snapshot_id}.bin", + f"data/{operator_name}/{partition}/{snapshot_id}.bin", data_to_store, - ).add_done_callback(self.async_snapshots.snapshot_completed_callback) + ).add_done_callback(self.async_snapshots.file_stored_callback(snapshot_id)) self.clear_delta_maps() def snapshotting_controller(self, data: bytes) -> None: diff --git a/worker/fault_tolerance/async_snapshots.py b/worker/fault_tolerance/async_snapshots.py index d99896fe..dfda3301 100644 --- a/worker/fault_tolerance/async_snapshots.py +++ b/worker/fault_tolerance/async_snapshots.py @@ -1,9 +1,12 @@ +from dataclasses import dataclass +import functools import os import socket import time from typing import TYPE_CHECKING import boto3 +from styx.common.logging import logging from styx.common.message_types import MessageType from styx.common.serialization import Serializer, zstd_msgpack_deserialization from styx.common.tcp_networking import NetworkingManager @@ -11,7 +14,7 @@ from worker.fault_tolerance.base_snapshoter import BaseSnapshotter if TYPE_CHECKING: - from collections.abc import Iterable + from collections.abc import Callable, Iterable from concurrent.futures import Future from botocore.client import BaseClient as S3Client @@ -55,6 +58,19 @@ def warm_s3_client() -> None: _get_s3_client() +@dataclass +class _SnapshotInProgress: + start_ms: float + input_offsets: dict[OperatorPartition, int] + output_offsets: dict[OperatorPartition, int] + epoch_counter: int + t_counter: int + files_left: int + is_migration_checkpoint: bool = False + size: int = 0 + failed: bool = False + + class AsyncSnapshotsS3(BaseSnapshotter): def __init__( self, @@ -63,16 +79,13 @@ def __init__( snapshot_id: int = 1, ) -> None: self.worker_id: int = worker_id + # Id of the next snapshot to start self.snapshot_id: int = snapshot_id self.n_assigned_partitions: int = n_assigned_partitions - self.completed_snapshots: int = 0 - self.snapshot_in_progress: bool = False - self.snapshot_start: float = 0.0 - self.current_input_offsets = None - self.current_output_offsets = None - self.current_epoch_counter = -1 - self.current_t_counter = -1 - self.total_snapshot_size = 0 + # Snapshots whose files are still uploading, by id. Each is a delta on top + # of the previous one, so they are reported to the coordinator strictly in + # id order, and never past one that failed to upload. + self._in_progress: dict[int, _SnapshotInProgress] = {} def update_n_assigned_partitions(self, n_assigned_partitions: int) -> None: self.n_assigned_partitions = n_assigned_partitions @@ -80,50 +93,86 @@ def update_n_assigned_partitions(self, n_assigned_partitions: int) -> None: def set_snapshot_id(self, snapshot_id: int) -> None: self.snapshot_id = snapshot_id + 1 - def snapshot_completed_callback(self, _: Future) -> None: - self.completed_snapshots += 1 - if self.completed_snapshots == self.n_assigned_partitions: - snapshot_end = time.time() * 1000 - msg = NetworkingManager.encode_message( - msg=( - self.worker_id, - self.snapshot_id, - self.snapshot_start, - snapshot_end, - self.current_input_offsets, - self.current_output_offsets, - self.current_epoch_counter, - self.current_t_counter, - self.total_snapshot_size, - ), - msg_type=MessageType.SnapID, - serializer=Serializer.MSGPACK, - ) - s = socket.socket() - s.connect((COORDINATOR_HOST, COORDINATOR_PORT)) - s.send(msg) - s.close() - self.snapshot_id += 1 - self.snapshot_in_progress = False - self.completed_snapshots = 0 - self.total_snapshot_size = 0 - def start_snapshotting( self, current_input_offsets: dict[OperatorPartition, int], current_output_offsets: dict[OperatorPartition, int], current_epoch_counter: int, current_t_counter: int, - ) -> None: - self.snapshot_in_progress = True - self.snapshot_start = time.time() * 1000 - self.current_input_offsets = current_input_offsets - self.current_output_offsets = current_output_offsets - self.current_epoch_counter = current_epoch_counter - self.current_t_counter = current_t_counter - - def register_size(self, file_size_bytes: int) -> None: - self.total_snapshot_size += file_size_bytes + n_files: int, + *, + is_migration_checkpoint: bool = False, + ) -> int: + """Start a snapshot of n_files files and return its id. + + A snapshot can start while earlier ones are still uploading; each keeps its + own id and metadata. + """ + snapshot_id = self.snapshot_id + self.snapshot_id += 1 + self._in_progress[snapshot_id] = _SnapshotInProgress( + start_ms=time.time() * 1000, + input_offsets=current_input_offsets, + output_offsets=current_output_offsets, + epoch_counter=current_epoch_counter, + t_counter=current_t_counter, + files_left=n_files, + is_migration_checkpoint=is_migration_checkpoint, + ) + # A snapshot with no files is complete already. + self._report_completed() + return snapshot_id + + def register_size(self, snapshot_id: int, file_size_bytes: int) -> None: + self._in_progress[snapshot_id].size += file_size_bytes + + def file_stored_callback(self, snapshot_id: int) -> Callable[[Future], None]: + return functools.partial(self._file_stored, snapshot_id) + + def _file_stored(self, snapshot_id: int, future: Future) -> None: + snapshot = self._in_progress[snapshot_id] + error = "cancelled" if future.cancelled() else future.exception() + if error is not None: + snapshot.failed = True + logging.error(f"Snapshot {snapshot_id}: a file failed to upload: {error}") + snapshot.files_left -= 1 + self._report_completed() + + def _report_completed(self) -> None: + while self._in_progress: + snapshot_id = min(self._in_progress) + snapshot = self._in_progress[snapshot_id] + if snapshot.files_left > 0: + return + if snapshot.failed: + logging.error( + f"Snapshot {snapshot_id} is incomplete in S3; neither it nor any later snapshot is reported", + ) + return + del self._in_progress[snapshot_id] + self._send_snapshot_id(snapshot_id, snapshot) + + def _send_snapshot_id(self, snapshot_id: int, snapshot: _SnapshotInProgress) -> None: + msg = NetworkingManager.encode_message( + msg=( + self.worker_id, + snapshot_id, + snapshot.start_ms, + time.time() * 1000, + snapshot.input_offsets, + snapshot.output_offsets, + snapshot.epoch_counter, + snapshot.t_counter, + snapshot.size, + snapshot.is_migration_checkpoint, + ), + msg_type=MessageType.SnapID, + serializer=Serializer.MSGPACK, + ) + s = socket.socket() + s.connect((COORDINATOR_HOST, COORDINATOR_PORT)) + s.send(msg) + s.close() @staticmethod def store_snapshot(snapshot_name: str, sn_data: bytes) -> bool: diff --git a/worker/transactional_protocols/aria.py b/worker/transactional_protocols/aria.py index 1b1587e9..63e40786 100644 --- a/worker/transactional_protocols/aria.py +++ b/worker/transactional_protocols/aria.py @@ -78,6 +78,7 @@ def __init__( # noqa: PLR0913 dedup_output_offsets: dict[OperatorPartition, int] | None = None, egress: StyxKafkaBatchEgress | None = None, kafka_consumer: AIOKafkaConsumer | None = None, + incarnation: int = 0, ) -> None: if topic_partition_offsets is None: topic_partition_offsets = {(tp.topic, tp.partition): -1 for tp in topic_partitions} @@ -85,6 +86,8 @@ def __init__( # noqa: PLR0913 output_offsets = {(tp.topic, tp.partition): -1 for tp in topic_partitions} self.id: int = worker_id + # Migration messages from another incarnation (before a recovery) are dropped. + self.incarnation: int = incarnation self.topic_partitions = topic_partitions self.networking = networking @@ -526,7 +529,14 @@ async def _handle_remote_wants_to_proceed(self, _: bytes) -> None: self.ingress.messages_available.set() async def _handle_async_migration(self, data: bytes) -> None: - operator_partition, batch, sender_id, nonce, batch_id = self.networking.decode_message(data) + operator_partition, batch, sender_id, nonce, batch_id, incarnation = self.networking.decode_message(data) + if incarnation != self.incarnation: + # Sent by a protocol from before a recovery: its keys are not ours to take. + logging.warning( + f"MIGRATION | Dropping batch for {operator_partition} from W{sender_id} " + f"of incarnation {incarnation} (now {self.incarnation})", + ) + return logging.debug(f"MIGRATION | Received migration batch for: {operator_partition}") self._apply_migration_batch(operator_partition, batch) @@ -614,7 +624,7 @@ async def _send_migration_batch( self.networking.send_message( worker[0], worker[2], - msg=(operator_partition, k_v_pairs, self.id, self._migration_nonce, batch_id), + msg=(operator_partition, k_v_pairs, self.id, self._migration_nonce, batch_id, self.incarnation), msg_type=MessageType.AsyncMigration, serializer=Serializer.MSGPACK, ), diff --git a/worker/worker_service.py b/worker/worker_service.py index c7a6b33c..ff226aeb 100644 --- a/worker/worker_service.py +++ b/worker/worker_service.py @@ -135,10 +135,13 @@ def __init__(self, thread_idx: int) -> None: self.migration_completed: asyncio.Event = asyncio.Event() self._migration_generation: int = 0 # bumped on recovery to invalidate zombie Phase B tasks + # Cluster incarnation from the coordinator (bumped on every recovery). Migration + # messages carry it, and those from another incarnation are dropped. + self.incarnation: int = 0 # Post-checkpoint recovery: workers exchange which keys each still holds # for the others' new partitions (all-to-all) before the protocol starts. - self._recovery_hash_tag: int = -1 # snapshot id of the recovery, -1 = no exchange + self._recovery_hash_tag: int = -1 # incarnation of the recovery, -1 = no exchange self._recovery_hashes_out: dict[OperatorPartition, dict[Any, tuple[int, int]]] = {} self._recovery_hash_senders: set[int] = set() self._recovery_hashes_received: asyncio.Event = asyncio.Event() @@ -274,6 +277,7 @@ async def _reset_protocol_networking(self) -> None: ) self.protocol_networking.set_worker_id(self.id) self.protocol_networking.set_peers(self.peers) + self.protocol_networking.incarnation = self.incarnation async def _send_snap_assigned(self, snapshot_id: int) -> None: await self.networking.send_message( @@ -350,6 +354,7 @@ async def _handle_receive_execution_plan(self, data: bytes, _: MessageType) -> N snapshotting_port=self.snapshotting_port, epoch_counter=epoch, t_counter=t_counter, + incarnation=self.incarnation, ) logging.warning("Starting Function Execution Protocol") self.function_execution_protocol.start() @@ -385,16 +390,16 @@ async def _handle_init_migration(self, data: bytes, _: MessageType) -> None: logging.error(f"Uncaught exception during migration Phase A: {e}") async def _migration_stop_protocol(self) -> None: - try: - await asyncio.wait_for( - self.function_execution_protocol.wait_stopped(), - timeout=15.0, - ) - except TimeoutError: - logging.warning( - "MIGRATION | Protocol did not stop within 15s via stop_gracefully — force-stopping", - ) - await self.function_execution_protocol.stop() + # The protocol stops at the next epoch boundary. Never cancel it here: a + # slow epoch would be cut mid-commit. If the stop message cannot reach this + # worker, the coordinator hands the worker to recovery, which stops it. + protocol = self.function_execution_protocol + while True: + try: + await asyncio.wait_for(asyncio.shield(protocol.wait_stopped()), timeout=15.0) + break + except TimeoutError: + logging.warning("MIGRATION | Protocol still running 15s after the stop request") logging.warning("MIGRATION | ARIA STOPPED") async def _migration_decode_and_apply_plan(self, data: bytes) -> None: @@ -573,6 +578,7 @@ async def _migration_rebuild_runtime(self) -> None: restart_after_migration=True, egress=old_egress if old_egress.is_running else None, kafka_consumer=old_consumer, + incarnation=self.incarnation, ) async def _handle_receive_migration_hashes( @@ -593,7 +599,10 @@ async def _handle_request_remote_key( # Lock-free: pop() is atomic in CPython; concurrent requests for the same # key are race-safe (the loser gets found=False, which the requester handles). del msg_type - operator_partition, key, old_partition, host, port = self.networking.decode_message(data) + operator_partition, key, old_partition, host, port, incarnation = self.networking.decode_message(data) + if incarnation != self.incarnation: + logging.warning(f"Dropping RequestRemoteKey for {key!r} from incarnation {incarnation}") + return found, value = self.local_state.get_key_to_migrate( operator_partition, key, @@ -602,7 +611,7 @@ async def _handle_request_remote_key( await self.networking.send_message( host, port, - msg=(operator_partition, key, found, value), + msg=(operator_partition, key, found, value, self.incarnation), msg_type=MessageType.ReceiveRemoteKey, serializer=Serializer.MSGPACK, ) @@ -614,7 +623,10 @@ async def _handle_receive_remote_key( ) -> None: # Lock-free: no awaits. del msg_type - operator_partition, key, found, value = self.networking.decode_message(data) + operator_partition, key, found, value, incarnation = self.networking.decode_message(data) + if incarnation != self.incarnation: + logging.warning(f"Dropping ReceiveRemoteKey for {key!r} from incarnation {incarnation}") + return if not found: # The source's async sender already took the key into a batch; the # batch wakes the waiting transaction when it lands here. @@ -693,6 +705,7 @@ async def _handle_migration_repartitioning_done( self.function_execution_protocol.egress.topic_partition_output_offsets, self.function_execution_protocol.sequencer.epoch_counter, self.function_execution_protocol.sequencer.t_counter, + True, # the migration checkpoint ), msg_type=MessageType.SnapTakeSnapshot, serializer=Serializer.MSGPACK, @@ -779,6 +792,7 @@ async def _handle_init_recovery(self, data: bytes, _: MessageType) -> None: self.operator_state_backend, snapshot_id, self.deployed_graph, + self.incarnation, ) = self.networking.decode_message(data) if self.function_execution_protocol is not None: @@ -794,6 +808,7 @@ async def _handle_init_recovery(self, data: bytes, _: MessageType) -> None: mode=MessagingMode.PROTOCOL_PROTOCOL, ) self.protocol_networking.set_worker_id(self.id) + self.protocol_networking.incarnation = self.incarnation self._update_peers(peers) self.protocol_networking.set_peers(self.peers) @@ -875,7 +890,7 @@ async def _handle_init_recovery(self, data: bytes, _: MessageType) -> None: f"[RECOVERY] W{self.id} Re-derived {sum(len(v) for v in keys_to_send.values())} keys to send", ) # Sent to the destinations on ReadyAfterRecovery, once every worker has loaded its snapshot. - self._reset_recovery_hash_exchange(tag=snapshot_id, hashes_out=hashes_out) + self._reset_recovery_hash_exchange(tag=self.incarnation, hashes_out=hashes_out) restart_after_migration = True request_id_to_t_id_map = await self.get_sequencer_assignments_before_failure( @@ -899,6 +914,7 @@ async def _handle_init_recovery(self, data: bytes, _: MessageType) -> None: restart_after_recovery=not restart_after_migration, restart_after_migration=restart_after_migration, dedup_output_offsets=all_dedup_out_offsets, + incarnation=self.incarnation, ) self.function_execution_protocol.start() From 3f87c3e45a51cced03446643caa4ef9431964dd4 Mon Sep 17 00:00:00 2001 From: Kyriakos Psarakis Date: Sat, 24 Oct 2026 20:18:34 +0200 Subject: [PATCH 11/20] Read the ingress source file in the except-clause test inspect.getsource() cannot read Cython-compiled functions, which the hot-path CI job uses. Co-Authored-By: Claude Opus 5.5 --- tests/unit/worker/test_kafka_ingress.py | 15 ++++++++++----- 1 file changed, 10 insertions(+), 5 deletions(-) diff --git a/tests/unit/worker/test_kafka_ingress.py b/tests/unit/worker/test_kafka_ingress.py index bb8ecb80..e77a0db1 100644 --- a/tests/unit/worker/test_kafka_ingress.py +++ b/tests/unit/worker/test_kafka_ingress.py @@ -381,14 +381,19 @@ def test_except_clause_catches_both_exception_types(self): or linters that might not catch this subtle syntax difference. """ import ast - import inspect - import textwrap + from pathlib import Path - source = textwrap.dedent(inspect.getsource(StyxKafkaIngress._create_kafka_consumer)) + # Parse the module file: the functions may be Cython-compiled, which + # inspect.getsource() cannot read. + source = (Path(__file__).parents[3] / "worker" / "ingress" / "styx_kafka_ingress.py").read_text() tree = ast.parse(source) + handlers = [node for node in ast.walk(tree) if isinstance(node, ast.ExceptHandler)] + assert any( + isinstance(h.type, ast.Tuple) and any(getattr(e, "id", None) == "KafkaConnectionError" for e in h.type.elts) + for h in handlers + ), "expected an except clause catching KafkaConnectionError alongside other errors" - # Find all except handlers in the function - for node in ast.walk(tree): + for node in handlers: if ( isinstance(node, ast.ExceptHandler) and node.type is not None From 54db9521e82479a3d731a536d90da60d9ab085ae Mon Sep 17 00:00:00 2001 From: Kyriakos Psarakis Date: Sat, 24 Oct 2026 20:51:08 +0200 Subject: [PATCH 12/20] Clean up: logging, compaction, test names, real-code tests, E2E state check - Remove the debug 10 s timeout on every barrier wait. - Lower per-message, per-snapshot and step-tracing warnings to debug. - Log a failed graph publication instead of dropping the future's error. - Compaction merged each partition's deltas in S3 listing order (10.bin before 2.bin), kept a partition's first delta when it was not 0.bin (so it was re-applied stale on load), and never dropped tombstones. It now merges in snapshot-id order into the partition's latest file, deletes the older ones, and only then drops tombstones (compaction is off by default). - FT E2E: names and docstrings now say which recovery path each kill hits. - Coordinator FT tests drive the real _perform_recovery instead of re-implementing its migration check. - FT E2E: after the run, every key's stored record is fingerprinted and compared with the loaded dataset, so a lost or emptied key fails the test. Co-Authored-By: Claude Opus 5.5 --- coordinator/coordinator_metadata.py | 19 ++- coordinator/coordinator_service.py | 48 +++---- coordinator/snapshot_compactor.py | 40 +++--- demo/demo-migration-ft-ycsb/client.py | 35 ++++- demo/demo-migration-ycsb/ycsb.py | 10 ++ tests/e2e/test_e2e_migration_ft_ycsb.py | 76 ++++++----- tests/unit/coordinator/service_helpers.py | 45 +++++++ .../test_migration_checkpoint_lifetime.py | 39 +----- .../test_migration_fault_tolerance.py | 124 +++++++----------- .../test_snapshot_compactor_coverage.py | 62 ++++++++- worker/egress/styx_kafka_batch_egress.py | 10 +- worker/transactional_protocols/aria.py | 10 +- worker/worker_service.py | 12 +- 13 files changed, 317 insertions(+), 213 deletions(-) create mode 100644 tests/unit/coordinator/service_helpers.py diff --git a/coordinator/coordinator_metadata.py b/coordinator/coordinator_metadata.py index 18246f31..2ee95743 100644 --- a/coordinator/coordinator_metadata.py +++ b/coordinator/coordinator_metadata.py @@ -38,6 +38,12 @@ KAFKA_URL: str = os.getenv("KAFKA_URL", None) +def _log_graph_publish_failure(future: asyncio.Future) -> None: + # Clients route by the published graph; a failed publish must not go unnoticed. + if not future.cancelled() and future.exception() is not None: + logging.error(f"Publishing the new graph to styx-metadata failed: {future.exception()}") + + class Coordinator: def __init__(self, networking: NetworkingManager, s3_client: S3Client) -> None: self.networking = networking @@ -233,7 +239,7 @@ def register_snapshot( self.completed_epoch_counter = epoch_counter self.completed_t_counter = t_counter current_completed_snapshot: int = self.get_current_completed_snapshot_id() - logging.warning( + logging.debug( f"register_snapshot | worker={worker_id} snap={snapshot_id} " f"cluster_min={current_completed_snapshot} prev={self.prev_completed_snapshot_id} " f"baseline={self._migration_checkpoint_baseline_snap_id} " @@ -305,7 +311,7 @@ async def update_stateflow_graph(self, new_stateflow_graph: StateflowGraph) -> N raise NotAStateflowGraphError # TODO the cluster was balanced by the previous deployment, if the graph is complex it might be # unbalanced after the update - logging.warning( + logging.debug( f"update_stateflow_graph | updating worker pool operators " f"(new partitions: { {n: op.n_partitions for n, op in new_stateflow_graph.nodes.items()} })", ) @@ -326,7 +332,7 @@ async def update_stateflow_graph(self, new_stateflow_graph: StateflowGraph) -> N # recovery, it harmlessly re-sets the field (which _perform_recovery # already handles via saved_migration_in_progress). self._pending_graph = new_stateflow_graph - logging.warning("update_stateflow_graph | _pending_graph set (before send)") + logging.debug("update_stateflow_graph | _pending_graph set (before send)") worker_assignments = self.worker_pool.get_worker_assignments() tasks = [ @@ -345,7 +351,7 @@ async def update_stateflow_graph(self, new_stateflow_graph: StateflowGraph) -> N for worker in self.worker_pool.get_participating_workers() ] await asyncio.gather(*tasks) - logging.warning("update_stateflow_graph | InitMigration sent to all workers") + logging.debug("update_stateflow_graph | InitMigration sent to all workers") def revert_worker_pool_to_submitted_graph(self) -> None: """Revert worker pool operators to match the current submitted_graph. @@ -373,9 +379,9 @@ def revert_worker_pool_to_submitted_graph(self) -> None: def finalize_graph_update(self) -> None: """Commit the deferred graph update after migration completes successfully.""" if self._pending_graph is None: - logging.warning("finalize_graph_update | no-op (_pending_graph is None)") + logging.debug("finalize_graph_update | no-op (_pending_graph is None)") return - logging.warning("finalize_graph_update | committing pending graph to submitted_graph") + logging.debug("finalize_graph_update | committing pending graph to submitted_graph") self.submitted_graph = self._pending_graph self._pending_graph = None # Write to Kafka via a background task; the producer is already running @@ -389,6 +395,7 @@ def finalize_graph_update(self) -> None: value=serialized_graph, ), ) + self._pending_kafka_write.add_done_callback(_log_graph_publish_failure) def set_migration_checkpoint( self, diff --git a/coordinator/coordinator_service.py b/coordinator/coordinator_service.py index 55b4eaac..c0b76b1b 100755 --- a/coordinator/coordinator_service.py +++ b/coordinator/coordinator_service.py @@ -347,7 +347,7 @@ async def _handle_migration_repartitioning_done( logging.warning("Migration repartitioning done received!") async with self.networking_locks[mt]: - logging.warning("MigrationRepartitioningDone | acquired lock") + logging.debug("MigrationRepartitioningDone | acquired lock") # Re-check inside the lock to close the TOCTOU window: # recovery may have cleared migration_in_progress between the # outer guard and lock acquisition. @@ -383,10 +383,10 @@ async def _handle_migration_repartitioning_done( f"MIGRATION | Checkpoint blob set, baseline snap_id={baseline}", ) - logging.warning("MigrationRepartitioningDone | calling finalize_migration_repartition") + logging.debug("MigrationRepartitioningDone | calling finalize_migration_repartition") await self.finalize_migration_repartition() await self.migration_metadata.cleanup(mt) - logging.warning("MigrationRepartitioningDone | done") + logging.debug("MigrationRepartitioningDone | done") async def _handle_migration_init_done( self, @@ -399,7 +399,7 @@ async def _handle_migration_init_done( return mt = MessageType.MigrationInitDone async with self.networking_locks[mt]: - logging.warning("MigrationInitDone | acquired lock") + logging.debug("MigrationInitDone | acquired lock") # Re-check inside the lock (TOCTOU protection). if not self.migration_in_progress: logging.warning("Dropping stale MigrationInitDone (recovery intervened)") @@ -427,9 +427,9 @@ async def _handle_migration_init_done( n_workers = len(self.coordinator.worker_pool.get_participating_workers()) self.aria_metadata = AriaSyncMetadata(n_workers) - logging.warning("MigrationInitDone | closing protocol connections") + logging.debug("MigrationInitDone | closing protocol connections") await self.protocol_networking.close_all_connections() - logging.warning( + logging.debug( f"MigrationInitDone | protocol connections closed, " f"current_snap={self.coordinator.get_current_completed_snapshot_id()} " f"baseline={self.coordinator._migration_checkpoint_baseline_snap_id}", # noqa: SLF001 @@ -445,7 +445,7 @@ async def _handle_migration_init_done( logging.warning("MIGRATION | Checkpoint snapshot already completed (race handled)") # Launch background task to wait for the checkpoint snapshot, then finalize - logging.warning("MigrationInitDone | launching checkpoint finalizer task") + logging.debug("MigrationInitDone | launching checkpoint finalizer task") self._migration_checkpoint_task = asyncio.ensure_future( self._wait_for_migration_checkpoint_and_finalize(mt), ) @@ -602,7 +602,7 @@ async def protocol_controller(self, data: bytes) -> None: ) return if self.migration_in_progress: - logging.warning(f"PROTOCOL_MSG during migration: {mt.name}") + logging.debug(f"PROTOCOL_MSG during migration: {mt.name}") await handler(data) # ------------------------ @@ -713,7 +713,7 @@ async def _handle_sync_cleanup(self, data: bytes) -> None: stop_flag = self.aria_metadata.stop_next_epoch if stop_flag: - logging.warning( + logging.debug( "SyncCleanup | sync_complete, sending stop_gracefully=True to all workers", ) await self.finalize_worker_sync( @@ -722,7 +722,7 @@ async def _handle_sync_cleanup(self, data: bytes) -> None: Serializer.MSGPACK, ) if stop_flag: - logging.warning("SyncCleanup | stop_gracefully sent successfully") + logging.debug("SyncCleanup | stop_gracefully sent successfully") self.aria_metadata.cleanup(epoch_end=True) def _record_epoch_metrics( @@ -793,7 +793,7 @@ async def _handle_migration_done(self, data: bytes) -> None: return mt = MessageType.MigrationDone async with self.networking_locks[mt]: - logging.warning("MigrationDone | acquired lock") + logging.debug("MigrationDone | acquired lock") # Re-check inside the lock (TOCTOU protection). if not self.migration_in_progress: logging.warning("Dropping stale MigrationDone (recovery intervened)") @@ -836,13 +836,13 @@ async def request_handler(reader: StreamReader, writer: StreamWriter) -> None: self.protocol_controller(await reader.readexactly(size)), ) except asyncio.IncompleteReadError as e: - logging.warning(f"Protocol client disconnected unexpectedly: {e}") + logging.info(f"Protocol client disconnected unexpectedly: {e}") except asyncio.CancelledError: pass except Exception as e: logging.warning(f"Protocol request_handler unexpected error: {e}") finally: - logging.warning("Protocol connection closing") + logging.info("Protocol connection closing") writer.close() await writer.wait_closed() @@ -957,10 +957,10 @@ async def send(worker: Worker) -> bool: async def _wait_for_migration_checkpoint_and_finalize(self, mt: MessageType) -> None: """Wait for the migration checkpoint snapshot to complete, then finalize migration.""" - logging.warning("checkpoint_finalizer | waiting for checkpoint snapshot event") + logging.debug("checkpoint_finalizer | waiting for checkpoint snapshot event") await self.coordinator._migration_checkpoint_snapshot_complete.wait() # noqa: SLF001 self.coordinator._migration_checkpoint_snapshot_complete.clear() # noqa: SLF001 - logging.warning( + logging.debug( f"checkpoint_finalizer | event fired, migration_in_progress={self.migration_in_progress}", ) @@ -979,14 +979,14 @@ async def _wait_for_migration_checkpoint_and_finalize(self, mt: MessageType) -> ) # Publish new graph to styx-metadata so ALL clients update partitioners - logging.warning("checkpoint_finalizer | calling finalize_graph_update") + logging.debug("checkpoint_finalizer | calling finalize_graph_update") self.coordinator.finalize_graph_update() # Send MigrationDone to workers - logging.warning("checkpoint_finalizer | sending MigrationDone to workers") + logging.debug("checkpoint_finalizer | sending MigrationDone to workers") await self.finalize_migration() await self.migration_metadata.cleanup(mt) - logging.warning("checkpoint_finalizer | done") + logging.debug("checkpoint_finalizer | done") async def finalize_worker_sync( self, @@ -1106,7 +1106,7 @@ async def _perform_recovery(self, workers_to_remove: set[Worker]) -> None: has_checkpoint_task = ( hasattr(self, "_migration_checkpoint_task") and self._migration_checkpoint_task is not None ) - logging.warning( + logging.debug( f"[RECOVERY] step 0 | cancel checkpoint task={has_checkpoint_task} " f"migration_in_progress={self.migration_in_progress}", ) @@ -1115,7 +1115,7 @@ async def _perform_recovery(self, workers_to_remove: set[Worker]) -> None: with contextlib.suppress(asyncio.CancelledError): await self._migration_checkpoint_task self._migration_checkpoint_task = None - logging.warning("[RECOVERY] step 0 | checkpoint task cancelled") + logging.debug("[RECOVERY] step 0 | checkpoint task cancelled") # 0b) Clear migration_in_progress immediately so any migration handler # that passed its guard check before recovery will see it is False @@ -1125,7 +1125,7 @@ async def _perform_recovery(self, workers_to_remove: set[Worker]) -> None: # failed mid-send). saved_migration_in_progress = self.migration_in_progress self.migration_in_progress = False - logging.warning( + logging.debug( f"[RECOVERY] step 0b | migration_in_progress was {saved_migration_in_progress}, now False", ) @@ -1189,7 +1189,7 @@ async def _perform_recovery(self, workers_to_remove: set[Worker]) -> None: self.coordinator.revert_worker_pool_to_submitted_graph() # Clear remaining migration state so stale messages are dropped. - logging.warning("[RECOVERY] step 1c | clearing remaining migration state") + logging.debug("[RECOVERY] step 1c | clearing remaining migration state") self.coordinator.pre_migration_snapshot_pending = False self.coordinator.post_migration_snapshot_pending = False self.coordinator._migration_checkpoint_baseline_snap_id = -1 # noqa: SLF001 @@ -1270,7 +1270,7 @@ async def heartbeat_monitor_coroutine(self) -> None: if ( workers_to_remove or self.workers_that_re_registered or self.unreachable_workers ) and self.recovery_state == RecoveryState.IDLE: - logging.warning( + logging.debug( f"heartbeat_monitor | entering recovery lock " f"dead={workers_to_remove} re_registered={self.workers_that_re_registered} " f"migration_in_progress={self.migration_in_progress}", @@ -1293,7 +1293,7 @@ async def heartbeat_monitor_coroutine(self) -> None: logging.error(f"Error during recovery: {e}") finally: self.recovery_state = RecoveryState.IDLE - logging.warning("heartbeat_monitor | recovery_state back to IDLE") + logging.debug("heartbeat_monitor | recovery_state back to IDLE") async def send_snapshot_marker(self) -> None: while True: diff --git a/coordinator/snapshot_compactor.py b/coordinator/snapshot_compactor.py index 1c22b46d..ca9eb5e6 100644 --- a/coordinator/snapshot_compactor.py +++ b/coordinator/snapshot_compactor.py @@ -88,35 +88,37 @@ def compact_deltas( sequencer_snapshots_to_delete = [s_sn for _, s_sn in sequencer_snapshots if s_sn != "sequencer/0.bin"] - # Partition Data - data: dict[OperatorPartition, KVPairs] = {} + # Partition Data. For each partition, merge its deltas in snapshot-id order + # into its latest file, then delete the older ones. A crash part-way leaves + # older files that the latest one fully overrides when loading in id order. snapshots_to_delete: list[str] = [] - - for operator_partition, sn_info in snapshots_per_operator_partition.items(): - for _, sn_name in sn_info: + compacted: dict[str, KVPairs] = {} + for sn_info in snapshots_per_operator_partition.values(): + ordered = sorted(sn_info) + merged: KVPairs = {} + for _, sn_name in ordered: partition_data = zstd_msgpack_deserialization( s3_get_bytes(s3_client, SNAPSHOT_BUCKET_NAME, sn_name), ) - if operator_partition in data: - snapshots_to_delete.append(sn_name) - if partition_data: - data[operator_partition].update(partition_data) - else: - data[operator_partition] = partition_data - - for operator_partition, partition_data in data.items(): - operator_name, partition = operator_partition - sn_data: bytes = zstd_msgpack_serialization(partition_data) - snapshot_name: str = f"data/{operator_name}/{partition}/0.bin" - - # Store the primary snapshot after compaction - s3_put_bytes(s3_client, SNAPSHOT_BUCKET_NAME, snapshot_name, sn_data) + if partition_data: + merged.update(partition_data) + latest_name = ordered[-1][1] + # Tombstones (None) still matter while older files exist. + s3_put_bytes(s3_client, SNAPSHOT_BUCKET_NAME, latest_name, zstd_msgpack_serialization(merged)) + snapshots_to_delete.extend(sn_name for _, sn_name in ordered[:-1]) + compacted[latest_name] = merged # Cleanup the compacted deltas logging.warning("Deleting data snapshots: %s", snapshots_to_delete) logging.warning("Deleting sequencer snapshots: %s", sequencer_snapshots_to_delete) cleanup_compacted_files(s3_client, sequencer_snapshots_to_delete, snapshots_to_delete) + # With the older files gone, tombstones are redundant: drop them. + for latest_name, merged in compacted.items(): + if any(v is None for v in merged.values()): + live = {k: v for k, v in merged.items() if v is not None} + s3_put_bytes(s3_client, SNAPSHOT_BUCKET_NAME, latest_name, zstd_msgpack_serialization(live)) + def cleanup_compacted_files( s3_client: S3Client, diff --git a/demo/demo-migration-ft-ycsb/client.py b/demo/demo-migration-ft-ycsb/client.py index c0e99e33..8e61a23f 100644 --- a/demo/demo-migration-ft-ycsb/client.py +++ b/demo/demo-migration-ft-ycsb/client.py @@ -10,6 +10,7 @@ - kill_at_sec: second at which to `docker kill styx-worker-1` (-1 = disabled) """ +import json import multiprocessing from multiprocessing import Pool, cpu_count import os @@ -20,6 +21,7 @@ import sys import time from timeit import default_timer as timer +import zlib import boto3 @@ -133,6 +135,36 @@ def ycsb_init(styx: SyncStyxClient, operator: Operator, num_workers: int = None) print("Data loaded") time.sleep(5) submit_graph(styx) + return partitions + + +def fingerprint(value) -> int | None: + # Same as ycsb.verify on the workers. + return None if value is None else zlib.crc32("\x1f".join(value).encode()) + + +def verify_final_state(partitions: dict[int, dict]) -> None: + """Compare every key's stored record with the loaded dataset. + + YCSB updates write back the record they read, so the final state must equal + the initial dataset: a lost, emptied or misplaced key shows up here even when + every request got exactly one response. + """ + expected = {key: fingerprint(value) for part in partitions.values() for key, value in part.items()} + styx = SyncStyxClient(STYX_HOST, STYX_PORT, kafka_url=KAFKA_URL) + styx.open(consume=True) + futures = {key: styx.send_event(operator=ycsb_operator, key=key, function="verify") for key in expected} + styx.flush() + mismatched = [] + for key, future in futures.items(): + response = future.get().response + stored = response[1] if response is not None else "no response" + if stored != expected[key]: + mismatched.append({"key": key, "stored": stored, "expected": expected[key]}) + styx.close() + print(f"State check: {len(mismatched)} of {len(expected)} keys differ from the loaded dataset") + with open(f"{SAVE_DIR}/state_check.json", "w", encoding="utf-8") as f: + json.dump({"keys": len(expected), "mismatched": len(mismatched), "examples": mismatched[:20]}, f) def transactional_ycsb_generator(operator: Operator): @@ -205,7 +237,7 @@ def main(): region_name=os.getenv("S3_REGION") or "us-east-1", ) styx_client = SyncStyxClient(STYX_HOST, STYX_PORT, kafka_url=KAFKA_URL, s3=s3) - ycsb_init(styx_client, ycsb_operator) + partitions = ycsb_init(styx_client, ycsb_operator) del styx_client time.sleep(5) @@ -224,6 +256,7 @@ def main(): ).sort_values("timestamp").to_csv(f"{SAVE_DIR}/client_requests.csv", index=False) print("Workload completed") + verify_final_state(partitions) if __name__ == "__main__": diff --git a/demo/demo-migration-ycsb/ycsb.py b/demo/demo-migration-ycsb/ycsb.py index 88fa5eba..ee4b170c 100755 --- a/demo/demo-migration-ycsb/ycsb.py +++ b/demo/demo-migration-ycsb/ycsb.py @@ -22,3 +22,13 @@ async def update(ctx: StatefulFunction): value: tuple = ctx.get() ctx.put(value) return ctx.key + + +@ycsb_operator.register +async def verify(ctx: StatefulFunction): + """Fingerprint of the stored record (None if missing), for end-of-run state + checks. Self-contained: it runs on the workers, where this module is absent.""" + import zlib + + value = ctx.get() + return ctx.key, None if value is None else zlib.crc32("\x1f".join(value).encode()) diff --git a/tests/e2e/test_e2e_migration_ft_ycsb.py b/tests/e2e/test_e2e_migration_ft_ycsb.py index 4279f8b0..7a079ae5 100644 --- a/tests/e2e/test_e2e_migration_ft_ycsb.py +++ b/tests/e2e/test_e2e_migration_ft_ycsb.py @@ -1,16 +1,20 @@ """E2E tests for fault tolerance during state migration. These tests verify that Styx maintains exactly-once semantics when a worker -crashes at different points during a migration cycle: - -1. Normal migration (baseline) — no crash -2. Kill during Phase A — crash while protocol is still running, before - repartitioning starts. Recovery should roll back to pre-migration snapshot. -3. Kill during Phase C — crash after repartitioning, during async data - transfer. Recovery should roll back to pre-migration snapshot (no - post-migration snapshot has completed yet). -4. Kill after migration completes — crash after MigrationDone. Recovery - should use the post-migration snapshot with new partition layout. +crashes at different points around a migration (4 -> 8 partitions, 10k keys). +With this dataset the whole migration takes about a second, so a kill at a +fixed second lands in a phase that depends on timing; the path noted for +each test is the one observed on a typical machine: + +1. Normal migration (baseline) — no crash. +2. Kill before the post-migration snapshot — the latest snapshot is the + checkpoint (or one cut during the transfer), so recovery resumes the + migration with the new layout. +3. Kill after the post-migration snapshot — ordinary recovery with the new + layout. +4. Kill well after the migration — ordinary recovery with the new layout. +5. External kill during the migration (Phase A) — no checkpoint yet, so + recovery reverts to the old layout and retries the migration. All tests validate: - No duplicate requests @@ -51,6 +55,15 @@ def _assert_metrics( assert metrics.get("exactly_once_output") is True, metrics assert int(metrics.get("missed messages")) == 0, metrics + # Every key must still hold the record it was loaded with (YCSB updates + # write back what they read), so a lost or emptied key fails here. + state_json = results_dir / "state_check.json" + assert state_json.exists(), f"Missing state check: {state_json}" + with state_json.open("r", encoding="utf-8") as f: + state = json.load(f) + log.info("State check: %s", state) + assert state["mismatched"] == 0, state + @dataclass(frozen=True) class _ClusterParams: @@ -256,20 +269,19 @@ def test_migration_ft_baseline(tmp_path: Path): # --------------------------------------------------------------------------- -# Test: kill during Phase A (protocol running, before repartitioning) +# Test: kill before the post-migration snapshot (resume from the checkpoint) # --------------------------------------------------------------------------- @pytest.mark.e2e_migration -def test_migration_ft_kill_during_phase_a(tmp_path: Path): - """Kill a worker during Phase A — crash while protocol is running. +def test_migration_ft_kill_before_post_migration_snapshot(tmp_path: Path): + """Kill a worker 5 s after the migration starts. - Migration starts at second 30. Kill at second 35: the protocol is still - running and rehash metadata may not have arrived yet. The coordinator - should detect the failure, recover to the pre-migration snapshot, and - retry the migration from scratch. + The transfer has usually finished, but the latest complete snapshot is the + checkpoint (or one cut while keys were still moving), so recovery takes the + new layout, re-derives the keys still to move and resumes the transfer. - Validates: recovery to pre-migration state, exactly-once output. + Validates: exactly-once output across a resumed migration. """ paths = _resolve_paths() results_dir = _make_results_dir(tmp_path) @@ -279,7 +291,7 @@ def test_migration_ft_kill_during_phase_a(tmp_path: Path): total_time=120, warmup_seconds=10, migration_at_sec=30, - kill_at_sec=35, # Shortly after migration starts (Phase A) + kill_at_sec=35, # 5 s after the migration starts ) env = make_test_env() @@ -299,20 +311,18 @@ def test_migration_ft_kill_during_phase_a(tmp_path: Path): # --------------------------------------------------------------------------- -# Test: kill during Phase C (async data transfer, before post-mig snapshot) +# Test: kill after the post-migration snapshot # --------------------------------------------------------------------------- @pytest.mark.e2e_migration -def test_migration_ft_kill_during_phase_c(tmp_path: Path): - """Kill a worker during Phase C — async data transfer in progress. +def test_migration_ft_kill_after_post_migration_snapshot(tmp_path: Path): + """Kill a worker 20 s after the migration starts. - Migration starts at second 30. Kill at second 50: repartitioning has - completed but async data transfer is likely still in progress. No - post-migration snapshot should have completed yet. Recovery should - roll back to pre-migration snapshot with OLD partition layout. + A snapshot without the checkpoint has completed by then, so recovery is an + ordinary one with the new layout. - Validates: recovery to pre-migration state, exactly-once output. + Validates: recovery with the new layout, exactly-once output. """ paths = _resolve_paths() results_dir = _make_results_dir(tmp_path) @@ -322,7 +332,7 @@ def test_migration_ft_kill_during_phase_c(tmp_path: Path): total_time=150, warmup_seconds=10, migration_at_sec=30, - kill_at_sec=50, # During async data transfer (Phase C) + kill_at_sec=50, # 20 s after the migration starts ) env = make_test_env() @@ -350,9 +360,8 @@ def test_migration_ft_kill_during_phase_c(tmp_path: Path): def test_migration_ft_kill_after_migration_done(tmp_path: Path): """Kill a worker after migration completes. - Migration starts at second 30. Kill at second 90: migration should have - completed and at least one post-migration snapshot should exist. Recovery - should use the post-migration snapshot with NEW partition layout. + Migration starts at second 30. Kill at second 90: several snapshots + with the new layout exist. Recovery is an ordinary one with that layout. Validates: recovery to post-migration state, exactly-once output. """ @@ -398,7 +407,10 @@ def test_migration_ft_external_kill_during_migration(tmp_path: Path): timing is not synchronized with client message processing. Migration starts at second 30 (inside client). - External kill fires 45 seconds after client starts (~second 45). + External kill fires 45 seconds after the client process starts, which + with the client's setup lands a few seconds into the migration (Phase A): + there is no checkpoint yet, so recovery reverts to the old layout and the + migration is retried after the first post-recovery snapshot. """ paths = _resolve_paths() results_dir = _make_results_dir(tmp_path) diff --git a/tests/unit/coordinator/service_helpers.py b/tests/unit/coordinator/service_helpers.py new file mode 100644 index 00000000..2b431fa2 --- /dev/null +++ b/tests/unit/coordinator/service_helpers.py @@ -0,0 +1,45 @@ +"""A CoordinatorService with its networking and recovery steps stubbed out, so +tests can drive the real migration and recovery handlers.""" + +import asyncio +from collections import defaultdict +from typing import TYPE_CHECKING +from unittest.mock import AsyncMock, MagicMock + +from styx.common.tcp_networking import NetworkingManager + +from coordinator.coordinator_service import CoordinatorService +from coordinator.migration_metadata import MigrationMetadata + +if TYPE_CHECKING: + from coordinator.coordinator_metadata import Coordinator + + +def stub_service(coordinator: Coordinator) -> CoordinatorService: + svc = CoordinatorService.__new__(CoordinatorService) + svc.coordinator = coordinator + svc.migration_in_progress = False + svc.migration_metadata = MigrationMetadata(2) + svc.networking_locks = defaultdict(asyncio.Lock) + svc._migration_checkpoint_task = None + svc._deferred_migration_graph = None + svc.unreachable_workers = [] + svc.networking = MagicMock(close_worker_connections=AsyncMock()) + svc.protocol_networking = MagicMock( + close_worker_connections=AsyncMock(), + close_all_connections=AsyncMock(), + decode_message=NetworkingManager.decode_message, + ) + svc.aio_task_scheduler = MagicMock(close=AsyncMock()) + svc._reset_after_recovery = AsyncMock() + svc._read_migration_blob_from_s3 = MagicMock(return_value=None) + coordinator.start_recovery_process = AsyncMock() + coordinator.wait_cluster_healthy = AsyncMock() + coordinator.notify_cluster_healthy = AsyncMock() + coordinator.revert_worker_pool_to_submitted_graph = MagicMock() + return svc + + +async def recover(svc: CoordinatorService) -> None: + """Run the real recovery for one dead worker.""" + await svc._perform_recovery({MagicMock(name="dead_worker")}) diff --git a/tests/unit/coordinator/test_migration_checkpoint_lifetime.py b/tests/unit/coordinator/test_migration_checkpoint_lifetime.py index 1b8e4e5e..8dd2001b 100644 --- a/tests/unit/coordinator/test_migration_checkpoint_lifetime.py +++ b/tests/unit/coordinator/test_migration_checkpoint_lifetime.py @@ -6,8 +6,6 @@ must keep tracking that resumed transfer. """ -import asyncio -from collections import defaultdict from unittest.mock import AsyncMock, MagicMock from styx.common.local_state_backends import LocalStateBackend @@ -18,8 +16,8 @@ from styx.common.tcp_networking import NetworkingManager from coordinator.coordinator_metadata import Coordinator -from coordinator.coordinator_service import CoordinatorService from coordinator.migration_metadata import MigrationMetadata +from tests.unit.coordinator.service_helpers import stub_service def _graph(n_partitions: int) -> StateflowGraph: @@ -110,29 +108,6 @@ async def test_transfer_done_epoch_is_the_latest_reported(self): assert m.transfer_done_epoch == -1 -def _service(coordinator: Coordinator) -> CoordinatorService: - svc = CoordinatorService.__new__(CoordinatorService) - svc.coordinator = coordinator - svc.migration_in_progress = False - svc.migration_metadata = MigrationMetadata(2) - svc.networking_locks = defaultdict(asyncio.Lock) - svc._migration_checkpoint_task = None - svc._deferred_migration_graph = None - svc.networking = MagicMock(close_worker_connections=AsyncMock()) - svc.protocol_networking = MagicMock( - close_worker_connections=AsyncMock(), - close_all_connections=AsyncMock(), - decode_message=NetworkingManager.decode_message, - ) - svc.aio_task_scheduler = MagicMock(close=AsyncMock()) - svc._reset_after_recovery = AsyncMock() - coordinator.start_recovery_process = AsyncMock() - coordinator.wait_cluster_healthy = AsyncMock() - coordinator.notify_cluster_healthy = AsyncMock() - coordinator.revert_worker_pool_to_submitted_graph = MagicMock() - return svc - - def _migration_done_msg(epoch: int) -> bytes: return NetworkingManager.encode_message( msg=(epoch,), @@ -144,7 +119,7 @@ def _migration_done_msg(epoch: int) -> bytes: class TestMigrationDoneHandler: async def test_all_workers_done_ends_checkpoint_after_latest_epoch(self): c = _coordinator_with_checkpoint() - svc = _service(c) + svc = stub_service(c) svc.migration_in_progress = True await svc._handle_migration_done(_migration_done_msg(41)) @@ -164,7 +139,7 @@ class TestRecoveryFromCheckpoint: async def test_resume_after_migration_done_keeps_tracking_the_transfer(self): c = Coordinator(MagicMock(), MagicMock()) c.worker_snapshot_ids = {0: 7, 1: 7} - svc = _service(c) + svc = stub_service(c) # MigrationDone already processed; the latest snapshot still has the checkpoint. svc._read_migration_blob_from_s3 = MagicMock(return_value=_blob(_graph(8))) state_when_workers_released = {} @@ -186,7 +161,7 @@ async def test_newer_migration_without_checkpoint_is_retried_after_resume(self): c.worker_snapshot_ids = {0: 7, 1: 7} newer = _graph(16) c._pending_graph = newer - svc = _service(c) + svc = stub_service(c) svc.migration_in_progress = True svc._read_migration_blob_from_s3 = MagicMock(return_value=_blob(_graph(8))) @@ -201,7 +176,7 @@ async def test_crash_before_checkpoint_reverts_and_retries(self): c.worker_snapshot_ids = {0: 7, 1: 7} pending = _graph(8) c._pending_graph = pending - svc = _service(c) + svc = stub_service(c) svc.migration_in_progress = True svc._read_migration_blob_from_s3 = MagicMock(return_value=None) @@ -215,7 +190,7 @@ async def test_crash_before_checkpoint_reverts_and_retries(self): async def test_ordinary_recovery_stays_ordinary(self): c = Coordinator(MagicMock(), MagicMock()) c.worker_snapshot_ids = {0: 7, 1: 7} - svc = _service(c) + svc = stub_service(c) svc._read_migration_blob_from_s3 = MagicMock(return_value=None) await svc._perform_recovery({MagicMock(name="dead_worker")}) @@ -228,7 +203,7 @@ async def test_ordinary_recovery_stays_ordinary(self): class TestUnreachableWorkers: def _service_with_workers(self, *results): c = Coordinator(MagicMock(), MagicMock()) - svc = _service(c) + svc = stub_service(c) svc.unreachable_workers = [] workers = [MagicMock(worker_id=i, worker_ip=f"w{i}", protocol_port=6000 + i) for i in range(len(results))] c.worker_pool = MagicMock(get_participating_workers=MagicMock(return_value=workers)) diff --git a/tests/unit/coordinator/test_migration_fault_tolerance.py b/tests/unit/coordinator/test_migration_fault_tolerance.py index 79f89613..30bb5905 100644 --- a/tests/unit/coordinator/test_migration_fault_tolerance.py +++ b/tests/unit/coordinator/test_migration_fault_tolerance.py @@ -17,6 +17,7 @@ from styx.common.stateflow_graph import StateflowGraph from coordinator.coordinator_metadata import Coordinator +from tests.unit.coordinator.service_helpers import recover, stub_service def _mock_coordinator(): @@ -221,66 +222,71 @@ async def test_finalize_happens_before_s3_write(self): class TestRecoveryConditionDetection: - """Tests that recovery correctly detects both mid-migration crashes AND - post-migration-pre-snapshot crashes. + """Recovery must notice a migration that had not reached its checkpoint, + including one whose MigrationDone arrived but whose post-migration snapshot + had not completed, and revert to the old layout and retry it. - Bug caught: if migration completed but no post-migration snapshot was taken, - recovery didn't detect this as migration-related, used the NEW graph with - the OLD snapshot, and lost half the data. + These drive the real CoordinatorService._perform_recovery. """ - def test_mid_migration_detected(self): - """migration_in_progress=True with _pending_graph set → needs revert.""" + async def test_mid_migration_reverts_and_retries(self): c = _mock_coordinator() - c._pending_graph = MagicMock() - migration_in_progress = True + pending = _make_graph(n_partitions=8) + c._pending_graph = pending + svc = stub_service(c) + svc.migration_in_progress = True - was_migrating = (migration_in_progress and c._pending_graph is not None) or c.post_migration_snapshot_pending - assert was_migrating is True + await recover(svc) - def test_post_migration_pre_snapshot_detected(self): - """migration_in_progress=False but post_migration_snapshot_pending=True - → needs revert. This is the bug scenario.""" + c.revert_worker_pool_to_submitted_graph.assert_called_once() + assert svc._deferred_migration_graph is pending + assert svc.migration_in_progress is False + + async def test_post_migration_pre_snapshot_reverts(self): c = _mock_coordinator() - c._pending_graph = MagicMock() + c._pending_graph = _make_graph(n_partitions=8) c.post_migration_snapshot_pending = True - migration_in_progress = False + svc = stub_service(c) + + await recover(svc) - was_migrating = (migration_in_progress and c._pending_graph is not None) or c.post_migration_snapshot_pending - assert was_migrating is True + c.revert_worker_pool_to_submitted_graph.assert_called_once() - def test_no_migration_not_detected(self): - """Normal state — no migration, no pending snapshot.""" + async def test_no_migration_is_an_ordinary_recovery(self): c = _mock_coordinator() - c._pending_graph = None - c.post_migration_snapshot_pending = False - migration_in_progress = False + svc = stub_service(c) - was_migrating = (migration_in_progress and c._pending_graph is not None) or c.post_migration_snapshot_pending - assert was_migrating is False + await recover(svc) - def test_completed_migration_with_snapshot_not_detected(self): - """Migration completed AND post-migration snapshot taken → no revert needed.""" + c.revert_worker_pool_to_submitted_graph.assert_not_called() + assert svc._deferred_migration_graph is None + + async def test_completed_migration_is_an_ordinary_recovery(self): c = _mock_coordinator() - c._pending_graph = None # finalize_graph_update() cleared it - c.post_migration_snapshot_pending = False - migration_in_progress = False + c.submitted_graph = _make_graph(n_partitions=8) # finalized + svc = stub_service(c) - was_migrating = (migration_in_progress and c._pending_graph is not None) or c.post_migration_snapshot_pending - assert was_migrating is False + await recover(svc) - def test_saved_pending_graph_preserved_for_retry(self): - """When was_migrating, the pending graph should be saved for auto-retry.""" + c.revert_worker_pool_to_submitted_graph.assert_not_called() + assert c.submitted_graph.nodes["ycsb"].n_partitions == 8 + + async def test_recovery_clears_migration_state(self): c = _mock_coordinator() - new_graph = MagicMock(name="new_graph") - c._pending_graph = new_graph + c._pending_graph = _make_graph(n_partitions=8) + c.pre_migration_snapshot_pending = True c.post_migration_snapshot_pending = True - migration_in_progress = False + c._migration_checkpoint_baseline_snap_id = 3 + svc = stub_service(c) + svc.migration_in_progress = True - was_migrating = (migration_in_progress and c._pending_graph is not None) or c.post_migration_snapshot_pending - saved_pending_graph = c._pending_graph if was_migrating else None + await recover(svc) - assert saved_pending_graph is new_graph + assert c.pre_migration_snapshot_pending is False + assert c.post_migration_snapshot_pending is False + assert c._pending_graph is None + assert c._migration_checkpoint_baseline_snap_id == -1 + assert c._migration_checkpoint_blob is None # --------------------------------------------------------------------------- @@ -330,11 +336,6 @@ def test_scenario_crash_after_migration_before_snapshot(self): # CRASH happens here — before post-migration snapshot - # Recovery check - migration_in_progress = False # already cleared by _handle_migration_done - was_migrating = (migration_in_progress and c._pending_graph is not None) or c.post_migration_snapshot_pending - assert was_migrating is True - # submitted_graph is still the OLD graph assert c.submitted_graph is old_graph assert c.submitted_graph.nodes["ycsb"].n_partitions == 4 @@ -382,10 +383,7 @@ async def test_scenario_crash_after_migration_after_snapshot(self): # CRASH happens here — after post-migration snapshot - # Recovery check: no migration recovery needed - migration_in_progress = False - was_migrating = (migration_in_progress and c._pending_graph is not None) or c.post_migration_snapshot_pending - assert was_migrating is False + # (Recovery from here is an ordinary one: TestRecoveryConditionDetection.) def test_scenario_crash_during_migration_phase_a(self): """Scenario: crash during Phase A (migration in progress, @@ -400,13 +398,9 @@ def test_scenario_crash_during_migration_phase_a(self): # Migration starts c._pending_graph = new_graph c.pre_migration_snapshot_pending = True - migration_in_progress = True # CRASH during Phase A — no pre-migration snapshot completed - was_migrating = (migration_in_progress and c._pending_graph is not None) or c.post_migration_snapshot_pending - assert was_migrating is True - # Old graph preserved assert c.submitted_graph is old_graph saved = c._pending_graph @@ -449,30 +443,6 @@ async def test_scenario_normal_migration_no_crash(self): c.register_snapshot(0, 3, {}, {}, 30, 300, pool) assert c.submitted_graph is new_graph # still new graph - def test_scenario_state_cleanup_on_recovery(self): - """Verify all migration state is properly cleared during recovery.""" - c = self._setup_coordinator(n_partitions=4) - new_graph = _make_graph("test-app", n_partitions=8, max_parallelism=8) - - # Set up mid-migration state - c._pending_graph = new_graph - c.pre_migration_snapshot_pending = True - c.pre_migration_snapshot_id = 1 - c.post_migration_snapshot_pending = True - - # Simulate recovery clearing state (mirrors _perform_recovery step 1c) - saved_pending_graph = c._pending_graph - c.pre_migration_snapshot_pending = False - c.post_migration_snapshot_pending = False - c._pending_graph = None - - # All state cleared - assert c.pre_migration_snapshot_pending is False - assert c.post_migration_snapshot_pending is False - assert c._pending_graph is None - # But we saved the graph for retry - assert saved_pending_graph is new_graph - # --------------------------------------------------------------------------- # Worker pool revert on recovery during migration diff --git a/tests/unit/coordinator/test_snapshot_compactor_coverage.py b/tests/unit/coordinator/test_snapshot_compactor_coverage.py index 929343a6..42517c60 100644 --- a/tests/unit/coordinator/test_snapshot_compactor_coverage.py +++ b/tests/unit/coordinator/test_snapshot_compactor_coverage.py @@ -85,9 +85,10 @@ def get_bytes_side_effect(Bucket=None, Key=None): mock_client.get_object.side_effect = get_bytes_side_effect compact_deltas(mock_client, files, max_snap_id=5) - # Should write compacted data + # The merged state goes into the partition's latest file; older ones go. put_calls = mock_client.put_object.call_args_list - assert any("data/users/0/0.bin" in str(c) for c in put_calls) + assert any("data/users/0/1.bin" in str(c) for c in put_calls) + mock_client.delete_object.assert_any_call(Bucket="styx-snapshots", Key="data/users/0/0.bin") def test_compact_sequencer_snapshots(self): mock_client = MagicMock() @@ -107,6 +108,63 @@ def get_bytes_side_effect(Bucket=None, Key=None): put_calls = mock_client.put_object.call_args_list assert any("sequencer/0.bin" in str(c) for c in put_calls) + def _s3(self, files: dict[str, dict]): + """In-memory S3: get/put/delete on a dict of key -> deserialized payload.""" + from styx.common.serialization import zstd_msgpack_deserialization + + store = {k: zstd_msgpack_serialization(v) for k, v in files.items()} + client = MagicMock() + + def get(Bucket=None, Key=None): + body = MagicMock() + body.read.return_value = store[Key] + return {"Body": body} + + def put(Bucket=None, Key=None, Body=None): + store[Key] = Body + + def delete(Bucket=None, Key=None): + del store[Key] + + client.get_object.side_effect = get + client.put_object.side_effect = put + client.delete_object.side_effect = delete + return client, store, zstd_msgpack_deserialization + + def test_merges_in_snapshot_id_order(self): + # S3 lists keys lexicographically: 10.bin before 2.bin. + client, store, load = self._s3( + { + "data/users/0/10.bin": {"k": "newest"}, + "data/users/0/2.bin": {"k": "old", "j": "v"}, + }, + ) + compact_deltas(client, list(store), max_snap_id=10) + assert set(store) == {"data/users/0/10.bin"} + assert load(store["data/users/0/10.bin"]) == {"k": "newest", "j": "v"} + + def test_partition_whose_first_file_is_not_zero(self): + # A partition created by a migration starts at a later snapshot. + client, store, load = self._s3( + { + "data/users/4/5.bin": {"k": "v5"}, + "data/users/4/6.bin": {"k": "v6"}, + }, + ) + compact_deltas(client, list(store), max_snap_id=6) + assert set(store) == {"data/users/4/6.bin"} + assert load(store["data/users/4/6.bin"]) == {"k": "v6"} + + def test_tombstones_are_dropped_once_older_files_are_gone(self): + client, store, load = self._s3( + { + "data/users/0/0.bin": {"moved": "v", "stays": "v"}, + "data/users/0/1.bin": {"moved": None}, + }, + ) + compact_deltas(client, list(store), max_snap_id=1) + assert load(store["data/users/0/1.bin"]) == {"stays": "v"} + def test_compact_empty_no_crash(self): mock_client = MagicMock() compact_deltas(mock_client, [], max_snap_id=5) diff --git a/worker/egress/styx_kafka_batch_egress.py b/worker/egress/styx_kafka_batch_egress.py index c21f696c..a92c8c09 100644 --- a/worker/egress/styx_kafka_batch_egress.py +++ b/worker/egress/styx_kafka_batch_egress.py @@ -152,7 +152,7 @@ async def send( else: self.messages_sent_before_recovery.discard(key) self._dedup_suppressed += 1 - logging.warning( + logging.debug( f"[DEDUP] W{getattr(self, 'worker_id', '?')} SUPPRESSED key={key!r} " f"target_partition={operator_name}--OUT/{partition}", ) @@ -194,7 +194,7 @@ async def send_immediate( else: self.messages_sent_before_recovery.discard(key) self._dedup_suppressed += 1 - logging.warning( + logging.debug( f"[DEDUP] W{getattr(self, 'worker_id', '?')} SUPPRESSED (immediate) key={key!r} " f"target_partition={operator_name}--OUT/{partition}", ) @@ -233,7 +233,7 @@ async def get_messages_sent_before_recovery(self) -> None: for tp in all_topic_partitions: offset = self.dedup_output_offsets[(tp.topic[:-5], tp.partition)] + 1 kafka_output_consumer.seek(tp, offset) - logging.warning(f"[DEDUP] Seeking {tp} to offset {offset}") + logging.debug(f"[DEDUP] Seeking {tp} to offset {offset}") except (UnknownTopicOrPartitionError, KafkaConnectionError): # fmt: skip await asyncio.sleep(1) logging.warning( @@ -243,14 +243,14 @@ async def get_messages_sent_before_recovery(self) -> None: break try: current_offsets = await kafka_output_consumer.end_offsets(all_topic_partitions) - logging.warning( + logging.debug( f"[DEDUP] End offsets: {current_offsets}", ) all_partitions_done: dict[TopicPartition, bool] = { tp: current_offsets[tp] == self.dedup_output_offsets[(tp.topic[:-5], tp.partition)] + 1 for tp in all_topic_partitions } - logging.warning(f"[DEDUP] Initial partitions_done: {all_partitions_done}") + logging.debug(f"[DEDUP] Initial partitions_done: {all_partitions_done}") per_partition_count: dict[TopicPartition, int] = dict.fromkeys(all_topic_partitions, 0) while not all(all_partitions_done.values()): result = await kafka_output_consumer.getmany(timeout_ms=EPOCH_INTERVAL_MS) diff --git a/worker/transactional_protocols/aria.py b/worker/transactional_protocols/aria.py index 63e40786..0240c259 100644 --- a/worker/transactional_protocols/aria.py +++ b/worker/transactional_protocols/aria.py @@ -1224,13 +1224,5 @@ async def sync_workers( msg_type=msg_type, serializer=serializer, ) - # Detect stuck barriers during migration debugging - try: - await asyncio.wait_for(self.sync_workers_event[msg_type].wait(), timeout=10.0) - except TimeoutError: - logging.warning( - f"Worker {self.id} | BARRIER STUCK for 10s on {msg_type.name} @epoch {self.sequencer.epoch_counter}", - ) - # Continue waiting (no timeout) - await self.sync_workers_event[msg_type].wait() + await self.sync_workers_event[msg_type].wait() self.sync_workers_event[msg_type].clear() diff --git a/worker/worker_service.py b/worker/worker_service.py index ff226aeb..827d12f3 100644 --- a/worker/worker_service.py +++ b/worker/worker_service.py @@ -832,13 +832,13 @@ async def _handle_init_recovery(self, data: bytes, _: MessageType) -> None: f"[RECOVERY] W{self.id} snapshot_id={snapshot_id} epoch={epoch} t_counter={t_counter}" f" migration_blob={'present' if migration_blob else 'none'}", ) - logging.warning( + logging.debug( f"[RECOVERY] W{self.id} registered_operators={list(self.registered_operators.keys())}", ) - logging.warning( + logging.debug( f"[RECOVERY] W{self.id} raw tp_offsets from snapshot: {tp_offsets}", ) - logging.warning( + logging.debug( f"[RECOVERY] W{self.id} raw tp_out_offsets from snapshot: {tp_out_offsets}", ) @@ -850,13 +850,13 @@ async def _handle_init_recovery(self, data: bytes, _: MessageType) -> None: tp_offsets = {k: v for k, v in tp_offsets.items() if k in self.registered_operators} tp_out_offsets = {k: v for k, v in tp_out_offsets.items() if k in self.registered_operators} - logging.warning( + logging.debug( f"[RECOVERY] W{self.id} filtered tp_offsets: {tp_offsets}", ) - logging.warning( + logging.debug( f"[RECOVERY] W{self.id} filtered tp_out_offsets: {tp_out_offsets}", ) - logging.warning( + logging.debug( f"[RECOVERY] W{self.id} dedup_output_offsets (all partitions): {all_dedup_out_offsets}", ) From 761163cf642b447c544ae31142a9d8849339748e Mon Sep 17 00:00:00 2001 From: Kyriakos Psarakis Date: Sat, 24 Oct 2026 21:32:44 +0200 Subject: [PATCH 13/20] Start partitions without a recovered offset from the beginning Recovery kept only the input offsets present in the snapshot. When the snapshot did not cover every assigned partition (a shadow partition, after recovering from the snapshot written at data load), that partition's consumer was never seeked and started at the latest offset, skipping the requests clients wrote to it before and during the failure. Seen in CI as 194 missed responses in the external-kill migration test. Recovery now defaults missing offsets to -1, and the ingress seeks every assigned partition, from the beginning when it has no offset. Co-Authored-By: Claude Opus 5.5 --- tests/unit/worker/test_kafka_ingress.py | 17 +++++++++++++++++ worker/ingress/styx_kafka_ingress.py | 6 ++++++ worker/worker_service.py | 8 ++++++-- 3 files changed, 29 insertions(+), 2 deletions(-) diff --git a/tests/unit/worker/test_kafka_ingress.py b/tests/unit/worker/test_kafka_ingress.py index e77a0db1..3b16e07a 100644 --- a/tests/unit/worker/test_kafka_ingress.py +++ b/tests/unit/worker/test_kafka_ingress.py @@ -465,3 +465,20 @@ async def test_stop_closes_consumer(self): consumer.stop.assert_awaited_once() assert ingress.take_consumer() is None + + @pytest.mark.asyncio + async def test_partition_without_offset_starts_from_the_beginning(self): + # E.g. a shadow partition after recovering from the data-load snapshot: + # left unseeked, the consumer would start at the latest offset. + ingress, *_ = _make_ingress() + consumer = MagicMock() + consumer.getmany = AsyncMock(side_effect=asyncio.CancelledError()) + consumer.stop = AsyncMock() + ingress.adopt_consumer(consumer) + tps = [TopicPartition("users", 0), TopicPartition("users", 4)] + + with pytest.raises(asyncio.CancelledError): + await ingress.start_kafka_consumer(tps, {("users", 0): 41}) + + consumer.seek.assert_any_call(TopicPartition("users", 0), 42) + consumer.seek.assert_any_call(TopicPartition("users", 4), 0) diff --git a/worker/ingress/styx_kafka_ingress.py b/worker/ingress/styx_kafka_ingress.py index f0168a2b..085b2b1a 100644 --- a/worker/ingress/styx_kafka_ingress.py +++ b/worker/ingress/styx_kafka_ingress.py @@ -157,6 +157,12 @@ async def start_kafka_consumer( topic_partitions: list[TopicPartition], topic_partition_offsets: dict[OperatorPartition, int], ) -> None: + # Every assigned partition gets a position: one without an offset starts from + # the beginning, never at the consumer's default (the latest offset). + topic_partition_offsets = { + (tp.topic, tp.partition): topic_partition_offsets.get((tp.topic, tp.partition), -1) + for tp in topic_partitions + } if self.kafka_consumer is not None: # Adopted from the previous protocol: already connected. Seeking every # assigned partition drops any records it had prefetched. diff --git a/worker/worker_service.py b/worker/worker_service.py index 827d12f3..a52436a6 100644 --- a/worker/worker_service.py +++ b/worker/worker_service.py @@ -847,8 +847,12 @@ async def _handle_init_recovery(self, data: bytes, _: MessageType) -> None: # that cross-worker dedup works: a request originally processed on # worker A may be replayed on worker B after recovery. all_dedup_out_offsets = dict(tp_out_offsets) - tp_offsets = {k: v for k, v in tp_offsets.items() if k in self.registered_operators} - tp_out_offsets = {k: v for k, v in tp_out_offsets.items() if k in self.registered_operators} + # A partition the snapshot has no offset for (e.g. a shadow partition and the + # snapshot written at data load) had nothing consumed by then: start it from + # the beginning. Left out, its consumer would start at the latest offset and + # skip what clients wrote to it before and during the failure. + tp_offsets = {k: tp_offsets.get(k, -1) for k in self.registered_operators} + tp_out_offsets = {k: tp_out_offsets.get(k, -1) for k in self.registered_operators} logging.debug( f"[RECOVERY] W{self.id} filtered tp_offsets: {tp_offsets}", From e8c0564545b79f3da7e39ec07cf48c97799248d1 Mon Sep 17 00:00:00 2001 From: Kyriakos Psarakis Date: Sat, 24 Oct 2026 22:49:10 +0200 Subject: [PATCH 14/20] Keep unanswered requests and the FT E2E results in the uploaded logs calculate_metrics writes missed_requests.csv, and the FT tests copy their results directory into logs/, which CI uploads, so a failure that only reproduces in CI can be traced to the requests involved. Co-Authored-By: Claude Opus 5.5 --- demo/demo-migration-ycsb/calculate_metrics.py | 1 + tests/e2e/test_e2e_migration_ft_ycsb.py | 8 ++++++++ 2 files changed, 9 insertions(+) diff --git a/demo/demo-migration-ycsb/calculate_metrics.py b/demo/demo-migration-ycsb/calculate_metrics.py index 12ab3dd2..15c0a4b7 100755 --- a/demo/demo-migration-ycsb/calculate_metrics.py +++ b/demo/demo-migration-ycsb/calculate_metrics.py @@ -48,6 +48,7 @@ def main( ) missed_requests = joined[joined["timestamp_output"].isna()] print(f"missed_requests: {missed_requests}") + missed_requests[["request_id", "timestamp_client", "op"]].to_csv(f"{save_dir}/missed_requests.csv", index=False) missed = len(missed_requests) print( f"Missed {missed} messages after the join for which no response " diff --git a/tests/e2e/test_e2e_migration_ft_ycsb.py b/tests/e2e/test_e2e_migration_ft_ycsb.py index 7a079ae5..76d91694 100644 --- a/tests/e2e/test_e2e_migration_ft_ycsb.py +++ b/tests/e2e/test_e2e_migration_ft_ycsb.py @@ -26,6 +26,7 @@ import json import logging from pathlib import Path +import shutil import pytest @@ -207,7 +208,14 @@ def _run_client( raise AssertionError(f"client.py failed (rc={rc}). Output:\n{out}") +def _keep_results(results_dir: Path) -> None: + """Copy the run's results next to the cluster logs, which CI uploads.""" + target = _resolve_paths().repo_root / "logs" / results_dir.parent.name + shutil.copytree(results_dir, target, dirs_exist_ok=True) + + def _assert_artifacts_and_metrics(results_dir: Path, client: _ClientParams) -> None: + _keep_results(results_dir) client_csv = results_dir / "client_requests.csv" output_csv = results_dir / "output.csv" log.info("Checking artifacts: %s and %s", client_csv, output_csv) From c7f86972b0f5dcfa0cd4c8384317ffddbbad30ca Mon Sep 17 00:00:00 2001 From: Kyriakos Psarakis Date: Sun, 25 Oct 2026 01:07:47 +0200 Subject: [PATCH 15/20] Run mis-routed requests as proxied calls from the consuming worker A request consumed for a key that migrated to a peer's partition was forwarded into the peer's epoch (WrongPartitionRequest). The ingress partition's offset was then advanced by two workers, each taking the max, and over a 16-socket pool the forwarded requests could run out of order: a snapshot or Phase B stop could record an offset past requests one of them had not run yet, and recovery or the rebuilt consumer skipped them. The request now stays in the consuming worker's epoch and runs on the owner as a proxied call (RunFunProxy). The owner runs its chain against itself and returns the whole outcome in one ProxyResult (response, abort, participants, fallback read/write-set change), which the root awaits before it can commit. Responses go to the request's ingress partition, which the root owns. Both messages carry the incarnation, since a replay reuses t_ids. Co-Authored-By: Claude Opus 5.5 --- styx-package/styx/common/message_types.py | 3 +- tests/unit/styx_package/test_message_types.py | 11 +- tests/unit/worker/test_kafka_ingress.py | 26 ++- tests/unit/worker/test_proxied_requests.py | 204 ++++++++++++++++++ worker/ingress/styx_kafka_ingress.py | 47 +--- worker/transactional_protocols/aria.py | 149 +++++++++++-- 6 files changed, 368 insertions(+), 72 deletions(-) create mode 100644 tests/unit/worker/test_proxied_requests.py diff --git a/styx-package/styx/common/message_types.py b/styx-package/styx/common/message_types.py index fb952567..8ca80af1 100644 --- a/styx-package/styx/common/message_types.py +++ b/styx-package/styx/common/message_types.py @@ -27,7 +27,6 @@ class MessageType(IntEnum): SyncCleanup = 22 RemoteWantsToProceed = 23 ChainAbort = 24 - WrongPartitionRequest = 26 InitMigration = 27 MigrationRepartitioningDone = 28 MigrationDone = 29 @@ -48,3 +47,5 @@ class MessageType(IntEnum): SnapMigrationReassign = 44 AsyncMigrationAck = 45 RecoveryMigrationHashes = 46 + RunFunProxy = 47 + ProxyResult = 48 diff --git a/tests/unit/styx_package/test_message_types.py b/tests/unit/styx_package/test_message_types.py index 7e56b6a1..9d5e97f2 100644 --- a/tests/unit/styx_package/test_message_types.py +++ b/tests/unit/styx_package/test_message_types.py @@ -10,17 +10,18 @@ def test_is_int_enum(self): assert issubclass(MessageType, IntEnum) def test_total_count(self): - assert len(MessageType) == 46 + assert len(MessageType) == 47 def test_all_values_unique(self): values = [m.value for m in MessageType] assert len(values) == len(set(values)) def test_values_in_expected_range(self): - # 25 (AckCache) was removed when the fallback cache was dropped; the - # remaining IDs are preserved to keep wire-format stability. + # 25 (AckCache) was removed when the fallback cache was dropped, and 26 + # (WrongPartitionRequest) when mis-routed requests became proxied calls; + # the remaining IDs are preserved to keep wire-format stability. values = sorted(m.value for m in MessageType) - assert values == [v for v in range(47) if v != 25] + assert values == [v for v in range(49) if v not in {25, 26}] def test_known_values(self): assert MessageType.RunFunRemote == 0 @@ -33,6 +34,8 @@ def test_known_values(self): assert MessageType.SnapMigrationReassign == 44 assert MessageType.AsyncMigrationAck == 45 assert MessageType.RecoveryMigrationHashes == 46 + assert MessageType.RunFunProxy == 47 + assert MessageType.ProxyResult == 48 def test_int_comparison(self): assert MessageType.ClientMsg == 16 diff --git a/tests/unit/worker/test_kafka_ingress.py b/tests/unit/worker/test_kafka_ingress.py index 3b16e07a..b396841b 100644 --- a/tests/unit/worker/test_kafka_ingress.py +++ b/tests/unit/worker/test_kafka_ingress.py @@ -68,7 +68,6 @@ def test_defaults(self): assert ingress.kafka_url == "localhost:9092" assert ingress.epoch_interval_ms == 1 assert ingress.sequence_max_size == 1000 - assert ingress.send_message_tasks == [] assert not ingress.started.is_set() assert ingress.kafka_consumer is None @@ -126,7 +125,9 @@ def test_key_not_in_state_wrong_partition_local(self): assert payload.partition == 1 # rerouted to correct partition def test_key_not_in_state_wrong_partition_remote(self): - """Key migrated to different partition on a remote worker.""" + """Key migrated to a partition on a remote worker: the request is still + sequenced here (it runs as a proxied call), so this worker alone advances + the ingress partition's offset.""" ingress, networking, sequencer, state = _make_ingress() state.exists.return_value = False networking.decode_message.return_value = ("users", "key1", "my_func", (1,), 0) @@ -134,10 +135,14 @@ def test_key_not_in_state_wrong_partition_remote(self): networking.in_the_same_network.return_value = False ingress.registered_operators[("users", 0)].dns = {"users": {1: ("remote_host", 5001, 6001)}} - msg = _make_msg() + msg = _make_msg(offset=7) ingress.handle_message_from_kafka(msg) - sequencer.sequence.assert_not_called() - assert len(ingress.send_message_tasks) == 1 + sequencer.sequence.assert_called_once() + payload = sequencer.sequence.call_args[0][0] + assert payload.partition == 1 + assert payload.kafka_ingress_partition == 0 + assert payload.kafka_offset == 7 + networking.send_message.assert_not_called() async def test_invalid_message_type(self): """Non-ClientMsg types should be logged as errors.""" @@ -229,8 +234,8 @@ def test_message_on_shadow_partition_redirected_locally(self): assert payload.partition == 0 # redirected to active partition def test_message_on_shadow_partition_redirected_remotely(self): - """Message on shadow partition redirected to remote worker via - WrongPartitionRequest.""" + """Message on shadow partition whose key is on a remote worker: sequenced + here with the active partition, to run there as a proxied call.""" ingress, networking, sequencer, _state = _make_ingress_with_shadow_partitions() networking.decode_message.return_value = ("ycsb", "test_key", "read", (), 5) ingress.registered_operators[("ycsb", 5)].which_partition = MagicMock(return_value=1) @@ -239,8 +244,11 @@ def test_message_on_shadow_partition_redirected_remotely(self): msg = _make_msg(partition=5) ingress.handle_message_from_kafka(msg) - sequencer.sequence.assert_not_called() - assert len(ingress.send_message_tasks) == 1 + sequencer.sequence.assert_called_once() + payload = sequencer.sequence.call_args[0][0] + assert payload.partition == 1 + assert payload.kafka_ingress_partition == 5 + networking.send_message.assert_not_called() def test_message_on_active_partition_not_redirected(self): """Message on active partition 0 with existing key is processed normally.""" diff --git a/tests/unit/worker/test_proxied_requests.py b/tests/unit/worker/test_proxied_requests.py new file mode 100644 index 00000000..e171289d --- /dev/null +++ b/tests/unit/worker/test_proxied_requests.py @@ -0,0 +1,204 @@ +"""Requests consumed by one worker for a key that migrated to a peer's partition. + +They stay in the consuming worker's epoch and run on the peer as a proxied call, +so an ingress partition's offset is only ever advanced by the worker consuming +it. The peer reports the whole outcome (response, abort, participants, fallback +read/write-set change) in one ProxyResult, which the root awaits before it can +commit. +""" + +import asyncio +from types import SimpleNamespace + +import pytest +from styx.common.message_types import MessageType +from styx.common.operator import Operator +from styx.common.run_func_payload import RunFuncPayload +from styx.common.serialization import Serializer +from styx.common.tcp_networking import NetworkingManager +from styx.common.util.aio_task_scheduler import AIOTaskScheduler + +from worker.operator_state.aria.in_memory_state import InMemoryOperatorState +from worker.transactional_protocols.aria import AriaProtocol + +OP = "op" +W1 = ("w1", 5001, 6001) # consumes ingress partition 0 +W2 = ("w2", 5002, 6002) # owns partition 1, where the key now lives +DNS = {OP: {0: W1, 1: W2}} +PEERS = {1: W1, 2: W2} + + +async def read(ctx): + return ctx.key, ctx.get() + + +async def fail(_ctx): + msg = "no such account" + raise ValueError(msg) + + +class _Link: + """Delivers protocol messages between fake workers, like the TCP readers.""" + + def __init__(self) -> None: + self.workers: dict[tuple[str, int], AriaProtocol] = {} + self.sent: list[MessageType] = [] + + async def deliver(self, host, port, msg, msg_type, serializer) -> bool: + self.sent.append(msg_type) + data = NetworkingManager.encode_message(msg=msg, msg_type=msg_type, serializer=serializer)[8:] + await asyncio.sleep(0) + await self.workers[(host, port)].protocol_handlers_map[msg_type](data) + return True + + +class _Networking(NetworkingManager): + def __init__(self, link: _Link, addr: tuple, worker_id: int) -> None: + super().__init__(addr[2]) + self.host_name = addr[0] + self.set_worker_id(worker_id) + self._link = link + + async def send_message(self, host, port, msg, msg_type, serializer=Serializer.MSGPACK) -> bool: + return await self._link.deliver(host, port, msg, msg_type, serializer) + + +def _worker(link: _Link, worker_id: int, addr: tuple, partitions: set[int]) -> AriaProtocol: + state = InMemoryOperatorState({(OP, p) for p in partitions}) + networking = _Networking(link, addr, worker_id) + registered = {} + for p in partitions: + operator = Operator(OP, n_partitions=2) + operator.register(read) + operator.register(fail) + operator.attach_state_networking(state, networking, DNS, None) + registered[(OP, p)] = operator + + proto = AriaProtocol.__new__(AriaProtocol) + proto.id = worker_id + proto.incarnation = 0 + proto.peers = {w: a for w, a in PEERS.items() if w != worker_id} + proto.dns = DNS + proto.local_state = state + proto.networking = networking + proto.registered_operators = registered + proto.background_functions = AIOTaskScheduler() + proto.waiting_on_transactions = {} + proto.fallback_locking_event_map = {} + proto._proxy_results = {} + proto._proxy_rw_changed = set() + proto.protocol_handlers_map = { + MessageType.RunFunProxy: proto._handle_run_fun_proxy, + MessageType.ProxyResult: proto._handle_proxy_result, + } + link.workers[(addr[0], addr[2])] = proto + return proto + + +def _cluster(): + link = _Link() + root = _worker(link, 1, W1, {0}) + owner = _worker(link, 2, W2, {1}) + owner.local_state.batch_insert({"k": "v"}, OP, 1) + return link, root, owner + + +def _misrouted(function_name: str = "read") -> RunFuncPayload: + # As the ingress sequences it: partition is where the key lives now, the + # ingress partition and offset are where it was consumed. + return RunFuncPayload( + request_id=b"rq", + key="k", + operator_name=OP, + partition=1, + function_name=function_name, + params=(), + kafka_offset=41, + kafka_ingress_partition=0, + ) + + +class TestProxiedRequest: + async def test_runs_on_owner_and_answers_at_the_root(self): + link, root, owner = _cluster() + + assert await root.run_function(7, _misrouted()) is True + + assert root.networking.client_responses[7] == ["k", "v"] + assert link.sent == [MessageType.RunFunProxy, MessageType.ProxyResult] + # The owner holds the transaction's reservation, for conflict detection, + # and is a chain participant, so a fallback run reaches it with Unlock. + assert "k" in owner.local_state.read_sets[(OP, 1)][7] + assert root.networking.chain_participants[7] == [2] + assert 7 not in root._proxy_results + + async def test_logic_abort_reaches_the_root(self): + _link, root, _owner = _cluster() + + assert await root.run_function(7, _misrouted("fail")) is False + + assert root.networking.aborted_events[7] == "no such account" + assert 7 in root.networking.logic_aborts_everywhere + assert 7 not in root.networking.client_responses + + async def test_fallback_read_write_set_change_reaches_the_root(self): + _link, root, owner = _cluster() + owner.local_state.has_fallback_rw_set_changed = lambda t_id: t_id == 7 + + await root.run_function(7, _misrouted(), fallback_mode=True) + + assert 7 in root._proxy_rw_changed + + async def test_response_goes_to_the_ingress_partition(self): + _link, root, _owner = _cluster() + payload = _misrouted() + assert root._response_partition(payload) == 0 + local = RunFuncPayload(request_id=b"r", key="x", operator_name=OP, partition=0, function_name="read", params=()) + assert root._response_partition(local) == 0 + + async def test_participant_call_for_a_missing_partition_is_not_proxied(self): + # Only a root (no ack payload) is proxied; a chain call must target a local partition. + link, root, _owner = _cluster() + payload = _misrouted() + payload.ack_payload = ("w9", 6009, 7, "1.0", []) + with pytest.raises(AttributeError): + await root.run_function(7, payload) + assert link.sent == [] + + +class TestStaleProxyMessages: + async def test_result_from_another_incarnation_is_dropped(self): + _link, root, _owner = _cluster() + pending = asyncio.get_running_loop().create_future() + root._proxy_results[7] = pending + data = NetworkingManager.encode_message( + msg=(7, "stale", None, [2], False, 99), + msg_type=MessageType.ProxyResult, + serializer=Serializer.MSGPACK, + )[8:] + + await root._handle_proxy_result(data) + + assert not pending.done() + + async def test_call_from_another_incarnation_is_not_run(self): + link, _root, owner = _cluster() + data = NetworkingManager.encode_message( + msg=(7, 1, b"rq", OP, "read", "k", 1, False, (), 99), + msg_type=MessageType.RunFunProxy, + serializer=Serializer.MSGPACK, + )[8:] + + await owner._handle_run_fun_proxy(data) + await asyncio.sleep(0) + + assert link.sent == [] + assert 7 not in owner.local_state.read_sets[(OP, 1)] + + +class TestIngressOffsetsStayWithTheConsumer: + def test_proxied_request_advances_the_consumers_ingress_offset(self): + proto = SimpleNamespace(concurrency_aborts_everywhere=set(), topic_partition_offsets={(OP, 0): 40}) + item = SimpleNamespace(t_id=7, payload=_misrouted()) + AriaProtocol._advance_offsets(proto, [item]) + assert proto.topic_partition_offsets == {(OP, 0): 41} diff --git a/worker/ingress/styx_kafka_ingress.py b/worker/ingress/styx_kafka_ingress.py index 085b2b1a..2a9ac908 100644 --- a/worker/ingress/styx_kafka_ingress.py +++ b/worker/ingress/styx_kafka_ingress.py @@ -7,7 +7,6 @@ from styx.common.logging import logging from styx.common.message_types import MessageType from styx.common.run_func_payload import RunFuncPayload -from styx.common.serialization import Serializer from worker.ingress.base_ingress import BaseIngress @@ -40,7 +39,6 @@ def __init__( self.kafka_url = kafka_url self.epoch_interval_ms = epoch_interval_ms self.sequence_max_size = sequence_max_size - self.send_message_tasks: list[typing.Coroutine] = [] self.started: asyncio.Event = asyncio.Event() self.messages_available: asyncio.Event = asyncio.Event() @@ -113,42 +111,12 @@ def handle_message_from_kafka(self, msg: ConsumerRecord) -> None: # Message received in the correct partition, but it was an insert operation (didn't exist in the state) self.sequencer.sequence(run_func_payload) else: - # In flight message during migration, currently the state belongs to another partition - dns = self.registered_operators[(operator_name, msg.partition)].dns - operator_host = dns[operator_name][true_partition][0] - operator_port = dns[operator_name][true_partition][2] - if self.networking.in_the_same_network(operator_host, operator_port): - run_func_payload = RunFuncPayload( - request_id=msg.key, - key=key, - operator_name=operator_name, - partition=true_partition, - function_name=fun_name, - kafka_offset=msg.offset, - params=params, - kafka_ingress_partition=msg.partition, - ) - self.sequencer.sequence(run_func_payload) - else: - payload = ( - msg.key, - operator_name, - fun_name, - key, - true_partition, - msg.partition, - msg.offset, - params, - ) - self.send_message_tasks.append( - self.networking.send_message( - operator_host, - operator_port, - msg=payload, - msg_type=MessageType.WrongPartitionRequest, - serializer=Serializer.MSGPACK, - ), - ) + # In flight message during migration: the key now belongs to another + # partition. It still runs in this worker's epoch, as a remote call + # if that partition is on a peer (AriaProtocol.run_function), so the + # ingress partition's offset only ever advances here, in order. + run_func_payload.partition = true_partition + self.sequencer.sequence(run_func_payload) else: logging.error(f"Invalid message type: {message_type} passed to KAFKA") @@ -230,9 +198,6 @@ async def _consume(self, topic_partition_offsets: dict[OperatorPartition, int]) s_messages = sorted(messages, key=lambda msg: msg.offset) for message in s_messages: self.handle_message_from_kafka(message) - if self.send_message_tasks: - await asyncio.gather(*self.send_message_tasks) - self.send_message_tasks = [] self.messages_available.set() finally: if not self._keep_consumer: diff --git a/worker/transactional_protocols/aria.py b/worker/transactional_protocols/aria.py index 0240c259..00f1fadc 100644 --- a/worker/transactional_protocols/aria.py +++ b/worker/transactional_protocols/aria.py @@ -185,10 +185,17 @@ def __init__( # noqa: PLR0913 self.migrating_state: bool = restart_after_migration self._init_async_migration_state() + # Requests consumed here whose key now lives on a peer run there as proxied + # calls: t_id -> the future their ProxyResult resolves. + self._proxy_results: dict[int, asyncio.Future] = {} + # Proxied transactions whose read/write set changed in their fallback run. + self._proxy_rw_changed: set[int] = set() + self.protocol_handlers_map: dict[MessageType, Callable[[bytes], Awaitable[None]]] = { MessageType.RunFunRemote: self._handle_run_fun_remote, MessageType.RunFunRemoteBatch: self._handle_run_fun_remote_batch, - MessageType.WrongPartitionRequest: self._handle_wrong_partition_request, + MessageType.RunFunProxy: self._handle_run_fun_proxy, + MessageType.ProxyResult: self._handle_proxy_result, MessageType.RunFunRqRsRemote: self._handle_deprecated_rqrs, MessageType.AriaCommit: self._handle_aria_commit, MessageType.AriaFallbackDone: self._handle_sync_event_only, @@ -280,7 +287,10 @@ async def run_function( # logging.info(f"Running function: {payload.function_name} with T_ID {t_id} with params {payload.params}" # f" and ack payload {payload.ack_payload}") - operator_partition = self.registered_operators[(payload.operator_name, payload.partition)] + operator_partition = self.registered_operators.get((payload.operator_name, payload.partition)) + if operator_partition is None and payload.ack_payload is None: + # A request consumed here for a key that migrated to a peer's partition. + return await self._run_proxied(t_id, payload, fallback_mode) success: bool = await operator_partition.run_function( payload.key, @@ -410,23 +420,66 @@ async def _handle_run_fun_remote_batch(self, data: bytes) -> None: else: self.background_functions.create_task(self.run_function(t_id, payload)) - async def _handle_wrong_partition_request(self, data: bytes) -> None: + async def _run_proxied(self, t_id: int, payload: RunFuncPayload, fallback_mode: bool) -> bool: + """Run a request consumed here on the peer whose partition now holds its key, + and take its outcome as this root's own. + + The request stays in this worker's epoch, so its ingress offset advances only + here and in consumption order: every snapshot or migration stop cuts each + ingress partition at one offset. (Forwarding it into the peer's epoch let two + workers advance one partition's offset, each by max, past requests the other + had not run yet.) + """ + host, _, protocol_port = self.dns[payload.operator_name][payload.partition] + result: asyncio.Future = asyncio.get_running_loop().create_future() + self._proxy_results[t_id] = result + await self.networking.send_message( + host, + protocol_port, + msg=( + t_id, + self.id, + payload.request_id, + payload.operator_name, + payload.function_name, + payload.key, + payload.partition, + fallback_mode, + payload.params, + self.incarnation, + ), + msg_type=MessageType.RunFunProxy, + serializer=Serializer.MSGPACK, + ) + response, abort, participants, rw_changed = await result + self.networking.add_chain_participants(t_id, participants) + if abort is not None: + self.networking.abort_chain(t_id, abort) + return False + if rw_changed: + self._proxy_rw_changed.add(t_id) + if response is not None: + self.networking.add_response(t_id, response) + return True + + async def _handle_run_fun_proxy(self, data: bytes) -> None: # Lock-free: no awaits. ( + t_id, + root_id, request_id, operator_name, function_name, key, partition, - kafka_ingress_partition, - kafka_offset, + fallback_mode, params, + incarnation, ) = self.networking.decode_message(data) - - logging.debug( - f"Aria WrongPartitionRequest: {request_id}:{operator_name}:{kafka_ingress_partition}", - ) - + if incarnation != self.incarnation: + # Sent before a recovery; the root replays it with the same t_id. + logging.warning(f"Dropping proxied call {t_id} from W{root_id} of incarnation {incarnation}") + return payload = RunFuncPayload( request_id=request_id, key=key, @@ -434,10 +487,71 @@ async def _handle_wrong_partition_request(self, data: bytes) -> None: partition=partition, function_name=function_name, params=params, - kafka_ingress_partition=kafka_ingress_partition, - kafka_offset=kafka_offset, ) - self.sequencer.sequence(payload) + # Unbounded: the call holds its task until its whole chain is done, and + # that chain may need bounded slots on this worker. + self.background_functions.create_unbounded_task( + self._run_proxied_call(t_id, root_id, payload, fallback_mode), + ) + + async def _run_proxied_call( + self, + t_id: int, + root_id: int, + payload: RunFuncPayload, + fallback_mode: bool, + ) -> None: + """Run a peer's request here on its behalf. The chain acknowledges, answers + and aborts to this worker, and one ProxyResult then carries the outcome to + the root, so the root cannot commit before it has the response.""" + networking = self.networking + networking.prepare_function_chain(t_id) + payload.ack_payload = (networking.host_name, networking.host_port, t_id, "1.0", []) + try: + if fallback_mode: + await self.run_fallback_function(t_id, payload, internal=True) + else: + await self.run_function(t_id, payload) + await networking.waited_ack_events[t_id].wait() + except Exception as e: + logging.exception(f"Proxied call {t_id} from W{root_id} failed") + networking.abort_chain(t_id, str(e)) + host, _, protocol_port = self.peers[root_id] + await networking.send_message( + host, + protocol_port, + msg=( + t_id, + networking.client_responses.get(t_id), + networking.aborted_events.get(t_id), + [self.id, *networking.chain_participants.get(t_id, ())], + fallback_mode and self.local_state.has_fallback_rw_set_changed(t_id), + self.incarnation, + ), + msg_type=MessageType.ProxyResult, + serializer=Serializer.MSGPACK, + ) + + async def _handle_proxy_result(self, data: bytes) -> None: + # Lock-free: no awaits. + t_id, response, abort, participants, rw_changed, incarnation = self.networking.decode_message(data) + if incarnation != self.incarnation: + # From before a recovery: a replayed request may have the same t_id. + logging.warning(f"Dropping ProxyResult for T_ID {t_id} of incarnation {incarnation}") + return + result = self._proxy_results.pop(t_id, None) + if result is None or result.done(): + logging.debug(f"Dropping ProxyResult for unknown T_ID {t_id}") + return + result.set_result((response, abort, participants, rw_changed)) + + def _response_partition(self, payload: RunFuncPayload) -> int: + """Output partition for a root's response: its own partition, or for a + proxied request (partition on a peer) the ingress partition it came from, + which this worker owns.""" + if (payload.operator_name, payload.partition) in self.registered_operators: + return payload.partition + return payload.kafka_ingress_partition async def _handle_deprecated_rqrs(self, _: bytes) -> None: logging.error("REQUEST RESPONSE HAS BEEN DEPRECATED") @@ -1042,6 +1156,7 @@ def cleanup_after_epoch(self) -> None: self.concurrency_aborts_everywhere.clear() self.t_ids_to_reschedule.clear() self.fallback_rescheduled_t_ids.clear() + self._proxy_rw_changed.clear() self.wait_responses_to_be_sent.clear() self.networking.cleanup_after_epoch() self.local_state.cleanup() @@ -1076,7 +1191,7 @@ async def run_fallback_function( # Per the Aria paper: if the rw-set changed during fallback # re-execution, the dependency graph is invalid — reschedule # the transaction to the next epoch instead of committing. - rw_changed = self.local_state.has_fallback_rw_set_changed(t_id) + rw_changed = self.local_state.has_fallback_rw_set_changed(t_id) or t_id in self._proxy_rw_changed if not transaction_failed and not rw_changed: self.local_state.commit_fallback_transaction(t_id) elif rw_changed: @@ -1103,14 +1218,14 @@ async def run_fallback_function( key=payload.request_id, value=msgpack_serialization(self.networking.aborted_events[t_id]), operator_name=payload.operator_name, - partition=payload.partition, + partition=self._response_partition(payload), ) elif t_id in self.networking.client_responses: await self.egress.send( key=payload.request_id, value=msgpack_serialization(self.networking.client_responses[t_id]), operator_name=payload.operator_name, - partition=payload.partition, + partition=self._response_partition(payload), ) async def run_fallback_strategy(self) -> None: @@ -1176,7 +1291,7 @@ async def send_responses( t_id = sequenced_item.t_id request_id = sequenced_item.payload.request_id operator_name = sequenced_item.payload.operator_name - partition = sequenced_item.payload.partition + partition = self._response_partition(sequenced_item.payload) if t_id in aborted_events: await self.egress.send( key=request_id, From c8bf91917b8968a0474796064b2789ac832db8e3 Mon Sep 17 00:00:00 2001 From: Kyriakos Psarakis Date: Sun, 25 Oct 2026 08:56:24 +0100 Subject: [PATCH 16/20] Test that the ingress offset never passes an unrun request; count proxied calls The regression test feeds Kafka records through the real ingress and sequencer into the consuming worker's epoch: a request for a migrated key, then a local one. The offset it records must not pass a request that has not run. It fails on the forwarding code (offset 41 recorded while request 40 waits in the owner's sequencer). Workers log how many proxied calls they sent in an epoch, when there were any, so the E2E worker logs show whether the path ran. Co-Authored-By: Claude Opus 5.5 --- tests/unit/worker/test_proxied_requests.py | 470 ++++++++++++--------- worker/transactional_protocols/aria.py | 6 + 2 files changed, 272 insertions(+), 204 deletions(-) diff --git a/tests/unit/worker/test_proxied_requests.py b/tests/unit/worker/test_proxied_requests.py index e171289d..1c2f8cab 100644 --- a/tests/unit/worker/test_proxied_requests.py +++ b/tests/unit/worker/test_proxied_requests.py @@ -1,204 +1,266 @@ -"""Requests consumed by one worker for a key that migrated to a peer's partition. - -They stay in the consuming worker's epoch and run on the peer as a proxied call, -so an ingress partition's offset is only ever advanced by the worker consuming -it. The peer reports the whole outcome (response, abort, participants, fallback -read/write-set change) in one ProxyResult, which the root awaits before it can -commit. -""" - -import asyncio -from types import SimpleNamespace - -import pytest -from styx.common.message_types import MessageType -from styx.common.operator import Operator -from styx.common.run_func_payload import RunFuncPayload -from styx.common.serialization import Serializer -from styx.common.tcp_networking import NetworkingManager -from styx.common.util.aio_task_scheduler import AIOTaskScheduler - -from worker.operator_state.aria.in_memory_state import InMemoryOperatorState -from worker.transactional_protocols.aria import AriaProtocol - -OP = "op" -W1 = ("w1", 5001, 6001) # consumes ingress partition 0 -W2 = ("w2", 5002, 6002) # owns partition 1, where the key now lives -DNS = {OP: {0: W1, 1: W2}} -PEERS = {1: W1, 2: W2} - - -async def read(ctx): - return ctx.key, ctx.get() - - -async def fail(_ctx): - msg = "no such account" - raise ValueError(msg) - - -class _Link: - """Delivers protocol messages between fake workers, like the TCP readers.""" - - def __init__(self) -> None: - self.workers: dict[tuple[str, int], AriaProtocol] = {} - self.sent: list[MessageType] = [] - - async def deliver(self, host, port, msg, msg_type, serializer) -> bool: - self.sent.append(msg_type) - data = NetworkingManager.encode_message(msg=msg, msg_type=msg_type, serializer=serializer)[8:] - await asyncio.sleep(0) - await self.workers[(host, port)].protocol_handlers_map[msg_type](data) - return True - - -class _Networking(NetworkingManager): - def __init__(self, link: _Link, addr: tuple, worker_id: int) -> None: - super().__init__(addr[2]) - self.host_name = addr[0] - self.set_worker_id(worker_id) - self._link = link - - async def send_message(self, host, port, msg, msg_type, serializer=Serializer.MSGPACK) -> bool: - return await self._link.deliver(host, port, msg, msg_type, serializer) - - -def _worker(link: _Link, worker_id: int, addr: tuple, partitions: set[int]) -> AriaProtocol: - state = InMemoryOperatorState({(OP, p) for p in partitions}) - networking = _Networking(link, addr, worker_id) - registered = {} - for p in partitions: - operator = Operator(OP, n_partitions=2) - operator.register(read) - operator.register(fail) - operator.attach_state_networking(state, networking, DNS, None) - registered[(OP, p)] = operator - - proto = AriaProtocol.__new__(AriaProtocol) - proto.id = worker_id - proto.incarnation = 0 - proto.peers = {w: a for w, a in PEERS.items() if w != worker_id} - proto.dns = DNS - proto.local_state = state - proto.networking = networking - proto.registered_operators = registered - proto.background_functions = AIOTaskScheduler() - proto.waiting_on_transactions = {} - proto.fallback_locking_event_map = {} - proto._proxy_results = {} - proto._proxy_rw_changed = set() - proto.protocol_handlers_map = { - MessageType.RunFunProxy: proto._handle_run_fun_proxy, - MessageType.ProxyResult: proto._handle_proxy_result, - } - link.workers[(addr[0], addr[2])] = proto - return proto - - -def _cluster(): - link = _Link() - root = _worker(link, 1, W1, {0}) - owner = _worker(link, 2, W2, {1}) - owner.local_state.batch_insert({"k": "v"}, OP, 1) - return link, root, owner - - -def _misrouted(function_name: str = "read") -> RunFuncPayload: - # As the ingress sequences it: partition is where the key lives now, the - # ingress partition and offset are where it was consumed. - return RunFuncPayload( - request_id=b"rq", - key="k", - operator_name=OP, - partition=1, - function_name=function_name, - params=(), - kafka_offset=41, - kafka_ingress_partition=0, - ) - - -class TestProxiedRequest: - async def test_runs_on_owner_and_answers_at_the_root(self): - link, root, owner = _cluster() - - assert await root.run_function(7, _misrouted()) is True - - assert root.networking.client_responses[7] == ["k", "v"] - assert link.sent == [MessageType.RunFunProxy, MessageType.ProxyResult] - # The owner holds the transaction's reservation, for conflict detection, - # and is a chain participant, so a fallback run reaches it with Unlock. - assert "k" in owner.local_state.read_sets[(OP, 1)][7] - assert root.networking.chain_participants[7] == [2] - assert 7 not in root._proxy_results - - async def test_logic_abort_reaches_the_root(self): - _link, root, _owner = _cluster() - - assert await root.run_function(7, _misrouted("fail")) is False - - assert root.networking.aborted_events[7] == "no such account" - assert 7 in root.networking.logic_aborts_everywhere - assert 7 not in root.networking.client_responses - - async def test_fallback_read_write_set_change_reaches_the_root(self): - _link, root, owner = _cluster() - owner.local_state.has_fallback_rw_set_changed = lambda t_id: t_id == 7 - - await root.run_function(7, _misrouted(), fallback_mode=True) - - assert 7 in root._proxy_rw_changed - - async def test_response_goes_to_the_ingress_partition(self): - _link, root, _owner = _cluster() - payload = _misrouted() - assert root._response_partition(payload) == 0 - local = RunFuncPayload(request_id=b"r", key="x", operator_name=OP, partition=0, function_name="read", params=()) - assert root._response_partition(local) == 0 - - async def test_participant_call_for_a_missing_partition_is_not_proxied(self): - # Only a root (no ack payload) is proxied; a chain call must target a local partition. - link, root, _owner = _cluster() - payload = _misrouted() - payload.ack_payload = ("w9", 6009, 7, "1.0", []) - with pytest.raises(AttributeError): - await root.run_function(7, payload) - assert link.sent == [] - - -class TestStaleProxyMessages: - async def test_result_from_another_incarnation_is_dropped(self): - _link, root, _owner = _cluster() - pending = asyncio.get_running_loop().create_future() - root._proxy_results[7] = pending - data = NetworkingManager.encode_message( - msg=(7, "stale", None, [2], False, 99), - msg_type=MessageType.ProxyResult, - serializer=Serializer.MSGPACK, - )[8:] - - await root._handle_proxy_result(data) - - assert not pending.done() - - async def test_call_from_another_incarnation_is_not_run(self): - link, _root, owner = _cluster() - data = NetworkingManager.encode_message( - msg=(7, 1, b"rq", OP, "read", "k", 1, False, (), 99), - msg_type=MessageType.RunFunProxy, - serializer=Serializer.MSGPACK, - )[8:] - - await owner._handle_run_fun_proxy(data) - await asyncio.sleep(0) - - assert link.sent == [] - assert 7 not in owner.local_state.read_sets[(OP, 1)] - - -class TestIngressOffsetsStayWithTheConsumer: - def test_proxied_request_advances_the_consumers_ingress_offset(self): - proto = SimpleNamespace(concurrency_aborts_everywhere=set(), topic_partition_offsets={(OP, 0): 40}) - item = SimpleNamespace(t_id=7, payload=_misrouted()) - AriaProtocol._advance_offsets(proto, [item]) - assert proto.topic_partition_offsets == {(OP, 0): 41} +"""Requests consumed by one worker for a key that migrated to a peer's partition. + +They stay in the consuming worker's epoch and run on the peer as a proxied call, +so an ingress partition's offset is only ever advanced by the worker consuming +it. The peer reports the whole outcome (response, abort, participants, fallback +read/write-set change) in one ProxyResult, which the root awaits before it can +commit. +""" + +import asyncio +from types import SimpleNamespace + +import pytest +from styx.common.base_networking import BaseNetworking +from styx.common.message_types import MessageType +from styx.common.operator import Operator +from styx.common.run_func_payload import RunFuncPayload +from styx.common.serialization import Serializer +from styx.common.tcp_networking import NetworkingManager +from styx.common.util.aio_task_scheduler import AIOTaskScheduler + +from worker.ingress.styx_kafka_ingress import StyxKafkaIngress +from worker.operator_state.aria.in_memory_state import InMemoryOperatorState +from worker.sequencer.sequencer import Sequencer +from worker.transactional_protocols.aria import AriaProtocol + +OP = "op" +W1 = ("w1", 5001, 6001) # consumes ingress partition 0 +W2 = ("w2", 5002, 6002) # owns partition 1, where the key now lives +DNS = {OP: {0: W1, 1: W2}} +PEERS = {1: W1, 2: W2} + + +async def read(ctx): + return ctx.key, ctx.get() + + +async def fail(_ctx): + msg = "no such account" + raise ValueError(msg) + + +class _Link: + """Delivers protocol messages between fake workers, like the TCP readers.""" + + def __init__(self) -> None: + self.workers: dict[tuple[str, int], AriaProtocol] = {} + self.sent: list[MessageType] = [] + + async def deliver(self, host, port, msg, msg_type, serializer) -> bool: + self.sent.append(msg_type) + data = NetworkingManager.encode_message(msg=msg, msg_type=msg_type, serializer=serializer)[8:] + await asyncio.sleep(0) + await self.workers[(host, port)].protocol_handlers_map[msg_type](data) + return True + + +class _Networking(NetworkingManager): + def __init__(self, link: _Link, addr: tuple, worker_id: int) -> None: + super().__init__(addr[2]) + self.host_name = addr[0] + self.set_worker_id(worker_id) + self._link = link + + async def send_message(self, host, port, msg, msg_type, serializer=Serializer.MSGPACK) -> bool: + return await self._link.deliver(host, port, msg, msg_type, serializer) + + +def _worker(link: _Link, worker_id: int, addr: tuple, partitions: set[int]) -> AriaProtocol: + state = InMemoryOperatorState({(OP, p) for p in partitions}) + networking = _Networking(link, addr, worker_id) + registered = {} + for p in partitions: + operator = Operator(OP, n_partitions=2) + operator.register(read) + operator.register(fail) + operator.attach_state_networking(state, networking, DNS, None) + registered[(OP, p)] = operator + + proto = AriaProtocol.__new__(AriaProtocol) + proto.id = worker_id + proto.incarnation = 0 + proto.peers = {w: a for w, a in PEERS.items() if w != worker_id} + proto.dns = DNS + proto.local_state = state + proto.networking = networking + proto.registered_operators = registered + proto.background_functions = AIOTaskScheduler() + proto.waiting_on_transactions = {} + proto.fallback_locking_event_map = {} + proto._proxy_results = {} + proto._proxy_rw_changed = set() + proto._proxied_calls = 0 + proto.protocol_handlers_map = { + MessageType.RunFunProxy: proto._handle_run_fun_proxy, + MessageType.ProxyResult: proto._handle_proxy_result, + } + link.workers[(addr[0], addr[2])] = proto + return proto + + +def _cluster(): + link = _Link() + root = _worker(link, 1, W1, {0}) + owner = _worker(link, 2, W2, {1}) + owner.local_state.batch_insert({"k": "v"}, OP, 1) + return link, root, owner + + +def _misrouted(function_name: str = "read") -> RunFuncPayload: + # As the ingress sequences it: partition is where the key lives now, the + # ingress partition and offset are where it was consumed. + return RunFuncPayload( + request_id=b"rq", + key="k", + operator_name=OP, + partition=1, + function_name=function_name, + params=(), + kafka_offset=41, + kafka_ingress_partition=0, + ) + + +class TestProxiedRequest: + async def test_runs_on_owner_and_answers_at_the_root(self): + link, root, owner = _cluster() + + assert await root.run_function(7, _misrouted()) is True + + assert root.networking.client_responses[7] == ["k", "v"] + assert link.sent == [MessageType.RunFunProxy, MessageType.ProxyResult] + # The owner holds the transaction's reservation, for conflict detection, + # and is a chain participant, so a fallback run reaches it with Unlock. + assert "k" in owner.local_state.read_sets[(OP, 1)][7] + assert root.networking.chain_participants[7] == [2] + assert 7 not in root._proxy_results + assert root._proxied_calls == 1 + assert owner._proxied_calls == 0 + + async def test_logic_abort_reaches_the_root(self): + _link, root, _owner = _cluster() + + assert await root.run_function(7, _misrouted("fail")) is False + + assert root.networking.aborted_events[7] == "no such account" + assert 7 in root.networking.logic_aborts_everywhere + assert 7 not in root.networking.client_responses + + async def test_fallback_read_write_set_change_reaches_the_root(self): + _link, root, owner = _cluster() + owner.local_state.has_fallback_rw_set_changed = lambda t_id: t_id == 7 + + await root.run_function(7, _misrouted(), fallback_mode=True) + + assert 7 in root._proxy_rw_changed + + async def test_response_goes_to_the_ingress_partition(self): + _link, root, _owner = _cluster() + payload = _misrouted() + assert root._response_partition(payload) == 0 + local = RunFuncPayload(request_id=b"r", key="x", operator_name=OP, partition=0, function_name="read", params=()) + assert root._response_partition(local) == 0 + + async def test_participant_call_for_a_missing_partition_is_not_proxied(self): + # Only a root (no ack payload) is proxied; a chain call must target a local partition. + link, root, _owner = _cluster() + payload = _misrouted() + payload.ack_payload = ("w9", 6009, 7, "1.0", []) + with pytest.raises(AttributeError): + await root.run_function(7, payload) + assert link.sent == [] + + +class TestStaleProxyMessages: + async def test_result_from_another_incarnation_is_dropped(self): + _link, root, _owner = _cluster() + pending = asyncio.get_running_loop().create_future() + root._proxy_results[7] = pending + data = NetworkingManager.encode_message( + msg=(7, "stale", None, [2], False, 99), + msg_type=MessageType.ProxyResult, + serializer=Serializer.MSGPACK, + )[8:] + + await root._handle_proxy_result(data) + + assert not pending.done() + + async def test_call_from_another_incarnation_is_not_run(self): + link, _root, owner = _cluster() + data = NetworkingManager.encode_message( + msg=(7, 1, b"rq", OP, "read", "k", 1, False, (), 99), + msg_type=MessageType.RunFunProxy, + serializer=Serializer.MSGPACK, + )[8:] + + await owner._handle_run_fun_proxy(data) + await asyncio.sleep(0) + + assert link.sent == [] + assert 7 not in owner.local_state.read_sets[(OP, 1)] + + +class TestIngressOffsetsStayWithTheConsumer: + def test_proxied_request_advances_the_consumers_ingress_offset(self): + proto = SimpleNamespace(concurrency_aborts_everywhere=set(), topic_partition_offsets={(OP, 0): 40}) + item = SimpleNamespace(t_id=7, payload=_misrouted()) + AriaProtocol._advance_offsets(proto, [item]) + assert proto.topic_partition_offsets == {(OP, 0): 41} + + +def _key_on(operator: Operator, partition: int) -> str: + return next(key for key in (f"key{i}" for i in range(1000)) if operator.which_partition(key) == partition) + + +def _record(offset: int, key: str) -> SimpleNamespace: + # As the client sent it before the migration: addressed to ingress partition 0. + value = BaseNetworking.encode_message( + msg=(OP, key, "read", (), 0), + msg_type=MessageType.ClientMsg, + serializer=Serializer.MSGPACK, + ) + return SimpleNamespace(key=f"rq{offset}".encode(), offset=offset, partition=0, value=value) + + +class TestOffsetNeverPassesAnUnrunRequest: + async def test_snapshot_offset_covers_only_requests_that_ran(self): + # Regression: the consuming worker used to forward a request for a migrated + # key into the owner's epoch and advance the ingress offset past it for its + # own later requests. A snapshot taken before the owner ran it recorded an + # offset past it, and recovery skipped the request. + _link, root, owner = _cluster() + operator = root.registered_operators[(OP, 0)] + migrated, local = _key_on(operator, 1), _key_on(operator, 0) + owner.local_state.batch_insert({migrated: "m"}, OP, 1) + root.local_state.batch_insert({local: "l"}, OP, 0) + + sequencer = Sequencer() + sequencer.set_sequencer_id([2], 1) + ingress = StyxKafkaIngress( + networking=root.networking, + sequencer=sequencer, + state=root.local_state, + registered_operators=root.registered_operators, + worker_id=1, + kafka_url="unused", + epoch_interval_ms=1, + sequence_max_size=100, + ) + records = [_record(40, migrated), _record(41, local)] + for record in records: + ingress.handle_message_from_kafka(record) + + # The consuming worker's epoch, up to its snapshot. + root.concurrency_aborts_everywhere = set() + root.topic_partition_offsets = {(OP, 0): 39} + sequence = sequencer.get_epoch() + for item in sequence: + await root.run_function(item.t_id, item.payload) + root._advance_offsets(sequence) + + recorded = root.topic_partition_offsets[(OP, 0)] + answered = {item.payload.kafka_offset for item in sequence if item.t_id in root.networking.client_responses} + assert recorded == 41 + assert {r.offset for r in records if r.offset <= recorded} <= answered diff --git a/worker/transactional_protocols/aria.py b/worker/transactional_protocols/aria.py index 00f1fadc..17f4f8d0 100644 --- a/worker/transactional_protocols/aria.py +++ b/worker/transactional_protocols/aria.py @@ -190,6 +190,8 @@ def __init__( # noqa: PLR0913 self._proxy_results: dict[int, asyncio.Future] = {} # Proxied transactions whose read/write set changed in their fallback run. self._proxy_rw_changed: set[int] = set() + # Proxied calls sent this epoch, logged when there were any. + self._proxied_calls: int = 0 self.protocol_handlers_map: dict[MessageType, Callable[[bytes], Awaitable[None]]] = { MessageType.RunFunRemote: self._handle_run_fun_remote, @@ -430,6 +432,7 @@ async def _run_proxied(self, t_id: int, payload: RunFuncPayload, fallback_mode: workers advance one partition's offset, each by max, past requests the other had not run yet.) """ + self._proxied_calls += 1 host, _, protocol_port = self.dns[payload.operator_name][payload.partition] result: asyncio.Future = asyncio.get_running_loop().create_future() self._proxy_results[t_id] = result @@ -1157,6 +1160,9 @@ def cleanup_after_epoch(self) -> None: self.t_ids_to_reschedule.clear() self.fallback_rescheduled_t_ids.clear() self._proxy_rw_changed.clear() + if self._proxied_calls: + logging.warning(f"{self.id} ||| Proxied calls this epoch: {self._proxied_calls}") + self._proxied_calls = 0 self.wait_responses_to_be_sent.clear() self.networking.cleanup_after_epoch() self.local_state.cleanup() From 069468f89fc64fe26dc847d2af72ffa5170c1d13 Mon Sep 17 00:00:00 2001 From: Kyriakos Psarakis Date: Sun, 25 Oct 2026 10:37:58 +0100 Subject: [PATCH 17/20] Log what a stalled epoch is waiting on A stalled epoch stalls every worker at the next barrier and nothing is logged: a CI run of the kill-before-post-migration-snapshot test went silent for two minutes after recovering. Each epoch wait (transactions, chains, the WAL write, a barrier, fallback, migration batch acks) now records what it waits on, and a watchdog warns every EPOCH_STALL_WARN_S (default 10 s) while one lasts, with the unfinished chains and pending proxied calls. A ProxyResult with no pending call, which would leave its call waiting forever, is now logged as a warning instead of at debug. Co-Authored-By: Claude Opus 5.5 --- .../test_migration_snapshot_consistency.py | 1 + tests/unit/worker/test_proxied_requests.py | 41 ++++++++++++++++ worker/transactional_protocols/aria.py | 48 ++++++++++++++++++- 3 files changed, 89 insertions(+), 1 deletion(-) diff --git a/tests/unit/worker/test_migration_snapshot_consistency.py b/tests/unit/worker/test_migration_snapshot_consistency.py index eb48a89d..0881bb8e 100644 --- a/tests/unit/worker/test_migration_snapshot_consistency.py +++ b/tests/unit/worker/test_migration_snapshot_consistency.py @@ -287,6 +287,7 @@ def _stoppable_protocol() -> AriaProtocol: proto.function_scheduler_task = asyncio.create_task(asyncio.sleep(10)) proto.communication_task = asyncio.create_task(asyncio.sleep(10)) proto.migration_sender_task = None + proto._stall_watchdog_task = None proto.topic_partition_offsets = {} return proto diff --git a/tests/unit/worker/test_proxied_requests.py b/tests/unit/worker/test_proxied_requests.py index 1c2f8cab..e7956620 100644 --- a/tests/unit/worker/test_proxied_requests.py +++ b/tests/unit/worker/test_proxied_requests.py @@ -22,6 +22,7 @@ from worker.ingress.styx_kafka_ingress import StyxKafkaIngress from worker.operator_state.aria.in_memory_state import InMemoryOperatorState from worker.sequencer.sequencer import Sequencer +from worker.transactional_protocols import aria as aria_module from worker.transactional_protocols.aria import AriaProtocol OP = "op" @@ -264,3 +265,43 @@ async def test_snapshot_offset_covers_only_requests_that_ran(self): answered = {item.payload.kafka_offset for item in sequence if item.t_id in root.networking.client_responses} assert recorded == 41 assert {r.offset for r in records if r.offset <= recorded} <= answered + + +class TestStallDiagnostics: + async def test_watchdog_reports_what_a_stalled_epoch_waits_on(self, monkeypatch): + _link, root, _owner = _cluster() + warnings: list[str] = [] + monkeypatch.setattr(aria_module, "EPOCH_STALL_WARN_S", 0.01) + monkeypatch.setattr(aria_module.logging, "warning", warnings.append) + root.sequencer = SimpleNamespace(epoch_counter=5) + root.networking.prepare_function_chain(7) + root._proxy_results[9] = asyncio.get_running_loop().create_future() + + root._begin_wait("chains") + watchdog = asyncio.create_task(root._stall_watchdog()) + await asyncio.sleep(0.05) + root._end_wait() + reported = len(warnings) + await asyncio.sleep(0.05) + watchdog.cancel() + + assert reported > 0 + assert "STALL @epoch 5: waiting on chains" in warnings[0] + assert "unfinished chains: 1 [7]" in warnings[0] + assert "proxied calls pending: 1 [9]" in warnings[0] + # Nothing once the wait is over. + assert len(warnings) == reported + + async def test_result_with_no_pending_call_is_reported(self, monkeypatch): + _link, root, _owner = _cluster() + warnings: list[str] = [] + monkeypatch.setattr(aria_module.logging, "warning", warnings.append) + data = NetworkingManager.encode_message( + msg=(7, "v", None, [2], False, 0), + msg_type=MessageType.ProxyResult, + serializer=Serializer.MSGPACK, + )[8:] + + await root._handle_proxy_result(data) + + assert warnings == ["Dropping ProxyResult for T_ID 7 with no pending proxied call"] diff --git a/worker/transactional_protocols/aria.py b/worker/transactional_protocols/aria.py index 17f4f8d0..97bc7b10 100644 --- a/worker/transactional_protocols/aria.py +++ b/worker/transactional_protocols/aria.py @@ -54,6 +54,10 @@ ASYNC_MIGRATION_KEYS_PER_EPOCH: int = int(os.getenv("ASYNC_MIGRATION_KEYS_PER_EPOCH", "3_000")) # With no epoch running, the budget refills after this long instead. ASYNC_MIGRATION_IDLE_REFILL_S: float = 0.01 +# An epoch wait (transactions, chains, a barrier) that lasts this long is logged +# with what the worker is still waiting for. A stalled epoch stalls every worker +# at the next barrier, and nothing else would be logged. +EPOCH_STALL_WARN_S: float = float(os.getenv("EPOCH_STALL_WARN_S", "10")) class AriaProtocol(BaseTransactionalProtocol): @@ -157,6 +161,10 @@ def __init__( # noqa: PLR0913 self.function_scheduler_task: asyncio.Task | None = None self.communication_task: asyncio.Task | None = None self.migration_sender_task: asyncio.Task | None = None + self._stall_watchdog_task: asyncio.Task | None = None + # What the epoch is waiting on, and since when (for _stall_watchdog). + self._waiting_on: str | None = None + self._waiting_since: float = 0.0 self.max_t_counter: int = -1 self.total_processed_seq_size: int = -1 @@ -260,6 +268,8 @@ async def stop(self, *, keep_kafka_clients: bool = False) -> None: self.communication_task.cancel() if self.migration_sender_task is not None: self.migration_sender_task.cancel() + if self._stall_watchdog_task is not None: + self._stall_watchdog_task.cancel() try: await self.function_scheduler_task await self.communication_task @@ -276,10 +286,35 @@ def start(self) -> None: self.communication_task = asyncio.create_task(self.communication_protocol()) if self.migrating_state and USE_ASYNC_MIGRATION: self.migration_sender_task = asyncio.create_task(self._continuous_migration_sender()) + self._stall_watchdog_task = asyncio.create_task(self._stall_watchdog()) logging.warning( f"Aria protocol started with operator partitions: {list(self.registered_operators.keys())}", ) + def _begin_wait(self, what: str) -> None: + self._waiting_on = what + self._waiting_since = timer() + + def _end_wait(self) -> None: + self._waiting_on = None + + async def _stall_watchdog(self) -> None: + while True: + await asyncio.sleep(EPOCH_STALL_WARN_S) + what = self._waiting_on + if what is None: + continue + waited = timer() - self._waiting_since + if waited < EPOCH_STALL_WARN_S: + continue + unacked = [t_id for t_id, ack in self.networking.waited_ack_events.items() if not ack.is_set()] + proxied = list(self._proxy_results) + logging.warning( + f"{self.id} ||| STALL @epoch {self.sequencer.epoch_counter}: waiting on {what} for {waited:.0f}s" + f" | unfinished chains: {len(unacked)} {unacked[:10]}" + f" | proxied calls pending: {len(proxied)} {proxied[:10]}", + ) + async def run_function( self, t_id: int, @@ -544,7 +579,8 @@ async def _handle_proxy_result(self, data: bytes) -> None: return result = self._proxy_results.pop(t_id, None) if result is None or result.done(): - logging.debug(f"Dropping ProxyResult for unknown T_ID {t_id}") + # Its call would never complete; worth seeing even at the default level. + logging.warning(f"Dropping ProxyResult for T_ID {t_id} with no pending proxied call") return result.set_result((response, abort, participants, rw_changed)) @@ -699,7 +735,9 @@ async def _drain_migration_batches(self) -> None: at its destination, so the coming snapshot cut sees each migrated key on exactly one side.""" self._migration_gate.clear() + self._begin_wait("migration batch acks") await self._migration_all_acked.wait() + self._end_wait() async def _write_to_wal(self, sequence: list[SequencedItem]) -> tuple[float, float]: start_wal = timer() @@ -942,6 +980,7 @@ async def _run_epoch_functions_and_chain( wal_task = asyncio.create_task(self._write_to_wal(sequence)) try: start_func = timer() + self._begin_wait("transactions") await asyncio.gather( *[self.run_function(item.t_id, item.payload) for item in sequence], ) @@ -949,12 +988,15 @@ async def _run_epoch_functions_and_chain( # Wait for chains to finish start_chain = timer() + self._begin_wait("chains") await asyncio.gather( *[ack.wait() for ack in self.networking.waited_ack_events.values()], ) end_chain = timer() finally: + self._begin_wait("WAL write") start_wal, end_wal = await wal_task + self._end_wait() return { "wal_ms": (end_wal - start_wal) * 1000, @@ -1262,7 +1304,9 @@ async def run_fallback_strategy(self) -> None: serializer=Serializer.MSGPACK, ) if fallback_tasks: + self._begin_wait("fallback") await asyncio.gather(*fallback_tasks) + self._end_wait() # Flush all fallback egress sends in one batch (each root above # used `egress.send`, which enqueues; `send_batch` awaits the # producer futures together). @@ -1345,5 +1389,7 @@ async def sync_workers( msg_type=msg_type, serializer=serializer, ) + self._begin_wait(f"barrier {msg_type.name}") await self.sync_workers_event[msg_type].wait() + self._end_wait() self.sync_workers_event[msg_type].clear() From e62994520fee283dc561df586f02e667f00965df Mon Sep 17 00:00:00 2001 From: Kyriakos Psarakis Date: Sun, 25 Oct 2026 10:53:13 +0100 Subject: [PATCH 18/20] Don't recover a dead unreachable worker as a participant; log stalled barriers A worker a barrier or migration message could not be delivered to was handed to recovery at once, but only the heartbeat check removes a worker from the pool and reschedules its partitions. A dead one therefore stayed a participant, got InitRecovery, and recovery waited on it forever (reproduced locally: a worker exiting at the Phase B stop, under CPU-capped workers, failed the SyncCleanup send first). Unreachable workers now wait HEARTBEAT_LIMIT: by then the heartbeat check has removed a dead one, and the ones still in the pool are alive and are recovered as participants. The coordinator also warns when no barrier completes for BARRIER_STALL_WARN_S (default 10 s) while some workers wait, naming the workers it waits on, and workers mark the remaining epoch waits (snapshot delta, responses to Kafka, snapshot request) for the stall watchdog. Co-Authored-By: Claude Opus 5.5 --- coordinator/aria_sync_metadata.py | 3 + coordinator/coordinator_service.py | 55 +- tests/unit/coordinator/service_helpers.py | 91 ++-- .../test_migration_checkpoint_lifetime.py | 503 ++++++++++-------- worker/transactional_protocols/aria.py | 6 + 5 files changed, 384 insertions(+), 274 deletions(-) diff --git a/coordinator/aria_sync_metadata.py b/coordinator/aria_sync_metadata.py index a6e71cf0..ce0ccdca 100644 --- a/coordinator/aria_sync_metadata.py +++ b/coordinator/aria_sync_metadata.py @@ -15,6 +15,8 @@ def __init__(self, n_workers: int) -> None: self.global_read_set: dict | None = None self.stop_next_epoch: bool = False self.take_snapshot: bool = False + # Barriers completed so far, so a stalled one can be told from a new one. + self.completed_barriers: int = 0 def check_distributed_barrier(self) -> bool: return len(self.arrived) == self.n_workers @@ -122,6 +124,7 @@ def __merge_rw_reservations( return output_dict def cleanup(self, epoch_end: bool = False, take_snapshot: bool = False) -> None: + self.completed_barriers += 1 self.arrived.clear() self.logic_aborts_everywhere: set[int] = set() self.sent_proceed_msg: bool = False diff --git a/coordinator/coordinator_service.py b/coordinator/coordinator_service.py index c0b76b1b..20dbde39 100755 --- a/coordinator/coordinator_service.py +++ b/coordinator/coordinator_service.py @@ -30,6 +30,7 @@ from styx.common.tcp_networking import MessagingMode, NetworkingManager from styx.common.util.aio_task_scheduler import AIOTaskScheduler import uvloop +from worker_pool import HEARTBEAT_LIMIT from coordinator.migration_metadata import MigrationMetadata @@ -50,6 +51,9 @@ SNAPSHOT_BUCKET_NAME: str = os.getenv("SNAPSHOT_BUCKET_NAME", "styx-snapshots") SNAPSHOT_FREQUENCY_SEC = int(os.getenv("SNAPSHOT_FREQUENCY_SEC", "30")) +# A protocol barrier that stays incomplete this long is logged with the workers +# it still waits on: they stall every other worker, and nothing else is logged. +BARRIER_STALL_WARN_S: float = float(os.getenv("BARRIER_STALL_WARN_S", "10")) HEARTBEAT_CHECK_INTERVAL: int = int( os.getenv("HEARTBEAT_CHECK_INTERVAL", "1000"), ) # 1000ms @@ -115,6 +119,8 @@ def __init__(self) -> None: # Workers a barrier or migration message could not be delivered to. The # protocol cannot progress without them, so they are handed to recovery. self.unreachable_workers: list[Worker] = [] + # When each unreachable worker was first seen by the heartbeat monitor. + self._unreachable_since: dict[int, float] = {} self.recovery_lock: asyncio.Lock = asyncio.Lock() self.recovery_state: RecoveryState = RecoveryState.IDLE @@ -1266,9 +1272,10 @@ async def heartbeat_monitor_coroutine(self) -> None: time_since_last_heartbeat_ms, ) + unreachable = self._unreachable_to_recover(heartbeat_check_time) # Add workers that re-registered (same IP/ports) to the failed set if ( - workers_to_remove or self.workers_that_re_registered or self.unreachable_workers + workers_to_remove or self.workers_that_re_registered or unreachable ) and self.recovery_state == RecoveryState.IDLE: logging.debug( f"heartbeat_monitor | entering recovery lock " @@ -1285,8 +1292,10 @@ async def heartbeat_monitor_coroutine(self) -> None: re_registered_set = set(self.workers_that_re_registered) workers_to_remove.update(re_registered_set) self.workers_that_re_registered = [] - workers_to_remove.update(self.unreachable_workers) - self.unreachable_workers = [] + workers_to_remove.update(unreachable) + self.unreachable_workers = [w for w in self.unreachable_workers if w not in unreachable] + for worker in unreachable: + self._unreachable_since.pop(worker.worker_id, None) await self._perform_recovery(workers_to_remove) except Exception as e: @@ -1295,6 +1304,28 @@ async def heartbeat_monitor_coroutine(self) -> None: self.recovery_state = RecoveryState.IDLE logging.debug("heartbeat_monitor | recovery_state back to IDLE") + def _unreachable_to_recover(self, now: float) -> list[Worker]: + """The unreachable workers still alive HEARTBEAT_LIMIT after they were seen. + + An unreachable worker may be dead. Recovered right away, it would stay a + participant (only the heartbeat check removes a worker from the pool and + reschedules its partitions), and recovery would wait on it forever. A + dead one is removed by the heartbeat check within HEARTBEAT_LIMIT, so the + ones still in the pool by then are alive and are recovered as participants. + """ + participating = {w.worker_id for w in self.coordinator.worker_pool.get_participating_workers()} + due: list[Worker] = [] + waiting: list[Worker] = [] + for worker in self.unreachable_workers: + if worker.worker_id not in participating: + # Removed as dead: recovered (and rescheduled) as a failed worker. + self._unreachable_since.pop(worker.worker_id, None) + continue + since = self._unreachable_since.setdefault(worker.worker_id, now) + (due if (now - since) * 1000 >= HEARTBEAT_LIMIT else waiting).append(worker) + self.unreachable_workers = waiting + due + return due + async def send_snapshot_marker(self) -> None: while True: await asyncio.sleep(SNAPSHOT_FREQUENCY_SEC) @@ -1333,6 +1364,23 @@ def init_snapshot_bucket(self) -> None: else: return + async def barrier_watchdog(self) -> None: + last: tuple[AriaSyncMetadata, int] | None = None + while True: + await asyncio.sleep(BARRIER_STALL_WARN_S) + meta = self.aria_metadata + if meta is None or not meta.arrived: + last = None + continue + now = (meta, meta.completed_barriers) + if now == last: + participating = {w.worker_id for w in self.coordinator.worker_pool.get_participating_workers()} + logging.warning( + f"BARRIER STALL: no barrier completed for {BARRIER_STALL_WARN_S:.0f}s;" + f" arrived {sorted(meta.arrived)}, waiting on {sorted(participating - meta.arrived)}", + ) + last = now + async def main(self) -> None: logging.warning("Coordinator Booted Successfully") self.init_snapshot_bucket() @@ -1345,6 +1393,7 @@ async def main(self) -> None: logging.warning("Coordinator Heartbeat Sentinel online") self.snapshotting_task = asyncio.create_task(self.send_snapshot_marker()) logging.warning("Coordinator Snapshotting online") + self.barrier_watchdog_task = asyncio.create_task(self.barrier_watchdog()) await self.tcp_service() diff --git a/tests/unit/coordinator/service_helpers.py b/tests/unit/coordinator/service_helpers.py index 2b431fa2..9be87c38 100644 --- a/tests/unit/coordinator/service_helpers.py +++ b/tests/unit/coordinator/service_helpers.py @@ -1,45 +1,46 @@ -"""A CoordinatorService with its networking and recovery steps stubbed out, so -tests can drive the real migration and recovery handlers.""" - -import asyncio -from collections import defaultdict -from typing import TYPE_CHECKING -from unittest.mock import AsyncMock, MagicMock - -from styx.common.tcp_networking import NetworkingManager - -from coordinator.coordinator_service import CoordinatorService -from coordinator.migration_metadata import MigrationMetadata - -if TYPE_CHECKING: - from coordinator.coordinator_metadata import Coordinator - - -def stub_service(coordinator: Coordinator) -> CoordinatorService: - svc = CoordinatorService.__new__(CoordinatorService) - svc.coordinator = coordinator - svc.migration_in_progress = False - svc.migration_metadata = MigrationMetadata(2) - svc.networking_locks = defaultdict(asyncio.Lock) - svc._migration_checkpoint_task = None - svc._deferred_migration_graph = None - svc.unreachable_workers = [] - svc.networking = MagicMock(close_worker_connections=AsyncMock()) - svc.protocol_networking = MagicMock( - close_worker_connections=AsyncMock(), - close_all_connections=AsyncMock(), - decode_message=NetworkingManager.decode_message, - ) - svc.aio_task_scheduler = MagicMock(close=AsyncMock()) - svc._reset_after_recovery = AsyncMock() - svc._read_migration_blob_from_s3 = MagicMock(return_value=None) - coordinator.start_recovery_process = AsyncMock() - coordinator.wait_cluster_healthy = AsyncMock() - coordinator.notify_cluster_healthy = AsyncMock() - coordinator.revert_worker_pool_to_submitted_graph = MagicMock() - return svc - - -async def recover(svc: CoordinatorService) -> None: - """Run the real recovery for one dead worker.""" - await svc._perform_recovery({MagicMock(name="dead_worker")}) +"""A CoordinatorService with its networking and recovery steps stubbed out, so +tests can drive the real migration and recovery handlers.""" + +import asyncio +from collections import defaultdict +from typing import TYPE_CHECKING +from unittest.mock import AsyncMock, MagicMock + +from styx.common.tcp_networking import NetworkingManager + +from coordinator.coordinator_service import CoordinatorService +from coordinator.migration_metadata import MigrationMetadata + +if TYPE_CHECKING: + from coordinator.coordinator_metadata import Coordinator + + +def stub_service(coordinator: Coordinator) -> CoordinatorService: + svc = CoordinatorService.__new__(CoordinatorService) + svc.coordinator = coordinator + svc.migration_in_progress = False + svc.migration_metadata = MigrationMetadata(2) + svc.networking_locks = defaultdict(asyncio.Lock) + svc._migration_checkpoint_task = None + svc._deferred_migration_graph = None + svc.unreachable_workers = [] + svc._unreachable_since = {} + svc.networking = MagicMock(close_worker_connections=AsyncMock()) + svc.protocol_networking = MagicMock( + close_worker_connections=AsyncMock(), + close_all_connections=AsyncMock(), + decode_message=NetworkingManager.decode_message, + ) + svc.aio_task_scheduler = MagicMock(close=AsyncMock()) + svc._reset_after_recovery = AsyncMock() + svc._read_migration_blob_from_s3 = MagicMock(return_value=None) + coordinator.start_recovery_process = AsyncMock() + coordinator.wait_cluster_healthy = AsyncMock() + coordinator.notify_cluster_healthy = AsyncMock() + coordinator.revert_worker_pool_to_submitted_graph = MagicMock() + return svc + + +async def recover(svc: CoordinatorService) -> None: + """Run the real recovery for one dead worker.""" + await svc._perform_recovery({MagicMock(name="dead_worker")}) diff --git a/tests/unit/coordinator/test_migration_checkpoint_lifetime.py b/tests/unit/coordinator/test_migration_checkpoint_lifetime.py index 8dd2001b..f200384e 100644 --- a/tests/unit/coordinator/test_migration_checkpoint_lifetime.py +++ b/tests/unit/coordinator/test_migration_checkpoint_lifetime.py @@ -1,226 +1,277 @@ -"""How long a migration checkpoint stays attached to snapshots. - -A snapshot cut before the last migrated key landed still has keys at their old -partitions, so its sequencer file must carry the checkpoint even if it completes -after MigrationDone. Recovery from it resumes the transfer, and the coordinator -must keep tracking that resumed transfer. -""" - -from unittest.mock import AsyncMock, MagicMock - -from styx.common.local_state_backends import LocalStateBackend -from styx.common.message_types import MessageType -from styx.common.operator import Operator -from styx.common.serialization import Serializer, cloudpickle_serialization, zstd_msgpack_deserialization -from styx.common.stateflow_graph import StateflowGraph -from styx.common.tcp_networking import NetworkingManager - -from coordinator.coordinator_metadata import Coordinator -from coordinator.migration_metadata import MigrationMetadata -from tests.unit.coordinator.service_helpers import stub_service - - -def _graph(n_partitions: int) -> StateflowGraph: - g = StateflowGraph("app", operator_state_backend=LocalStateBackend.DICT, max_operator_parallelism=8) - g.add_operators(Operator("ycsb", n_partitions=n_partitions)) - return g - - -def _complete_snapshot(c: Coordinator, snapshot_id: int, epoch: int, *, checkpoint: bool = False): - """All (two) workers register snapshot_id; returns the sequencer tuple written to S3.""" - c.s3_client.put_object.reset_mock() - for worker_id in (0, 1): - c.register_snapshot(worker_id, snapshot_id, {}, {}, epoch, 0, MagicMock(), is_migration_checkpoint=checkpoint) - body = c.s3_client.put_object.call_args.kwargs["Body"] - return zstd_msgpack_deserialization(body) - - -def _coordinator_with_checkpoint() -> Coordinator: - c = Coordinator(MagicMock(), MagicMock()) - c.s3_client = MagicMock() - c.worker_snapshot_ids = {0: 2, 1: 2} - c.prev_completed_snapshot_id = 2 - c.set_migration_checkpoint(_graph(8), {}, from_snapshot_id=3) - return c - - -class TestCheckpointLifetime: - def test_snapshot_cut_before_last_key_landed_keeps_checkpoint(self): - c = _coordinator_with_checkpoint() - c.end_migration_checkpoint_after(transfer_done_epoch=50) - - # Cut at the end of epoch 49 (carries 50), completes after MigrationDone. - assert _complete_snapshot(c, 3, epoch=50)[4] is not None - # Cut at the end of epoch 50 (carries 51): every key had landed. - assert _complete_snapshot(c, 4, epoch=51)[4] is None - assert _complete_snapshot(c, 5, epoch=60)[4] is None - - def test_checkpoint_kept_while_transfer_runs(self): - c = _coordinator_with_checkpoint() - assert _complete_snapshot(c, 3, epoch=500)[4] is not None - - def test_new_checkpoint_resets_the_end(self): - c = _coordinator_with_checkpoint() - c.end_migration_checkpoint_after(transfer_done_epoch=50) - c.set_migration_checkpoint(_graph(16), {}, from_snapshot_id=3) - assert _complete_snapshot(c, 3, epoch=60)[4] is not None - - -class TestCheckpointTagging: - def test_snapshot_cut_before_phase_b_gets_no_checkpoint(self): - c = Coordinator(MagicMock(), MagicMock()) - c.s3_client = MagicMock() - c.worker_snapshot_ids = {0: 2, 1: 2} - c.prev_completed_snapshot_id = 2 - c.set_migration_checkpoint(_graph(8), {}) - c._migration_checkpoint_baseline_snap_id = 2 # waiting for the checkpoint - - # Started before Phase B, completes after the checkpoint metadata was set. - assert _complete_snapshot(c, 3, epoch=40)[4] is None - assert not c._migration_checkpoint_snapshot_complete.is_set() - - # The checkpoint itself, flagged by the workers. - assert _complete_snapshot(c, 4, epoch=41, checkpoint=True)[4] is not None - assert c._migration_checkpoint_snapshot_complete.is_set() - # Snapshots after it keep the checkpoint until the transfer ends. - assert _complete_snapshot(c, 5, epoch=45)[4] is not None - - def test_checkpoint_not_complete_until_every_worker_has_it(self): - c = Coordinator(MagicMock(), MagicMock()) - c.s3_client = MagicMock() - c.worker_snapshot_ids = {0: 2, 1: 2} - c.prev_completed_snapshot_id = 2 - c.set_migration_checkpoint(_graph(8), {}) - c._migration_checkpoint_baseline_snap_id = 2 - c.register_snapshot(0, 3, {}, {}, 41, 0, MagicMock(), is_migration_checkpoint=True) - assert not c.migration_checkpoint_completed() - c.register_snapshot(1, 3, {}, {}, 41, 0, MagicMock(), is_migration_checkpoint=True) - assert c.migration_checkpoint_completed() - - -class TestMigrationDoneEpochs: - async def test_transfer_done_epoch_is_the_latest_reported(self): - m = MigrationMetadata(2) - assert await m.migration_done(12) is False - assert await m.migration_done(10) is True - assert m.transfer_done_epoch == 12 - await m.cleanup(MessageType.MigrationDone) - assert m.transfer_done_epoch == -1 - - -def _migration_done_msg(epoch: int) -> bytes: - return NetworkingManager.encode_message( - msg=(epoch,), - msg_type=MessageType.MigrationDone, - serializer=Serializer.MSGPACK, - )[8:] - - -class TestMigrationDoneHandler: - async def test_all_workers_done_ends_checkpoint_after_latest_epoch(self): - c = _coordinator_with_checkpoint() - svc = stub_service(c) - svc.migration_in_progress = True - - await svc._handle_migration_done(_migration_done_msg(41)) - assert svc.migration_in_progress is True - await svc._handle_migration_done(_migration_done_msg(43)) - - assert svc.migration_in_progress is False - assert c._migration_checkpoint_blob is not None # kept for snapshots cut up to 43 - assert c._migration_checkpoint_until_epoch == 43 - - -def _blob(graph: StateflowGraph) -> bytes: - return cloudpickle_serialization({"new_graph": graph, "operator_partition_locations": {}}) - - -class TestRecoveryFromCheckpoint: - async def test_resume_after_migration_done_keeps_tracking_the_transfer(self): - c = Coordinator(MagicMock(), MagicMock()) - c.worker_snapshot_ids = {0: 7, 1: 7} - svc = stub_service(c) - # MigrationDone already processed; the latest snapshot still has the checkpoint. - svc._read_migration_blob_from_s3 = MagicMock(return_value=_blob(_graph(8))) - state_when_workers_released = {} - - async def notify() -> None: - state_when_workers_released["in_progress"] = svc.migration_in_progress - state_when_workers_released["blob"] = c._migration_checkpoint_blob - - c.notify_cluster_healthy = AsyncMock(side_effect=notify) - await svc._perform_recovery({MagicMock(name="dead_worker")}) - - svc._read_migration_blob_from_s3.assert_called_once_with(7) - assert c.submitted_graph.nodes["ycsb"].n_partitions == 8 - assert state_when_workers_released["in_progress"] is True - assert state_when_workers_released["blob"] is not None - - async def test_newer_migration_without_checkpoint_is_retried_after_resume(self): - c = Coordinator(MagicMock(), MagicMock()) - c.worker_snapshot_ids = {0: 7, 1: 7} - newer = _graph(16) - c._pending_graph = newer - svc = stub_service(c) - svc.migration_in_progress = True - svc._read_migration_blob_from_s3 = MagicMock(return_value=_blob(_graph(8))) - - await svc._perform_recovery({MagicMock(name="dead_worker")}) - - assert svc._deferred_migration_graph is newer - assert c.submitted_graph.nodes["ycsb"].n_partitions == 8 - c.revert_worker_pool_to_submitted_graph.assert_called_once() - - async def test_crash_before_checkpoint_reverts_and_retries(self): - c = Coordinator(MagicMock(), MagicMock()) - c.worker_snapshot_ids = {0: 7, 1: 7} - pending = _graph(8) - c._pending_graph = pending - svc = stub_service(c) - svc.migration_in_progress = True - svc._read_migration_blob_from_s3 = MagicMock(return_value=None) - - await svc._perform_recovery({MagicMock(name="dead_worker")}) - - assert svc.migration_in_progress is False - assert svc._deferred_migration_graph is pending - assert c._migration_checkpoint_blob is None - c.revert_worker_pool_to_submitted_graph.assert_called_once() - - async def test_ordinary_recovery_stays_ordinary(self): - c = Coordinator(MagicMock(), MagicMock()) - c.worker_snapshot_ids = {0: 7, 1: 7} - svc = stub_service(c) - svc._read_migration_blob_from_s3 = MagicMock(return_value=None) - - await svc._perform_recovery({MagicMock(name="dead_worker")}) - - assert svc.migration_in_progress is False - assert svc._deferred_migration_graph is None - c.revert_worker_pool_to_submitted_graph.assert_not_called() - - -class TestUnreachableWorkers: - def _service_with_workers(self, *results): - c = Coordinator(MagicMock(), MagicMock()) - svc = stub_service(c) - svc.unreachable_workers = [] - workers = [MagicMock(worker_id=i, worker_ip=f"w{i}", protocol_port=6000 + i) for i in range(len(results))] - c.worker_pool = MagicMock(get_participating_workers=MagicMock(return_value=workers)) - svc.protocol_networking.send_message = AsyncMock(side_effect=list(results)) - return svc, workers - - async def test_worker_a_barrier_reply_cannot_reach_is_handed_to_recovery(self): - svc, workers = self._service_with_workers(True, False) - await svc.finalize_worker_sync(MessageType.SyncCleanup, (True,)) - assert svc.unreachable_workers == [workers[1]] - - async def test_connection_error_counts_as_undelivered(self): - svc, workers = self._service_with_workers(ConnectionRefusedError(), True) - await svc.finalize_worker_sync(MessageType.SyncCleanup, (True,)) - assert svc.unreachable_workers == [workers[0]] - - async def test_all_delivered(self): - svc, _ = self._service_with_workers(True, True) - await svc.finalize_worker_sync(MessageType.SyncCleanup, (False,)) - assert svc.unreachable_workers == [] +"""How long a migration checkpoint stays attached to snapshots. + +A snapshot cut before the last migrated key landed still has keys at their old +partitions, so its sequencer file must carry the checkpoint even if it completes +after MigrationDone. Recovery from it resumes the transfer, and the coordinator +must keep tracking that resumed transfer. +""" + +import asyncio +from unittest.mock import AsyncMock, MagicMock + +from styx.common.local_state_backends import LocalStateBackend +from styx.common.message_types import MessageType +from styx.common.operator import Operator +from styx.common.serialization import Serializer, cloudpickle_serialization, zstd_msgpack_deserialization +from styx.common.stateflow_graph import StateflowGraph +from styx.common.tcp_networking import NetworkingManager + +from coordinator import coordinator_service +from coordinator.aria_sync_metadata import AriaSyncMetadata +from coordinator.coordinator_metadata import Coordinator +from coordinator.migration_metadata import MigrationMetadata +from coordinator.worker_pool import HEARTBEAT_LIMIT +from tests.unit.coordinator.service_helpers import stub_service + + +def _graph(n_partitions: int) -> StateflowGraph: + g = StateflowGraph("app", operator_state_backend=LocalStateBackend.DICT, max_operator_parallelism=8) + g.add_operators(Operator("ycsb", n_partitions=n_partitions)) + return g + + +def _complete_snapshot(c: Coordinator, snapshot_id: int, epoch: int, *, checkpoint: bool = False): + """All (two) workers register snapshot_id; returns the sequencer tuple written to S3.""" + c.s3_client.put_object.reset_mock() + for worker_id in (0, 1): + c.register_snapshot(worker_id, snapshot_id, {}, {}, epoch, 0, MagicMock(), is_migration_checkpoint=checkpoint) + body = c.s3_client.put_object.call_args.kwargs["Body"] + return zstd_msgpack_deserialization(body) + + +def _coordinator_with_checkpoint() -> Coordinator: + c = Coordinator(MagicMock(), MagicMock()) + c.s3_client = MagicMock() + c.worker_snapshot_ids = {0: 2, 1: 2} + c.prev_completed_snapshot_id = 2 + c.set_migration_checkpoint(_graph(8), {}, from_snapshot_id=3) + return c + + +class TestCheckpointLifetime: + def test_snapshot_cut_before_last_key_landed_keeps_checkpoint(self): + c = _coordinator_with_checkpoint() + c.end_migration_checkpoint_after(transfer_done_epoch=50) + + # Cut at the end of epoch 49 (carries 50), completes after MigrationDone. + assert _complete_snapshot(c, 3, epoch=50)[4] is not None + # Cut at the end of epoch 50 (carries 51): every key had landed. + assert _complete_snapshot(c, 4, epoch=51)[4] is None + assert _complete_snapshot(c, 5, epoch=60)[4] is None + + def test_checkpoint_kept_while_transfer_runs(self): + c = _coordinator_with_checkpoint() + assert _complete_snapshot(c, 3, epoch=500)[4] is not None + + def test_new_checkpoint_resets_the_end(self): + c = _coordinator_with_checkpoint() + c.end_migration_checkpoint_after(transfer_done_epoch=50) + c.set_migration_checkpoint(_graph(16), {}, from_snapshot_id=3) + assert _complete_snapshot(c, 3, epoch=60)[4] is not None + + +class TestCheckpointTagging: + def test_snapshot_cut_before_phase_b_gets_no_checkpoint(self): + c = Coordinator(MagicMock(), MagicMock()) + c.s3_client = MagicMock() + c.worker_snapshot_ids = {0: 2, 1: 2} + c.prev_completed_snapshot_id = 2 + c.set_migration_checkpoint(_graph(8), {}) + c._migration_checkpoint_baseline_snap_id = 2 # waiting for the checkpoint + + # Started before Phase B, completes after the checkpoint metadata was set. + assert _complete_snapshot(c, 3, epoch=40)[4] is None + assert not c._migration_checkpoint_snapshot_complete.is_set() + + # The checkpoint itself, flagged by the workers. + assert _complete_snapshot(c, 4, epoch=41, checkpoint=True)[4] is not None + assert c._migration_checkpoint_snapshot_complete.is_set() + # Snapshots after it keep the checkpoint until the transfer ends. + assert _complete_snapshot(c, 5, epoch=45)[4] is not None + + def test_checkpoint_not_complete_until_every_worker_has_it(self): + c = Coordinator(MagicMock(), MagicMock()) + c.s3_client = MagicMock() + c.worker_snapshot_ids = {0: 2, 1: 2} + c.prev_completed_snapshot_id = 2 + c.set_migration_checkpoint(_graph(8), {}) + c._migration_checkpoint_baseline_snap_id = 2 + c.register_snapshot(0, 3, {}, {}, 41, 0, MagicMock(), is_migration_checkpoint=True) + assert not c.migration_checkpoint_completed() + c.register_snapshot(1, 3, {}, {}, 41, 0, MagicMock(), is_migration_checkpoint=True) + assert c.migration_checkpoint_completed() + + +class TestMigrationDoneEpochs: + async def test_transfer_done_epoch_is_the_latest_reported(self): + m = MigrationMetadata(2) + assert await m.migration_done(12) is False + assert await m.migration_done(10) is True + assert m.transfer_done_epoch == 12 + await m.cleanup(MessageType.MigrationDone) + assert m.transfer_done_epoch == -1 + + +def _migration_done_msg(epoch: int) -> bytes: + return NetworkingManager.encode_message( + msg=(epoch,), + msg_type=MessageType.MigrationDone, + serializer=Serializer.MSGPACK, + )[8:] + + +class TestMigrationDoneHandler: + async def test_all_workers_done_ends_checkpoint_after_latest_epoch(self): + c = _coordinator_with_checkpoint() + svc = stub_service(c) + svc.migration_in_progress = True + + await svc._handle_migration_done(_migration_done_msg(41)) + assert svc.migration_in_progress is True + await svc._handle_migration_done(_migration_done_msg(43)) + + assert svc.migration_in_progress is False + assert c._migration_checkpoint_blob is not None # kept for snapshots cut up to 43 + assert c._migration_checkpoint_until_epoch == 43 + + +def _blob(graph: StateflowGraph) -> bytes: + return cloudpickle_serialization({"new_graph": graph, "operator_partition_locations": {}}) + + +class TestRecoveryFromCheckpoint: + async def test_resume_after_migration_done_keeps_tracking_the_transfer(self): + c = Coordinator(MagicMock(), MagicMock()) + c.worker_snapshot_ids = {0: 7, 1: 7} + svc = stub_service(c) + # MigrationDone already processed; the latest snapshot still has the checkpoint. + svc._read_migration_blob_from_s3 = MagicMock(return_value=_blob(_graph(8))) + state_when_workers_released = {} + + async def notify() -> None: + state_when_workers_released["in_progress"] = svc.migration_in_progress + state_when_workers_released["blob"] = c._migration_checkpoint_blob + + c.notify_cluster_healthy = AsyncMock(side_effect=notify) + await svc._perform_recovery({MagicMock(name="dead_worker")}) + + svc._read_migration_blob_from_s3.assert_called_once_with(7) + assert c.submitted_graph.nodes["ycsb"].n_partitions == 8 + assert state_when_workers_released["in_progress"] is True + assert state_when_workers_released["blob"] is not None + + async def test_newer_migration_without_checkpoint_is_retried_after_resume(self): + c = Coordinator(MagicMock(), MagicMock()) + c.worker_snapshot_ids = {0: 7, 1: 7} + newer = _graph(16) + c._pending_graph = newer + svc = stub_service(c) + svc.migration_in_progress = True + svc._read_migration_blob_from_s3 = MagicMock(return_value=_blob(_graph(8))) + + await svc._perform_recovery({MagicMock(name="dead_worker")}) + + assert svc._deferred_migration_graph is newer + assert c.submitted_graph.nodes["ycsb"].n_partitions == 8 + c.revert_worker_pool_to_submitted_graph.assert_called_once() + + async def test_crash_before_checkpoint_reverts_and_retries(self): + c = Coordinator(MagicMock(), MagicMock()) + c.worker_snapshot_ids = {0: 7, 1: 7} + pending = _graph(8) + c._pending_graph = pending + svc = stub_service(c) + svc.migration_in_progress = True + svc._read_migration_blob_from_s3 = MagicMock(return_value=None) + + await svc._perform_recovery({MagicMock(name="dead_worker")}) + + assert svc.migration_in_progress is False + assert svc._deferred_migration_graph is pending + assert c._migration_checkpoint_blob is None + c.revert_worker_pool_to_submitted_graph.assert_called_once() + + async def test_ordinary_recovery_stays_ordinary(self): + c = Coordinator(MagicMock(), MagicMock()) + c.worker_snapshot_ids = {0: 7, 1: 7} + svc = stub_service(c) + svc._read_migration_blob_from_s3 = MagicMock(return_value=None) + + await svc._perform_recovery({MagicMock(name="dead_worker")}) + + assert svc.migration_in_progress is False + assert svc._deferred_migration_graph is None + c.revert_worker_pool_to_submitted_graph.assert_not_called() + + +class TestUnreachableWorkers: + def _service_with_workers(self, *results): + c = Coordinator(MagicMock(), MagicMock()) + svc = stub_service(c) + svc.unreachable_workers = [] + workers = [MagicMock(worker_id=i, worker_ip=f"w{i}", protocol_port=6000 + i) for i in range(len(results))] + c.worker_pool = MagicMock(get_participating_workers=MagicMock(return_value=workers)) + svc.protocol_networking.send_message = AsyncMock(side_effect=list(results)) + return svc, workers + + async def test_worker_a_barrier_reply_cannot_reach_is_handed_to_recovery(self): + svc, workers = self._service_with_workers(True, False) + await svc.finalize_worker_sync(MessageType.SyncCleanup, (True,)) + assert svc.unreachable_workers == [workers[1]] + + async def test_connection_error_counts_as_undelivered(self): + svc, workers = self._service_with_workers(ConnectionRefusedError(), True) + await svc.finalize_worker_sync(MessageType.SyncCleanup, (True,)) + assert svc.unreachable_workers == [workers[0]] + + async def test_all_delivered(self): + svc, _ = self._service_with_workers(True, True) + await svc.finalize_worker_sync(MessageType.SyncCleanup, (False,)) + assert svc.unreachable_workers == [] + + def test_dead_unreachable_worker_is_left_to_the_heartbeat_check(self): + # Recovered as a participant, a dead worker would get InitRecovery and + # recovery would wait on it forever; the heartbeat check removes it instead. + svc, workers = self._service_with_workers(True, True) + svc.unreachable_workers = [workers[1]] + svc.coordinator.worker_pool.get_participating_workers.return_value = [workers[0]] + + assert svc._unreachable_to_recover(100.0) == [] + assert svc.unreachable_workers == [] + assert svc._unreachable_since == {} + + def test_live_unreachable_worker_is_recovered_after_the_heartbeat_limit(self): + svc, workers = self._service_with_workers(True, True) + svc.unreachable_workers = [workers[1]] + limit_s = HEARTBEAT_LIMIT / 1000 + + assert svc._unreachable_to_recover(100.0) == [] + assert svc._unreachable_to_recover(100.0 + limit_s / 2) == [] + assert svc._unreachable_to_recover(100.0 + limit_s) == [workers[1]] + + +class TestBarrierWatchdog: + async def test_names_the_workers_a_stalled_barrier_waits_on(self, monkeypatch): + c = Coordinator(MagicMock(), MagicMock()) + svc = stub_service(c) + workers = [MagicMock(worker_id=i) for i in (1, 2, 3)] + c.worker_pool = MagicMock(get_participating_workers=MagicMock(return_value=workers)) + svc.aria_metadata = AriaSyncMetadata(3) + svc.aria_metadata.arrived.update({1, 3}) + warnings: list[str] = [] + monkeypatch.setattr(coordinator_service, "BARRIER_STALL_WARN_S", 0.01) + monkeypatch.setattr(coordinator_service.logging, "warning", warnings.append) + + watchdog = asyncio.create_task(svc.barrier_watchdog()) + await asyncio.sleep(0.05) + stalled = list(warnings) + # Barriers completing again: nothing more is reported. + for _ in range(10): + svc.aria_metadata.cleanup() + svc.aria_metadata.arrived.add(1) + await asyncio.sleep(0.01) + watchdog.cancel() + + assert stalled + assert "arrived [1, 3], waiting on [2]" in stalled[0] + assert len(warnings) <= len(stalled) + 1 diff --git a/worker/transactional_protocols/aria.py b/worker/transactional_protocols/aria.py index 97bc7b10..84cca132 100644 --- a/worker/transactional_protocols/aria.py +++ b/worker/transactional_protocols/aria.py @@ -924,7 +924,9 @@ async def _process_epoch(self, sequence: list[SequencedItem]) -> None: # HERE WE KNOW ALL THE CONCURRENCY ABORTS commit_start = timer() n_committed = self._commit_and_prepare_responses(sequence) + self._begin_wait("snapshot delta") await self.send_delta_to_snapshotting_proc() + self._end_wait() commit_end = timer() # Share of this worker's own transactions aborted anywhere in the cluster. @@ -939,11 +941,15 @@ async def _process_epoch(self, sequence: list[SequencedItem]) -> None: self._advance_offsets(sequence) self.sequencer.increment_epoch(self.max_t_counter, self.t_ids_to_reschedule) + self._begin_wait("responses to Kafka") await self.wait_responses_to_be_sent.wait() + self._end_wait() self.cleanup_after_epoch() snap_start = timer() + self._begin_wait("snapshot request") await self.take_snapshot() + self._end_wait() snap_end = timer() epoch_end = timer() From f006ad2a31cbdddcfb8141df856631326c1f766c Mon Sep 17 00:00:00 2001 From: Kyriakos Psarakis Date: Sun, 25 Oct 2026 13:00:45 +0100 Subject: [PATCH 19/20] Carry rescheduled requests across snapshots and migration stops An epoch's ingress offset passed requests rescheduled to a later epoch (a fallback read/write-set change, or a concurrency abort with fallback off): with fallback the offset advanced past them, without it the max over the committed ones did. They were only held in the sequencer, so a recovery from a snapshot cut in between, or the protocol rebuilt after a Phase B stop, never ran them. Offsets now advance past every request of the epoch, and the requests still in the sequencer at or below their partition's recorded offset are pending. A snapshot carries them (SnapTakeSnapshot, SnapID, a sixth element of the sequencer file), as do MigrationInitDone and MigrationDone. The protocol that takes over (after recovery, possibly on another worker, or after a migration) sequences the ones of its ingress partitions before consuming anything, routed to the partition that holds the key now. The coordinator now writes each sequencer file from its own snapshot's reports. It used to fold every report into the completed offsets on arrival, so a worker's report of snapshot N+1 that came before another worker's report of N put N+1's offsets in N's file, past requests N's state does not contain. Co-Authored-By: Claude Opus 5.5 --- coordinator/coordinator_metadata.py | 47 ++++-- coordinator/coordinator_service.py | 10 +- coordinator/migration_metadata.py | 6 + .../coordinator/test_coordinator_metadata.py | 55 ++++++- tests/unit/worker/test_async_snapshots.py | 54 +++++-- tests/unit/worker/test_pending_requests.py | 150 ++++++++++++++++++ worker/async_snapshotting.py | 2 + worker/fault_tolerance/async_snapshots.py | 40 ++++- worker/ingress/styx_kafka_ingress.py | 40 +++-- worker/sequencer/pending_requests.py | 81 ++++++++++ worker/transactional_protocols/aria.py | 29 +++- worker/worker_service.py | 15 +- 12 files changed, 478 insertions(+), 51 deletions(-) create mode 100644 tests/unit/worker/test_pending_requests.py create mode 100644 worker/sequencer/pending_requests.py diff --git a/coordinator/coordinator_metadata.py b/coordinator/coordinator_metadata.py index 2ee95743..3e77cf59 100644 --- a/coordinator/coordinator_metadata.py +++ b/coordinator/coordinator_metadata.py @@ -54,6 +54,14 @@ def __init__(self, networking: NetworkingManager, s3_client: S3Client) -> None: self.completed_out_offsets: dict[OperatorPartition, int] = {} self.completed_epoch_counter: int = 0 self.completed_t_counter: int = 0 + # Requests behind completed_input_offsets that had not committed, of the + # last completed snapshot (worker/sequencer/pending_requests.py). + self.completed_pending_requests: list = [] + # Worker reports of snapshots not yet complete on every worker, by snapshot + # id: (input offsets, output offsets, epoch, t_counter, pending requests). + # A sequencer file is written from its own snapshot's reports only, never + # pairing that snapshot's state with a later one's offsets. + self.snapshot_reports: dict[int, list[tuple]] = {} self.worker_pool = WorkerPool() self.submitted_graph: StateflowGraph | None = None @@ -217,6 +225,7 @@ def register_snapshot( pool: concurrent.futures.ProcessPoolExecutor, *, is_migration_checkpoint: bool = False, + pending_requests: list | None = None, ) -> None: self.worker_snapshot_ids[worker_id] = snapshot_id if ( @@ -226,19 +235,12 @@ def register_snapshot( ): self._migration_checkpoint_snap_id = snapshot_id logging.warning(f"Migration checkpoint is snapshot {snapshot_id}") - for operator_partition, pio in partial_input_offsets.items(): - self.completed_input_offsets[operator_partition] = max( - pio, - self.completed_input_offsets.get(operator_partition, pio), - ) - for operator_partition, poo in partial_output_offsets.items(): - self.completed_out_offsets[operator_partition] = max( - poo, - self.completed_out_offsets.get(operator_partition, poo), - ) - self.completed_epoch_counter = epoch_counter - self.completed_t_counter = t_counter + self.snapshot_reports.setdefault(snapshot_id, []).append( + (partial_input_offsets, partial_output_offsets, epoch_counter, t_counter, pending_requests or []), + ) current_completed_snapshot: int = self.get_current_completed_snapshot_id() + if current_completed_snapshot != self.prev_completed_snapshot_id: + self._merge_snapshot_reports(current_completed_snapshot) logging.debug( f"register_snapshot | worker={worker_id} snap={snapshot_id} " f"cluster_min={current_completed_snapshot} prev={self.prev_completed_snapshot_id} " @@ -278,6 +280,7 @@ def register_snapshot( self.completed_epoch_counter, self.completed_t_counter, self._migration_checkpoint_blob if self.migration_checkpoint_completed() else None, + self.completed_pending_requests, ), ) self.s3_client.put_object( @@ -294,6 +297,26 @@ def register_snapshot( current_completed_snapshot, ) + def _merge_snapshot_reports(self, completed_snapshot_id: int) -> None: + """Fold the reports of every snapshot up to the completed one into the + completed offsets and counters; its pending requests are that snapshot's.""" + for snapshot_id in sorted(s for s in self.snapshot_reports if s <= completed_snapshot_id): + reports = self.snapshot_reports.pop(snapshot_id) + for input_offsets, output_offsets, _, _, _ in reports: + for operator_partition, pio in input_offsets.items(): + self.completed_input_offsets[operator_partition] = max( + pio, + self.completed_input_offsets.get(operator_partition, pio), + ) + for operator_partition, poo in output_offsets.items(): + self.completed_out_offsets[operator_partition] = max( + poo, + self.completed_out_offsets.get(operator_partition, poo), + ) + self.completed_epoch_counter = max(report[2] for report in reports) + self.completed_t_counter = max(report[3] for report in reports) + self.completed_pending_requests = [fields for report in reports for fields in report[4]] + def migration_checkpoint_completed(self) -> bool: """Whether the cluster has completed the checkpoint snapshot (or a later one).""" return ( diff --git a/coordinator/coordinator_service.py b/coordinator/coordinator_service.py index 20dbde39..c299fd04 100755 --- a/coordinator/coordinator_service.py +++ b/coordinator/coordinator_service.py @@ -411,13 +411,16 @@ async def _handle_migration_init_done( logging.warning("Dropping stale MigrationInitDone (recovery intervened)") return - epoch_counter, t_counter, input_offsets, output_offsets = self.networking.decode_message(data) + epoch_counter, t_counter, input_offsets, output_offsets, pending_requests = self.networking.decode_message( + data, + ) sync_complete: bool = await self.migration_metadata.init_done( epoch_counter, t_counter, input_offsets, output_offsets, + pending_requests, ) logging.warning(f"MIGRATION | MigrationInitDone | {self.migration_metadata.sync_sum}") @@ -512,6 +515,7 @@ async def _handle_snap_id( t_counter, sn_size, is_migration_checkpoint, + pending_requests, ) = self.networking.decode_message(data) snapshot_time = end - start @@ -536,6 +540,7 @@ async def _handle_snap_id( t_counter, pool, is_migration_checkpoint=is_migration_checkpoint, + pending_requests=pending_requests, ) # Check for deferred migration re-trigger: fire once the first @@ -927,6 +932,7 @@ async def finalize_migration(self) -> None: self.migration_metadata.t_counter, self.migration_metadata.input_offsets, self.migration_metadata.output_offsets, + self.migration_metadata.pending_requests, ), MessageType.MigrationDone, Serializer.MSGPACK, @@ -1048,6 +1054,8 @@ async def _reset_after_recovery(self) -> None: self.coordinator.completed_out_offsets.clear() self.coordinator.completed_epoch_counter = 0 self.coordinator.completed_t_counter = 0 + self.coordinator.completed_pending_requests = [] + self.coordinator.snapshot_reports.clear() self.coordinator.prev_completed_snapshot_id = -1 # All workers will effectively need to rebuild their snapshot IDs self.coordinator.worker_snapshot_ids = {worker.worker_id: -1 for worker in participating_workers} diff --git a/coordinator/migration_metadata.py b/coordinator/migration_metadata.py index cc3e8e40..4fc45a74 100644 --- a/coordinator/migration_metadata.py +++ b/coordinator/migration_metadata.py @@ -17,6 +17,9 @@ def __init__(self, n_workers: int) -> None: self.output_offsets: dict[OperatorPartition, int] = {} self.epoch_counter: int = -1 self.t_counter: int = -1 + # Every worker's requests behind its stop offsets that had not committed; + # the rebuilt protocols sequence them (worker/sequencer/pending_requests.py). + self.pending_requests: list = [] # Latest epoch in which a worker's last migrated key landed self.transfer_done_epoch: int = -1 @@ -35,10 +38,12 @@ async def init_done( t_counter: int, input_offsets: dict[OperatorPartition, int], output_offsets: dict[OperatorPartition, int], + pending_requests: list | None = None, ) -> bool: """Collect epoch/offset counters from each worker during Phase B.""" async with self.lock: self.sync_sum[MessageType.MigrationInitDone] += 1 + self.pending_requests.extend(pending_requests or []) for operator_partition, pio in input_offsets.items(): self.input_offsets[operator_partition] = max( pio, @@ -73,5 +78,6 @@ async def cleanup(self, msg_type: MessageType) -> None: self.t_counter = -1 self.input_offsets.clear() self.output_offsets.clear() + self.pending_requests = [] elif msg_type == MessageType.MigrationDone: self.transfer_done_epoch = -1 diff --git a/tests/unit/coordinator/test_coordinator_metadata.py b/tests/unit/coordinator/test_coordinator_metadata.py index 84bccd7a..d9d14578 100644 --- a/tests/unit/coordinator/test_coordinator_metadata.py +++ b/tests/unit/coordinator/test_coordinator_metadata.py @@ -7,6 +7,7 @@ from styx.common.exceptions import NotAStateflowGraphError from styx.common.local_state_backends import LocalStateBackend from styx.common.operator import Operator +from styx.common.serialization import zstd_msgpack_deserialization from styx.common.stateflow_graph import StateflowGraph from coordinator.coordinator_metadata import Coordinator @@ -187,11 +188,61 @@ def test_updates_snapshot_id(self): pool=pool, ) assert c.worker_snapshot_ids[0] == 1 - assert c.completed_input_offsets[("users", 0)] == 10 - assert c.completed_out_offsets[("users", 0)] == 5 + # Snapshot 1 is not complete on worker 1 yet: its offsets wait for it. + assert c.completed_input_offsets == {} + + c.register_snapshot( + worker_id=1, + snapshot_id=1, + partial_input_offsets={("users", 1): 7}, + partial_output_offsets={("users", 1): 3}, + epoch_counter=1, + t_counter=100, + pool=pool, + ) + assert c.completed_input_offsets == {("users", 0): 10, ("users", 1): 7} + assert c.completed_out_offsets == {("users", 0): 5, ("users", 1): 3} assert c.completed_epoch_counter == 1 assert c.completed_t_counter == 100 + def test_sequencer_file_never_takes_a_later_snapshots_report(self): + # Worker 0 reports snapshot 2 before worker 1 reports snapshot 1. Merging + # it into snapshot 1's file would pair snapshot 1's state with offsets past + # requests it does not contain. + c = _coordinator() + c.worker_snapshot_ids = {0: 0, 1: 0} + c.prev_completed_snapshot_id = 0 + pool = MagicMock() + + def report(worker_id, snapshot_id, offset, pending): + c.register_snapshot( + worker_id=worker_id, + snapshot_id=snapshot_id, + partial_input_offsets={("users", worker_id): offset}, + partial_output_offsets={}, + epoch_counter=snapshot_id, + t_counter=snapshot_id, + pool=pool, + pending_requests=pending, + ) + + report(0, 1, 10, []) + report(0, 2, 20, [["rq2"]]) + report(1, 1, 11, [["rq1"]]) + + assert c.completed_input_offsets == {("users", 0): 10, ("users", 1): 11} + assert c.completed_epoch_counter == 1 + assert c.completed_pending_requests == [["rq1"]] + written = zstd_msgpack_deserialization(c.s3_client.put_object.call_args.kwargs["Body"]) + assert written[0] == {("users", 0): 10, ("users", 1): 11} + assert [list(p) for p in written[5]] == [["rq1"]] + assert 2 in c.snapshot_reports + + report(1, 2, 21, []) + assert c.completed_input_offsets == {("users", 0): 20, ("users", 1): 21} + assert c.completed_pending_requests == [["rq2"]] + assert c.snapshot_reports == {} + def test_partial_offsets_take_max(self): c = _coordinator() c.worker_snapshot_ids = {0: 0, 1: 0} diff --git a/tests/unit/worker/test_async_snapshots.py b/tests/unit/worker/test_async_snapshots.py index 7877a968..dee5d00b 100644 --- a/tests/unit/worker/test_async_snapshots.py +++ b/tests/unit/worker/test_async_snapshots.py @@ -91,10 +91,20 @@ def test_reports_once_every_file_is_stored(self, mock_socket): msg = mock_socket.socket.return_value.send.call_args.args[0] from styx.common.tcp_networking import NetworkingManager - (_w, _sid, _start, _end, _in, _out, epoch, t_counter, size, checkpoint) = NetworkingManager.decode_message( - msg[8:] + (_w, _sid, _start, _end, _in, _out, epoch, t_counter, size, checkpoint, pending) = ( + NetworkingManager.decode_message(msg[8:]) ) assert (epoch, t_counter, size, checkpoint) == (10, 100, 100, False) + assert list(pending) == [] + + def test_reports_pending_requests(self, mock_socket): + from styx.common.tcp_networking import NetworkingManager + + s = _snap(snapshot_id=5) + pending = [[b"rq", "k", "op", 1, "f", [], 40, 0]] + s.start_snapshotting({("op", 0): 41}, {}, 10, 100, n_files=0, pending_requests=pending) + msg = mock_socket.socket.return_value.send.call_args.args[0] + assert [list(p) for p in NetworkingManager.decode_message(msg[8:])[10]] == pending def test_overlapping_snapshots_keep_their_own_ids_and_metadata(self, mock_socket): s = _snap(snapshot_id=1) @@ -155,7 +165,8 @@ def test_puts_object(self, mock_mk): class TestRetrieveSnapshot: def test_snapshot_id_negative_one_returns_empty(self): s = _snap() - data, tp_off, tp_out_off, epoch, t_counter, migration_blob = s.retrieve_snapshot(-1, []) + data, tp_off, tp_out_off, epoch, t_counter, migration_blob, pending = s.retrieve_snapshot(-1, []) + assert pending == [] assert data == {} assert tp_off == {} assert tp_out_off == {} @@ -174,7 +185,7 @@ def test_retrieve_with_no_files(self, mock_mk): mock_s3.get_paginator.return_value = paginator s = _snap() - data, _tp_off, _tp_out_off, _epoch, _t_counter, _migration_blob = s.retrieve_snapshot( + data, _tp_off, _tp_out_off, _epoch, _t_counter, _migration_blob, _pending = s.retrieve_snapshot( 0, [("users", 0)], ) @@ -352,12 +363,35 @@ def test_5_element_tuple_with_migration_blob(self, mock_mk): blob, ) - tp_off, tp_out, epoch, t_counter, migration_blob = s._load_sequencer_state(mock_s3, 1) + tp_off, tp_out, epoch, t_counter, migration_blob, pending = s._load_sequencer_state(mock_s3, 1) assert tp_off == {("users", 0): 10} assert tp_out == {("users", 0): 5} assert epoch == 3 assert t_counter == 42 assert migration_blob == blob + assert pending == [] + + @patch("worker.fault_tolerance.async_snapshots._get_s3_client") + def test_6_element_tuple_with_pending_requests(self, mock_mk): + """Sequencer file with the requests behind its offsets that had not committed.""" + mock_s3 = MagicMock() + mock_mk.return_value = mock_s3 + + s = _snap() + paginator = MagicMock() + paginator.paginate.return_value = [{"Contents": [{"Key": "sequencer/1.bin"}, {"Key": "sequencer/2.bin"}]}] + mock_s3.get_paginator.return_value = paginator + + files = { + "sequencer/1.bin": ({("users", 0): 9}, {}, 2, 40, None, [[b"old", "k", "users", 0, "f", [], 8, 0]]), + "sequencer/2.bin": ({("users", 0): 10}, {}, 3, 42, None, [[b"rq", "k", "users", 0, "f", [], 9, 0]]), + } + s._get_zstd_msgpack = lambda _s3_client, key: files[key] + + tp_off, _tp_out, epoch, _t, migration_blob, pending = s._load_sequencer_state(mock_s3, 2) + assert (tp_off, epoch, migration_blob) == ({("users", 0): 10}, 3, None) + # Only the recovered snapshot's own pending requests. + assert pending == [[b"rq", "k", "users", 0, "f", [], 9, 0]] @patch("worker.fault_tolerance.async_snapshots._get_s3_client") def test_4_element_tuple_legacy_format(self, mock_mk): @@ -377,12 +411,13 @@ def test_4_element_tuple_legacy_format(self, mock_mk): 42, ) - tp_off, tp_out, epoch, t_counter, migration_blob = s._load_sequencer_state(mock_s3, 1) + tp_off, tp_out, epoch, t_counter, migration_blob, pending = s._load_sequencer_state(mock_s3, 1) assert tp_off == {("users", 0): 10} assert tp_out == {("users", 0): 5} assert epoch == 3 assert t_counter == 42 assert migration_blob is None + assert pending == [] @patch("worker.fault_tolerance.async_snapshots._get_s3_client") def test_5_element_with_none_migration_blob(self, mock_mk): @@ -403,7 +438,7 @@ def test_5_element_with_none_migration_blob(self, mock_mk): None, ) - *_, migration_blob = s._load_sequencer_state(mock_s3, 1) + *_, migration_blob, _pending = s._load_sequencer_state(mock_s3, 1) assert migration_blob is None @patch("worker.fault_tolerance.async_snapshots._get_s3_client") @@ -417,12 +452,13 @@ def test_no_sequencer_files_returns_defaults(self, mock_mk): paginator.paginate.return_value = [{"Contents": []}] mock_s3.get_paginator.return_value = paginator - tp_off, tp_out, epoch, t_counter, migration_blob = s._load_sequencer_state(mock_s3, 1) + tp_off, tp_out, epoch, t_counter, migration_blob, pending = s._load_sequencer_state(mock_s3, 1) assert tp_off == {} assert tp_out == {} assert epoch == 0 assert t_counter == 0 assert migration_blob is None + assert pending == [] @patch("worker.fault_tolerance.async_snapshots.socket") @@ -432,4 +468,4 @@ def test_checkpoint_flag_is_reported(mock_socket): s = _snap(snapshot_id=7) s.start_snapshotting({}, {}, 1, 1, n_files=0, is_migration_checkpoint=True) msg = mock_socket.socket.return_value.send.call_args.args[0] - assert NetworkingManager.decode_message(msg[8:])[-1] is True + assert NetworkingManager.decode_message(msg[8:])[9] is True diff --git a/tests/unit/worker/test_pending_requests.py b/tests/unit/worker/test_pending_requests.py new file mode 100644 index 00000000..2cec8a4d --- /dev/null +++ b/tests/unit/worker/test_pending_requests.py @@ -0,0 +1,150 @@ +"""Requests behind their ingress partition's recorded offset that have not committed. + +An epoch's ingress offsets advance past all its requests, the ones rescheduled to a +later epoch included. Those are not consumed from Kafka again, so a snapshot or a +migration stop carries them, and the protocol that takes over sequences them before +anything it consumes. +""" + +from types import SimpleNamespace + +from styx.common.operator import Operator +from styx.common.run_func_payload import RunFuncPayload +from styx.common.serialization import msgpack_deserialization, msgpack_serialization + +from worker.ingress.styx_kafka_ingress import StyxKafkaIngress +from worker.operator_state.aria.in_memory_state import InMemoryOperatorState +from worker.sequencer import pending_requests +from worker.sequencer.sequencer import Sequencer +from worker.transactional_protocols.aria import AriaProtocol + +OP = "op" + + +def _payload(offset: int, key: str = "k", partition: int = 0, ingress_partition: int = 0) -> RunFuncPayload: + return RunFuncPayload( + request_id=f"rq{offset}".encode(), + key=key, + operator_name=OP, + partition=partition, + function_name="f", + params=(1, "a"), + kafka_offset=offset, + kafka_ingress_partition=ingress_partition, + ) + + +def _protocol(offsets: dict | None = None) -> AriaProtocol: + """A protocol with a real sequencer and ingress over two partitions of one worker.""" + state = InMemoryOperatorState({(OP, 0), (OP, 1)}) + operators = {(OP, p): Operator(OP, n_partitions=2) for p in (0, 1)} + proto = AriaProtocol.__new__(AriaProtocol) + proto.id = 1 + proto.local_state = state + proto.registered_operators = operators + proto.sequencer = Sequencer() + proto.sequencer.set_sequencer_id([], 1) + proto.topic_partition_offsets = dict(offsets or {}) + proto.ingress = StyxKafkaIngress( + networking=SimpleNamespace(), + sequencer=proto.sequencer, + state=state, + registered_operators=operators, + worker_id=1, + kafka_url="unused", + epoch_interval_ms=1, + sequence_max_size=100, + ) + proto._restored_requests = [] + return proto + + +def _run_epoch(proto: AriaProtocol, rescheduled_offsets: set[int]) -> None: + """The offset bookkeeping of an epoch whose requests at `rescheduled_offsets` + are rescheduled (fallback read/write-set change or concurrency abort).""" + sequence = proto.sequencer.get_epoch() + rescheduled = {item.t_id for item in sequence if item.payload.kafka_offset in rescheduled_offsets} + proto._advance_offsets(sequence) + proto.sequencer.increment_epoch(max_t_counter=proto.sequencer.t_counter, t_ids_to_reschedule=rescheduled) + + +class TestPendingAtTheCut: + def test_rescheduled_request_is_behind_the_offset_and_pending(self): + proto = _protocol({(OP, 0): 39}) + proto.sequencer.sequence(_payload(40)) + proto.sequencer.sequence(_payload(41)) + + _run_epoch(proto, rescheduled_offsets={40}) + + # The offset passes 40, which only runs next epoch: the cut must carry it. + assert proto.topic_partition_offsets[(OP, 0)] == 41 + assert [fields[6] for fields in proto.pending_requests()] == [40] + + def test_requests_consumed_after_the_offset_are_not_pending(self): + proto = _protocol({(OP, 0): 39}) + proto.sequencer.sequence(_payload(40)) + _run_epoch(proto, rescheduled_offsets=set()) + # Consumed after the epoch: Kafka still has it past the recorded offset. + proto.sequencer.sequence(_payload(41)) + + assert proto.topic_partition_offsets[(OP, 0)] == 40 + assert proto.pending_requests() == [] + + def test_committed_epoch_leaves_nothing_pending(self): + proto = _protocol({(OP, 0): 39}) + proto.sequencer.sequence(_payload(40)) + _run_epoch(proto, rescheduled_offsets=set()) + assert proto.pending_requests() == [] + + +class TestRestoring: + def test_restored_request_runs_once_before_newly_consumed_ones(self): + # A cut at offset 41 with 40 rescheduled; the next protocol consumes from 42. + before = _protocol({(OP, 0): 39}) + before.sequencer.sequence(_payload(40)) + before.sequencer.sequence(_payload(41)) + _run_epoch(before, rescheduled_offsets={40}) + offsets = dict(before.topic_partition_offsets) + carried = msgpack_deserialization(msgpack_serialization(before.pending_requests())) + + after = _protocol(offsets) + after._restored_requests = pending_requests.consumed_by(carried, after.registered_operators) + after._sequence_restored_requests() + after.sequencer.sequence(_payload(offsets[(OP, 0)] + 1)) + + sequence = after.sequencer.get_epoch() + assert [item.payload.kafka_offset for item in sequence] == [40, 42] + restored = sequence[0].payload + assert (restored.request_id, restored.params, restored.kafka_ingress_partition) == (b"rq40", (1, "a"), 0) + assert after._restored_requests == [] + + def test_restored_request_goes_to_the_partition_holding_its_key_now(self): + proto = _protocol() + key = next( + f"key{i}" for i in range(1000) if proto.registered_operators[(OP, 0)].which_partition(f"key{i}") == 1 + ) + proto.local_state.batch_insert({key: "v"}, OP, 1) + # Sequenced at partition 0 before a migration moved the key to partition 1. + proto._restored_requests = [pending_requests.encode(_payload(40, key=key, partition=0))] + + proto._sequence_restored_requests() + + (item,) = proto.sequencer.get_epoch() + assert (item.payload.partition, item.payload.kafka_ingress_partition) == (1, 0) + + +class TestConsumedBy: + def test_keeps_this_workers_ingress_partitions_in_offset_order_once(self): + mine, other = (OP, 0), (OP, 1) + carried = [ + pending_requests.encode(_payload(42)), + pending_requests.encode(_payload(40)), + pending_requests.encode(_payload(7, ingress_partition=1)), + # The same request from the checkpoint and from MigrationDone. + pending_requests.encode(_payload(40)), + ] + + kept = pending_requests.consumed_by(carried, {mine}) + + assert [fields[6] for fields in kept] == [40, 42] + assert [fields[6] for fields in pending_requests.consumed_by(carried, {other})] == [7] diff --git a/worker/async_snapshotting.py b/worker/async_snapshotting.py index 1a1cc1df..c3969063 100644 --- a/worker/async_snapshotting.py +++ b/worker/async_snapshotting.py @@ -84,6 +84,7 @@ def take_snapshot(self, metadata: tuple) -> None: *rest, ) = metadata is_migration_checkpoint = bool(rest and rest[0]) + pending = list(rest[1]) if len(rest) > 1 else [] logging.warning(f"ASYNC_SN | Starting Snapshot at epoch: {epoch_counter}") snapshot_id = self.async_snapshots.start_snapshotting( topic_partition_offsets, @@ -92,6 +93,7 @@ def take_snapshot(self, metadata: tuple) -> None: t_counter, n_files=len(self.delta_maps), is_migration_checkpoint=is_migration_checkpoint, + pending_requests=pending, ) for operator_partition, delta_map in self.delta_maps.items(): operator_name, partition = operator_partition diff --git a/worker/fault_tolerance/async_snapshots.py b/worker/fault_tolerance/async_snapshots.py index dfda3301..ad8f8026 100644 --- a/worker/fault_tolerance/async_snapshots.py +++ b/worker/fault_tolerance/async_snapshots.py @@ -1,4 +1,4 @@ -from dataclasses import dataclass +from dataclasses import dataclass, field import functools import os import socket @@ -31,7 +31,10 @@ SNAPSHOT_BUCKET_NAME: str = os.getenv("SNAPSHOT_BUCKET_NAME", "styx-snapshots") SEQUENCER_PREFIX = "sequencer/" -# Number of elements in sequencer tuple (with/without migration checkpoint blob) +# Number of elements in sequencer tuple: offsets, output offsets, epoch, t_counter, +# then the migration checkpoint blob, then the pending requests. Older files lack +# the trailing elements. +_SEQUENCER_TUPLE_LEN_WITH_PENDING = 6 _SEQUENCER_TUPLE_LEN_WITH_MIGRATION = 5 _SEQUENCER_TUPLE_LEN_LEGACY = 4 @@ -67,6 +70,8 @@ class _SnapshotInProgress: t_counter: int files_left: int is_migration_checkpoint: bool = False + # Requests behind input_offsets that had not committed (worker/sequencer/pending_requests.py). + pending_requests: list = field(default_factory=list) size: int = 0 failed: bool = False @@ -102,6 +107,7 @@ def start_snapshotting( n_files: int, *, is_migration_checkpoint: bool = False, + pending_requests: list | None = None, ) -> int: """Start a snapshot of n_files files and return its id. @@ -118,6 +124,7 @@ def start_snapshotting( t_counter=current_t_counter, files_left=n_files, is_migration_checkpoint=is_migration_checkpoint, + pending_requests=pending_requests or [], ) # A snapshot with no files is complete already. self._report_completed() @@ -165,6 +172,7 @@ def _send_snapshot_id(self, snapshot_id: int, snapshot: _SnapshotInProgress) -> snapshot.t_counter, snapshot.size, snapshot.is_migration_checkpoint, + snapshot.pending_requests, ), msg_type=MessageType.SnapID, serializer=Serializer.MSGPACK, @@ -194,15 +202,19 @@ def retrieve_snapshot( int, int, bytes | None, + list, ]: self.snapshot_id = snapshot_id + 1 if snapshot_id == -1: - return {}, {}, {}, 0, 0, None + return {}, {}, {}, 0, 0, None, [] s3 = _get_s3_client() data = self._load_operator_state(s3, snapshot_id, registered_operators) - tp_offsets, tp_out_offsets, epoch, t_counter, migration_blob = self._load_sequencer_state(s3, snapshot_id) - return data, tp_offsets, tp_out_offsets, epoch, t_counter, migration_blob + tp_offsets, tp_out_offsets, epoch, t_counter, migration_blob, pending = self._load_sequencer_state( + s3, + snapshot_id, + ) + return data, tp_offsets, tp_out_offsets, epoch, t_counter, migration_blob, pending def _load_operator_state( self, @@ -234,16 +246,28 @@ def _load_sequencer_state( self, s3: S3Client, snapshot_id: int, - ) -> tuple[dict[OperatorPartition, int], dict[OperatorPartition, int], int, int, bytes | None]: + ) -> tuple[dict[OperatorPartition, int], dict[OperatorPartition, int], int, int, bytes | None, list]: topic_partition_offsets: dict[OperatorPartition, int] = {} topic_partition_output_offsets: dict[OperatorPartition, int] = {} epoch = 0 t_counter = 0 migration_blob: bytes | None = None + pending: list = [] for _, key in self._iter_snapshot_files(s3, SEQUENCER_PREFIX, snapshot_id): loaded = self._get_zstd_msgpack(s3, key) - if isinstance(loaded, (tuple, list)) and len(loaded) == _SEQUENCER_TUPLE_LEN_WITH_MIGRATION: + pending = [] + if isinstance(loaded, (tuple, list)) and len(loaded) == _SEQUENCER_TUPLE_LEN_WITH_PENDING: + ( + topic_partition_offsets, + topic_partition_output_offsets, + epoch, + t_counter, + migration_blob, + pending, + ) = loaded + pending = list(pending) + elif isinstance(loaded, (tuple, list)) and len(loaded) == _SEQUENCER_TUPLE_LEN_WITH_MIGRATION: ( topic_partition_offsets, topic_partition_output_offsets, @@ -260,7 +284,7 @@ def _load_sequencer_state( ) = loaded migration_blob = None - return topic_partition_offsets, topic_partition_output_offsets, epoch, t_counter, migration_blob + return topic_partition_offsets, topic_partition_output_offsets, epoch, t_counter, migration_blob, pending def _iter_snapshot_files( self, diff --git a/worker/ingress/styx_kafka_ingress.py b/worker/ingress/styx_kafka_ingress.py index 2a9ac908..c31a778f 100644 --- a/worker/ingress/styx_kafka_ingress.py +++ b/worker/ingress/styx_kafka_ingress.py @@ -102,24 +102,34 @@ def handle_message_from_kafka(self, msg: ConsumerRecord) -> None: params=params, kafka_ingress_partition=msg.partition, ) - if key is None or self.state.exists(key, operator_name, partition): - # Message received in the correct partition (Normal operation) - self.sequencer.sequence(run_func_payload) - elif ( - true_partition := self.registered_operators[(operator_name, msg.partition)].which_partition(key) - ) == partition: - # Message received in the correct partition, but it was an insert operation (didn't exist in the state) - self.sequencer.sequence(run_func_payload) - else: - # In flight message during migration: the key now belongs to another - # partition. It still runs in this worker's epoch, as a remote call - # if that partition is on a peer (AriaProtocol.run_function), so the - # ingress partition's offset only ever advances here, in order. - run_func_payload.partition = true_partition - self.sequencer.sequence(run_func_payload) + self.route_and_sequence(run_func_payload) else: logging.error(f"Invalid message type: {message_type} passed to KAFKA") + def route_and_sequence(self, run_func_payload: RunFuncPayload) -> None: + """Sequence a request consumed from its ingress partition, at the partition + that holds its key now.""" + key = run_func_payload.key + operator_name = run_func_payload.operator_name + partition = run_func_payload.partition + if key is None or self.state.exists(key, operator_name, partition): + # Message received in the correct partition (Normal operation) + self.sequencer.sequence(run_func_payload) + elif ( + true_partition := self.registered_operators[ + (operator_name, run_func_payload.kafka_ingress_partition) + ].which_partition(key) + ) == partition: + # Message received in the correct partition, but it was an insert operation (didn't exist in the state) + self.sequencer.sequence(run_func_payload) + else: + # In flight message during migration: the key now belongs to another + # partition. It still runs in this worker's epoch, as a remote call + # if that partition is on a peer (AriaProtocol.run_function), so the + # ingress partition's offset only ever advances here, in order. + run_func_payload.partition = true_partition + self.sequencer.sequence(run_func_payload) + async def start_kafka_consumer( self, topic_partitions: list[TopicPartition], diff --git a/worker/sequencer/pending_requests.py b/worker/sequencer/pending_requests.py new file mode 100644 index 00000000..d1c45677 --- /dev/null +++ b/worker/sequencer/pending_requests.py @@ -0,0 +1,81 @@ +"""Requests already behind their ingress partition's recorded offset that have not +committed yet. + +The ingress offset of an epoch advances past every request in it, including the +ones rescheduled to a later epoch (a fallback read/write-set change, or a +concurrency abort without fallback). Those are not consumed from Kafka again, so +a snapshot or a migration stop cut at that offset carries them, and the protocol +that takes over sequences them before it consumes anything. + +They travel as plain lists (msgpack) so the snapshot process and the coordinator +can pass them on without knowing the payload type. +""" + +from typing import TYPE_CHECKING + +from styx.common.run_func_payload import RunFuncPayload + +if TYPE_CHECKING: + from collections.abc import Iterable + + from styx.common.run_func_payload import SequencedItem + from styx.common.types import OperatorPartition + +# request_id, key, operator_name, partition, function_name, params, kafka_offset, +# kafka_ingress_partition +PendingRequest = list + + +def encode(payload: RunFuncPayload) -> PendingRequest: + return [ + payload.request_id, + payload.key, + payload.operator_name, + payload.partition, + payload.function_name, + payload.params, + payload.kafka_offset, + payload.kafka_ingress_partition, + ] + + +def decode(fields: PendingRequest) -> RunFuncPayload: + request_id, key, operator_name, partition, function_name, params, kafka_offset, ingress_partition = fields + return RunFuncPayload( + request_id=request_id, + key=key, + operator_name=operator_name, + partition=partition, + function_name=function_name, + params=tuple(params), + kafka_offset=kafka_offset, + kafka_ingress_partition=ingress_partition, + ) + + +def behind_offsets( + items: Iterable[SequencedItem], + ingress_offsets: dict[OperatorPartition, int], +) -> list[PendingRequest]: + """The sequenced requests at or below their ingress partition's recorded offset.""" + behind = [] + for item in items: + payload = item.payload + recorded = ingress_offsets.get((payload.operator_name, payload.kafka_ingress_partition), -1) + if 0 <= payload.kafka_offset <= recorded: + behind.append(encode(payload)) + return behind + + +def consumed_by( + pending: Iterable[PendingRequest], + ingress_partitions: Iterable[OperatorPartition], +) -> list[PendingRequest]: + """The pending requests of the given ingress partitions, in offset order and + without duplicates.""" + ours = set(ingress_partitions) + by_request: dict[bytes, PendingRequest] = {} + for fields in pending: + if (fields[2], fields[7]) in ours: + by_request.setdefault(fields[0], fields) + return sorted(by_request.values(), key=lambda fields: (fields[2], fields[7], fields[6])) diff --git a/worker/transactional_protocols/aria.py b/worker/transactional_protocols/aria.py index 84cca132..36527f14 100644 --- a/worker/transactional_protocols/aria.py +++ b/worker/transactional_protocols/aria.py @@ -21,6 +21,7 @@ from worker.operator_state.aria.conflict_detection_types import ( AriaConflictDetectionType, ) +from worker.sequencer import pending_requests from worker.sequencer.sequencer import Sequencer if TYPE_CHECKING: @@ -83,6 +84,7 @@ def __init__( # noqa: PLR0913 egress: StyxKafkaBatchEgress | None = None, kafka_consumer: AIOKafkaConsumer | None = None, incarnation: int = 0, + restored_requests: list[pending_requests.PendingRequest] | None = None, ) -> None: if topic_partition_offsets is None: topic_partition_offsets = {(tp.topic, tp.partition): -1 for tp in topic_partitions} @@ -129,6 +131,9 @@ def __init__( # noqa: PLR0913 ) self.sequencer.set_sequencer_id(list(self.peers.keys()), self.id) self.sequencer.set_wal_values_after_recovery(request_id_to_t_id_map) + # Requests behind the ingress offsets this protocol starts from that had not + # committed (pending_requests); sequenced before anything is consumed. + self._restored_requests: list[pending_requests.PendingRequest] = restored_requests or [] self.ingress: StyxKafkaIngress = StyxKafkaIngress( networking=self.networking, @@ -355,6 +360,8 @@ async def take_snapshot(self) -> None: self.egress.topic_partition_output_offsets, self.sequencer.epoch_counter, self.sequencer.t_counter, + False, # not the migration checkpoint + self.pending_requests(), ), msg_type=MessageType.SnapTakeSnapshot, serializer=Serializer.MSGPACK, @@ -362,12 +369,29 @@ async def take_snapshot(self) -> None: self.snapshot_marker_received = False async def communication_protocol(self) -> None: + self._sequence_restored_requests() await self.ingress.start(self.topic_partitions, self.topic_partition_offsets) logging.warning("Ingress started") await self.egress.start(self.id) logging.warning("Egress started") await self.started.wait() + def _sequence_restored_requests(self) -> None: + if not self._restored_requests: + return + for fields in self._restored_requests: + self.ingress.route_and_sequence(pending_requests.decode(fields)) + logging.warning( + f"{self.id} ||| Sequenced {len(self._restored_requests)} requests rescheduled before the cut", + ) + self._restored_requests = [] + + def pending_requests(self) -> list[pending_requests.PendingRequest]: + """The sequenced requests behind the recorded ingress offsets: rescheduled + ones that have not committed. A cut at these offsets must carry them, since + they are not consumed again.""" + return pending_requests.behind_offsets(self.sequencer.distributed_log, self.topic_partition_offsets) + async def protocol_tcp_controller(self, data: bytes) -> None: """Legacy entry point retained for symmetry with the rest of the codebase. @@ -1134,10 +1158,9 @@ async def _maybe_run_fallback(self) -> float: return abort_rate def _advance_offsets(self, sequence: list[SequencedItem]) -> None: + # Past every request of the epoch, the rescheduled ones too: those are back in + # the sequencer, and a cut at these offsets carries them (pending_requests). for item in sequence: - if item.t_id in self.concurrency_aborts_everywhere: - continue - payload = item.payload tpo_key = (payload.operator_name, payload.kafka_ingress_partition) diff --git a/worker/worker_service.py b/worker/worker_service.py index a52436a6..4aa7e911 100644 --- a/worker/worker_service.py +++ b/worker/worker_service.py @@ -44,6 +44,7 @@ from worker.fault_tolerance.async_snapshots import AsyncSnapshotsS3 from worker.operator_state.aria.in_memory_state import InMemoryOperatorState from worker.operator_state.stateless import Stateless +from worker.sequencer import pending_requests from worker.transactional_protocols.aria import AriaProtocol from worker.util.container_monitor import ContainerMonitor @@ -148,6 +149,7 @@ def __init__(self, thread_idx: int) -> None: self.m_input_offsets: dict[OperatorPartition, int] = {} self.m_output_offsets: dict[OperatorPartition, int] = {} + self.m_pending_requests: list[pending_requests.PendingRequest] = [] self.m_epoch_counter: int = -1 self.m_t_counter: int = -1 @@ -579,6 +581,7 @@ async def _migration_rebuild_runtime(self) -> None: egress=old_egress if old_egress.is_running else None, kafka_consumer=old_consumer, incarnation=self.incarnation, + restored_requests=pending_requests.consumed_by(self.m_pending_requests, self.registered_operators), ) async def _handle_receive_migration_hashes( @@ -693,6 +696,10 @@ async def _handle_migration_repartitioning_done( f"MIGRATION | PHASE B Hash metadata sent | took: {t_hash_send_end - t_hash_send_start}", ) + # Requests behind the stop offsets that had not committed: the checkpoint + # and the rebuilt protocols carry them, since they are not consumed again. + stop_pending = self.function_execution_protocol.pending_requests() + # 5.5 Trigger a migration checkpoint snapshot. # This captures the state at the exact moment before data transfer begins, # enabling checkpoint-and-resume recovery instead of revert-and-retry. @@ -706,6 +713,7 @@ async def _handle_migration_repartitioning_done( self.function_execution_protocol.sequencer.epoch_counter, self.function_execution_protocol.sequencer.t_counter, True, # the migration checkpoint + stop_pending, ), msg_type=MessageType.SnapTakeSnapshot, serializer=Serializer.MSGPACK, @@ -721,6 +729,7 @@ async def _handle_migration_repartitioning_done( self.function_execution_protocol.sequencer.t_counter, self.function_execution_protocol.topic_partition_offsets, self.function_execution_protocol.egress.topic_partition_output_offsets, + stop_pending, ), msg_type=MessageType.MigrationInitDone, serializer=Serializer.MSGPACK, @@ -768,6 +777,7 @@ async def _handle_migration_done(self, data: bytes, _: MessageType) -> None: self.m_t_counter, self.m_input_offsets, self.m_output_offsets, + self.m_pending_requests, ) = self.networking.decode_message(data) self.migration_completed.set() await asyncio.sleep(0) @@ -821,7 +831,7 @@ async def _handle_init_recovery(self, data: bytes, _: MessageType) -> None: ) await self._send_snap_assigned(snapshot_id=snapshot_id) - (snap_data, tp_offsets, tp_out_offsets, epoch, t_counter, migration_blob) = ( + (snap_data, tp_offsets, tp_out_offsets, epoch, t_counter, migration_blob, snapshot_pending) = ( self.async_snapshots.retrieve_snapshot( snapshot_id, self.registered_operators.keys(), @@ -919,6 +929,9 @@ async def _handle_init_recovery(self, data: bytes, _: MessageType) -> None: restart_after_migration=restart_after_migration, dedup_output_offsets=all_dedup_out_offsets, incarnation=self.incarnation, + # Rescheduled requests behind the snapshot's offsets, possibly from a + # worker that is gone: this worker consumes their ingress partitions now. + restored_requests=pending_requests.consumed_by(snapshot_pending, self.registered_operators), ) self.function_execution_protocol.start() From 657f01c12e714d3debf8881716f3de5454f9f636 Mon Sep 17 00:00:00 2001 From: Kyriakos Psarakis Date: Sun, 25 Oct 2026 13:30:21 +0100 Subject: [PATCH 20/20] Unpack the pending requests where a worker starts from its first snapshot retrieve_snapshot returns them as a seventh value; the startup path and the S3 integration tests still unpacked six. Co-Authored-By: Claude Opus 5.5 --- tests/integration/test_s3_snapshots.py | 6 +++--- worker/worker_service.py | 8 +++++--- 2 files changed, 8 insertions(+), 6 deletions(-) diff --git a/tests/integration/test_s3_snapshots.py b/tests/integration/test_s3_snapshots.py index 6e4b7a3a..9d9a5b07 100644 --- a/tests/integration/test_s3_snapshots.py +++ b/tests/integration/test_s3_snapshots.py @@ -38,7 +38,7 @@ def test_retrieve_snapshot_negative_id_returns_empty(self, s3_bucket): from worker.fault_tolerance.async_snapshots import AsyncSnapshotsS3 snap = AsyncSnapshotsS3(worker_id=0, n_assigned_partitions=1) - data, off_in, off_out, ep, tc, _migration_blob = snap.retrieve_snapshot(-1, [("any_op", 0)]) + data, off_in, off_out, ep, tc, _migration_blob, _pending = snap.retrieve_snapshot(-1, [("any_op", 0)]) assert data == {} assert off_in == {} assert off_out == {} @@ -68,7 +68,7 @@ def test_store_and_retrieve_round_trip(self, s3_bucket): # Retrieve snap = AsyncSnapshotsS3(worker_id=0, n_assigned_partitions=1) - data, _off_in, _off_out, ep, tc, _migration_blob = snap.retrieve_snapshot( + data, _off_in, _off_out, ep, tc, _migration_blob, _pending = snap.retrieve_snapshot( snapshot_id, [(operator_name, partition)] ) assert data[(operator_name, partition)] == state_data @@ -103,7 +103,7 @@ def test_retrieve_merges_incremental_snapshots(self, s3_bucket): ) snap = AsyncSnapshotsS3(worker_id=0, n_assigned_partitions=1) - data, _off_in, _off_out, ep, tc, _migration_blob = snap.retrieve_snapshot(2, [(op, partition)]) + data, _off_in, _off_out, ep, tc, _migration_blob, _pending = snap.retrieve_snapshot(2, [(op, partition)]) # b should be overwritten by snapshot 2 assert data[(op, partition)]["a"] == 1 diff --git a/worker/worker_service.py b/worker/worker_service.py index 4aa7e911..9fce304f 100644 --- a/worker/worker_service.py +++ b/worker/worker_service.py @@ -339,9 +339,11 @@ async def _handle_receive_execution_plan(self, data: bytes, _: MessageType) -> N await self._send_snap_assigned(snapshot_id=snapshot_id) logging.warning("Retrieving Snapshot from S3") - (snap_data, _in_off, _out_off, epoch, t_counter, _migration_blob) = self.async_snapshots.retrieve_snapshot( - snapshot_id, - self.registered_operators.keys(), + (snap_data, _in_off, _out_off, epoch, t_counter, _migration_blob, _pending) = ( + self.async_snapshots.retrieve_snapshot( + snapshot_id, + self.registered_operators.keys(), + ) ) self.attach_state_to_operators_after_snapshot(snap_data) logging.warning("State attached to operators")