[kPsarakis]

Summary

Makes state migration fault tolerant. Before this, a worker crash during a migration reverted to the pre-migration snapshot and retried from scratch. Now migration progress is checkpointed into snapshots: a crash after the checkpoint resumes the transfer with the new layout, a crash before it reverts to the old layout and retries the migration. The branch is rebased onto main and also fixes several correctness and latency problems in the migration path, some of which predate it.

Migration algorithm

A migration starts when a client submits a graph whose operators have different partition counts. The coordinator runs one migration at a time. It forces a snapshot at the next epoch boundary (the revert point if the migration fails before its checkpoint), assigns the new partitions to workers, and sends every worker InitMigration with the new graph, its assignment and the partition locations. Partitions are pre-created up to the maximum parallelism; those at or above an operator's partition count are shadow partitions. The migration then runs in three phases. Steps marked (new) are added by this PR; the rest is on main.

Phase A: rehash, protocol running. Each worker rehashes the keys of its partitions with the new partitioner, in a process pool and without network I/O, and records which keys must move and to which partition. It does not apply the new plan yet. Each worker reports MigrationRepartitioningDone. Once all have, the coordinator sets the checkpoint metadata (new) and asks the protocol to stop at the next epoch boundary.

Phase B: switch over, protocol stopped. Each worker:

1. Waits for the protocol to stop at the epoch boundary, with every response flushed. It is never cancelled mid-epoch (new).
2. Applies the plan (DNS, peers, assignments) and creates its new partitions.
3. Rehashes the keys created during Phase A.
4. Sends each destination the metadata of the keys it will receive (key → source worker and old partition). The destination keeps it as remote_keys, so it knows where to find a key that hasn't arrived yet.
5. Takes the migration checkpoint snapshot (new).
6. Reports MigrationInitDone with its epoch and transaction counters and its ingress and egress offsets.

The coordinator merges the counters and offsets (max per partition). Once the checkpoint is complete on every worker (new), it publishes the new graph, so the other clients switch partitioners (the client that submitted the migration switched when it submitted it), and sends MigrationDone with the merged counters. Each worker then rebuilds the protocol for its new assignment, reusing its Kafka clients (new), and starts it.

Phase C: transfer, protocol running. A background sender on each worker moves its keys to their new partitions in batches, alongside normal processing. A batch to a partition on the same worker is applied in place (new); a batch to a peer stays unacknowledged until the peer has applied it. Each epoch drains the in-flight batches before its processing-done barrier and caps how many keys the sender moves (ASYNC_MIGRATION_KEYS_PER_EPOCH) (new). Transactions don't wait for the transfer:

• A key that hasn't arrived is fetched on demand: moved directly from a local partition, or requested from its source worker (RequestRemoteKey / ReceiveRemoteKey). If the sender already took it into a batch, the transaction waits for that batch (new).
• A request addressed to a partition that doesn't hold its key is rerouted by the consuming worker to the partition that does, as a proxied call if that partition is on a peer (new). This covers requests routed by the old graph and still in Kafka, and requests the submitting client routes by the new graph before the switch-over, which land on shadow partitions.

When a worker has sent every key and every batch is acknowledged, it reports MigrationDone with the epoch its last key landed in. Once every worker has, the migration is finished. Snapshots keep the checkpoint metadata until one is cut after that epoch (new).

Fault tolerance

Checkpoint. In Phase B, after the protocol stops and the hash metadata is exchanged, every worker takes a snapshot flagged as the migration checkpoint. The coordinator attaches the migration metadata (new graph, partition locations) to the sequencer file of that snapshot and of every later one, until a snapshot is cut after the last migrated key landed. It publishes the new graph and sends MigrationDone only once the checkpoint is complete on every worker.

Recovery. On recovery the coordinator reads the latest complete snapshot:

• With the checkpoint metadata: recover with the new layout. Each worker re-derives the keys it still has to send from its snapshot state, and the workers exchange the matching remote_keys (all-to-all, on ReadyAfterRecovery) before the protocol starts. The coordinator keeps tracking the resumed transfer, so a second crash during it resumes again.
• Without it, while a migration was running: revert to the old layout and retry the migration after the first post-recovery snapshot.

Snapshots during the transfer. Migrated keys are recorded as tombstones at the source and as inserts at the destination. A batch never straddles a snapshot cut: each epoch drains in-flight batches (acknowledged by the destination once applied) before its processing-done barrier, and the sender resumes after the cleanup barrier. The flush to the snapshot process swaps the delta map before sending, so writes landing during the flush are not lost.

Exactly-once output. After a recovery, every worker scans all output partitions (a replayed request may now run on a different worker) and suppresses responses already sent. A clean migration stop needs no scan: it stops at an epoch boundary with every response flushed.

Requests for migrated keys. A request consumed for a key that now belongs to a partition on another worker stays in the consuming worker's epoch and runs on the owner as a proxied call. The owner returns the whole outcome (response, abort, chain participants, fallback read/write-set change) in one ProxyResult, which the consuming worker awaits before it can commit, and the response goes to the request's ingress partition. So an ingress partition's offset is only ever advanced by the worker consuming it, in order.

Stale messages. The coordinator bumps an incarnation number on every recovery. Migration batches, remote-key requests/replies, proxied calls and their results, and the recovery hash exchange carry it, and messages from another incarnation are dropped.

Fixes beyond the original change

Correctness:

• Keys could be lost across a snapshot cut: async batches were in flight across it, and the delta flush cleared writes that landed while it awaited.
• After a resume, destinations had no remote_keys, so transactions read keys still at their source as missing.
• A transaction could read a migrating key as missing when the async sender had already taken it into a batch (local move, or a None reply to a remote fetch); it now waits for the batch.
• Keys whose value is None were dropped by the transfer.
• The checkpoint metadata was cleared on MigrationDone, so a snapshot cut before the last key landed but completed afterwards lost it; and any snapshot completing after Phase B began was taken as the checkpoint.
• A snapshot started while the previous one was uploading wrote its files under the same id, and failed uploads were reported as complete. Snapshots now get their id at start and are reported strictly in order.
• Phase B force-stopped a protocol still running after 15 s, possibly mid-commit. A barrier or migration message the coordinator cannot deliver now hands the worker to recovery instead, after one heartbeat limit: by then the heartbeat check has removed a dead worker (and rescheduled its partitions), so only a live one is recovered as a participant. Recovered at once, a dead worker stayed a participant and recovery waited on it forever.
• A destination could receive a batch for a partition it had not created yet.
• Requests consumed for a migrated key were forwarded into the owner's epoch, so two workers advanced one ingress partition's offset, each by max, and forwarded requests could run out of order over the socket pool. A snapshot or the Phase B stop could then record an offset past requests not yet run, which recovery or the rebuilt consumer skipped (the likely cause of 27 missed responses in one CI run of the external-kill test). They now run as proxied calls from the consuming worker.
• A request rescheduled to a later epoch (a fallback read/write-set change, or a concurrency abort with fallback off) was already behind the epoch's ingress offset, and only the sequencer held it. A recovery from a snapshot cut in between, or the protocol rebuilt after a Phase B stop, never ran it. Snapshots, MigrationInitDone and MigrationDone now carry these pending requests, and the protocol that takes over sequences them before consuming anything, routed to the partition that holds the key now (possibly on another worker, after a recovery).
• The coordinator folded each worker's snapshot report into the completed offsets on arrival, so a worker's report of snapshot N+1 arriving before another worker's report of N put N+1's offsets in N's sequencer file, past requests N's state does not contain. Each sequencer file is now written from its own snapshot's reports.
• Recovery started a partition without an offset in the snapshot (e.g. a shadow partition, recovering from the data-load snapshot) at the latest offset, skipping requests written to it before and during the failure.
• Compaction (off by default) merged deltas in S3 listing order, kept a stale first delta, and never dropped tombstones.
• The migration YCSB / TPC-C and scalability demos never sent InitDataComplete, so those E2E tests ran on an empty store.

Latency (the stall after Phase B):

• The rebuilt protocol reuses the egress producer and the ingress consumer instead of creating new Kafka clients, and skips the dedup scan after a clean stop.
• exists() scanned all pending migration keys for every incoming message (about 3 ms per call with 60k keys pending, which kept the event loop busy for seconds during Phase C). It is now O(1).
• Local batches are applied in place instead of being sent over TCP to the same worker, and the transfer is spread over epochs (ASYNC_MIGRATION_KEYS_PER_EPOCH, default 3000) so no single epoch serializes the whole share into its snapshot delta.

Performance

YCSB, 4000 tx/s, 500k keys, 4 → 8 partitions, interleaved runs on the same machine:

	main	this PR
p99 / max latency during the migration	1199–1793 / 1301–1898 ms	540–947 / 620–1081 ms
lowest throughput around the migration	239–556 tx/s	2409–3634 tx/s
steady-state p50 / p99	unchanged	unchanged

The remaining spike is mostly the Phase B stop, which includes waiting for the checkpoint snapshot (about 100 ms).

Diagnostics

A stalled epoch stalls every worker at the next barrier and nothing is logged. Workers now warn every EPOCH_STALL_WARN_S (10 s) while one epoch wait lasts, naming the wait (transactions, chains, WAL write, a barrier, fallback, migration batch acks, snapshot delta, responses to Kafka, snapshot request), the unfinished chains and the pending proxied calls. The coordinator warns when no barrier completes for BARRIER_STALL_WARN_S (10 s) while some workers wait, naming the workers it waits on.

Wire format changes

New message types SnapMigrationReassign (44), AsyncMigrationAck (45), RecoveryMigrationHashes (46), RunFunProxy (47), ProxyResult (48); WrongPartitionRequest (26) is gone. AsyncMigration, RequestRemoteKey, ReceiveRemoteKey, MigrationInitDone, MigrationDone, SnapID, SnapTakeSnapshot and InitRecovery carry extra fields. The sequencer file has a fifth element (the migration metadata) and a sixth (the pending requests); 4- and 5-element files still load.

Testing

• Unit tests for each fix above, including the coordinator recovery paths through the real _perform_recovery.
• A regression test for the forwarding bug runs Kafka records through the real ingress and sequencer and checks that the recorded ingress offset never passes a request that has not run; it fails on the forwarding code. Workers log Proxied calls this epoch: N when they sent any, so the E2E worker logs show whether that path ran. In a local run of the migration E2E suite, 236 of the 239 proxied calls came from the external-kill test: its recovery put shadow partition 5 and old partition 1, which holds partition 5's keys under the old layout, on different workers, so every request the submitting client sent to partition 5 crossed workers. On the forwarding code, that is where two workers advanced one partition's offset.
• E2E (tests/e2e/test_e2e_migration_ft_ycsb.py): baseline migration, a kill before the post-migration snapshot (resume from the checkpoint), a kill after it, a kill well after the migration, and an external kill during Phase A (revert and retry). Besides duplicates and missed responses, every key's stored record is now compared with the loaded dataset at the end.
• Unit tests for the pending requests: a rescheduled request is behind the offset and carried at the cut, it runs once and before newly consumed requests after a restore, it goes to the partition that holds its key now, and a sequencer file never takes a later snapshot's report. The E2E runs did not reschedule a request across a cut, so that path is covered by the unit tests only.
• Which phase a kill lands in depends on timing; the docstrings note the path observed on a typical machine.

🤖 Generated with Claude Code

[codecov]

Codecov Report

❌ Patch coverage is 93.50000% with 26 lines in your changes missing coverage. Please review.
✅ Project coverage is 91.12%. Comparing base (0979587) to head (657f01c).

Files with missing lines	Patch %	Lines
worker/egress/styx_kafka_batch_egress.py	69.56%	21 Missing ⚠️
coordinator/coordinator_metadata.py	97.80%	2 Missing ⚠️
styx-package/styx/common/tcp_networking.py	83.33%	2 Missing ⚠️
styx-package/styx/common/stateful_function.py	95.83%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main      #35      +/-   ##
==========================================
+ Coverage   89.73%   91.12%   +1.39%     
==========================================
  Files          45       46       +1     
  Lines        2707     2952     +245     
==========================================
+ Hits         2429     2690     +261     
+ Misses        278      262      -16     

Flag	Coverage Δ
coordinator	96.34% <98.37%> (+2.91%)	⬆️
integration	9.89% <22.75%> (+0.98%)	⬆️
styx-package	87.88% <93.18%> (+0.02%)	⬆️
worker	90.92% <88.84%> (+7.15%)	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

Files with missing lines	Coverage Δ
coordinator/aria_sync_metadata.py	100.00% <100.00%> (ø)
coordinator/migration_metadata.py	100.00% <100.00%> (ø)
coordinator/snapshot_compactor.py	100.00% <100.00%> (ø)
coordinator/worker_pool.py	94.35% <100.00%> (+0.13%)	⬆️
styx-package/styx/common/logging.py	100.00% <100.00%> (ø)
styx-package/styx/common/message_types.py	100.00% <100.00%> (ø)
worker/fault_tolerance/async_snapshots.py	98.65% <100.00%> (-0.48%)	⬇️
worker/ingress/styx_kafka_ingress.py	100.00% <100.00%> (+6.97%)	⬆️
worker/operator_state/aria/in_memory_state.py	100.00% <100.00%> (+1.34%)	⬆️
worker/sequencer/pending_requests.py	100.00% <100.00%> (ø)
... and 4 more

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.