diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 15690d63..2ba15cd8 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 @@ -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 @@ -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/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_metadata.py b/coordinator/coordinator_metadata.py index f0ed5c37..3e77cf59 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 @@ -48,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 @@ -58,6 +72,31 @@ 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 + + # 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 + 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 # >= 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: self.kafka_metadata_producer = AIOKafkaProducer( bootstrap_servers=[KAFKA_URL], @@ -110,9 +149,13 @@ 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} + 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( @@ -133,6 +176,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, ), @@ -179,23 +223,55 @@ def register_snapshot( epoch_counter: int, t_counter: int, pool: concurrent.futures.ProcessPoolExecutor, + *, + is_migration_checkpoint: bool = False, + pending_requests: list | None = None, ) -> None: self.worker_snapshot_ids[worker_id] = 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 + 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}") + 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} " + 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 + 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 {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}") + 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( ( @@ -203,6 +279,8 @@ def register_snapshot( self.completed_out_offsets, 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( @@ -219,6 +297,33 @@ 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 ( + 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()) @@ -229,6 +334,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.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()} })", + ) for _, operator in iter(new_stateflow_graph): for partition in range(self.max_operator_parallelism): operator_copy = deepcopy(operator) @@ -239,6 +348,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.debug("update_stateflow_graph | _pending_graph set (before send)") + worker_assignments = self.worker_pool.get_worker_assignments() tasks = [ self.networking.send_message( @@ -256,14 +374,83 @@ 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.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. + + 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.debug("finalize_graph_update | no-op (_pending_graph is None)") + return + 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 + # 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, + ), + ) + self._pending_kafka_write.add_done_callback(_log_graph_publish_failure) + + def set_migration_checkpoint( + self, + new_graph: StateflowGraph, + operator_partition_locations: dict, + from_snapshot_id: int | None = None, + ) -> None: + """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, + "operator_partition_locations": operator_partition_locations, + }, ) + 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 + 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 + 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 a414585c..c299fd04 100755 --- a/coordinator/coordinator_service.py +++ b/coordinator/coordinator_service.py @@ -22,10 +22,15 @@ 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 +from worker_pool import HEARTBEAT_LIMIT from coordinator.migration_metadata import MigrationMetadata @@ -46,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 @@ -108,6 +116,11 @@ 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] = [] + # 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 @@ -192,6 +205,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 +299,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 +346,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.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. + 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 +368,31 @@ 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 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() + 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.debug("MigrationRepartitioningDone | calling finalize_migration_repartition") await self.finalize_migration_repartition() await self.migration_metadata.cleanup(mt) + logging.debug("MigrationRepartitioningDone | done") async def _handle_migration_init_done( self, @@ -325,28 +400,64 @@ 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]: - epoch_counter, t_counter, input_offsets, output_offsets = self.networking.decode_message(data) + 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)") + return + + 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}") 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.debug("MigrationInitDone | closing protocol connections") await self.protocol_networking.close_all_connections() - await self.finalize_migration() - await self.migration_metadata.cleanup(mt) + 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 + ) + + # 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.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)") + + # Launch background task to wait for the checkpoint snapshot, then finalize + logging.debug("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, @@ -403,6 +514,8 @@ async def _handle_snap_id( epoch_counter, t_counter, sn_size, + is_migration_checkpoint, + pending_requests, ) = self.networking.decode_message(data) snapshot_time = end - start @@ -426,8 +539,28 @@ async def _handle_snap_id( epoch_counter, t_counter, pool, + is_migration_checkpoint=is_migration_checkpoint, + pending_requests=pending_requests, ) + # 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 +612,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.debug(f"PROTOCOL_MSG during migration: {mt.name}") await handler(data) # ------------------------ @@ -587,11 +722,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.debug( + "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.debug("SyncCleanup | stop_gracefully sent successfully") self.aria_metadata.cleanup(epoch_end=True) def _record_epoch_metrics( @@ -656,24 +798,44 @@ 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 mt = MessageType.MigrationDone async with self.networking_locks[mt]: - sync_complete: bool = await self.migration_metadata.set_empty_sync_done(mt) + 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)") + return + (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 - logging.warning("Restarting the snapshotting mechanism") - self.snapshotting_task = asyncio.create_task(self.send_snapshot_marker()) + # 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.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 + # 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 +847,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.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.info("Closing the connection") + logging.info("Protocol connection closing") writer.close() await writer.wait_closed() @@ -752,36 +916,89 @@ 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, + self.migration_metadata.pending_requests, + ), + 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.""" + 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.debug( + 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.debug("checkpoint_finalizer | calling finalize_graph_update") + self.coordinator.finalize_graph_update() + + # Send MigrationDone to workers + logging.debug("checkpoint_finalizer | sending MigrationDone to workers") + await self.finalize_migration() + await self.migration_metadata.cleanup(mt) + logging.debug("checkpoint_finalizer | done") async def finalize_worker_sync( self, @@ -789,17 +1006,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: @@ -832,13 +1044,18 @@ 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() 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} @@ -860,6 +1077,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 +1113,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.debug( + 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.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 + # 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.debug( + 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 +1157,69 @@ 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 + resumed_migration: dict | None = 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}", + ) + # 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.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 + 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) @@ -903,6 +1231,17 @@ 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"], + from_snapshot_id=snap_id, + ) + self.migration_in_progress = True + # 5) Close all protocol connections (workers will reconnect clean) await self.protocol_networking.close_all_connections() @@ -912,6 +1251,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: @@ -926,8 +1280,16 @@ 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) and self.recovery_state == RecoveryState.IDLE: + if ( + 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 " + 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 @@ -938,12 +1300,39 @@ 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(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: logging.error(f"Error during recovery: {e}") finally: 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: @@ -983,6 +1372,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() @@ -995,6 +1401,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/coordinator/migration_metadata.py b/coordinator/migration_metadata.py index 7f41f675..4fc45a74 100644 --- a/coordinator/migration_metadata.py +++ b/coordinator/migration_metadata.py @@ -17,6 +17,11 @@ 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 def check_sum(self, msg_type: MessageType) -> bool: return self.sync_sum[msg_type] == self.n_workers @@ -33,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, @@ -51,6 +58,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 +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/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/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..8e61a23f --- /dev/null +++ b/demo/demo-migration-ft-ycsb/client.py @@ -0,0 +1,270 @@ +"""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 json +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 zlib + +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) + styx.notify_init_data_complete() + 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): + 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) + partitions = 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") + verify_final_state(partitions) + + +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/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/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/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-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/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/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/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/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..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 @@ -45,3 +44,8 @@ class MessageType(IntEnum): UpdateExecutionGraph = 41 RunFunRemoteBatch = 42 AckBatch = 43 + SnapMigrationReassign = 44 + AsyncMigrationAck = 45 + RecoveryMigrationHashes = 46 + RunFunProxy = 47 + ProxyResult = 48 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..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, @@ -403,10 +411,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/e2e/test_e2e_migration_ft_ycsb.py b/tests/e2e/test_e2e_migration_ft_ycsb.py new file mode 100644 index 00000000..76d91694 --- /dev/null +++ b/tests/e2e/test_e2e_migration_ft_ycsb.py @@ -0,0 +1,455 @@ +"""E2E tests for fault tolerance during state migration. + +These tests verify that Styx maintains exactly-once semantics when a worker +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 +- Exactly-once output (no duplicate responses) +- Zero missed messages +""" + +from dataclasses import dataclass +import json +import logging +from pathlib import Path +import shutil + +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 + + # 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: + 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 _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) + + 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 before the post-migration snapshot (resume from the checkpoint) +# --------------------------------------------------------------------------- + + +@pytest.mark.e2e_migration +def test_migration_ft_kill_before_post_migration_snapshot(tmp_path: Path): + """Kill a worker 5 s after the migration starts. + + 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: exactly-once output across a resumed migration. + """ + 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, # 5 s after the migration starts + ) + + 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 after the post-migration snapshot +# --------------------------------------------------------------------------- + + +@pytest.mark.e2e_migration +def test_migration_ft_kill_after_post_migration_snapshot(tmp_path: Path): + """Kill a worker 20 s after the migration starts. + + A snapshot without the checkpoint has completed by then, so recovery is an + ordinary one with the new layout. + + Validates: recovery with the new layout, 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, # 20 s after the migration starts + ) + + 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: several snapshots + with the new layout exist. Recovery is an ordinary one with that 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 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) + + 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..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 = 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,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, _pending = 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, _pending = 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/service_helpers.py b/tests/unit/coordinator/service_helpers.py new file mode 100644 index 00000000..9be87c38 --- /dev/null +++ b/tests/unit/coordinator/service_helpers.py @@ -0,0 +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._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_coordinator_metadata.py b/tests/unit/coordinator/test_coordinator_metadata.py index 61cf7848..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} @@ -286,3 +337,206 @@ 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_checkpoint_snapshot_completes(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, 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() + 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_checkpoint_lifetime.py b/tests/unit/coordinator/test_migration_checkpoint_lifetime.py new file mode 100644 index 00000000..f200384e --- /dev/null +++ b/tests/unit/coordinator/test_migration_checkpoint_lifetime.py @@ -0,0 +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. +""" + +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/tests/unit/coordinator/test_migration_fault_tolerance.py b/tests/unit/coordinator/test_migration_fault_tolerance.py new file mode 100644 index 00000000..30bb5905 --- /dev/null +++ b/tests/unit/coordinator/test_migration_fault_tolerance.py @@ -0,0 +1,497 @@ +"""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 +from tests.unit.coordinator.service_helpers import recover, stub_service + + +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: + """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. + + These drive the real CoordinatorService._perform_recovery. + """ + + async def test_mid_migration_reverts_and_retries(self): + c = _mock_coordinator() + pending = _make_graph(n_partitions=8) + c._pending_graph = pending + svc = stub_service(c) + svc.migration_in_progress = True + + await recover(svc) + + 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 = _make_graph(n_partitions=8) + c.post_migration_snapshot_pending = True + svc = stub_service(c) + + await recover(svc) + + c.revert_worker_pool_to_submitted_graph.assert_called_once() + + async def test_no_migration_is_an_ordinary_recovery(self): + c = _mock_coordinator() + svc = stub_service(c) + + await recover(svc) + + 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.submitted_graph = _make_graph(n_partitions=8) # finalized + svc = stub_service(c) + + await recover(svc) + + 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() + c._pending_graph = _make_graph(n_partitions=8) + c.pre_migration_snapshot_pending = True + c.post_migration_snapshot_pending = True + c._migration_checkpoint_baseline_snap_id = 3 + svc = stub_service(c) + svc.migration_in_progress = True + + await recover(svc) + + 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 + + +# --------------------------------------------------------------------------- +# 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 + + # 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 from here is an ordinary one: TestRecoveryConditionDetection.) + + 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 + + # CRASH during Phase A — no pre-migration snapshot completed + + # 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 + + +# --------------------------------------------------------------------------- +# 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_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/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..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) == 43 + 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(44) 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 @@ -30,6 +31,11 @@ def test_known_values(self): assert MessageType.UpdateExecutionGraph == 41 assert MessageType.RunFunRemoteBatch == 42 assert MessageType.AckBatch == 43 + 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/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_async_snapshots.py b/tests/unit/worker/test_async_snapshots.py index afc77dda..dee5d00b 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,84 @@ 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 +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] -# --------------------------------------------------------------------------- -# snapshot_completed_callback -# --------------------------------------------------------------------------- - -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, 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) + 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] # --------------------------------------------------------------------------- @@ -158,12 +165,14 @@ 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, pending = s.retrieve_snapshot(-1, []) + assert pending == [] 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 +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 = s.retrieve_snapshot( + data, _tp_off, _tp_out_off, _epoch, _t_counter, _migration_blob, _pending = s.retrieve_snapshot( 0, [("users", 0)], ) @@ -267,3 +276,196 @@ 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, 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): + """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, 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): + """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, _pending = 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, 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") +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:])[9] is True diff --git a/tests/unit/worker/test_in_memory_state.py b/tests/unit/worker/test_in_memory_state.py index ca086764..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 @@ -440,11 +476,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 @@ -478,10 +523,28 @@ 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_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" + 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_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..ebadddc8 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,241 @@ 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 + + +# --------------------------------------------------------------------------- +# 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_kafka_ingress.py b/tests/unit/worker/test_kafka_ingress.py index 0e0b3475..b396841b 100644 --- a/tests/unit/worker/test_kafka_ingress.py +++ b/tests/unit/worker/test_kafka_ingress.py @@ -1,7 +1,11 @@ """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 import TopicPartition +from aiokafka.errors import KafkaConnectionError, UnknownTopicOrPartitionError +import pytest from styx.common.message_types import MessageType from styx.common.run_func_payload import RunFuncPayload @@ -64,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 @@ -122,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) @@ -130,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.""" @@ -142,3 +151,342 @@ 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 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) + networking.in_the_same_network.return_value = False + + msg = _make_msg(partition=5) + ingress.handle_message_from_kafka(msg) + + 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.""" + 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 + from pathlib import Path + + # 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" + + for node in handlers: + 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.", + ) + + +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 + + @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/tests/unit/worker/test_migration_fault_tolerance.py b/tests/unit/worker/test_migration_fault_tolerance.py new file mode 100644 index 00000000..abbbc178 --- /dev/null +++ b/tests/unit/worker/test_migration_fault_tolerance.py @@ -0,0 +1,222 @@ +"""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_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 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) + 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 == (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 == (False, 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_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 + 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/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..962b8ed3 --- /dev/null +++ b/tests/unit/worker/test_migration_key_in_flight.py @@ -0,0 +1,70 @@ +"""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, 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, incarnation), + msg_type=MessageType.ReceiveRemoteKey, + serializer=Serializer.MSGPACK, + )[8:] + + +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()}} + w.local_state = InMemoryOperatorState({DST}) + w.local_state.add_remote_keys(DST, {"k1": (2, 0)}) + return w + + +class TestRemoteKeyResponse: + 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", 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", 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) + + 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 new file mode 100644 index 00000000..0881bb8e --- /dev/null +++ b/tests/unit/worker/test_migration_snapshot_consistency.py @@ -0,0 +1,346 @@ +"""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.incarnation = 0 + 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.currently_processing = False + proto.registered_operators = dict.fromkeys(state.operator_partitions) + 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() + + +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) + + 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) + + +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(), close_consumer=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._stall_watchdog_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_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() + + 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_kafka_clients=True) + await proto.stop() + proto.ingress.close_consumer.assert_awaited_once() + 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 + + +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/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/tests/unit/worker/test_proxied_requests.py b/tests/unit/worker/test_proxied_requests.py new file mode 100644 index 00000000..e7956620 --- /dev/null +++ b/tests/unit/worker/test_proxied_requests.py @@ -0,0 +1,307 @@ +"""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 import aria as aria_module +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 + + +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/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/async_snapshotting.py b/worker/async_snapshotting.py index 0ca6878c..c3969063 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() ( @@ -70,24 +81,30 @@ 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]) + pending = list(rest[1]) if len(rest) > 1 else [] 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, + pending_requests=pending, ) 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: @@ -110,6 +127,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..a92c8c09 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,48 +24,111 @@ 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() + + @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 + 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() + # 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"Getting messages sent before recovery: {self.topic_partition_output_offsets}", + 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("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 - # 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 + logging.warning( + f"[DEDUP] Done scanning. Dedup set size: {len(self.messages_sent_before_recovery)}", + ) 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() - await self.kafka_egress_producer.stop() + 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)}", + ) + producer, self.kafka_egress_producer = self.kafka_egress_producer, None + await producer.stop() async def send( self, @@ -77,8 +139,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 +148,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.debug( + 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 +179,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 +190,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.debug( + 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 +216,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.debug(f"[DEDUP] Seeking {tp} to offset {offset}") except (UnknownTopicOrPartitionError, KafkaConnectionError): # fmt: skip await asyncio.sleep(1) logging.warning( @@ -171,38 +242,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, - ) - logging.warning( - f"Reading from output from: {self.topic_partition_output_offsets} + 1 to {current_offsets}", + current_offsets = await kafka_output_consumer.end_offsets(all_topic_partitions) + logging.debug( + f"[DEDUP] End offsets: {current_offsets}", ) - all_partitions_done: dict[TopicPartition, bool] = dict.fromkeys( - output_topic_partitions, - False, - ) - 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.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) 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 +275,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..ad8f8026 100644 --- a/worker/fault_tolerance/async_snapshots.py +++ b/worker/fault_tolerance/async_snapshots.py @@ -1,9 +1,12 @@ +from dataclasses import dataclass, field +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 @@ -28,6 +31,12 @@ SNAPSHOT_BUCKET_NAME: str = os.getenv("SNAPSHOT_BUCKET_NAME", "styx-snapshots") SEQUENCER_PREFIX = "sequencer/" +# 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 # Per-process cached S3 client (avoids re-creating boto3 client on every call) _s3_client: S3Client | None = None @@ -52,6 +61,21 @@ 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 + # 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 + + class AsyncSnapshotsS3(BaseSnapshotter): def __init__( self, @@ -60,16 +84,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 @@ -77,50 +98,89 @@ 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 + 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. + + 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, + pending_requests=pending_requests or [], + ) + # 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 register_size(self, file_size_bytes: int) -> None: - self.total_snapshot_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, + snapshot.pending_requests, + ), + 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: @@ -141,15 +201,20 @@ def retrieve_snapshot( dict[OperatorPartition, int], int, int, + bytes | None, + list, ]: 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, 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, @@ -165,7 +230,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 +246,45 @@ 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, 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): - ( - 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) + 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, + 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, pending def _iter_snapshot_files( self, diff --git a/worker/ingress/styx_kafka_ingress.py b/worker/ingress/styx_kafka_ingress.py index cfbda3cd..c31a778f 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,13 +39,26 @@ 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() 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 +70,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) @@ -81,58 +102,62 @@ 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, 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, - ), - ) + 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], 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. + 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 +185,8 @@ async def start_kafka_consumer( ) continue break + + async def _consume(self, topic_partition_offsets: dict[OperatorPartition, int]) -> None: try: # Consume messages logging.warning( @@ -181,9 +208,7 @@ async def start_kafka_consumer( 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: - await self.kafka_consumer.stop() + if not self._keep_consumer: + await self.close_consumer() diff --git a/worker/operator_state/aria/in_memory_state.py b/worker/operator_state/aria/in_memory_state.py index 12b69059..c4f02063 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] @@ -33,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) @@ -45,33 +66,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 - # 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, @@ -79,13 +104,13 @@ 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), - ) + 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 @@ -110,10 +135,13 @@ 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 c += 1 if not keys: @@ -132,6 +160,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: @@ -209,6 +238,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(): @@ -269,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) 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 8552bc12..36527f14 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 @@ -19,12 +21,13 @@ 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: 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 @@ -45,11 +48,21 @@ 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")) +# 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 +# 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): - def __init__( + def __init__( # noqa: PLR0913 self, worker_id: int, peers: dict[int, tuple[str, int, int]], @@ -66,6 +79,12 @@ 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, + 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} @@ -73,6 +92,8 @@ def __init__( 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 @@ -110,6 +131,9 @@ def __init__( ) 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, @@ -121,15 +145,31 @@ def __init__( sequence_max_size=SEQUENCE_MAX_SIZE, epoch_interval_ms=100, ) - - self.egress: StyxKafkaBatchEgress = StyxKafkaBatchEgress( - output_offsets, - restart_after_recovery or restart_after_migration, - ) + 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. + 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 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 @@ -150,16 +190,27 @@ 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 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() + # 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, 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, @@ -174,14 +225,47 @@ def __init__( 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) + 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() - async def stop(self) -> None: - await self.ingress.stop() - await self.egress.stop() + 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_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(keep_consumer=keep_kafka_clients) + if keep_kafka_clients: + 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() @@ -189,6 +273,8 @@ async def stop(self) -> 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 @@ -205,10 +291,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, @@ -218,7 +329,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, @@ -246,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, @@ -253,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. @@ -348,23 +481,67 @@ 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.) + """ + 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 + 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, @@ -372,10 +549,72 @@ 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(): + # 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)) + + 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") @@ -401,6 +640,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() @@ -466,17 +706,62 @@ 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, 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.local_state.set_batch_data_from_migration(operator_partition, batch) - # Unblock any transactions waiting for these keys via RequestRemoteKey + 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: + 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, + ) + + def _apply_migration_batch(self, operator_partition: OperatorPartition, batch: dict) -> None: op = tuple(operator_partition) - if op in self.networking.wait_remote_key_event: + 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 self.networking.wait_remote_key_event.get(op, {}): - self.networking.wait_remote_key_event[op][key].set() + 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) + 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() + 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() @@ -498,50 +783,69 @@ 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. + + 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] - 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, self.incarnation), + 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. 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}" @@ -553,7 +857,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) @@ -561,35 +872,65 @@ 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") - 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() + graceful = False + 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() + 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") + 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: + # Only a migration stops the protocol gracefully; its rebuilt + # 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() 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() @@ -607,7 +948,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. @@ -622,11 +965,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() @@ -644,6 +991,10 @@ 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() + self._migration_epoch_ended.set() async def _run_epoch_functions_and_chain( self, @@ -659,6 +1010,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], ) @@ -666,12 +1018,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, @@ -803,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) @@ -862,7 +1216,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, @@ -871,12 +1225,15 @@ 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() 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() @@ -911,7 +1268,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: @@ -938,14 +1295,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: @@ -976,7 +1333,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). @@ -1011,7 +1370,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, @@ -1059,5 +1418,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() diff --git a/worker/worker_service.py b/worker/worker_service.py index d6310615..9fce304f 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 @@ -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 @@ -134,9 +135,21 @@ 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 + # 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 # 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() 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 @@ -176,6 +189,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 @@ -265,6 +279,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( @@ -275,6 +290,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,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) = 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") @@ -332,6 +358,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() @@ -367,7 +394,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() + # 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: @@ -381,6 +417,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: @@ -508,8 +548,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() @@ -521,7 +561,11 @@ 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 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, @@ -536,6 +580,10 @@ async def _migration_rebuild_runtime(self) -> None: epoch_counter=self.m_epoch_counter, t_counter=self.m_t_counter, restart_after_migration=True, + 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( @@ -554,10 +602,13 @@ 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( + 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, old_partition, @@ -565,7 +616,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, self.incarnation), msg_type=MessageType.ReceiveRemoteKey, serializer=Serializer.MSGPACK, ) @@ -577,7 +628,14 @@ async def _handle_receive_remote_key( ) -> None: # Lock-free: no awaits. del msg_type - operator_partition, key, 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. + 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. @@ -594,6 +652,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 +661,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 +698,30 @@ 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. + 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, + True, # the migration checkpoint + stop_pending, + ), + 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, @@ -640,6 +731,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, @@ -649,6 +741,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 +758,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 @@ -677,6 +779,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) @@ -684,6 +787,15 @@ 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._reset_recovery_hash_exchange(tag=-1, hashes_out={}) + ( self.id, self.worker_operators, @@ -692,6 +804,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: @@ -707,6 +820,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) @@ -719,16 +833,82 @@ 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, snapshot_pending) = ( + 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.debug( + f"[RECOVERY] W{self.id} registered_operators={list(self.registered_operators.keys())}", + ) + logging.debug( + f"[RECOVERY] W{self.id} raw tp_offsets from snapshot: {tp_offsets}", + ) + logging.debug( + f"[RECOVERY] W{self.id} raw tp_out_offsets from snapshot: {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} + # 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) + # 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}", + ) + logging.debug( + f"[RECOVERY] W{self.id} filtered tp_out_offsets: {tp_out_offsets}", + ) + logging.debug( + 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, 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: + 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=self.incarnation, hashes_out=hashes_out) + restart_after_migration = True + request_id_to_t_id_map = await self.get_sequencer_assignments_before_failure( epoch, ) @@ -747,7 +927,13 @@ 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, + 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() @@ -766,12 +952,88 @@ async def _handle_init_recovery(self, data: bytes, _: MessageType) -> None: ) async def _handle_ready_after_recovery(self, _data: bytes, _: MessageType) -> None: - self.function_execution_protocol.started.set() + 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 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)