From f1692cafdc0875a6d42d1f4194b7b9139df1e63d Mon Sep 17 00:00:00 2001 From: Migorithm Date: Tue, 21 Jul 2026 22:29:45 +0400 Subject: [PATCH 01/41] preliminary doc update before d10 - idempotent write --- docs/clients/c2_producer.md | 32 ++++++++++++++++++--------- docs/clients/client_roadmap.md | 5 +++-- docs/data-plane/data_plane_roadmap.md | 11 +++++++-- 3 files changed, 34 insertions(+), 14 deletions(-) diff --git a/docs/clients/c2_producer.md b/docs/clients/c2_producer.md index 8df7f3a2..ff6df5ad 100644 --- a/docs/clients/c2_producer.md +++ b/docs/clients/c2_producer.md @@ -117,15 +117,24 @@ this is; the codec tag says *how to read* it. --- -## Idempotency hook (future) +## Delivery contract and idempotency seam -Today produce is **at-least-once**: a produce that times out after the leader -committed but before the client saw the ack will be retried and stored twice. Exactly- -once needs a producer session id + per-record sequence numbers carried on the produce, -deduplicated at the segment leader — the server-side half is its own backlog item. C2 -should leave a clean seam for it (a place to stamp session/sequence on each produce) -even before the server enforces it, so enabling exactly-once later doesn't reshape the -producer API. +Today produce is explicitly **at-least-once**: a produce that times out after the +leader committed but before the client saw the acknowledgment can be retried and stored +twice. The UUID and per-record counter currently allocated by the producer do not cross +the wire and do not participate in a broker decision. + +The future protocol is specified in +[D10: Idempotent Production](../data-plane/d10_idempotent_production.md). It assigns a +sequence to each immutable broker batch, orders independently routed batches in lanes, +uses metadata-backed incarnations for fencing, and retains range-scoped deduplication +frontiers across rolls, failover, and lineage changes. A bounded recent-result window +returns exact positions for normal retries without retaining one position forever per +request. This is intentionally described as idempotent production rather than end-to- +end exactly-once processing. + +Until D10 is implemented end to end, callers must not interpret construction with a +stable producer UUID as a delivery guarantee. --- @@ -142,8 +151,9 @@ producer API. 5. **Compression** — optional client-side codec over the (batched) records, with a cleartext codec tag prefixing the compressed block so the consumer can decompress; the broker stays opaque and never decompresses. -6. **Idempotency seam** — a no-op-today hook to stamp producer session + sequence, - ready for the server-side dedup work. +6. **Idempotency seam** — a no-op-today identity allocation point. D10 will replace the + per-record counter with immutable batch identity and carry it through the broker + protocol. 7. **Tests** — against the simulated cluster: produce to the right leader from a warm cache (one hop), correction after a roll/failover (`NotWriteLeader` follow), correction from a cold/stale cache (`ShardNotLocal` follow), `TopicNotFound` @@ -156,4 +166,6 @@ producer API. - `d6_produce_consume_api.md` — the server's produce routing and redirect contract. - `d1_storage_engine.md` — the broker-opaque entry payload and end-to-end compression this stamps a codec into. +- `../data-plane/d10_idempotent_production.md` — session fencing, ordered lanes, + durable range ledgers, topology handoff, and bounded retry guarantees. - `client_roadmap.md` — the idempotency / batching / compression backlog context. diff --git a/docs/clients/client_roadmap.md b/docs/clients/client_roadmap.md index 4e0f09df..18c6735b 100644 --- a/docs/clients/client_roadmap.md +++ b/docs/clients/client_roadmap.md @@ -121,8 +121,9 @@ C2 (Producer) C3 (Consumer) --- ## Backlog (beyond a working client) -- **Idempotent / exactly-once produce** — producer session + sequence numbers, dedup - at the segment leader (the server-side half is its own backlog item). +- **Idempotent produce** — the complete client/server protocol is designed in + [data-plane D10](../data-plane/d10_idempotent_production.md): immutable batch + identities, independently ordered lanes, metadata-backed fencing, recovery journal, and bounded retry lifetime. Cross-range transactions and end-to-end exactly-once processing remain separate future work. - **Client-side batching & compression** — amortize round-trips and bytes. - **Adaptive replica selection** — pick the fastest/nearest replica per range from observed latency, not just `replica_set` order. diff --git a/docs/data-plane/data_plane_roadmap.md b/docs/data-plane/data_plane_roadmap.md index eaa99bcc..469472c5 100644 --- a/docs/data-plane/data_plane_roadmap.md +++ b/docs/data-plane/data_plane_roadmap.md @@ -243,6 +243,7 @@ Port layout: | [D6: Produce/Consume API](d6_produce_consume_api.md) | Server-side produce/consume routing via redirects | D4, D5 | Server routing (client SDK: see clients/) | | [D7: Retention GC](d7_retention_gc.md) | Optional per-topic **age** retention; expire sealed segments oldest-first, reclaim files | D3, D5 | Opt-in, logical (time); keep-forever is the default. Disk capacity is a separate node-level concern | | [D8: Consumer Offset Management](d8_consumer_offset_management.md) | Generation-fenced offset tracking and consumer-group work distribution | D2, D4, D6 | Raft-backed assignment + shared-WAL offset replication | +| [D10: Idempotent Production](d10_idempotent_production.md) | Session fencing, ordered producer lanes, and crash-durable deduplication | D2, D3, D5, D6, client C2 | Retry one logical batch without appending it twice | D1 defines the storage primitives (WAL, segment files, sparse index) and the threading model that drives them: DataPlaneActor on a dedicated OS thread (WAL + cache publish), lock-free per-segment `SegmentRingBuffer` (concurrent consumer reads without locking), and I/O thread pools (checkpoint writes + cold reads). D2–D5 extend the D1 foundation with replication, metadata integration, consumer tracking, and crash recovery. D6 adds the client-facing protocol layer (produce/consume wire format, connection management) — consumer tasks on tokio read directly from `SegmentRingBuffer`. @@ -267,6 +268,11 @@ D6 (Produce API) D7 (Retention GC) ← D7 also depends on D3 | v D8 (Consumer Groups) + +D2 + D3 + D5 + D6 + client C2 + | + v + D10 (Idempotent Produce) ``` D6 completes the **server-side** routing. The **client SDK** (producer, consumer, @@ -290,8 +296,9 @@ admin) that consumes those redirects is its own track — see ## Backlog -### Exactly-Once Semantics -Producer idempotency keys, deduplication at segment leader. Requires producer session tracking. +### Transactions / End-to-End Exactly-Once +Idempotent append is designed in D10. Atomic writes across ranges and atomic coupling of +produced entries to consumer-offset commits remain future transaction work. From 4e630a433580ced61fcec3d0f92314bc98d16fc7 Mon Sep 17 00:00:00 2001 From: Migorithm Date: Wed, 22 Jul 2026 02:15:09 +0400 Subject: [PATCH 02/41] doc plans out range auxiliary state --- docs/clients/c2_producer.md | 4 +- docs/clients/client_roadmap.md | 4 +- .../d8_consumer_offset_management.md | 51 ++++++++++++++++++- .../d9_offset_placement_graduation.md | 32 ++++++++++++ docs/data-plane/data_plane_roadmap.md | 2 +- 5 files changed, 87 insertions(+), 6 deletions(-) diff --git a/docs/clients/c2_producer.md b/docs/clients/c2_producer.md index ff6df5ad..dfd78be2 100644 --- a/docs/clients/c2_producer.md +++ b/docs/clients/c2_producer.md @@ -126,7 +126,7 @@ the wire and do not participate in a broker decision. The future protocol is specified in [D10: Idempotent Production](../data-plane/d10_idempotent_production.md). It assigns a -sequence to each immutable broker batch, orders independently routed batches in lanes, +sequence to each immutable per-range broker batch and orders each producer-range stream, uses metadata-backed incarnations for fencing, and retains range-scoped deduplication frontiers across rolls, failover, and lineage changes. A bounded recent-result window returns exact positions for normal retries without retaining one position forever per @@ -166,6 +166,6 @@ stable producer UUID as a delivery guarantee. - `d6_produce_consume_api.md` — the server's produce routing and redirect contract. - `d1_storage_engine.md` — the broker-opaque entry payload and end-to-end compression this stamps a codec into. -- `../data-plane/d10_idempotent_production.md` — session fencing, ordered lanes, +- `../data-plane/d10_idempotent_production.md` — session fencing, ordered producer-range streams, durable range ledgers, topology handoff, and bounded retry guarantees. - `client_roadmap.md` — the idempotency / batching / compression backlog context. diff --git a/docs/clients/client_roadmap.md b/docs/clients/client_roadmap.md index 18c6735b..490d42e1 100644 --- a/docs/clients/client_roadmap.md +++ b/docs/clients/client_roadmap.md @@ -123,7 +123,9 @@ C2 (Producer) C3 (Consumer) ## Backlog (beyond a working client) - **Idempotent produce** — the complete client/server protocol is designed in [data-plane D10](../data-plane/d10_idempotent_production.md): immutable batch - identities, independently ordered lanes, metadata-backed fencing, recovery journal, and bounded retry lifetime. Cross-range transactions and end-to-end exactly-once processing remain separate future work. + identities, independently ordered producer-range streams, metadata-backed fencing, + recovery journal, and bounded retry lifetime. Cross-range transactions and end-to-end + exactly-once processing remain separate future work. - **Client-side batching & compression** — amortize round-trips and bytes. - **Adaptive replica selection** — pick the fastest/nearest replica per range from observed latency, not just `replica_set` order. diff --git a/docs/data-plane/d8_consumer_offset_management.md b/docs/data-plane/d8_consumer_offset_management.md index aa007b43..a836b67d 100644 --- a/docs/data-plane/d8_consumer_offset_management.md +++ b/docs/data-plane/d8_consumer_offset_management.md @@ -168,12 +168,59 @@ offset commit fsync Snapshot creation is driven by reclamation pressure. Many commits therefore collapse into a single asynchronous snapshot. WAL deletion uses the minimum safe watermark across segment -checkpoints and consumer-offset checkpoints; neither state family can delete recovery data -still needed by the other. +checkpoints and the auxiliary-state checkpoint, which contains only consumer offsets today; +neither state family can delete recovery data still needed by the other. On restart, recovery loads the latest snapshot and replays later offset records from the shared WAL. Recovery then writes the consolidated snapshot before removing the replayed WAL. +### Consolidation with producer frontier state + +The implemented snapshot currently contains consumer-offset state only. D10 adds another +range-scoped, end-state ledger for producer retry frontiers. It must extend this checkpoint +into one **range-auxiliary-state snapshot**, rather than introducing a second independently +fsynced snapshot: + +```text +shared WAL + ├── consumer offset records + └── produced entries carrying producer sequence identity + │ reclamation reaches uncovered state + ▼ +one auxiliary-state snapshot at one LSN boundary + ├── current consumer offsets + placement readiness + └── current producer frontiers + bounded recent results + │ durable rename completes + ▼ +advance one auxiliary checkpoint watermark +``` + +This follows the metadata log-to-snapshot model: replayable log records are authoritative +until one complete current-state image safely replaces their covered prefix. The snapshot +is asynchronous and reclamation-driven, so its single fsync is not part of offset-commit or +produce latency. Segment payload checkpointing remains separate because segment files retain +application data, whereas both auxiliary ledgers retain only their latest state. + +### Migration from the consumer-only checkpoint + +The D10 implementation should replace, rather than wrap, the current consumer-only snapshot +pipeline: + +1. Replace the consumer-offset-only snapshot image with the typed auxiliary-state image. +2. Replace its dedicated checkpoint job, completion signal, in-flight flag, and watermark + with one auxiliary-state checkpoint lifecycle. +3. Keep offset commits and graduation imports in the shared WAL; their WAL fsync is still the + acknowledgement durability boundary. +4. On upgrade, load the legacy consumer-offset snapshot when no consolidated snapshot exists, + combine it with producer frontier state recovered from WAL, and publish a durable + consolidated snapshot before reclaiming either the legacy file or covered WAL. +5. Remove the legacy snapshot only after the consolidated snapshot's durable rename and + directory sync complete. + +This cleanup removes the independent consumer-snapshot fsync once D10 adds producer state. It +does not remove the shared-WAL fsync required before acknowledging an offset commit, nor the +single asynchronous fsync that makes a replacement snapshot safe for WAL truncation. + --- ## Range Lifecycle and Placement diff --git a/docs/data-plane/d9_offset_placement_graduation.md b/docs/data-plane/d9_offset_placement_graduation.md index 168105e8..8a2de8fe 100644 --- a/docs/data-plane/d9_offset_placement_graduation.md +++ b/docs/data-plane/d9_offset_placement_graduation.md @@ -179,6 +179,38 @@ positions while ordinary commits may still arrive. Sharing the reconciliation sh a second recovery lifecycle without pretending the two kinds of state have the same merge rules. +### Generalization for range auxiliary state + +D10 extends this implemented offset-graduation path into a shared **range auxiliary-state +graduation** framework. A placement announcement, successor-segment token, ready-source +selection, snapshot transport, shared-WAL completion marker, joining/ready transition, +acknowledgement drain, and reconciliation retry are placement mechanics; they should not be +implemented again for each range-scoped ledger. + +The transferred snapshot is an envelope with typed sections: + +| Section | Merge authority | Meaning | +|---|---|---| +| Consumer offsets | Group generation and committed position | Where each group may safely resume | +| Producer frontiers | Session incarnation and highest committed producer-range sequence | Which producer retries must not append again | + +The envelope and graduation state machine are shared, while each section owns its validation, +monotonic merge, memory index, and snapshot encoding. A replica becomes authoritative for a +placement only after every required section and the placement completion marker are durable. +This avoids duplicate lifecycle machinery without hiding materially different ledger rules +inside an untyped generic map. + +The same envelope becomes the node's consolidated range-auxiliary-state checkpoint. Offset +updates and producer identities first become durable through typed records in the shared WAL. +When reclamation needs to replace those records, one checkpoint job snapshots every auxiliary +section at one LSN boundary, performs one asynchronous snapshot fsync, and advances one +auxiliary watermark. Producer frontiers therefore do not add an independent snapshot file, +fsync, or reclamation gate. + +The current implementation graduates consumer offsets only. Producer-frontier participation is +part of D10 and must preserve the existing offset behavior while extending the snapshot envelope +and readiness check. + --- ## Leadership Changes diff --git a/docs/data-plane/data_plane_roadmap.md b/docs/data-plane/data_plane_roadmap.md index 469472c5..9240ba95 100644 --- a/docs/data-plane/data_plane_roadmap.md +++ b/docs/data-plane/data_plane_roadmap.md @@ -243,7 +243,7 @@ Port layout: | [D6: Produce/Consume API](d6_produce_consume_api.md) | Server-side produce/consume routing via redirects | D4, D5 | Server routing (client SDK: see clients/) | | [D7: Retention GC](d7_retention_gc.md) | Optional per-topic **age** retention; expire sealed segments oldest-first, reclaim files | D3, D5 | Opt-in, logical (time); keep-forever is the default. Disk capacity is a separate node-level concern | | [D8: Consumer Offset Management](d8_consumer_offset_management.md) | Generation-fenced offset tracking and consumer-group work distribution | D2, D4, D6 | Raft-backed assignment + shared-WAL offset replication | -| [D10: Idempotent Production](d10_idempotent_production.md) | Session fencing, ordered producer lanes, and crash-durable deduplication | D2, D3, D5, D6, client C2 | Retry one logical batch without appending it twice | +| [D10: Idempotent Production](d10_idempotent_production.md) | Session fencing, ordered producer-range streams, and crash-durable deduplication | D2, D3, D5, D6, D9, client C2 | Retry one logical batch without appending it twice | D1 defines the storage primitives (WAL, segment files, sparse index) and the threading model that drives them: DataPlaneActor on a dedicated OS thread (WAL + cache publish), lock-free per-segment `SegmentRingBuffer` (concurrent consumer reads without locking), and I/O thread pools (checkpoint writes + cold reads). D2–D5 extend the D1 foundation with replication, metadata integration, consumer tracking, and crash recovery. D6 adds the client-facing protocol layer (produce/consume wire format, connection management) — consumer tasks on tokio read directly from `SegmentRingBuffer`. From fd03f8ca94f2c946775e9fb075a6e33659df1f08 Mon Sep 17 00:00:00 2001 From: Migorithm Date: Wed, 22 Jul 2026 02:30:42 +0400 Subject: [PATCH 03/41] d10 Transparent Retry Deduplication --- .../d10_transparent_retry_deduplication.md | 344 ++++++++++++++++++ 1 file changed, 344 insertions(+) create mode 100644 docs/data-plane/d10_transparent_retry_deduplication.md diff --git a/docs/data-plane/d10_transparent_retry_deduplication.md b/docs/data-plane/d10_transparent_retry_deduplication.md new file mode 100644 index 00000000..fce8445e --- /dev/null +++ b/docs/data-plane/d10_transparent_retry_deduplication.md @@ -0,0 +1,344 @@ +# D10 — Transparent Retry Deduplication (Write Idempotency) + +**Goal:** Prevent a transparent producer retry from storing the same logical batch +twice, while preserving EastGuard's direct-to-data-leader write path across crashes, +segment rolls, and range lineage changes. + +**Depends on:** [D2: Segment Replication](d2_segment_replication.md), +[D3: Segment Lifecycle Integration](d3_segment_roll_integration.md), +[D5: Crash Recovery](d5_crash_recovery.md), +[D6: Produce/Consume API](d6_produce_consume_api.md), +[D9: Offset Placement Graduation](d9_offset_placement_graduation.md), and +[C2: Producer](../clients/c2_producer.md). + +--- + +## Delivery contract + +Within a live producer session: + +- the first successful request appends one entry; +- retrying it never appends another entry; +- a retry inside the bounded result window returns its original range and entry id; +- an older retry is still recognized as committed, but its position may have aged out; +- activating a newer producer incarnation fences the older one; +- expired and unknown sessions reject produce requests and are never auto-created. + +This is transparent **producer retry deduplication**, not application-level idempotency. +Arbitrary business keys that must deduplicate across producers or ranges need a separate +index with its own routing, conflict, retention, and compaction contract. + +It is also not end-to-end exactly-once processing. Applications can create duplicates by +sending the same records under a new request identity or by starting a new session instead +of recovering an ambiguous one. Cross-range transactions and atomic coupling to consumer- +offset commits remain outside this phase. + +--- + +## Two protocol layers + +The control plane owns producer authority; the data plane owns append ordering and +durability. + +| Layer | State and operations | Frequency | +|---|---|---| +| Topic Raft group | Open session, recover session and increment incarnation, renew lease, expire session, activate fences | Low | +| Range data leader | Check sequence, append entry plus producer identity, replicate, advance producer state, answer duplicates | Every produce | + +No produce request enters the Raft log. The existing primary-backup data replication and +all-replica fsync still apply; “no consensus per produce” means no metadata-Raft proposal, +not no cross-node durability. + +``` +control plane: session lifecycle ── fence ──► active range leaders + +data plane: client ── produce ──► range leader ── replicate ──► followers + WAL + fsync WAL + fsync +``` + +--- + +## Request identity and batching + +A request is identified by: + +| Part | Purpose | +|---|---| +| Producer id | Broker-issued random session identity persisted by the client for recovery | +| Incarnation | Metadata-allocated fencing epoch | +| Range id | Selects the producer-specific ordered stream | +| Sequence | Zero-based contiguous batch number for this producer and range | +| Payload digest | Detects reuse of a recent identity with different content | + +Sequence belongs to the immutable broker batch, not to each application record. The +current per-record counter is only a client-side seam and must be replaced. + +One client flush can contain records that a refreshed routing snapshot places in different +ranges. Before assigning sequences, the producer divides it into one immutable broker +batch per destination range: + +```text +client flush [a, b, c, d] + │ refresh routing + ├── range 4: [a, c] → producer/range-4 sequence 17 → node A + └── range 9: [b, d] → producer/range-9 sequence 17 → node D +``` + +The two sub-batches may both use sequence 17 because range id is part of their identity. +A stored entry never spans ranges or nodes. Within one range, batch closure allocates the +next sequence atomically and a per-range dispatch queue preserves order. Different ranges +continue concurrently. Once assigned, a retry preserves both the sequence and the exact +serialized bytes. + +A segment roll preserves the range and its sequence stream. A split or merge closes the +parent streams and starts sequence zero for each successor range. + +--- + +## The range producer state: one integer plus a small ring + +For each live `(producer session, range)`, the range stores only two things: + +1. **Committed frontier.** The highest durably committed sequence. This one integer + decides whether a request is next, duplicate, or unexpectedly ahead. +2. **Recent results.** A fixed-size ring of recent sequence, digest, and committed- + position tuples. It exists only to return positions for in-flight retries and detect + recent identity conflicts. + +The frontier prevents duplicates for the entire session. The ring does not grow with the +number of requests and must be at least as large as the producer's maximum in-flight +batches per range. + +### Range-leader verification pipeline + +```text +request(session, incarnation, range, sequence, digest) +│ +├── session unknown or expired +│ └── SessionExpired +├── incarnation older than installed fence +│ └── ProducerFenced +├── dedup state or placement still recovering +│ └── DedupStateRecovering +├── sequence == frontier + 1 +│ ├── same request already staged ── RequestInFlight +│ └── otherwise ACCEPT +│ └── WAL(entry + producer identity) → replicate → fsync all → advance +├── sequence <= frontier +│ ├── recent + same digest ──────── Duplicate(position) +│ ├── recent + different digest ─── RequestIdentityConflict +│ └── aged out of ring ──────────── DuplicatePositionUnavailable +└── sequence > frontier + 1 + └── SequenceGap { expected: frontier + 1 } +``` + +An older request outside the ring is rejected as a duplicate regardless of its current +payload. The broker can no longer diagnose a digest conflict or return the exact position, +but it never appends the request again. + +The normal producer API keeps unresolved requests inside the result window, so its send +future still completes with a committed position. A low-level historical retry may receive +successful deduplication without one. + +--- + +## Session lifecycle and fencing + +A producer session is scoped to one topic, whose metadata already lives in one shard +group. Opening creates a random producer id and incarnation zero. Recovering an existing +session commits the next incarnation through Raft. + +After the increment commits, the coordinator broadcasts the new fence to every active +range leader and waits for older in-flight writes to drain or fail. The recovered producer +cannot send until that activation completes. Assignments created by a concurrent roll, +split, or merge install the current fence before accepting writes. + +``` +restarting producer metadata coordinator range leaders + │ │ │ + ├── recover ──────────────►│ │ + │ ├── Raft: incarnation+1 │ + │ ├── install fence ──────►│ + │ │◄── old writes drained ─┤ + │◄── active incarnation ───┤ │ +``` + +The client journals its producer id, incarnation, unresolved immutable batches, and the +next sequence for every touched range. Recovery reconciles those journaled ranges with +their durable frontiers before retrying ambiguous batches. + +### Inactivity and deletion + +Activity renews a coarse session lease, not one metadata record per produce. The metadata +leader proposes expiration after the lease and grace period pass; only committed +expiration invalidates the session. + +Range leaders install the expiration fence before producer state becomes garbage- +collectable. Reopening later creates a new random producer id, so sequence zero cannot +collide with a delayed request from the old session. Explicit close may accelerate the +same process, but correctness does not depend on clean client shutdown. + +--- + +## Durability and recovery + +The frontier is an in-memory index, but the facts used to rebuild it are durable: + +- the producer identity and sequence ride in the existing entry WAL record; +- followers receive and fsync the same identity with the entry; +- the leader advances the frontier only after the entry reaches the normal D2 commit + boundary; +- recovery loads the consolidated auxiliary-state snapshot and replays later committed + entries up to the recovered commit boundary. + +This adds no second per-produce WAL record or fsync. Consumer offsets and producer +frontiers use the single consolidated end-state snapshot defined with D8/D9: one typed +snapshot, one asynchronous save pipeline, and one auxiliary-state WAL checkpoint +watermark. Producer state is another typed section in that design, not an independent +snapshot lifecycle. + +Segment checkpointing remains separate because segment files retain application payload +and sparse seek structure. The auxiliary snapshot retains only current consumer offsets, +placement readiness, producer frontiers, and bounded recent results. + +### Leader failure + +If an entry becomes durable and its response is lost, a surviving replica rebuilds the +same frontier from the auxiliary snapshot and committed WAL suffix. When recovery changes +the replica set, D9 graduation installs the auxiliary state on joining replicas before +they become authoritative. A retry therefore resolves as a duplicate after leader +failover. + +### Segment roll + +A roll retains the range identity, so its producer state continues unchanged. If the +replica set changes, producer state uses the same D9 placement trigger, successor-segment +token, ready source, snapshot envelope, durable completion marker, joining/ready states, +and retry machinery as consumer offsets. The transport and readiness framework are +shared; each typed section keeps its own monotonic merge rules. + +--- + +## Split and merge: freeze the parent + +When a range splits or merges, closure first settles every staged request against the +committed boundary; its producer state then becomes a complete, read-only frozen snapshot. +Successor ranges start new producer-range streams at sequence zero; they do not copy the +parent frontier. + +Retries issued against the parent before it closed continue to target that frozen state: + +| Frozen-parent result | Action | +|---|---| +| Sequence at or below the frontier, result retained | Return duplicate with the parent position | +| Sequence at or below the frontier, result aged out | Return duplicate without position | +| Sequence exactly after the frontier | Return `RangeClosedNotCommitted`; the client may repartition into successor batches | +| Frozen state unavailable or incomplete | Return a retriable recovery error; never infer non-commit from missing state | + +Routing metadata directs the retry to a replica holding the frozen parent state. The +client may form new child requests only after the complete frozen frontier proves that +the parent request did not commit. This closes the acknowledgment-loss window across a +split or merge. + +Frozen parent state remains available until its referenced producer sessions expire. If +repair moves the sealed parent data, its auxiliary-state section moves through the same +D9 graduation mechanism. + +--- + +## Bounds and public API + +The session contract fixes: + +| Bound | Effect | +|---|---| +| Session lease / maximum retry age | Expired requests receive `SessionExpired`, never a new append | +| Maximum live sessions | Bounds topic- and cluster-wide producer state | +| Maximum touched ranges per session | Bounds the number of frontiers | +| Recent results per producer-range | Bounds exact position and digest history | + +The consolidated snapshot contains the same bounded end state. Request identities after +its LSN boundary are replayed from WAL; older WAL can be reclaimed after both segment and +auxiliary checkpoint watermarks permit it. + +Produce requests add producer id, incarnation, producer-range sequence, and digest. +Successful responses distinguish newly appended, duplicate with position, and duplicate +without position. + +Keep a clearly named raw/at-least-once append only if internal callers need it. The public +producer should default to idempotent sessions only after this protocol is implemented +end to end. Until then EastGuard's public delivery contract remains at-least-once. + +--- + +## Validation + +### Sequence and deduplication + +- Normal next sequence appends once and advances the frontier after durability. +- Losing an acknowledgment and retrying returns the same position with one stored entry. +- A recent identity with different bytes returns an identity conflict. +- A sequence gap returns the expected sequence without appending. +- An aged-out sequence returns duplicate without position and does not append. +- Concurrent range batches advance independently without false gaps. + +### Session lifecycle and fencing + +- Activating a newer incarnation drains the boundary and fences every later old write. +- Expired and unknown sessions reject writes and are never implicitly recreated. +- Admission, touched-range, and recent-result limits keep state bounded. + +### Failover and recovery + +- Crash the write leader after follower fsync but before the client response; the + successor returns the original position for the unresolved retry. +- Restart all replicas and reconstruct the same frontier and recent window from the + consolidated snapshot plus committed WAL suffix. +- Migrate a legacy consumer-only snapshot before reclaiming covered WAL. + +### Topology transitions + +- Retry across a segment roll and changed replica placement without appending twice. +- Retry against a frozen split or merge parent and return duplicate or proven non-commit. +- Move frozen parent state during repair and keep it queryable. + +Deterministic data-plane tests should cover sequence decisions and state transfer. +Multi-node turmoil tests should pin RNG and node ids and use size-based rolls, following +the repository's data-plane test rules. + +--- + +## Implementation phases + +1. **Storage and WAL framing.** Extend entry identity, generalize the consumer-only + checkpoint into the typed D9 auxiliary-state snapshot, and support legacy snapshot + migration. +2. **Local deduplication.** Add the frontier plus recent-result ring and the range-leader + verification pipeline; change client batching to assign immutable per-range sequences. +3. **Session lifecycle and fencing.** Add Raft-backed open, recover, renew, expire, and + active-range fence installation; add the durable client recovery journal. +4. **Topology transitions.** Reuse D9 graduation for rolls and repair, and add frozen- + parent lookup for split and merge retries. + +## Design rules + +1. **One immutable request identity names one broker batch.** A retry never changes its + sequence or serialized bytes. +2. **Ordering is per producer and range.** Different ranges progress concurrently without + creating false gaps. +3. **The durable frontier decides append versus duplicate.** The recent ring only enriches + duplicate responses; evicting it cannot reopen an append. +4. **Raft owns session authority, not produce throughput.** Produce requests use the data + replica protocol and never enter the metadata log. +5. **A newer incarnation becomes usable only after fencing completes.** An older process + cannot write after recovery activates its successor. +6. **Rolls preserve producer state; split and merge freeze it.** Successors start new + streams while parents remain available for retry decisions. +7. **Missing state never proves non-commit.** Recovery uncertainty delays a retry instead + of turning it into a new append. +8. **Expired or unknown sessions reject every produce.** They are never auto-created. +9. **Acknowledgment follows full data-replica durability.** Entry and request identity are + fsynced together before success. + +These rules describe the proposed phase, not current behavior. Current production remains +at-least-once until every layer is implemented and validated. From 356a9723cb52df27ccd8e336e3a2042ff18cdf47 Mon Sep 17 00:00:00 2001 From: Migorithm Date: Wed, 22 Jul 2026 12:43:32 +0400 Subject: [PATCH 04/41] sdk default to protocol write idempotency --- .../d10_transparent_retry_deduplication.md | 36 +++++++++++++++---- 1 file changed, 30 insertions(+), 6 deletions(-) diff --git a/docs/data-plane/d10_transparent_retry_deduplication.md b/docs/data-plane/d10_transparent_retry_deduplication.md index fce8445e..d90c640d 100644 --- a/docs/data-plane/d10_transparent_retry_deduplication.md +++ b/docs/data-plane/d10_transparent_retry_deduplication.md @@ -265,9 +265,33 @@ Produce requests add producer id, incarnation, producer-range sequence, and dige Successful responses distinguish newly appended, duplicate with position, and duplicate without position. -Keep a clearly named raw/at-least-once append only if internal callers need it. The public -producer should default to idempotent sessions only after this protocol is implemented -end to end. Until then EastGuard's public delivery contract remains at-least-once. +The SDK `Producer` uses this protocol by default. Low-level `Client::produce` remains an +explicit raw append and therefore remains at-least-once. + +The implemented delivery contract is: + +- Within a live producer lease, retrying the same producer-range sequence and payload + appends at most once and returns its original entry id while that result remains in the + 16-entry recent window. +- After the result leaves that window, the frontier still prevents another append, but + the broker returns `DuplicatePositionUnavailable` because the original entry id is no + longer retained. +- Reusing a sequence with a different payload checksum, skipping a sequence, using an + expired session, or writing with a fenced incarnation is rejected without appending. +- A restarted producer using the same producer id opens with a new session nonce. Metadata + advances its incarnation and fences the older process. Retrying the open itself is + idempotent. +- Request identity is part of the replicated data entry and the shared WAL batch. There is + no producer-only WAL record or additional fsync. Consumer offsets and producer ledgers + share `auxiliary-state.snapshot`; startup migrates the legacy consumer-only snapshot. +- A retry against a sealed split/merge parent is routed to the frozen parent ledger. A + committed request returns its old position; a proven next sequence returns stale range + so the SDK can repartition it to children. + +The remaining limits are deliberate. Exact historical positions are bounded by the recent +window, payload conflict detection uses the stored checksum rather than retaining payload +bytes, and a frozen parent lookup requires one of the parent's retained replicas to remain +reachable. An expired request is never accepted merely because its ledger was collected. --- @@ -308,7 +332,7 @@ the repository's data-plane test rules. --- -## Implementation phases +## Implementation shape 1. **Storage and WAL framing.** Extend entry identity, generalize the consumer-only checkpoint into the typed D9 auxiliary-state snapshot, and support legacy snapshot @@ -340,5 +364,5 @@ the repository's data-plane test rules. 9. **Acknowledgment follows full data-replica durability.** Entry and request identity are fsynced together before success. -These rules describe the proposed phase, not current behavior. Current production remains -at-least-once until every layer is implemented and validated. +These rules describe the SDK producer path. The low-level raw append path remains +at-least-once by design. From 071a553366dea5abe0760313ad1fb695073e8847 Mon Sep 17 00:00:00 2001 From: Migorithm Date: Wed, 22 Jul 2026 12:45:06 +0400 Subject: [PATCH 05/41] move now_ms() --- src/lib.rs | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/src/lib.rs b/src/lib.rs index a73140c1..a45dd1d7 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -189,3 +189,10 @@ impl StartUp { } } } + +pub(crate) fn now_ms() -> u64 { + std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap_or_default() + .as_millis() as u64 +} From 15ff20389dcf5cfeb5e36152574d2b3bfb2f7f15 Mon Sep 17 00:00:00 2001 From: Migorithm Date: Wed, 22 Jul 2026 12:45:57 +0400 Subject: [PATCH 06/41] wire produce error and ProducerAppendIdentity --- src/client/error.rs | 4 ++++ src/data_plane/mod.rs | 35 +++++++++++++++++++++++++++++++++++ 2 files changed, 39 insertions(+) diff --git a/src/client/error.rs b/src/client/error.rs index bde42853..607602c8 100644 --- a/src/client/error.rs +++ b/src/client/error.rs @@ -3,6 +3,7 @@ use std::time::Duration; use crate::client::RangeId; use crate::control_plane::metadata::consumer_group::GenerationId; +use crate::data_plane::ProduceError; /// Errors a caller decides on. Redirect-following, reconnect, and retry-within-deadline /// are handled internally. @@ -38,6 +39,9 @@ pub enum ClientError { #[error("stale range routing")] StaleRange, + #[error("produce rejected: {0}")] + ProduceRejected(ProduceError), + #[error( "consumer group generation {request_generation:?} is stale; data layer sealed at {sealed_generation:?}" )] diff --git a/src/data_plane/mod.rs b/src/data_plane/mod.rs index 39556be9..1d80047d 100644 --- a/src/data_plane/mod.rs +++ b/src/data_plane/mod.rs @@ -1,6 +1,7 @@ use borsh::{BorshDeserialize, BorshSerialize}; use bytes::Bytes; use std::path::{Path, PathBuf}; +use uuid::Uuid; use crate::control_plane::metadata::{EntryId, RangeId, SegmentId, TopicId}; use crate::impl_new_struct_wrapper; @@ -10,6 +11,7 @@ pub(crate) mod checkpoint; pub(crate) mod cold_read; pub(crate) mod consumer_offset_management; pub(crate) mod messages; +pub(crate) mod producer_ledger; pub(crate) mod recovery; pub(crate) mod segment_writer; pub(crate) mod sparse_index; @@ -28,6 +30,39 @@ pub struct SegmentKey { pub segment_id: SegmentId, } +/// Stable identity of one producer append. The range is supplied by the +/// enclosing `SegmentKey`, so it is deliberately not duplicated here. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, BorshSerialize, BorshDeserialize)] +pub struct ProducerAppendIdentity { + pub producer_id: Uuid, + pub incarnation: u32, + pub expires_at: u64, + pub sequence: u64, + pub digest: u32, +} + +#[derive(Debug, Clone, PartialEq, Eq, BorshSerialize, BorshDeserialize, thiserror::Error)] +pub enum ProduceError { + #[error("not the write leader")] + NotLeader, + #[error("segment not found")] + SegmentNotFound, + #[error("producer incarnation was fenced")] + ProducerFenced, + #[error("producer session expired or is unknown")] + SessionExpired, + #[error("sequence was reused with a different payload")] + RequestIdentityConflict, + #[error("duplicate position is outside the retained result window")] + DuplicatePositionUnavailable, + #[error("sequence gap; expected {0}")] + SequenceGap(u64), + #[error("the same producer request is already in flight")] + RequestInFlight, + #[error("internal produce failure: {0}")] + Internal(String), +} + #[derive(Debug, Clone, BorshSerialize, BorshDeserialize)] pub struct EntryPayload(Bytes); From 7e352a8050840ea478bf8dbc7fef7aeda9b1603c Mon Sep 17 00:00:00 2001 From: Migorithm Date: Wed, 22 Jul 2026 12:50:53 +0400 Subject: [PATCH 07/41] e2e test: duplicate_unknown_and_fenced_sessions_are_end_to_end_visible --- src/it/e2e/producer_deduplication.rs | 172 +++++++++++++++++++++++++++ 1 file changed, 172 insertions(+) create mode 100644 src/it/e2e/producer_deduplication.rs diff --git a/src/it/e2e/producer_deduplication.rs b/src/it/e2e/producer_deduplication.rs new file mode 100644 index 00000000..f2188b46 --- /dev/null +++ b/src/it/e2e/producer_deduplication.rs @@ -0,0 +1,172 @@ +//! End-to-end validation of D10's client-to-broker producer protocol. + +use std::net::SocketAddr; +use std::sync::Arc; +use std::time::Duration; + +use turmoil::Builder; +use uuid::Uuid; + +use super::{NODES, host_cluster}; +use crate::client::{Client, ClientError, PartitionStrategy, StoragePolicy}; +use crate::connections::protocol::{ + ClientDataPlaneRequest, ClientRequest, ClientResponse, DataPlaneResponse, + OpenProducerSessionRequest, ProduceRequest, +}; +use crate::data_plane::{ProduceError, ProducerAppendIdentity}; +use crate::it::helpers::send_request; + +fn seeds() -> Vec { + NODES + .iter() + .map(|(host, port, _)| SocketAddr::new(turmoil::lookup(*host), *port)) + .collect() +} + +fn one_record(key: &[u8], value: &[u8]) -> Vec { + let mut payload = vec![0]; // CompressionCodec::None + payload.extend_from_slice(&(key.len() as u32).to_be_bytes()); + payload.extend_from_slice(key); + payload.extend_from_slice(&(value.len() as u32).to_be_bytes()); + payload.extend_from_slice(value); + payload +} + +#[test] +#[serial_test::serial] +fn duplicate_unknown_and_fenced_sessions_are_end_to_end_visible() -> turmoil::Result { + let mut sim = Builder::new() + .tick_duration(Duration::from_millis(100)) + .simulation_duration(Duration::from_secs(90)) + .tcp_capacity(4096) + .rng_seed(1) + .build(); + host_cluster(&mut sim, &NODES, |env| { + env.node_id_suffix = Some("producer-dedup".to_string()); + env.vnodes_per_node = 16; + }); + + sim.client("producer-dedup-client", async { + let client = Arc::new(Client::connect(seeds()).expect("client connects")); + let topic = "producer-dedup"; + client + .create_topic( + topic, + StoragePolicy { + retention_ms: Some(3_600_000), + replication_factor: 3, + partition_strategy: PartitionStrategy::AutoSplit, + }, + ) + .await + .expect("create topic"); + let detail = client.resolve_topic(topic).await.expect("resolve topic"); + let range_id = detail.ranges[0].range_id; + let producer_id = Uuid::new_v4(); + + let open_session_request1 = OpenProducerSessionRequest { + topic_name: topic.to_string(), + producer_id, + session_nonce: Uuid::new_v4(), + }; + let session = client + .open_producer_session(open_session_request1) + .await + .expect("open producer session"); + tokio::time::sleep(Duration::from_secs(1)).await; + let payload = one_record(b"k", b"only-once"); + let append_identity = ProducerAppendIdentity { + producer_id, + incarnation: session.incarnation, + expires_at: session.expires_at, + sequence: 0, + digest: crc32fast::hash(&payload), + }; + + let first = client + .produce_to_range( + topic, + range_id, + b"k", + payload.clone(), + 1, + Some(append_identity), + ) + .await + .expect("first append"); + let retry = client + .produce_to_range( + topic, + range_id, + b"k", + payload.clone(), + 1, + Some(append_identity), + ) + .await + .expect("retry resolves as duplicate"); + assert_eq!(retry, first, "retry must return the original position"); + + let unknown = ProducerAppendIdentity { + producer_id: Uuid::new_v4(), + ..append_identity + }; + let unknown_wire = + ClientRequest::DataPlane(ClientDataPlaneRequest::Produce(ProduceRequest { + topic_name: topic.to_string(), + range_id, + routing_key: b"k".to_vec(), + data: payload.clone(), + record_count: 1, + producer: Some(unknown), + })); + let mut rejected = false; + for (host, port, _) in NODES { + if matches!( + send_request(host, port, unknown_wire.clone()).await, + ClientResponse::DataPlane(DataPlaneResponse::ProduceRejected( + ProduceError::SessionExpired + )) + ) { + rejected = true; + break; + } + } + assert!( + rejected, + "the write leader did not reject the unknown session" + ); + + let open_session_request2 = OpenProducerSessionRequest { + topic_name: topic.to_string(), + producer_id, + session_nonce: Uuid::new_v4(), + }; + let newer = client + .open_producer_session(open_session_request2) + .await + .expect("recover producer session"); + assert!(newer.incarnation > session.incarnation); + tokio::time::sleep(Duration::from_secs(1)).await; + let newer_payload = one_record(b"k", b"new-incarnation"); + let newer_request = ProducerAppendIdentity { + incarnation: newer.incarnation, + expires_at: newer.expires_at, + digest: crc32fast::hash(&newer_payload), + ..append_identity + }; + client + .produce_to_range(topic, range_id, b"k", newer_payload, 1, Some(newer_request)) + .await + .expect("new incarnation becomes active on range leader"); + assert!(matches!( + client + .produce_to_range(topic, range_id, b"k", payload, 1, Some(append_identity)) + .await, + Err(ClientError::ProduceRejected(ProduceError::ProducerFenced)) + )); + Ok(()) + }); + + sim.run() +} From a92ab042b0e554bf9642d4ea895366654b59beb5 Mon Sep 17 00:00:00 2001 From: Migorithm Date: Wed, 22 Jul 2026 12:58:35 +0400 Subject: [PATCH 08/41] fix test --- src/it/e2e/client_protocol.rs | 1 + src/it/e2e/client_sdk.rs | 53 +++++++++++++++++++++++++---------- src/it/e2e/mod.rs | 1 + src/it/sim/scenario.rs | 1 + 4 files changed, 41 insertions(+), 15 deletions(-) diff --git a/src/it/e2e/client_protocol.rs b/src/it/e2e/client_protocol.rs index fba3c8f8..56a537a6 100644 --- a/src/it/e2e/client_protocol.rs +++ b/src/it/e2e/client_protocol.rs @@ -1832,6 +1832,7 @@ async fn produce_once(topic: &str, payload: &[u8], node: (&str, u16)) -> Produce routing_key: b"k".to_vec(), data: payload.to_vec(), record_count: 1, + producer: None, })); match send_request(node.0, node.1, req).await { ClientResponse::DataPlane(DataPlaneResponse::Produced { .. }) => ProduceOutcome::Acked, diff --git a/src/it/e2e/client_sdk.rs b/src/it/e2e/client_sdk.rs index eba16eaf..34a85406 100644 --- a/src/it/e2e/client_sdk.rs +++ b/src/it/e2e/client_sdk.rs @@ -120,11 +120,25 @@ fn client_produce_routes_to_write_leader() -> turmoil::Result { // The SDK absorbs the post-create segment-assignment lag inside produce. let first = client - .produce("produce-route", b"k", b"rec-0".to_vec(), 1) + .produce_to_range( + "produce-route", + RangeId(0), + b"k", + b"rec-0".to_vec(), + 1, + None, + ) .await .expect("first produce acked"); let second = client - .produce("produce-route", b"k", b"rec-1".to_vec(), 1) + .produce_to_range( + "produce-route", + RangeId(0), + b"k", + b"rec-1".to_vec(), + 1, + None, + ) .await .expect("second produce acked"); assert!( @@ -150,16 +164,16 @@ fn client_multiplexes_concurrent_requests() -> turmoil::Result { client.create_topic("mux", policy()).await.expect("create"); // Warm the cache + segment assignment so the burst below routes directly. let warm = client - .produce("mux", b"k", b"warm".to_vec(), 1) + .produce_to_range("mux", RangeId(0), b"k", b"warm".to_vec(), 1, None) .await .expect("warmup produce"); // Four produces + two describes concurrently over the shared connection. let (a, b, c, d, e, f) = tokio::join!( - client.produce("mux", b"k", b"p1".to_vec(), 1), - client.produce("mux", b"k", b"p2".to_vec(), 1), - client.produce("mux", b"k", b"p3".to_vec(), 1), - client.produce("mux", b"k", b"p4".to_vec(), 1), + client.produce_to_range("mux", RangeId(0), b"k", b"p1".to_vec(), 1, None), + client.produce_to_range("mux", RangeId(0), b"k", b"p2".to_vec(), 1, None), + client.produce_to_range("mux", RangeId(0), b"k", b"p3".to_vec(), 1, None), + client.produce_to_range("mux", RangeId(0), b"k", b"p4".to_vec(), 1, None), client.resolve_topic("mux"), client.resolve_topic("mux"), ); @@ -205,7 +219,7 @@ fn client_reconnects_after_node_crash() -> turmoil::Result { .await .expect("create"); client - .produce("reconnect", b"k", b"rec-0".to_vec(), 1) + .produce_to_range("reconnect", RangeId(0), b"k", b"rec-0".to_vec(), 1, None) .await .expect("produce"); @@ -286,7 +300,7 @@ fn client_corrects_stale_write_leader() -> turmoil::Result { // Starts at `wrong`, gets NotWriteLeader{real}, follows once, commits. client - .produce("stale", b"k", b"rec-1".to_vec(), 1) + .produce_to_range("stale", RangeId(0), b"k", b"rec-1".to_vec(), 1, None) .await .expect("produce self-corrects through the redirect"); @@ -296,7 +310,7 @@ fn client_corrects_stale_write_leader() -> turmoil::Result { "stale entry dropped after the corrected produce" ); client - .produce("stale", b"k", b"rec-2".to_vec(), 1) + .produce_to_range("stale", RangeId(0), b"k", b"rec-2".to_vec(), 1, None) .await .expect("produce after re-resolve"); assert_eq!( @@ -842,13 +856,15 @@ fn consumer_basic_consume_earliest() -> turmoil::Result { let mut sim = build_sim(90); host_cluster(&mut sim, &NODES, |env| { sim_cluster(env); - // Force fast size-based rolls - env.segment_size_limit_bytes = 10; + // The five encoded records cross this once near the end. Rolling every + // individual record turns this consumer test into a topology-race test. + env.segment_size_limit_bytes = 80; }); sim.client("test-client", async { let client = Arc::new(Client::connect(client_seeds()).expect("client connects")); - let custom_policy = policy(); + let mut custom_policy = policy(); + custom_policy.partition_strategy = PartitionStrategy::Fixed; client .create_topic("basic-consume", custom_policy) .await @@ -880,7 +896,7 @@ fn consumer_basic_consume_earliest() -> turmoil::Result { producer .send(key.as_bytes(), val) .await - .expect("produce record"); + .unwrap_or_else(|error| panic!("produce record {i}: {error:?}")); } // Consume them with Earliest @@ -1942,7 +1958,14 @@ fn producer_split_fence_retry() -> turmoil::Result { for i in 0..3 { client_admin - .produce(topic, b"k", b"longer_than_ten_bytes".to_vec(), 1) + .produce_to_range( + topic, + RangeId(0), + b"k", + b"longer_than_ten_bytes".to_vec(), + 1, + None, + ) .await .expect("produce trigger record"); wait_for_segment_roll(&client_admin, topic, 0, (i + 1) as u64).await; diff --git a/src/it/e2e/mod.rs b/src/it/e2e/mod.rs index a7a0d463..17efcc4e 100644 --- a/src/it/e2e/mod.rs +++ b/src/it/e2e/mod.rs @@ -1,6 +1,7 @@ mod client_protocol; mod client_sdk; mod consumer_group_test; +mod producer_deduplication; use crate::StartUp; use crate::config::Environment; diff --git a/src/it/sim/scenario.rs b/src/it/sim/scenario.rs index 2ea4f519..5aedc540 100644 --- a/src/it/sim/scenario.rs +++ b/src/it/sim/scenario.rs @@ -403,6 +403,7 @@ async fn produce_until_acked( routing_key: b"campaign-key".to_vec(), data: payload.to_vec(), record_count: 1, + producer: None, })); loop { for (host, port) in nodes { From 342b578b68a823f37ad4d6cb980a65a5a1c00198 Mon Sep 17 00:00:00 2001 From: Migorithm Date: Wed, 22 Jul 2026 12:59:02 +0400 Subject: [PATCH 09/41] doc update --- docs/clients/c2_producer.md | 16 ++++++++-------- .../d10_transparent_retry_deduplication.md | 5 +++-- 2 files changed, 11 insertions(+), 10 deletions(-) diff --git a/docs/clients/c2_producer.md b/docs/clients/c2_producer.md index dfd78be2..1dd4a57d 100644 --- a/docs/clients/c2_producer.md +++ b/docs/clients/c2_producer.md @@ -117,15 +117,15 @@ this is; the codec tag says *how to read* it. --- -## Delivery contract and idempotency seam +## Delivery contract -Today produce is explicitly **at-least-once**: a produce that times out after the -leader committed but before the client saw the acknowledgment can be retried and stored -twice. The UUID and per-record counter currently allocated by the producer do not cross -the wire and do not participate in a broker decision. +The SDK producer uses transparent retry deduplication. If an acknowledgment is lost after +commit, retrying the same immutable batch returns its original position without storing it +again. Producer identity, incarnation, sequence, and payload digest participate in the +broker decision and are durable with the append. -The future protocol is specified in -[D10: Idempotent Production](../data-plane/d10_idempotent_production.md). It assigns a +The protocol is specified in +[D10: Transparent Retry Deduplication](../data-plane/d10_transparent_retry_deduplication.md). It assigns a sequence to each immutable per-range broker batch and orders each producer-range stream, uses metadata-backed incarnations for fencing, and retains range-scoped deduplication frontiers across rolls, failover, and lineage changes. A bounded recent-result window @@ -166,6 +166,6 @@ stable producer UUID as a delivery guarantee. - `d6_produce_consume_api.md` — the server's produce routing and redirect contract. - `d1_storage_engine.md` — the broker-opaque entry payload and end-to-end compression this stamps a codec into. -- `../data-plane/d10_idempotent_production.md` — session fencing, ordered producer-range streams, +- `../data-plane/d10_transparent_retry_deduplication.md` — session fencing, ordered producer-range streams, durable range ledgers, topology handoff, and bounded retry guarantees. - `client_roadmap.md` — the idempotency / batching / compression backlog context. diff --git a/docs/data-plane/d10_transparent_retry_deduplication.md b/docs/data-plane/d10_transparent_retry_deduplication.md index d90c640d..c027ae7d 100644 --- a/docs/data-plane/d10_transparent_retry_deduplication.md +++ b/docs/data-plane/d10_transparent_retry_deduplication.md @@ -265,8 +265,9 @@ Produce requests add producer id, incarnation, producer-range sequence, and dige Successful responses distinguish newly appended, duplicate with position, and duplicate without position. -The SDK `Producer` uses this protocol by default. Low-level `Client::produce` remains an -explicit raw append and therefore remains at-least-once. +The SDK `Producer` is the production write API and uses this protocol by default. Routing +and topology tests can exercise the internal range-write path without producer identity; +that path is not exposed as a production client API. The implemented delivery contract is: From cb5ab928e8c23e0bbb1017e5efc0b9d7ad4697f4 Mon Sep 17 00:00:00 2001 From: Migorithm Date: Wed, 22 Jul 2026 13:09:47 +0400 Subject: [PATCH 10/41] command, wires --- src/client/producer/buffers.rs | 4 +-- src/connections/protocol/control_plane.rs | 18 ++++++++++++- src/connections/protocol/data_plane.rs | 3 +++ src/control_plane/consensus/messages/actor.rs | 2 +- src/control_plane/consensus/multi_raft.rs | 4 +-- src/control_plane/consensus/raft/command.rs | 8 +++--- src/control_plane/metadata/command.rs | 26 +++++++++++++++---- src/data_plane/checkpoint.rs | 2 +- src/data_plane/cold_read.rs | 1 + src/data_plane/messages/pending.rs | 2 +- src/data_plane/states/segment/cache.rs | 4 ++- 11 files changed, 56 insertions(+), 18 deletions(-) diff --git a/src/client/producer/buffers.rs b/src/client/producer/buffers.rs index a5c92b21..3b4b4ae3 100644 --- a/src/client/producer/buffers.rs +++ b/src/client/producer/buffers.rs @@ -5,7 +5,6 @@ use crate::{client::error::ClientError, control_plane::metadata::EntryId}; use dashmap::DashMap; use std::time::{Duration, Instant}; use tokio::sync::oneshot; -use uuid::Uuid; pub(crate) type Records = Box<[(Vec, Vec)]>; @@ -49,10 +48,9 @@ impl RangeBuffer { } pub struct PendingRecord { + pub order: u64, pub key: Vec, pub value: Vec, - pub producer_id: Uuid, - pub sequence_number: u32, pub tx: oneshot::Sender>, } impl PendingRecord { diff --git a/src/connections/protocol/control_plane.rs b/src/connections/protocol/control_plane.rs index 641ec1f8..b5b0819e 100644 --- a/src/connections/protocol/control_plane.rs +++ b/src/connections/protocol/control_plane.rs @@ -39,13 +39,22 @@ pub enum ControlPlaneRequest { name: String, }, SyncConsumerGroup(SyncConsumerGroupRequest), + OpenProducerSession(OpenProducerSessionRequest), } impl_from_variant!( ControlPlaneRequest, - SyncConsumerGroup(SyncConsumerGroupRequest) + SyncConsumerGroup(SyncConsumerGroupRequest), + OpenProducerSession(OpenProducerSessionRequest) ); +#[derive(Debug, Clone, BorshSerialize, BorshDeserialize)] +pub struct OpenProducerSessionRequest { + pub topic_name: String, + pub producer_id: uuid::Uuid, + pub session_nonce: uuid::Uuid, +} + #[derive(Debug, Clone, Copy, PartialEq, Eq, BorshSerialize, BorshDeserialize)] pub enum ConsumerGroupSyncAction { Heartbeat, @@ -77,10 +86,17 @@ pub enum ControlPlaneResponse { }, ConsumerGroupAssignment(ConsumerGroupAssignmentResponse), ConsumerGroupLeft, + ProducerSessionOpened(ProducerSessionOpened), // All control plane operations InternalError(String), } +#[derive(Debug, Clone, Copy, BorshSerialize, BorshDeserialize)] +pub struct ProducerSessionOpened { + pub incarnation: u32, + pub expires_at: u64, +} + #[derive(Debug, BorshSerialize, BorshDeserialize)] pub struct ConsumerGroupAssignmentResponse { pub generation: GenerationId, diff --git a/src/connections/protocol/data_plane.rs b/src/connections/protocol/data_plane.rs index c07740c6..7508fb74 100644 --- a/src/connections/protocol/data_plane.rs +++ b/src/connections/protocol/data_plane.rs @@ -17,6 +17,7 @@ use crate::{ }, }, data_plane::{ + ProduceError, ProducerAppendIdentity, consumer_offset_management::ledger::{ ConsumerOffsetKey, ConsumerOffsetPosition, ConsumerOffsetUpdate, }, @@ -60,6 +61,7 @@ pub struct ProduceRequest { /// The broker stamps an entry_id and stores/replicates this opaque payload as-is. pub data: Vec, pub record_count: u32, + pub producer_append_id: Option, } #[derive(Debug, Clone, BorshSerialize, BorshDeserialize)] @@ -115,6 +117,7 @@ pub enum DataPlaneResponse { Produced { entry_id: EntryId, }, + ProduceRejected(ProduceError), // Fetch Fetched { entries: Box<[Entry]>, diff --git a/src/control_plane/consensus/messages/actor.rs b/src/control_plane/consensus/messages/actor.rs index 7124a5d3..d123a6de 100644 --- a/src/control_plane/consensus/messages/actor.rs +++ b/src/control_plane/consensus/messages/actor.rs @@ -110,6 +110,6 @@ pub(crate) enum DeferredReply { ), GetTopics(oneshot::Sender>, Box<[String]>), GetTopicStats(oneshot::Sender>, Box<[TopicStats]>), - GetTopicMetadata(oneshot::Sender>, Option), + GetTopicMetadata(oneshot::Sender>, Box>), GetConsumerGroupAssignment(DeferredConsumerGroupAssignment), } diff --git a/src/control_plane/consensus/multi_raft.rs b/src/control_plane/consensus/multi_raft.rs index d83d4f2c..4fe31c65 100644 --- a/src/control_plane/consensus/multi_raft.rs +++ b/src/control_plane/consensus/multi_raft.rs @@ -259,7 +259,7 @@ impl MultiRaft { MultiRaftActorCommand::GetTopicMetadata { topic_name, reply } => { let meta = self.get_topic_metadata(&topic_name); self.deferred - .push(DeferredReply::GetTopicMetadata(reply, meta)); + .push(DeferredReply::GetTopicMetadata(reply, Box::new(meta))); } MultiRaftActorCommand::GetConsumerGroupAssignment(query) => { let value = self.get_consumer_group_assignment( @@ -309,7 +309,7 @@ impl MultiRaft { let _ = sender.send(v); } DeferredReply::GetTopicMetadata(sender, v) => { - let _ = sender.send(v); + let _ = sender.send(*v); } DeferredReply::GetConsumerGroupAssignment(deferred) => { let _ = deferred.reply.send(deferred.value); diff --git a/src/control_plane/consensus/raft/command.rs b/src/control_plane/consensus/raft/command.rs index ed2e933d..6d67e313 100644 --- a/src/control_plane/consensus/raft/command.rs +++ b/src/control_plane/consensus/raft/command.rs @@ -3,8 +3,8 @@ use borsh::{BorshDeserialize, BorshSerialize}; use crate::control_plane::NodeId; use crate::control_plane::metadata::ReassignSegment; use crate::control_plane::metadata::command::{ - CreateTopic, DeleteSegments, DeleteTopic, MergeRange, MetadataCommand, RollSegment, SplitRange, - SyncConsumerGroup, + CreateTopic, DeleteSegments, DeleteTopic, ExpireProducerSessions, MergeRange, MetadataCommand, + OpenProducerSession, RollSegment, SplitRange, SyncConsumerGroup, }; use crate::{impl_from_variant, impl_from_variant_via}; @@ -31,5 +31,7 @@ impl_from_variant_via!( DeleteTopic, ReassignSegment, DeleteSegments, - SyncConsumerGroup + SyncConsumerGroup, + OpenProducerSession, + ExpireProducerSessions ); diff --git a/src/control_plane/metadata/command.rs b/src/control_plane/metadata/command.rs index ae4e77e8..de8395b4 100644 --- a/src/control_plane/metadata/command.rs +++ b/src/control_plane/metadata/command.rs @@ -91,6 +91,21 @@ pub struct SyncConsumerGroup { pub session_timeout_ms: u64, } +#[derive(Debug, Clone, PartialEq, Eq, BorshSerialize, BorshDeserialize)] +pub struct OpenProducerSession { + pub topic_name: String, + pub producer_id: Uuid, + pub session_nonce: Uuid, + pub observed_at: u64, + pub session_timeout_ms: u64, +} + +#[derive(Debug, Clone, PartialEq, Eq, BorshSerialize, BorshDeserialize)] +pub struct ExpireProducerSessions { + pub topic_id: TopicId, + pub observed_at: u64, +} + #[derive(Debug, Clone, PartialEq, Eq, BorshSerialize, BorshDeserialize)] pub struct SyncConsumerGroupRequest { pub topic_name: String, @@ -110,10 +125,7 @@ impl Deref for SyncConsumerGroup { impl SyncConsumerGroup { pub(crate) fn new(req: SyncConsumerGroupRequest) -> Self { const SESSION_TIMEOUT_MS: u64 = 10_000; - let observed_at = std::time::SystemTime::now() - .duration_since(std::time::UNIX_EPOCH) - .unwrap_or_default() - .as_millis() as u64; + let observed_at = crate::now_ms(); SyncConsumerGroup { req, observed_at, @@ -132,6 +144,8 @@ pub enum MetadataCommand { ReassignSegment(ReassignSegment), DeleteSegments(DeleteSegments), SyncConsumerGroup(SyncConsumerGroup), + OpenProducerSession(OpenProducerSession), + ExpireProducerSessions(ExpireProducerSessions), } impl_from_variant!( @@ -143,5 +157,7 @@ impl_from_variant!( DeleteTopic, ReassignSegment, DeleteSegments, - SyncConsumerGroup + SyncConsumerGroup, + OpenProducerSession, + ExpireProducerSessions ); diff --git a/src/data_plane/checkpoint.rs b/src/data_plane/checkpoint.rs index 48a0918c..bb90d983 100644 --- a/src/data_plane/checkpoint.rs +++ b/src/data_plane/checkpoint.rs @@ -136,7 +136,7 @@ pub(crate) enum CheckpointTask { DeleteSegmentIndex(SegmentKey), /// Snapshot the consumer-offset cache after its corresponding shared-WAL /// batch is durable. This runs off the data-plane write path. - ConsumerOffsets(OffsetCheckpointJob), + ConsumerOffsets(Box), } pub struct CheckpointJob { diff --git a/src/data_plane/cold_read.rs b/src/data_plane/cold_read.rs index 386aab84..2d3f7d5a 100644 --- a/src/data_plane/cold_read.rs +++ b/src/data_plane/cold_read.rs @@ -252,6 +252,7 @@ impl ColdReadPool { // lsn is a live-WAL concept; cold reads serve durable // data, so it carries no meaningful LSN. lsn: 0, + producer_append_id: None, })); if bytes_read >= req.max_bytes { diff --git a/src/data_plane/messages/pending.rs b/src/data_plane/messages/pending.rs index 6ba1eab4..cb9cf57b 100644 --- a/src/data_plane/messages/pending.rs +++ b/src/data_plane/messages/pending.rs @@ -101,7 +101,7 @@ impl DataPlaneOutputs { pub(crate) fn store_offset_checkpoint(&mut self, job: OffsetCheckpointJob) { self.checkpoint_tasks - .push(CheckpointTask::ConsumerOffsets(job)); + .push(CheckpointTask::ConsumerOffsets(Box::new(job))); } pub(crate) fn store_put_anchors(&mut self, anchors: Vec) { diff --git a/src/data_plane/states/segment/cache.rs b/src/data_plane/states/segment/cache.rs index 81a68291..b3c7ce25 100644 --- a/src/data_plane/states/segment/cache.rs +++ b/src/data_plane/states/segment/cache.rs @@ -7,7 +7,7 @@ use arc_swap::ArcSwapOption; use tokio::sync::Notify; use crate::control_plane::metadata::EntryId; -use crate::data_plane::EntryPayload; +use crate::data_plane::{EntryPayload, ProducerAppendIdentity}; #[derive(Debug, Clone)] pub(crate) struct CachedEntry { @@ -15,6 +15,7 @@ pub(crate) struct CachedEntry { pub(crate) record_count: u32, pub(crate) entry_id: EntryId, pub(crate) lsn: u64, + pub(crate) producer_append_id: Option, } impl CachedEntry { @@ -331,6 +332,7 @@ mod tests { record_count: 1, entry_id: EntryId(entry_id), lsn, + producer_append_id: None, }) } From c58bc393bb761a5c73901bbecbc3ce42af45c18a Mon Sep 17 00:00:00 2001 From: Migorithm Date: Wed, 22 Jul 2026 13:11:43 +0400 Subject: [PATCH 11/41] doc: no legacy snapshot --- docs/data-plane/d10_transparent_retry_deduplication.md | 6 ++---- docs/data-plane/d8_consumer_offset_management.md | 7 ++----- 2 files changed, 4 insertions(+), 9 deletions(-) diff --git a/docs/data-plane/d10_transparent_retry_deduplication.md b/docs/data-plane/d10_transparent_retry_deduplication.md index c027ae7d..0d41e5f8 100644 --- a/docs/data-plane/d10_transparent_retry_deduplication.md +++ b/docs/data-plane/d10_transparent_retry_deduplication.md @@ -284,7 +284,7 @@ The implemented delivery contract is: idempotent. - Request identity is part of the replicated data entry and the shared WAL batch. There is no producer-only WAL record or additional fsync. Consumer offsets and producer ledgers - share `auxiliary-state.snapshot`; startup migrates the legacy consumer-only snapshot. + share `auxiliary-state.snapshot`. - A retry against a sealed split/merge parent is routed to the frozen parent ledger. A committed request returns its old position; a proven next sequence returns stale range so the SDK can repartition it to children. @@ -319,7 +319,6 @@ reachable. An expired request is never accepted merely because its ledger was co successor returns the original position for the unresolved retry. - Restart all replicas and reconstruct the same frontier and recent window from the consolidated snapshot plus committed WAL suffix. -- Migrate a legacy consumer-only snapshot before reclaiming covered WAL. ### Topology transitions @@ -336,8 +335,7 @@ the repository's data-plane test rules. ## Implementation shape 1. **Storage and WAL framing.** Extend entry identity, generalize the consumer-only - checkpoint into the typed D9 auxiliary-state snapshot, and support legacy snapshot - migration. + checkpoint into the typed D9 auxiliary-state snapshot. 2. **Local deduplication.** Add the frontier plus recent-result ring and the range-leader verification pipeline; change client batching to assign immutable per-range sequences. 3. **Session lifecycle and fencing.** Add Raft-backed open, recover, renew, expire, and diff --git a/docs/data-plane/d8_consumer_offset_management.md b/docs/data-plane/d8_consumer_offset_management.md index a836b67d..b4cdbcbe 100644 --- a/docs/data-plane/d8_consumer_offset_management.md +++ b/docs/data-plane/d8_consumer_offset_management.md @@ -211,11 +211,8 @@ pipeline: with one auxiliary-state checkpoint lifecycle. 3. Keep offset commits and graduation imports in the shared WAL; their WAL fsync is still the acknowledgement durability boundary. -4. On upgrade, load the legacy consumer-offset snapshot when no consolidated snapshot exists, - combine it with producer frontier state recovered from WAL, and publish a durable - consolidated snapshot before reclaiming either the legacy file or covered WAL. -5. Remove the legacy snapshot only after the consolidated snapshot's durable rename and - directory sync complete. +4. Reclaim covered WAL only after the consolidated snapshot's durable rename and directory + sync complete. This cleanup removes the independent consumer-snapshot fsync once D10 adds producer state. It does not remove the shared-WAL fsync required before acknowledging an offset commit, nor the From 5072eaa57c1ca948bb1d378c3aa577e7a1bb7eec Mon Sep 17 00:00:00 2001 From: Migorithm Date: Wed, 22 Jul 2026 13:14:43 +0400 Subject: [PATCH 12/41] producer -> producer_append_id, remove raw produce --- src/client/mod.rs | 37 +++++++++++++++++----------- src/it/e2e/client_protocol.rs | 2 +- src/it/e2e/producer_deduplication.rs | 2 +- src/it/sim/scenario.rs | 2 +- 4 files changed, 25 insertions(+), 18 deletions(-) diff --git a/src/client/mod.rs b/src/client/mod.rs index ad32e4da..6efa845f 100644 --- a/src/client/mod.rs +++ b/src/client/mod.rs @@ -45,9 +45,10 @@ use uuid::Uuid; use crate::connections::protocol::{ ClientDataPlaneRequest, ClientRequest, ClientResponse, CommitConsumerOffsetRequest, ConsumerGroupAssignmentResponse, ConsumerGroupSyncAction, ControlPlaneRequest, - ControlPlaneResponse, DataPlaneResponse, FetchConsumerOffsetRequest, ProduceRequest, - RangeOffsetRequest, + ControlPlaneResponse, DataPlaneResponse, FetchConsumerOffsetRequest, + OpenProducerSessionRequest, ProduceRequest, ProducerSessionOpened, RangeOffsetRequest, }; +use crate::data_plane::{ProduceError, ProducerAppendIdentity}; use futures::{StreamExt, stream::FuturesUnordered}; pub use redirect::RetryPolicy; @@ -311,20 +312,19 @@ impl Client { } } - pub(crate) async fn produce( + pub(crate) async fn open_producer_session( &self, - topic: &str, - routing_key: &[u8], - data: Vec, - record_count: u32, - ) -> Result { - let routing = self.resolve_topic_if_missing(topic).await?; - let range_id = routing - .range_id(routing_key) - .ok_or(ClientError::TopicNotFound)?; - - self.produce_to_range(topic, range_id, routing_key, data, record_count) - .await + req: OpenProducerSessionRequest, + ) -> Result { + let served = self + .call(self.next_known_node(), ControlPlaneRequest::from(req)) + .await?; + match served.response { + ClientResponse::ControlPlane(ControlPlaneResponse::ProducerSessionOpened(session)) => { + Ok(session) + } + _ => Err(ClientError::UnexpectedResponse), + } } /// Produce one entry under `routing_key`, returning the committed `entry_id`. @@ -337,6 +337,7 @@ impl Client { routing_key: &[u8], data: Vec, record_count: u32, + producer: Option, ) -> Result { // Describe once to seed the cache (gives the first hop). let routing = self.resolve_topic_if_missing(topic).await?; @@ -354,6 +355,7 @@ impl Client { routing_key: routing_key.to_vec(), data, record_count, + producer_append_id: producer, }; let served = self.call(start, request).await?; @@ -372,6 +374,9 @@ impl Client { ClientResponse::DataPlane(DataPlaneResponse::StaleRange) => { Err(ClientError::StaleRange) } + ClientResponse::DataPlane(DataPlaneResponse::ProduceRejected(error)) => { + Err(ClientError::ProduceRejected(error)) + } _ => Err(ClientError::UnexpectedResponse), } } @@ -489,6 +494,7 @@ impl Client { | ControlPlaneResponse::TopicDeleted | ControlPlaneResponse::ConsumerGroupAssignment(_) | ControlPlaneResponse::ConsumerGroupLeft + | ControlPlaneResponse::ProducerSessionOpened(_) | ControlPlaneResponse::TopicList { .. } | ControlPlaneResponse::TopicDetail(_) => Redirect::Done, }, @@ -506,6 +512,7 @@ impl Client { DataPlaneResponse::InternalError(_) => Redirect::Reresolve, DataPlaneResponse::SegmentNotLocal | DataPlaneResponse::Produced { .. } + | DataPlaneResponse::ProduceRejected(..) | DataPlaneResponse::Fetched { .. } | DataPlaneResponse::EntryIdOutOfRange | DataPlaneResponse::KeyspaceBoundNarrowed diff --git a/src/it/e2e/client_protocol.rs b/src/it/e2e/client_protocol.rs index 56a537a6..ae8445c4 100644 --- a/src/it/e2e/client_protocol.rs +++ b/src/it/e2e/client_protocol.rs @@ -1832,7 +1832,7 @@ async fn produce_once(topic: &str, payload: &[u8], node: (&str, u16)) -> Produce routing_key: b"k".to_vec(), data: payload.to_vec(), record_count: 1, - producer: None, + producer_append_id: None, })); match send_request(node.0, node.1, req).await { ClientResponse::DataPlane(DataPlaneResponse::Produced { .. }) => ProduceOutcome::Acked, diff --git a/src/it/e2e/producer_deduplication.rs b/src/it/e2e/producer_deduplication.rs index f2188b46..3f2ea61b 100644 --- a/src/it/e2e/producer_deduplication.rs +++ b/src/it/e2e/producer_deduplication.rs @@ -118,7 +118,7 @@ fn duplicate_unknown_and_fenced_sessions_are_end_to_end_visible() -> turmoil::Re routing_key: b"k".to_vec(), data: payload.clone(), record_count: 1, - producer: Some(unknown), + producer_append_id: Some(unknown), })); let mut rejected = false; for (host, port, _) in NODES { diff --git a/src/it/sim/scenario.rs b/src/it/sim/scenario.rs index 5aedc540..2e1f7ae9 100644 --- a/src/it/sim/scenario.rs +++ b/src/it/sim/scenario.rs @@ -403,7 +403,7 @@ async fn produce_until_acked( routing_key: b"campaign-key".to_vec(), data: payload.to_vec(), record_count: 1, - producer: None, + producer_append_id: None, })); loop { for (host, port) in nodes { From f2215525eeb69e8364c6e12519e14d80da2838c8 Mon Sep 17 00:00:00 2001 From: Migorithm Date: Wed, 22 Jul 2026 18:28:16 +0400 Subject: [PATCH 13/41] merge into Auxiliary state --- src/client/consumer/group.rs | 2 +- src/client/consumer/mod.rs | 2 +- src/client/consumer/range_fetcher.rs | 2 +- src/client/consumer/topic_fetch_manager.rs | 2 +- src/client/mod.rs | 2 +- src/connections/protocol/data_plane.rs | 2 +- src/control_plane/metadata/event.rs | 2 +- .../auxiliary_states/consumer_offsets/mod.rs | 208 +++++++++++ .../consumer_offsets}/replication.rs | 4 +- .../consumer_offsets/state.rs} | 55 ++- .../consumer_offsets}/types.rs | 4 +- .../manager.rs | 339 ++++++------------ src/data_plane/auxiliary_states/mod.rs | 3 + src/data_plane/checkpoint.rs | 12 +- .../consumer_offset_management/mod.rs | 6 - src/data_plane/messages/pending.rs | 6 +- src/data_plane/messages/query.rs | 2 +- src/data_plane/mod.rs | 2 +- src/data_plane/recovery/inventory.rs | 6 +- src/it/e2e/client_sdk.rs | 2 +- 20 files changed, 394 insertions(+), 269 deletions(-) create mode 100644 src/data_plane/auxiliary_states/consumer_offsets/mod.rs rename src/data_plane/{consumer_offset_management => auxiliary_states/consumer_offsets}/replication.rs (95%) rename src/data_plane/{consumer_offset_management/ledger.rs => auxiliary_states/consumer_offsets/state.rs} (83%) rename src/data_plane/{consumer_offset_management => auxiliary_states/consumer_offsets}/types.rs (95%) rename src/data_plane/{consumer_offset_management => auxiliary_states}/manager.rs (77%) create mode 100644 src/data_plane/auxiliary_states/mod.rs delete mode 100644 src/data_plane/consumer_offset_management/mod.rs diff --git a/src/client/consumer/group.rs b/src/client/consumer/group.rs index aee91e70..2d9026e8 100644 --- a/src/client/consumer/group.rs +++ b/src/client/consumer/group.rs @@ -11,7 +11,7 @@ use crate::client::{Client, ClientError}; use crate::connections::protocol::{ClientResponse, ConsumerGroupSyncAction, ControlPlaneResponse}; use crate::control_plane::metadata::consumer_group::GenerationId; use crate::control_plane::metadata::{EntryId, RangeId, SyncConsumerGroupRequest, TopicId}; -use crate::data_plane::consumer_offset_management::ledger::{ +use crate::data_plane::auxiliary_states::consumer_offsets::state::{ ConsumerOffsetKey, ConsumerOffsetPosition, ConsumerOffsetUpdate, }; diff --git a/src/client/consumer/mod.rs b/src/client/consumer/mod.rs index 33e26404..c6122aa3 100644 --- a/src/client/consumer/mod.rs +++ b/src/client/consumer/mod.rs @@ -14,7 +14,7 @@ use crate::connections::protocol::{ RangeOffsetRequest, RangeProgressSignal, RangeTransition, SegmentDetail, TopicDetail, }; use crate::control_plane::metadata::{EntryId, RangeId, RangeState, TopicId}; -use crate::data_plane::consumer_offset_management::ledger::ConsumerOffsetPosition; +use crate::data_plane::auxiliary_states::consumer_offsets::state::ConsumerOffsetPosition; pub mod config; pub use config::*; diff --git a/src/client/consumer/range_fetcher.rs b/src/client/consumer/range_fetcher.rs index 25b01335..2ff15084 100644 --- a/src/client/consumer/range_fetcher.rs +++ b/src/client/consumer/range_fetcher.rs @@ -10,7 +10,7 @@ use crate::connections::protocol::{ ClientResponse, DataPlaneResponse, RangeProgressSignal, RangeTransition, SegmentDetail, }; use crate::control_plane::metadata::{EntryId, RangeId}; -use crate::data_plane::consumer_offset_management::ledger::ConsumerOffsetPosition; +use crate::data_plane::auxiliary_states::consumer_offsets::state::ConsumerOffsetPosition; const EMPTY_FETCHES_BEFORE_REFRESH: u8 = 20; diff --git a/src/client/consumer/topic_fetch_manager.rs b/src/client/consumer/topic_fetch_manager.rs index 2b814854..428520ee 100644 --- a/src/client/consumer/topic_fetch_manager.rs +++ b/src/client/consumer/topic_fetch_manager.rs @@ -10,7 +10,7 @@ use crate::client::consumer::{ use crate::client::{ClientError, CommitMode, ConsumerConfig}; use crate::connections::protocol::{RangeDetail, RangeTransition, RebalancePlan}; use crate::control_plane::metadata::{EntryId, RangeId, RangeState}; -use crate::data_plane::consumer_offset_management::ledger::ConsumerOffsetPosition; +use crate::data_plane::auxiliary_states::consumer_offsets::state::ConsumerOffsetPosition; use crate::impl_from_variant; #[cfg(any(test, debug_assertions))] diff --git a/src/client/mod.rs b/src/client/mod.rs index 6efa845f..4c5fdc04 100644 --- a/src/client/mod.rs +++ b/src/client/mod.rs @@ -29,7 +29,7 @@ use crate::control_plane::metadata::consumer_group::GenerationId; pub use crate::control_plane::metadata::strategy::{PartitionStrategy, StoragePolicy}; pub use crate::control_plane::metadata::{EntryId, RangeId}; use crate::control_plane::metadata::{SyncConsumerGroupRequest, TopicId}; -use crate::data_plane::consumer_offset_management::ledger::{ +use crate::data_plane::auxiliary_states::consumer_offsets::state::{ ConsumerOffsetKey, ConsumerOffsetPosition, }; pub use codec::CompressionCodec; diff --git a/src/connections/protocol/data_plane.rs b/src/connections/protocol/data_plane.rs index 7508fb74..11845fb0 100644 --- a/src/connections/protocol/data_plane.rs +++ b/src/connections/protocol/data_plane.rs @@ -18,7 +18,7 @@ use crate::{ }, data_plane::{ ProduceError, ProducerAppendIdentity, - consumer_offset_management::ledger::{ + auxiliary_states::consumer_offsets::state::{ ConsumerOffsetKey, ConsumerOffsetPosition, ConsumerOffsetUpdate, }, messages::query::{FetchResult, ListOffsetsResult}, diff --git a/src/control_plane/metadata/event.rs b/src/control_plane/metadata/event.rs index 3bf36d0e..0765f076 100644 --- a/src/control_plane/metadata/event.rs +++ b/src/control_plane/metadata/event.rs @@ -4,7 +4,7 @@ use crate::control_plane::membership::ShardGroupId; use crate::control_plane::metadata::consumer_group::GenerationId; use crate::control_plane::metadata::{EntryId, RangeId, SegmentId, TopicId}; use crate::data_plane::SegmentKey; -use crate::data_plane::consumer_offset_management::ledger::{ConsumerOffsetKey, EpochSeal}; +use crate::data_plane::auxiliary_states::consumer_offsets::state::{ConsumerOffsetKey, EpochSeal}; use crate::data_plane::messages::command::{ AssignSegmentCatchUp, DataPlanePeerMessage, DeleteSegments, PlaceSegment, SegmentMetaSealed, SegmentRollCommitted, diff --git a/src/data_plane/auxiliary_states/consumer_offsets/mod.rs b/src/data_plane/auxiliary_states/consumer_offsets/mod.rs new file mode 100644 index 00000000..da9cb059 --- /dev/null +++ b/src/data_plane/auxiliary_states/consumer_offsets/mod.rs @@ -0,0 +1,208 @@ +pub(crate) mod replication; +pub(crate) mod state; +pub(crate) mod types; +use crate::control_plane::NodeId; +use crate::control_plane::Replicas; +use crate::control_plane::membership::ShardGroupId; +use crate::control_plane::metadata::RangeId; +use crate::control_plane::metadata::TopicId; +use crate::control_plane::metadata::consumer_group::GenerationId; +use crate::data_plane::SegmentKey; +use crate::data_plane::messages::ConsumerOffsetCommitAck; +use crate::data_plane::messages::ConsumerOffsetReplicated; +use replication::OffsetReplicationState; +use state::*; +use std::collections::HashMap; +use std::collections::HashSet; +use tokio::sync::oneshot; +use types::*; + +#[derive(Default)] +pub(crate) struct ConsumerOffsetCoordination { + pub(crate) placements: HashMap<(TopicId, RangeId), OffsetPlacement>, + pending_offset_mutations: Vec, + future_offset_commits: HashMap>, + offset_replication: OffsetReplicationState, +} + +impl ConsumerOffsetCoordination { + pub(crate) fn future_entry(&mut self, key: ConsumerOffsetKey) -> &mut Vec { + self.future_offset_commits.entry(key).or_default() + } + + pub(crate) fn push_pending(&mut self, entry: PendingOffsetMutation) { + self.pending_offset_mutations.push(entry); + } + + pub(crate) fn has_pending_replication_for(&self, node: &NodeId) -> bool { + self.offset_replication.has_pending_replication_for(node) + } + + pub(crate) fn ack_replication(&mut self, ack: ConsumerOffsetReplicated) { + self.offset_replication.process_ack(ack); + } + + pub(crate) fn latest_generation(&self, key: &ConsumerOffsetKey) -> Option { + self.pending_offset_mutations + .iter() + .rev() // Look backwards through the pending queue (newest first) + .find_map(|pending| match &pending.record { + OffsetRecord::EpochSeal(EpochSeal { + key: pending_key, + generation, + }) if pending_key == key => Some(*generation), + OffsetRecord::BootstrapEntry(snapshot) if &snapshot.key == key => { + Some(snapshot.generation) + } + _ => None, + }) + } + + pub(crate) fn encode_offsets(&mut self) -> Result>, std::io::Error> { + match self + .pending_offset_mutations + .iter() + .map(|pending| borsh::to_vec(&pending.record)) + .collect() + { + Ok(encoded) => Ok(encoded), + Err(error) => { + tracing::error!(?error, "consumer offset WAL encoding failed"); + for pending in std::mem::take(&mut self.pending_offset_mutations) { + if let OffsetMutationCompletion::LeaderCommit(commit) = pending.completion { + let _ = commit + .reply + .send(ConsumerOffsetCommitAck::InternalError(error.to_string())); + } + } + Err(error) + } + } + } + + pub(crate) fn begin_replication( + &mut self, + followers: HashSet, + required: HashSet, + update: ConsumerOffsetUpdate, + reply: oneshot::Sender, + ) -> u64 { + self.offset_replication + .begin(followers, required, update, reply) + } + + pub(crate) fn fail_all(&mut self, error: &str) { + self.offset_replication.fail_all(error); + + for pending in self.take_pending() { + if let OffsetMutationCompletion::LeaderCommit(commit) = pending.completion { + let _ = commit + .reply + .send(ConsumerOffsetCommitAck::InternalError(error.to_string())); + } + } + + for parked in self + .future_offset_commits + .drain() + .flat_map(|(_, commits)| commits) + { + if let FutureOffsetCommit::Client(commit) = parked { + let _ = commit + .reply + .send(ConsumerOffsetCommitAck::InternalError(error.to_string())); + } + } + } + + // ? why per one consumer offset key, multiple future offset commit + pub(crate) fn take_future_commits( + &mut self, + key: &ConsumerOffsetKey, + ) -> Vec { + self.future_offset_commits.remove(key).unwrap_or_default() + } + + pub(crate) fn take_pending(&mut self) -> Vec { + std::mem::take(&mut self.pending_offset_mutations) + } + + pub(crate) fn has_pending_offsets(&self) -> bool { + !self.pending_offset_mutations.is_empty() + } + + #[cfg(test)] + pub fn is_empty(&self) -> bool { + self.pending_offset_mutations.is_empty() && self.future_offset_commits.is_empty() + } +} + +pub(crate) struct OffsetPlacement { + pub(crate) segment_key: SegmentKey, + pub(crate) shard_group_id: ShardGroupId, + // Desired data replicas for this placement. This also preserves their + // ordering, including which replica is the leader. + pub(crate) replicas: Replicas, + // Consumer-offset readiness for every member of `replicas`. + pub(crate) replica_states: HashMap, + pub(crate) placement_ack_sent: bool, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum OffsetReplicaState { + Joining, + SnapshotInstalled, + Ready, +} + +impl OffsetPlacement { + pub(crate) fn leader(&self) -> &NodeId { + self.replicas + .leader() + .expect("an offset placement must have a leader") + } + + pub(crate) fn can_bootstrap_replicas(&self, segment_key: &SegmentKey) -> bool { + self.segment_key == *segment_key && self.is_replica_ready(self.leader()) + } + + pub(crate) fn is_replica_ready(&self, node: &NodeId) -> bool { + self.replica_states.get(node) == Some(&OffsetReplicaState::Ready) + } + + pub(crate) fn all_replicas_ready(&self) -> bool { + self.replica_states + .values() + .all(|state| *state == OffsetReplicaState::Ready) + } + + pub(crate) fn compare(&self, other_key: &SegmentKey, other: &Replicas) -> PlacementObservation { + if self.segment_key == *other_key { + if &self.replicas != other { + tracing::warn!( + ?other_key, + current_replicas = ?self.replicas, + observed_replicas = ?other, + "ignored conflicting replica set for an existing segment" + ); + } + return if &self.replicas == other { + PlacementObservation::Unchanged + } else { + PlacementObservation::Conflict + }; + } + if self.segment_key.segment_id >= other_key.segment_id { + return PlacementObservation::Stale; + } + PlacementObservation::Accepted + } + + pub(crate) fn unready_followers(&self) -> Box<[NodeId]> { + self.replicas + .followers() + .filter(|node| !self.is_replica_ready(node)) + .cloned() + .collect() + } +} diff --git a/src/data_plane/consumer_offset_management/replication.rs b/src/data_plane/auxiliary_states/consumer_offsets/replication.rs similarity index 95% rename from src/data_plane/consumer_offset_management/replication.rs rename to src/data_plane/auxiliary_states/consumer_offsets/replication.rs index 5be7a558..5ea0b90d 100644 --- a/src/data_plane/consumer_offset_management/replication.rs +++ b/src/data_plane/auxiliary_states/consumer_offsets/replication.rs @@ -1,5 +1,5 @@ +use super::state::ConsumerOffsetUpdate; use crate::control_plane::NodeId; -use crate::data_plane::consumer_offset_management::ledger::ConsumerOffsetUpdate; use crate::data_plane::messages::command::{ ConsumerOffsetCommitAck, ConsumerOffsetReplicated, ConsumerOffsetReplicationResult, }; @@ -98,7 +98,7 @@ impl OffsetReplicationState { } } - pub(crate) fn has_pending_for(&self, node_id: &NodeId) -> bool { + pub(crate) fn has_pending_replication_for(&self, node_id: &NodeId) -> bool { self.pending .values() .any(|pending| pending.pending_acks.contains(node_id)) diff --git a/src/data_plane/consumer_offset_management/ledger.rs b/src/data_plane/auxiliary_states/consumer_offsets/state.rs similarity index 83% rename from src/data_plane/consumer_offset_management/ledger.rs rename to src/data_plane/auxiliary_states/consumer_offsets/state.rs index eb4bebd8..6c35c30f 100644 --- a/src/data_plane/consumer_offset_management/ledger.rs +++ b/src/data_plane/auxiliary_states/consumer_offsets/state.rs @@ -9,10 +9,12 @@ use borsh::{BorshDeserialize, BorshSerialize}; use crate::client::RangeId; use crate::control_plane::metadata::consumer_group::GenerationId; use crate::control_plane::metadata::{EntryId, TopicId}; -use crate::data_plane::SegmentKey; +use crate::data_plane::messages::command::AuthorizedProducerIdentity; +use crate::data_plane::producer_ledger::{AppendKey, ProducerDecision, ProducerLedger}; +use crate::data_plane::{ProducerAppendIdentity, SegmentKey}; -const SNAPSHOT_FILE: &str = "consumer-offsets.snapshot"; -const SNAPSHOT_TEMP_FILE: &str = "consumer-offsets.snapshot.tmp"; +const SNAPSHOT_FILE: &str = "auxiliary-state.snapshot"; +const SNAPSHOT_TEMP_FILE: &str = "auxiliary-state.snapshot.tmp"; #[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash, BorshSerialize, BorshDeserialize)] pub(crate) struct ConsumerOffsetKey { @@ -81,20 +83,47 @@ pub(crate) enum OffsetRecord { /// data-plane WAL; this type is only the in-memory cache and its asynchronous /// WAL-reclamation snapshot. #[derive(Debug, Clone, Default, BorshSerialize, BorshDeserialize)] -pub(crate) struct OffsetLedger { +pub(crate) struct AuxiliaryState { epochs: HashMap, offsets: HashMap, installed_placements: HashMap<(TopicId, RangeId), SegmentKey>, + producer_ledger: ProducerLedger, } -impl OffsetLedger { +impl AuxiliaryState { + pub(crate) fn verify_producer( + &mut self, + segment_key: SegmentKey, + producer: AuthorizedProducerIdentity, + ) -> ProducerDecision { + self.producer_ledger.verify(segment_key, producer) + } + + pub(crate) fn unstage_producer(&mut self, append_key: &AppendKey) { + self.producer_ledger.unstage(append_key); + } + + pub(crate) fn apply_producer( + &mut self, + segment_key: SegmentKey, + producer: ProducerAppendIdentity, + entry_id: EntryId, + ) { + self.producer_ledger.commit(segment_key, producer, entry_id); + } + + #[cfg(any(test, debug_assertions))] + pub(crate) fn assert_producer_invariants(&self) { + crate::test_traits::TAssertInvariant::assert_invariants(&self.producer_ledger); + } pub(crate) fn load_snapshot(data_dir: &Path) -> io::Result { let path = data_dir.join(SNAPSHOT_FILE); - match fs::read(path) { - Ok(bytes) => Self::try_from_slice(&bytes).map_err(io::Error::other), - Err(error) if error.kind() == io::ErrorKind::NotFound => Ok(Self::default()), - Err(error) => Err(error), - } + let bytes = match fs::read(path) { + Ok(bytes) => bytes, + Err(error) if error.kind() == io::ErrorKind::NotFound => return Ok(Self::default()), + Err(error) => return Err(error), + }; + Self::try_from_slice(&bytes).map_err(io::Error::other) } pub(crate) fn write_snapshot(&self, data_dir: &Path) -> io::Result<()> { @@ -222,7 +251,7 @@ mod tests { #[test] fn placement_readiness_and_graduation_are_distinct() { let current = SegmentKey::new(TopicId(1), RangeId(2), SegmentId(4)); - let mut ledger = OffsetLedger::default(); + let mut ledger = AuxiliaryState::default(); ledger.apply(OffsetRecord::PlacementInstalled(current)); assert!(ledger.has_installed_placement(¤t)); @@ -239,7 +268,7 @@ mod tests { batch_offset: 2, absolute_offset: 9, }; - let mut ledger = OffsetLedger::default(); + let mut ledger = AuxiliaryState::default(); ledger.apply(OffsetRecord::EpochSeal(EpochSeal { key: key(), generation: GenerationId(3), @@ -251,7 +280,7 @@ mod tests { })); ledger.write_snapshot(dir.path()).unwrap(); - let recovered = OffsetLedger::load_snapshot(dir.path()).unwrap(); + let recovered = AuxiliaryState::load_snapshot(dir.path()).unwrap(); assert_eq!(*recovered.generation(&key()), 3); assert_eq!(recovered.offset(&key()), Some(position)); } diff --git a/src/data_plane/consumer_offset_management/types.rs b/src/data_plane/auxiliary_states/consumer_offsets/types.rs similarity index 95% rename from src/data_plane/consumer_offset_management/types.rs rename to src/data_plane/auxiliary_states/consumer_offsets/types.rs index d673b75e..d7dd332c 100644 --- a/src/data_plane/consumer_offset_management/types.rs +++ b/src/data_plane/auxiliary_states/consumer_offsets/types.rs @@ -1,10 +1,8 @@ use std::collections::HashSet; +use super::state::{ConsumerOffsetUpdate, EpochSeal, OffsetRecord}; use crate::control_plane::{NodeId, Replicas}; use crate::data_plane::SegmentKey; -use crate::data_plane::consumer_offset_management::ledger::{ - ConsumerOffsetUpdate, EpochSeal, OffsetRecord, -}; use crate::data_plane::messages::command::{ CommitConsumerOffset, ConsumerOffsetCommitAck, ConsumerOffsetSnapshotInstalled, }; diff --git a/src/data_plane/consumer_offset_management/manager.rs b/src/data_plane/auxiliary_states/manager.rs similarity index 77% rename from src/data_plane/consumer_offset_management/manager.rs rename to src/data_plane/auxiliary_states/manager.rs index 5b1f87ad..deccea62 100644 --- a/src/data_plane/consumer_offset_management/manager.rs +++ b/src/data_plane/auxiliary_states/manager.rs @@ -1,139 +1,87 @@ +use super::consumer_offsets::state::{ + AuxiliaryState, ConsumerOffsetKey, ConsumerOffsetPosition, EpochSeal, OffsetRecord, StaleEpoch, +}; use crate::control_plane::membership::ShardGroupId; use crate::control_plane::metadata::consumer_group::GenerationId; use crate::control_plane::metadata::{RangeId, SegmentId, TopicId}; use crate::control_plane::{NodeId, Replicas}; use crate::data_plane::SegmentKey; -use crate::data_plane::checkpoint::OffsetCheckpointJob; -use crate::data_plane::consumer_offset_management::ledger::{ - ConsumerOffsetKey, ConsumerOffsetPosition, EpochSeal, OffsetLedger, OffsetRecord, StaleEpoch, + +use crate::data_plane::auxiliary_states::consumer_offsets::{ + ConsumerOffsetCoordination, OffsetPlacement, OffsetReplicaState, }; -use crate::data_plane::consumer_offset_management::replication::OffsetReplicationState; +use crate::data_plane::checkpoint::AuxiliaryCheckpointJob; use crate::data_plane::messages::command::{ - CommitConsumerOffset, ConsumerOffsetCommitAck, ConsumerOffsetReplicated, - ConsumerOffsetReplicationResult, ConsumerOffsetSnapshot, ConsumerOffsetSnapshotInstalled, + AuthorizedProducerIdentity, CommitConsumerOffset, ConsumerOffsetCommitAck, + ConsumerOffsetReplicated, ConsumerOffsetReplicationResult, ConsumerOffsetSnapshot, + ConsumerOffsetSnapshotInstalled, }; use crate::data_plane::messages::{RequestConsumerOffsetSnapshot, SegmentPlaced}; +use crate::data_plane::ProducerAppendIdentity; +use crate::data_plane::producer_ledger::{AppendKey, ProducerDecision}; use crate::data_plane::transport::command::DataTransportCommand; use std::collections::{HashMap, HashSet, VecDeque}; use std::path::PathBuf; -use super::types::*; +use super::consumer_offsets::types::*; #[derive(Default)] -pub(crate) struct ConsumerOffsetManager { - offset_ledger: OffsetLedger, // durable offsets, epochs, and local placement readiness +pub(crate) struct AuxiliaryStateManager { + durable: AuxiliaryState, coordination: ConsumerOffsetCoordination, // live placements, pending mutations, parked commits, and replication tracking - checkpoint: OffsetCheckpointState, // WAL reclamation and checkpoint progress -} - -#[derive(Default)] -struct ConsumerOffsetCoordination { - placements: HashMap<(TopicId, RangeId), OffsetPlacement>, - pending_offset_mutations: Vec, - future_offset_commits: HashMap>, - offset_replication: OffsetReplicationState, -} - -struct OffsetPlacement { - segment_key: SegmentKey, - shard_group_id: ShardGroupId, - // Desired data replicas for this placement. This also preserves their - // ordering, including which replica is the leader. - replicas: Replicas, - // Consumer-offset readiness for every member of `replicas`. - replica_states: HashMap, - placement_ack_sent: bool, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -enum OffsetReplicaState { - Joining, - SnapshotInstalled, - Ready, -} - -impl OffsetPlacement { - fn leader(&self) -> &NodeId { - self.replicas - .leader() - .expect("an offset placement must have a leader") - } - - fn can_bootstrap_replicas(&self, segment_key: &SegmentKey) -> bool { - self.segment_key == *segment_key && self.is_replica_ready(self.leader()) - } - - fn is_replica_ready(&self, node: &NodeId) -> bool { - self.replica_states.get(node) == Some(&OffsetReplicaState::Ready) - } - - fn all_replicas_ready(&self) -> bool { - self.replica_states - .values() - .all(|state| *state == OffsetReplicaState::Ready) - } - - fn compare(&self, other_key: &SegmentKey, other: &Replicas) -> PlacementObservation { - if self.segment_key == *other_key { - if &self.replicas != other { - tracing::warn!( - ?other_key, - current_replicas = ?self.replicas, - observed_replicas = ?other, - "ignored conflicting replica set for an existing segment" - ); - } - return if &self.replicas == other { - PlacementObservation::Unchanged - } else { - PlacementObservation::Conflict - }; - } - if self.segment_key.segment_id >= other_key.segment_id { - return PlacementObservation::Stale; - } - PlacementObservation::Accepted - } - - fn unready_followers(&self) -> Box<[NodeId]> { - self.replicas - .followers() - .filter(|node| !self.is_replica_ready(node)) - .cloned() - .collect() - } + checkpoint: AuxiliaryCheckpointState, // WAL reclamation and checkpoint progress } #[derive(Default)] -struct OffsetCheckpointState { +struct AuxiliaryCheckpointState { uncheckpointed_lsns: VecDeque, checkpoint_lsn: u64, in_flight: bool, } -impl ConsumerOffsetManager { - pub(crate) fn new(offset_ledger: OffsetLedger) -> Self { +impl AuxiliaryStateManager { + pub(crate) fn new(durable: AuxiliaryState) -> Self { Self { - offset_ledger, + durable, ..Default::default() } } - pub(crate) fn future_entry(&mut self, key: ConsumerOffsetKey) -> &mut Vec { - self.coordination - .future_offset_commits - .entry(key) - .or_default() + pub(crate) fn verify_producer( + &mut self, + segment_key: SegmentKey, + producer: AuthorizedProducerIdentity, + ) -> ProducerDecision { + self.durable.verify_producer(segment_key, producer) + } + + pub(crate) fn unstage_producer(&mut self, append_key: &AppendKey) { + self.durable.unstage_producer(append_key); + } + + pub(crate) fn commit_producer( + &mut self, + segment_key: SegmentKey, + producer: ProducerAppendIdentity, + entry_id: crate::control_plane::metadata::EntryId, + lsn: u64, + ) { + self.durable.apply_producer(segment_key, producer, entry_id); + self.mark_auxiliary_state_dirty(lsn); + } + + pub(crate) fn push_future(&mut self, key: ConsumerOffsetKey, future: FutureOffsetCommit) { + self.coordination.future_entry(key).push(future); } pub(crate) fn offset(&self, key: &ConsumerOffsetKey) -> Option { - self.offset_ledger.offset(key) + self.durable.offset(key) } pub(crate) fn durable_generation(&self, key: &ConsumerOffsetKey) -> GenerationId { - self.offset_ledger.generation(key) + self.durable.generation(key) } pub(crate) fn commit_consumer_offset( @@ -161,8 +109,7 @@ impl ConsumerOffsetManager { return false; } if cmd.update.generation > sealed_generation { - self.future_entry(cmd.update.key.clone()) - .push(FutureOffsetCommit::Client(cmd)); + self.push_future(cmd.update.key.clone(), FutureOffsetCommit::Client(cmd)); return false; } @@ -179,7 +126,7 @@ impl ConsumerOffsetManager { true } - pub(crate) fn handle_offset_checkpoint_complete(&mut self, checkpointed_lsn: u64) { + pub(crate) fn handle_auxiliary_checkpoint_complete(&mut self, checkpointed_lsn: u64) { self.checkpoint.checkpoint_lsn = self.checkpoint.checkpoint_lsn.max(checkpointed_lsn); // Clear now-safely-persisted LSNs from the uncheckpointed tracking queue @@ -194,8 +141,14 @@ impl ConsumerOffsetManager { self.checkpoint.in_flight = false; } + pub(crate) fn mark_auxiliary_state_dirty(&mut self, lsn: u64) { + if self.checkpoint.uncheckpointed_lsns.back() != Some(&lsn) { + self.checkpoint.uncheckpointed_lsns.push_back(lsn); + } + } + // Looks at the oldest uncheckpointed consumer offset commit. If there are consumer offsets in memory that haven't been - // written to the persistent offset_ledger file, their WAL entries must be preserved. + // included in the persistent auxiliary snapshot, their WAL entries must be preserved. pub(crate) fn reclamation_watermark(&self, reclaimable_lsn: u64) -> u64 { self.checkpoint .uncheckpointed_lsns @@ -212,19 +165,24 @@ impl ConsumerOffsetManager { self.checkpoint.uncheckpointed_lsns.back() } - pub(crate) fn raise_offset_checkpoint_job( + pub(crate) fn raise_auxiliary_checkpoint_job( &mut self, data_dir: PathBuf, - ) -> Option { + ) -> Option { self.checkpoint.in_flight = true; - Some(OffsetCheckpointJob { + Some(AuxiliaryCheckpointJob { checkpointed_lsn: *self.latest_uncheckpointed_lsn()?, - offsets: self.offset_ledger.clone(), + offsets: self.durable.clone(), data_dir, }) } - pub(crate) fn offset_checkpoint_in_flight(&self) -> bool { + #[cfg(any(test, debug_assertions))] + pub(crate) fn assert_producer_invariants(&self) { + self.durable.assert_producer_invariants(); + } + + pub(crate) fn auxiliary_checkpoint_in_flight(&self) -> bool { self.checkpoint.in_flight } @@ -241,7 +199,7 @@ impl ConsumerOffsetManager { } // need to place! - let local_placement_installed = self.offset_ledger.has_installed_placement(&segment_key); + let local_placement_installed = self.durable.has_installed_placement(&segment_key); let replica_states = replicas .iter() .cloned() @@ -358,7 +316,7 @@ impl ConsumerOffsetManager { .collect(); if segment_key.segment_id == SegmentId(0) - || self.offset_ledger.has_installed_placement(&segment_key) + || self.durable.has_installed_placement(&segment_key) { replica_states.insert(leader.clone(), OffsetReplicaState::Ready); } @@ -398,23 +356,19 @@ impl ConsumerOffsetManager { let mut bootstrapped_keys = HashSet::new(); for entry in cmd.entries { bootstrapped_keys.insert(entry.key.clone()); - self.coordination - .pending_offset_mutations - .push(PendingOffsetMutation::new( - OffsetRecord::BootstrapEntry(entry), - OffsetMutationCompletion::EpochSeal, - )); - } - self.coordination - .pending_offset_mutations - .push(PendingOffsetMutation::new( - OffsetRecord::PlacementInstalled(cmd.segment_key), - ConsumerOffsetSnapshotInstalled { - segment_key: cmd.segment_key, - from: node_id.clone(), - leader, - }, + self.coordination.push_pending(PendingOffsetMutation::new( + OffsetRecord::BootstrapEntry(entry), + OffsetMutationCompletion::EpochSeal, )); + } + self.coordination.push_pending(PendingOffsetMutation::new( + OffsetRecord::PlacementInstalled(cmd.segment_key), + ConsumerOffsetSnapshotInstalled { + segment_key: cmd.segment_key, + from: node_id.clone(), + leader, + }, + )); bootstrapped_keys .into_iter() .flat_map(|key| self.take_future_commits(key)) @@ -434,11 +388,7 @@ impl ConsumerOffsetManager { }); // Reject only when the source is neither live-ready nor durably eligible. - if !live_ready - && !self - .offset_ledger - .can_source_snapshot_for(&request.segment_key) - { + if !live_ready && !self.durable.can_source_snapshot_for(&request.segment_key) { return None; } Some(DataTransportCommand::send_to_targets( @@ -454,7 +404,7 @@ impl ConsumerOffsetManager { ) -> ConsumerOffsetSnapshot { ConsumerOffsetSnapshot { entries: self - .offset_ledger + .durable .snapshot_range(segment_key.topic_id, segment_key.range_id), segment_key, replica_set, @@ -496,7 +446,7 @@ impl ConsumerOffsetManager { } fn mark_offset_replica_ready_if_caught_up(&mut self, node: &NodeId) -> bool { - if self.coordination.offset_replication.has_pending_for(node) { + if self.coordination.has_pending_replication_for(node) { return false; } let mut readiness_changed = false; @@ -550,16 +500,16 @@ impl ConsumerOffsetManager { } pub(crate) fn has_pending_offsets(&self) -> bool { - !self.coordination.pending_offset_mutations.is_empty() + self.coordination.has_pending_offsets() } pub(crate) fn push_pending_offset_mutation(&mut self, cmd: impl Into) { - self.coordination.pending_offset_mutations.push(cmd.into()); + self.coordination.push_pending(cmd.into()); } pub(crate) fn handle_replica_offset_ack(&mut self, ack: ConsumerOffsetReplicated) -> bool { let from = ack.from.clone(); - self.coordination.offset_replication.process_ack(ack); + self.coordination.ack_replication(ack); self.mark_offset_replica_ready_if_caught_up(&from) } @@ -578,73 +528,21 @@ impl ConsumerOffsetManager { pub(crate) fn latest_generation(&self, key: &ConsumerOffsetKey) -> GenerationId { self.coordination - .pending_offset_mutations - .iter() - .rev() // Look backwards through the pending queue (newest first) - .find_map(|pending| match &pending.record { - OffsetRecord::EpochSeal(EpochSeal { - key: pending_key, - generation, - }) if pending_key == key => Some(*generation), - OffsetRecord::BootstrapEntry(snapshot) if &snapshot.key == key => { - Some(snapshot.generation) - } - _ => None, - }) + .latest_generation(key) // No pending EpochSeals, fall back to the actual ledger's generation - .unwrap_or_else(|| self.offset_ledger.generation(key)) + .unwrap_or_else(|| self.durable.generation(key)) } pub(crate) fn encode_offsets(&mut self) -> Result>, std::io::Error> { - match self - .coordination - .pending_offset_mutations - .iter() - .map(|pending| borsh::to_vec(&pending.record)) - .collect() - { - Ok(encoded) => Ok(encoded), - Err(error) => { - tracing::error!(?error, "consumer offset WAL encoding failed"); - for pending in std::mem::take(&mut self.coordination.pending_offset_mutations) { - if let OffsetMutationCompletion::LeaderCommit(commit) = pending.completion { - let _ = commit - .reply - .send(ConsumerOffsetCommitAck::InternalError(error.to_string())); - } - } - Err(error) - } - } + self.coordination.encode_offsets() } pub(crate) fn take_pending(&mut self) -> Vec { - std::mem::take(&mut self.coordination.pending_offset_mutations) + self.coordination.take_pending() } - pub(crate) fn fail_all(&mut self, error: &str) { - self.coordination.offset_replication.fail_all(error); - - for pending in self.take_pending() { - if let OffsetMutationCompletion::LeaderCommit(commit) = pending.completion { - let _ = commit - .reply - .send(ConsumerOffsetCommitAck::InternalError(error.to_string())); - } - } - - for parked in self - .coordination - .future_offset_commits - .drain() - .flat_map(|(_, commits)| commits) - { - if let FutureOffsetCommit::Client(commit) = parked { - let _ = commit - .reply - .send(ConsumerOffsetCommitAck::InternalError(error.to_string())); - } - } + pub(crate) fn fail_inflight_coordination(&mut self, error: &str) { + self.coordination.fail_all(error); } pub(crate) fn flush_batch( @@ -659,7 +557,7 @@ impl ConsumerOffsetManager { } for PendingOffsetMutation { record, completion } in pending { - self.offset_ledger.apply(record.clone()); + self.durable.apply(record.clone()); match completion { OffsetMutationCompletion::EpochSeal => {} OffsetMutationCompletion::LeaderCommit(commit) => { @@ -674,7 +572,7 @@ impl ConsumerOffsetManager { continue; } - let seq = self.coordination.offset_replication.begin( + let seq = self.coordination.begin_replication( followers.clone().into_iter().collect(), commit.required_followers, offset_commit.clone(), @@ -723,10 +621,7 @@ impl ConsumerOffsetManager { } fn take_future_commits(&mut self, key: ConsumerOffsetKey) -> Vec { - self.coordination - .future_offset_commits - .remove(&key) - .unwrap_or_default() + self.coordination.take_future_commits(&key) } #[cfg(test)] @@ -735,7 +630,7 @@ impl ConsumerOffsetManager { } #[cfg(test)] - pub(crate) fn offset_checkpoint_lsn(&self) -> u64 { + pub(crate) fn auxiliary_checkpoint_lsn(&self) -> u64 { self.checkpoint.checkpoint_lsn } } @@ -746,7 +641,7 @@ mod tests { use super::*; use crate::control_plane::metadata::EntryId; - use crate::data_plane::consumer_offset_management::ledger::{ + use crate::data_plane::auxiliary_states::consumer_offsets::state::{ ConsumerOffsetPosition, ConsumerOffsetUpdate, }; use crate::data_plane::messages::DataPlanePeerMessage; @@ -776,9 +671,9 @@ mod tests { fn observed_placement_restores_local_durable_readiness() { let local = NodeId::new("b"); let current = segment(2); - let mut ledger = OffsetLedger::default(); + let mut ledger = AuxiliaryState::default(); ledger.apply(OffsetRecord::PlacementInstalled(current)); - let mut manager = ConsumerOffsetManager::new(ledger); + let mut manager = AuxiliaryStateManager::new(ledger); let replicas = Replicas::new(vec![NodeId::new("a"), local.clone(), NodeId::new("c")]); manager.observe_placement(current, ShardGroupId(0), &replicas, &local); @@ -797,7 +692,7 @@ mod tests { let leader = NodeId::new("a"); let replicas = Replicas::new(vec![leader, local.clone()]); let current = segment(2); - let mut manager = ConsumerOffsetManager::new(OffsetLedger::default()); + let mut manager = AuxiliaryStateManager::new(AuxiliaryState::default()); manager.observe_placement(current, ShardGroupId(0), &replicas, &local); let live_placement = manager .coordination @@ -830,7 +725,7 @@ mod tests { let stale_replicas = Replicas::new(vec![NodeId::new("old"), local.clone()]); let current = segment(2); let stale = segment(1); - let mut manager = ConsumerOffsetManager::new(OffsetLedger::default()); + let mut manager = AuxiliaryStateManager::new(AuxiliaryState::default()); manager.observe_placement(current, ShardGroupId(0), ¤t_replicas, &local); manager.observe_placement(stale, ShardGroupId(0), &stale_replicas, &local); @@ -850,7 +745,7 @@ mod tests { let follower = NodeId::new("b"); let current = segment(2); let current_replicas = Replicas::new(vec![leader.clone(), follower]); - let mut manager = ConsumerOffsetManager::new(OffsetLedger::default()); + let mut manager = AuxiliaryStateManager::new(AuxiliaryState::default()); manager.install_leader_placement(current, ShardGroupId(7), current_replicas.clone()); @@ -887,7 +782,7 @@ mod tests { let old_leader = NodeId::new("a"); let retained = NodeId::new("b"); let new_leader = NodeId::new("d"); - let mut manager = ConsumerOffsetManager::new(OffsetLedger::default()); + let mut manager = AuxiliaryStateManager::new(AuxiliaryState::default()); manager.install_leader_placement( segment(1), ShardGroupId(7), @@ -949,14 +844,14 @@ mod tests { requester: new_leader.clone(), }; - let empty = ConsumerOffsetManager::new(OffsetLedger::default()); + let empty = AuxiliaryStateManager::new(AuxiliaryState::default()); assert!( empty .create_offset_snapshot(request.clone(), &NodeId::new("empty")) .is_none() ); - let mut ledger = OffsetLedger::default(); + let mut ledger = AuxiliaryState::default(); ledger.apply(OffsetRecord::PlacementInstalled(old_segment)); ledger.apply(OffsetRecord::EpochSeal(EpochSeal { key: key(), @@ -967,7 +862,7 @@ mod tests { generation: GenerationId(1), position: position(), })); - let source_manager = ConsumerOffsetManager::new(ledger); + let source_manager = AuxiliaryStateManager::new(ledger); let skipped_request = RequestConsumerOffsetSnapshot { segment_key: segment(3), @@ -1002,7 +897,7 @@ mod tests { let retained = NodeId::new("b"); let removed = NodeId::new("c"); let joining = NodeId::new("d"); - let mut ledger = OffsetLedger::default(); + let mut ledger = AuxiliaryState::default(); ledger.apply(OffsetRecord::EpochSeal(EpochSeal { key: key(), generation: GenerationId(1), @@ -1012,7 +907,7 @@ mod tests { generation: GenerationId(1), position: position(), })); - let mut manager = ConsumerOffsetManager::new(ledger); + let mut manager = AuxiliaryStateManager::new(ledger); manager.coordination.placements.insert( (TopicId(1), RangeId(2)), OffsetPlacement { @@ -1096,12 +991,11 @@ mod tests { #[test] fn fail_all_resolves_pending_in_flight_and_future_clients() { - let mut manager = ConsumerOffsetManager::new(OffsetLedger::default()); + let mut manager = AuxiliaryStateManager::new(AuxiliaryState::default()); let (pending_reply, mut pending_response) = oneshot::channel(); manager .coordination - .pending_offset_mutations - .push(PendingOffsetMutation::new( + .push_pending(PendingOffsetMutation::new( OffsetRecord::OffsetCommit(ConsumerOffsetUpdate { key: key(), generation: GenerationId(1), @@ -1115,20 +1009,20 @@ mod tests { )); let (future_reply, mut future_response) = oneshot::channel(); - manager.coordination.future_offset_commits.insert( + manager.push_future( key(), - vec![FutureOffsetCommit::Client(CommitConsumerOffset { + FutureOffsetCommit::Client(CommitConsumerOffset { update: ConsumerOffsetUpdate { key: key(), generation: GenerationId(2), position: position(), }, reply: future_reply, - })], + }), ); let (replication_reply, mut replication_response) = oneshot::channel(); - manager.coordination.offset_replication.begin( + manager.coordination.begin_replication( HashSet::from_iter([NodeId::new("follower")]), HashSet::from_iter([NodeId::new("follower")]), ConsumerOffsetUpdate { @@ -1139,13 +1033,12 @@ mod tests { replication_reply, ); - manager.fail_all("WAL failed"); + manager.fail_inflight_coordination("WAL failed"); let expected = ConsumerOffsetCommitAck::InternalError("WAL failed".into()); assert_eq!(pending_response.try_recv().unwrap(), expected); assert_eq!(future_response.try_recv().unwrap(), expected); assert_eq!(replication_response.try_recv().unwrap(), expected); - assert!(manager.coordination.pending_offset_mutations.is_empty()); - assert!(manager.coordination.future_offset_commits.is_empty()); + assert!(manager.coordination.is_empty()); } } diff --git a/src/data_plane/auxiliary_states/mod.rs b/src/data_plane/auxiliary_states/mod.rs new file mode 100644 index 00000000..37534470 --- /dev/null +++ b/src/data_plane/auxiliary_states/mod.rs @@ -0,0 +1,3 @@ +pub(crate) mod consumer_offsets; +pub(crate) mod manager; +pub(crate) use manager::AuxiliaryStateManager; diff --git a/src/data_plane/checkpoint.rs b/src/data_plane/checkpoint.rs index bb90d983..c58a3ed4 100644 --- a/src/data_plane/checkpoint.rs +++ b/src/data_plane/checkpoint.rs @@ -4,7 +4,7 @@ use std::sync::Arc; use flume::Receiver; use crate::data_plane::actor::DataPlaneSender; -use crate::data_plane::consumer_offset_management::ledger::OffsetLedger; +use crate::data_plane::auxiliary_states::consumer_offsets::state::AuxiliaryState; use crate::data_plane::states::segment::cache::SegmentRingBuffer; use crate::data_plane::wal::WalRecord; @@ -69,12 +69,12 @@ impl CheckpointWorker { } // A snapshot is created only when WAL reclamation would cross an uncheckpointed offset record. - CheckpointTask::ConsumerOffsets(job) => { + CheckpointTask::AuxiliaryState(job) => { if let Err(e) = job.offsets.write_snapshot(&job.data_dir) { tracing::error!("consumer offset checkpoint failed: {e}"); continue; } - let completion: DataPlaneCommand = OffsetCheckpointComplete { + let completion: DataPlaneCommand = AuxiliaryCheckpointComplete { checkpointed_lsn: job.checkpointed_lsn, } .into(); @@ -136,7 +136,7 @@ pub(crate) enum CheckpointTask { DeleteSegmentIndex(SegmentKey), /// Snapshot the consumer-offset cache after its corresponding shared-WAL /// batch is durable. This runs off the data-plane write path. - ConsumerOffsets(Box), + AuxiliaryState(Box), } pub struct CheckpointJob { @@ -145,8 +145,8 @@ pub struct CheckpointJob { pub segment_file_path: PathBuf, } -pub struct OffsetCheckpointJob { - pub offsets: OffsetLedger, +pub struct AuxiliaryCheckpointJob { + pub offsets: AuxiliaryState, pub data_dir: PathBuf, pub checkpointed_lsn: u64, } diff --git a/src/data_plane/consumer_offset_management/mod.rs b/src/data_plane/consumer_offset_management/mod.rs deleted file mode 100644 index 80f397c0..00000000 --- a/src/data_plane/consumer_offset_management/mod.rs +++ /dev/null @@ -1,6 +0,0 @@ -pub(crate) mod ledger; -pub(crate) mod manager; -pub(crate) mod replication; -pub(crate) mod types; - -pub(crate) use manager::ConsumerOffsetManager; diff --git a/src/data_plane/messages/pending.rs b/src/data_plane/messages/pending.rs index cb9cf57b..98f3ba7f 100644 --- a/src/data_plane/messages/pending.rs +++ b/src/data_plane/messages/pending.rs @@ -6,7 +6,7 @@ use crate::control_plane::consensus::actor::MutlRaftSender; use crate::control_plane::consensus::messages::MultiRaftActorCommand; use crate::control_plane::metadata::EntryId; use crate::data_plane::SegmentKey; -use crate::data_plane::checkpoint::{CheckpointJob, CheckpointTask, OffsetCheckpointJob}; +use crate::data_plane::checkpoint::{AuxiliaryCheckpointJob, CheckpointJob, CheckpointTask}; use crate::data_plane::messages::command::ProduceAck; use crate::data_plane::sparse_index::SparseEntry; use crate::data_plane::timer::{BatchFlushTimer, ReplicationTimer}; @@ -99,9 +99,9 @@ impl DataPlaneOutputs { self.checkpoint_tasks.push(CheckpointTask::Checkpoint(job)); } - pub(crate) fn store_offset_checkpoint(&mut self, job: OffsetCheckpointJob) { + pub(crate) fn store_auxiliary_checkpoint(&mut self, job: AuxiliaryCheckpointJob) { self.checkpoint_tasks - .push(CheckpointTask::ConsumerOffsets(Box::new(job))); + .push(CheckpointTask::AuxiliaryState(Box::new(job))); } pub(crate) fn store_put_anchors(&mut self, anchors: Vec) { diff --git a/src/data_plane/messages/query.rs b/src/data_plane/messages/query.rs index a7454806..c74b4d70 100644 --- a/src/data_plane/messages/query.rs +++ b/src/data_plane/messages/query.rs @@ -13,7 +13,7 @@ use crate::connections::protocol::{ConsumerOffsetGenerationMismatch, RangeProgre use crate::control_plane::NodeId; use crate::control_plane::metadata::consumer_group::GenerationId; use crate::control_plane::metadata::{EntryId, RangeId, TopicId}; -use crate::data_plane::consumer_offset_management::ledger::{ +use crate::data_plane::auxiliary_states::consumer_offsets::state::{ ConsumerOffsetKey, ConsumerOffsetPosition, }; use crate::data_plane::states::segment::cache::CachedEntry; diff --git a/src/data_plane/mod.rs b/src/data_plane/mod.rs index 1d80047d..07a33127 100644 --- a/src/data_plane/mod.rs +++ b/src/data_plane/mod.rs @@ -7,9 +7,9 @@ use crate::control_plane::metadata::{EntryId, RangeId, SegmentId, TopicId}; use crate::impl_new_struct_wrapper; pub(crate) mod actor; +pub(crate) mod auxiliary_states; pub(crate) mod checkpoint; pub(crate) mod cold_read; -pub(crate) mod consumer_offset_management; pub(crate) mod messages; pub(crate) mod producer_ledger; pub(crate) mod recovery; diff --git a/src/data_plane/recovery/inventory.rs b/src/data_plane/recovery/inventory.rs index 8a3ce46d..a5eb3948 100644 --- a/src/data_plane/recovery/inventory.rs +++ b/src/data_plane/recovery/inventory.rs @@ -4,7 +4,7 @@ use std::path::PathBuf; use super::segment_scan::RecoveredSegments; use crate::control_plane::metadata::EntryId; use crate::data_plane::SegmentKey; -use crate::data_plane::consumer_offset_management::ledger::OffsetLedger; +use crate::data_plane::auxiliary_states::consumer_offsets::state::AuxiliaryState; /// What recovery verified for one on-disk segment: its `start_offset` (the filename base, /// needed to locate the file) and `verified_end` (the highest entry id a CRC-complete batch @@ -70,7 +70,7 @@ impl LocalInventory { pub(crate) struct RecoveryOutput { pub(crate) inventory: LocalInventory, pub(crate) data_dir: PathBuf, - pub(crate) offsets: OffsetLedger, + pub(crate) offsets: AuxiliaryState, } #[cfg(test)] @@ -126,7 +126,7 @@ mod tests { let output = RecoveryOutput { inventory: inv, data_dir: PathBuf::from("/var/lib/eastguard/data"), - offsets: OffsetLedger::default(), + offsets: AuxiliaryState::default(), }; assert_eq!(output.inventory.orphan_candidates().count(), 0); assert_eq!(output.data_dir, PathBuf::from("/var/lib/eastguard/data")); diff --git a/src/it/e2e/client_sdk.rs b/src/it/e2e/client_sdk.rs index 34a85406..90221a15 100644 --- a/src/it/e2e/client_sdk.rs +++ b/src/it/e2e/client_sdk.rs @@ -22,7 +22,7 @@ use crate::connections::protocol::ClientResponse; use crate::control_plane::metadata::consumer_group::GenerationId; use crate::control_plane::metadata::{EntryId, RangeId}; use crate::control_plane::{NodeAddress, NodeAddressInfo, NodeId}; -use crate::data_plane::consumer_offset_management::ledger::ConsumerOffsetKey; +use crate::data_plane::auxiliary_states::consumer_offsets::state::ConsumerOffsetKey; use crate::it::e2e::{NODES, NODES_4}; /// Per-test cluster tweaks for the SDK suite: pinned node-id suffix + low vnode From 5e854461c7b07a75899802dc9f2676c723c83af6 Mon Sep 17 00:00:00 2001 From: Migorithm Date: Wed, 22 Jul 2026 20:01:00 +0400 Subject: [PATCH 14/41] unify snapshotting --- .../consensus/raft/states/metadata_state.rs | 42 +++++- src/control_plane/metadata/mod.rs | 1 + src/control_plane/metadata/topic.rs | 11 ++ .../consumer_offsets/state.rs | 80 ++--------- src/data_plane/auxiliary_states/manager.rs | 130 ++++++++++-------- src/data_plane/auxiliary_states/mod.rs | 2 + .../auxiliary_states/producer/mod.rs | 5 + .../auxiliary_states/producer/types.rs | 35 +++++ src/data_plane/checkpoint.rs | 8 +- src/data_plane/mod.rs | 35 +---- src/data_plane/recovery/inventory.rs | 9 +- src/data_plane/recovery/mod.rs | 90 ++++++++++-- src/data_plane/states/segment/tracker.rs | 72 +++++++--- 13 files changed, 319 insertions(+), 201 deletions(-) create mode 100644 src/data_plane/auxiliary_states/producer/mod.rs create mode 100644 src/data_plane/auxiliary_states/producer/types.rs diff --git a/src/control_plane/consensus/raft/states/metadata_state.rs b/src/control_plane/consensus/raft/states/metadata_state.rs index e064f77d..ec4b19fa 100644 --- a/src/control_plane/consensus/raft/states/metadata_state.rs +++ b/src/control_plane/consensus/raft/states/metadata_state.rs @@ -207,12 +207,51 @@ impl MetadataState { ReassignSegment(cmd) => self.reassign_segment(cmd)?, DeleteSegments(cmd) => self.delete_segments(cmd)?, SyncConsumerGroup(cmd) => self.sync_consumer_group(cmd)?, + OpenProducerSession(cmd) => self.open_producer_session(cmd)?, + ExpireProducerSessions(cmd) => self.expire_producer_sessions(cmd)?, } #[cfg(any(test, debug_assertions))] self.assert_invariants(); Ok(()) } + fn open_producer_session(&mut self, cmd: OpenProducerSession) -> Result<(), MetadataError> { + let topic_id = self + .topic_name_index + .get(&cmd.topic_name) + .copied() + .ok_or_else(|| MetadataError::TopicNameNotFound(cmd.topic_name.clone()))?; + let topic = self + .topics + .get_mut(&topic_id) + .ok_or(MetadataError::TopicNotFound(topic_id))?; + topic.producer_sessions.open_producer_session( + cmd.producer_id, + cmd.session_nonce, + cmd.observed_at, + cmd.session_timeout_ms, + ); + Ok(()) + } + + fn expire_producer_sessions( + &mut self, + cmd: ExpireProducerSessions, + ) -> Result<(), MetadataError> { + self.get_active_topic_mut(cmd.topic_id)? + .producer_sessions + .expire_producer_sessions(cmd.observed_at); + Ok(()) + } + + pub(crate) fn topics_with_expired_producer_sessions(&self, now: u64) -> Box<[TopicId]> { + self.topics + .values() + .filter(|topic| topic.producer_sessions.has_expired_producer_sessions(now)) + .map(|topic| topic.id) + .collect() + } + fn create_topic(&mut self, cmd: CreateTopic) -> Result<(), MetadataError> { if self.topic_name_index.contains_key(&cmd.name) { return Err(TopicNameAlreadyExists(cmd.name)); @@ -539,8 +578,9 @@ impl crate::test_traits::TAssertInvariant for MetadataState { assert_eq!(&topic.name, name); } - for id in self.topics.keys() { + for (id, topic) in &self.topics { assert!(id.0 < self.next_topic_id, "topic ID >= next_topic_id"); + assert_eq!(*id, topic.id, "topic map key does not match topic identity"); } for topic in self.topics.values() { topic.assert_invariants(); diff --git a/src/control_plane/metadata/mod.rs b/src/control_plane/metadata/mod.rs index 6255f737..3b5c1693 100644 --- a/src/control_plane/metadata/mod.rs +++ b/src/control_plane/metadata/mod.rs @@ -6,6 +6,7 @@ pub mod error; pub(crate) mod event; pub(crate) mod range; +mod producer_sessions; pub mod strategy; pub(crate) mod topic; diff --git a/src/control_plane/metadata/topic.rs b/src/control_plane/metadata/topic.rs index 55b9f893..7418601f 100644 --- a/src/control_plane/metadata/topic.rs +++ b/src/control_plane/metadata/topic.rs @@ -8,6 +8,7 @@ use crate::{ metadata::{ error::MetadataError, event::ConsumerGroupEpochSnapshot, + producer_sessions::ProducerSessions, strategy::{PartitionStrategy, StoragePolicy}, }, }, @@ -32,7 +33,9 @@ pub struct TopicMeta { pub ranges: HashMap, pub next_range_id: u64, pub(crate) consumer_groups: HashMap, + pub(crate) producer_sessions: ProducerSessions, } + impl TopicMeta { pub(crate) fn new( name: String, @@ -59,6 +62,7 @@ impl TopicMeta { ranges: HashMap::from([(range_id, range)]), next_range_id: 1, consumer_groups: HashMap::new(), + producer_sessions: Default::default(), } } @@ -495,6 +499,7 @@ impl TopicMeta { pub(crate) fn delete(&mut self) { self.state = TopicState::Deleted; self.active_ranges.clear(); + self.producer_sessions.clear(); for range in self.ranges.values_mut() { range.delete(); @@ -529,6 +534,12 @@ pub mod props { for group in self.consumer_groups.values() { group.assert_assignments(&self.active_ranges); } + if self.state == TopicState::Deleted { + assert!( + self.producer_sessions.is_empty(), + "deleted topic retains producer sessions" + ); + } } } diff --git a/src/data_plane/auxiliary_states/consumer_offsets/state.rs b/src/data_plane/auxiliary_states/consumer_offsets/state.rs index 6c35c30f..5f019576 100644 --- a/src/data_plane/auxiliary_states/consumer_offsets/state.rs +++ b/src/data_plane/auxiliary_states/consumer_offsets/state.rs @@ -1,20 +1,12 @@ use std::cmp::Ordering; use std::collections::HashMap; -use std::fs::{self, File, OpenOptions}; -use std::io::{self, Write}; -use std::path::Path; use borsh::{BorshDeserialize, BorshSerialize}; use crate::client::RangeId; use crate::control_plane::metadata::consumer_group::GenerationId; use crate::control_plane::metadata::{EntryId, TopicId}; -use crate::data_plane::messages::command::AuthorizedProducerIdentity; -use crate::data_plane::producer_ledger::{AppendKey, ProducerDecision, ProducerLedger}; -use crate::data_plane::{ProducerAppendIdentity, SegmentKey}; - -const SNAPSHOT_FILE: &str = "auxiliary-state.snapshot"; -const SNAPSHOT_TEMP_FILE: &str = "auxiliary-state.snapshot.tmp"; +use crate::data_plane::SegmentKey; #[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash, BorshSerialize, BorshDeserialize)] pub(crate) struct ConsumerOffsetKey { @@ -83,65 +75,13 @@ pub(crate) enum OffsetRecord { /// data-plane WAL; this type is only the in-memory cache and its asynchronous /// WAL-reclamation snapshot. #[derive(Debug, Clone, Default, BorshSerialize, BorshDeserialize)] -pub(crate) struct AuxiliaryState { +pub(crate) struct ConsumerOffsets { epochs: HashMap, offsets: HashMap, installed_placements: HashMap<(TopicId, RangeId), SegmentKey>, - producer_ledger: ProducerLedger, } -impl AuxiliaryState { - pub(crate) fn verify_producer( - &mut self, - segment_key: SegmentKey, - producer: AuthorizedProducerIdentity, - ) -> ProducerDecision { - self.producer_ledger.verify(segment_key, producer) - } - - pub(crate) fn unstage_producer(&mut self, append_key: &AppendKey) { - self.producer_ledger.unstage(append_key); - } - - pub(crate) fn apply_producer( - &mut self, - segment_key: SegmentKey, - producer: ProducerAppendIdentity, - entry_id: EntryId, - ) { - self.producer_ledger.commit(segment_key, producer, entry_id); - } - - #[cfg(any(test, debug_assertions))] - pub(crate) fn assert_producer_invariants(&self) { - crate::test_traits::TAssertInvariant::assert_invariants(&self.producer_ledger); - } - pub(crate) fn load_snapshot(data_dir: &Path) -> io::Result { - let path = data_dir.join(SNAPSHOT_FILE); - let bytes = match fs::read(path) { - Ok(bytes) => bytes, - Err(error) if error.kind() == io::ErrorKind::NotFound => return Ok(Self::default()), - Err(error) => return Err(error), - }; - Self::try_from_slice(&bytes).map_err(io::Error::other) - } - - pub(crate) fn write_snapshot(&self, data_dir: &Path) -> io::Result<()> { - fs::create_dir_all(data_dir)?; - let bytes = borsh::to_vec(self).map_err(io::Error::other)?; - let temporary = data_dir.join(SNAPSHOT_TEMP_FILE); - let final_path = data_dir.join(SNAPSHOT_FILE); - let mut file = OpenOptions::new() - .create(true) - .write(true) - .truncate(true) - .open(&temporary)?; - file.write_all(&bytes)?; - file.sync_all()?; - fs::rename(temporary, final_path)?; - File::open(data_dir)?.sync_all() - } - +impl ConsumerOffsets { pub(crate) fn generation(&self, key: &ConsumerOffsetKey) -> GenerationId { self.epochs.get(key).copied().unwrap_or(GenerationId(0)) } @@ -251,7 +191,7 @@ mod tests { #[test] fn placement_readiness_and_graduation_are_distinct() { let current = SegmentKey::new(TopicId(1), RangeId(2), SegmentId(4)); - let mut ledger = AuxiliaryState::default(); + let mut ledger = ConsumerOffsets::default(); ledger.apply(OffsetRecord::PlacementInstalled(current)); assert!(ledger.has_installed_placement(¤t)); @@ -261,14 +201,13 @@ mod tests { } #[test] - fn snapshot_survives_restart() { - let dir = tempfile::tempdir().unwrap(); + fn offsets_advance_monotonically() { let position = ConsumerOffsetPosition { entry_id: EntryId(4), batch_offset: 2, absolute_offset: 9, }; - let mut ledger = AuxiliaryState::default(); + let mut ledger = ConsumerOffsets::default(); ledger.apply(OffsetRecord::EpochSeal(EpochSeal { key: key(), generation: GenerationId(3), @@ -278,10 +217,7 @@ mod tests { generation: GenerationId(3), position, })); - ledger.write_snapshot(dir.path()).unwrap(); - - let recovered = AuxiliaryState::load_snapshot(dir.path()).unwrap(); - assert_eq!(*recovered.generation(&key()), 3); - assert_eq!(recovered.offset(&key()), Some(position)); + assert_eq!(*ledger.generation(&key()), 3); + assert_eq!(ledger.offset(&key()), Some(position)); } } diff --git a/src/data_plane/auxiliary_states/manager.rs b/src/data_plane/auxiliary_states/manager.rs index deccea62..0e5ceb04 100644 --- a/src/data_plane/auxiliary_states/manager.rs +++ b/src/data_plane/auxiliary_states/manager.rs @@ -1,6 +1,8 @@ use super::consumer_offsets::state::{ - AuxiliaryState, ConsumerOffsetKey, ConsumerOffsetPosition, EpochSeal, OffsetRecord, StaleEpoch, + ConsumerOffsetKey, ConsumerOffsetPosition, ConsumerOffsets, EpochSeal, OffsetRecord, StaleEpoch, }; +use super::state::AuxiliarySnapshot; +use crate::client::EntryId; use crate::control_plane::membership::ShardGroupId; use crate::control_plane::metadata::consumer_group::GenerationId; use crate::control_plane::metadata::{RangeId, SegmentId, TopicId}; @@ -20,7 +22,7 @@ use crate::data_plane::messages::command::{ use crate::data_plane::messages::{RequestConsumerOffsetSnapshot, SegmentPlaced}; use crate::data_plane::ProducerAppendIdentity; -use crate::data_plane::producer_ledger::{AppendKey, ProducerDecision}; +use crate::data_plane::auxiliary_states::producer::{AppendKey, ProducerDecision, ProducerTracker}; use crate::data_plane::transport::command::DataTransportCommand; use std::collections::{HashMap, HashSet, VecDeque}; use std::path::PathBuf; @@ -29,9 +31,10 @@ use super::consumer_offsets::types::*; #[derive(Default)] pub(crate) struct AuxiliaryStateManager { - durable: AuxiliaryState, - coordination: ConsumerOffsetCoordination, // live placements, pending mutations, parked commits, and replication tracking - checkpoint: AuxiliaryCheckpointState, // WAL reclamation and checkpoint progress + offsets: ConsumerOffsets, + consumer_coord: ConsumerOffsetCoordination, // live placements, pending mutations, parked commits, and replication tracking + producer: ProducerTracker, + checkpoint: AuxiliaryCheckpointState, // WAL reclamation and checkpoint progress } #[derive(Default)] @@ -42,9 +45,15 @@ struct AuxiliaryCheckpointState { } impl AuxiliaryStateManager { - pub(crate) fn new(durable: AuxiliaryState) -> Self { + #[cfg(test)] + pub(crate) fn new(offsets: ConsumerOffsets) -> Self { + Self::from_recovery(offsets, ProducerTracker::default()) + } + + pub(crate) fn from_recovery(offsets: ConsumerOffsets, producer: ProducerTracker) -> Self { Self { - durable, + offsets, + producer, ..Default::default() } } @@ -54,34 +63,34 @@ impl AuxiliaryStateManager { segment_key: SegmentKey, producer: AuthorizedProducerIdentity, ) -> ProducerDecision { - self.durable.verify_producer(segment_key, producer) + self.producer.verify(segment_key, producer) } pub(crate) fn unstage_producer(&mut self, append_key: &AppendKey) { - self.durable.unstage_producer(append_key); + self.producer.unstage(append_key); } - pub(crate) fn commit_producer( + pub(crate) fn advance_producer( &mut self, segment_key: SegmentKey, producer: ProducerAppendIdentity, - entry_id: crate::control_plane::metadata::EntryId, + entry_id: EntryId, lsn: u64, ) { - self.durable.apply_producer(segment_key, producer, entry_id); + self.producer.advance(segment_key, producer, entry_id); self.mark_auxiliary_state_dirty(lsn); } pub(crate) fn push_future(&mut self, key: ConsumerOffsetKey, future: FutureOffsetCommit) { - self.coordination.future_entry(key).push(future); + self.consumer_coord.future_entry(key).push(future); } pub(crate) fn offset(&self, key: &ConsumerOffsetKey) -> Option { - self.durable.offset(key) + self.offsets.offset(key) } pub(crate) fn durable_generation(&self, key: &ConsumerOffsetKey) -> GenerationId { - self.durable.generation(key) + self.offsets.generation(key) } pub(crate) fn commit_consumer_offset( @@ -172,14 +181,17 @@ impl AuxiliaryStateManager { self.checkpoint.in_flight = true; Some(AuxiliaryCheckpointJob { checkpointed_lsn: *self.latest_uncheckpointed_lsn()?, - offsets: self.durable.clone(), + snapshot: AuxiliarySnapshot { + consumer_offsets: self.offsets.clone(), + producer_sessions: self.producer.sessions(), + }, data_dir, }) } #[cfg(any(test, debug_assertions))] pub(crate) fn assert_producer_invariants(&self) { - self.durable.assert_producer_invariants(); + self.producer.assert_invariants(); } pub(crate) fn auxiliary_checkpoint_in_flight(&self) -> bool { @@ -199,7 +211,7 @@ impl AuxiliaryStateManager { } // need to place! - let local_placement_installed = self.durable.has_installed_placement(&segment_key); + let local_placement_installed = self.offsets.has_installed_placement(&segment_key); let replica_states = replicas .iter() .cloned() @@ -212,7 +224,7 @@ impl AuxiliaryStateManager { (node, state) }) .collect(); - self.coordination.placements.insert( + self.consumer_coord.placements.insert( segment_key.placement_key(), OffsetPlacement { segment_key, @@ -230,7 +242,7 @@ impl AuxiliaryStateManager { segment_key: &SegmentKey, replicas: &Replicas, ) -> PlacementObservation { - self.coordination + self.consumer_coord .placements .get(&(segment_key.topic_id, segment_key.range_id)) .map(|placement| placement.compare(segment_key, replicas)) @@ -316,7 +328,7 @@ impl AuxiliaryStateManager { .collect(); if segment_key.segment_id == SegmentId(0) - || self.durable.has_installed_placement(&segment_key) + || self.offsets.has_installed_placement(&segment_key) { replica_states.insert(leader.clone(), OffsetReplicaState::Ready); } @@ -324,7 +336,7 @@ impl AuxiliaryStateManager { } fn get_placement(&self, placement_key: &(TopicId, RangeId)) -> Option<&OffsetPlacement> { - self.coordination.placements.get(placement_key) + self.consumer_coord.placements.get(placement_key) } fn get_placement_followers(&self, cmd: &CommitConsumerOffset) -> HashSet { @@ -344,7 +356,7 @@ impl AuxiliaryStateManager { &mut self, placement_key: &(TopicId, RangeId), ) -> Option<&mut OffsetPlacement> { - self.coordination.placements.get_mut(placement_key) + self.consumer_coord.placements.get_mut(placement_key) } pub(crate) fn install_consumer_offsets( @@ -356,12 +368,12 @@ impl AuxiliaryStateManager { let mut bootstrapped_keys = HashSet::new(); for entry in cmd.entries { bootstrapped_keys.insert(entry.key.clone()); - self.coordination.push_pending(PendingOffsetMutation::new( + self.consumer_coord.push_pending(PendingOffsetMutation::new( OffsetRecord::BootstrapEntry(entry), OffsetMutationCompletion::EpochSeal, )); } - self.coordination.push_pending(PendingOffsetMutation::new( + self.consumer_coord.push_pending(PendingOffsetMutation::new( OffsetRecord::PlacementInstalled(cmd.segment_key), ConsumerOffsetSnapshotInstalled { segment_key: cmd.segment_key, @@ -388,7 +400,7 @@ impl AuxiliaryStateManager { }); // Reject only when the source is neither live-ready nor durably eligible. - if !live_ready && !self.durable.can_source_snapshot_for(&request.segment_key) { + if !live_ready && !self.offsets.can_source_snapshot_for(&request.segment_key) { return None; } Some(DataTransportCommand::send_to_targets( @@ -404,7 +416,7 @@ impl AuxiliaryStateManager { ) -> ConsumerOffsetSnapshot { ConsumerOffsetSnapshot { entries: self - .durable + .offsets .snapshot_range(segment_key.topic_id, segment_key.range_id), segment_key, replica_set, @@ -446,11 +458,11 @@ impl AuxiliaryStateManager { } fn mark_offset_replica_ready_if_caught_up(&mut self, node: &NodeId) -> bool { - if self.coordination.has_pending_replication_for(node) { + if self.consumer_coord.has_pending_replication_for(node) { return false; } let mut readiness_changed = false; - for placement in self.coordination.placements.values_mut() { + for placement in self.consumer_coord.placements.values_mut() { if placement.replica_states.get(node) == Some(&OffsetReplicaState::SnapshotInstalled) { placement .replica_states @@ -462,7 +474,7 @@ impl AuxiliaryStateManager { } pub(crate) fn drain_ready_placements(&mut self, from: &NodeId) -> Vec { - self.coordination + self.consumer_coord .placements .values_mut() .filter_map(|placement| { @@ -500,16 +512,16 @@ impl AuxiliaryStateManager { } pub(crate) fn has_pending_offsets(&self) -> bool { - self.coordination.has_pending_offsets() + self.consumer_coord.has_pending_offsets() } pub(crate) fn push_pending_offset_mutation(&mut self, cmd: impl Into) { - self.coordination.push_pending(cmd.into()); + self.consumer_coord.push_pending(cmd.into()); } pub(crate) fn handle_replica_offset_ack(&mut self, ack: ConsumerOffsetReplicated) -> bool { let from = ack.from.clone(); - self.coordination.ack_replication(ack); + self.consumer_coord.ack_replication(ack); self.mark_offset_replica_ready_if_caught_up(&from) } @@ -527,22 +539,22 @@ impl AuxiliaryStateManager { } pub(crate) fn latest_generation(&self, key: &ConsumerOffsetKey) -> GenerationId { - self.coordination + self.consumer_coord .latest_generation(key) // No pending EpochSeals, fall back to the actual ledger's generation - .unwrap_or_else(|| self.durable.generation(key)) + .unwrap_or_else(|| self.offsets.generation(key)) } pub(crate) fn encode_offsets(&mut self) -> Result>, std::io::Error> { - self.coordination.encode_offsets() + self.consumer_coord.encode_offsets() } pub(crate) fn take_pending(&mut self) -> Vec { - self.coordination.take_pending() + self.consumer_coord.take_pending() } pub(crate) fn fail_inflight_coordination(&mut self, error: &str) { - self.coordination.fail_all(error); + self.consumer_coord.fail_all(error); } pub(crate) fn flush_batch( @@ -557,7 +569,7 @@ impl AuxiliaryStateManager { } for PendingOffsetMutation { record, completion } in pending { - self.durable.apply(record.clone()); + self.offsets.apply(record.clone()); match completion { OffsetMutationCompletion::EpochSeal => {} OffsetMutationCompletion::LeaderCommit(commit) => { @@ -572,7 +584,7 @@ impl AuxiliaryStateManager { continue; } - let seq = self.coordination.begin_replication( + let seq = self.consumer_coord.begin_replication( followers.clone().into_iter().collect(), commit.required_followers, offset_commit.clone(), @@ -621,7 +633,7 @@ impl AuxiliaryStateManager { } fn take_future_commits(&mut self, key: ConsumerOffsetKey) -> Vec { - self.coordination.take_future_commits(&key) + self.consumer_coord.take_future_commits(&key) } #[cfg(test)] @@ -671,7 +683,7 @@ mod tests { fn observed_placement_restores_local_durable_readiness() { let local = NodeId::new("b"); let current = segment(2); - let mut ledger = AuxiliaryState::default(); + let mut ledger = ConsumerOffsets::default(); ledger.apply(OffsetRecord::PlacementInstalled(current)); let mut manager = AuxiliaryStateManager::new(ledger); @@ -679,7 +691,7 @@ mod tests { manager.observe_placement(current, ShardGroupId(0), &replicas, &local); let placement = manager - .coordination + .consumer_coord .placements .get(&(current.topic_id, current.range_id)) .unwrap(); @@ -692,10 +704,10 @@ mod tests { let leader = NodeId::new("a"); let replicas = Replicas::new(vec![leader, local.clone()]); let current = segment(2); - let mut manager = AuxiliaryStateManager::new(AuxiliaryState::default()); + let mut manager = AuxiliaryStateManager::new(ConsumerOffsets::default()); manager.observe_placement(current, ShardGroupId(0), &replicas, &local); let live_placement = manager - .coordination + .consumer_coord .placements .get_mut(&(current.topic_id, current.range_id)) .unwrap(); @@ -707,7 +719,7 @@ mod tests { manager.observe_placement(current, ShardGroupId(0), &replicas, &local); let placement = manager - .coordination + .consumer_coord .placements .get(&(current.topic_id, current.range_id)) .unwrap(); @@ -725,13 +737,13 @@ mod tests { let stale_replicas = Replicas::new(vec![NodeId::new("old"), local.clone()]); let current = segment(2); let stale = segment(1); - let mut manager = AuxiliaryStateManager::new(AuxiliaryState::default()); + let mut manager = AuxiliaryStateManager::new(ConsumerOffsets::default()); manager.observe_placement(current, ShardGroupId(0), ¤t_replicas, &local); manager.observe_placement(stale, ShardGroupId(0), &stale_replicas, &local); let placement = manager - .coordination + .consumer_coord .placements .get(&(current.topic_id, current.range_id)) .unwrap(); @@ -745,7 +757,7 @@ mod tests { let follower = NodeId::new("b"); let current = segment(2); let current_replicas = Replicas::new(vec![leader.clone(), follower]); - let mut manager = AuxiliaryStateManager::new(AuxiliaryState::default()); + let mut manager = AuxiliaryStateManager::new(ConsumerOffsets::default()); manager.install_leader_placement(current, ShardGroupId(7), current_replicas.clone()); @@ -769,7 +781,7 @@ mod tests { ); let placement = manager - .coordination + .consumer_coord .placements .get(&(current.topic_id, current.range_id)) .unwrap(); @@ -782,7 +794,7 @@ mod tests { let old_leader = NodeId::new("a"); let retained = NodeId::new("b"); let new_leader = NodeId::new("d"); - let mut manager = AuxiliaryStateManager::new(AuxiliaryState::default()); + let mut manager = AuxiliaryStateManager::new(ConsumerOffsets::default()); manager.install_leader_placement( segment(1), ShardGroupId(7), @@ -844,14 +856,14 @@ mod tests { requester: new_leader.clone(), }; - let empty = AuxiliaryStateManager::new(AuxiliaryState::default()); + let empty = AuxiliaryStateManager::new(ConsumerOffsets::default()); assert!( empty .create_offset_snapshot(request.clone(), &NodeId::new("empty")) .is_none() ); - let mut ledger = AuxiliaryState::default(); + let mut ledger = ConsumerOffsets::default(); ledger.apply(OffsetRecord::PlacementInstalled(old_segment)); ledger.apply(OffsetRecord::EpochSeal(EpochSeal { key: key(), @@ -897,7 +909,7 @@ mod tests { let retained = NodeId::new("b"); let removed = NodeId::new("c"); let joining = NodeId::new("d"); - let mut ledger = AuxiliaryState::default(); + let mut ledger = ConsumerOffsets::default(); ledger.apply(OffsetRecord::EpochSeal(EpochSeal { key: key(), generation: GenerationId(1), @@ -908,7 +920,7 @@ mod tests { position: position(), })); let mut manager = AuxiliaryStateManager::new(ledger); - manager.coordination.placements.insert( + manager.consumer_coord.placements.insert( (TopicId(1), RangeId(2)), OffsetPlacement { segment_key: segment(1), @@ -991,10 +1003,10 @@ mod tests { #[test] fn fail_all_resolves_pending_in_flight_and_future_clients() { - let mut manager = AuxiliaryStateManager::new(AuxiliaryState::default()); + let mut manager = AuxiliaryStateManager::new(ConsumerOffsets::default()); let (pending_reply, mut pending_response) = oneshot::channel(); manager - .coordination + .consumer_coord .push_pending(PendingOffsetMutation::new( OffsetRecord::OffsetCommit(ConsumerOffsetUpdate { key: key(), @@ -1022,7 +1034,7 @@ mod tests { ); let (replication_reply, mut replication_response) = oneshot::channel(); - manager.coordination.begin_replication( + manager.consumer_coord.begin_replication( HashSet::from_iter([NodeId::new("follower")]), HashSet::from_iter([NodeId::new("follower")]), ConsumerOffsetUpdate { @@ -1039,6 +1051,6 @@ mod tests { assert_eq!(pending_response.try_recv().unwrap(), expected); assert_eq!(future_response.try_recv().unwrap(), expected); assert_eq!(replication_response.try_recv().unwrap(), expected); - assert!(manager.coordination.is_empty()); + assert!(manager.consumer_coord.is_empty()); } } diff --git a/src/data_plane/auxiliary_states/mod.rs b/src/data_plane/auxiliary_states/mod.rs index 37534470..bcc0557a 100644 --- a/src/data_plane/auxiliary_states/mod.rs +++ b/src/data_plane/auxiliary_states/mod.rs @@ -1,3 +1,5 @@ pub(crate) mod consumer_offsets; pub(crate) mod manager; +pub(crate) mod producer; +pub(crate) mod state; pub(crate) use manager::AuxiliaryStateManager; diff --git a/src/data_plane/auxiliary_states/producer/mod.rs b/src/data_plane/auxiliary_states/producer/mod.rs new file mode 100644 index 00000000..2a08ec76 --- /dev/null +++ b/src/data_plane/auxiliary_states/producer/mod.rs @@ -0,0 +1,5 @@ +mod state; +mod types; +use crate::data_plane::messages::command::AuthorizedProducerIdentity; +pub(crate) use state::{AppendKey, ProducerDecision, ProducerSessions, ProducerTracker}; +pub use types::{ProduceError, ProducerAppendIdentity}; diff --git a/src/data_plane/auxiliary_states/producer/types.rs b/src/data_plane/auxiliary_states/producer/types.rs new file mode 100644 index 00000000..9501401e --- /dev/null +++ b/src/data_plane/auxiliary_states/producer/types.rs @@ -0,0 +1,35 @@ +use borsh::{BorshDeserialize, BorshSerialize}; +use uuid::Uuid; + +/// Stable identity of one producer append. The range is supplied by the +/// enclosing segment key, so it is deliberately not duplicated here. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, BorshSerialize, BorshDeserialize)] +pub struct ProducerAppendIdentity { + pub producer_id: Uuid, + pub incarnation: u32, + pub expires_at: u64, + pub sequence: u64, + pub digest: u32, +} + +#[derive(Debug, Clone, PartialEq, Eq, BorshSerialize, BorshDeserialize, thiserror::Error)] +pub enum ProduceError { + #[error("not the write leader")] + NotLeader, + #[error("segment not found")] + SegmentNotFound, + #[error("producer incarnation was fenced")] + ProducerFenced, + #[error("producer session expired or is unknown")] + SessionExpired, + #[error("sequence was reused with a different payload")] + RequestIdentityConflict, + #[error("duplicate position is outside the retained result window")] + DuplicatePositionUnavailable, + #[error("sequence gap; expected {0}")] + SequenceGap(u64), + #[error("the same producer request is already in flight")] + RequestInFlight, + #[error("internal produce failure: {0}")] + Internal(String), +} diff --git a/src/data_plane/checkpoint.rs b/src/data_plane/checkpoint.rs index c58a3ed4..d56faca7 100644 --- a/src/data_plane/checkpoint.rs +++ b/src/data_plane/checkpoint.rs @@ -4,7 +4,7 @@ use std::sync::Arc; use flume::Receiver; use crate::data_plane::actor::DataPlaneSender; -use crate::data_plane::auxiliary_states::consumer_offsets::state::AuxiliaryState; +use crate::data_plane::auxiliary_states::state::AuxiliarySnapshot; use crate::data_plane::states::segment::cache::SegmentRingBuffer; use crate::data_plane::wal::WalRecord; @@ -70,8 +70,8 @@ impl CheckpointWorker { // A snapshot is created only when WAL reclamation would cross an uncheckpointed offset record. CheckpointTask::AuxiliaryState(job) => { - if let Err(e) = job.offsets.write_snapshot(&job.data_dir) { - tracing::error!("consumer offset checkpoint failed: {e}"); + if let Err(e) = job.snapshot.write(&job.data_dir) { + tracing::error!("auxiliary state checkpoint failed: {e}"); continue; } let completion: DataPlaneCommand = AuxiliaryCheckpointComplete { @@ -146,7 +146,7 @@ pub struct CheckpointJob { } pub struct AuxiliaryCheckpointJob { - pub offsets: AuxiliaryState, + pub snapshot: AuxiliarySnapshot, pub data_dir: PathBuf, pub checkpointed_lsn: u64, } diff --git a/src/data_plane/mod.rs b/src/data_plane/mod.rs index 07a33127..88570193 100644 --- a/src/data_plane/mod.rs +++ b/src/data_plane/mod.rs @@ -1,7 +1,6 @@ use borsh::{BorshDeserialize, BorshSerialize}; use bytes::Bytes; use std::path::{Path, PathBuf}; -use uuid::Uuid; use crate::control_plane::metadata::{EntryId, RangeId, SegmentId, TopicId}; use crate::impl_new_struct_wrapper; @@ -11,7 +10,6 @@ pub(crate) mod auxiliary_states; pub(crate) mod checkpoint; pub(crate) mod cold_read; pub(crate) mod messages; -pub(crate) mod producer_ledger; pub(crate) mod recovery; pub(crate) mod segment_writer; pub(crate) mod sparse_index; @@ -30,38 +28,7 @@ pub struct SegmentKey { pub segment_id: SegmentId, } -/// Stable identity of one producer append. The range is supplied by the -/// enclosing `SegmentKey`, so it is deliberately not duplicated here. -#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, BorshSerialize, BorshDeserialize)] -pub struct ProducerAppendIdentity { - pub producer_id: Uuid, - pub incarnation: u32, - pub expires_at: u64, - pub sequence: u64, - pub digest: u32, -} - -#[derive(Debug, Clone, PartialEq, Eq, BorshSerialize, BorshDeserialize, thiserror::Error)] -pub enum ProduceError { - #[error("not the write leader")] - NotLeader, - #[error("segment not found")] - SegmentNotFound, - #[error("producer incarnation was fenced")] - ProducerFenced, - #[error("producer session expired or is unknown")] - SessionExpired, - #[error("sequence was reused with a different payload")] - RequestIdentityConflict, - #[error("duplicate position is outside the retained result window")] - DuplicatePositionUnavailable, - #[error("sequence gap; expected {0}")] - SequenceGap(u64), - #[error("the same producer request is already in flight")] - RequestInFlight, - #[error("internal produce failure: {0}")] - Internal(String), -} +pub use auxiliary_states::producer::{ProduceError, ProducerAppendIdentity}; #[derive(Debug, Clone, BorshSerialize, BorshDeserialize)] pub struct EntryPayload(Bytes); diff --git a/src/data_plane/recovery/inventory.rs b/src/data_plane/recovery/inventory.rs index a5eb3948..d319d4c4 100644 --- a/src/data_plane/recovery/inventory.rs +++ b/src/data_plane/recovery/inventory.rs @@ -4,7 +4,8 @@ use std::path::PathBuf; use super::segment_scan::RecoveredSegments; use crate::control_plane::metadata::EntryId; use crate::data_plane::SegmentKey; -use crate::data_plane::auxiliary_states::consumer_offsets::state::AuxiliaryState; +use crate::data_plane::auxiliary_states::consumer_offsets::state::ConsumerOffsets; +use crate::data_plane::auxiliary_states::producer::ProducerTracker; /// What recovery verified for one on-disk segment: its `start_offset` (the filename base, /// needed to locate the file) and `verified_end` (the highest entry id a CRC-complete batch @@ -70,7 +71,8 @@ impl LocalInventory { pub(crate) struct RecoveryOutput { pub(crate) inventory: LocalInventory, pub(crate) data_dir: PathBuf, - pub(crate) offsets: AuxiliaryState, + pub(crate) offsets: ConsumerOffsets, + pub(crate) producer: ProducerTracker, } #[cfg(test)] @@ -126,7 +128,8 @@ mod tests { let output = RecoveryOutput { inventory: inv, data_dir: PathBuf::from("/var/lib/eastguard/data"), - offsets: AuxiliaryState::default(), + offsets: ConsumerOffsets::default(), + producer: ProducerTracker::default(), }; assert_eq!(output.inventory.orphan_candidates().count(), 0); assert_eq!(output.data_dir, PathBuf::from("/var/lib/eastguard/data")); diff --git a/src/data_plane/recovery/mod.rs b/src/data_plane/recovery/mod.rs index f80b193f..13942264 100644 --- a/src/data_plane/recovery/mod.rs +++ b/src/data_plane/recovery/mod.rs @@ -17,7 +17,9 @@ use self::inventory::{LocalInventory, RecoveryOutput}; use self::replay::ReplayWriter; use self::segment_scan::RecoveredSegments; use self::wal_scan::{ScanError, WalScanner}; -use crate::data_plane::consumer_offset_management::ledger::{OffsetLedger, OffsetRecord}; +use crate::data_plane::auxiliary_states::consumer_offsets::state::{ConsumerOffsets, OffsetRecord}; +use crate::data_plane::auxiliary_states::producer::ProducerTracker; +use crate::data_plane::auxiliary_states::state::AuxiliarySnapshot; use crate::data_plane::sparse_index::SparseIndex; use crate::data_plane::states::segment::record::RoutingHeader; use crate::data_plane::wal::WalRecordType; @@ -56,7 +58,8 @@ struct Replaying { data_dir: PathBuf, writer: ReplayWriter, scanner: WalScanner, - offsets: OffsetLedger, + offsets: ConsumerOffsets, + producer: ProducerTracker, } impl Replaying { @@ -64,12 +67,13 @@ impl Replaying { let recovered = RecoveredSegments::scan_data_dir(&data_dir)?; let writer = ReplayWriter::new(data_dir.clone(), recovered); let scanner = WalScanner::open(&data_dir)?; - let offsets = OffsetLedger::load_snapshot(&data_dir)?; + let snapshot = AuxiliarySnapshot::load(&data_dir)?; Ok(Self { data_dir, writer, scanner, - offsets, + offsets: snapshot.consumer_offsets, + producer: ProducerTracker::from_sessions(snapshot.producer_sessions), }) } @@ -90,6 +94,13 @@ impl Replaying { let (header, entry_data) = RoutingHeader::split_wal_payload(&record.payload)?; self.writer.replay(&header, entry_data)?; + if let Some(producer) = header.producer { + self.producer.advance( + header.segment_key(), + producer, + header.entry_id, + ); + } } WalRecordType::ConsumerOffset => { let offset = OffsetRecord::try_from_slice(&record.payload) @@ -124,6 +135,7 @@ impl Replaying { scanner: self.scanner, recovered, offsets: self.offsets, + producer: self.producer, }) } } @@ -132,7 +144,8 @@ struct Durable { data_dir: PathBuf, scanner: WalScanner, recovered: RecoveredSegments, - offsets: OffsetLedger, + offsets: ConsumerOffsets, + producer: ProducerTracker, } impl Durable { @@ -154,12 +167,17 @@ impl Durable { } let inventory = LocalInventory::from_recovered(&self.recovered); - self.offsets.write_snapshot(&self.data_dir)?; + AuxiliarySnapshot { + consumer_offsets: self.offsets.clone(), + producer_sessions: self.producer.sessions(), + } + .write(&self.data_dir)?; self.scanner.delete_files()?; Ok(RecoveryOutput { inventory, data_dir: self.data_dir, offsets: self.offsets, + producer: self.producer, }) } } @@ -174,11 +192,13 @@ mod tests { use super::segment_scan::scan_segment_file; use super::*; use crate::control_plane::metadata::{EntryId, RangeId, SegmentId, TopicId}; - use crate::data_plane::SegmentKey; - use crate::data_plane::consumer_offset_management::ledger::{ + use crate::data_plane::auxiliary_states::consumer_offsets::state::{ ConsumerOffsetKey, ConsumerOffsetPosition, ConsumerOffsetUpdate, EpochSeal, OffsetRecord, }; + use crate::data_plane::auxiliary_states::producer::ProducerDecision; + use crate::data_plane::messages::AuthorizedProducerIdentity; use crate::data_plane::wal::{WalRecord, WalStorage, WalWriter}; + use crate::data_plane::{ProducerAppendIdentity, SegmentKey}; /// One bare-payload `(Data, BatchEnd)` batch — the segment-file framing the /// checkpoint worker writes (no routing header). @@ -199,6 +219,17 @@ mod tests { WalRecord::data(wal_payload, record_count) } + fn wal_producer_data( + key: SegmentKey, + entry_id: u64, + producer: ProducerAppendIdentity, + ) -> WalRecord { + let wal_payload = RoutingHeader::new(key, EntryId(entry_id), 1) + .with_producer(Some(producer)) + .build_wal_payload(b"data"); + WalRecord::data(wal_payload, 1) + } + fn wal_offset(record: OffsetRecord) -> WalRecord { WalRecord::consumer_offset(borsh::to_vec(&record).unwrap().into()) } @@ -316,6 +347,49 @@ mod tests { assert_eq!(restarted.offsets.offset(&key), Some(position)); } + #[test] + fn recovers_producer_deduplication_from_the_shared_wal() { + let tmp = tempfile::tempdir().unwrap(); + let data_dir = tmp.path().to_path_buf(); + let key = SegmentKey::new(TopicId(1), RangeId(2), SegmentId(0)); + let producer = ProducerAppendIdentity { + producer_id: uuid::Uuid::new_v4(), + incarnation: 1, + expires_at: u64::MAX, + sequence: 0, + digest: 9, + }; + write_wal( + &data_dir, + 1, + &wal_batch(&[wal_producer_data(key, 7, producer)]), + ); + + let (_db_dir, db) = open_db(); + let output = run(data_dir.clone(), &db).unwrap(); + let mut ledger = crate::data_plane::auxiliary_states::AuxiliaryStateManager::from_recovery( + output.offsets, + output.producer, + ); + assert_eq!( + ledger.verify_producer(key, AuthorizedProducerIdentity::ExistingOnly(producer)), + ProducerDecision::Duplicate(EntryId(7)) + ); + + let restarted = run(data_dir, &db).unwrap(); + let mut restarted_tracker = + crate::data_plane::auxiliary_states::AuxiliaryStateManager::from_recovery( + restarted.offsets, + restarted.producer, + ); + assert_eq!( + restarted_tracker + .verify_producer(key, AuthorizedProducerIdentity::ExistingOnly(producer)), + ProducerDecision::Duplicate(EntryId(7)), + "the auxiliary snapshot must recover producer sessions after WAL reclamation" + ); + } + #[test] fn re_running_recovery_converges() { let tmp = tempfile::tempdir().unwrap(); diff --git a/src/data_plane/states/segment/tracker.rs b/src/data_plane/states/segment/tracker.rs index 6a89f497..7855a5ca 100644 --- a/src/data_plane/states/segment/tracker.rs +++ b/src/data_plane/states/segment/tracker.rs @@ -3,6 +3,7 @@ use std::{path::PathBuf, sync::Arc, time::Duration}; use crate::control_plane::membership::ShardGroupId; use crate::control_plane::metadata::EntryId; use crate::control_plane::{NodeId, Replicas}; +use crate::data_plane::ProducerAppendIdentity; use crate::data_plane::{EntryPayload, SegmentKey, checkpoint::CheckpointJob, wal::WalRecord}; use super::cache::CachedEntry; @@ -106,7 +107,8 @@ impl SegmentTracker { pub(crate) fn stage_to_wal(&mut self, wal_buf: &mut Vec) { for (i, staged) in self.staged_entries.iter().enumerate() { let entry_id = self.next_entry_id + i as u64; - let header = RoutingHeader::new(staged.segment_key, entry_id, staged.record_count); + let header = RoutingHeader::new(staged.segment_key, entry_id, staged.record_count) + .with_producer(staged.producer); let _ = WalRecord::data(header.build_wal_payload(&staged.data), staged.record_count) .encode_to(wal_buf); } @@ -116,8 +118,11 @@ impl SegmentTracker { !self.staged_entries.is_empty() } - pub(crate) fn abort_staged(&mut self) { - self.staged_entries.clear(); + pub(crate) fn abort_staged(&mut self) -> Box<[crate::data_plane::ProducerAppendIdentity]> { + self.staged_entries + .drain(..) + .filter_map(|entry| entry.producer) + .collect() } pub(crate) fn publish_staged(&mut self, lsn: u64) -> Box<[Arc]> { @@ -131,6 +136,7 @@ impl SegmentTracker { record_count: s.record_count, entry_id, lsn, + producer_append_id: s.producer, }); self.cache.publish(entry.clone()); entry @@ -138,13 +144,19 @@ impl SegmentTracker { .collect() } - pub(crate) fn uncommitted_entries(&self) -> impl Iterator + '_ { + pub(crate) fn uncommitted_entries( + &self, + ) -> impl Iterator)> + '_ { let commit = self.cache.load_read_cursor(); let tail = self.cache.load_write_cursor(); (commit..tail).filter_map(|pos| { - self.cache - .load_published(pos) - .map(|entry| (entry.data.clone(), entry.record_count)) + self.cache.load_published(pos).map(|entry| { + ( + entry.data.clone(), + entry.record_count, + entry.producer_append_id, + ) + }) }) } @@ -153,17 +165,19 @@ impl SegmentTracker { /// the WAL, so a seal must carry them forward too — and they are already /// counted in the data plane's `buffer_byte_count`, so the caller must not /// re-count them. - pub(crate) fn staged_for_replay(&self) -> impl Iterator + '_ { + pub(crate) fn staged_for_replay( + &self, + ) -> impl Iterator)> + '_ { self.staged_entries .iter() - .map(|s| (s.data.clone(), s.record_count)) + .map(|s| (s.data.clone(), s.record_count, s.producer)) } // carrying over uncommitted data! pub(crate) fn replayable_bytes(&self) -> usize { self.uncommitted_entries() .chain(self.staged_for_replay()) - .map(|(data, _)| data.len()) + .map(|(data, _, _)| data.len()) .sum() } @@ -258,12 +272,13 @@ impl SegmentTracker { self.last_committed_entry_id() == self.durable_end_entry_id() } - pub(crate) fn stage_entry( + pub(crate) fn stage_producer_entry( &mut self, segment_key: SegmentKey, data: EntryPayload, record_count: u32, entry_id: EntryId, + producer_append_id: Option, ) { let expected = self.next_staged_entry_id(); if entry_id < expected { @@ -274,8 +289,12 @@ impl SegmentTracker { "entry_id gap: expected {expected}, got {entry_id}", ); self.size_bytes += data.len() as u64; - self.staged_entries - .push(StagedEntry::new(data, record_count, segment_key)); + self.staged_entries.push(StagedEntry::new( + data, + record_count, + segment_key, + producer_append_id, + )); } pub(crate) fn next_staged_entry_id(&self) -> EntryId { @@ -383,7 +402,7 @@ pub mod tests { #[test] fn stage_entry_tracks_size() { let mut t = make_tracker(SegmentRole::Leader); - t.stage_entry(test_key(), Bytes::from("abcde").into(), 2, EntryId(0)); + t.stage_producer_entry(test_key(), Bytes::from("abcde").into(), 2, EntryId(0), None); assert_eq!(t.size_bytes, 5); assert!(t.has_staged()); @@ -399,7 +418,7 @@ pub mod tests { ShardGroupId(1), EntryId(5), ); - t.stage_entry(test_key(), Bytes::from("data").into(), 1, EntryId(5)); + t.stage_producer_entry(test_key(), Bytes::from("data").into(), 1, EntryId(5), None); assert!(t.has_staged()); // Publish to advance next_entry_id @@ -408,7 +427,7 @@ pub mod tests { t.publish_staged(1); // Duplicate entry_id (5) should be skipped since next is now 6 - t.stage_entry(test_key(), Bytes::from("dup").into(), 1, EntryId(5)); + t.stage_producer_entry(test_key(), Bytes::from("dup").into(), 1, EntryId(5), None); assert!(!t.has_staged()); assert_eq!(t.next_entry_id, EntryId(6)); t.assert_invariants(); @@ -420,11 +439,12 @@ pub mod tests { let mut wal_buf = Vec::new(); for i in 0..3u64 { - t.stage_entry( + t.stage_producer_entry( test_key(), Bytes::from(format!("entry-{i}")).into(), 1, EntryId(i), + None, ); t.stage_to_wal(&mut wal_buf); t.publish_staged(i + 1); @@ -454,7 +474,13 @@ pub mod tests { ShardGroupId(1), EntryId(2), ); - t.stage_entry(test_key(), Bytes::from("entry-2").into(), 1, EntryId(2)); + t.stage_producer_entry( + test_key(), + Bytes::from("entry-2").into(), + 1, + EntryId(2), + None, + ); t.publish_staged(1); t.commit_entry(EntryId(1)); @@ -467,7 +493,13 @@ pub mod tests { fn stage_then_publish_drains() { let mut t = make_tracker(SegmentRole::Leader); let mut wal_buf = Vec::new(); - t.stage_entry(test_key(), Bytes::from("payload").into(), 3, EntryId(0)); + t.stage_producer_entry( + test_key(), + Bytes::from("payload").into(), + 3, + EntryId(0), + None, + ); t.stage_to_wal(&mut wal_buf); assert!(!wal_buf.is_empty()); @@ -527,7 +559,7 @@ pub mod tests { t.last_activity_at = tokio::time::Instant::now() - Duration::from_secs(60); assert!(t.idle_timeout_reached(Duration::from_secs(30))); - t.stage_entry(test_key(), Bytes::from("data").into(), 1, EntryId(0)); + t.stage_producer_entry(test_key(), Bytes::from("data").into(), 1, EntryId(0), None); t.publish_staged(1); t.commit_entry(EntryId(0)); From ebeef821c97982e771c069a923b5e8f70eecb215 Mon Sep 17 00:00:00 2001 From: Migorithm Date: Wed, 22 Jul 2026 20:16:08 +0400 Subject: [PATCH 15/41] producer sessions --- .../metadata/producer_sessions.rs | 70 +++++++++++++++++++ 1 file changed, 70 insertions(+) create mode 100644 src/control_plane/metadata/producer_sessions.rs diff --git a/src/control_plane/metadata/producer_sessions.rs b/src/control_plane/metadata/producer_sessions.rs new file mode 100644 index 00000000..866480fc --- /dev/null +++ b/src/control_plane/metadata/producer_sessions.rs @@ -0,0 +1,70 @@ +use std::collections::HashMap; + +use borsh::{BorshDeserialize, BorshSerialize}; + +use crate::impl_new_struct_wrapper; + +#[derive(Debug, Clone, Copy, PartialEq, Eq, BorshSerialize, BorshDeserialize)] +pub(crate) struct ProducerSessionMeta { + pub(crate) incarnation: u32, + pub(crate) expires_at: u64, + session_nonce: uuid::Uuid, +} + +#[derive(Debug, Default, Clone, PartialEq, Eq, BorshSerialize, BorshDeserialize)] +pub(crate) struct ProducerSessions(HashMap); +impl_new_struct_wrapper!(ProducerSessions,HashMap); + +impl ProducerSessions { + pub(crate) fn open_producer_session( + &mut self, + producer_id: uuid::Uuid, + session_nonce: uuid::Uuid, + observed_at: u64, + session_timeout_ms: u64, + ) -> ProducerSessionMeta { + self.expire_producer_sessions(observed_at); + + let incarnation = match self.get(&producer_id) { + Some(session) if session.session_nonce == session_nonce => session.incarnation, + Some(session) => session.incarnation.saturating_add(1), + None => 0, + }; + let expires_at = observed_at.saturating_add(session_timeout_ms); + let session = ProducerSessionMeta { + incarnation, + expires_at, + session_nonce, + }; + self.insert(producer_id, session); + session + } + + pub(crate) fn has_expired_producer_sessions(&self, observed_at: u64) -> bool { + self.values() + .any(|session| session.expires_at < observed_at) + } + + pub(crate) fn expire_producer_sessions(&mut self, observed_at: u64) { + self.retain(|_, session| session.expires_at >= observed_at); + } +} + +#[test] +fn producer_session_recovery_bumps_incarnation_and_expiry_removes_it() { + let mut producer_sessions = ProducerSessions::default(); + let producer_id = uuid::Uuid::new_v4(); + let first_nonce = uuid::Uuid::new_v4(); + + let session = producer_sessions.open_producer_session(producer_id, first_nonce, 10, 100); + assert_eq!((session.incarnation, session.expires_at), (0, 110)); + let session2 = producer_sessions.open_producer_session(producer_id, first_nonce, 20, 100); + assert_eq!((session2.incarnation, session2.expires_at), (0, 120)); + let session3 = + producer_sessions.open_producer_session(producer_id, uuid::Uuid::new_v4(), 20, 100); + assert_eq!((session3.incarnation, session3.expires_at), (1, 120)); + assert!(producer_sessions.has_expired_producer_sessions(121)); + + producer_sessions.expire_producer_sessions(121); + assert!(!producer_sessions.contains_key(&producer_id)); +} From 6c4995dab810c41fbb2240da28b1adc873fefdc5 Mon Sep 17 00:00:00 2001 From: Migorithm Date: Thu, 23 Jul 2026 01:11:44 +0400 Subject: [PATCH 16/41] AuxiliarySnapshot holding ConsumerOffsets, ProducerSessions --- src/data_plane/auxiliary_states/manager.rs | 2 +- src/data_plane/auxiliary_states/snapshot.rs | 45 +++++++++++++++++++++ src/data_plane/checkpoint.rs | 2 +- src/data_plane/recovery/mod.rs | 2 +- 4 files changed, 48 insertions(+), 3 deletions(-) create mode 100644 src/data_plane/auxiliary_states/snapshot.rs diff --git a/src/data_plane/auxiliary_states/manager.rs b/src/data_plane/auxiliary_states/manager.rs index 0e5ceb04..316d01d7 100644 --- a/src/data_plane/auxiliary_states/manager.rs +++ b/src/data_plane/auxiliary_states/manager.rs @@ -1,7 +1,7 @@ use super::consumer_offsets::state::{ ConsumerOffsetKey, ConsumerOffsetPosition, ConsumerOffsets, EpochSeal, OffsetRecord, StaleEpoch, }; -use super::state::AuxiliarySnapshot; +use super::snapshot::AuxiliarySnapshot; use crate::client::EntryId; use crate::control_plane::membership::ShardGroupId; use crate::control_plane::metadata::consumer_group::GenerationId; diff --git a/src/data_plane/auxiliary_states/snapshot.rs b/src/data_plane/auxiliary_states/snapshot.rs new file mode 100644 index 00000000..ce13f82e --- /dev/null +++ b/src/data_plane/auxiliary_states/snapshot.rs @@ -0,0 +1,45 @@ +use std::fs::{self, File, OpenOptions}; +use std::io::{self, Write}; +use std::path::Path; + +use borsh::{BorshDeserialize, BorshSerialize}; + +use super::consumer_offsets::state::ConsumerOffsets; +use super::producer::ProducerSessions; + +const SNAPSHOT_FILE: &str = "auxiliary-state.snapshot"; +const SNAPSHOT_TEMP_FILE: &str = "auxiliary-state.snapshot.tmp"; + +/// Crash-durable end state reconstructed from the shared WAL. +#[derive(Debug, Clone, Default, BorshSerialize, BorshDeserialize)] +pub(crate) struct AuxiliarySnapshot { + pub(crate) consumer_offsets: ConsumerOffsets, + pub(crate) producer_sessions: ProducerSessions, +} + +impl AuxiliarySnapshot { + pub(crate) fn load(data_dir: &Path) -> io::Result { + let bytes = match fs::read(data_dir.join(SNAPSHOT_FILE)) { + Ok(bytes) => bytes, + Err(error) if error.kind() == io::ErrorKind::NotFound => return Ok(Self::default()), + Err(error) => return Err(error), + }; + Self::try_from_slice(&bytes).map_err(io::Error::other) + } + + pub(crate) fn write(&self, data_dir: &Path) -> io::Result<()> { + fs::create_dir_all(data_dir)?; + let bytes = borsh::to_vec(self).map_err(io::Error::other)?; + let temporary = data_dir.join(SNAPSHOT_TEMP_FILE); + let final_path = data_dir.join(SNAPSHOT_FILE); + let mut file = OpenOptions::new() + .create(true) + .write(true) + .truncate(true) + .open(&temporary)?; + file.write_all(&bytes)?; + file.sync_all()?; + fs::rename(temporary, final_path)?; + File::open(data_dir)?.sync_all() + } +} diff --git a/src/data_plane/checkpoint.rs b/src/data_plane/checkpoint.rs index d56faca7..96501453 100644 --- a/src/data_plane/checkpoint.rs +++ b/src/data_plane/checkpoint.rs @@ -4,7 +4,7 @@ use std::sync::Arc; use flume::Receiver; use crate::data_plane::actor::DataPlaneSender; -use crate::data_plane::auxiliary_states::state::AuxiliarySnapshot; +use crate::data_plane::auxiliary_states::snapshot::AuxiliarySnapshot; use crate::data_plane::states::segment::cache::SegmentRingBuffer; use crate::data_plane::wal::WalRecord; diff --git a/src/data_plane/recovery/mod.rs b/src/data_plane/recovery/mod.rs index 13942264..d9b87fb8 100644 --- a/src/data_plane/recovery/mod.rs +++ b/src/data_plane/recovery/mod.rs @@ -19,7 +19,7 @@ use self::segment_scan::RecoveredSegments; use self::wal_scan::{ScanError, WalScanner}; use crate::data_plane::auxiliary_states::consumer_offsets::state::{ConsumerOffsets, OffsetRecord}; use crate::data_plane::auxiliary_states::producer::ProducerTracker; -use crate::data_plane::auxiliary_states::state::AuxiliarySnapshot; +use crate::data_plane::auxiliary_states::snapshot::AuxiliarySnapshot; use crate::data_plane::sparse_index::SparseIndex; use crate::data_plane::states::segment::record::RoutingHeader; use crate::data_plane::wal::WalRecordType; From 30da57cd09aa9b6ba101b7cc2b9584cede9e8faa Mon Sep 17 00:00:00 2001 From: Migorithm Date: Thu, 23 Jul 2026 01:17:27 +0400 Subject: [PATCH 17/41] append key moved to types --- src/data_plane/auxiliary_states/manager.rs | 23 ++++++++------ src/data_plane/auxiliary_states/mod.rs | 2 +- .../auxiliary_states/producer/mod.rs | 4 +-- .../auxiliary_states/producer/types.rs | 31 +++++++++++++++++++ 4 files changed, 48 insertions(+), 12 deletions(-) diff --git a/src/data_plane/auxiliary_states/manager.rs b/src/data_plane/auxiliary_states/manager.rs index 316d01d7..742e163c 100644 --- a/src/data_plane/auxiliary_states/manager.rs +++ b/src/data_plane/auxiliary_states/manager.rs @@ -12,6 +12,7 @@ use crate::data_plane::SegmentKey; use crate::data_plane::auxiliary_states::consumer_offsets::{ ConsumerOffsetCoordination, OffsetPlacement, OffsetReplicaState, }; +use crate::data_plane::auxiliary_states::producer::types::AppendKey; use crate::data_plane::checkpoint::AuxiliaryCheckpointJob; use crate::data_plane::messages::command::{ @@ -22,7 +23,7 @@ use crate::data_plane::messages::command::{ use crate::data_plane::messages::{RequestConsumerOffsetSnapshot, SegmentPlaced}; use crate::data_plane::ProducerAppendIdentity; -use crate::data_plane::auxiliary_states::producer::{AppendKey, ProducerDecision, ProducerTracker}; +use crate::data_plane::auxiliary_states::producer::{ProducerDecision, ProducerTracker}; use crate::data_plane::transport::command::DataTransportCommand; use std::collections::{HashMap, HashSet, VecDeque}; use std::path::PathBuf; @@ -33,7 +34,7 @@ use super::consumer_offsets::types::*; pub(crate) struct AuxiliaryStateManager { offsets: ConsumerOffsets, consumer_coord: ConsumerOffsetCoordination, // live placements, pending mutations, parked commits, and replication tracking - producer: ProducerTracker, + producer_tracker: ProducerTracker, checkpoint: AuxiliaryCheckpointState, // WAL reclamation and checkpoint progress } @@ -50,10 +51,13 @@ impl AuxiliaryStateManager { Self::from_recovery(offsets, ProducerTracker::default()) } - pub(crate) fn from_recovery(offsets: ConsumerOffsets, producer: ProducerTracker) -> Self { + pub(crate) fn from_recovery( + offsets: ConsumerOffsets, + producer_tracker: ProducerTracker, + ) -> Self { Self { offsets, - producer, + producer_tracker, ..Default::default() } } @@ -63,11 +67,11 @@ impl AuxiliaryStateManager { segment_key: SegmentKey, producer: AuthorizedProducerIdentity, ) -> ProducerDecision { - self.producer.verify(segment_key, producer) + self.producer_tracker.verify(segment_key, producer) } pub(crate) fn unstage_producer(&mut self, append_key: &AppendKey) { - self.producer.unstage(append_key); + self.producer_tracker.unstage(append_key); } pub(crate) fn advance_producer( @@ -77,7 +81,8 @@ impl AuxiliaryStateManager { entry_id: EntryId, lsn: u64, ) { - self.producer.advance(segment_key, producer, entry_id); + self.producer_tracker + .advance(segment_key, producer, entry_id); self.mark_auxiliary_state_dirty(lsn); } @@ -183,7 +188,7 @@ impl AuxiliaryStateManager { checkpointed_lsn: *self.latest_uncheckpointed_lsn()?, snapshot: AuxiliarySnapshot { consumer_offsets: self.offsets.clone(), - producer_sessions: self.producer.sessions(), + producer_sessions: self.producer_tracker.sessions(), }, data_dir, }) @@ -191,7 +196,7 @@ impl AuxiliaryStateManager { #[cfg(any(test, debug_assertions))] pub(crate) fn assert_producer_invariants(&self) { - self.producer.assert_invariants(); + self.producer_tracker.assert_invariants(); } pub(crate) fn auxiliary_checkpoint_in_flight(&self) -> bool { diff --git a/src/data_plane/auxiliary_states/mod.rs b/src/data_plane/auxiliary_states/mod.rs index bcc0557a..68055102 100644 --- a/src/data_plane/auxiliary_states/mod.rs +++ b/src/data_plane/auxiliary_states/mod.rs @@ -1,5 +1,5 @@ pub(crate) mod consumer_offsets; pub(crate) mod manager; pub(crate) mod producer; -pub(crate) mod state; +pub(crate) mod snapshot; pub(crate) use manager::AuxiliaryStateManager; diff --git a/src/data_plane/auxiliary_states/producer/mod.rs b/src/data_plane/auxiliary_states/producer/mod.rs index 2a08ec76..18b760b6 100644 --- a/src/data_plane/auxiliary_states/producer/mod.rs +++ b/src/data_plane/auxiliary_states/producer/mod.rs @@ -1,5 +1,5 @@ mod state; -mod types; +pub(crate) mod types; use crate::data_plane::messages::command::AuthorizedProducerIdentity; -pub(crate) use state::{AppendKey, ProducerDecision, ProducerSessions, ProducerTracker}; +pub(crate) use state::{ProducerDecision, ProducerSessions, ProducerTracker}; pub use types::{ProduceError, ProducerAppendIdentity}; diff --git a/src/data_plane/auxiliary_states/producer/types.rs b/src/data_plane/auxiliary_states/producer/types.rs index 9501401e..8a9e5517 100644 --- a/src/data_plane/auxiliary_states/producer/types.rs +++ b/src/data_plane/auxiliary_states/producer/types.rs @@ -1,6 +1,12 @@ use borsh::{BorshDeserialize, BorshSerialize}; use uuid::Uuid; +use crate::{ + client::RangeId, + control_plane::metadata::TopicId, + data_plane::{SegmentKey, auxiliary_states::producer::state::ProducerKey}, +}; + /// Stable identity of one producer append. The range is supplied by the /// enclosing segment key, so it is deliberately not duplicated here. #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, BorshSerialize, BorshDeserialize)] @@ -33,3 +39,28 @@ pub enum ProduceError { #[error("internal produce failure: {0}")] Internal(String), } + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub(crate) struct AppendKey { + pub(super) topic_id: TopicId, + pub(super) range_id: RangeId, + pub(super) producer_id: Uuid, + pub(super) incarnation: u32, + pub(super) sequence: u64, +} + +impl AppendKey { + pub(crate) fn new(segment: SegmentKey, request: &ProducerAppendIdentity) -> Self { + Self { + topic_id: segment.topic_id, + range_id: segment.range_id, + producer_id: request.producer_id, + incarnation: request.incarnation, + sequence: request.sequence, + } + } + + pub(crate) fn producer_key(&self) -> ProducerKey { + (self.topic_id, self.range_id, self.producer_id) + } +} From 0ddc81a0bfc7aa67be313d9d8105b6d6ec194d23 Mon Sep 17 00:00:00 2001 From: Migorithm Date: Thu, 23 Jul 2026 01:31:05 +0400 Subject: [PATCH 18/41] rm ProducerDecision --- src/data_plane/auxiliary_states/manager.rs | 5 +++-- src/data_plane/auxiliary_states/producer/mod.rs | 2 +- src/data_plane/auxiliary_states/producer/types.rs | 4 +++- src/data_plane/recovery/mod.rs | 7 +++---- 4 files changed, 10 insertions(+), 8 deletions(-) diff --git a/src/data_plane/auxiliary_states/manager.rs b/src/data_plane/auxiliary_states/manager.rs index 742e163c..42ae3491 100644 --- a/src/data_plane/auxiliary_states/manager.rs +++ b/src/data_plane/auxiliary_states/manager.rs @@ -22,8 +22,9 @@ use crate::data_plane::messages::command::{ }; use crate::data_plane::messages::{RequestConsumerOffsetSnapshot, SegmentPlaced}; +use crate::data_plane::ProduceError; use crate::data_plane::ProducerAppendIdentity; -use crate::data_plane::auxiliary_states::producer::{ProducerDecision, ProducerTracker}; +use crate::data_plane::auxiliary_states::producer::ProducerTracker; use crate::data_plane::transport::command::DataTransportCommand; use std::collections::{HashMap, HashSet, VecDeque}; use std::path::PathBuf; @@ -66,7 +67,7 @@ impl AuxiliaryStateManager { &mut self, segment_key: SegmentKey, producer: AuthorizedProducerIdentity, - ) -> ProducerDecision { + ) -> Result<(), ProduceError> { self.producer_tracker.verify(segment_key, producer) } diff --git a/src/data_plane/auxiliary_states/producer/mod.rs b/src/data_plane/auxiliary_states/producer/mod.rs index 18b760b6..ca0a4d5e 100644 --- a/src/data_plane/auxiliary_states/producer/mod.rs +++ b/src/data_plane/auxiliary_states/producer/mod.rs @@ -1,5 +1,5 @@ mod state; pub(crate) mod types; use crate::data_plane::messages::command::AuthorizedProducerIdentity; -pub(crate) use state::{ProducerDecision, ProducerSessions, ProducerTracker}; +pub(crate) use state::{ProducerSessions, ProducerTracker}; pub use types::{ProduceError, ProducerAppendIdentity}; diff --git a/src/data_plane/auxiliary_states/producer/types.rs b/src/data_plane/auxiliary_states/producer/types.rs index 8a9e5517..f399cdd1 100644 --- a/src/data_plane/auxiliary_states/producer/types.rs +++ b/src/data_plane/auxiliary_states/producer/types.rs @@ -3,7 +3,7 @@ use uuid::Uuid; use crate::{ client::RangeId, - control_plane::metadata::TopicId, + control_plane::metadata::{EntryId, TopicId}, data_plane::{SegmentKey, auxiliary_states::producer::state::ProducerKey}, }; @@ -38,6 +38,8 @@ pub enum ProduceError { RequestInFlight, #[error("internal produce failure: {0}")] Internal(String), + #[error("duplicate sequence produce: {0:?}")] + Duplicate(EntryId), } #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] diff --git a/src/data_plane/recovery/mod.rs b/src/data_plane/recovery/mod.rs index d9b87fb8..4037c689 100644 --- a/src/data_plane/recovery/mod.rs +++ b/src/data_plane/recovery/mod.rs @@ -195,10 +195,9 @@ mod tests { use crate::data_plane::auxiliary_states::consumer_offsets::state::{ ConsumerOffsetKey, ConsumerOffsetPosition, ConsumerOffsetUpdate, EpochSeal, OffsetRecord, }; - use crate::data_plane::auxiliary_states::producer::ProducerDecision; use crate::data_plane::messages::AuthorizedProducerIdentity; use crate::data_plane::wal::{WalRecord, WalStorage, WalWriter}; - use crate::data_plane::{ProducerAppendIdentity, SegmentKey}; + use crate::data_plane::{ProduceError, ProducerAppendIdentity, SegmentKey}; /// One bare-payload `(Data, BatchEnd)` batch — the segment-file framing the /// checkpoint worker writes (no routing header). @@ -373,7 +372,7 @@ mod tests { ); assert_eq!( ledger.verify_producer(key, AuthorizedProducerIdentity::ExistingOnly(producer)), - ProducerDecision::Duplicate(EntryId(7)) + Err(ProduceError::Duplicate(EntryId(7))) ); let restarted = run(data_dir, &db).unwrap(); @@ -385,7 +384,7 @@ mod tests { assert_eq!( restarted_tracker .verify_producer(key, AuthorizedProducerIdentity::ExistingOnly(producer)), - ProducerDecision::Duplicate(EntryId(7)), + Err(ProduceError::Duplicate(EntryId(7))), "the auxiliary snapshot must recover producer sessions after WAL reclamation" ); } From 4526810cf13de39e51b18d426c08810089ef4309 Mon Sep 17 00:00:00 2001 From: Migorithm Date: Thu, 23 Jul 2026 01:47:29 +0400 Subject: [PATCH 19/41] received time passing --- src/data_plane/actor.rs | 14 ++--- src/data_plane/auxiliary_states/manager.rs | 4 +- src/data_plane/messages/command.rs | 60 ++++++++++++++++------ src/data_plane/recovery/mod.rs | 9 ++-- 4 files changed, 60 insertions(+), 27 deletions(-) diff --git a/src/data_plane/actor.rs b/src/data_plane/actor.rs index 3c660ace..7439b247 100644 --- a/src/data_plane/actor.rs +++ b/src/data_plane/actor.rs @@ -154,17 +154,17 @@ fn run_orphan_gc(internal: Duration, tx: &Sender) { pub(crate) struct DataPlaneSender(pub flume::Sender); impl DataPlaneSender { - pub fn send( - &self, - msg: impl Into, - ) -> Result<(), flume::SendError> { - self.0.send(msg.into()) + pub fn send(&self, msg: impl Into) -> Result<(), flume::SendError<()>> { + self.0.send(msg.into()).map_err(|_| flume::SendError(())) } pub async fn send_async( &self, msg: impl Into, - ) -> Result<(), flume::SendError> { - self.0.send_async(msg.into()).await + ) -> Result<(), flume::SendError<()>> { + self.0 + .send_async(msg.into()) + .await + .map_err(|_| flume::SendError(())) } } diff --git a/src/data_plane/auxiliary_states/manager.rs b/src/data_plane/auxiliary_states/manager.rs index 42ae3491..20ad1558 100644 --- a/src/data_plane/auxiliary_states/manager.rs +++ b/src/data_plane/auxiliary_states/manager.rs @@ -67,8 +67,10 @@ impl AuxiliaryStateManager { &mut self, segment_key: SegmentKey, producer: AuthorizedProducerIdentity, + received_at_ms: u64, ) -> Result<(), ProduceError> { - self.producer_tracker.verify(segment_key, producer) + self.producer_tracker + .verify(segment_key, producer, received_at_ms) } pub(crate) fn unstage_producer(&mut self, append_key: &AppendKey) { diff --git a/src/data_plane/messages/command.rs b/src/data_plane/messages/command.rs index 675e61d0..ef3561d9 100644 --- a/src/data_plane/messages/command.rs +++ b/src/data_plane/messages/command.rs @@ -1,17 +1,15 @@ use crate::control_plane::Replicas; use crate::control_plane::metadata::SegmentRollIntent; -use crate::data_plane::consumer_offset_management::ledger::ConsumerOffsetEntry; -use crate::data_plane::consumer_offset_management::ledger::ConsumerOffsetUpdate; -use crate::data_plane::consumer_offset_management::ledger::EpochSeal; -use crate::data_plane::consumer_offset_management::ledger::StaleEpoch; -use crate::data_plane::consumer_offset_management::types::ReplicateConsumerOffset; +use crate::data_plane::ProduceError; +use crate::data_plane::ProducerAppendIdentity; +use crate::data_plane::auxiliary_states::consumer_offsets::state::ConsumerOffsetEntry; +use crate::data_plane::auxiliary_states::consumer_offsets::state::ConsumerOffsetUpdate; +use crate::data_plane::auxiliary_states::consumer_offsets::state::EpochSeal; +use crate::data_plane::auxiliary_states::consumer_offsets::state::StaleEpoch; +use crate::data_plane::auxiliary_states::consumer_offsets::types::ReplicateConsumerOffset; +use crate::data_plane::timer::{BatchFlushCallback, ReplicationCallback, SegmentIdleCallback}; use crate::impl_from_variant; use crate::impl_from_variant_via; -use borsh::{BorshDeserialize, BorshSerialize}; -use std::borrow::Borrow; -use std::sync::Arc; -use tokio::sync::oneshot; - use crate::{ control_plane::NodeId, control_plane::membership::ShardGroupId, @@ -19,11 +17,15 @@ use crate::{ data_plane::states::segment::cache::CachedEntry, data_plane::{EntryPayload, SegmentKey, timer::DataPlaneTimeoutCallback}, }; +use borsh::{BorshDeserialize, BorshSerialize}; +use std::borrow::Borrow; +use std::sync::Arc; +use tokio::sync::oneshot; pub enum DataPlaneCommand { Produce(Produce), SegmentCheckpointComplete(SegmentCheckpointComplete), - OffsetCheckpointComplete(OffsetCheckpointComplete), + AuxiliaryCheckpointComplete(AuxiliaryCheckpointComplete), DataPlaneTimeoutCallback(DataPlaneTimeoutCallback), ReceivePeerMessage(ReceivePeerMessage), /// Internal (not a wire message): the cold-read pool's reply for a catch-up @@ -56,9 +58,36 @@ pub struct Produce { pub segment_key: SegmentKey, pub data: EntryPayload, pub record_count: u32, + pub received_at_ms: u64, + pub producer_identity: Option, pub reply: oneshot::Sender, } +#[derive(Debug, Clone, Copy)] +pub enum AuthorizedProducerIdentity { + // Local metadata confirmed the session, so the producer tracker may register it. + MetadataVerified(ProducerAppendIdentity), + // Local metadata missed, so authority falls back to the tracker's recovered sessions. + //ONLY IF it already has an existing, non-expired session on disk for this producer. + //If the ledger doesn't know this session either, it is rejected. + ExistingOnly(ProducerAppendIdentity), +} + +impl AuthorizedProducerIdentity { + pub(crate) fn request(self) -> ProducerAppendIdentity { + match self { + Self::MetadataVerified(request) | Self::ExistingOnly(request) => request, + } + } + + pub(crate) fn destruct(self) -> (ProducerAppendIdentity, bool) { + match self { + AuthorizedProducerIdentity::MetadataVerified(req) => (req, true), + AuthorizedProducerIdentity::ExistingOnly(req) => (req, false), + } + } +} + pub struct CommitConsumerOffset { pub update: ConsumerOffsetUpdate, pub reply: oneshot::Sender, @@ -94,6 +123,7 @@ pub struct ReplicateSegmentEntries { pub data: EntryPayload, pub record_count: u32, pub entry_id: EntryId, + pub producer_identity: Option, } #[derive(Debug, Clone, BorshSerialize, BorshDeserialize)] @@ -338,7 +368,7 @@ pub enum ProduceAck { Ok { entry_id: EntryId, }, - Err(String), + Err(ProduceError), } pub struct SegmentCheckpointComplete { @@ -348,7 +378,7 @@ pub struct SegmentCheckpointComplete { pub checkpointed_bytes: u64, } -pub struct OffsetCheckpointComplete { +pub struct AuxiliaryCheckpointComplete { pub checkpointed_lsn: u64, } @@ -366,7 +396,7 @@ impl_from_variant!( DataPlaneCommand, Produce, SegmentCheckpointComplete, - OffsetCheckpointComplete, + AuxiliaryCheckpointComplete, CatchUpReadComplete, OrphanGcCheck, CommitConsumerOffset, @@ -374,8 +404,6 @@ impl_from_variant!( ReceivePeerMessage, ); -use crate::data_plane::timer::{BatchFlushCallback, ReplicationCallback, SegmentIdleCallback}; - impl_from_variant_via!( DataPlaneCommand, DataPlaneTimeoutCallback, diff --git a/src/data_plane/recovery/mod.rs b/src/data_plane/recovery/mod.rs index 4037c689..bcc799c1 100644 --- a/src/data_plane/recovery/mod.rs +++ b/src/data_plane/recovery/mod.rs @@ -371,7 +371,7 @@ mod tests { output.producer, ); assert_eq!( - ledger.verify_producer(key, AuthorizedProducerIdentity::ExistingOnly(producer)), + ledger.verify_producer(key, AuthorizedProducerIdentity::ExistingOnly(producer), 1,), Err(ProduceError::Duplicate(EntryId(7))) ); @@ -382,8 +382,11 @@ mod tests { restarted.producer, ); assert_eq!( - restarted_tracker - .verify_producer(key, AuthorizedProducerIdentity::ExistingOnly(producer)), + restarted_tracker.verify_producer( + key, + AuthorizedProducerIdentity::ExistingOnly(producer), + 1, + ), Err(ProduceError::Duplicate(EntryId(7))), "the auxiliary snapshot must recover producer sessions after WAL reclamation" ); From 95d6e74883c2a2090d497e91564b39d0437d007d Mon Sep 17 00:00:00 2001 From: Migorithm Date: Thu, 23 Jul 2026 02:05:52 +0400 Subject: [PATCH 20/41] TopicMeta::verify_producer_session --- src/client/mod.rs | 4 ++-- src/connections/protocol/data_plane.rs | 2 +- src/control_plane/metadata/topic.rs | 33 +++++++++++++++++++++++++- src/it/e2e/client_protocol.rs | 2 +- src/it/e2e/producer_deduplication.rs | 2 +- src/it/sim/scenario.rs | 2 +- 6 files changed, 38 insertions(+), 7 deletions(-) diff --git a/src/client/mod.rs b/src/client/mod.rs index 4c5fdc04..886bb083 100644 --- a/src/client/mod.rs +++ b/src/client/mod.rs @@ -337,7 +337,7 @@ impl Client { routing_key: &[u8], data: Vec, record_count: u32, - producer: Option, + producer_identity: Option, ) -> Result { // Describe once to seed the cache (gives the first hop). let routing = self.resolve_topic_if_missing(topic).await?; @@ -355,7 +355,7 @@ impl Client { routing_key: routing_key.to_vec(), data, record_count, - producer_append_id: producer, + producer_identity, }; let served = self.call(start, request).await?; diff --git a/src/connections/protocol/data_plane.rs b/src/connections/protocol/data_plane.rs index 11845fb0..90ebd4c7 100644 --- a/src/connections/protocol/data_plane.rs +++ b/src/connections/protocol/data_plane.rs @@ -61,7 +61,7 @@ pub struct ProduceRequest { /// The broker stamps an entry_id and stores/replicates this opaque payload as-is. pub data: Vec, pub record_count: u32, - pub producer_append_id: Option, + pub producer_identity: Option, } #[derive(Debug, Clone, BorshSerialize, BorshDeserialize)] diff --git a/src/control_plane/metadata/topic.rs b/src/control_plane/metadata/topic.rs index 7418601f..3d88811b 100644 --- a/src/control_plane/metadata/topic.rs +++ b/src/control_plane/metadata/topic.rs @@ -12,7 +12,10 @@ use crate::{ strategy::{PartitionStrategy, StoragePolicy}, }, }, - data_plane::SegmentKey, + data_plane::{ + ProduceError, ProducerAppendIdentity, SegmentKey, + messages::command::AuthorizedProducerIdentity, + }, }; use borsh::{BorshDeserialize, BorshSerialize}; @@ -505,6 +508,34 @@ impl TopicMeta { range.delete(); } } + + pub(crate) fn verify_producer_session( + &self, + q: ProducerAppendIdentity, + received_at_ms: u64, + ) -> Result { + let Some(session) = self.producer_sessions.get(&q.producer_id) else { + // Data replica placement is independent from metadata-Raft placement. + // A restarted data leader can therefore serve before its local metadata + // view contains this session. The durable range ledger still enforces + // the lease, sequence, and incarnation; do not strand recovery on a + // non-authoritative local metadata miss. + return Ok(AuthorizedProducerIdentity::ExistingOnly(q)); + }; + + if q.incarnation < session.incarnation { + return Err(ProduceError::ProducerFenced); + }; + + if q.incarnation == session.incarnation + && q.expires_at <= session.expires_at + && q.expires_at >= received_at_ms + { + return Ok(AuthorizedProducerIdentity::MetadataVerified(q)); + } + + Err(ProduceError::SessionExpired) + } } pub struct TopicStats { diff --git a/src/it/e2e/client_protocol.rs b/src/it/e2e/client_protocol.rs index ae8445c4..713fc222 100644 --- a/src/it/e2e/client_protocol.rs +++ b/src/it/e2e/client_protocol.rs @@ -1832,7 +1832,7 @@ async fn produce_once(topic: &str, payload: &[u8], node: (&str, u16)) -> Produce routing_key: b"k".to_vec(), data: payload.to_vec(), record_count: 1, - producer_append_id: None, + producer_identity: None, })); match send_request(node.0, node.1, req).await { ClientResponse::DataPlane(DataPlaneResponse::Produced { .. }) => ProduceOutcome::Acked, diff --git a/src/it/e2e/producer_deduplication.rs b/src/it/e2e/producer_deduplication.rs index 3f2ea61b..fbf49119 100644 --- a/src/it/e2e/producer_deduplication.rs +++ b/src/it/e2e/producer_deduplication.rs @@ -118,7 +118,7 @@ fn duplicate_unknown_and_fenced_sessions_are_end_to_end_visible() -> turmoil::Re routing_key: b"k".to_vec(), data: payload.clone(), record_count: 1, - producer_append_id: Some(unknown), + producer_identity: Some(unknown), })); let mut rejected = false; for (host, port, _) in NODES { diff --git a/src/it/sim/scenario.rs b/src/it/sim/scenario.rs index 2e1f7ae9..9f227cb0 100644 --- a/src/it/sim/scenario.rs +++ b/src/it/sim/scenario.rs @@ -403,7 +403,7 @@ async fn produce_until_acked( routing_key: b"campaign-key".to_vec(), data: payload.to_vec(), record_count: 1, - producer_append_id: None, + producer_identity: None, })); loop { for (host, port) in nodes { From 35e42527bf854bb7034fc687d59c6dbb1bec446b Mon Sep 17 00:00:00 2001 From: Migorithm Date: Thu, 23 Jul 2026 02:34:02 +0400 Subject: [PATCH 21/41] clear out the boundary between Controller vs Data Response --- src/client/mod.rs | 17 ++++++++------- src/client/redirect.rs | 17 ++++++--------- src/connections/protocol/control_plane.rs | 26 +++++++++++++++++++++-- src/connections/protocol/data_plane.rs | 8 +------ src/data_plane/recovery/mod.rs | 4 ++-- src/data_plane/states/segment/tracker.rs | 8 +++---- src/it/e2e/client_protocol.rs | 6 +++--- 7 files changed, 50 insertions(+), 36 deletions(-) diff --git a/src/client/mod.rs b/src/client/mod.rs index 886bb083..5fbeb9b3 100644 --- a/src/client/mod.rs +++ b/src/client/mod.rs @@ -47,6 +47,7 @@ use crate::connections::protocol::{ ConsumerGroupAssignmentResponse, ConsumerGroupSyncAction, ControlPlaneRequest, ControlPlaneResponse, DataPlaneResponse, FetchConsumerOffsetRequest, OpenProducerSessionRequest, ProduceRequest, ProducerSessionOpened, RangeOffsetRequest, + ShardNotLocal, }; use crate::data_plane::{ProduceError, ProducerAppendIdentity}; @@ -432,9 +433,9 @@ impl Client { ClientResponse::DataPlane(DataPlaneResponse::SegmentNotLocal) => { Redirect::Done } - ClientResponse::DataPlane(DataPlaneResponse::ShardNotLocal { - hint_node: None, - }) => Redirect::Done, + ClientResponse::ControlPlane(ControlPlaneResponse::ShardNotLocal( + ShardNotLocal { hint_node: None }, + )) => Redirect::Done, _ => Self::redirect_target(&response), } } else { @@ -488,6 +489,11 @@ impl Client { None => Redirect::Reresolve, }, ControlPlaneResponse::TopicNotFound => Redirect::NotFound, + ControlPlaneResponse::ShardNotLocal(ShardNotLocal { hint_node }) => match hint_node + { + Some(addr) => Redirect::Follow(addr.client_addr()), + None => Redirect::Reresolve, + }, ControlPlaneResponse::InternalError(_) => Redirect::Reresolve, ControlPlaneResponse::TopicCreated | ControlPlaneResponse::AlreadyExists @@ -503,11 +509,6 @@ impl Client { Some(addr) => Redirect::Follow(addr.client_addr()), None => Redirect::Reresolve, }, - DataPlaneResponse::ShardNotLocal { hint_node } => match hint_node { - Some(addr) => Redirect::Follow(addr.client_addr()), - None => Redirect::Reresolve, - }, - DataPlaneResponse::TopicNotFound => Redirect::NotFound, DataPlaneResponse::StaleRange => Redirect::Done, DataPlaneResponse::InternalError(_) => Redirect::Reresolve, DataPlaneResponse::SegmentNotLocal diff --git a/src/client/redirect.rs b/src/client/redirect.rs index 0294e033..8ce98af4 100644 --- a/src/client/redirect.rs +++ b/src/client/redirect.rs @@ -15,6 +15,7 @@ use crate::client::Client; use crate::client::error::ClientError; use crate::connections::protocol::{ ClientDataPlaneRequest, ClientRequest, ClientResponse, ControlPlaneResponse, DataPlaneResponse, + ShardNotLocal, }; /// Consecutive hint-follows before forcing a backoff — breaks a redirect cycle @@ -175,6 +176,12 @@ mod tests { )), "notfound" ); + assert_eq!( + classify(&ClientResponse::ControlPlane( + ControlPlaneResponse::ShardNotLocal(ShardNotLocal { hint_node: None }) + )), + "reresolve" + ); assert_eq!( classify(&ClientResponse::ControlPlane( ControlPlaneResponse::TopicCreated @@ -193,16 +200,6 @@ mod tests { )), "follow" ); - assert_eq!( - classify(&ClientResponse::DataPlane( - DataPlaneResponse::ShardNotLocal { hint_node: None } - )), - "reresolve" - ); - assert_eq!( - classify(&ClientResponse::DataPlane(DataPlaneResponse::TopicNotFound)), - "notfound" - ); assert_eq!( classify(&ClientResponse::DataPlane( DataPlaneResponse::SegmentNotLocal diff --git a/src/connections/protocol/control_plane.rs b/src/connections/protocol/control_plane.rs index b5b0819e..22f7d409 100644 --- a/src/connections/protocol/control_plane.rs +++ b/src/connections/protocol/control_plane.rs @@ -21,8 +21,8 @@ use borsh::{BorshDeserialize, BorshSerialize}; use std::collections::{HashMap, HashSet}; use crate::control_plane::metadata::{ - EntryId, RangeId, RangeMeta, RangeState, SegmentId, SegmentMeta, SegmentMetaState, - SyncConsumerGroupRequest, TopicId, TopicMeta, TopicState, + EntryId, OpenProducerSession, RangeId, RangeMeta, RangeState, SegmentId, SegmentMeta, + SegmentMetaState, SyncConsumerGroupRequest, TopicId, TopicMeta, TopicState, }; #[derive(Debug, Clone, BorshSerialize, BorshDeserialize)] @@ -55,6 +55,20 @@ pub struct OpenProducerSessionRequest { pub session_nonce: uuid::Uuid, } +impl OpenProducerSessionRequest { + pub fn into_command(self) -> OpenProducerSession { + const SESSION_TIMEOUT_MS: u64 = 60_000; + let observed_at = crate::now_ms(); + OpenProducerSession { + topic_name: self.topic_name, + producer_id: self.producer_id, + session_nonce: self.session_nonce, + observed_at, + session_timeout_ms: SESSION_TIMEOUT_MS, + } + } +} + #[derive(Debug, Clone, Copy, PartialEq, Eq, BorshSerialize, BorshDeserialize)] pub enum ConsumerGroupSyncAction { Heartbeat, @@ -69,6 +83,9 @@ pub enum ControlPlaneResponse { // DeleteTopic TopicDeleted, TopicNotFound, + /// This node does not own the topic's metadata shard. The hint is absent + /// while membership or shard routing is still converging. + ShardNotLocal(ShardNotLocal), // ListHostedTopics TopicList { topics: Box<[TopicSummary]>, @@ -91,6 +108,11 @@ pub enum ControlPlaneResponse { InternalError(String), } +#[derive(Debug, BorshSerialize, BorshDeserialize)] +pub struct ShardNotLocal { + pub hint_node: Option, +} + #[derive(Debug, Clone, Copy, BorshSerialize, BorshDeserialize)] pub struct ProducerSessionOpened { pub incarnation: u32, diff --git a/src/connections/protocol/data_plane.rs b/src/connections/protocol/data_plane.rs index 90ebd4c7..1b88ef17 100644 --- a/src/connections/protocol/data_plane.rs +++ b/src/connections/protocol/data_plane.rs @@ -137,11 +137,7 @@ pub enum DataPlaneResponse { NotWriteLeader { leader_addr: Option, }, - ShardNotLocal { - hint_node: Option, - }, StaleRange, - TopicNotFound, SegmentNotLocal, ConsumerOffsetCommitted, @@ -155,9 +151,7 @@ impl DataPlaneResponse { pub fn is_routing_error(&self) -> bool { matches!( self, - DataPlaneResponse::SegmentNotLocal - | DataPlaneResponse::ShardNotLocal { .. } - | DataPlaneResponse::NotWriteLeader { .. } + DataPlaneResponse::SegmentNotLocal | DataPlaneResponse::NotWriteLeader { .. } ) } diff --git a/src/data_plane/recovery/mod.rs b/src/data_plane/recovery/mod.rs index bcc799c1..2b81940f 100644 --- a/src/data_plane/recovery/mod.rs +++ b/src/data_plane/recovery/mod.rs @@ -94,10 +94,10 @@ impl Replaying { let (header, entry_data) = RoutingHeader::split_wal_payload(&record.payload)?; self.writer.replay(&header, entry_data)?; - if let Some(producer) = header.producer { + if let Some(producer_identity) = header.producer_identity { self.producer.advance( header.segment_key(), - producer, + producer_identity, header.entry_id, ); } diff --git a/src/data_plane/states/segment/tracker.rs b/src/data_plane/states/segment/tracker.rs index 7855a5ca..df68015b 100644 --- a/src/data_plane/states/segment/tracker.rs +++ b/src/data_plane/states/segment/tracker.rs @@ -108,7 +108,7 @@ impl SegmentTracker { for (i, staged) in self.staged_entries.iter().enumerate() { let entry_id = self.next_entry_id + i as u64; let header = RoutingHeader::new(staged.segment_key, entry_id, staged.record_count) - .with_producer(staged.producer); + .with_producer(staged.producer_identity); let _ = WalRecord::data(header.build_wal_payload(&staged.data), staged.record_count) .encode_to(wal_buf); } @@ -121,7 +121,7 @@ impl SegmentTracker { pub(crate) fn abort_staged(&mut self) -> Box<[crate::data_plane::ProducerAppendIdentity]> { self.staged_entries .drain(..) - .filter_map(|entry| entry.producer) + .filter_map(|entry| entry.producer_identity) .collect() } @@ -136,7 +136,7 @@ impl SegmentTracker { record_count: s.record_count, entry_id, lsn, - producer_append_id: s.producer, + producer_append_id: s.producer_identity, }); self.cache.publish(entry.clone()); entry @@ -170,7 +170,7 @@ impl SegmentTracker { ) -> impl Iterator)> + '_ { self.staged_entries .iter() - .map(|s| (s.data.clone(), s.record_count, s.producer)) + .map(|s| (s.data.clone(), s.record_count, s.producer_identity)) } // carrying over uncommitted data! diff --git a/src/it/e2e/client_protocol.rs b/src/it/e2e/client_protocol.rs index 713fc222..1c5ba041 100644 --- a/src/it/e2e/client_protocol.rs +++ b/src/it/e2e/client_protocol.rs @@ -7,7 +7,7 @@ use crate::StartUp; use crate::connections::protocol::{ AdminRequest, AdminResponse, ClientDataPlaneRequest, ClientRequest, ClientResponse, ControlPlaneRequest, ControlPlaneResponse, DataPlaneResponse, FetchByIdRequest, FetchRequest, - NodeState, ProduceRequest, RangeProgressSignal, TopicDetail, + NodeState, ProduceRequest, RangeProgressSignal, ShardNotLocal, TopicDetail, }; use crate::control_plane::membership::ShardGroupId; use crate::control_plane::metadata::strategy::{PartitionStrategy, StoragePolicy}; @@ -1839,9 +1839,9 @@ async fn produce_once(topic: &str, payload: &[u8], node: (&str, u16)) -> Produce ClientResponse::DataPlane(DataPlaneResponse::NotWriteLeader { leader_addr: Some(addr), }) => ProduceOutcome::Redirect(addr.client_addr()), - ClientResponse::DataPlane(DataPlaneResponse::ShardNotLocal { + ClientResponse::ControlPlane(ControlPlaneResponse::ShardNotLocal(ShardNotLocal { hint_node: Some(addr), - }) => ProduceOutcome::Redirect(addr.client_addr()), + })) => ProduceOutcome::Redirect(addr.client_addr()), _ => ProduceOutcome::Retry, } } From 338b9d4da276c12e31bd68ea453b7ebf03087252 Mon Sep 17 00:00:00 2001 From: Migorithm Date: Thu, 23 Jul 2026 02:42:29 +0400 Subject: [PATCH 22/41] header creation --- src/data_plane/states/segment/record.rs | 143 +++++++++++++++++++++--- 1 file changed, 129 insertions(+), 14 deletions(-) diff --git a/src/data_plane/states/segment/record.rs b/src/data_plane/states/segment/record.rs index 8808b295..79cecf78 100644 --- a/src/data_plane/states/segment/record.rs +++ b/src/data_plane/states/segment/record.rs @@ -3,9 +3,24 @@ use std::io; use bytes::{BufMut, Bytes, BytesMut}; use crate::control_plane::metadata::{EntryId, RangeId, SegmentId, TopicId}; -use crate::data_plane::{EntryPayload, SegmentKey}; - -const ROUTING_HEADER_SIZE: usize = 36; +use crate::data_plane::{EntryPayload, ProducerAppendIdentity, SegmentKey}; + +const ROUTING_HEADER_MAGIC: u32 = 0x4547_5032; +const ROUTING_HEADER_SIZE: usize = 81; +// Current frame: magic | legacy routing fields | producer-present | +// producer UUID | incarnation | expiry | sequence | digest. +const MAGIC_SIZE: usize = 4; +const TOPIC_OFFSET: usize = 0; +const RANGE_OFFSET: usize = 8; +const SEGMENT_OFFSET: usize = 16; +const ENTRY_OFFSET: usize = 24; +const RECORD_COUNT_OFFSET: usize = 32; +const PRODUCER_PRESENT_OFFSET: usize = 36; +const PRODUCER_ID_OFFSET: usize = 37; +const PRODUCER_INCARNATION_OFFSET: usize = 53; +const PRODUCER_EXPIRY_OFFSET: usize = 57; +const PRODUCER_SEQUENCE_OFFSET: usize = 65; +const PRODUCER_DIGEST_OFFSET: usize = 73; #[derive(Debug, Clone, PartialEq, Eq)] pub(crate) struct RoutingHeader { @@ -14,6 +29,7 @@ pub(crate) struct RoutingHeader { segment_id: SegmentId, pub(crate) entry_id: EntryId, pub(crate) record_count: u32, + pub(crate) producer_identity: Option, } impl RoutingHeader { @@ -24,19 +40,40 @@ impl RoutingHeader { segment_id: key.segment_id, entry_id, record_count, + producer_identity: None, } } + pub(crate) fn with_producer(mut self, producer: Option) -> Self { + self.producer_identity = producer; + self + } + pub(crate) fn segment_key(&self) -> SegmentKey { SegmentKey::new(self.topic_id, self.range_id, self.segment_id) } fn encode(&self, buf: &mut BytesMut) { + buf.put_u32(ROUTING_HEADER_MAGIC); buf.put_u64(self.topic_id.0); buf.put_u64(*self.range_id); buf.put_u64(self.segment_id.0); buf.put_u64(*self.entry_id); buf.put_u32(self.record_count); + match self.producer_identity { + Some(producer) => { + buf.put_u8(1); + buf.put_slice(producer.producer_id.as_bytes()); + buf.put_u32(producer.incarnation); + buf.put_u64(producer.expires_at); + buf.put_u64(producer.sequence); + buf.put_u32(producer.digest); + } + None => { + buf.put_u8(0); + buf.put_bytes(0, 40); + } + } } pub(crate) fn build_wal_payload(&self, entry_data: &[u8]) -> Bytes { @@ -46,20 +83,62 @@ impl RoutingHeader { buf.freeze() } - pub(crate) fn decode(data: &[u8]) -> io::Result { + fn decode_with_size(data: &[u8]) -> io::Result<(Self, usize)> { if data.len() < ROUTING_HEADER_SIZE { return Err(io::Error::new( io::ErrorKind::UnexpectedEof, "routing header too short", )); } - Ok(RoutingHeader { - topic_id: TopicId(u64::from_be_bytes(data[0..8].try_into().unwrap())), - range_id: RangeId(u64::from_be_bytes(data[8..16].try_into().unwrap())), - segment_id: SegmentId(u64::from_be_bytes(data[16..24].try_into().unwrap())), - entry_id: EntryId(u64::from_be_bytes(data[24..32].try_into().unwrap())), - record_count: u32::from_be_bytes(data[32..36].try_into().unwrap()), - }) + + // Verify magic header + if u32::from_be_bytes(data[0..4].try_into().unwrap()) != ROUTING_HEADER_MAGIC { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + "invalid routing header magic", + )); + } + + // Parse fixed fields + let body = &data[MAGIC_SIZE..]; + let topic_id = TopicId(Self::parse_u64(&body[TOPIC_OFFSET..])); + let range_id = RangeId(Self::parse_u64(&body[RANGE_OFFSET..])); + let segment_id = SegmentId(Self::parse_u64(&body[SEGMENT_OFFSET..])); + let entry_id = EntryId(Self::parse_u64(&body[ENTRY_OFFSET..])); + let record_count = Self::parse_u32(&body[RECORD_COUNT_OFFSET..]); + + // Parse optional producer identity + let producer_identity = + (body[PRODUCER_PRESENT_OFFSET] == 1).then(|| ProducerAppendIdentity { + producer_id: uuid::Uuid::from_slice( + &body[PRODUCER_ID_OFFSET..PRODUCER_INCARNATION_OFFSET], + ) + .unwrap(), + incarnation: Self::parse_u32(&body[PRODUCER_INCARNATION_OFFSET..]), + expires_at: Self::parse_u64(&body[PRODUCER_EXPIRY_OFFSET..]), + sequence: Self::parse_u64(&body[PRODUCER_SEQUENCE_OFFSET..]), + digest: Self::parse_u32(&body[PRODUCER_DIGEST_OFFSET..]), + }); + + Ok(( + RoutingHeader { + topic_id, + range_id, + segment_id, + entry_id, + record_count, + producer_identity, + }, + ROUTING_HEADER_SIZE, + )) + } + + fn parse_u32(src: &[u8]) -> u32 { + u32::from_be_bytes(src[..4].try_into().unwrap()) + } + + fn parse_u64(src: &[u8]) -> u64 { + u64::from_be_bytes(src[..8].try_into().unwrap()) } /// Splits a WAL record payload (as built by [`RoutingHeader::build_wal_payload`]) @@ -67,8 +146,13 @@ impl RoutingHeader { /// inverse of `build_wal_payload`. Recovery replay uses it to route a WAL /// record by its header while writing only the bare data the segment keeps. pub(crate) fn split_wal_payload(payload: &[u8]) -> io::Result<(Self, &[u8])> { - let header = Self::decode(payload)?; - Ok((header, &payload[ROUTING_HEADER_SIZE..])) + let (header, size) = Self::decode_with_size(payload)?; + Ok((header, &payload[size..])) + } + + #[cfg(test)] + pub(crate) fn decode(data: &[u8]) -> io::Result { + Self::decode_with_size(data).map(|(header, _)| header) } } @@ -76,14 +160,21 @@ pub(crate) struct StagedEntry { pub(crate) data: EntryPayload, pub(crate) record_count: u32, pub(crate) segment_key: SegmentKey, + pub(crate) producer_identity: Option, } impl StagedEntry { - pub(crate) fn new(data: EntryPayload, record_count: u32, segment_key: SegmentKey) -> Self { + pub(crate) fn new( + data: EntryPayload, + record_count: u32, + segment_key: SegmentKey, + producer_identity: Option, + ) -> Self { Self { data, record_count, segment_key, + producer_identity, } } @@ -104,6 +195,7 @@ mod tests { segment_id: SegmentId(3), entry_id: EntryId(999), record_count: 50, + producer_identity: None, }; let mut buf = BytesMut::new(); header.encode(&mut buf); @@ -120,6 +212,7 @@ mod tests { segment_id: SegmentId(3), entry_id: EntryId(100), record_count: 5, + producer_identity: None, }; let entry_data = b"opaque entry payload"; let payload = header.build_wal_payload(entry_data); @@ -130,4 +223,26 @@ mod tests { assert_eq!(header, decoded_header); assert_eq!(&decoded_data[..], entry_data); } + + #[test] + fn producer_identity_roundtrips_inside_the_data_wal_record() { + let producer = ProducerAppendIdentity { + producer_id: uuid::Uuid::new_v4(), + incarnation: 3, + expires_at: 99, + sequence: 7, + digest: 42, + }; + let header = RoutingHeader::new( + SegmentKey::new(TopicId(1), RangeId(2), SegmentId(3)), + EntryId(4), + 5, + ) + .with_producer(Some(producer)); + let payload = header.build_wal_payload(b"data"); + let (decoded, data) = RoutingHeader::split_wal_payload(&payload).unwrap(); + + assert_eq!(decoded.producer_identity, Some(producer)); + assert_eq!(data, b"data"); + } } From ecfafeba475629bf9a26015dbc132f7ac8d3b7df Mon Sep 17 00:00:00 2001 From: Migorithm Date: Thu, 23 Jul 2026 02:47:45 +0400 Subject: [PATCH 23/41] move stale range from data response to controller response --- src/client/mod.rs | 6 +++--- src/connections/protocol/control_plane.rs | 1 + src/connections/protocol/data_plane.rs | 1 - 3 files changed, 4 insertions(+), 4 deletions(-) diff --git a/src/client/mod.rs b/src/client/mod.rs index 5fbeb9b3..7a5dd6a2 100644 --- a/src/client/mod.rs +++ b/src/client/mod.rs @@ -365,14 +365,14 @@ impl Client { if served.redirected || matches!( &served.response, - ClientResponse::DataPlane(DataPlaneResponse::StaleRange) + ClientResponse::ControlPlane(ControlPlaneResponse::StaleRange) ) { self.cache.invalidate(topic); } match served.response { ClientResponse::DataPlane(DataPlaneResponse::Produced { entry_id }) => Ok(entry_id), - ClientResponse::DataPlane(DataPlaneResponse::StaleRange) => { + ClientResponse::ControlPlane(ControlPlaneResponse::StaleRange) => { Err(ClientError::StaleRange) } ClientResponse::DataPlane(DataPlaneResponse::ProduceRejected(error)) => { @@ -496,6 +496,7 @@ impl Client { }, ControlPlaneResponse::InternalError(_) => Redirect::Reresolve, ControlPlaneResponse::TopicCreated + | ControlPlaneResponse::StaleRange | ControlPlaneResponse::AlreadyExists | ControlPlaneResponse::TopicDeleted | ControlPlaneResponse::ConsumerGroupAssignment(_) @@ -509,7 +510,6 @@ impl Client { Some(addr) => Redirect::Follow(addr.client_addr()), None => Redirect::Reresolve, }, - DataPlaneResponse::StaleRange => Redirect::Done, DataPlaneResponse::InternalError(_) => Redirect::Reresolve, DataPlaneResponse::SegmentNotLocal | DataPlaneResponse::Produced { .. } diff --git a/src/connections/protocol/control_plane.rs b/src/connections/protocol/control_plane.rs index 22f7d409..b485061f 100644 --- a/src/connections/protocol/control_plane.rs +++ b/src/connections/protocol/control_plane.rs @@ -83,6 +83,7 @@ pub enum ControlPlaneResponse { // DeleteTopic TopicDeleted, TopicNotFound, + StaleRange, /// This node does not own the topic's metadata shard. The hint is absent /// while membership or shard routing is still converging. ShardNotLocal(ShardNotLocal), diff --git a/src/connections/protocol/data_plane.rs b/src/connections/protocol/data_plane.rs index 1b88ef17..f86f8ea6 100644 --- a/src/connections/protocol/data_plane.rs +++ b/src/connections/protocol/data_plane.rs @@ -137,7 +137,6 @@ pub enum DataPlaneResponse { NotWriteLeader { leader_addr: Option, }, - StaleRange, SegmentNotLocal, ConsumerOffsetCommitted, From 18a7736355c0cb3180ef0663c650d3ba65565ede Mon Sep 17 00:00:00 2001 From: Migorithm Date: Thu, 23 Jul 2026 10:14:21 +0400 Subject: [PATCH 24/41] refactor: flattening responses We have completed the refactoring of the wire protocol to flatten responses directly into ClientResponse::Ok(ClientSuccess) and ClientResponse::Err(ServerError) across the codebase: 1. Wire Protocol Structure: - Removed ControlPlaneResponse, DataPlaneResponse, and AdminResponse sub-protocol wrapper enums. - Updated ClientSuccess in mod.rs to directly contain all 16 success variants (TopicCreated, TopicDeleted, TopicList, TopicDetail, ConsumerGroupAssignment, ConsumerGroupLeft, ProducerSessionOpened, Produced, Fetched, RangeOffset, ConsumerOffsetCommitted, ConsumerOffset, ClusterInfo, TopicStats, ShardInfo, ShardLeader). - Standardized ClientResponse to consist of: - ClientResponse::Ok(ClientSuccess) - ClientResponse::Err(ServerError) - ClientResponse::Stop 2. Server Controller Integration: - Refactored handler methods in controller.rs (handle_create_topic, handle_delete_topic, handle_produce, handle_fetch, etc.) to return Result. 3. Client SDK & Integration Tests --- src/client/consumer/group.rs | 8 +- src/client/consumer/mod.rs | 4 +- src/client/consumer/range_fetcher.rs | 209 +++++++++++----------- src/client/mod.rs | 107 +++++------ src/client/redirect.rs | 53 +++--- src/connections/protocol/admin.rs | 25 +-- src/connections/protocol/control_plane.rs | 51 +----- src/connections/protocol/data_plane.rs | 118 ++---------- src/connections/protocol/error.rs | 72 ++++++++ src/connections/protocol/mod.rs | 126 +++++++++++-- src/it/e2e/client_protocol.rs | 129 ++++++------- src/it/e2e/client_sdk.rs | 30 ++-- src/it/e2e/producer_deduplication.rs | 10 +- src/it/helpers.rs | 7 +- src/it/sim/invariants.rs | 10 +- src/it/sim/properties.rs | 7 +- src/it/sim/scenario.rs | 27 ++- 17 files changed, 475 insertions(+), 518 deletions(-) create mode 100644 src/connections/protocol/error.rs diff --git a/src/client/consumer/group.rs b/src/client/consumer/group.rs index 2d9026e8..384f4b90 100644 --- a/src/client/consumer/group.rs +++ b/src/client/consumer/group.rs @@ -7,8 +7,8 @@ use arc_swap::ArcSwap; use dashmap::DashMap; use uuid::Uuid; -use crate::client::{Client, ClientError}; -use crate::connections::protocol::{ClientResponse, ConsumerGroupSyncAction, ControlPlaneResponse}; +use crate::client::{Client, ClientError, ClientSuccess}; +use crate::connections::protocol::{ClientResponse, ConsumerGroupSyncAction}; use crate::control_plane::metadata::consumer_group::GenerationId; use crate::control_plane::metadata::{EntryId, RangeId, SyncConsumerGroupRequest, TopicId}; use crate::data_plane::auxiliary_states::consumer_offsets::state::{ @@ -201,9 +201,7 @@ impl ConsumerGroup { .await? .response { - ClientResponse::ControlPlane(ControlPlaneResponse::ConsumerGroupAssignment( - assignment, - )) => assignment, + ClientResponse::Ok(ClientSuccess::ConsumerGroupAssignment(assignment)) => assignment, _ => return Err(ClientError::UnexpectedResponse), }; diff --git a/src/client/consumer/mod.rs b/src/client/consumer/mod.rs index c6122aa3..c74c23ec 100644 --- a/src/client/consumer/mod.rs +++ b/src/client/consumer/mod.rs @@ -10,8 +10,8 @@ use crate::client::consumer::topic_fetch_manager::{ use crate::client::redirect::Served; use crate::client::{Client, ClientError}; use crate::connections::protocol::{ - ClientDataPlaneRequest, ClientResponse, DataPlaneResponse, FetchByIdRequest, RangeDetail, - RangeOffsetRequest, RangeProgressSignal, RangeTransition, SegmentDetail, TopicDetail, + ClientDataPlaneRequest, ClientResponse, FetchByIdRequest, RangeDetail, RangeOffsetRequest, + RangeProgressSignal, RangeTransition, SegmentDetail, TopicDetail, }; use crate::control_plane::metadata::{EntryId, RangeId, RangeState, TopicId}; use crate::data_plane::auxiliary_states::consumer_offsets::state::ConsumerOffsetPosition; diff --git a/src/client/consumer/range_fetcher.rs b/src/client/consumer/range_fetcher.rs index 2ff15084..a54c6343 100644 --- a/src/client/consumer/range_fetcher.rs +++ b/src/client/consumer/range_fetcher.rs @@ -5,9 +5,9 @@ use super::{ConsumerContext, ConsumerRecord, RangeCursor}; use crate::client::consumer::context::RangeLookupResult; use crate::client::consumer::messages::RangeDrained; use crate::client::redirect::Served; -use crate::client::{ClientError, CompressionCodec}; +use crate::client::{ClientError, ClientSuccess, CompressionCodec, ServerError}; use crate::connections::protocol::{ - ClientResponse, DataPlaneResponse, RangeProgressSignal, RangeTransition, SegmentDetail, + ClientResponse, EntryPayload, RangeProgressSignal, RangeTransition, SegmentDetail, }; use crate::control_plane::metadata::{EntryId, RangeId}; use crate::data_plane::auxiliary_states::consumer_offsets::state::ConsumerOffsetPosition; @@ -137,115 +137,16 @@ impl RangeFetchActor { .await? }; - let ClientResponse::DataPlane(dp_response) = served.response else { - return Err(ClientError::UnexpectedResponse); - }; - - self.handle_data_plane_response(dp_response).await - } - - async fn handle_data_plane_response( - &mut self, - dp_response: DataPlaneResponse, - ) -> Result { - match dp_response { - DataPlaneResponse::Fetched { + match served.response { + ClientResponse::Ok(ClientSuccess::Fetched { entries, next_entry_id, progress_signal, - } => { - if entries.is_empty() { - // ! short-polling backoff mechanism to prevent the consumer from unintentionally DDoS-ing the data plane server - tokio::time::sleep(Duration::from_millis(50)).await; - if self.group_fetch { - self.consecutive_empty_fetches = - self.consecutive_empty_fetches.saturating_add(1); - if self.consecutive_empty_fetches >= EMPTY_FETCHES_BEFORE_REFRESH { - self.ctx.refresh_metadata().await?; - self.consecutive_empty_fetches = 0; - } - } - } else { - self.consecutive_empty_fetches = 0; - } - - for entry in entries.into_iter() { - if self.should_skip_entry_by_absolute_offset(entry.record_count) { - continue; - } - - let records = CompressionCodec::decode_payload(&entry.data, entry.record_count) - .map_err(|e| { - tracing::error!("Failed to decompress entry payload: {}", e); - ClientError::UnexpectedResponse - })?; - - let skip_batch_offsets_below = if entry.entry_id == self.cursor.next_entry_id { - self.cursor.skip_batch_offsets_below.take() - } else { - None - }; - - for (i, (key, value)) in records.into_iter().enumerate() { - if skip_batch_offsets_below.is_some_and(|skip| i as u64 <= skip) { - continue; - } - - let absolute_offset = self.cursor.next_absolute_offset; - self.cursor.next_absolute_offset = - self.cursor.next_absolute_offset.saturating_add(1); - - if self - .cursor - .skip_absolute_offsets_below - .is_some_and(|target| absolute_offset < target) - { - continue; - } - self.cursor.skip_absolute_offsets_below = None; - - let consumer_rec = ConsumerRecord { - topic: self.ctx.topic.clone(), - range_id: self.cursor.range_id, - position: ConsumerOffsetPosition { - batch_offset: i as u64, - entry_id: entry.entry_id, - absolute_offset, - }, - key, - value, - }; - - if self.record_tx.send(Ok(consumer_rec)).is_err() { - return Ok(false); // Disconnected - } - } - } - - // Data plane returns Fetched response with progress_signal, - // If the range is permanently sealed, it attaches Sealed signal. - // end_entry_id suggest hard stop point. - // - // next_entry_id is offset of the next record the consumer needs to fetch. - // So the end_entry_id could be 999 while next_entry_id is only 501, suggesting that you still have more to fetch in this range - if let RangeProgressSignal::Sealed { - end_entry_id, - transition, - } = progress_signal - && next_entry_id > end_entry_id - { - self.cursor.next_entry_id = next_entry_id; - let _ = self.ctx.cursor_tx.send(RangeDrained { - cursor: self.cursor.clone(), - transition, - }); - return Ok(false); - } - - self.cursor.next_entry_id = next_entry_id; - Ok(true) + }) => { + self.handle_fetched(entries, next_entry_id, progress_signal) + .await } - DataPlaneResponse::EntryIdOutOfRange => { + ClientResponse::Err(ServerError::EntryIdOutOfRange) => { let prev_entry_id = self.cursor.next_entry_id; let (start_id, _) = self .ctx @@ -262,8 +163,7 @@ impl RangeFetchActor { } Ok(true) } - - dp if dp.is_routing_error() => { + ClientResponse::Err(err) if err.is_routing_error() => { tokio::time::sleep(Duration::from_millis(50)).await; self.ctx.refresh_metadata().await?; Ok(true) @@ -272,6 +172,97 @@ impl RangeFetchActor { } } + async fn handle_fetched( + &mut self, + entries: Box<[EntryPayload]>, + next_entry_id: EntryId, + progress_signal: RangeProgressSignal, + ) -> Result { + if entries.is_empty() { + // ! short-polling backoff mechanism to prevent the consumer from unintentionally DDoS-ing the data plane server + tokio::time::sleep(Duration::from_millis(50)).await; + if self.group_fetch { + self.consecutive_empty_fetches = self.consecutive_empty_fetches.saturating_add(1); + if self.consecutive_empty_fetches >= EMPTY_FETCHES_BEFORE_REFRESH { + self.ctx.refresh_metadata().await?; + self.consecutive_empty_fetches = 0; + } + } + } else { + self.consecutive_empty_fetches = 0; + } + + for entry in entries.into_vec() { + if self.should_skip_entry_by_absolute_offset(entry.record_count) { + continue; + } + + let records = CompressionCodec::decode_payload(&entry.data, entry.record_count) + .map_err(|e| { + tracing::error!("Failed to decompress entry payload: {}", e); + ClientError::UnexpectedResponse + })?; + + let skip_batch_offsets_below = if entry.entry_id == self.cursor.next_entry_id { + self.cursor.skip_batch_offsets_below.take() + } else { + None + }; + + for (i, (key, value)) in records.into_iter().enumerate() { + if skip_batch_offsets_below.is_some_and(|skip| i as u64 <= skip) { + continue; + } + + let absolute_offset = self.cursor.next_absolute_offset; + self.cursor.next_absolute_offset = + self.cursor.next_absolute_offset.saturating_add(1); + + if self + .cursor + .skip_absolute_offsets_below + .is_some_and(|target| absolute_offset < target) + { + continue; + } + self.cursor.skip_absolute_offsets_below = None; + + let consumer_rec = ConsumerRecord { + topic: self.ctx.topic.clone(), + range_id: self.cursor.range_id, + position: ConsumerOffsetPosition { + batch_offset: i as u64, + entry_id: entry.entry_id, + absolute_offset, + }, + key, + value, + }; + + if self.record_tx.send(Ok(consumer_rec)).is_err() { + return Ok(false); // Disconnected + } + } + } + + if let RangeProgressSignal::Sealed { + end_entry_id, + transition, + } = progress_signal + && next_entry_id > end_entry_id + { + self.cursor.next_entry_id = next_entry_id; + let _ = self.ctx.cursor_tx.send(RangeDrained { + cursor: self.cursor.clone(), + transition, + }); + return Ok(false); + } + + self.cursor.next_entry_id = next_entry_id; + Ok(true) + } + fn should_skip_entry_by_absolute_offset(&mut self, record_count: u32) -> bool { let Some(target) = self.cursor.skip_absolute_offsets_below else { return false; diff --git a/src/client/mod.rs b/src/client/mod.rs index 7a5dd6a2..ea299516 100644 --- a/src/client/mod.rs +++ b/src/client/mod.rs @@ -22,7 +22,7 @@ mod routing; use crate::client::nodes::KnownNodes; use crate::client::routing::TopicRouting; -pub use crate::connections::protocol::{TopicDetail, TopicSummary}; +pub use crate::connections::protocol::{ClientSuccess, ServerError, TopicDetail, TopicSummary}; #[cfg(test)] use crate::control_plane::NodeAddressInfo; use crate::control_plane::metadata::consumer_group::GenerationId; @@ -45,9 +45,8 @@ use uuid::Uuid; use crate::connections::protocol::{ ClientDataPlaneRequest, ClientRequest, ClientResponse, CommitConsumerOffsetRequest, ConsumerGroupAssignmentResponse, ConsumerGroupSyncAction, ControlPlaneRequest, - ControlPlaneResponse, DataPlaneResponse, FetchConsumerOffsetRequest, - OpenProducerSessionRequest, ProduceRequest, ProducerSessionOpened, RangeOffsetRequest, - ShardNotLocal, + FetchConsumerOffsetRequest, OpenProducerSessionRequest, ProduceRequest, ProducerSessionOpened, + RangeOffsetRequest, }; use crate::data_plane::{ProduceError, ProducerAppendIdentity}; @@ -157,11 +156,11 @@ impl Client { let served = self.call(addr, req).await?; match served.response { - ClientResponse::DataPlane(DataPlaneResponse::RangeOffset { + ClientResponse::Ok(ClientSuccess::RangeOffset { start_entry_id, next_entry_id, }) => return Ok((start_entry_id, next_entry_id)), - ClientResponse::DataPlane(DataPlaneResponse::SegmentNotLocal) => { + ClientResponse::Err(ServerError::SegmentNotLocal) => { last_error = Some("range offsets segment not local".to_string()); } _ => return Err(ClientError::UnexpectedResponse), @@ -192,8 +191,8 @@ impl Client { self.cache.invalidate(topic_name); } match served.response { - ClientResponse::DataPlane(DataPlaneResponse::ConsumerOffsetCommitted) => Ok(()), - ClientResponse::DataPlane(DataPlaneResponse::StaleConsumerGroupEpoch(stale)) => { + ClientResponse::Ok(ClientSuccess::ConsumerOffsetCommitted) => Ok(()), + ClientResponse::Err(ServerError::StaleConsumerGroupEpoch(stale)) => { Err(ClientError::StaleConsumerGroupEpoch { request_generation, sealed_generation: stale, @@ -231,12 +230,12 @@ impl Client { self.cache.invalidate(topic_name); } match served.response { - ClientResponse::DataPlane(DataPlaneResponse::ConsumerOffset(position)) => { + ClientResponse::Ok(ClientSuccess::ConsumerOffset(position)) => { Ok(position.map(|p| (range_id, p))) } - ClientResponse::DataPlane( - DataPlaneResponse::ConsumerOffsetGenerationMismatch(mismatch), - ) if mismatch.observed_generation > generation => { + ClientResponse::Err(ServerError::ConsumerOffsetGenerationMismatch( + mismatch, + )) if mismatch.observed_generation > generation => { Err(ClientError::StaleConsumerGroupEpoch { request_generation: generation, sealed_generation: mismatch.observed_generation, @@ -262,8 +261,8 @@ impl Client { }; let served = self.call(self.next_known_node(), request).await?; match served.response { - ClientResponse::ControlPlane(ControlPlaneResponse::TopicCreated) => Ok(true), - ClientResponse::ControlPlane(ControlPlaneResponse::AlreadyExists) => Ok(false), + ClientResponse::Ok(ClientSuccess::TopicCreated) => Ok(true), + ClientResponse::Err(ServerError::AlreadyExists) => Ok(false), _ => Err(ClientError::UnexpectedResponse), } } @@ -277,7 +276,7 @@ impl Client { // The redirect loop already turned a `TopicNotFound` response into an error. self.cache.invalidate(name); match served.response { - ClientResponse::ControlPlane(ControlPlaneResponse::TopicDeleted) => Ok(()), + ClientResponse::Ok(ClientSuccess::TopicDeleted) => Ok(()), _ => Err(ClientError::UnexpectedResponse), } } @@ -301,7 +300,7 @@ impl Client { }; let served = self.call(self.next_known_node(), request).await?; match served.response { - ClientResponse::ControlPlane(ControlPlaneResponse::TopicDetail(detail)) => { + ClientResponse::Ok(ClientSuccess::TopicDetail(detail)) => { self.remember_nodes(&detail); self.cache.insert(&detail); Ok(detail) @@ -321,9 +320,7 @@ impl Client { .call(self.next_known_node(), ControlPlaneRequest::from(req)) .await?; match served.response { - ClientResponse::ControlPlane(ControlPlaneResponse::ProducerSessionOpened(session)) => { - Ok(session) - } + ClientResponse::Ok(ClientSuccess::ProducerSessionOpened(session)) => Ok(session), _ => Err(ClientError::UnexpectedResponse), } } @@ -365,17 +362,15 @@ impl Client { if served.redirected || matches!( &served.response, - ClientResponse::ControlPlane(ControlPlaneResponse::StaleRange) + ClientResponse::Err(ServerError::StaleRange) ) { self.cache.invalidate(topic); } match served.response { - ClientResponse::DataPlane(DataPlaneResponse::Produced { entry_id }) => Ok(entry_id), - ClientResponse::ControlPlane(ControlPlaneResponse::StaleRange) => { - Err(ClientError::StaleRange) - } - ClientResponse::DataPlane(DataPlaneResponse::ProduceRejected(error)) => { + ClientResponse::Ok(ClientSuccess::Produced { entry_id }) => Ok(entry_id), + ClientResponse::Err(ServerError::StaleRange) => Err(ClientError::StaleRange), + ClientResponse::Err(ServerError::ProduceRejected(error)) => { Err(ClientError::ProduceRejected(error)) } _ => Err(ClientError::UnexpectedResponse), @@ -430,12 +425,10 @@ impl Client { let redirect = if is_direct_to_node { match &response { - ClientResponse::DataPlane(DataPlaneResponse::SegmentNotLocal) => { + ClientResponse::Err(ServerError::SegmentNotLocal) => Redirect::Done, + ClientResponse::Err(ServerError::ShardNotLocal { hint_node: None }) => { Redirect::Done } - ClientResponse::ControlPlane(ControlPlaneResponse::ShardNotLocal( - ShardNotLocal { hint_node: None }, - )) => Redirect::Done, _ => Self::redirect_target(&response), } } else { @@ -480,51 +473,31 @@ impl Client { /// won't compile until handled. pub(crate) fn redirect_target(response: &ClientResponse) -> Redirect { match response { - ClientResponse::ControlPlane(cp) => match cp { - ControlPlaneResponse::TopicMetadataRedirect { owner } => { - Redirect::Follow(owner.client_addr()) - } - ControlPlaneResponse::NotRaftLeader { leader_addr } => match leader_addr { - Some(addr) => Redirect::Follow(addr.client_addr()), - None => Redirect::Reresolve, - }, - ControlPlaneResponse::TopicNotFound => Redirect::NotFound, - ControlPlaneResponse::ShardNotLocal(ShardNotLocal { hint_node }) => match hint_node - { + ClientResponse::Ok(_) => Redirect::Done, + ClientResponse::Err(err) => match err { + ServerError::ShardNotLocal { hint_node } => match hint_node { Some(addr) => Redirect::Follow(addr.client_addr()), None => Redirect::Reresolve, }, - ControlPlaneResponse::InternalError(_) => Redirect::Reresolve, - ControlPlaneResponse::TopicCreated - | ControlPlaneResponse::StaleRange - | ControlPlaneResponse::AlreadyExists - | ControlPlaneResponse::TopicDeleted - | ControlPlaneResponse::ConsumerGroupAssignment(_) - | ControlPlaneResponse::ConsumerGroupLeft - | ControlPlaneResponse::ProducerSessionOpened(_) - | ControlPlaneResponse::TopicList { .. } - | ControlPlaneResponse::TopicDetail(_) => Redirect::Done, - }, - ClientResponse::DataPlane(dp) => match dp { - DataPlaneResponse::NotWriteLeader { leader_addr } => match leader_addr { + ServerError::NotRaftLeader { leader_addr } + | ServerError::NotWriteLeader { leader_addr } => match leader_addr { Some(addr) => Redirect::Follow(addr.client_addr()), None => Redirect::Reresolve, }, - DataPlaneResponse::InternalError(_) => Redirect::Reresolve, - DataPlaneResponse::SegmentNotLocal - | DataPlaneResponse::Produced { .. } - | DataPlaneResponse::ProduceRejected(..) - | DataPlaneResponse::Fetched { .. } - | DataPlaneResponse::EntryIdOutOfRange - | DataPlaneResponse::KeyspaceBoundNarrowed - | DataPlaneResponse::RangeOffset { .. } => Redirect::Done, - DataPlaneResponse::ConsumerOffsetCommitted - | DataPlaneResponse::StaleConsumerGroupEpoch(_) - | DataPlaneResponse::ConsumerOffset(_) - | DataPlaneResponse::ConsumerOffsetGenerationMismatch(_) => Redirect::Done, + ServerError::TopicMetadataRedirect { owner } => { + Redirect::Follow(owner.client_addr()) + } + ServerError::TopicNotFound => Redirect::NotFound, + ServerError::SegmentNotLocal | ServerError::Internal(_) => Redirect::Reresolve, + ServerError::AlreadyExists + | ServerError::StaleRange + | ServerError::ProduceRejected(_) + | ServerError::EntryIdOutOfRange + | ServerError::KeyspaceBoundNarrowed + | ServerError::StaleConsumerGroupEpoch(_) + | ServerError::ConsumerOffsetGenerationMismatch(_) + | ServerError::InvalidSplitPoint => Redirect::Done, }, - ClientResponse::Admin(_) => Redirect::Done, - // The server's writer-loop sentinel; a client never legitimately reads it. ClientResponse::Stop => Redirect::Reresolve, } } diff --git a/src/client/redirect.rs b/src/client/redirect.rs index 8ce98af4..c025f7fb 100644 --- a/src/client/redirect.rs +++ b/src/client/redirect.rs @@ -14,8 +14,7 @@ use tokio::time::Instant; use crate::client::Client; use crate::client::error::ClientError; use crate::connections::protocol::{ - ClientDataPlaneRequest, ClientRequest, ClientResponse, ControlPlaneResponse, DataPlaneResponse, - ShardNotLocal, + ClientDataPlaneRequest, ClientRequest, ClientResponse, ClientSuccess, ServerError, }; /// Consecutive hint-follows before forcing a backoff — breaks a redirect cycle @@ -151,41 +150,35 @@ mod tests { #[test] fn control_plane_redirects_map_to_actions() { assert_eq!( - classify(&ClientResponse::ControlPlane( - ControlPlaneResponse::TopicMetadataRedirect { owner: info(8081) } - )), + classify(&ClientResponse::Err(ServerError::TopicMetadataRedirect { + owner: info(8081) + })), "follow" ); assert_eq!( - classify(&ClientResponse::ControlPlane( - ControlPlaneResponse::NotRaftLeader { - leader_addr: Some(info(8082)) - } - )), + classify(&ClientResponse::Err(ServerError::NotRaftLeader { + leader_addr: Some(info(8082)) + })), "follow" ); assert_eq!( - classify(&ClientResponse::ControlPlane( - ControlPlaneResponse::NotRaftLeader { leader_addr: None } - )), + classify(&ClientResponse::Err(ServerError::NotRaftLeader { + leader_addr: None + })), "reresolve" ); assert_eq!( - classify(&ClientResponse::ControlPlane( - ControlPlaneResponse::TopicNotFound - )), + classify(&ClientResponse::Err(ServerError::TopicNotFound)), "notfound" ); assert_eq!( - classify(&ClientResponse::ControlPlane( - ControlPlaneResponse::ShardNotLocal(ShardNotLocal { hint_node: None }) - )), + classify(&ClientResponse::Err(ServerError::ShardNotLocal { + hint_node: None + })), "reresolve" ); assert_eq!( - classify(&ClientResponse::ControlPlane( - ControlPlaneResponse::TopicCreated - )), + classify(&ClientResponse::Ok(ClientSuccess::TopicCreated)), "done" ); } @@ -193,21 +186,17 @@ mod tests { #[test] fn data_plane_redirects_map_to_actions() { assert_eq!( - classify(&ClientResponse::DataPlane( - DataPlaneResponse::NotWriteLeader { - leader_addr: Some(info(8083)) - } - )), + classify(&ClientResponse::Err(ServerError::NotWriteLeader { + leader_addr: Some(info(8083)) + })), "follow" ); assert_eq!( - classify(&ClientResponse::DataPlane( - DataPlaneResponse::SegmentNotLocal - )), - "done" + classify(&ClientResponse::Err(ServerError::SegmentNotLocal)), + "reresolve" ); assert_eq!( - classify(&ClientResponse::DataPlane(DataPlaneResponse::Produced { + classify(&ClientResponse::Ok(ClientSuccess::Produced { entry_id: 7.into() })), "done" diff --git a/src/connections/protocol/admin.rs b/src/connections/protocol/admin.rs index e9207374..746f9aed 100644 --- a/src/connections/protocol/admin.rs +++ b/src/connections/protocol/admin.rs @@ -17,24 +17,7 @@ pub enum AdminRequest { GetShardLeader { shard_group_id: ShardGroupId }, } -#[derive(Debug, BorshSerialize, BorshDeserialize)] -pub enum AdminResponse { - // DescribeCluster - ClusterInfo { nodes: Box<[NodeInfo]> }, - // ListHostedTopicsWithStats - TopicStats { topics: Box<[TopicStats]> }, - // SplitRange - RangeSplit, - InvalidSplitPoint, - // GetShardInfo - ShardInfo { detail: Option }, - // GetShardLeader - ShardLeader { leader: Option }, - // All admin operations - InternalError(String), -} - -#[derive(Debug, BorshSerialize, BorshDeserialize)] +#[derive(Debug, Clone, PartialEq, Eq, BorshSerialize, BorshDeserialize)] pub struct ShardDetail { pub shard_group_id: ShardGroupId, pub leader_node_id: Option, @@ -42,21 +25,21 @@ pub struct ShardDetail { pub member_node_ids: Box<[String]>, } -#[derive(Debug, Clone, BorshSerialize, BorshDeserialize)] +#[derive(Debug, Clone, PartialEq, Eq, BorshSerialize, BorshDeserialize)] pub struct NodeInfo { pub node_id: String, pub addr: SocketAddr, pub state: NodeState, } -#[derive(Debug, Clone, BorshSerialize, BorshDeserialize)] +#[derive(Debug, Clone, PartialEq, Eq, BorshSerialize, BorshDeserialize)] pub enum NodeState { Alive, Suspect, Dead, } -#[derive(Debug, Clone, BorshSerialize, BorshDeserialize)] +#[derive(Debug, Clone, PartialEq, Eq, BorshSerialize, BorshDeserialize)] pub struct TopicStats { pub name: String, pub range_count: u32, diff --git a/src/connections/protocol/control_plane.rs b/src/connections/protocol/control_plane.rs index b485061f..9392a83a 100644 --- a/src/connections/protocol/control_plane.rs +++ b/src/connections/protocol/control_plane.rs @@ -75,58 +75,19 @@ pub enum ConsumerGroupSyncAction { Leave, } -#[derive(Debug, BorshSerialize, BorshDeserialize)] -pub enum ControlPlaneResponse { - // CreateTopic - TopicCreated, - AlreadyExists, - // DeleteTopic - TopicDeleted, - TopicNotFound, - StaleRange, - /// This node does not own the topic's metadata shard. The hint is absent - /// while membership or shard routing is still converging. - ShardNotLocal(ShardNotLocal), - // ListHostedTopics - TopicList { - topics: Box<[TopicSummary]>, - }, - // DescribeTopic — when this node owns the topic's metadata - TopicDetail(TopicDetail), - // DescribeTopic — when this node does not own the topic's metadata. - TopicMetadataRedirect { - owner: NodeAddressInfo, - }, - // Write landed on a member that isn't the metadata Raft leader — client retries - // there. `None` when the leader isn't yet known. - NotRaftLeader { - leader_addr: Option, - }, - ConsumerGroupAssignment(ConsumerGroupAssignmentResponse), - ConsumerGroupLeft, - ProducerSessionOpened(ProducerSessionOpened), - // All control plane operations - InternalError(String), -} - -#[derive(Debug, BorshSerialize, BorshDeserialize)] -pub struct ShardNotLocal { - pub hint_node: Option, -} - -#[derive(Debug, Clone, Copy, BorshSerialize, BorshDeserialize)] +#[derive(Debug, Clone, Copy, PartialEq, Eq, BorshSerialize, BorshDeserialize)] pub struct ProducerSessionOpened { pub incarnation: u32, pub expires_at: u64, } -#[derive(Debug, BorshSerialize, BorshDeserialize)] +#[derive(Debug, Clone, PartialEq, Eq, BorshSerialize, BorshDeserialize)] pub struct ConsumerGroupAssignmentResponse { pub generation: GenerationId, pub ranges: Box<[RangeId]>, } -#[derive(Debug, Clone, BorshSerialize, BorshDeserialize)] +#[derive(Debug, Clone, PartialEq, Eq, BorshSerialize, BorshDeserialize)] pub struct TopicSummary { pub name: String, pub range_count: u32, @@ -136,7 +97,7 @@ pub struct TopicSummary { /// Full metadata snapshot for a topic, returned by `DescribeTopic`. Includes every /// range — active, sealed, and deleting — so a consumer can walk lineage from any /// historical ancestor forward through splits and merges. -#[derive(Debug, BorshSerialize, BorshDeserialize)] +#[derive(Debug, Clone, PartialEq, Eq, BorshSerialize, BorshDeserialize)] pub struct TopicDetail { /// Resolved id, so a consumer can fetch by id directly from a holding replica /// (no name re-resolution on the data node). See `FetchByIdRequest`. @@ -265,7 +226,7 @@ pub struct RebalancePlan { /// Per-range metadata. Lineage fields (`split_into`, `merged_into`, `merged_from`) /// let the consumer reconstruct the DAG; only the predecessor that owned a given /// key needs to be drained before its successor. See d4_consumer_range_tracking.md. -#[derive(Debug, Clone, BorshSerialize, BorshDeserialize)] +#[derive(Debug, Clone, PartialEq, Eq, BorshSerialize, BorshDeserialize)] pub struct RangeDetail { pub range_id: RangeId, pub keyspace_start: Vec, @@ -363,7 +324,7 @@ impl RangeDetail { /// Per-segment metadata exposed to consumers. The replica set carries resolved /// client addresses so the consumer can pick a node to send fetches to without a /// separate address-resolution round trip. -#[derive(Debug, Clone, BorshSerialize, BorshDeserialize)] +#[derive(Debug, Clone, PartialEq, Eq, BorshSerialize, BorshDeserialize)] pub struct SegmentDetail { pub segment_id: SegmentId, pub start_entry_id: EntryId, diff --git a/src/connections/protocol/data_plane.rs b/src/connections/protocol/data_plane.rs index f86f8ea6..b075b274 100644 --- a/src/connections/protocol/data_plane.rs +++ b/src/connections/protocol/data_plane.rs @@ -9,23 +9,23 @@ use borsh::{BorshDeserialize, BorshSerialize}; use crate::{ - control_plane::{ - NodeAddressInfo, - metadata::{ - EntryId, RangeId, RangeMeta, RangeState, TopicId, TopicMeta, - consumer_group::GenerationId, - }, + control_plane::metadata::{ + EntryId, RangeId, RangeMeta, RangeState, TopicId, TopicMeta, consumer_group::GenerationId, }, data_plane::{ - ProduceError, ProducerAppendIdentity, - auxiliary_states::consumer_offsets::state::{ - ConsumerOffsetKey, ConsumerOffsetPosition, ConsumerOffsetUpdate, - }, - messages::query::{FetchResult, ListOffsetsResult}, + ProducerAppendIdentity, + auxiliary_states::consumer_offsets::state::{ConsumerOffsetKey, ConsumerOffsetUpdate}, }, impl_from_variant, impl_new_struct_wrapper, }; +#[derive(Debug, Clone, PartialEq, Eq, BorshSerialize, BorshDeserialize)] +pub struct EntryPayload { + pub entry_id: EntryId, + pub record_count: u32, + pub data: Vec, +} + /// Client → broker data-plane request. Every variant carries a `Client*Request` /// struct so the carried fields are named in one place (consistent with the /// project's tuple-variant + named-struct enum pattern) and the `Client` @@ -106,101 +106,11 @@ pub struct FetchConsumerOffsetRequest { pub generation: GenerationId, } -#[derive(Debug, Clone, Copy, BorshSerialize, BorshDeserialize)] +#[derive(Debug, Clone, Copy, PartialEq, Eq, BorshSerialize, BorshDeserialize)] pub struct ConsumerOffsetGenerationMismatch { pub observed_generation: GenerationId, } -#[derive(Debug, BorshSerialize, BorshDeserialize)] -pub enum DataPlaneResponse { - // Produce - Produced { - entry_id: EntryId, - }, - ProduceRejected(ProduceError), - // Fetch - Fetched { - entries: Box<[Entry]>, - next_entry_id: EntryId, - progress_signal: RangeProgressSignal, - }, - EntryIdOutOfRange, - // keyspace_bound was set but narrower than the target range's keyspace. - KeyspaceBoundNarrowed, - - RangeOffset { - start_entry_id: EntryId, - next_entry_id: EntryId, - }, - // `NotWriteLeader` is the segment's data-replica write leader (`replica_set[0]`), - // distinct from the metadata Raft leader (`ControlPlaneResponse::NotRaftLeader`). - NotWriteLeader { - leader_addr: Option, - }, - SegmentNotLocal, - - ConsumerOffsetCommitted, - StaleConsumerGroupEpoch(GenerationId), - ConsumerOffset(Option), - ConsumerOffsetGenerationMismatch(ConsumerOffsetGenerationMismatch), - InternalError(String), -} - -impl DataPlaneResponse { - pub fn is_routing_error(&self) -> bool { - matches!( - self, - DataPlaneResponse::SegmentNotLocal | DataPlaneResponse::NotWriteLeader { .. } - ) - } - - pub(crate) fn from_list_offset_result(value: ListOffsetsResult) -> Self { - match value { - ListOffsetsResult::RangeOffsets { - start_entry_id, - next_entry_id, - } => DataPlaneResponse::RangeOffset { - start_entry_id, - next_entry_id, - }, - ListOffsetsResult::SegmentNotLocal => DataPlaneResponse::SegmentNotLocal, - } - } - - pub(crate) fn from_fetch_result(result: FetchResult) -> Self { - match result { - FetchResult::Records { - entries, - next_entry_id, - progress_signal, - } => { - // Single Bytes → Vec copy per entry at the wire boundary - // — borsh's owned-byte encoding requires Vec. The - // intermediate `FetchedEntry` step that used to live in the - // data plane is gone; the Arc rides straight through the - // channel from the cache. - let wire_entries = entries - .into_iter() - .map(|cached| Entry { - entry_id: cached.entry_id, - // Auto-deref EntryPayload → Bytes → [u8], then to_vec. - data: cached.data.to_vec(), - record_count: cached.record_count, - }) - .collect(); - DataPlaneResponse::Fetched { - entries: wire_entries, - next_entry_id, - progress_signal, - } - } - FetchResult::EntryIdOutOfRange => DataPlaneResponse::EntryIdOutOfRange, - FetchResult::SegmentNotLocal => DataPlaneResponse::SegmentNotLocal, - FetchResult::InternalError(s) => DataPlaneResponse::InternalError(s), - } - } -} - #[derive(Clone, Debug, BorshSerialize, BorshDeserialize)] pub struct KeyspaceBound { pub start: Vec, @@ -223,7 +133,7 @@ pub struct Entry { /// payload) — tells the consumer whether to keep fetching this range or drain /// to `end_offset` and follow the lineage transition to its successor(s). /// See (d4_consumer_range_tracking.md, "Range Transitions".) -#[derive(Debug, BorshSerialize, BorshDeserialize)] +#[derive(Debug, Clone, PartialEq, Eq, BorshSerialize, BorshDeserialize)] pub enum RangeProgressSignal { Active, Sealed { @@ -232,7 +142,7 @@ pub enum RangeProgressSignal { }, } -#[derive(Debug, BorshSerialize, BorshDeserialize)] +#[derive(Debug, Clone, PartialEq, Eq, BorshSerialize, BorshDeserialize)] pub enum RangeTransition { Split { left_range_id: RangeId, diff --git a/src/connections/protocol/error.rs b/src/connections/protocol/error.rs new file mode 100644 index 00000000..1222a966 --- /dev/null +++ b/src/connections/protocol/error.rs @@ -0,0 +1,72 @@ +use borsh::{BorshDeserialize, BorshSerialize}; + +use crate::connections::protocol::data_plane::ConsumerOffsetGenerationMismatch; +use crate::control_plane::NodeAddressInfo; +use crate::control_plane::metadata::consumer_group::GenerationId; +use crate::data_plane::ProduceError; + +#[derive(Debug, Clone, PartialEq, Eq, BorshSerialize, BorshDeserialize, thiserror::Error)] +pub enum ServerError { + #[error("shard not local")] + ShardNotLocal { hint_node: Option }, + + #[error("not metadata raft leader")] + NotRaftLeader { + leader_addr: Option, + }, + + #[error("not write leader")] + NotWriteLeader { + leader_addr: Option, + }, + + #[error("topic metadata redirect")] + TopicMetadataRedirect { owner: NodeAddressInfo }, + + #[error("topic not found")] + TopicNotFound, + + #[error("topic already exists")] + AlreadyExists, + + #[error("stale range")] + StaleRange, + + #[error("produce rejected: {0}")] + ProduceRejected(ProduceError), + + #[error("entry id out of range")] + EntryIdOutOfRange, + + #[error("keyspace bound narrowed")] + KeyspaceBoundNarrowed, + + #[error("segment not local")] + SegmentNotLocal, + + #[error("stale consumer group epoch: {0:?}")] + StaleConsumerGroupEpoch(GenerationId), + + #[error("consumer offset generation mismatch")] + ConsumerOffsetGenerationMismatch(ConsumerOffsetGenerationMismatch), + + #[error("invalid split point")] + InvalidSplitPoint, + + #[error("internal server error: {0}")] + Internal(String), +} + +impl ServerError { + pub fn is_routing_error(&self) -> bool { + matches!( + self, + ServerError::ShardNotLocal { .. } + | ServerError::NotRaftLeader { .. } + | ServerError::NotWriteLeader { .. } + | ServerError::TopicMetadataRedirect { .. } + | ServerError::SegmentNotLocal + | ServerError::StaleRange + ) + } +} diff --git a/src/connections/protocol/mod.rs b/src/connections/protocol/mod.rs index 7d8e08ad..44a97ab9 100644 --- a/src/connections/protocol/mod.rs +++ b/src/connections/protocol/mod.rs @@ -20,15 +20,22 @@ mod admin; mod control_plane; mod data_plane; +mod error; pub use admin::*; pub use control_plane::*; pub use data_plane::*; +pub use error::*; use borsh::{BorshDeserialize, BorshSerialize}; use crate::{ - control_plane::metadata::SyncConsumerGroupRequest, impl_from_variant, impl_from_variant_via, + control_plane::metadata::{EntryId, SyncConsumerGroupRequest}, + data_plane::{ + auxiliary_states::consumer_offsets::state::ConsumerOffsetPosition, + messages::query::{FetchResult, ListOffsetsResult}, + }, + impl_from_variant, impl_from_variant_via, }; // ── Top-level dispatch ───────────────────────────────────────────────────── @@ -42,12 +49,116 @@ pub enum ClientRequest { #[derive(Debug, BorshSerialize, BorshDeserialize)] pub enum ClientResponse { - ControlPlane(ControlPlaneResponse), - DataPlane(DataPlaneResponse), - Admin(AdminResponse), + Ok(ClientSuccess), + Err(ServerError), Stop, } +#[derive(Debug, Clone, PartialEq, Eq, BorshSerialize, BorshDeserialize)] +pub enum ClientSuccess { + // Control Plane + TopicCreated, + TopicDeleted, + TopicList { + topics: Box<[TopicSummary]>, + }, + TopicDetail(TopicDetail), + ConsumerGroupAssignment(ConsumerGroupAssignmentResponse), + ConsumerGroupLeft, + ProducerSessionOpened(ProducerSessionOpened), + + // Data Plane + Produced { + entry_id: EntryId, + }, + Fetched { + entries: Box<[EntryPayload]>, + next_entry_id: EntryId, + progress_signal: RangeProgressSignal, + }, + RangeOffset { + start_entry_id: EntryId, + next_entry_id: EntryId, + }, + ConsumerOffsetCommitted, + ConsumerOffset(Option), + + // Admin + ClusterInfo { + nodes: Box<[NodeInfo]>, + }, + TopicStats { + topics: Box<[TopicStats]>, + }, + ShardInfo { + detail: Option, + }, + ShardLeader { + leader: Option, + }, +} + +impl ClientSuccess { + // TODO refactor - remove ListOffsetsResult and return Result + pub(crate) fn from_list_offset_result(value: ListOffsetsResult) -> Result { + match value { + ListOffsetsResult::RangeOffsets { + start_entry_id, + next_entry_id, + } => Ok(ClientSuccess::RangeOffset { + start_entry_id, + next_entry_id, + }), + ListOffsetsResult::SegmentNotLocal => Err(ServerError::SegmentNotLocal), + } + } + + // TODO refactor - remove FetchResult and return Result + pub(crate) fn from_fetch_result(result: FetchResult) -> Result { + match result { + FetchResult::Records { + entries, + next_entry_id, + progress_signal, + } => { + let wire_entries = entries + .into_iter() + .map(|cached| EntryPayload { + entry_id: cached.entry_id, + record_count: cached.record_count, + data: cached.data.to_vec(), + }) + .collect::>() + .into_boxed_slice(); + + Ok(ClientSuccess::Fetched { + entries: wire_entries, + next_entry_id, + progress_signal, + }) + } + FetchResult::SegmentNotLocal => Err(ServerError::SegmentNotLocal), + FetchResult::EntryIdOutOfRange => Err(ServerError::EntryIdOutOfRange), + FetchResult::InternalError(s) => Err(ServerError::Internal(s)), + } + } +} + +impl From> for ClientResponse { + fn from(res: Result) -> Self { + match res { + Ok(ok) => ClientResponse::Ok(ok), + Err(err) => ClientResponse::Err(err), + } + } +} + +impl From for ClientResponse { + fn from(ok: ClientSuccess) -> Self { + ClientResponse::Ok(ok) + } +} + impl_from_variant!( ClientRequest, ControlPlane(ControlPlaneRequest), @@ -55,13 +166,6 @@ impl_from_variant!( Admin(AdminRequest), ); -impl_from_variant!( - ClientResponse, - ControlPlane(ControlPlaneResponse), - DataPlane(DataPlaneResponse), - Admin(AdminResponse), -); - impl_from_variant_via!( ClientRequest, ClientDataPlaneRequest, diff --git a/src/it/e2e/client_protocol.rs b/src/it/e2e/client_protocol.rs index 1c5ba041..1e26dbb4 100644 --- a/src/it/e2e/client_protocol.rs +++ b/src/it/e2e/client_protocol.rs @@ -5,9 +5,9 @@ use turmoil::Builder; use crate::StartUp; use crate::connections::protocol::{ - AdminRequest, AdminResponse, ClientDataPlaneRequest, ClientRequest, ClientResponse, - ControlPlaneRequest, ControlPlaneResponse, DataPlaneResponse, FetchByIdRequest, FetchRequest, - NodeState, ProduceRequest, RangeProgressSignal, ShardNotLocal, TopicDetail, + AdminRequest, ClientDataPlaneRequest, ClientRequest, ClientResponse, ClientSuccess, + ControlPlaneRequest, FetchByIdRequest, FetchRequest, NodeState, ProduceRequest, + RangeProgressSignal, ServerError, TopicDetail, }; use crate::control_plane::membership::ShardGroupId; use crate::control_plane::metadata::strategy::{PartitionStrategy, StoragePolicy}; @@ -62,11 +62,11 @@ fn create_topic_and_describe_cluster() -> turmoil::Result { tokio::time::sleep(Duration::from_secs(2)).await; for (host, port) in [("node-1", 8081u16), ("node-2", 8082), ("node-3", 8083)] { match send_request(host, port, req.clone()).await { - ClientResponse::ControlPlane(ControlPlaneResponse::TopicCreated) => { + ClientResponse::Ok(ClientSuccess::TopicCreated) => { created = true; break 'outer; } - ClientResponse::ControlPlane(ControlPlaneResponse::AlreadyExists) => { + ClientResponse::Err(ServerError::AlreadyExists) => { created = true; break 'outer; } @@ -85,7 +85,7 @@ fn create_topic_and_describe_cluster() -> turmoil::Result { ) .await; match resp { - ClientResponse::Admin(AdminResponse::ClusterInfo { nodes }) => { + ClientResponse::Ok(ClientSuccess::ClusterInfo { nodes }) => { assert!(!nodes.is_empty(), "{host} returned empty node list"); } other => panic!("{host}: unexpected DescribeCluster response: {other:?}"), @@ -128,8 +128,8 @@ fn delete_topic() -> turmoil::Result { tokio::time::sleep(Duration::from_secs(2)).await; for (host, port) in [("node-1", 8081u16), ("node-2", 8082), ("node-3", 8083)] { match send_request(host, port, create_req.clone()).await { - ClientResponse::ControlPlane(ControlPlaneResponse::TopicCreated) - | ClientResponse::ControlPlane(ControlPlaneResponse::AlreadyExists) => { + ClientResponse::Ok(ClientSuccess::TopicCreated) + | ClientResponse::Err(ServerError::AlreadyExists) => { acked = Some((host, port)); break 'outer; } @@ -146,8 +146,7 @@ fn delete_topic() -> turmoil::Result { ClientRequest::ControlPlane(ControlPlaneRequest::ListHostedTopics), ) .await; - let ClientResponse::ControlPlane(ControlPlaneResponse::TopicList { topics }) = list_resp - else { + let ClientResponse::Ok(ClientSuccess::TopicList { topics }) = list_resp else { panic!("expected TopicList, got {list_resp:?}"); }; assert!( @@ -165,10 +164,7 @@ fn delete_topic() -> turmoil::Result { ) .await; assert!( - matches!( - del_resp, - ClientResponse::ControlPlane(ControlPlaneResponse::TopicDeleted) - ), + matches!(del_resp, ClientResponse::Ok(ClientSuccess::TopicDeleted)), "expected TopicDeleted, got {del_resp:?}" ); @@ -177,14 +173,13 @@ fn delete_topic() -> turmoil::Result { tokio::time::sleep(Duration::from_millis(500)).await; let mut still_present = false; for (h, p) in [("node-1", 8081u16), ("node-2", 8082), ("node-3", 8083)] { - if let ClientResponse::ControlPlane(ControlPlaneResponse::TopicList { - topics: listed, - }) = send_request( - h, - p, - ClientRequest::ControlPlane(ControlPlaneRequest::ListHostedTopics), - ) - .await + if let ClientResponse::Ok(ClientSuccess::TopicList { topics: listed }) = + send_request( + h, + p, + ClientRequest::ControlPlane(ControlPlaneRequest::ListHostedTopics), + ) + .await && listed.iter().any(|t| t.name == "del-test") { still_present = true; @@ -196,9 +191,7 @@ fn delete_topic() -> turmoil::Result { } for (h, p) in [("node-1", 8081u16), ("node-2", 8082), ("node-3", 8083)] { - if let ClientResponse::ControlPlane(ControlPlaneResponse::TopicList { - topics: listed, - }) = send_request( + if let ClientResponse::Ok(ClientSuccess::TopicList { topics: listed }) = send_request( h, p, ClientRequest::ControlPlane(ControlPlaneRequest::ListHostedTopics), @@ -248,8 +241,8 @@ fn list_topic_stats_after_create() -> turmoil::Result { tokio::time::sleep(Duration::from_secs(2)).await; for (host, port) in [("node-1", 8081u16), ("node-2", 8082), ("node-3", 8083)] { match send_request(host, port, create_req.clone()).await { - ClientResponse::ControlPlane(ControlPlaneResponse::TopicCreated) - | ClientResponse::ControlPlane(ControlPlaneResponse::AlreadyExists) => { + ClientResponse::Ok(ClientSuccess::TopicCreated) + | ClientResponse::Err(ServerError::AlreadyExists) => { created = true; break 'outer; } @@ -262,7 +255,7 @@ fn list_topic_stats_after_create() -> turmoil::Result { // At least one node must report stats for "stats-test". let mut found = false; for (host, port) in [("node-1", 8081u16), ("node-2", 8082), ("node-3", 8083)] { - if let ClientResponse::Admin(AdminResponse::TopicStats { topics }) = send_request( + if let ClientResponse::Ok(ClientSuccess::TopicStats { topics }) = send_request( host, port, ClientRequest::Admin(AdminRequest::ListHostedTopicsWithStats), @@ -321,8 +314,8 @@ fn describe_topic_returns_topic_metadata() -> turmoil::Result { tokio::time::sleep(Duration::from_secs(2)).await; for (host, port) in [("node-1", 8081u16), ("node-2", 8082), ("node-3", 8083)] { match send_request(host, port, create_req.clone()).await { - ClientResponse::ControlPlane(ControlPlaneResponse::TopicCreated) - | ClientResponse::ControlPlane(ControlPlaneResponse::AlreadyExists) => { + ClientResponse::Ok(ClientSuccess::TopicCreated) + | ClientResponse::Err(ServerError::AlreadyExists) => { created = true; break 'outer; } @@ -342,13 +335,11 @@ fn describe_topic_returns_topic_metadata() -> turmoil::Result { for (host, port) in [("node-1", 8081u16), ("node-2", 8082), ("node-3", 8083)] { let resp = send_request(host, port, describe.clone()).await; match resp { - ClientResponse::ControlPlane(ControlPlaneResponse::TopicDetail(detail)) => { + ClientResponse::Ok(ClientSuccess::TopicDetail(detail)) => { found_detail = Some(detail); break; } - ClientResponse::ControlPlane(ControlPlaneResponse::TopicMetadataRedirect { - owner, - }) => { + ClientResponse::Err(ServerError::TopicMetadataRedirect { owner }) => { // Follow the redirect — owner.client_addr is the SWIM-resolved // address; for turmoil tests we use the well-known port map. let owner_port = match &*owner.node_id { @@ -359,9 +350,7 @@ fn describe_topic_returns_topic_metadata() -> turmoil::Result { }; let owner_host = &*owner.node_id; let resp2 = send_request(owner_host, owner_port, describe.clone()).await; - if let ClientResponse::ControlPlane(ControlPlaneResponse::TopicDetail(d)) = - resp2 - { + if let ClientResponse::Ok(ClientSuccess::TopicDetail(d)) = resp2 { found_detail = Some(d); break; } @@ -463,7 +452,7 @@ fn produce_then_fetch_hot() -> turmoil::Result { max_bytes: 1 << 20, keyspace_bound: None, })); - if let ClientResponse::DataPlane(DataPlaneResponse::Fetched { + if let ClientResponse::Ok(ClientSuccess::Fetched { entries: batch, next_entry_id: next, progress_signal: signal, @@ -837,7 +826,7 @@ fn sealed_segment_repair_catches_up_the_spare() -> turmoil::Result { entry_id: EntryId(0), max_bytes: 1 << 20, })); - if let ClientResponse::DataPlane(DataPlaneResponse::Fetched { entries, .. }) = + if let ClientResponse::Ok(ClientSuccess::Fetched { entries, .. }) = send_request(spare.0, spare.1, fetch).await && !entries.is_empty() && entries[0].data[..] == rec(0)[..] @@ -1020,7 +1009,7 @@ fn leader_crash_seals_active_segment_at_min_and_continues() -> turmoil::Result { entry_id: EntryId(0), max_bytes: 1 << 20, })); - if let ClientResponse::DataPlane(DataPlaneResponse::Fetched { entries, .. }) = + if let ClientResponse::Ok(ClientSuccess::Fetched { entries, .. }) = send_request(host, port, fetch).await && !entries.is_empty() && entries[0].data[..] == rec0[..] @@ -1215,7 +1204,7 @@ fn sealed_repair_survives_coordinator_crash() -> turmoil::Result { entry_id: EntryId(0), max_bytes: 1 << 20, })); - if let ClientResponse::DataPlane(DataPlaneResponse::Fetched { entries, .. }) = + if let ClientResponse::Ok(ClientSuccess::Fetched { entries, .. }) = send_request(spare.0, spare.1, fetch).await && !entries.is_empty() && entries[0].data[..] == rec(0)[..] @@ -1380,7 +1369,7 @@ fn catch_up_redrive_recovers_dropped_assignment() -> turmoil::Result { // partition didn't actually drop the assignment — the test would be vacuous.) tokio::time::sleep(Duration::from_secs(15)).await; for _ in 0..5 { - if let ClientResponse::DataPlane(DataPlaneResponse::Fetched { entries, .. }) = + if let ClientResponse::Ok(ClientSuccess::Fetched { entries, .. }) = send_request(spare.0, spare.1, fetch.clone()).await { assert!( @@ -1397,7 +1386,7 @@ fn catch_up_redrive_recovers_dropped_assignment() -> turmoil::Result { let mut served = false; for _ in 0..60 { tokio::time::sleep(Duration::from_secs(1)).await; - if let ClientResponse::DataPlane(DataPlaneResponse::Fetched { entries, .. }) = + if let ClientResponse::Ok(ClientSuccess::Fetched { entries, .. }) = send_request(spare.0, spare.1, fetch.clone()).await && !entries.is_empty() && entries[0].data[..] == rec(0)[..] @@ -1573,7 +1562,7 @@ fn restarted_node_reuses_recovered_segment_on_reassignment() -> turmoil::Result let mut served = false; for _ in 0..180 { tokio::time::sleep(Duration::from_secs(1)).await; - if let Some(ClientResponse::DataPlane(DataPlaneResponse::Fetched { entries, .. })) = + if let Some(ClientResponse::Ok(ClientSuccess::Fetched { entries, .. })) = try_send_request(VICTIM.0, VICTIM.1, fetch.clone()).await && !entries.is_empty() && entries[0].data[..] == rec(0)[..] @@ -1663,7 +1652,7 @@ async fn wait_for_cluster(nodes: &[(&str, u16)], expected: usize) { ) .await { - ClientResponse::Admin(AdminResponse::ClusterInfo { nodes: members }) => members + ClientResponse::Ok(ClientSuccess::ClusterInfo { nodes: members }) => members .iter() .filter(|m| matches!(m.state, NodeState::Alive)) .count(), @@ -1704,18 +1693,16 @@ async fn create_topic_anywhere(topic: &str, nodes: &[(&str, u16)], replication_f tokio::time::sleep(Duration::from_secs(2)).await; for &(host, port) in nodes { match send_request(host, port, req.clone()).await { - ClientResponse::ControlPlane(ControlPlaneResponse::TopicCreated) - | ClientResponse::ControlPlane(ControlPlaneResponse::AlreadyExists) => return, - ClientResponse::ControlPlane(ControlPlaneResponse::NotRaftLeader { + ClientResponse::Ok(ClientSuccess::TopicCreated) + | ClientResponse::Err(ServerError::AlreadyExists) => return, + ClientResponse::Err(ServerError::NotRaftLeader { leader_addr: Some(addr), }) => { if let Some((lh, lp)) = redirect_addr_to_node(addr.client_addr(), nodes) && matches!( send_request(lh, lp, req.clone()).await, - ClientResponse::ControlPlane( - ControlPlaneResponse::TopicCreated - | ControlPlaneResponse::AlreadyExists - ) + ClientResponse::Ok(ClientSuccess::TopicCreated) + | ClientResponse::Err(ServerError::AlreadyExists) ) { return; @@ -1812,17 +1799,19 @@ async fn produce_once(topic: &str, payload: &[u8], node: (&str, u16)) -> Produce ) .await { - ClientResponse::ControlPlane(ControlPlaneResponse::TopicDetail(d)) => match d - .ranges - .iter() - .filter(|r| { - r.state == RangeState::Active && r.keyspace_start.as_slice() <= b"k".as_slice() - }) - .max_by_key(|r| &r.keyspace_start) - { - Some(range) => range.range_id, - None => return ProduceOutcome::Retry, - }, + ClientResponse::Ok(ClientSuccess::TopicDetail(d)) => { + match d + .ranges + .iter() + .filter(|r| { + r.state == RangeState::Active && r.keyspace_start.as_slice() <= b"k".as_slice() + }) + .max_by_key(|r| &r.keyspace_start) + { + Some(range) => range.range_id, + None => return ProduceOutcome::Retry, + } + } _ => return ProduceOutcome::Retry, }; @@ -1835,13 +1824,13 @@ async fn produce_once(topic: &str, payload: &[u8], node: (&str, u16)) -> Produce producer_identity: None, })); match send_request(node.0, node.1, req).await { - ClientResponse::DataPlane(DataPlaneResponse::Produced { .. }) => ProduceOutcome::Acked, - ClientResponse::DataPlane(DataPlaneResponse::NotWriteLeader { + ClientResponse::Ok(ClientSuccess::Produced { .. }) => ProduceOutcome::Acked, + ClientResponse::Err(ServerError::NotWriteLeader { leader_addr: Some(addr), }) => ProduceOutcome::Redirect(addr.client_addr()), - ClientResponse::ControlPlane(ControlPlaneResponse::ShardNotLocal(ShardNotLocal { + ClientResponse::Err(ServerError::ShardNotLocal { hint_node: Some(addr), - })) => ProduceOutcome::Redirect(addr.client_addr()), + }) => ProduceOutcome::Redirect(addr.client_addr()), _ => ProduceOutcome::Retry, } } @@ -1852,11 +1841,11 @@ async fn describe_topic(topic: &str, nodes: &[(&str, u16)]) -> Option return Some(d), - ClientResponse::ControlPlane(ControlPlaneResponse::TopicMetadataRedirect { owner }) => { + ClientResponse::Ok(ClientSuccess::TopicDetail(d)) => return Some(d), + ClientResponse::Err(ServerError::TopicMetadataRedirect { owner }) => { if let Some(&(owner_host, owner_port)) = nodes.iter().find(|(n, _)| owner.node_id.starts_with(n)) - && let ClientResponse::ControlPlane(ControlPlaneResponse::TopicDetail(d)) = + && let ClientResponse::Ok(ClientSuccess::TopicDetail(d)) = send_request(owner_host, owner_port, req.clone()).await { return Some(d); diff --git a/src/it/e2e/client_sdk.rs b/src/it/e2e/client_sdk.rs index 90221a15..ed2e3725 100644 --- a/src/it/e2e/client_sdk.rs +++ b/src/it/e2e/client_sdk.rs @@ -13,9 +13,9 @@ use turmoil::Builder; use super::host_cluster; use crate::client::{ - BufferConfig, Client, ClientError, CompressionCodec, Consumer, ConsumerConfig, KeyInterest, - PartitionStrategy, Producer, ProducerConfig, RetryPolicy, StartPolicy, StoragePolicy, - TopicDetail, + BufferConfig, Client, ClientError, ClientSuccess, CompressionCodec, Consumer, ConsumerConfig, + KeyInterest, PartitionStrategy, Producer, ProducerConfig, RetryPolicy, StartPolicy, + StoragePolicy, TopicDetail, }; use crate::config::Environment; use crate::connections::protocol::ClientResponse; @@ -459,9 +459,7 @@ fn client_discovers_nodes_beyond_the_seed() -> turmoil::Result { #[test] #[serial_test::serial] fn producer_compression_lz4_end_to_end() -> turmoil::Result { - use crate::connections::protocol::{ - ClientDataPlaneRequest, DataPlaneResponse, FetchByIdRequest, - }; + use crate::connections::protocol::{ClientDataPlaneRequest, FetchByIdRequest}; let mut sim = build_sim(90); host_cluster(&mut sim, &NODES, sim_cluster); @@ -524,9 +522,7 @@ fn producer_compression_lz4_end_to_end() -> turmoil::Result { .call(leader_addr, ClientDataPlaneRequest::FetchById(fetch_req)) .await .expect("fetch"); - if let ClientResponse::DataPlane(DataPlaneResponse::Fetched { entries, .. }) = - served.response - { + if let ClientResponse::Ok(ClientSuccess::Fetched { entries, .. }) = served.response { assert_eq!(entries.len(), 1); let entry = &entries[0]; assert_eq!(entry.entry_id, entry_id); @@ -562,9 +558,7 @@ fn producer_compression_lz4_end_to_end() -> turmoil::Result { #[test] #[serial_test::serial] fn producer_compression_zstd_end_to_end() -> turmoil::Result { - use crate::connections::protocol::{ - ClientDataPlaneRequest, DataPlaneResponse, FetchByIdRequest, - }; + use crate::connections::protocol::{ClientDataPlaneRequest, FetchByIdRequest}; let mut sim = build_sim(90); host_cluster(&mut sim, &NODES, sim_cluster); @@ -621,9 +615,7 @@ fn producer_compression_zstd_end_to_end() -> turmoil::Result { .call(leader_addr, ClientDataPlaneRequest::FetchById(fetch_req)) .await .expect("fetch"); - if let ClientResponse::DataPlane(DataPlaneResponse::Fetched { entries, .. }) = - served.response - { + if let ClientResponse::Ok(ClientSuccess::Fetched { entries, .. }) = served.response { assert_eq!(entries.len(), 1); let entry = &entries[0]; assert_eq!(entry.entry_id, entry_id); @@ -1908,7 +1900,7 @@ async fn wait_for_retention_deletion(client: &Client, topic: &str, range_id: u64 #[test] #[serial_test::serial] fn producer_split_fence_retry() -> turmoil::Result { - let mut sim = build_sim(90); + let mut sim = build_sim(150); host_cluster(&mut sim, &NODES, |env| { sim_cluster(env); env.segment_size_limit_bytes = 10; @@ -1979,8 +1971,10 @@ fn producer_split_fence_retry() -> turmoil::Result { ); producer.flush().await; - assert!(left.await.expect("left sender completes").is_ok()); - assert!(right.await.expect("right sender completes").is_ok()); + let res_left = left.await.expect("left sender completes"); + let res_right = right.await.expect("right sender completes"); + assert!(res_left.is_ok(), "res_left failed: {:?}", res_left); + assert!(res_right.is_ok(), "res_right failed: {:?}", res_right); let refreshed = client1.resolve_topic(topic).await.expect("refresh routing"); let left_range = refreshed diff --git a/src/it/e2e/producer_deduplication.rs b/src/it/e2e/producer_deduplication.rs index fbf49119..dd3670bd 100644 --- a/src/it/e2e/producer_deduplication.rs +++ b/src/it/e2e/producer_deduplication.rs @@ -8,10 +8,10 @@ use turmoil::Builder; use uuid::Uuid; use super::{NODES, host_cluster}; -use crate::client::{Client, ClientError, PartitionStrategy, StoragePolicy}; +use crate::client::{Client, ClientError, PartitionStrategy, ServerError, StoragePolicy}; use crate::connections::protocol::{ - ClientDataPlaneRequest, ClientRequest, ClientResponse, DataPlaneResponse, - OpenProducerSessionRequest, ProduceRequest, + ClientDataPlaneRequest, ClientRequest, ClientResponse, OpenProducerSessionRequest, + ProduceRequest, }; use crate::data_plane::{ProduceError, ProducerAppendIdentity}; use crate::it::helpers::send_request; @@ -124,9 +124,7 @@ fn duplicate_unknown_and_fenced_sessions_are_end_to_end_visible() -> turmoil::Re for (host, port, _) in NODES { if matches!( send_request(host, port, unknown_wire.clone()).await, - ClientResponse::DataPlane(DataPlaneResponse::ProduceRejected( - ProduceError::SessionExpired - )) + ClientResponse::Err(ServerError::ProduceRejected(ProduceError::SessionExpired)) ) { rejected = true; break; diff --git a/src/it/helpers.rs b/src/it/helpers.rs index ad548629..9ad2cc3c 100644 --- a/src/it/helpers.rs +++ b/src/it/helpers.rs @@ -1,7 +1,6 @@ +use crate::client::ClientSuccess; use crate::config::Environment; -use crate::connections::protocol::{ - AdminRequest, AdminResponse, ClientRequest, ClientResponse, NodeState, -}; +use crate::connections::protocol::{AdminRequest, ClientRequest, ClientResponse, NodeState}; use crate::connections::reader::ClientStreamReader; use crate::connections::writer::ClientRawWriter; use crate::control_plane::{NodeAddress, SwimNode}; @@ -60,7 +59,7 @@ pub async fn get_members(host: &str, port: u16) -> turmoil::Result .await?; let (_, response): (_, ClientResponse) = reader.read_request().await?; let nodes = match response { - ClientResponse::Admin(AdminResponse::ClusterInfo { nodes }) => nodes, + ClientResponse::Ok(ClientSuccess::ClusterInfo { nodes }) => nodes, other => { return Err(std::io::Error::new( std::io::ErrorKind::InvalidData, diff --git a/src/it/sim/invariants.rs b/src/it/sim/invariants.rs index 702e7846..03752568 100644 --- a/src/it/sim/invariants.rs +++ b/src/it/sim/invariants.rs @@ -1,6 +1,6 @@ +use crate::client::ClientSuccess; use crate::connections::protocol::{ - AdminRequest, AdminResponse, ClientRequest, ClientResponse, ControlPlaneRequest, - ControlPlaneResponse, ShardDetail, TopicSummary, + AdminRequest, ClientRequest, ClientResponse, ControlPlaneRequest, ShardDetail, TopicSummary, }; use crate::connections::reader::ClientStreamReader; use crate::connections::writer::ClientRawWriter; @@ -224,7 +224,7 @@ pub(in crate::it) async fn query_shard_leader( let (_, response): (_, ClientResponse) = tokio::time::timeout(Duration::from_secs(3), reader.read_request()).await??; match response { - ClientResponse::Admin(AdminResponse::ShardLeader { leader }) => Ok(leader), + ClientResponse::Ok(ClientSuccess::ShardLeader { leader }) => Ok(leader), _ => Ok(None), } } @@ -248,7 +248,7 @@ pub(in crate::it) async fn query_shard_info( let (_, response): (_, ClientResponse) = tokio::time::timeout(Duration::from_secs(3), reader.read_request()).await??; match response { - ClientResponse::Admin(AdminResponse::ShardInfo { detail }) => Ok(detail), + ClientResponse::Ok(ClientSuccess::ShardInfo { detail }) => Ok(detail), _ => Ok(None), } } @@ -268,7 +268,7 @@ pub(super) async fn query_topics(host: &str, port: u16) -> turmoil::Result Ok(topics), + ClientResponse::Ok(ClientSuccess::TopicList { topics }) => Ok(topics), _ => Ok(Box::new([])), } } diff --git a/src/it/sim/properties.rs b/src/it/sim/properties.rs index 9cb1131d..afdffcb6 100644 --- a/src/it/sim/properties.rs +++ b/src/it/sim/properties.rs @@ -5,7 +5,8 @@ use std::time::Duration; use turmoil::Builder; use crate::StartUp; -use crate::connections::protocol::ControlPlaneResponse; +use crate::client::ClientSuccess; +use crate::connections::protocol::ClientResponse; use crate::control_plane::membership::ShardGroupId; use crate::it::helpers::{check_alive_count, check_dead_or_not_exist, default_env}; use crate::it::sim::invariants::{ @@ -113,7 +114,7 @@ fn metadata_visible() -> turmoil::Result { for _ in 0..30u32 { tokio::time::sleep(Duration::from_secs(1)).await; for i in 1..=3u8 { - if let Some(ControlPlaneResponse::TopicCreated) = + if let Some(ClientResponse::Ok(ClientSuccess::TopicCreated)) = try_propose(&node_name(i), client_port(i), &req).await { acked = true; @@ -208,7 +209,7 @@ fn leader_elects_after_kill() -> turmoil::Result { } if !topic_acked { for i in 1..=3u8 { - if let Some(ControlPlaneResponse::TopicCreated) = + if let Some(ClientResponse::Ok(ClientSuccess::TopicCreated)) = try_propose(&node_name(i), client_port(i), &req).await { topic_acked = true; diff --git a/src/it/sim/scenario.rs b/src/it/sim/scenario.rs index 9f227cb0..67e59671 100644 --- a/src/it/sim/scenario.rs +++ b/src/it/sim/scenario.rs @@ -8,12 +8,12 @@ use rand::{SeedableRng, rngs::StdRng, seq::IndexedRandom}; use serde::{Deserialize, Serialize}; use turmoil::Builder; -use crate::StartUp; -use crate::impl_from_variant; +use crate::{StartUp, client::ClientSuccess}; +use crate::{client::ServerError, impl_from_variant}; use crate::connections::protocol::{ - ClientDataPlaneRequest, ClientRequest, ClientResponse, ControlPlaneRequest, - ControlPlaneResponse, DataPlaneResponse, FetchByIdRequest, ProduceRequest, TopicDetail, + ClientDataPlaneRequest, ClientRequest, ClientResponse, ControlPlaneRequest, FetchByIdRequest, + ProduceRequest, TopicDetail, }; use crate::connections::reader::ClientStreamReader; use crate::connections::writer::ClientRawWriter; @@ -326,7 +326,7 @@ pub(super) async fn try_propose( host: &str, port: u16, req: &ClientRequest, -) -> Option { +) -> Option { let stream = tokio::time::timeout(Duration::from_secs(2), TcpStream::connect((host, port))) .await .ok() @@ -340,10 +340,7 @@ pub(super) async fn try_propose( .await .ok() .and_then(|r| r.ok())?; - match response { - ClientResponse::ControlPlane(cp) => Some(cp), - _ => None, - } + Some(response) } async fn send_request( @@ -374,7 +371,7 @@ async fn describe_topic_anywhere( loop { for (host, port) in nodes { match send_request(host, *port, &request).await { - Ok(ClientResponse::ControlPlane(ControlPlaneResponse::TopicDetail(detail))) => { + Ok(ClientResponse::Ok(ClientSuccess::TopicDetail(detail))) => { return detail; } Ok(_) => {} @@ -408,7 +405,7 @@ async fn produce_until_acked( loop { for (host, port) in nodes { match send_request(host, *port, &request).await { - Ok(ClientResponse::DataPlane(DataPlaneResponse::Produced { entry_id })) => { + Ok(ClientResponse::Ok(ClientSuccess::Produced { entry_id })) => { return entry_id; } Ok(_) => {} @@ -441,7 +438,7 @@ async fn assert_record_readable( loop { for (host, port) in nodes { match send_request(host, *port, &request).await { - Ok(ClientResponse::DataPlane(DataPlaneResponse::Fetched { entries, .. })) + Ok(ClientResponse::Ok(ClientSuccess::Fetched { entries, .. })) if entries .iter() .any(|entry| entry.entry_id == entry_id && entry.data == expected) => @@ -555,10 +552,8 @@ pub fn run_for_scenario(scenario: &SimScenario) -> turmoil::Result { try_propose(&node_name(node), client_port(node), &req).await; if matches!( response, - Some( - ControlPlaneResponse::TopicCreated - | ControlPlaneResponse::AlreadyExists - ) + Some(ClientResponse::Ok(ClientSuccess::TopicCreated)) + | Some(ClientResponse::Err(ServerError::AlreadyExists)) ) { break 'await_commit; } From e9be08336032a283d05d592fa5124161ae4abf70 Mon Sep 17 00:00:00 2001 From: Migorithm Date: Thu, 23 Jul 2026 10:41:45 +0400 Subject: [PATCH 25/41] unify with ServerError --- src/control_plane/membership/actor.rs | 43 +++++++++++++++++---------- src/data_plane/actor.rs | 14 ++++----- 2 files changed, 35 insertions(+), 22 deletions(-) diff --git a/src/control_plane/membership/actor.rs b/src/control_plane/membership/actor.rs index 0da40ee0..ff3fda0b 100644 --- a/src/control_plane/membership/actor.rs +++ b/src/control_plane/membership/actor.rs @@ -5,6 +5,7 @@ use std::sync::Arc; use arc_swap::ArcSwap; use crate::channels::BatchSender; +use crate::client::ServerError; use crate::control_plane::NodeAddress; use crate::control_plane::NodeAddressInfo; use crate::control_plane::NodeId; @@ -162,49 +163,61 @@ impl SwimSender { pub(crate) async fn resolve_shard_group( &self, resource_key: Vec, - ) -> anyhow::Result> { + ) -> Result, ServerError> { let (send, recv) = tokio::sync::oneshot::channel(); self.send(QueryCommand::ResolveShardGroup { key: resource_key, reply: send, }) - .await?; - Ok(recv.await?) + .await + .map_err(|err| ServerError::Internal(err.to_string()))?; + + recv.await + .map_err(|err| ServerError::Internal(err.to_string())) } pub(crate) async fn get_shard_info( &self, key: Vec, - ) -> anyhow::Result)>> { + ) -> Result)>, ServerError> { let (send, recv) = tokio::sync::oneshot::channel(); self.send(QueryCommand::GetShardInfo { key, reply: send }) - .await?; - Ok(recv.await?) + .await + .map_err(|err| ServerError::Internal(err.to_string()))?; + recv.await + .map_err(|err| ServerError::Internal(err.to_string())) } - pub(crate) async fn get_members(&self) -> anyhow::Result> { + pub(crate) async fn get_members(&self) -> Result, ServerError> { let (send, recv) = tokio::sync::oneshot::channel(); - self.send(QueryCommand::GetMembers { reply: send }).await?; - Ok(recv.await?) + self.send(QueryCommand::GetMembers { reply: send }) + .await + .map_err(|err| ServerError::Internal(err.to_string()))?; + + recv.await + .map_err(|err| ServerError::Internal(err.to_string())) } pub(crate) async fn resolve_address( &self, node_id: NodeId, - ) -> anyhow::Result> { + ) -> Result, ServerError> { let (send, recv) = tokio::sync::oneshot::channel(); self.send(QueryCommand::ResolveAddress { node_id, reply: send, }) - .await?; - Ok(recv.await?) + .await + .map_err(|err| ServerError::Internal(err.to_string()))?; + + recv.await + .map_err(|err| ServerError::Internal(err.to_string())) } pub(crate) async fn resolve_any( &self, members: &[NodeId], - ) -> anyhow::Result> { + ) -> Result, ServerError> { for member in members { if let Some(addr) = self.resolve_address(member.clone()).await? { return Ok(Some(NodeAddressInfo { @@ -222,7 +235,7 @@ impl SwimSender { &self, key: Vec, node_id: &NodeId, - ) -> anyhow::Result { + ) -> Result { let Some(group) = self.resolve_shard_group(key).await? else { return Ok(ShardRouting::Redirect(None)); }; @@ -235,7 +248,7 @@ impl SwimSender { pub(crate) async fn list_all_node_addresses( &self, - ) -> anyhow::Result> { + ) -> Result, ServerError> { Ok(self .get_members() .await? diff --git a/src/data_plane/actor.rs b/src/data_plane/actor.rs index 7439b247..d005e84f 100644 --- a/src/data_plane/actor.rs +++ b/src/data_plane/actor.rs @@ -7,6 +7,7 @@ use super::state::DataPlane; use super::timer::SegmentIdleTimer; use super::wal::WalWriter; use crate::channels::BatchSender; +use crate::client::ServerError; use crate::config::DataNodeConfig; use crate::control_plane::NodeId; use crate::control_plane::consensus::actor::MutlRaftSender; @@ -154,17 +155,16 @@ fn run_orphan_gc(internal: Duration, tx: &Sender) { pub(crate) struct DataPlaneSender(pub flume::Sender); impl DataPlaneSender { - pub fn send(&self, msg: impl Into) -> Result<(), flume::SendError<()>> { - self.0.send(msg.into()).map_err(|_| flume::SendError(())) + pub fn send(&self, msg: impl Into) -> Result<(), ServerError> { + self.0 + .send(msg.into()) + .map_err(|err| ServerError::Internal(format!("Channel dropped: {}", err))) } - pub async fn send_async( - &self, - msg: impl Into, - ) -> Result<(), flume::SendError<()>> { + pub async fn send_async(&self, msg: impl Into) -> Result<(), ServerError> { self.0 .send_async(msg.into()) .await - .map_err(|_| flume::SendError(())) + .map_err(|err| ServerError::Internal(format!("Channel dropped: {}", err))) } } From 2e1f00cb511f2145b07443da66b66758d65243be Mon Sep 17 00:00:00 2001 From: Migorithm Date: Thu, 23 Jul 2026 11:14:44 +0400 Subject: [PATCH 26/41] refact: remove redundant wrapper enums refactored FetchResult and ListOffsetsResult to leverage standard Result<..., ServerError> types. --- src/client/consumer/mod.rs | 5 +- src/client/consumer/range_fetcher.rs | 7 +-- src/client/consumer/topic_fetch_manager.rs | 8 +++- src/client/mod.rs | 8 ++-- src/connections/protocol/data_plane.rs | 22 +++++++++ src/connections/protocol/mod.rs | 53 +--------------------- src/data_plane/cold_read.rs | 11 +++-- src/data_plane/messages/command.rs | 8 ++-- src/data_plane/messages/query.rs | 42 ++++++----------- src/data_plane/mod.rs | 6 +-- src/data_plane/states/segment/cache.rs | 8 ++-- src/data_plane/states/segment/record.rs | 6 +-- src/data_plane/states/segment/tracker.rs | 8 ++-- 13 files changed, 77 insertions(+), 115 deletions(-) diff --git a/src/client/consumer/mod.rs b/src/client/consumer/mod.rs index c74c23ec..25a0f26c 100644 --- a/src/client/consumer/mod.rs +++ b/src/client/consumer/mod.rs @@ -15,6 +15,7 @@ use crate::connections::protocol::{ }; use crate::control_plane::metadata::{EntryId, RangeId, RangeState, TopicId}; use crate::data_plane::auxiliary_states::consumer_offsets::state::ConsumerOffsetPosition; +use crate::data_plane::messages::query::RangeOffsets; pub mod config; pub use config::*; @@ -85,10 +86,10 @@ impl Consumer { if matches!(config.start_policy, StartPolicy::Latest) { for cursor in cursors.iter_mut() { - let (_, tail_entry_id) = client + let RangeOffsets { next_entry_id, .. } = client .fetch_range_entry_ids(&topic, cursor.range_id) .await?; - cursor.next_entry_id = tail_entry_id; + cursor.next_entry_id = next_entry_id; } } diff --git a/src/client/consumer/range_fetcher.rs b/src/client/consumer/range_fetcher.rs index a54c6343..b2168415 100644 --- a/src/client/consumer/range_fetcher.rs +++ b/src/client/consumer/range_fetcher.rs @@ -11,6 +11,7 @@ use crate::connections::protocol::{ }; use crate::control_plane::metadata::{EntryId, RangeId}; use crate::data_plane::auxiliary_states::consumer_offsets::state::ConsumerOffsetPosition; +use crate::data_plane::messages::query::RangeOffsets; const EMPTY_FETCHES_BEFORE_REFRESH: u8 = 20; @@ -148,14 +149,14 @@ impl RangeFetchActor { } ClientResponse::Err(ServerError::EntryIdOutOfRange) => { let prev_entry_id = self.cursor.next_entry_id; - let (start_id, _) = self + let RangeOffsets { start_entry_id, .. } = self .ctx .client .fetch_range_entry_ids(&self.ctx.topic, self.cursor.range_id) .await?; - if self.cursor.next_entry_id < start_id { - self.cursor.next_entry_id = start_id; + if self.cursor.next_entry_id < start_entry_id { + self.cursor.next_entry_id = start_entry_id; } if self.cursor.next_entry_id == prev_entry_id { tokio::time::sleep(Duration::from_millis(50)).await; diff --git a/src/client/consumer/topic_fetch_manager.rs b/src/client/consumer/topic_fetch_manager.rs index 428520ee..3834f174 100644 --- a/src/client/consumer/topic_fetch_manager.rs +++ b/src/client/consumer/topic_fetch_manager.rs @@ -11,6 +11,7 @@ use crate::client::{ClientError, CommitMode, ConsumerConfig}; use crate::connections::protocol::{RangeDetail, RangeTransition, RebalancePlan}; use crate::control_plane::metadata::{EntryId, RangeId, RangeState}; use crate::data_plane::auxiliary_states::consumer_offsets::state::ConsumerOffsetPosition; +use crate::data_plane::messages::query::RangeOffsets; use crate::impl_from_variant; #[cfg(any(test, debug_assertions))] @@ -380,8 +381,11 @@ impl TopicFetchManagerState { ) -> Result { // 1. Fetch the latest boundary from the replica if the start policy is Latest. if self.config.start_policy == StartPolicy::Latest { - let (_, tail_entry_id) = ctx.client.fetch_range_entry_ids(&ctx.topic, range).await?; - return Ok(tail_entry_id); + let RangeOffsets { + start_entry_id: _, + next_entry_id, + } = ctx.client.fetch_range_entry_ids(&ctx.topic, range).await?; + return Ok(next_entry_id); } // 2. Fallback to the local cursor's next entry ID. diff --git a/src/client/mod.rs b/src/client/mod.rs index ea299516..9f3fe6b0 100644 --- a/src/client/mod.rs +++ b/src/client/mod.rs @@ -32,6 +32,7 @@ use crate::control_plane::metadata::{SyncConsumerGroupRequest, TopicId}; use crate::data_plane::auxiliary_states::consumer_offsets::state::{ ConsumerOffsetKey, ConsumerOffsetPosition, }; +use crate::data_plane::messages::query::RangeOffsets; pub use codec::CompressionCodec; pub use consumer::{ CommitMode, Consumer, ConsumerConfig, ConsumerRecord, DeliverySemantic, KeyInterest, @@ -121,7 +122,7 @@ impl Client { &self, topic: &str, range_id: RangeId, - ) -> Result<(EntryId, EntryId), ClientError> { + ) -> Result { let deadline = Instant::now() + self.retry.deadline; let mut backoff = self.retry.initial_backoff; let mut last_error = None; @@ -156,10 +157,7 @@ impl Client { let served = self.call(addr, req).await?; match served.response { - ClientResponse::Ok(ClientSuccess::RangeOffset { - start_entry_id, - next_entry_id, - }) => return Ok((start_entry_id, next_entry_id)), + ClientResponse::Ok(ClientSuccess::RangeOffset(res)) => return Ok(res), ClientResponse::Err(ServerError::SegmentNotLocal) => { last_error = Some("range offsets segment not local".to_string()); } diff --git a/src/connections/protocol/data_plane.rs b/src/connections/protocol/data_plane.rs index b075b274..5b5549c4 100644 --- a/src/connections/protocol/data_plane.rs +++ b/src/connections/protocol/data_plane.rs @@ -9,12 +9,14 @@ use borsh::{BorshDeserialize, BorshSerialize}; use crate::{ + client::ClientSuccess, control_plane::metadata::{ EntryId, RangeId, RangeMeta, RangeState, TopicId, TopicMeta, consumer_group::GenerationId, }, data_plane::{ ProducerAppendIdentity, auxiliary_states::consumer_offsets::state::{ConsumerOffsetKey, ConsumerOffsetUpdate}, + messages::query::FetchedRecords, }, impl_from_variant, impl_new_struct_wrapper, }; @@ -26,6 +28,26 @@ pub struct EntryPayload { pub data: Vec, } +impl EntryPayload { + pub fn from_fetched_record(r: FetchedRecords) -> ClientSuccess { + let wire_entries = r + .entries + .into_iter() + .map(|cached| EntryPayload { + entry_id: cached.entry_id, + record_count: cached.record_count, + data: cached.data.to_vec(), + }) + .collect(); + + ClientSuccess::Fetched { + entries: wire_entries, + next_entry_id: r.next_entry_id, + progress_signal: r.progress_signal, + } + } +} + /// Client → broker data-plane request. Every variant carries a `Client*Request` /// struct so the carried fields are named in one place (consistent with the /// project's tuple-variant + named-struct enum pattern) and the `Client` diff --git a/src/connections/protocol/mod.rs b/src/connections/protocol/mod.rs index 44a97ab9..611a2d1e 100644 --- a/src/connections/protocol/mod.rs +++ b/src/connections/protocol/mod.rs @@ -33,7 +33,7 @@ use crate::{ control_plane::metadata::{EntryId, SyncConsumerGroupRequest}, data_plane::{ auxiliary_states::consumer_offsets::state::ConsumerOffsetPosition, - messages::query::{FetchResult, ListOffsetsResult}, + messages::query::RangeOffsets, }, impl_from_variant, impl_from_variant_via, }; @@ -76,10 +76,7 @@ pub enum ClientSuccess { next_entry_id: EntryId, progress_signal: RangeProgressSignal, }, - RangeOffset { - start_entry_id: EntryId, - next_entry_id: EntryId, - }, + RangeOffset(RangeOffsets), ConsumerOffsetCommitted, ConsumerOffset(Option), @@ -98,52 +95,6 @@ pub enum ClientSuccess { }, } -impl ClientSuccess { - // TODO refactor - remove ListOffsetsResult and return Result - pub(crate) fn from_list_offset_result(value: ListOffsetsResult) -> Result { - match value { - ListOffsetsResult::RangeOffsets { - start_entry_id, - next_entry_id, - } => Ok(ClientSuccess::RangeOffset { - start_entry_id, - next_entry_id, - }), - ListOffsetsResult::SegmentNotLocal => Err(ServerError::SegmentNotLocal), - } - } - - // TODO refactor - remove FetchResult and return Result - pub(crate) fn from_fetch_result(result: FetchResult) -> Result { - match result { - FetchResult::Records { - entries, - next_entry_id, - progress_signal, - } => { - let wire_entries = entries - .into_iter() - .map(|cached| EntryPayload { - entry_id: cached.entry_id, - record_count: cached.record_count, - data: cached.data.to_vec(), - }) - .collect::>() - .into_boxed_slice(); - - Ok(ClientSuccess::Fetched { - entries: wire_entries, - next_entry_id, - progress_signal, - }) - } - FetchResult::SegmentNotLocal => Err(ServerError::SegmentNotLocal), - FetchResult::EntryIdOutOfRange => Err(ServerError::EntryIdOutOfRange), - FetchResult::InternalError(s) => Err(ServerError::Internal(s)), - } - } -} - impl From> for ClientResponse { fn from(res: Result) -> Self { match res { diff --git a/src/data_plane/cold_read.rs b/src/data_plane/cold_read.rs index 2d3f7d5a..5b80d14c 100644 --- a/src/data_plane/cold_read.rs +++ b/src/data_plane/cold_read.rs @@ -10,11 +10,12 @@ use super::actor::DataPlaneSender; use super::sparse_index::SparseIndex; use super::states::segment::cache::CachedEntry; use super::wal::{WalRecord, WalRecordType}; +use crate::client::ServerError; use crate::connections::protocol::RangeProgressSignal; use crate::control_plane::NodeId; use crate::control_plane::metadata::EntryId; use crate::data_plane::messages::command::CatchUpReadComplete; -use crate::data_plane::messages::query::FetchResult; +use crate::data_plane::messages::query::FetchedRecords; pub const DEFAULT_POOL_SIZE: usize = 4; @@ -43,7 +44,7 @@ pub(crate) struct ColdReadRequest { /// Where a finished cold read is delivered. pub(crate) enum ColdReadReply { Consumer { - reply: oneshot::Sender, + reply: oneshot::Sender>, progress_signal: RangeProgressSignal, }, CatchUp(CatchUpReadReply), @@ -137,11 +138,11 @@ impl ColdReadPool { reply, progress_signal, } => { - let _ = reply.send(FetchResult::Records { + let _ = reply.send(Ok(FetchedRecords { entries: records.entries, next_entry_id: records.next_offset, progress_signal, - }); + })); } ColdReadReply::CatchUp(cu) => { let done = CatchUpReadComplete { @@ -167,7 +168,7 @@ impl ColdReadPool { ColdReadReply::Consumer { reply, .. } => { tracing::warn!("cold read failed for {segment_key:?}: {e}"); let err_msg = format!("cold read failed: {e}"); - let _ = reply.send(FetchResult::InternalError(err_msg)); + let _ = reply.send(Err(ServerError::Internal(err_msg))); } ColdReadReply::CatchUp(_) => { // Source-side read failed; nothing durable changed. The diff --git a/src/data_plane/messages/command.rs b/src/data_plane/messages/command.rs index ef3561d9..52616168 100644 --- a/src/data_plane/messages/command.rs +++ b/src/data_plane/messages/command.rs @@ -15,7 +15,7 @@ use crate::{ control_plane::membership::ShardGroupId, control_plane::metadata::{EntryId, SegmentId}, data_plane::states::segment::cache::CachedEntry, - data_plane::{EntryPayload, SegmentKey, timer::DataPlaneTimeoutCallback}, + data_plane::{PayloadBytes, SegmentKey, timer::DataPlaneTimeoutCallback}, }; use borsh::{BorshDeserialize, BorshSerialize}; use std::borrow::Borrow; @@ -56,7 +56,7 @@ pub struct OrphanGcCheck { pub struct Produce { pub segment_key: SegmentKey, - pub data: EntryPayload, + pub data: PayloadBytes, pub record_count: u32, pub received_at_ms: u64, pub producer_identity: Option, @@ -120,7 +120,7 @@ pub struct SegmentPlaced { pub struct ReplicateSegmentEntries { pub segment_key: SegmentKey, pub replicas: Replicas, - pub data: EntryPayload, + pub data: PayloadBytes, pub record_count: u32, pub entry_id: EntryId, pub producer_identity: Option, @@ -241,7 +241,7 @@ pub struct CatchUpEntries { #[derive(Debug, Clone, BorshSerialize, BorshDeserialize)] pub struct CatchUpEntry { pub entry_id: EntryId, - pub data: EntryPayload, + pub data: PayloadBytes, pub record_count: u32, } impl CatchUpEntry { diff --git a/src/data_plane/messages/query.rs b/src/data_plane/messages/query.rs index c74b4d70..f4e672b4 100644 --- a/src/data_plane/messages/query.rs +++ b/src/data_plane/messages/query.rs @@ -7,6 +7,7 @@ use std::sync::Arc; +use borsh::{BorshDeserialize, BorshSerialize}; use tokio::sync::oneshot; use crate::connections::protocol::{ConsumerOffsetGenerationMismatch, RangeProgressSignal}; @@ -31,6 +32,8 @@ impl_from_variant!(DataPlaneQuery, Fetch, ListOffsets, ReadConsumerOffset); /// `topic_name → topic_id` and the range's current `progress_signal` via a /// metadata query before dispatching, so the data plane never reaches into the /// metadata RSM during its sync apply loop. +use crate::client::ServerError; + pub struct Fetch { pub topic_id: TopicId, pub range_id: RangeId, @@ -41,30 +44,14 @@ pub struct Fetch { /// in-band on the fetch that delivers the range's final records (see /// d4_consumer_range_tracking.md §Range Transitions). pub progress_signal: RangeProgressSignal, - pub reply: oneshot::Sender, + pub reply: oneshot::Sender>, } #[derive(Debug)] -pub enum FetchResult { - /// Records read from the segment (may be empty if the consumer is at the - /// tail). `next_entry_id` is what the consumer should request next. - /// - /// Carries `Arc`s straight from the tail cache — no - /// intermediate struct construction. The broker layer (wire translation) - /// does the single `Bytes → Vec` copy at serialization time; the - /// `lsn` field rides along through the channel but isn't on the wire. - Records { - entries: Vec>, - next_entry_id: EntryId, - progress_signal: RangeProgressSignal, - }, - /// `entry_id` was past the last committed offset on this node (e.g. caller - /// asked for an offset that hasn't been produced/committed yet). - EntryIdOutOfRange, - /// This node does not host any segment of the requested `(topic, range)`. - /// The consumer should re-resolve via `DescribeTopic` and retry. - SegmentNotLocal, - InternalError(String), +pub struct FetchedRecords { + pub(crate) entries: Vec>, + pub next_entry_id: EntryId, + pub progress_signal: RangeProgressSignal, } /// Offset-bounds query for a range. Returns the start and currently-committed @@ -72,16 +59,13 @@ pub enum FetchResult { pub struct ListOffsets { pub topic_id: TopicId, pub range_id: RangeId, - pub reply: oneshot::Sender, + pub reply: oneshot::Sender>, } -#[derive(Debug)] -pub enum ListOffsetsResult { - RangeOffsets { - start_entry_id: EntryId, - next_entry_id: EntryId, - }, - SegmentNotLocal, +#[derive(Debug, Clone, PartialEq, Eq, BorshSerialize, BorshDeserialize)] +pub struct RangeOffsets { + pub start_entry_id: EntryId, + pub next_entry_id: EntryId, } pub struct ReadConsumerOffset { diff --git a/src/data_plane/mod.rs b/src/data_plane/mod.rs index 88570193..20478da0 100644 --- a/src/data_plane/mod.rs +++ b/src/data_plane/mod.rs @@ -31,11 +31,11 @@ pub struct SegmentKey { pub use auxiliary_states::producer::{ProduceError, ProducerAppendIdentity}; #[derive(Debug, Clone, BorshSerialize, BorshDeserialize)] -pub struct EntryPayload(Bytes); +pub struct PayloadBytes(Bytes); -impl_new_struct_wrapper!(EntryPayload, Bytes); +impl_new_struct_wrapper!(PayloadBytes, Bytes); -impl From> for EntryPayload { +impl From> for PayloadBytes { fn from(v: Vec) -> Self { Self(Bytes::from(v)) } diff --git a/src/data_plane/states/segment/cache.rs b/src/data_plane/states/segment/cache.rs index b3c7ce25..96a76e88 100644 --- a/src/data_plane/states/segment/cache.rs +++ b/src/data_plane/states/segment/cache.rs @@ -7,11 +7,11 @@ use arc_swap::ArcSwapOption; use tokio::sync::Notify; use crate::control_plane::metadata::EntryId; -use crate::data_plane::{EntryPayload, ProducerAppendIdentity}; +use crate::data_plane::{PayloadBytes, ProducerAppendIdentity}; #[derive(Debug, Clone)] pub(crate) struct CachedEntry { - pub(crate) data: EntryPayload, + pub(crate) data: PayloadBytes, pub(crate) record_count: u32, pub(crate) entry_id: EntryId, pub(crate) lsn: u64, @@ -322,13 +322,13 @@ impl TAssertInvariant for SegmentRingBuffer { #[cfg(test)] mod tests { - use crate::data_plane::EntryPayload; + use crate::data_plane::PayloadBytes; use bytes::Bytes; use super::*; fn make_entry(entry_id: u64, lsn: u64) -> Arc { Arc::new(CachedEntry { - data: EntryPayload::from(Bytes::from("test")), + data: PayloadBytes::from(Bytes::from("test")), record_count: 1, entry_id: EntryId(entry_id), lsn, diff --git a/src/data_plane/states/segment/record.rs b/src/data_plane/states/segment/record.rs index 79cecf78..23656e29 100644 --- a/src/data_plane/states/segment/record.rs +++ b/src/data_plane/states/segment/record.rs @@ -3,7 +3,7 @@ use std::io; use bytes::{BufMut, Bytes, BytesMut}; use crate::control_plane::metadata::{EntryId, RangeId, SegmentId, TopicId}; -use crate::data_plane::{EntryPayload, ProducerAppendIdentity, SegmentKey}; +use crate::data_plane::{PayloadBytes, ProducerAppendIdentity, SegmentKey}; const ROUTING_HEADER_MAGIC: u32 = 0x4547_5032; const ROUTING_HEADER_SIZE: usize = 81; @@ -157,7 +157,7 @@ impl RoutingHeader { } pub(crate) struct StagedEntry { - pub(crate) data: EntryPayload, + pub(crate) data: PayloadBytes, pub(crate) record_count: u32, pub(crate) segment_key: SegmentKey, pub(crate) producer_identity: Option, @@ -165,7 +165,7 @@ pub(crate) struct StagedEntry { impl StagedEntry { pub(crate) fn new( - data: EntryPayload, + data: PayloadBytes, record_count: u32, segment_key: SegmentKey, producer_identity: Option, diff --git a/src/data_plane/states/segment/tracker.rs b/src/data_plane/states/segment/tracker.rs index df68015b..13e6ce36 100644 --- a/src/data_plane/states/segment/tracker.rs +++ b/src/data_plane/states/segment/tracker.rs @@ -4,7 +4,7 @@ use crate::control_plane::membership::ShardGroupId; use crate::control_plane::metadata::EntryId; use crate::control_plane::{NodeId, Replicas}; use crate::data_plane::ProducerAppendIdentity; -use crate::data_plane::{EntryPayload, SegmentKey, checkpoint::CheckpointJob, wal::WalRecord}; +use crate::data_plane::{PayloadBytes, SegmentKey, checkpoint::CheckpointJob, wal::WalRecord}; use super::cache::CachedEntry; @@ -146,7 +146,7 @@ impl SegmentTracker { pub(crate) fn uncommitted_entries( &self, - ) -> impl Iterator)> + '_ { + ) -> impl Iterator)> + '_ { let commit = self.cache.load_read_cursor(); let tail = self.cache.load_write_cursor(); (commit..tail).filter_map(|pos| { @@ -167,7 +167,7 @@ impl SegmentTracker { /// re-count them. pub(crate) fn staged_for_replay( &self, - ) -> impl Iterator)> + '_ { + ) -> impl Iterator)> + '_ { self.staged_entries .iter() .map(|s| (s.data.clone(), s.record_count, s.producer_identity)) @@ -275,7 +275,7 @@ impl SegmentTracker { pub(crate) fn stage_producer_entry( &mut self, segment_key: SegmentKey, - data: EntryPayload, + data: PayloadBytes, record_count: u32, entry_id: EntryId, producer_append_id: Option, From f332628ddeccbd96b802daad7ac0b8c9e37df2b5 Mon Sep 17 00:00:00 2001 From: Migorithm Date: Thu, 23 Jul 2026 12:11:57 +0400 Subject: [PATCH 27/41] replication factors in producer indentity --- src/data_plane/cold_read.rs | 2 +- src/data_plane/states/replication.rs | 64 +++++++++++++++++++++--- src/data_plane/states/segment/cache.rs | 4 +- src/data_plane/states/segment/tracker.rs | 4 +- 4 files changed, 61 insertions(+), 13 deletions(-) diff --git a/src/data_plane/cold_read.rs b/src/data_plane/cold_read.rs index 5b80d14c..5ec7335f 100644 --- a/src/data_plane/cold_read.rs +++ b/src/data_plane/cold_read.rs @@ -253,7 +253,7 @@ impl ColdReadPool { // lsn is a live-WAL concept; cold reads serve durable // data, so it carries no meaningful LSN. lsn: 0, - producer_append_id: None, + producer_identity: None, })); if bytes_read >= req.max_bytes { diff --git a/src/data_plane/states/replication.rs b/src/data_plane/states/replication.rs index 9087fdbe..5cf7eae3 100644 --- a/src/data_plane/states/replication.rs +++ b/src/data_plane/states/replication.rs @@ -8,6 +8,7 @@ use crate::control_plane::{NodeId, Replicas}; use crate::data_plane::SegmentKey; use crate::data_plane::messages::command::{ProduceAck, ReplicateSegmentEntries}; use crate::data_plane::states::segment::cache::CachedEntry; +use crate::data_plane::{ProduceError, ProducerAppendIdentity}; #[derive(Default)] pub(crate) struct ReplicationState { @@ -21,12 +22,16 @@ pub(crate) struct PendingBatch { pub(crate) replies: Vec>, pub(crate) pending_acks: HashSet, pub(crate) entry_id: EntryId, + pub(crate) producer_append_id: Option, + pub(crate) lsn: u64, } pub(crate) struct AckCommitted { pub entry_id: EntryId, pub replies: Vec>, pub reset_timer_seq: Option, + pub producer_append_id: Option, + pub lsn: u64, } pub(crate) struct PendingReplicationBatch { @@ -45,6 +50,7 @@ impl PendingReplicationBatch { data: self.entry.data.clone(), record_count: self.entry.record_count, entry_id: self.entry.entry_id, + producer_identity: self.entry.producer_identity, }; (targets, message) } @@ -93,13 +99,14 @@ impl ReplicationState { } pub(crate) fn fail_all(&mut self, err: String) { + self.timer_seqs.clear(); for reply in self.drain_all_pending_replies() { - let _ = reply.send(ProduceAck::Err(err.clone())); + let _ = reply.send(ProduceAck::Err(ProduceError::Internal(err.clone()))); } for batch in self.drain_all_in_flight() { for reply in batch.replies { - let _ = reply.send(ProduceAck::Err(err.clone())); + let _ = reply.send(ProduceAck::Err(ProduceError::Internal(err.clone()))); } } } @@ -119,6 +126,8 @@ impl ReplicationState { replies, pending_acks, entry_id: pending_repl.entry.entry_id, + producer_append_id: pending_repl.entry.producer_identity, + lsn: pending_repl.entry.lsn, }, ); if is_first { @@ -144,25 +153,32 @@ impl ReplicationState { ) -> Option { let deque = self.in_flight.get_mut(segment_key)?; let front = deque.front_mut()?; + front.pending_acks.remove(from); if !front.pending_acks.is_empty() { return None; } - let batch = deque.pop_front().unwrap(); + // All ACKs collected: commit the batch + let batch = deque.pop_front()?; self.timer_seqs.remove(segment_key); - let reset_timer_seq = (!deque.is_empty()).then_some({ + let reset_timer_seq = if deque.is_empty() { + self.in_flight.remove(segment_key); + None + } else { let seq = self.alloc_timer_seq(); self.set_timer_seq(*segment_key, seq); - seq - }); + Some(seq) + }; Some(AckCommitted { entry_id: batch.entry_id, replies: batch.replies, reset_timer_seq, + producer_append_id: batch.producer_append_id, + lsn: batch.lsn, }) } @@ -191,6 +207,7 @@ impl ReplicationState { } pub(crate) fn segment_handoff(&mut self, old: SegmentKey, new: SegmentKey) { + self.timer_seqs.remove(&old); if let Some(replies) = self.pending_replies.remove(&old) { self.pending_replies.entry(new).or_default().extend(replies); } @@ -205,6 +222,34 @@ impl ReplicationState { } } +#[cfg(any(test, debug_assertions))] +impl crate::test_traits::TAssertInvariant for ReplicationState { + fn assert_invariants(&self) { + for (key, replies) in &self.pending_replies { + assert!(!replies.is_empty(), "empty pending reply queue for {key:?}"); + } + for (key, batches) in &self.in_flight { + assert!(!batches.is_empty(), "empty replication queue for {key:?}"); + assert!( + self.timer_seqs.contains_key(key), + "in-flight replication has no active timer for {key:?}" + ); + for batch in batches { + assert!( + !batch.pending_acks.is_empty(), + "in-flight replication batch has no pending acknowledgments" + ); + } + } + for key in self.timer_seqs.keys() { + assert!( + self.in_flight.contains_key(key), + "replication timer exists without an in-flight batch for {key:?}" + ); + } + } +} + #[cfg(test)] mod tests { use super::*; @@ -213,7 +258,7 @@ mod tests { use bytes::Bytes; use crate::control_plane::metadata::{RangeId, SegmentId, TopicId}; - use crate::data_plane::EntryPayload; + use crate::data_plane::PayloadBytes; use crate::data_plane::states::segment::cache::CachedEntry; fn key(seg: u64) -> SegmentKey { @@ -226,10 +271,11 @@ mod tests { fn cached_entry(entry_id: EntryId) -> Arc { Arc::new(CachedEntry { - data: EntryPayload::from(Bytes::from("data")), + data: PayloadBytes::from(Bytes::from("data")), record_count: 1, entry_id, lsn: 1, + producer_identity: None, }) } @@ -327,6 +373,8 @@ mod tests { assert!(committed.reset_timer_seq.is_none()); assert!(!state.is_active_timer(&sk, 0)); + assert!(!state.in_flight.contains_key(&sk)); + crate::test_traits::TAssertInvariant::assert_invariants(&state); let _ = committed .replies diff --git a/src/data_plane/states/segment/cache.rs b/src/data_plane/states/segment/cache.rs index 96a76e88..62a0069f 100644 --- a/src/data_plane/states/segment/cache.rs +++ b/src/data_plane/states/segment/cache.rs @@ -15,7 +15,7 @@ pub(crate) struct CachedEntry { pub(crate) record_count: u32, pub(crate) entry_id: EntryId, pub(crate) lsn: u64, - pub(crate) producer_append_id: Option, + pub(crate) producer_identity: Option, } impl CachedEntry { @@ -332,7 +332,7 @@ mod tests { record_count: 1, entry_id: EntryId(entry_id), lsn, - producer_append_id: None, + producer_identity: None, }) } diff --git a/src/data_plane/states/segment/tracker.rs b/src/data_plane/states/segment/tracker.rs index 13e6ce36..999b9d15 100644 --- a/src/data_plane/states/segment/tracker.rs +++ b/src/data_plane/states/segment/tracker.rs @@ -136,7 +136,7 @@ impl SegmentTracker { record_count: s.record_count, entry_id, lsn, - producer_append_id: s.producer_identity, + producer_identity: s.producer_identity, }); self.cache.publish(entry.clone()); entry @@ -154,7 +154,7 @@ impl SegmentTracker { ( entry.data.clone(), entry.record_count, - entry.producer_append_id, + entry.producer_identity, ) }) }) From 98ea803d4280122d500d6f912e95c144297db6af Mon Sep 17 00:00:00 2001 From: Migorithm Date: Thu, 23 Jul 2026 12:32:35 +0400 Subject: [PATCH 28/41] RecoveryOutput return auxiliary state --- .../auxiliary_states/consumer_offsets/mod.rs | 3 ++- .../consumer_offsets/replication.rs | 3 ++- .../consumer_offsets/types.rs | 4 ++++ src/data_plane/auxiliary_states/manager.rs | 4 ++-- src/data_plane/messages/command.rs | 1 + src/data_plane/recovery/inventory.rs | 19 +++++++++------ src/data_plane/recovery/mod.rs | 24 +++++++------------ 7 files changed, 32 insertions(+), 26 deletions(-) diff --git a/src/data_plane/auxiliary_states/consumer_offsets/mod.rs b/src/data_plane/auxiliary_states/consumer_offsets/mod.rs index da9cb059..ad3ed680 100644 --- a/src/data_plane/auxiliary_states/consumer_offsets/mod.rs +++ b/src/data_plane/auxiliary_states/consumer_offsets/mod.rs @@ -17,7 +17,7 @@ use std::collections::HashSet; use tokio::sync::oneshot; use types::*; -#[derive(Default)] +#[derive(Debug, Default)] pub(crate) struct ConsumerOffsetCoordination { pub(crate) placements: HashMap<(TopicId, RangeId), OffsetPlacement>, pending_offset_mutations: Vec, @@ -137,6 +137,7 @@ impl ConsumerOffsetCoordination { } } +#[derive(Debug)] pub(crate) struct OffsetPlacement { pub(crate) segment_key: SegmentKey, pub(crate) shard_group_id: ShardGroupId, diff --git a/src/data_plane/auxiliary_states/consumer_offsets/replication.rs b/src/data_plane/auxiliary_states/consumer_offsets/replication.rs index 5ea0b90d..c5b456ee 100644 --- a/src/data_plane/auxiliary_states/consumer_offsets/replication.rs +++ b/src/data_plane/auxiliary_states/consumer_offsets/replication.rs @@ -6,12 +6,13 @@ use crate::data_plane::messages::command::{ use std::collections::{HashMap, HashSet}; use tokio::sync::oneshot; -#[derive(Default)] +#[derive(Debug, Default)] pub(crate) struct OffsetReplicationState { seq: u64, pending: HashMap, } +#[derive(Debug)] struct PendingOffsetReplication { update: ConsumerOffsetUpdate, pending_acks: HashSet, diff --git a/src/data_plane/auxiliary_states/consumer_offsets/types.rs b/src/data_plane/auxiliary_states/consumer_offsets/types.rs index d7dd332c..112bf5f9 100644 --- a/src/data_plane/auxiliary_states/consumer_offsets/types.rs +++ b/src/data_plane/auxiliary_states/consumer_offsets/types.rs @@ -11,6 +11,7 @@ use crate::impl_from_variant; use borsh::{BorshDeserialize, BorshSerialize}; use tokio::sync::oneshot; +#[derive(Debug)] pub(crate) struct PendingOffsetMutation { pub(crate) record: OffsetRecord, pub(crate) completion: OffsetMutationCompletion, @@ -41,12 +42,14 @@ impl From for PendingOffsetMutation { } } +#[derive(Debug)] pub(crate) struct LeaderOffsetCommitApplied { pub(crate) replica_set: Replicas, pub(crate) required_followers: HashSet, pub(crate) reply: oneshot::Sender, } +#[derive(Debug)] pub(crate) enum OffsetMutationCompletion { EpochSeal, LeaderCommit(LeaderOffsetCommitApplied), @@ -69,6 +72,7 @@ pub struct ReplicateConsumerOffset { pub update: ConsumerOffsetUpdate, } +#[derive(Debug)] pub(crate) enum FutureOffsetCommit { Client(CommitConsumerOffset), Replica(ReplicateConsumerOffset), diff --git a/src/data_plane/auxiliary_states/manager.rs b/src/data_plane/auxiliary_states/manager.rs index 20ad1558..d9f6b1d1 100644 --- a/src/data_plane/auxiliary_states/manager.rs +++ b/src/data_plane/auxiliary_states/manager.rs @@ -31,7 +31,7 @@ use std::path::PathBuf; use super::consumer_offsets::types::*; -#[derive(Default)] +#[derive(Debug, Default)] pub(crate) struct AuxiliaryStateManager { offsets: ConsumerOffsets, consumer_coord: ConsumerOffsetCoordination, // live placements, pending mutations, parked commits, and replication tracking @@ -39,7 +39,7 @@ pub(crate) struct AuxiliaryStateManager { checkpoint: AuxiliaryCheckpointState, // WAL reclamation and checkpoint progress } -#[derive(Default)] +#[derive(Debug, Default)] struct AuxiliaryCheckpointState { uncheckpointed_lsns: VecDeque, checkpoint_lsn: u64, diff --git a/src/data_plane/messages/command.rs b/src/data_plane/messages/command.rs index 52616168..8277f1e8 100644 --- a/src/data_plane/messages/command.rs +++ b/src/data_plane/messages/command.rs @@ -88,6 +88,7 @@ impl AuthorizedProducerIdentity { } } +#[derive(Debug)] pub struct CommitConsumerOffset { pub update: ConsumerOffsetUpdate, pub reply: oneshot::Sender, diff --git a/src/data_plane/recovery/inventory.rs b/src/data_plane/recovery/inventory.rs index d319d4c4..6c1c2861 100644 --- a/src/data_plane/recovery/inventory.rs +++ b/src/data_plane/recovery/inventory.rs @@ -4,8 +4,7 @@ use std::path::PathBuf; use super::segment_scan::RecoveredSegments; use crate::control_plane::metadata::EntryId; use crate::data_plane::SegmentKey; -use crate::data_plane::auxiliary_states::consumer_offsets::state::ConsumerOffsets; -use crate::data_plane::auxiliary_states::producer::ProducerTracker; +use crate::data_plane::auxiliary_states::AuxiliaryStateManager; /// What recovery verified for one on-disk segment: its `start_offset` (the filename base, /// needed to locate the file) and `verified_end` (the highest entry id a CRC-complete batch @@ -71,14 +70,18 @@ impl LocalInventory { pub(crate) struct RecoveryOutput { pub(crate) inventory: LocalInventory, pub(crate) data_dir: PathBuf, - pub(crate) offsets: ConsumerOffsets, - pub(crate) producer: ProducerTracker, + pub(crate) auxiliary_state: AuxiliaryStateManager, } #[cfg(test)] mod tests { use super::*; - use crate::control_plane::metadata::{RangeId, SegmentId, TopicId}; + use crate::{ + control_plane::metadata::{RangeId, SegmentId, TopicId}, + data_plane::auxiliary_states::{ + consumer_offsets::state::ConsumerOffsets, producer::ProducerTracker, + }, + }; fn seg(topic: u64, segment: u64) -> SegmentKey { SegmentKey::new(TopicId(topic), RangeId(0), SegmentId(segment)) @@ -128,8 +131,10 @@ mod tests { let output = RecoveryOutput { inventory: inv, data_dir: PathBuf::from("/var/lib/eastguard/data"), - offsets: ConsumerOffsets::default(), - producer: ProducerTracker::default(), + auxiliary_state: AuxiliaryStateManager::from_recovery( + ConsumerOffsets::default(), + ProducerTracker::default(), + ), }; assert_eq!(output.inventory.orphan_candidates().count(), 0); assert_eq!(output.data_dir, PathBuf::from("/var/lib/eastguard/data")); diff --git a/src/data_plane/recovery/mod.rs b/src/data_plane/recovery/mod.rs index 2b81940f..72cd6029 100644 --- a/src/data_plane/recovery/mod.rs +++ b/src/data_plane/recovery/mod.rs @@ -17,6 +17,7 @@ use self::inventory::{LocalInventory, RecoveryOutput}; use self::replay::ReplayWriter; use self::segment_scan::RecoveredSegments; use self::wal_scan::{ScanError, WalScanner}; +use crate::data_plane::auxiliary_states::AuxiliaryStateManager; use crate::data_plane::auxiliary_states::consumer_offsets::state::{ConsumerOffsets, OffsetRecord}; use crate::data_plane::auxiliary_states::producer::ProducerTracker; use crate::data_plane::auxiliary_states::snapshot::AuxiliarySnapshot; @@ -176,8 +177,7 @@ impl Durable { Ok(RecoveryOutput { inventory, data_dir: self.data_dir, - offsets: self.offsets, - producer: self.producer, + auxiliary_state: AuxiliaryStateManager::from_recovery(self.offsets, self.producer), }) } } @@ -338,12 +338,12 @@ mod tests { let (_db_dir, db) = open_db(); let output = run(data_dir.clone(), &db).unwrap(); - assert_eq!(*output.offsets.generation(&key), 4); - assert_eq!(output.offsets.offset(&key), Some(position)); + assert_eq!(*output.auxiliary_state.durable_generation(&key), 4); + assert_eq!(output.auxiliary_state.offset(&key), Some(position)); let restarted = run(data_dir, &db).unwrap(); - assert_eq!(*restarted.offsets.generation(&key), 4); - assert_eq!(restarted.offsets.offset(&key), Some(position)); + assert_eq!(*restarted.auxiliary_state.durable_generation(&key), 4); + assert_eq!(restarted.auxiliary_state.offset(&key), Some(position)); } #[test] @@ -366,21 +366,15 @@ mod tests { let (_db_dir, db) = open_db(); let output = run(data_dir.clone(), &db).unwrap(); - let mut ledger = crate::data_plane::auxiliary_states::AuxiliaryStateManager::from_recovery( - output.offsets, - output.producer, - ); + let mut ledger = output.auxiliary_state; assert_eq!( ledger.verify_producer(key, AuthorizedProducerIdentity::ExistingOnly(producer), 1,), Err(ProduceError::Duplicate(EntryId(7))) ); let restarted = run(data_dir, &db).unwrap(); - let mut restarted_tracker = - crate::data_plane::auxiliary_states::AuxiliaryStateManager::from_recovery( - restarted.offsets, - restarted.producer, - ); + let mut restarted_tracker = restarted.auxiliary_state; + assert_eq!( restarted_tracker.verify_producer( key, From 45e807ed5b606597401dc4447a0f88409ae90e56 Mon Sep 17 00:00:00 2001 From: Migorithm Date: Thu, 23 Jul 2026 13:15:32 +0400 Subject: [PATCH 29/41] feat(data_plane): add ProducerTracker and ProducerSession state machine Introduce ProducerTracker to manage producer session verification, sequence tracking, and append deduplication: - Track session incarnation, expiration, and sequence frontiers in ProducerSession - Enforce strict sequence continuity and detect sequence gaps or identity conflicts - Maintain a sliding deduplication window (ProducerDeduplicationEntry) to handle idempotent retries - Track in-flight appends to prevent concurrent duplicate processing --- .../auxiliary_states/producer/state.rs | 387 ++++++++++++++++++ 1 file changed, 387 insertions(+) create mode 100644 src/data_plane/auxiliary_states/producer/state.rs diff --git a/src/data_plane/auxiliary_states/producer/state.rs b/src/data_plane/auxiliary_states/producer/state.rs new file mode 100644 index 00000000..c41b45f0 --- /dev/null +++ b/src/data_plane/auxiliary_states/producer/state.rs @@ -0,0 +1,387 @@ +use std::collections::{HashMap, HashSet, VecDeque}; + +use borsh::{BorshDeserialize, BorshSerialize}; +use uuid::Uuid; + +use super::{AuthorizedProducerIdentity, ProduceError, ProducerAppendIdentity}; +use crate::control_plane::metadata::{EntryId, RangeId, TopicId}; +use crate::data_plane::SegmentKey; +use crate::data_plane::auxiliary_states::producer::types::AppendKey; +use crate::impl_new_struct_wrapper; + +const RECENT_RESULT_LIMIT: usize = 16; + +pub type ProducerKey = (TopicId, RangeId, Uuid); + +#[derive(Debug, Clone, BorshSerialize, BorshDeserialize)] +struct ProducerDeduplicationEntry { + sequence: u64, + digest: u32, + entry_id: EntryId, +} +impl ProducerDeduplicationEntry { + fn is_exact_retry(&self, digest: u32) -> bool { + self.digest == digest + } +} + +#[derive(Debug, Clone, BorshSerialize, BorshDeserialize)] +pub(crate) struct ProducerSession { + incarnation: u32, + expires_at: u64, + frontier: Option, + dedup_entries: VecDeque, +} + +impl ProducerSession { + pub(crate) fn new(incarnation: u32, expires_at: u64) -> Self { + Self { + incarnation, + expires_at, + frontier: None, + dedup_entries: VecDeque::new(), + } + } + + /// Verify an incoming append request against this session's incarnation and sequence history. + pub(crate) fn verify( + &mut self, + identity: &ProducerAppendIdentity, + allow_session_creation: bool, + ) -> Result<(), ProduceError> { + if identity.incarnation < self.incarnation { + return Err(ProduceError::ProducerFenced); + } + if identity.incarnation > self.incarnation { + if !allow_session_creation { + return Err(ProduceError::SessionExpired); + } + // * frontier and dedup_entries must be reset rather than reserved because: + // * dedup_entries stores entries indexed by sequence number and if old dedup were kept, sending seq=0 in new incarnation could match previous one, + // * leading to cross session contamination + *self = Self::new(identity.incarnation, identity.expires_at); + } + + let Some(frontier) = self.frontier else { + // ! If self.frontier is none, no append has been made so the first message in any session must start at 0 + if identity.sequence != 0 { + return Err(ProduceError::SequenceGap(0)); + } + return Ok(()); + }; + + if identity.sequence <= frontier { + // * Idempotent retry case + let Some(entry) = self + .dedup_entries + .iter() + .find(|e| e.sequence == identity.sequence) + else { + return Err(ProduceError::DuplicatePositionUnavailable); + }; + + if entry.is_exact_retry(identity.digest) { + return Err(ProduceError::Duplicate(entry.entry_id)); + } + + // Conflict or payload reuse case + return Err(ProduceError::RequestIdentityConflict); + } + + if identity.sequence != frontier + 1 { + return Err(ProduceError::SequenceGap(frontier + 1)); + } + Ok(()) + } + + /// Record a committed entry for a sequence number and advance the session frontier. + pub(crate) fn advance(&mut self, p_id: &ProducerAppendIdentity, entry_id: EntryId) { + if p_id.incarnation < self.incarnation + || self.frontier.is_some_and(|seq| p_id.sequence <= seq) + { + return; + } + self.incarnation = p_id.incarnation; + self.expires_at = p_id.expires_at; + self.frontier = Some(p_id.sequence); + self.dedup_entries.push_back(ProducerDeduplicationEntry { + sequence: p_id.sequence, + digest: p_id.digest, + entry_id, + }); + while self.dedup_entries.len() > RECENT_RESULT_LIMIT { + self.dedup_entries.pop_front(); + } + } +} + +/// Crash-durable producer end state. In-flight appends are intentionally not +/// part of this snapshot image. +#[derive(Debug, Clone, Default, BorshSerialize, BorshDeserialize)] +pub(crate) struct ProducerSessions(HashMap); +impl_new_struct_wrapper!(ProducerSessions,HashMap); + +/// Transient producer append coordination. This state is deliberately absent +/// from snapshots and is rebuilt empty after recovery. +#[derive(Debug, Clone, Default)] +pub(crate) struct ProducerTracker { + sessions: ProducerSessions, + in_flight_appends: HashSet, + last_cleanup_ms: u64, +} + +impl ProducerTracker { + pub(crate) fn from_sessions(sessions: ProducerSessions) -> Self { + Self { + sessions, + ..Self::default() + } + } + + pub(crate) fn sessions(&self) -> ProducerSessions { + self.sessions.clone() + } + + pub(crate) fn verify( + &mut self, + segment: SegmentKey, + identity: AuthorizedProducerIdentity, + received_at_ms: u64, + ) -> Result<(), ProduceError> { + let (producer_identity, allow_session_creation) = identity.destruct(); + + if producer_identity.expires_at < received_at_ms { + return Err(ProduceError::SessionExpired); + } + + self.cleanup_expired_sessions(received_at_ms); + let append_key = AppendKey::new(segment, &producer_identity); + let producer_key = append_key.producer_key(); + + if !allow_session_creation && !self.sessions.contains_key(&producer_key) { + return Err(ProduceError::SessionExpired); + } + + let session = self.sessions.entry(producer_key).or_insert_with(|| { + ProducerSession::new(producer_identity.incarnation, producer_identity.expires_at) + }); + + session.verify(&producer_identity, allow_session_creation)?; + + if !self.in_flight_appends.insert(append_key) { + return Err(ProduceError::RequestInFlight); + } + Ok(()) + } + + pub(crate) fn advance( + &mut self, + segment: SegmentKey, + p_id: ProducerAppendIdentity, + entry_id: EntryId, + ) { + let append_key = AppendKey::new(segment, &p_id); + self.unstage(&append_key); + + let session = self + .sessions + .entry((segment.topic_id, segment.range_id, p_id.producer_id)) + .or_insert_with(|| ProducerSession::new(p_id.incarnation, p_id.expires_at)); + session.advance(&p_id, entry_id); + } + + pub(crate) fn unstage(&mut self, append_key: &AppendKey) { + self.in_flight_appends.remove(append_key); + } + + fn cleanup_expired_sessions(&mut self, now: u64) { + const CLEANUP_INTERVAL_MS: u64 = 10_000; + + if now.saturating_sub(self.last_cleanup_ms) >= CLEANUP_INTERVAL_MS { + self.sessions.retain(|_, state| state.expires_at >= now); + self.last_cleanup_ms = now; + } + } + + #[cfg(any(test, debug_assertions))] + pub(crate) fn assert_invariants(&self) { + for write_key in &self.in_flight_appends { + if let Some(state) = self.sessions.get(&write_key.producer_key()) { + assert!( + write_key.incarnation >= state.incarnation, + "a fenced producer incarnation remains staged" + ); + } + } + } +} + +#[cfg(any(test, debug_assertions))] +impl crate::test_traits::TAssertInvariant for ProducerTracker { + fn assert_invariants(&self) { + for state in self.sessions.values() { + assert!( + state.dedup_entries.len() <= RECENT_RESULT_LIMIT, + "producer result window exceeds its fixed bound" + ); + let mut previous = None; + for result in &state.dedup_entries { + assert!( + state + .frontier + .is_some_and(|frontier| result.sequence <= frontier), + "producer result is ahead of its frontier" + ); + assert!( + previous.is_none_or(|sequence| sequence < result.sequence), + "producer result window is not strictly ordered" + ); + previous = Some(result.sequence); + } + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::control_plane::metadata::SegmentId; + + fn verify( + coordination: &mut ProducerTracker, + identity: AuthorizedProducerIdentity, + ) -> Result<(), ProduceError> { + coordination.verify(key(), identity, 1) + } + + fn key() -> SegmentKey { + SegmentKey::new(TopicId(1), RangeId(2), SegmentId(3)) + } + + fn request( + producer_id: Uuid, + incarnation: u32, + sequence: u64, + digest: u32, + ) -> ProducerAppendIdentity { + ProducerAppendIdentity { + producer_id, + incarnation, + expires_at: u64::MAX, + sequence, + digest, + } + } + + fn verified(req: ProducerAppendIdentity) -> AuthorizedProducerIdentity { + AuthorizedProducerIdentity::MetadataVerified(req) + } + + fn existing(req: ProducerAppendIdentity) -> AuthorizedProducerIdentity { + AuthorizedProducerIdentity::ExistingOnly(req) + } + + #[test] + fn committed_retry_returns_original_position_without_advancing() { + let id = Uuid::new_v4(); + let mut coordination = ProducerTracker::default(); + let first = request(id, 0, 0, 7); + + assert_eq!(verify(&mut coordination, verified(first)), Ok(())); + coordination.advance(key(), first, EntryId(11)); + assert_eq!( + verify(&mut coordination, existing(first)), + Err(ProduceError::Duplicate(EntryId(11))) + ); + assert_eq!( + verify(&mut coordination, existing(request(id, 0, 1, 8)),), + Ok(()) + ); + } + + #[test] + fn conflict_gap_in_flight_and_fencing_are_explicit() { + let id = Uuid::new_v4(); + + let mut coordination = ProducerTracker::default(); + let first = request(id, 2, 0, 7); + assert_eq!(verify(&mut coordination, verified(first)), Ok(())); + assert_eq!( + verify(&mut coordination, existing(first)), + Err(ProduceError::RequestInFlight) + ); + coordination.advance(key(), first, EntryId(4)); + + assert_eq!( + verify(&mut coordination, existing(request(id, 2, 0, 9))), + Err(ProduceError::RequestIdentityConflict) + ); + assert_eq!( + verify(&mut coordination, existing(request(id, 2, 2, 9))), + Err(ProduceError::SequenceGap(1)) + ); + assert_eq!( + verify(&mut coordination, verified(request(id, 3, 0, 9))), + Ok(()) + ); + assert_eq!( + verify(&mut coordination, existing(request(id, 2, 1, 9))), + Err(ProduceError::ProducerFenced) + ); + } + + #[test] + fn result_window_is_bounded_and_frontier_still_rejects_old_retries() { + let id = Uuid::new_v4(); + + let mut coordination = ProducerTracker::default(); + for sequence in 0..=RECENT_RESULT_LIMIT as u64 { + let item = request(id, 0, sequence, sequence as u32); + assert_eq!(verify(&mut coordination, verified(item)), Ok(())); + coordination.advance(key(), item, EntryId(sequence)); + } + assert_eq!( + verify(&mut coordination, existing(request(id, 0, 0, 0))), + Err(ProduceError::DuplicatePositionUnavailable) + ); + } + + #[test] + fn unknown_session_cannot_create_frontier_without_metadata_authority() { + let id = Uuid::new_v4(); + + let mut coordination = ProducerTracker::default(); + assert_eq!( + verify(&mut coordination, existing(request(id, 0, 0, 7))), + Err(ProduceError::SessionExpired) + ); + assert_eq!( + verify(&mut coordination, verified(request(id, 0, 0, 7))), + Ok(()) + ); + } + + #[test] + fn expired_and_aborted_requests_have_distinct_retry_behavior() { + let id = Uuid::new_v4(); + + let mut coordination = ProducerTracker::default(); + let mut expired = request(id, 0, 0, 7); + expired.expires_at = 0; + assert_eq!( + verify(&mut coordination, verified(expired)), + Err(ProduceError::SessionExpired) + ); + + let live = request(id, 0, 0, 7); + assert_eq!(verify(&mut coordination, verified(live)), Ok(())); + + let append_key = AppendKey::new(key(), &live); + coordination.unstage(&append_key); + assert_eq!( + verify(&mut coordination, existing(live)), + Ok(()), + "an aborted append must not remain permanently in flight" + ); + } +} From 6a0ac030daa5c60342c765a433d95bce5b2f3523 Mon Sep 17 00:00:00 2001 From: Migorithm Date: Thu, 23 Jul 2026 14:07:34 +0400 Subject: [PATCH 30/41] ProduceAck wraps EntryId for Ok case --- src/data_plane/messages/command.rs | 4 +--- src/data_plane/messages/pending.rs | 2 +- src/data_plane/states/replication.rs | 8 ++------ 3 files changed, 4 insertions(+), 10 deletions(-) diff --git a/src/data_plane/messages/command.rs b/src/data_plane/messages/command.rs index 8277f1e8..dadb8252 100644 --- a/src/data_plane/messages/command.rs +++ b/src/data_plane/messages/command.rs @@ -366,9 +366,7 @@ pub enum ProduceAck { /// `entry_id` is the committed offset for this produce. Exact for a producer /// with one request in flight at a time (the common case); under pipelining /// it is the segment's committed highwater at ack time (an upper bound). - Ok { - entry_id: EntryId, - }, + Ok(EntryId), Err(ProduceError), } diff --git a/src/data_plane/messages/pending.rs b/src/data_plane/messages/pending.rs index 98f3ba7f..2c331a00 100644 --- a/src/data_plane/messages/pending.rs +++ b/src/data_plane/messages/pending.rs @@ -58,7 +58,7 @@ impl DataPlaneOutputs { let _ = self.coordinator_tx.send(cmd).await; } for (entry_id, reply) in self.produce_replies.drain(..) { - let _ = reply.send(ProduceAck::Ok { entry_id }); + let _ = reply.send(ProduceAck::Ok(entry_id)); } let checkpoint_tasks = std::mem::take(&mut self.checkpoint_tasks); diff --git a/src/data_plane/states/replication.rs b/src/data_plane/states/replication.rs index 5cf7eae3..346683f9 100644 --- a/src/data_plane/states/replication.rs +++ b/src/data_plane/states/replication.rs @@ -381,14 +381,10 @@ mod tests { .into_iter() .next() .unwrap() - .send(ProduceAck::Ok { - entry_id: EntryId(42), - }); + .send(ProduceAck::Ok(EntryId(42))); assert!(matches!( rx.blocking_recv().unwrap(), - ProduceAck::Ok { - entry_id: EntryId(42) - } + ProduceAck::Ok(EntryId(42)) )); } From 0997cf519a37b1e15336de611b5c7763f21e1783 Mon Sep 17 00:00:00 2001 From: Migorithm Date: Thu, 23 Jul 2026 14:44:12 +0400 Subject: [PATCH 31/41] extend segment tracker --- src/data_plane/states/segment_store.rs | 27 ++++++++++++++++++++++++-- 1 file changed, 25 insertions(+), 2 deletions(-) diff --git a/src/data_plane/states/segment_store.rs b/src/data_plane/states/segment_store.rs index 778c2788..f43a5e40 100644 --- a/src/data_plane/states/segment_store.rs +++ b/src/data_plane/states/segment_store.rs @@ -21,11 +21,11 @@ //! never originates a side effect, never talks to the network, never emits //! events. `DataPlane` keeps doing all of that. -use super::segment::tracker::SegmentTracker; +use super::segment::tracker::{SegmentRole, SegmentTracker}; use std::collections::{BTreeMap, HashMap}; use crate::control_plane::metadata::{EntryId, RangeId, SegmentId, TopicId}; -use crate::data_plane::SegmentKey; +use crate::data_plane::{PayloadBytes, ProduceError, ProducerAppendIdentity, SegmentKey}; /// Where a segment's data currently lives once its tracker has been dropped. /// Returned from the resolver as part of a `Sealed` outcome so the cold-read @@ -81,6 +81,29 @@ impl SegmentStore { self.by_key.get_mut(key) } + /// Validates that the target segment exists and this node is the write leader. + pub(crate) fn verify_write_leader(&self, key: &SegmentKey) -> Result<(), ProduceError> { + match self.get(key) { + Some(t) if t.role() == SegmentRole::Follower => Err(ProduceError::NotLeader), + None => Err(ProduceError::SegmentNotFound), + _ => Ok(()), + } + } + + /// Stages a producer entry into the segment tracker, returning the assigned entry_id. + pub(crate) fn stage_producer_entry( + &mut self, + key: SegmentKey, + data: PayloadBytes, + record_count: u32, + producer_identity: Option, + ) -> Option { + let tracker = self.get_mut(&key)?; + let entry_id = tracker.next_staged_entry_id(); + tracker.stage_producer_entry(key, data, record_count, entry_id, producer_identity); + Some(entry_id) + } + pub(crate) fn values(&self) -> impl Iterator { self.by_key.values() } From 0f529baf019e9038c7ef9a3ff1b54477cf185ba2 Mon Sep 17 00:00:00 2001 From: Migorithm Date: Thu, 23 Jul 2026 15:04:05 +0400 Subject: [PATCH 32/41] feat(data_plane): integrate producer deduplication state into AuxiliaryStateManager - Integrate ProducerTracker into AuxiliaryStateManager alongside consumer offsets - Refactor DataPlane::handle_produce to delegate segment validation and staging to SegmentStore - Rename SegmentTracker::stage_producer_entry to stage_entry - Standardize producer_identity across ReplicationState and commit callbacks --- .../auxiliary_states/consumer_offsets/mod.rs | 4 +- .../consumer_offsets/replication.rs | 2 +- src/data_plane/auxiliary_states/manager.rs | 8 +- src/data_plane/state.rs | 427 +++++++++++++----- src/data_plane/states/replication.rs | 4 +- src/data_plane/states/segment/tracker.rs | 16 +- src/data_plane/states/segment_store.rs | 2 +- 7 files changed, 331 insertions(+), 132 deletions(-) diff --git a/src/data_plane/auxiliary_states/consumer_offsets/mod.rs b/src/data_plane/auxiliary_states/consumer_offsets/mod.rs index ad3ed680..1226a8f5 100644 --- a/src/data_plane/auxiliary_states/consumer_offsets/mod.rs +++ b/src/data_plane/auxiliary_states/consumer_offsets/mod.rs @@ -91,8 +91,8 @@ impl ConsumerOffsetCoordination { .begin(followers, required, update, reply) } - pub(crate) fn fail_all(&mut self, error: &str) { - self.offset_replication.fail_all(error); + pub(crate) fn drop_all(&mut self, error: &str) { + self.offset_replication.drop_all(error); for pending in self.take_pending() { if let OffsetMutationCompletion::LeaderCommit(commit) = pending.completion { diff --git a/src/data_plane/auxiliary_states/consumer_offsets/replication.rs b/src/data_plane/auxiliary_states/consumer_offsets/replication.rs index c5b456ee..784ab79c 100644 --- a/src/data_plane/auxiliary_states/consumer_offsets/replication.rs +++ b/src/data_plane/auxiliary_states/consumer_offsets/replication.rs @@ -105,7 +105,7 @@ impl OffsetReplicationState { .any(|pending| pending.pending_acks.contains(node_id)) } - pub(crate) fn fail_all(&mut self, error: &str) { + pub(crate) fn drop_all(&mut self, error: &str) { for (_, mut pending) in self.pending.drain() { if let Some(reply) = pending.reply.take() { let _ = reply.send(ConsumerOffsetCommitAck::InternalError(error.to_string())); diff --git a/src/data_plane/auxiliary_states/manager.rs b/src/data_plane/auxiliary_states/manager.rs index d9f6b1d1..3696446c 100644 --- a/src/data_plane/auxiliary_states/manager.rs +++ b/src/data_plane/auxiliary_states/manager.rs @@ -533,7 +533,7 @@ impl AuxiliaryStateManager { self.mark_offset_replica_ready_if_caught_up(&from) } - pub(crate) fn handle_consumer_group_epoch_seal( + pub(crate) fn unpark_future_offset_commits( &mut self, cmd: EpochSeal, ) -> Option> { @@ -561,8 +561,8 @@ impl AuxiliaryStateManager { self.consumer_coord.take_pending() } - pub(crate) fn fail_inflight_coordination(&mut self, error: &str) { - self.consumer_coord.fail_all(error); + pub(crate) fn drop_coordination(&mut self, error: &str) { + self.consumer_coord.drop_all(error); } pub(crate) fn flush_batch( @@ -1053,7 +1053,7 @@ mod tests { replication_reply, ); - manager.fail_inflight_coordination("WAL failed"); + manager.drop_coordination("WAL failed"); let expected = ConsumerOffsetCommitAck::InternalError("WAL failed".into()); assert_eq!(pending_response.try_recv().unwrap(), expected); diff --git a/src/data_plane/state.rs b/src/data_plane/state.rs index c3989d46..a3fb7ce9 100644 --- a/src/data_plane/state.rs +++ b/src/data_plane/state.rs @@ -6,12 +6,13 @@ use super::messages::command::DataPlanePeerMessage; use super::messages::command::*; use super::messages::pending::DataPlaneOutputs; use super::messages::query::{ - DataPlaneQuery, Fetch, FetchResult, ListOffsets, ListOffsetsResult, ReadConsumerOffsetResult, + DataPlaneQuery, Fetch, FetchedRecords, ListOffsets, RangeOffsets, ReadConsumerOffsetResult, }; use super::recovery::inventory::{LocalInventory, RecoveryOutput}; use super::recovery::orphan::OrphanCandidate; use super::recovery::segment_scan::scan_segment_file; use super::segment_writer::SegmentAppender; +use crate::client::ServerError; use super::states::replication::PendingReplicationBatch; use super::states::replication::ReplicationState; @@ -28,11 +29,12 @@ use crate::control_plane::membership::ShardGroupId; use crate::control_plane::metadata::EntryId; use crate::control_plane::metadata::command::SegmentRollIntent; use crate::control_plane::{NodeId, Replicas}; -use crate::data_plane::EntryPayload; -use crate::data_plane::consumer_offset_management::ConsumerOffsetManager; -use crate::data_plane::consumer_offset_management::ledger::{EpochSeal, StaleEpoch}; +use crate::data_plane::auxiliary_states::AuxiliaryStateManager; +use crate::data_plane::auxiliary_states::consumer_offsets::state::{EpochSeal, StaleEpoch}; +use crate::data_plane::auxiliary_states::producer::types::AppendKey; +use crate::data_plane::{PayloadBytes, ProduceError, ProducerAppendIdentity}; -use crate::data_plane::consumer_offset_management::types::*; +use crate::data_plane::auxiliary_states::consumer_offsets::types::*; use crate::data_plane::messages::query::ReadConsumerOffset; use crate::data_plane::states::segment::tracker::{SegmentRole, SegmentTracker}; use crate::data_plane::states::segment_store::SegmentReadState; @@ -92,7 +94,7 @@ pub struct DataPlane { recovered: LocalInventory, pending_boundary_corrections: std::collections::HashSet, pressure_checkpoints_in_flight: BTreeSet, - consumer_offsets: ConsumerOffsetManager, + auxiliary_state: AuxiliaryStateManager, } impl DataPlane { @@ -108,7 +110,7 @@ impl DataPlane { let RecoveryOutput { inventory: recovered, data_dir: _, - offsets: offset_ledger, + auxiliary_state, } = recovery; DataPlane { node_id, @@ -127,7 +129,7 @@ impl DataPlane { pending_roll_requests: PendingSegmentRollRequests::default(), pending_boundary_corrections: std::collections::HashSet::new(), pressure_checkpoints_in_flight: BTreeSet::new(), - consumer_offsets: ConsumerOffsetManager::new(offset_ledger), + auxiliary_state, } } @@ -145,8 +147,8 @@ impl DataPlane { DataPlaneCommand::SegmentCheckpointComplete(complete) => { self.handle_checkpoint_complete(complete); } - DataPlaneCommand::OffsetCheckpointComplete(complete) => { - self.handle_offset_checkpoint_complete(complete); + DataPlaneCommand::AuxiliaryCheckpointComplete(complete) => { + self.handle_auxiliary_checkpoint_complete(complete); } DataPlaneCommand::DataPlaneTimeoutCallback(callback) => { self.handle_timeout(callback); @@ -185,26 +187,48 @@ impl DataPlane { } fn handle_produce(&mut self, cmd: Produce) { - let reject = match self.segments.get(&cmd.segment_key) { - Some(t) if t.role() == SegmentRole::Follower => Some("not leader"), - None => Some("segment not found"), - _ => None, - }; - if let Some(reason) = reject { - let _ = cmd.reply.send(ProduceAck::Err(reason.into())); + if let Some(identity) = cmd.producer_identity + && let Err(err) = + self.auxiliary_state + .verify_producer(cmd.segment_key, identity, cmd.received_at_ms) + { + let ack = match err { + ProduceError::Duplicate(entry_id) => ProduceAck::Ok(entry_id), + err => ProduceAck::Err(err), + }; + let _ = cmd.reply.send(ack); return; } - self.replication.enqueue_reply(cmd.segment_key, cmd.reply); + let producer_identity = cmd + .producer_identity + .map(AuthorizedProducerIdentity::request); - let Some(tracker) = self.segments.get_mut(&cmd.segment_key) else { + if let Err(reason) = self.segments.verify_write_leader(&cmd.segment_key) { + if let Some(p_id) = producer_identity { + let append_key = AppendKey::new(cmd.segment_key, &p_id); + self.auxiliary_state.unstage_producer(&append_key); + } + let _ = cmd.reply.send(ProduceAck::Err(reason)); return; - }; + } + + self.replication.enqueue_reply(cmd.segment_key, cmd.reply); self.buffer_byte_count += cmd.data.len(); - let entry_id = tracker.next_staged_entry_id(); - tracker.stage_entry(cmd.segment_key, cmd.data, cmd.record_count, entry_id); - self.dirty_segments.insert(cmd.segment_key); + + if self + .segments + .stage_producer_entry( + cmd.segment_key, + cmd.data, + cmd.record_count, + producer_identity, + ) + .is_some() + { + self.dirty_segments.insert(cmd.segment_key); + } self.out .store_batch_produce_timer(TimerCommand::SetSchedule { @@ -226,7 +250,7 @@ impl DataPlane { .segments .resolve(cmd.topic_id, cmd.range_id, cmd.entry_id) else { - let _ = cmd.reply.send(FetchResult::SegmentNotLocal); + let _ = cmd.reply.send(Err(ServerError::SegmentNotLocal)); return; }; @@ -244,7 +268,7 @@ impl DataPlane { fn handle_hot_fetch(&self, cmd: Fetch, key: SegmentKey) { let Some(tracker) = self.segments.get(&key) else { // Resolver and store invariants should keep these in sync(just a guard) - let _ = cmd.reply.send(FetchResult::SegmentNotLocal); + let _ = cmd.reply.send(Err(ServerError::SegmentNotLocal)); return; }; @@ -260,11 +284,11 @@ impl DataPlane { // checking against the cache read cursor. if *cache_position >= read_cursor { // Past the tail — nothing to return. - let _ = cmd.reply.send(FetchResult::Records { + let _ = cmd.reply.send(Ok(FetchedRecords { entries: Vec::new(), next_entry_id: cmd.entry_id, progress_signal: cmd.progress_signal, - }); + })); return; } @@ -290,11 +314,11 @@ impl DataPlane { } } - let _ = cmd.reply.send(FetchResult::Records { + let _ = cmd.reply.send(Ok(FetchedRecords { entries, next_entry_id, progress_signal: cmd.progress_signal, - }); + })); } /// Hand a sealed-segment read off to the cold-read pool. The pool owns the @@ -307,7 +331,7 @@ impl DataPlane { end_entry_id: EntryId, ) { if cmd.entry_id > end_entry_id { - let _ = cmd.reply.send(FetchResult::EntryIdOutOfRange); + let _ = cmd.reply.send(Err(ServerError::EntryIdOutOfRange)); return; } @@ -325,9 +349,9 @@ impl DataPlane { if let Err(flume::SendError(req)) = self.cold_read_handoff_sender.send(req) && let ColdReadReply::Consumer { reply, .. } = req.reply { - let _ = reply.send(FetchResult::InternalError( + let _ = reply.send(Err(ServerError::Internal( "cold-read pool unavailable".into(), - )); + ))); } } @@ -342,29 +366,29 @@ impl DataPlane { self.segments .resolve(cmd.topic_id, cmd.range_id, EntryId(u64::MAX)) else { - let _ = cmd.reply.send(ListOffsetsResult::SegmentNotLocal); + let _ = cmd.reply.send(Err(ServerError::SegmentNotLocal)); return; }; let Some(tracker) = self.segments.get(&key) else { - let _ = cmd.reply.send(ListOffsetsResult::SegmentNotLocal); + let _ = cmd.reply.send(Err(ServerError::SegmentNotLocal)); return; }; - let _ = cmd.reply.send(ListOffsetsResult::RangeOffsets { + let _ = cmd.reply.send(Ok(RangeOffsets { start_entry_id: tracker.start_entry_id(), next_entry_id: tracker.successor_start_entry_id(), - }); + })); } fn read_consumer_offset(&self, query: ReadConsumerOffset) { let result = match self - .consumer_offsets + .auxiliary_state .get_replica_set_if_leader(&query.key, &self.node_id) { Ok(_) => { - let observed_generation = self.consumer_offsets.durable_generation(&query.key); + let observed_generation = self.auxiliary_state.durable_generation(&query.key); if observed_generation == query.generation { - ReadConsumerOffsetResult::Offset(self.consumer_offsets.offset(&query.key)) + ReadConsumerOffsetResult::Offset(self.auxiliary_state.offset(&query.key)) } else { ReadConsumerOffsetResult::GenerationMismatch(ConsumerOffsetGenerationMismatch { observed_generation, @@ -389,9 +413,9 @@ impl DataPlane { self.reclaim_checkpointed_wal(); } - fn handle_offset_checkpoint_complete(&mut self, complete: OffsetCheckpointComplete) { - self.consumer_offsets - .handle_offset_checkpoint_complete(complete.checkpointed_lsn); + fn handle_auxiliary_checkpoint_complete(&mut self, complete: AuxiliaryCheckpointComplete) { + self.auxiliary_state + .handle_auxiliary_checkpoint_complete(complete.checkpointed_lsn); self.reclaim_checkpointed_wal(); } @@ -404,13 +428,13 @@ impl DataPlane { self.wal.delete_below(watermark); } - if self.consumer_offsets.offset_checkpoint_in_flight() { + if self.auxiliary_state.auxiliary_checkpoint_in_flight() { return; } let Some(reclaimable_lsn) = self.wal.reclaimable_lsn() else { return; }; - let Some(oldest_offset_lsn) = self.consumer_offsets.oldest_uncheckpointed_lsn() else { + let Some(oldest_offset_lsn) = self.auxiliary_state.oldest_uncheckpointed_lsn() else { return; }; if self.segments.reclamation_watermark(reclaimable_lsn) < *oldest_offset_lsn { @@ -419,10 +443,10 @@ impl DataPlane { // At this point, lsn must exist in consumer_offset so the following is just safeguard. if let Some(job) = self - .consumer_offsets - .raise_offset_checkpoint_job(self.config.data_dir.clone()) + .auxiliary_state + .raise_auxiliary_checkpoint_job(self.config.data_dir.clone()) { - self.out.store_offset_checkpoint(job); + self.out.store_auxiliary_checkpoint(job); } } @@ -561,8 +585,7 @@ impl DataPlane { } fn handle_consumer_group_epoch_seal(&mut self, cmd: EpochSeal) { - let Some(parked_commits) = self.consumer_offsets.handle_consumer_group_epoch_seal(cmd) - else { + let Some(parked_commits) = self.auxiliary_state.unpark_future_offset_commits(cmd) else { return; }; @@ -949,6 +972,14 @@ impl DataPlane { }; self.commit_segment(cmd.segment_key, committed.entry_id); + if let Some(p_id) = committed.producer_identity { + self.auxiliary_state.advance_producer( + cmd.segment_key, + p_id, + committed.entry_id, + committed.lsn, + ); + } self.out.produce_replies.extend( committed .replies @@ -967,7 +998,7 @@ impl DataPlane { fn commit_consumer_offset(&mut self, cmd: CommitConsumerOffset) { self.needs_flush |= self - .consumer_offsets + .auxiliary_state .commit_consumer_offset(cmd, &self.node_id); } @@ -980,7 +1011,7 @@ impl DataPlane { return; }; - let sealed_generation = self.consumer_offsets.latest_generation(&cmd.update.key); + let sealed_generation = self.auxiliary_state.latest_generation(&cmd.update.key); if cmd.update.generation < sealed_generation { let transport = DataTransportCommand::send_to_targets( vec![leader], @@ -1011,20 +1042,19 @@ impl DataPlane { self.out.store_transport_cmd(transport) } - self.consumer_offsets - .future_entry(cmd.update.key.clone()) - .push(FutureOffsetCommit::Replica(cmd)); + self.auxiliary_state + .push_future(cmd.update.key.clone(), FutureOffsetCommit::Replica(cmd)); return; } - self.consumer_offsets.push_pending_offset_mutation(cmd); + self.auxiliary_state.push_pending_offset_mutation(cmd); self.needs_flush = true; } fn handle_replica_offset_ack(&mut self, cmd: ConsumerOffsetReplicated) { // ! Hot PATH without the following condition, // every consumer-offset replica acknowledgement currently will trigger a full placement scan - if self.consumer_offsets.handle_replica_offset_ack(cmd) { + if self.auxiliary_state.handle_replica_offset_ack(cmd) { self.ack_ready_offset_placements(); } } @@ -1041,7 +1071,7 @@ impl DataPlane { }; let parked = self - .consumer_offsets + .auxiliary_state .install_consumer_offsets(cmd, &self.node_id); self.resume_future_offset_commits(parked); @@ -1061,7 +1091,7 @@ impl DataPlane { } if let Some(bootstrap) = self - .consumer_offsets + .auxiliary_state .create_offset_snapshot(cmd, &self.node_id) { self.out.store_transport_cmd(bootstrap); @@ -1073,10 +1103,10 @@ impl DataPlane { return; } - self.consumer_offsets.handle_snapshot_installed_ack(&cmd); + self.auxiliary_state.handle_snapshot_installed_ack(&cmd); if let Some(bootstrap) = self - .consumer_offsets + .auxiliary_state .create_offset_snapshot_for_unready_replicas(cmd.segment_key) { self.out.store_transport_cmd(bootstrap); @@ -1104,7 +1134,7 @@ impl DataPlane { return; } - if let Some(transport) = self.consumer_offsets.install_leader_placement( + if let Some(transport) = self.auxiliary_state.install_leader_placement( cmd.segment_key, cmd.shard_group_id, cmd.replica_set.clone(), @@ -1131,7 +1161,7 @@ impl DataPlane { /// have successfully bootstrapped and are up-to-date) /// If then sends [`SegmentPlaced`](self::SegmentPlaced) to coordinator responsible for the shard fn ack_ready_offset_placements(&mut self) { - for ack in self.consumer_offsets.drain_ready_placements(&self.node_id) { + for ack in self.auxiliary_state.drain_ready_placements(&self.node_id) { self.out .store_transport_cmd(DataTransportSendToCoordinator { shard_group_id: ack.shard_group_id, @@ -1150,7 +1180,7 @@ impl DataPlane { } if matches!( - self.consumer_offsets.observe_placement( + self.auxiliary_state.observe_placement( cmd.segment_key, ShardGroupId(0), &cmd.replicas, @@ -1179,7 +1209,13 @@ impl DataPlane { return; }; - tracker.stage_entry(cmd.segment_key, cmd.data, cmd.record_count, cmd.entry_id); + tracker.stage_entry( + cmd.segment_key, + cmd.data, + cmd.record_count, + cmd.entry_id, + cmd.producer_identity, + ); self.dirty_segments.insert(cmd.segment_key); self.needs_flush = true; @@ -1201,7 +1237,7 @@ impl DataPlane { let uncommitted_entry = old_tracker .uncommitted_entries() .chain(old_tracker.staged_for_replay()) - .collect::>(); + .collect::)]>>(); if !old_tracker.followers().is_empty() { self.out @@ -1250,13 +1286,14 @@ impl DataPlane { .segments .get_mut(&cmd.old_segment_key.with_segment_id(cmd.new_segment_id)) { - for (data, record_count) in uncommitted_entry { + for (data, record_count, producer) in uncommitted_entry { let entry_id = successor.next_staged_entry_id(); successor.stage_entry( cmd.old_segment_key.with_segment_id(cmd.new_segment_id), data, record_count, entry_id, + producer, ); } } @@ -1356,7 +1393,7 @@ impl DataPlane { if tracker.role() == SegmentRole::Leader { let staged_bytes: usize = tracker .staged_for_replay() - .map(|(data, _)| data.len()) + .map(|(data, _, _)| data.len()) .sum(); self.buffer_byte_count -= staged_bytes; } @@ -1472,7 +1509,7 @@ impl DataPlane { } fn has_staged_data(&self) -> bool { - self.consumer_offsets.has_pending_offsets() + self.auxiliary_state.has_pending_offsets() || self .dirty_segments .iter() @@ -1480,7 +1517,7 @@ impl DataPlane { } fn flush_batch(&mut self) { - let Ok(encoded_offsets) = self.consumer_offsets.encode_offsets() else { + let Ok(encoded_offsets) = self.auxiliary_state.encode_offsets() else { self.needs_flush = !self.dirty_segments.is_empty(); return; }; @@ -1509,7 +1546,7 @@ impl DataPlane { } }; - for c in self.consumer_offsets.flush_batch(&self.node_id, lsn) { + for c in self.auxiliary_state.flush_batch(&self.node_id, lsn) { self.out.store_transport_cmd(c); } @@ -1525,10 +1562,16 @@ impl DataPlane { if tracker.followers().is_empty() { // No followers → committed the moment it's durable. tracker.commit_entry(entry.entry_id); + if let Some(p_id) = entry.producer_identity { + self.auxiliary_state.advance_producer( + key, + p_id, + entry.entry_id, + lsn, + ); + } if let Some(reply) = self.replication.pop_pending_reply(&key) { - let _ = reply.send(ProduceAck::Ok { - entry_id: entry.entry_id, - }); + let _ = reply.send(ProduceAck::Ok(entry.entry_id)); } } else { segment_batches.push(PendingReplicationBatch { @@ -1540,6 +1583,10 @@ impl DataPlane { } } SegmentRole::Follower => { + if let Some(p_id) = entry.producer_identity { + self.auxiliary_state + .advance_producer(key, p_id, entry.entry_id, lsn); + } self.out .store_transport_cmd(DataTransportCommand::send_to_targets( vec![tracker.leader_node()], @@ -1587,12 +1634,15 @@ impl DataPlane { for key in std::mem::take(&mut self.dirty_segments) { if let Some(tracker) = self.segments.get_mut(&key) { - tracker.abort_staged(); + for append_identity in tracker.abort_staged() { + let append_key = AppendKey::new(key, &append_identity); + self.auxiliary_state.unstage_producer(&append_key); + } } } self.replication.fail_all(e.to_string()); - self.consumer_offsets.fail_all(&e.to_string()); + self.auxiliary_state.drop_coordination(&e.to_string()); self.needs_flush = false; } @@ -1603,7 +1653,7 @@ impl DataPlane { }; let segment_watermark = self.segments.reclamation_watermark(reclaimable_lsn); - let offset_watermark = self.consumer_offsets.reclamation_watermark(reclaimable_lsn); + let offset_watermark = self.auxiliary_state.reclamation_watermark(reclaimable_lsn); segment_watermark.min(offset_watermark) } @@ -1708,25 +1758,27 @@ impl TAssertInvariant for DataPlane { ); self.segments.assert_invariants(); + self.replication.assert_invariants(); + self.auxiliary_state.assert_producer_invariants(); } } #[cfg(test)] mod tests { + use super::*; use crate::control_plane::Replicas; use crate::control_plane::metadata::consumer_group::GenerationId; - use crate::data_plane::consumer_offset_management::ledger::{ - ConsumerOffsetKey, ConsumerOffsetPosition, ConsumerOffsetUpdate, EpochSeal, OffsetLedger, + use crate::data_plane::auxiliary_states::consumer_offsets::state::{ + ConsumerOffsetEntry, ConsumerOffsetKey, ConsumerOffsetPosition, ConsumerOffsetUpdate, }; - use std::path::PathBuf; - use super::*; use crate::control_plane::membership::ShardGroupId; use crate::control_plane::metadata::{RangeId, SegmentId, TopicId}; use crate::data_plane::checkpoint::CheckpointTask; use crate::data_plane::cold_read::DEFAULT_POOL_SIZE; use crate::data_plane::recovery::segment_scan::RecoveredSegments; use crate::data_plane::wal::WalWriter; + use std::path::PathBuf; use tokio::sync::oneshot; use bytes::Bytes; @@ -1896,11 +1948,11 @@ mod tests { ConsumerOffsetCommitAck::Committed ); assert_eq!( - dp.consumer_offsets.offset(&consumer_offset_key()), + dp.auxiliary_state.offset(&consumer_offset_key()), Some(position) ); assert_eq!( - dp.consumer_offsets + dp.auxiliary_state .uncheckpointed_offset_lsns() .iter() .copied() @@ -1950,14 +2002,14 @@ mod tests { dp.reclaim_checkpointed_wal(); assert!(matches!( dp.out.checkpoint_tasks.as_slice(), - [CheckpointTask::ConsumerOffsets(job)] if job.checkpointed_lsn == 2 + [CheckpointTask::AuxiliaryState(job)] if job.checkpointed_lsn == 2 )); - dp.handle_command(OffsetCheckpointComplete { + dp.handle_command(AuxiliaryCheckpointComplete { checkpointed_lsn: 2, }); - assert!(dp.consumer_offsets.uncheckpointed_offset_lsns().is_empty()); - assert_eq!(dp.consumer_offsets.offset_checkpoint_lsn(), 2); + assert!(dp.auxiliary_state.uncheckpointed_offset_lsns().is_empty()); + assert_eq!(dp.auxiliary_state.auxiliary_checkpoint_lsn(), 2); assert_eq!(dp.wal.reclaimable_lsn(), None); } @@ -2025,7 +2077,7 @@ mod tests { ConsumerOffsetCommitAck::Committed ); assert_eq!( - dp.consumer_offsets.offset(&consumer_offset_key()), + dp.auxiliary_state.offset(&consumer_offset_key()), Some(position) ); } @@ -2152,7 +2204,7 @@ mod tests { dp.flush_batch(); assert_eq!( - dp.consumer_offsets.offset(&consumer_offset_key()), + dp.auxiliary_state.offset(&consumer_offset_key()), Some(position) ); assert!(dp.out.transport_cmds.iter().any(|command| matches!( @@ -2184,13 +2236,11 @@ mod tests { ConsumerOffsetSnapshot { segment_key: test_key(), replica_set: Replicas::new(vec![leader.clone(), test_node_id()]), - entries: vec![ - crate::data_plane::consumer_offset_management::ledger::ConsumerOffsetEntry { - key: consumer_offset_key(), - generation: 3.into(), - position: Some(position), - }, - ] + entries: vec![ConsumerOffsetEntry { + key: consumer_offset_key(), + generation: 3.into(), + position: Some(position), + }] .into_boxed_slice(), }, )); @@ -2204,7 +2254,7 @@ mod tests { dp.flush_batch(); assert_eq!( - dp.consumer_offsets.offset(&consumer_offset_key()), + dp.auxiliary_state.offset(&consumer_offset_key()), Some(position) ); assert!(dp.out.transport_cmds.iter().any(|cmd| matches!( @@ -2238,7 +2288,7 @@ mod tests { update: update.clone(), }, )); - assert!(!dp.consumer_offsets.has_pending_offsets()); + assert!(!dp.auxiliary_state.has_pending_offsets()); assert!(dp.out.transport_cmds.iter().any(|cmd| matches!( cmd, DataTransportCommand::SendToTargets(send) @@ -2250,20 +2300,18 @@ mod tests { ConsumerOffsetSnapshot { segment_key: test_key(), replica_set: Replicas::new(vec![leader.clone(), test_node_id()]), - entries: vec![ - crate::data_plane::consumer_offset_management::ledger::ConsumerOffsetEntry { - key: consumer_offset_key(), - generation: 2.into(), - position: None, - }, - ] + entries: vec![ConsumerOffsetEntry { + key: consumer_offset_key(), + generation: 2.into(), + position: None, + }] .into_boxed_slice(), }, )); dp.flush_batch(); assert_eq!( - dp.consumer_offsets.offset(&consumer_offset_key()), + dp.auxiliary_state.offset(&consumer_offset_key()), Some(position) ); assert!(dp.out.transport_cmds.iter().any(|cmd| matches!( @@ -2329,7 +2377,10 @@ mod tests { RecoveryOutput { inventory: recovered, data_dir: dir.path().to_path_buf(), - offsets: OffsetLedger::default(), + auxiliary_state: AuxiliaryStateManager::from_recovery( + Default::default(), + Default::default(), + ), }, ) } @@ -2384,6 +2435,8 @@ mod tests { segment_key: key, data: Bytes::from("data").into(), record_count: 1, + received_at_ms: 0, + producer_identity: None, reply: tx, }; (cmd.into(), rx) @@ -2462,6 +2515,8 @@ mod tests { segment_key: test_key(), data: Bytes::from("hello world!").into(), record_count: 2, + received_at_ms: 0, + producer_identity: None, reply: tx, }); @@ -2470,6 +2525,125 @@ mod tests { assert_eq!(tracker.next_entry_id(), EntryId(0)); } + #[test] + fn committed_producer_retry_returns_same_entry_without_appending() { + let dir = tempfile::tempdir().unwrap(); + let mut dp = make_data_plane(&dir); + dp.handle_command(assign_segment(test_key(), vec![test_node_id()])); + let producer = ProducerAppendIdentity { + producer_id: uuid::Uuid::new_v4(), + incarnation: 0, + expires_at: u64::MAX, + sequence: 0, + digest: 7, + }; + + let send = |data_plane: &mut DataPlane<_>| { + let (reply, ack) = oneshot::channel(); + data_plane.handle_command(Produce { + segment_key: test_key(), + data: Bytes::from_static(b"one").into(), + record_count: 1, + received_at_ms: 0, + producer_identity: Some(AuthorizedProducerIdentity::MetadataVerified(producer)), + reply, + }); + data_plane.flush_batch(); + ack.blocking_recv().unwrap() + }; + + assert!(matches!(send(&mut dp), ProduceAck::Ok(EntryId(0)))); + assert!(matches!(send(&mut dp), ProduceAck::Ok(EntryId(0)))); + assert_eq!( + dp.segments.get(&test_key()).unwrap().next_staged_entry_id(), + EntryId(1) + ); + } + + #[test] + fn committed_retry_survives_segment_removal_but_new_append_does_not() { + let dir = tempfile::tempdir().unwrap(); + let mut dp = make_data_plane(&dir); + dp.handle_command(assign_segment(test_key(), vec![test_node_id()])); + let producer_id = uuid::Uuid::new_v4(); + let request = ProducerAppendIdentity { + producer_id, + incarnation: 0, + expires_at: u64::MAX, + sequence: 0, + digest: 7, + }; + let (reply, ack) = oneshot::channel(); + dp.handle_command(Produce { + segment_key: test_key(), + data: Bytes::from_static(b"one").into(), + record_count: 1, + received_at_ms: 0, + producer_identity: Some(AuthorizedProducerIdentity::MetadataVerified(request)), + reply, + }); + dp.flush_batch(); + assert!(matches!( + ack.blocking_recv().unwrap(), + ProduceAck::Ok { .. } + )); + dp.segments.take_active_and_seal(test_key()); + + let (duplicate_reply, duplicate) = oneshot::channel(); + dp.handle_command(Produce { + segment_key: test_key(), + data: Bytes::from_static(b"one").into(), + record_count: 1, + received_at_ms: 0, + producer_identity: Some(AuthorizedProducerIdentity::ExistingOnly(request)), + reply: duplicate_reply, + }); + assert!(matches!( + duplicate.blocking_recv().unwrap(), + ProduceAck::Ok(EntryId(0)) + )); + + let (not_committed_reply, not_committed) = oneshot::channel(); + dp.handle_command(Produce { + segment_key: test_key(), + data: Bytes::from_static(b"two").into(), + record_count: 1, + received_at_ms: 0, + producer_identity: Some(AuthorizedProducerIdentity::ExistingOnly( + ProducerAppendIdentity { + sequence: 1, + digest: 8, + ..request + }, + )), + reply: not_committed_reply, + }); + assert!(matches!( + not_committed.blocking_recv().unwrap(), + ProduceAck::Err(ProduceError::SegmentNotFound) + )); + + let (retry_reply, retry) = oneshot::channel(); + dp.handle_command(Produce { + segment_key: test_key(), + data: Bytes::from_static(b"two").into(), + record_count: 1, + received_at_ms: 0, + producer_identity: Some(AuthorizedProducerIdentity::ExistingOnly( + ProducerAppendIdentity { + sequence: 1, + digest: 8, + ..request + }, + )), + reply: retry_reply, + }); + assert!(matches!( + retry.blocking_recv().unwrap(), + ProduceAck::Err(ProduceError::SegmentNotFound) + )); + } + #[test] fn segment_assignment_creates_tracker() { let dir = tempfile::tempdir().unwrap(); @@ -2566,6 +2740,8 @@ mod tests { segment_key: test_key(), data: Bytes::from(vec![0u8; TEST_BATCH_MAX_BYTES]).into(), record_count: 1, + received_at_ms: 0, + producer_identity: None, reply: oneshot::channel().0, }); @@ -2584,6 +2760,8 @@ mod tests { segment_key: test_key(), data: Bytes::from(vec![0u8; TEST_BATCH_MAX_BYTES]).into(), record_count: 1, + received_at_ms: 0, + producer_identity: None, reply: oneshot::channel().0, }); assert!(dp.buffer_byte_count >= TEST_BATCH_MAX_BYTES); @@ -2605,6 +2783,8 @@ mod tests { segment_key: test_key(), data: Bytes::from(vec![0u8; 16]).into(), record_count: 1, + received_at_ms: 0, + producer_identity: None, reply: oneshot::channel().0, }); dp.flush_batch(); @@ -2633,12 +2813,16 @@ mod tests { segment_key: small, data: Bytes::from(vec![0u8; 10]).into(), record_count: 1, + received_at_ms: 0, + producer_identity: None, reply: oneshot::channel().0, }); dp.handle_command(Produce { segment_key: large, data: Bytes::from(vec![0u8; 45]).into(), record_count: 1, + received_at_ms: 0, + producer_identity: None, reply: oneshot::channel().0, }); dp.flush_batch(); @@ -2666,6 +2850,8 @@ mod tests { segment_key: test_key(), data: Bytes::from(vec![0u8; 8]).into(), record_count: 1, + received_at_ms: 0, + producer_identity: None, reply: oneshot::channel().0, }); dp.flush_batch(); @@ -2693,6 +2879,8 @@ mod tests { segment_key: test_key(), data: Bytes::from(vec![0u8; 8]).into(), record_count: 1, + received_at_ms: 0, + producer_identity: None, reply: oneshot::channel().0, }); dp.flush_batch(); @@ -2966,6 +3154,7 @@ mod tests { data: b"setup".to_vec().into(), record_count: 1, entry_id: EntryId(0), + producer_identity: None, } .into(), )); @@ -2991,6 +3180,7 @@ mod tests { data: b"data".to_vec().into(), record_count: 1, entry_id: EntryId(0), + producer_identity: None, } .into(), )); @@ -3016,6 +3206,7 @@ mod tests { data: b"data".to_vec().into(), record_count: 1, entry_id: EntryId(0), + producer_identity: None, }, )); @@ -3038,6 +3229,7 @@ mod tests { data: b"data".to_vec().into(), record_count: 1, entry_id: EntryId(0), + producer_identity: None, } .into(), )); @@ -3058,6 +3250,7 @@ mod tests { data: b"data".to_vec().into(), record_count: 1, entry_id: EntryId(0), + producer_identity: None, } .into(), )); @@ -3181,6 +3374,7 @@ mod tests { data: b"data".to_vec().into(), record_count: 1, entry_id: EntryId(0), + producer_identity: None, } .into(), )); @@ -3311,6 +3505,7 @@ mod tests { data: b"data".to_vec().into(), record_count: 1, entry_id: EntryId(0), + producer_identity: None, } .into(), )); @@ -3508,7 +3703,7 @@ mod tests { RecoveryOutput { inventory: empty_inventory(), data_dir: dir.path().to_path_buf(), - offsets: OffsetLedger::default(), + auxiliary_state: AuxiliaryStateManager::default(), }, ); @@ -3521,6 +3716,8 @@ mod tests { segment_key: test_key(), data: Bytes::copy_from_slice(payload).into(), record_count, + received_at_ms: 0, + producer_identity: None, reply: tx, }); } @@ -3574,14 +3771,11 @@ mod tests { let result = reply_rx .blocking_recv() .expect("cold-read pool dropped the reply"); - let FetchResult::Records { + let FetchedRecords { entries, next_entry_id, .. - } = result - else { - panic!("expected Records from cold read, got {result:?}"); - }; + } = result.expect("expected Records from cold read"); let got: Vec<(u64, u32, Vec)> = entries .iter() .map(|e| (*e.entry_id, e.record_count, e.data.to_vec())) @@ -3620,7 +3814,7 @@ mod tests { RecoveryOutput { inventory: empty_inventory(), data_dir: dir.path().to_path_buf(), - offsets: OffsetLedger::default(), + auxiliary_state: AuxiliaryStateManager::default(), }, ); let mut tracker = SegmentTracker::new_with_start_entry_id( @@ -3712,18 +3906,21 @@ mod tests { record_count: 1, entry_id: EntryId(0), lsn: 0, + producer_identity: None, }), Arc::new(CachedEntry { data: Bytes::copy_from_slice(b"bb").into(), record_count: 2, entry_id: EntryId(1), lsn: 0, + producer_identity: None, }), Arc::new(CachedEntry { data: Bytes::copy_from_slice(b"ccc").into(), record_count: 3, entry_id: EntryId(2), lsn: 0, + producer_identity: None, }), ]; dp.handle_command(DataPlaneCommand::CatchUpReadComplete(CatchUpReadComplete { @@ -3790,7 +3987,7 @@ mod tests { RecoveryOutput { inventory: empty_inventory(), data_dir: dir.path().to_path_buf(), - offsets: OffsetLedger::default(), + auxiliary_state: AuxiliaryStateManager::default(), }, ); let requester = NodeId::new("replacement"); @@ -3805,6 +4002,7 @@ mod tests { record_count: 1, entry_id: EntryId(4), lsn: 0, + producer_identity: None, })], next_offset: EntryId(5), })); @@ -4386,6 +4584,7 @@ mod tests { data: b"x".to_vec().into(), record_count: 1, entry_id: EntryId(entry_id), + producer_identity: None, } .into(), )); diff --git a/src/data_plane/states/replication.rs b/src/data_plane/states/replication.rs index 346683f9..16ced5ad 100644 --- a/src/data_plane/states/replication.rs +++ b/src/data_plane/states/replication.rs @@ -30,7 +30,7 @@ pub(crate) struct AckCommitted { pub entry_id: EntryId, pub replies: Vec>, pub reset_timer_seq: Option, - pub producer_append_id: Option, + pub producer_identity: Option, pub lsn: u64, } @@ -177,7 +177,7 @@ impl ReplicationState { entry_id: batch.entry_id, replies: batch.replies, reset_timer_seq, - producer_append_id: batch.producer_append_id, + producer_identity: batch.producer_append_id, lsn: batch.lsn, }) } diff --git a/src/data_plane/states/segment/tracker.rs b/src/data_plane/states/segment/tracker.rs index 999b9d15..5aa13076 100644 --- a/src/data_plane/states/segment/tracker.rs +++ b/src/data_plane/states/segment/tracker.rs @@ -272,7 +272,7 @@ impl SegmentTracker { self.last_committed_entry_id() == self.durable_end_entry_id() } - pub(crate) fn stage_producer_entry( + pub(crate) fn stage_entry( &mut self, segment_key: SegmentKey, data: PayloadBytes, @@ -402,7 +402,7 @@ pub mod tests { #[test] fn stage_entry_tracks_size() { let mut t = make_tracker(SegmentRole::Leader); - t.stage_producer_entry(test_key(), Bytes::from("abcde").into(), 2, EntryId(0), None); + t.stage_entry(test_key(), Bytes::from("abcde").into(), 2, EntryId(0), None); assert_eq!(t.size_bytes, 5); assert!(t.has_staged()); @@ -418,7 +418,7 @@ pub mod tests { ShardGroupId(1), EntryId(5), ); - t.stage_producer_entry(test_key(), Bytes::from("data").into(), 1, EntryId(5), None); + t.stage_entry(test_key(), Bytes::from("data").into(), 1, EntryId(5), None); assert!(t.has_staged()); // Publish to advance next_entry_id @@ -427,7 +427,7 @@ pub mod tests { t.publish_staged(1); // Duplicate entry_id (5) should be skipped since next is now 6 - t.stage_producer_entry(test_key(), Bytes::from("dup").into(), 1, EntryId(5), None); + t.stage_entry(test_key(), Bytes::from("dup").into(), 1, EntryId(5), None); assert!(!t.has_staged()); assert_eq!(t.next_entry_id, EntryId(6)); t.assert_invariants(); @@ -439,7 +439,7 @@ pub mod tests { let mut wal_buf = Vec::new(); for i in 0..3u64 { - t.stage_producer_entry( + t.stage_entry( test_key(), Bytes::from(format!("entry-{i}")).into(), 1, @@ -474,7 +474,7 @@ pub mod tests { ShardGroupId(1), EntryId(2), ); - t.stage_producer_entry( + t.stage_entry( test_key(), Bytes::from("entry-2").into(), 1, @@ -493,7 +493,7 @@ pub mod tests { fn stage_then_publish_drains() { let mut t = make_tracker(SegmentRole::Leader); let mut wal_buf = Vec::new(); - t.stage_producer_entry( + t.stage_entry( test_key(), Bytes::from("payload").into(), 3, @@ -559,7 +559,7 @@ pub mod tests { t.last_activity_at = tokio::time::Instant::now() - Duration::from_secs(60); assert!(t.idle_timeout_reached(Duration::from_secs(30))); - t.stage_producer_entry(test_key(), Bytes::from("data").into(), 1, EntryId(0), None); + t.stage_entry(test_key(), Bytes::from("data").into(), 1, EntryId(0), None); t.publish_staged(1); t.commit_entry(EntryId(0)); diff --git a/src/data_plane/states/segment_store.rs b/src/data_plane/states/segment_store.rs index f43a5e40..74f71f2b 100644 --- a/src/data_plane/states/segment_store.rs +++ b/src/data_plane/states/segment_store.rs @@ -100,7 +100,7 @@ impl SegmentStore { ) -> Option { let tracker = self.get_mut(&key)?; let entry_id = tracker.next_staged_entry_id(); - tracker.stage_producer_entry(key, data, record_count, entry_id, producer_identity); + tracker.stage_entry(key, data, record_count, entry_id, producer_identity); Some(entry_id) } From 18dabb084c891d7659600b0513d2c9976117026a Mon Sep 17 00:00:00 2001 From: Migorithm Date: Thu, 23 Jul 2026 15:09:59 +0400 Subject: [PATCH 33/41] evict_expired_producer_sessions --- src/control_plane/consensus/raft/state.rs | 29 +++++++++++++++++++---- 1 file changed, 25 insertions(+), 4 deletions(-) diff --git a/src/control_plane/consensus/raft/state.rs b/src/control_plane/consensus/raft/state.rs index 60c42b68..9c7e5304 100644 --- a/src/control_plane/consensus/raft/state.rs +++ b/src/control_plane/consensus/raft/state.rs @@ -13,7 +13,7 @@ use crate::control_plane::consensus::raft::storage::{ }; use crate::control_plane::consensus::raft::{compute_replacement_replica_set, now_ms}; use crate::control_plane::membership::{ShardGroupId, TopologyReader}; -use crate::control_plane::metadata::command::DeleteSegments; +use crate::control_plane::metadata::command::{DeleteSegments, ExpireProducerSessions}; use crate::control_plane::metadata::event::MetadataEvent; use crate::control_plane::metadata::{ ConsumerGroupAssignment, ConsumerMemberId, MetadataCommand, ReassignSegment, RollSegment, @@ -202,6 +202,7 @@ impl Raft { // Ring members to assert as peers (see doc above); `None` skips the add step. target_members: Option>, ) -> bool { + let now = now_ms(); let mut changed = false; let live_set: HashSet = topology.live_nodes().into_iter().collect(); if let Some(members) = target_members { @@ -221,7 +222,8 @@ impl Raft { changed |= self.reconcile_peers(topology, &live_set); changed |= self.reconcile_segments(&live_set); changed |= self.refill_under_replicated_segments(topology); - changed |= self.reconcile_retention_deletes(); + changed |= self.reconcile_retention_deletes(now); + changed |= self.evict_expired_producer_sessions(now); // Record a ring observation every run, whether routine RingCheck or // leadership change. During a leadership change, evictions are paused @@ -573,8 +575,7 @@ impl Raft { /// topic's `retention_ms` (against the wall clock, the one age-based decision — /// leader-only so cross-node skew can't diverge replicas), and propose /// `DeleteSegments`. Topics with no retention policy contribute nothing. - fn reconcile_retention_deletes(&mut self) -> bool { - let now = now_ms(); + fn reconcile_retention_deletes(&mut self, now: u64) -> bool { let targets = self.metadata.expipred_segments(now); let mut changed = false; @@ -599,6 +600,25 @@ impl Raft { changed } + fn evict_expired_producer_sessions(&mut self, now: u64) -> bool { + let topics = self.metadata.topics_with_expired_producer_sessions(now); + let mut changed = false; + for topic_id in topics { + if let Err(error) = self.propose( + ExpireProducerSessions { + topic_id, + observed_at: now, + } + .into(), + ) { + tracing::warn!(?topic_id, ?error, "producer session expiry rejected"); + continue; + } + changed = true; + } + changed + } + pub(crate) fn handle_node_death( &mut self, dead_node_id: &NodeId, @@ -1839,6 +1859,7 @@ impl crate::test_traits::TAssertInvariant for Raft { !self.peers.contains(&self.node_id), "self found in peers set", ); + self.metadata.assert_invariants(); } } From e83463838e9ad4d50a37368905a670e41d1079a7 Mon Sep 17 00:00:00 2001 From: Migorithm Date: Thu, 23 Jul 2026 15:30:52 +0400 Subject: [PATCH 34/41] option -> result --- src/control_plane/consensus/actor.rs | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/src/control_plane/consensus/actor.rs b/src/control_plane/consensus/actor.rs index 5a303c93..7679972f 100644 --- a/src/control_plane/consensus/actor.rs +++ b/src/control_plane/consensus/actor.rs @@ -3,6 +3,7 @@ use std::collections::BTreeMap; use crate::channels::BatchSender; +use crate::client::ServerError; use crate::control_plane::NodeId; use crate::control_plane::consensus::messages::*; use crate::control_plane::consensus::multi_raft::MultiRaft; @@ -245,13 +246,18 @@ impl MutlRaftSender { /// Look up a topic's full metadata on this node. Returns `None` when this /// node is not in the owning shard group (caller should redirect) or when /// the actor channel is dead. - pub(crate) async fn get_topic_metadata(&self, topic_name: String) -> Option { + pub(crate) async fn get_topic_metadata( + &self, + topic_name: String, + ) -> Result { let (reply, recv) = tokio::sync::oneshot::channel(); let _ = self .send(MultiRaftActorCommand::GetTopicMetadata { topic_name, reply }) .await; - recv.await.unwrap_or_default() + recv.await + .map_err(|e| ServerError::Internal(e.to_string()))? + .ok_or(ServerError::TopicNotFound) } pub(crate) async fn get_consumer_group_assignment( From 6df346cf375f94a8dbc45b3635e108b6f8751754 Mon Sep 17 00:00:00 2001 From: Migorithm Date: Thu, 23 Jul 2026 15:47:23 +0400 Subject: [PATCH 35/41] error mapping --- src/connections/protocol/error.rs | 15 +++++++++++++++ src/control_plane/metadata/error.rs | 14 +++++++++++++- src/control_plane/metadata/topic.rs | 6 ++++++ 3 files changed, 34 insertions(+), 1 deletion(-) diff --git a/src/connections/protocol/error.rs b/src/connections/protocol/error.rs index 1222a966..11c225fc 100644 --- a/src/connections/protocol/error.rs +++ b/src/connections/protocol/error.rs @@ -70,3 +70,18 @@ impl ServerError { ) } } + +impl From for ServerError { + fn from(err: crate::control_plane::metadata::error::MetadataError) -> Self { + use crate::control_plane::metadata::error::MetadataError::*; + + match err { + TopicNotFound(_) | TopicNameNotFound(_) => ServerError::TopicNotFound, + TopicNameAlreadyExists(_) => ServerError::AlreadyExists, + TopicNotActive(_) | RangeNotFound | RangeNotActive => ServerError::StaleRange, + SegmentNotFound | SegmentNotActive | SegmentNotSealed => ServerError::SegmentNotLocal, + InvalidSplitPoint => ServerError::InvalidSplitPoint, + SplitNotAllowed(_) | RangesNotAdjacent => ServerError::Internal(err.to_string()), + } + } +} diff --git a/src/control_plane/metadata/error.rs b/src/control_plane/metadata/error.rs index 089e4684..f6e63147 100644 --- a/src/control_plane/metadata/error.rs +++ b/src/control_plane/metadata/error.rs @@ -1,17 +1,29 @@ use crate::control_plane::metadata::TopicId; -#[derive(Debug, Clone, PartialEq, Eq)] +#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)] pub enum MetadataError { + #[error("topic not found: {0:?}")] TopicNotFound(TopicId), + #[error("topic name not found: {0}")] TopicNameNotFound(String), + #[error("topic name already exists: {0}")] TopicNameAlreadyExists(String), + #[error("topic not active: {0:?}")] TopicNotActive(TopicId), + #[error("range not found")] RangeNotFound, + #[error("range not active")] RangeNotActive, + #[error("segment not found")] SegmentNotFound, + #[error("segment not active")] SegmentNotActive, + #[error("segment not sealed")] SegmentNotSealed, + #[error("split not allowed for topic: {0:?}")] SplitNotAllowed(TopicId), + #[error("ranges not adjacent")] RangesNotAdjacent, + #[error("invalid split point")] InvalidSplitPoint, } diff --git a/src/control_plane/metadata/topic.rs b/src/control_plane/metadata/topic.rs index 3d88811b..528a4aca 100644 --- a/src/control_plane/metadata/topic.rs +++ b/src/control_plane/metadata/topic.rs @@ -85,6 +85,12 @@ impl TopicMeta { Ok(vec_ranges.try_into().unwrap()) } + pub(crate) fn get_range(&self, range_id: &RangeId) -> Result<&RangeMeta, MetadataError> { + self.ranges + .get(range_id) + .ok_or(MetadataError::RangeNotFound) + } + pub(crate) fn get_range_mut( &mut self, range_id: &RangeId, From 947be5e71c1536837ede705a337d137110399a52 Mon Sep 17 00:00:00 2001 From: Migorithm Date: Thu, 23 Jul 2026 16:10:40 +0400 Subject: [PATCH 36/41] vec -> bytes --- src/client/mod.rs | 6 +++--- src/connections/protocol/data_plane.rs | 4 ++-- src/control_plane/metadata/range.rs | 21 +++++++++++++++------ src/control_plane/metadata/segment.rs | 6 +++++- src/control_plane/metadata/topic.rs | 17 +++++++++++++++++ src/it/e2e/client_protocol.rs | 2 +- src/it/e2e/producer_deduplication.rs | 6 +++--- src/it/sim/scenario.rs | 2 +- 8 files changed, 47 insertions(+), 17 deletions(-) diff --git a/src/client/mod.rs b/src/client/mod.rs index 9f3fe6b0..e370b653 100644 --- a/src/client/mod.rs +++ b/src/client/mod.rs @@ -49,7 +49,7 @@ use crate::connections::protocol::{ FetchConsumerOffsetRequest, OpenProducerSessionRequest, ProduceRequest, ProducerSessionOpened, RangeOffsetRequest, }; -use crate::data_plane::{ProduceError, ProducerAppendIdentity}; +use crate::data_plane::{PayloadBytes, ProduceError, ProducerAppendIdentity}; use futures::{StreamExt, stream::FuturesUnordered}; pub use redirect::RetryPolicy; @@ -331,7 +331,7 @@ impl Client { topic: &str, range_id: RangeId, routing_key: &[u8], - data: Vec, + data: impl Into, record_count: u32, producer_identity: Option, ) -> Result { @@ -349,7 +349,7 @@ impl Client { topic_name: topic.to_string(), range_id, routing_key: routing_key.to_vec(), - data, + data: data.into(), record_count, producer_identity, }; diff --git a/src/connections/protocol/data_plane.rs b/src/connections/protocol/data_plane.rs index 5b5549c4..4e8efeee 100644 --- a/src/connections/protocol/data_plane.rs +++ b/src/connections/protocol/data_plane.rs @@ -14,7 +14,7 @@ use crate::{ EntryId, RangeId, RangeMeta, RangeState, TopicId, TopicMeta, consumer_group::GenerationId, }, data_plane::{ - ProducerAppendIdentity, + PayloadBytes, ProducerAppendIdentity, auxiliary_states::consumer_offsets::state::{ConsumerOffsetKey, ConsumerOffsetUpdate}, messages::query::FetchedRecords, }, @@ -81,7 +81,7 @@ pub struct ProduceRequest { /// Pre-serialized blob produced by the client: a leading 1-byte cleartext /// codec tag (none/lz4/zstd) followed by the optionally-compressed records. /// The broker stamps an entry_id and stores/replicates this opaque payload as-is. - pub data: Vec, + pub data: PayloadBytes, pub record_count: u32, pub producer_identity: Option, } diff --git a/src/control_plane/metadata/range.rs b/src/control_plane/metadata/range.rs index 947999f8..c845bbb8 100644 --- a/src/control_plane/metadata/range.rs +++ b/src/control_plane/metadata/range.rs @@ -2,12 +2,15 @@ use super::constants::*; use borsh::{BorshDeserialize, BorshSerialize}; use std::collections::HashMap; -use crate::control_plane::{ - Replicas, - metadata::{ - EntryId, RangeId, SegmentId, SegmentMeta, SegmentMetaState, SplitRange, - command::{RollSegment, SegmentRollIntent}, - error::MetadataError, +use crate::{ + client::ServerError, + control_plane::{ + Replicas, + metadata::{ + EntryId, RangeId, SegmentId, SegmentMeta, SegmentMetaState, SplitRange, + command::{RollSegment, SegmentRollIntent}, + error::MetadataError, + }, }, }; @@ -93,6 +96,12 @@ impl RangeMeta { Ok((left, right)) } + pub(crate) fn active_write_segment(&self) -> Result<&SegmentMeta, ServerError> { + self.active_segment + .map(|id| self.segments.get(&id).unwrap()) + .ok_or(ServerError::StaleRange) + } + pub(crate) fn validate_active(&self) -> Result { if self.state != RangeState::Active { return Err(MetadataError::RangeNotActive); diff --git a/src/control_plane/metadata/segment.rs b/src/control_plane/metadata/segment.rs index 200849c5..253061a5 100644 --- a/src/control_plane/metadata/segment.rs +++ b/src/control_plane/metadata/segment.rs @@ -1,7 +1,7 @@ use borsh::{BorshDeserialize, BorshSerialize}; use crate::control_plane::{ - Replicas, + NodeId, Replicas, metadata::{EntryId, SegmentId, error::MetadataError}, }; @@ -79,4 +79,8 @@ impl SegmentMeta { && self.end_entry_id.is_some() && self.replica_set.len() < replication_factor } + + pub(crate) fn is_shard_leader(&self, node_id: &NodeId) -> bool { + self.replica_set.first() == Some(node_id) + } } diff --git a/src/control_plane/metadata/topic.rs b/src/control_plane/metadata/topic.rs index 528a4aca..87f6f82f 100644 --- a/src/control_plane/metadata/topic.rs +++ b/src/control_plane/metadata/topic.rs @@ -3,6 +3,7 @@ use std::collections::{HashMap, HashSet}; use super::constants::*; use super::*; use crate::{ + client::ServerError, control_plane::{ NodeId, Replicas, metadata::{ @@ -91,6 +92,22 @@ impl TopicMeta { .ok_or(MetadataError::RangeNotFound) } + /// Resolves the active target range for a write operation. + /// Returns `ServerError::StaleRange` if the range is sealed/split or not active for the key. + pub(crate) fn resolve_writable_range( + &self, + range_id: RangeId, + routing_key: &[u8], + ) -> Result<&RangeMeta, ServerError> { + let range = self + .route_active_range(routing_key) + .ok_or(ServerError::TopicNotFound)?; + if range.range_id != range_id { + return Err(ServerError::StaleRange); + } + Ok(range) + } + pub(crate) fn get_range_mut( &mut self, range_id: &RangeId, diff --git a/src/it/e2e/client_protocol.rs b/src/it/e2e/client_protocol.rs index 1e26dbb4..637157b7 100644 --- a/src/it/e2e/client_protocol.rs +++ b/src/it/e2e/client_protocol.rs @@ -1819,7 +1819,7 @@ async fn produce_once(topic: &str, payload: &[u8], node: (&str, u16)) -> Produce topic_name: topic.into(), range_id, routing_key: b"k".to_vec(), - data: payload.to_vec(), + data: payload.to_vec().into(), record_count: 1, producer_identity: None, })); diff --git a/src/it/e2e/producer_deduplication.rs b/src/it/e2e/producer_deduplication.rs index dd3670bd..19412cff 100644 --- a/src/it/e2e/producer_deduplication.rs +++ b/src/it/e2e/producer_deduplication.rs @@ -13,7 +13,7 @@ use crate::connections::protocol::{ ClientDataPlaneRequest, ClientRequest, ClientResponse, OpenProducerSessionRequest, ProduceRequest, }; -use crate::data_plane::{ProduceError, ProducerAppendIdentity}; +use crate::data_plane::{PayloadBytes, ProduceError, ProducerAppendIdentity}; use crate::it::helpers::send_request; fn seeds() -> Vec { @@ -23,13 +23,13 @@ fn seeds() -> Vec { .collect() } -fn one_record(key: &[u8], value: &[u8]) -> Vec { +fn one_record(key: &[u8], value: &[u8]) -> PayloadBytes { let mut payload = vec![0]; // CompressionCodec::None payload.extend_from_slice(&(key.len() as u32).to_be_bytes()); payload.extend_from_slice(key); payload.extend_from_slice(&(value.len() as u32).to_be_bytes()); payload.extend_from_slice(value); - payload + payload.into() } #[test] diff --git a/src/it/sim/scenario.rs b/src/it/sim/scenario.rs index 67e59671..dfa91fe4 100644 --- a/src/it/sim/scenario.rs +++ b/src/it/sim/scenario.rs @@ -398,7 +398,7 @@ async fn produce_until_acked( topic_name: topic.to_string(), range_id, routing_key: b"campaign-key".to_vec(), - data: payload.to_vec(), + data: payload.to_vec().into(), record_count: 1, producer_identity: None, })); From 0ae2dae790cc1b4001997376b4f23b2d8081e021 Mon Sep 17 00:00:00 2001 From: Migorithm Date: Thu, 23 Jul 2026 16:16:51 +0400 Subject: [PATCH 37/41] feat(controller): simplify ClientController handlers with ServerError and domain range resolution - Change ClientController handlers to return Result directly - Map MetadataError into ServerError to unify control-plane error handling - Encapsulate range routing (resolve_writable_range) and segment lookup (active_write_segment) into TopicMeta/RangeMeta - Flatten ClientResponse into Ok(ClientSuccess) and Err(ServerError) --- src/client/mod.rs | 2 +- src/client/redirect.rs | 4 +- src/connections/controller.rs | 457 +++++++++++++++----------------- src/connections/protocol/mod.rs | 4 +- src/it/sim/scenario.rs | 2 +- 5 files changed, 216 insertions(+), 253 deletions(-) diff --git a/src/client/mod.rs b/src/client/mod.rs index e370b653..1a48029b 100644 --- a/src/client/mod.rs +++ b/src/client/mod.rs @@ -366,7 +366,7 @@ impl Client { self.cache.invalidate(topic); } match served.response { - ClientResponse::Ok(ClientSuccess::Produced { entry_id }) => Ok(entry_id), + ClientResponse::Ok(ClientSuccess::Produced(entry_id)) => Ok(entry_id), ClientResponse::Err(ServerError::StaleRange) => Err(ClientError::StaleRange), ClientResponse::Err(ServerError::ProduceRejected(error)) => { Err(ClientError::ProduceRejected(error)) diff --git a/src/client/redirect.rs b/src/client/redirect.rs index c025f7fb..3b219ee5 100644 --- a/src/client/redirect.rs +++ b/src/client/redirect.rs @@ -196,9 +196,7 @@ mod tests { "reresolve" ); assert_eq!( - classify(&ClientResponse::Ok(ClientSuccess::Produced { - entry_id: 7.into() - })), + classify(&ClientResponse::Ok(ClientSuccess::Produced(7.into()))), "done" ); } diff --git a/src/connections/controller.rs b/src/connections/controller.rs index 464e6a4c..732cbbeb 100644 --- a/src/connections/controller.rs +++ b/src/connections/controller.rs @@ -4,7 +4,7 @@ use crate::connections::{protocol::*, run_client_writer}; use crate::control_plane::NodeAddressInfo; use crate::control_plane::consensus::raft::errors::ProposalError; use crate::control_plane::metadata::{ - RangeMeta, SyncConsumerGroup, SyncConsumerGroupRequest, TopicState, + OpenProducerSession, RangeMeta, SyncConsumerGroup, SyncConsumerGroupRequest, TopicState, }; use crate::control_plane::{ NodeId, SwimNodeState, @@ -18,10 +18,10 @@ use crate::control_plane::{ strategy::StoragePolicy, }, }; -use crate::data_plane::EntryPayload; +use crate::data_plane::ProduceError; use crate::data_plane::SegmentKey; use crate::data_plane::actor::DataPlaneSender; -use crate::data_plane::consumer_offset_management::ledger::StaleEpoch; +use crate::data_plane::auxiliary_states::consumer_offsets::state::StaleEpoch; use crate::data_plane::messages::command::{ CommitConsumerOffset, ConsumerOffsetCommitAck, Produce, ProduceAck, }; @@ -29,26 +29,8 @@ use crate::data_plane::messages::query::{ DataPlaneQuery, Fetch, ListOffsets, ReadConsumerOffset, ReadConsumerOffsetResult, }; use crate::net::TcpStream; -use anyhow::Context; use tokio::sync::mpsc; -enum MetadataProposalOutcome { - Committed, - Redirect(Option), -} - -impl MetadataProposalOutcome { - fn on_ok(self, committed: ControlPlaneResponse) -> ClientResponse { - match self { - Self::Committed => committed.into(), - Self::Redirect(redirect) => ControlPlaneResponse::NotRaftLeader { - leader_addr: redirect, - } - .into(), - } - } -} - /// # Client ↔ Server request_id protocol /// /// 1. **Client assigns** — before sending each request, the client obtains a @@ -136,44 +118,67 @@ impl ClientController { name, storage_policy, } => self.handle_create_topic(name, storage_policy).await, - DeleteTopic { name } => self.handle_delete_topic(name).await, - ListHostedTopics => return self.handle_list_hosted_topics().await, - DescribeTopic { name } => self.handle_describe_topic(name).await, - SyncConsumerGroup(req) => self.handle_sync_consumer_group(req).await, + DeleteTopic { name } => self.delete_topic(name).await, + ListHostedTopics => self.list_hosted_topics().await, + DescribeTopic { name } => self.describe_topic(name).await, + SyncConsumerGroup(req) => self.sync_consumer_group(req).await, + OpenProducerSession(req) => self.open_producer_session(req).await, }; - match res { - Ok(resp) => resp, - Err(e) => { - tracing::error!("client control plane dispatch error: {e}"); - ClientResponse::ControlPlane(ControlPlaneResponse::InternalError(e.to_string())) - } - } + res.into() } - async fn handle_sync_consumer_group( + async fn open_producer_session( &self, - req: SyncConsumerGroupRequest, - ) -> anyhow::Result { - let group = match self.route(req.topic_name.as_bytes().to_vec()).await? { + req: OpenProducerSessionRequest, + ) -> Result { + let command: OpenProducerSession = req.into_command(); + + let group = match self.route(command.topic_name.as_bytes().to_vec()).await? { ShardRouting::Local(group) => group, ShardRouting::Redirect(member) => { - return Ok(self.control_plane_redirect(member).into()); + return Err(self.control_plane_redirect(member)); } }; - let proposal = self - .propose_topic_write(group.id, SyncConsumerGroup::new(req.clone()).into()) + self.propose_topic_write(group.id, command.clone()).await?; + + let topic_meta = self + .raft_sender + .get_topic_metadata(command.topic_name) .await?; - if let MetadataProposalOutcome::Redirect(redirect) = proposal { - return Ok(ControlPlaneResponse::NotRaftLeader { - leader_addr: redirect, + let session = topic_meta + .producer_sessions + .get(&command.producer_id) + .copied() + .ok_or_else(|| { + ServerError::Internal("committed producer session is unavailable".into()) + })?; + + Ok(ClientSuccess::ProducerSessionOpened( + ProducerSessionOpened { + incarnation: session.incarnation, + expires_at: session.expires_at, + }, + )) + } + + async fn sync_consumer_group( + &self, + req: SyncConsumerGroupRequest, + ) -> Result { + let group = match self.route(req.topic_name.as_bytes().to_vec()).await? { + ShardRouting::Local(group) => group, + ShardRouting::Redirect(member) => { + return Err(self.control_plane_redirect(member)); } - .into()); }; + self.propose_topic_write(group.id, SyncConsumerGroup::new(req.clone())) + .await?; + if req.action == ConsumerGroupSyncAction::Leave { - return Ok(ControlPlaneResponse::ConsumerGroupLeft.into()); + return Ok(ClientSuccess::ConsumerGroupLeft); } let Some(assignment) = self @@ -181,46 +186,42 @@ impl ClientController { .get_consumer_group_assignment(req.topic_name, req.group_id, req.member_id) .await else { - return Ok(ControlPlaneResponse::InternalError( + return Err(ServerError::Internal( "committed consumer group assignment is unavailable".into(), - ) - .into()); + )); }; - Ok( - ControlPlaneResponse::ConsumerGroupAssignment(ConsumerGroupAssignmentResponse { + Ok(ClientSuccess::ConsumerGroupAssignment( + ConsumerGroupAssignmentResponse { generation: assignment.generation, ranges: assignment.ranges, - }) - .into(), - ) + }, + )) } /// If this node belongs to the topic's owning shard group, answers directly; /// otherwise it redirects to a member so the consumer retries against the right /// node — no server-side proxying. - async fn handle_describe_topic(&self, topic_name: String) -> anyhow::Result { + async fn describe_topic(&self, topic_name: String) -> Result { if let ShardRouting::Redirect(member) = self.route(topic_name.as_bytes().to_vec()).await? { - return Ok(self.control_plane_redirect(member).into()); + return Err(self.control_plane_redirect(member)); } - let Some(meta) = self.raft_sender.get_topic_metadata(topic_name).await else { - return Ok(ControlPlaneResponse::TopicNotFound.into()); - }; + let topic = self.raft_sender.get_topic_metadata(topic_name).await?; let addresses = self.swim_sender.list_all_node_addresses().await?; - let detail = TopicDetail::from_meta(meta, &addresses); - Ok(ControlPlaneResponse::TopicDetail(detail).into()) + let detail = TopicDetail::from_meta(topic, &addresses); + Ok(ClientSuccess::TopicDetail(detail)) } async fn handle_create_topic( &self, name: String, storage_policy: StoragePolicy, - ) -> anyhow::Result { + ) -> Result { let group = match self.route(name.as_bytes().to_vec()).await? { ShardRouting::Local(group) => group, ShardRouting::Redirect(member) => { - return Ok(self.control_plane_redirect(member).into()); + return Err(self.control_plane_redirect(member)); } }; @@ -235,30 +236,26 @@ impl ClientController { replica_set: group.replicas, created_at, }; - Ok(self - .propose_topic_write(group.id, cmd.into()) - .await? - .on_ok(ControlPlaneResponse::TopicCreated)) + self.propose_topic_write(group.id, cmd).await?; + Ok(ClientSuccess::TopicCreated) } - async fn handle_delete_topic(&self, topic_name: String) -> anyhow::Result { + async fn delete_topic(&self, topic_name: String) -> Result { let group = match self.route(topic_name.as_bytes().to_vec()).await? { ShardRouting::Local(group) => group, ShardRouting::Redirect(member) => { - return Ok(self.control_plane_redirect(member).into()); + return Err(self.control_plane_redirect(member)); } }; let cmd = DeleteTopic { name: topic_name }; - Ok(self - .propose_topic_write(group.id, cmd.into()) - .await? - .on_ok(ControlPlaneResponse::TopicDeleted)) + self.propose_topic_write(group.id, cmd).await?; + Ok(ClientSuccess::TopicDeleted) } /// Route a key relative to this node — the single SWIM call behind every /// membership-first handler. - async fn route(&self, key: Vec) -> anyhow::Result { + async fn route(&self, key: Vec) -> Result { self.swim_sender .resolve_shard_routing(key, &self.node_id) .await @@ -266,38 +263,35 @@ impl ClientController { /// Structural redirect for a control-plane op that isn't local: to the member if /// one resolves, else a retriable error (no member's address known here yet). - fn control_plane_redirect(&self, member: Option) -> ControlPlaneResponse { + fn control_plane_redirect(&self, member: Option) -> ServerError { match member { - Some(node_addr) => ControlPlaneResponse::TopicMetadataRedirect { owner: node_addr }, - None => { - ControlPlaneResponse::InternalError("no reachable member of the shard group".into()) - } + Some(node_addr) => ServerError::TopicMetadataRedirect { owner: node_addr }, + None => ServerError::Internal("no reachable member of the shard group".into()), } } - /// Propose a control-plane write to the group's Raft leader. A committed - /// proposal is distinct from the retriable wire response returned when the - /// addressed node cannot currently accept it. + /// Propose a control-plane write to the group's Raft leader. + /// Returns `ServerError::NotRaftLeader` when the write must be retried at the leader. async fn propose_topic_write( &self, group_id: ShardGroupId, - cmd: MetadataCommand, - ) -> anyhow::Result { - let leader = match self + cmd: impl Into, + ) -> Result<(), ServerError> { + let Err(err) = self .raft_sender - .submit_metadata_command(group_id, cmd) + .submit_metadata_command(group_id, cmd.into()) .await - { - Ok(()) => { - return Ok(MetadataProposalOutcome::Committed); - } - Err(ProposalError::NotLeader(leader)) => leader, - Err(ProposalError::ShardNotFound | ProposalError::ShardGroupRemoved) => None, + else { + return Ok(()); }; - Ok(MetadataProposalOutcome::Redirect( - self.resolve_id_to_addr(leader).await, - )) + let leader = match err { + ProposalError::NotLeader(leader) => leader, + ProposalError::ShardNotFound | ProposalError::ShardGroupRemoved => None, + }; + + let leader_addr = self.resolve_id_to_addr(leader).await; + Err(ServerError::NotRaftLeader { leader_addr }) } async fn resolve_id_to_addr(&self, node_id: Option) -> Option { @@ -309,7 +303,7 @@ impl ClientController { .map(|addr| NodeAddressInfo { node_id, addr }) } - async fn handle_list_hosted_topics(&self) -> ClientResponse { + async fn list_hosted_topics(&self) -> Result { let topics = self .raft_sender .get_topics() @@ -321,60 +315,56 @@ impl ClientController { state: TopicState::Active, }) .collect(); - ClientResponse::ControlPlane(ControlPlaneResponse::TopicList { topics }) + Ok(ClientSuccess::TopicList { topics }) } async fn handle_data_plane(&self, request: ClientDataPlaneRequest) -> ClientResponse { use ClientDataPlaneRequest::*; - let res = match request { - Produce(req) => self.handle_produce(req).await, - Fetch(req) => self.handle_fetch(req).await, - FetchById(req) => self.handle_fetch_by_id(req).await, - ListOffsets(req) => self.handle_list_offsets(req).await, + let res: Result = match request { + Produce(req) => self.produce_entry(req).await, + Fetch(req) => self.fetch(req).await, + FetchById(req) => self.fetch_by_id(req).await, + ListOffsets(req) => self.list_offsets(req).await, CommitConsumerOffset(req) => self.handle_commit_consumer_offset(req).await, FetchConsumerOffset(req) => self.handle_fetch_consumer_offset(req).await, }; - match res { - Ok(resp) => resp, - Err(e) => { - tracing::error!("client data plane dispatch error: {e}"); - ClientResponse::DataPlane(DataPlaneResponse::InternalError(e.to_string())) - } - } + res.into() } async fn handle_commit_consumer_offset( &self, req: CommitConsumerOffsetRequest, - ) -> anyhow::Result { - let (reply, response) = tokio::sync::oneshot::channel(); + ) -> Result { + let (tx, recv) = tokio::sync::oneshot::channel(); self.data_plane_tx .send_async(CommitConsumerOffset { update: req.0, - reply, + reply: tx, }) .await?; - Ok(match response.await? { - ConsumerOffsetCommitAck::Committed => DataPlaneResponse::ConsumerOffsetCommitted.into(), - ConsumerOffsetCommitAck::NotWriteLeader(leader) => DataPlaneResponse::NotWriteLeader { + + let res = recv + .await + .map_err(|e| ServerError::Internal(e.to_string()))?; + + match res { + ConsumerOffsetCommitAck::Committed => Ok(ClientSuccess::ConsumerOffsetCommitted), + ConsumerOffsetCommitAck::NotWriteLeader(leader) => Err(ServerError::NotWriteLeader { leader_addr: self.resolve_id_to_addr(leader).await, - } - .into(), + }), ConsumerOffsetCommitAck::StaleEpoch(StaleEpoch(sealed_generation)) => { - DataPlaneResponse::StaleConsumerGroupEpoch(sealed_generation).into() - } - ConsumerOffsetCommitAck::InternalError(error) => { - DataPlaneResponse::InternalError(error).into() + Err(ServerError::StaleConsumerGroupEpoch(sealed_generation)) } - }) + ConsumerOffsetCommitAck::InternalError(error) => Err(ServerError::Internal(error)), + } } async fn handle_fetch_consumer_offset( &self, req: FetchConsumerOffsetRequest, - ) -> anyhow::Result { - let (reply, response) = tokio::sync::oneshot::channel(); + ) -> Result { + let (reply, recv) = tokio::sync::oneshot::channel(); self.data_plane_tx .send_async(ReadConsumerOffset { key: req.key, @@ -383,81 +373,77 @@ impl ClientController { }) .await?; - Ok(match response.await? { - ReadConsumerOffsetResult::Offset(offset) => { - DataPlaneResponse::ConsumerOffset(offset).into() - } + let res = recv + .await + .map_err(|e| ServerError::Internal(e.to_string()))?; + + match res { + ReadConsumerOffsetResult::Offset(offset) => Ok(ClientSuccess::ConsumerOffset(offset)), ReadConsumerOffsetResult::GenerationMismatch(mismatch) => { - DataPlaneResponse::ConsumerOffsetGenerationMismatch(mismatch).into() + Err(ServerError::ConsumerOffsetGenerationMismatch(mismatch)) } - ReadConsumerOffsetResult::NotLeader(leader) => DataPlaneResponse::NotWriteLeader { + ReadConsumerOffsetResult::NotLeader(leader) => Err(ServerError::NotWriteLeader { leader_addr: self.resolve_id_to_addr(leader).await, - } - .into(), - }) + }), + } } /// Routes a produce to the active segment's write leader (`replica_set[0]`), /// redirecting if this node isn't it. Membership-first (like `handle_describe_topic`) /// so we can tell "topic absent" from "not hosted here". - async fn handle_produce(&self, req: ProduceRequest) -> anyhow::Result { + async fn produce_entry(&self, req: ProduceRequest) -> Result { + let received_at_ms = crate::now_ms(); // Not local (ring unconverged or this node isn't a member) → retriable // redirect; the hint is best-effort, absent until SWIM converges. if let ShardRouting::Redirect(hint_node) = self.route(req.topic_name.as_bytes().to_vec()).await? { - return Ok(DataPlaneResponse::ShardNotLocal { hint_node }.into()); + return Err(ServerError::ShardNotLocal { hint_node }); } - let Some(meta) = self.raft_sender.get_topic_metadata(req.topic_name).await else { - return Ok(DataPlaneResponse::TopicNotFound.into()); - }; - - let Some(range) = meta.route_active_range(&req.routing_key) else { - return Ok(DataPlaneResponse::TopicNotFound.into()); - }; + let topic = self.raft_sender.get_topic_metadata(req.topic_name).await?; - if range.range_id != req.range_id { - return Ok(DataPlaneResponse::StaleRange.into()); - } + let producer_identity = req + .producer_identity + .map(|id| topic.verify_producer_session(id, received_at_ms)) + .transpose() + .map_err(ServerError::ProduceRejected)?; - // Invariant: an active range has exactly one active segment. - let Some(seg) = range.active_segment.and_then(|id| range.segments.get(&id)) else { - return Ok(DataPlaneResponse::InternalError( - "active range has no active segment".into(), - ) - .into()); - }; + let range = topic.resolve_writable_range(req.range_id, &req.routing_key)?; + let seg = range.active_write_segment()?; - if seg.replica_set.first() != Some(&self.node_id) { + if !seg.is_shard_leader(&self.node_id) { let leader_addr = self .resolve_id_to_addr(seg.replica_set.first().cloned()) .await; - return Ok(DataPlaneResponse::NotWriteLeader { leader_addr }.into()); + return Err(ServerError::NotWriteLeader { leader_addr }); } - let segment_key = SegmentKey::new(meta.id, range.range_id, seg.segment_id); + let segment_key = SegmentKey::new(topic.id, range.range_id, seg.segment_id); let (reply_tx, reply_rx) = tokio::sync::oneshot::channel(); + let cmd = Produce { segment_key, - data: EntryPayload::from(req.data), + data: req.data, record_count: req.record_count, + received_at_ms, + producer_identity, reply: reply_tx, }; let _ = self.data_plane_tx.send_async(cmd).await; let ack = reply_rx .await - .map_err(|_| anyhow::anyhow!("data plane dropped produce reply"))?; + .map_err(|_| ServerError::Internal("data plane dropped produce reply".into()))?; - Ok(match ack { - ProduceAck::Ok { entry_id } => DataPlaneResponse::Produced { entry_id }.into(), + match ack { + ProduceAck::Ok(entry_id) => Ok(ClientSuccess::Produced(entry_id)), // Metadata named us leader but the assignment hasn't applied yet — retriable. - ProduceAck::Err(reason) if reason == "not leader" || reason == "segment not found" => { - DataPlaneResponse::NotWriteLeader { leader_addr: None }.into() + ProduceAck::Err(ProduceError::NotLeader | ProduceError::SegmentNotFound) => { + Err(ServerError::NotWriteLeader { leader_addr: None }) } - ProduceAck::Err(reason) => DataPlaneResponse::InternalError(reason).into(), - }) + ProduceAck::Err(error) => Err(ServerError::ProduceRejected(error)), + } } /// Consume read path. @@ -467,21 +453,19 @@ impl ClientController { /// 4. dispatches a `DataPlaneQuery::Fetch` carrying everything the sync data-plane /// /// state machine needs to answer without any further I/O. - async fn handle_fetch(&self, req: FetchRequest) -> anyhow::Result { - let Some(meta) = self.raft_sender.get_topic_metadata(req.topic_name).await else { - return Ok(DataPlaneResponse::TopicNotFound.into()); - }; - let Some(range) = meta.ranges.get(&req.range_id) else { - return Ok(DataPlaneResponse::TopicNotFound.into()); - }; + async fn fetch(&self, req: FetchRequest) -> Result { + let topic = self.raft_sender.get_topic_metadata(req.topic_name).await?; + + let range = topic.get_range(&req.range_id)?; + if !keyspace_bound_matches_range(&req.keyspace_bound, range) { - return Ok(DataPlaneResponse::KeyspaceBoundNarrowed.into()); + return Err(ServerError::KeyspaceBoundNarrowed); } - let progress_signal = RangeProgressSignal::compute_progress_signal(range, &meta); + let progress_signal = RangeProgressSignal::compute_progress_signal(range, &topic); let (reply_tx, reply_rx) = tokio::sync::oneshot::channel(); let query = Fetch { - topic_id: meta.id, + topic_id: topic.id, range_id: req.range_id, entry_id: req.entry_id, max_bytes: req.max_bytes, @@ -489,13 +473,12 @@ impl ClientController { reply: reply_tx, }; - if self.data_plane_tx.send_async(query).await.is_err() { - return Ok(DataPlaneResponse::InternalError("data plane closed".into()).into()); - } - let result = reply_rx + self.data_plane_tx.send_async(query).await?; + let fetched_records = reply_rx .await - .map_err(|_| anyhow::anyhow!("data plane dropped reply"))?; - Ok(DataPlaneResponse::from_fetch_result(result).into()) + .map_err(|_| ServerError::Internal("data plane dropped reply".into()))??; + + Ok(EntryPayload::from_fetched_record(fetched_records)) } /// Consume read path addressed by resolved `topic_id` (the client resolved the @@ -503,7 +486,7 @@ impl ClientController { /// store, so a replica that holds the segment but isn't a metadata peer can /// serve it. No proxying: a miss returns `SegmentNotLocal` and the client /// retries another replica. - async fn handle_fetch_by_id(&self, req: FetchByIdRequest) -> anyhow::Result { + async fn fetch_by_id(&self, req: FetchByIdRequest) -> Result { let (reply_tx, reply_rx) = tokio::sync::oneshot::channel(); let query = Fetch { topic_id: req.topic_id, @@ -515,56 +498,51 @@ impl ClientController { progress_signal: RangeProgressSignal::Active, reply: reply_tx, }; - if self.data_plane_tx.send_async(query).await.is_err() { - return Ok(DataPlaneResponse::InternalError("data plane closed".into()).into()); - } - let result = reply_rx + self.data_plane_tx.send_async(query).await?; + + let fetched_records = reply_rx .await - .map_err(|_| anyhow::anyhow!("data plane dropped reply"))?; - Ok(DataPlaneResponse::from_fetch_result(result).into()) + .map_err(|_| ServerError::Internal("data plane dropped reply".into()))??; + Ok(EntryPayload::from_fetched_record(fetched_records)) } /// Offset-bounds query — used by consumers to bound their read window /// for the range's currently-active segment on this node. - async fn handle_list_offsets(&self, req: RangeOffsetRequest) -> anyhow::Result { - let Some(meta) = self.raft_sender.get_topic_metadata(req.topic_name).await else { - return Ok(DataPlaneResponse::TopicNotFound.into()); - }; + async fn list_offsets(&self, req: RangeOffsetRequest) -> Result { + let topic = self.raft_sender.get_topic_metadata(req.topic_name).await?; + let (reply_tx, reply_rx) = tokio::sync::oneshot::channel(); let query: DataPlaneQuery = ListOffsets { - topic_id: meta.id, + topic_id: topic.id, range_id: req.range_id, reply: reply_tx, } .into(); - let _ = self.data_plane_tx.send_async(query).await; + self.data_plane_tx.send_async(query).await?; - let list_offset_result = reply_rx.await.context("data plane dropped reply")?; - let result = DataPlaneResponse::from_list_offset_result(list_offset_result); - Ok(result.into()) + let range_offset = reply_rx + .await + .map_err(|_| ServerError::Internal("data plane dropped reply".into()))??; + + Ok(ClientSuccess::RangeOffset(range_offset)) } async fn handle_admin(&self, request: AdminRequest) -> ClientResponse { + use AdminRequest::*; + let res = match request { - AdminRequest::DescribeCluster => self.handle_describe_cluster().await, - AdminRequest::ListHostedTopicsWithStats => { - self.handle_list_hosted_topics_with_stats().await - } + DescribeCluster => self.describe_cluster().await, + ListHostedTopicsWithStats => self.list_hosted_topics_with_stats().await, - AdminRequest::GetShardInfo { key } => self.handle_get_shard_info(key).await, - AdminRequest::GetShardLeader { shard_group_id } => { - self.handle_get_shard_leader(shard_group_id).await - } + GetShardInfo { key } => self.get_shard_info(key).await, + GetShardLeader { shard_group_id } => self.handle_get_shard_leader(shard_group_id).await, }; - res.unwrap_or_else(|e| { - tracing::error!("client admin dispatch error: {e}"); - ClientResponse::Admin(AdminResponse::InternalError(e.to_string())) - }) + res.into() } - async fn handle_describe_cluster(&self) -> anyhow::Result { + async fn describe_cluster(&self) -> Result { let members = self.swim_sender.get_members().await?; let nodes = members .into_iter() @@ -578,10 +556,10 @@ impl ClientController { }, }) .collect(); - Ok(ClientResponse::Admin(AdminResponse::ClusterInfo { nodes })) + Ok(ClientSuccess::ClusterInfo { nodes }) } - async fn handle_list_hosted_topics_with_stats(&self) -> anyhow::Result { + async fn list_hosted_topics_with_stats(&self) -> Result { let topics = self .raft_sender .get_topic_stats() @@ -593,10 +571,10 @@ impl ClientController { total_bytes: s.total_bytes, }) .collect(); - Ok(ClientResponse::Admin(AdminResponse::TopicStats { topics })) + Ok(ClientSuccess::TopicStats { topics }) } - async fn handle_get_shard_info(&self, key: Vec) -> anyhow::Result { + async fn get_shard_info(&self, key: Vec) -> Result { let detail = self .swim_sender .get_shard_info(key) @@ -607,19 +585,19 @@ impl ClientController { leader_addr: leader.map(|e| e.leader.client_addr()), member_node_ids: group.replicas.iter().map(|n| n.to_string()).collect(), }); - Ok(ClientResponse::Admin(AdminResponse::ShardInfo { detail })) + Ok(ClientSuccess::ShardInfo { detail }) } async fn handle_get_shard_leader( &self, shard_group_id: ShardGroupId, - ) -> anyhow::Result { + ) -> Result { let leader = self .raft_sender .get_leader(shard_group_id) .await .map(|n| n.to_string()); - Ok(ClientResponse::Admin(AdminResponse::ShardLeader { leader })) + Ok(ClientSuccess::ShardLeader { leader }) } } @@ -659,8 +637,8 @@ pub async fn handle_client_stream( mod tests { use super::*; use crate::connections::protocol::{ - AdminRequest, AdminResponse, ClientDataPlaneRequest, ClientRequest, ClientResponse, - ControlPlaneRequest, ControlPlaneResponse, DataPlaneResponse, NodeState, ProduceRequest, + AdminRequest, ClientDataPlaneRequest, ClientRequest, ClientResponse, ClientSuccess, + ControlPlaneRequest, NodeState, ProduceRequest, ServerError, }; use crate::control_plane::consensus::actor::MultiRaftActor; use crate::control_plane::consensus::messages::MultiRaftActorCommand; @@ -730,7 +708,7 @@ mod tests { tokio::spawn(async move { while let Ok(msg) = rx.recv_async().await { if let DataPlaneMessage::Command(DataPlaneCommand::Produce(p)) = msg { - let _ = p.reply.send(ProduceAck::Ok { entry_id: 7.into() }); + let _ = p.reply.send(ProduceAck::Ok(7.into())); } } }); @@ -742,8 +720,9 @@ mod tests { topic_name: "t1".into(), range_id: RangeId(0), routing_key: b"k".to_vec(), - data: b"hello".to_vec(), + data: b"hello".to_vec().into(), record_count: 1, + producer_identity: None, })) } @@ -778,7 +757,7 @@ mod tests { assert!( matches!( resp, - ClientResponse::DataPlane(DataPlaneResponse::ShardNotLocal { hint_node: None }) + ClientResponse::Err(ServerError::ShardNotLocal { hint_node: None }) ), "expected ShardNotLocal{{None}}, got {resp:?}" ); @@ -807,7 +786,7 @@ mod tests { ClientController::new(node_id("self"), swim, raft_sender_with(|_| {}), dp_stub()) .dispatch(produce_req()) .await; - let ClientResponse::DataPlane(DataPlaneResponse::ShardNotLocal { hint_node }) = resp else { + let ClientResponse::Err(ServerError::ShardNotLocal { hint_node }) = resp else { panic!("expected ShardNotLocal, got {resp:?}"); }; assert_eq!(hint_node.unwrap().client_addr(), addr(8083)); @@ -841,7 +820,7 @@ mod tests { let resp = ClientController::new(me, swim, raft, dp_stub()) .dispatch(produce_req()) .await; - let ClientResponse::DataPlane(DataPlaneResponse::NotWriteLeader { + let ClientResponse::Err(ServerError::NotWriteLeader { leader_addr: Some(a), }) = resp else { @@ -871,7 +850,7 @@ mod tests { let resp = ClientController::new(me, swim, raft, dp_acking()) .dispatch(produce_req()) .await; - let ClientResponse::DataPlane(DataPlaneResponse::Produced { entry_id }) = resp else { + let ClientResponse::Ok(ClientSuccess::Produced(entry_id)) = resp else { panic!("expected Produced, got {resp:?}"); }; assert_eq!(entry_id, 7.into()); @@ -909,7 +888,7 @@ mod tests { }, )) .await; - let ClientResponse::ControlPlane(ControlPlaneResponse::TopicMetadataRedirect { + let ClientResponse::Err(ServerError::TopicMetadataRedirect { owner: redirect_owner, }) = resp else { @@ -944,10 +923,7 @@ mod tests { )) .await; assert!( - matches!( - resp, - ClientResponse::ControlPlane(ControlPlaneResponse::TopicCreated) - ), + matches!(resp, ClientResponse::Ok(ClientSuccess::TopicCreated)), "expected TopicCreated, got {resp:?}" ); } @@ -973,10 +949,7 @@ mod tests { )) .await; assert!( - matches!( - resp, - ClientResponse::ControlPlane(ControlPlaneResponse::InternalError(_)) - ), + matches!(resp, ClientResponse::Err(ServerError::Internal(_))), "expected InternalError, got {resp:?}" ); } @@ -999,10 +972,7 @@ mod tests { )) .await; assert!( - matches!( - resp, - ClientResponse::ControlPlane(ControlPlaneResponse::TopicDeleted) - ), + matches!(resp, ClientResponse::Ok(ClientSuccess::TopicDeleted)), "expected TopicDeleted, got {resp:?}" ); } @@ -1020,7 +990,7 @@ mod tests { ControlPlaneRequest::ListHostedTopics, )) .await; - let ClientResponse::ControlPlane(ControlPlaneResponse::TopicList { topics }) = resp else { + let ClientResponse::Ok(ClientSuccess::TopicList { topics }) = resp else { panic!("expected TopicList, got {resp:?}"); }; assert_eq!(topics.len(), 2); @@ -1060,7 +1030,7 @@ mod tests { }, )) .await; - let ClientResponse::ControlPlane(ControlPlaneResponse::TopicMetadataRedirect { + let ClientResponse::Err(ServerError::TopicMetadataRedirect { owner: redirect_owner, }) = resp else { @@ -1098,10 +1068,7 @@ mod tests { )) .await; assert!( - matches!( - resp, - ClientResponse::ControlPlane(ControlPlaneResponse::TopicNotFound) - ), + matches!(resp, ClientResponse::Err(ServerError::TopicNotFound)), "expected TopicNotFound, got {resp:?}" ); } @@ -1135,7 +1102,7 @@ mod tests { ClientController::new(node_id("self"), swim, raft_sender_with(|_| {}), dp_stub()) .dispatch(ClientRequest::Admin(AdminRequest::DescribeCluster)) .await; - let ClientResponse::Admin(AdminResponse::ClusterInfo { nodes: info }) = resp else { + let ClientResponse::Ok(ClientSuccess::ClusterInfo { nodes: info }) = resp else { panic!("expected ClusterInfo, got {resp:?}"); }; assert_eq!(info.len(), 2); @@ -1167,7 +1134,7 @@ mod tests { key: b"any".to_vec(), })) .await; - let ClientResponse::Admin(AdminResponse::ShardInfo { detail: Some(d) }) = resp else { + let ClientResponse::Ok(ClientSuccess::ShardInfo { detail: Some(d) }) = resp else { panic!("expected ShardInfo with detail, got {resp:?}"); }; assert_eq!(*d.shard_group_id, 42); @@ -1193,7 +1160,7 @@ mod tests { assert!( matches!( resp, - ClientResponse::Admin(AdminResponse::ShardInfo { detail: None }) + ClientResponse::Ok(ClientSuccess::ShardInfo { detail: None }) ), "expected ShardInfo{{None}}, got {resp:?}" ); @@ -1212,7 +1179,7 @@ mod tests { shard_group_id: ShardGroupId(42), })) .await; - let ClientResponse::Admin(AdminResponse::ShardLeader { leader }) = resp else { + let ClientResponse::Ok(ClientSuccess::ShardLeader { leader }) = resp else { panic!("expected ShardLeader, got {resp:?}"); }; assert_eq!(leader.as_deref(), Some("n1")); @@ -1235,7 +1202,7 @@ mod tests { AdminRequest::ListHostedTopicsWithStats, )) .await; - let ClientResponse::Admin(AdminResponse::TopicStats { topics }) = resp else { + let ClientResponse::Ok(ClientSuccess::TopicStats { topics }) = resp else { panic!("expected TopicStats, got {resp:?}"); }; assert_eq!(topics.len(), 1); diff --git a/src/connections/protocol/mod.rs b/src/connections/protocol/mod.rs index 611a2d1e..9e8b2c56 100644 --- a/src/connections/protocol/mod.rs +++ b/src/connections/protocol/mod.rs @@ -68,9 +68,7 @@ pub enum ClientSuccess { ProducerSessionOpened(ProducerSessionOpened), // Data Plane - Produced { - entry_id: EntryId, - }, + Produced(EntryId), Fetched { entries: Box<[EntryPayload]>, next_entry_id: EntryId, diff --git a/src/it/sim/scenario.rs b/src/it/sim/scenario.rs index dfa91fe4..4a3adf09 100644 --- a/src/it/sim/scenario.rs +++ b/src/it/sim/scenario.rs @@ -405,7 +405,7 @@ async fn produce_until_acked( loop { for (host, port) in nodes { match send_request(host, *port, &request).await { - Ok(ClientResponse::Ok(ClientSuccess::Produced { entry_id })) => { + Ok(ClientResponse::Ok(ClientSuccess::Produced(entry_id))) => { return entry_id; } Ok(_) => {} From 5e524620cbfbff7a530ac977e736ceccebcfe76d Mon Sep 17 00:00:00 2001 From: Migorithm Date: Thu, 23 Jul 2026 21:05:31 +0400 Subject: [PATCH 38/41] feat(client): implement producer session manager and per-range sequence tracking - Add ClientProducerSessionManager and RangeSequences in session.rs for per-range sequence generation and topology tracking. - Retain in-flight sequence counters during range retirement and handle session lease renewals safely. - Consolidate transparent retry logic in producer flush for SessionExpired and Request InFlight errors. - Add end-to-end SDK tests for producer deduplication. --- .../d10_transparent_retry_deduplication.md | 6 +- src/client/producer/buffers.rs | 6 + src/client/producer/mod.rs | 239 +++++++------ src/client/producer/session.rs | 319 ++++++++++++++++++ src/client/routing.rs | 4 + src/control_plane/metadata/topic.rs | 6 +- .../auxiliary_states/producer/state.rs | 6 +- .../auxiliary_states/producer/types.rs | 2 + src/it/e2e/client_sdk.rs | 57 ++++ src/it/e2e/producer_deduplication.rs | 4 +- 10 files changed, 542 insertions(+), 107 deletions(-) create mode 100644 src/client/producer/session.rs diff --git a/docs/data-plane/d10_transparent_retry_deduplication.md b/docs/data-plane/d10_transparent_retry_deduplication.md index 0d41e5f8..bb5ab20e 100644 --- a/docs/data-plane/d10_transparent_retry_deduplication.md +++ b/docs/data-plane/d10_transparent_retry_deduplication.md @@ -114,8 +114,10 @@ batches per range. ```text request(session, incarnation, range, sequence, digest) │ -├── session unknown or expired +├── lease timestamp expired │ └── SessionExpired +├── unexpired session not installed on this replica +│ └── SessionNotInstalled ├── incarnation older than installed fence │ └── ProducerFenced ├── dedup state or placement still recovering @@ -279,6 +281,8 @@ The implemented delivery contract is: longer retained. - Reusing a sequence with a different payload checksum, skipping a sequence, using an expired session, or writing with a fenced incarnation is rejected without appending. + An unexpired session that has not reached the serving replica is distinguished from + expiration so the producer can retry the unchanged identity while metadata converges. - A restarted producer using the same producer id opens with a new session nonce. Metadata advances its incarnation and fences the older process. Retrying the open itself is idempotent. diff --git a/src/client/producer/buffers.rs b/src/client/producer/buffers.rs index 3b4b4ae3..2f7380f6 100644 --- a/src/client/producer/buffers.rs +++ b/src/client/producer/buffers.rs @@ -54,6 +54,12 @@ pub struct PendingRecord { pub tx: oneshot::Sender>, } impl PendingRecord { + pub fn complete_all(records: Vec, result: Result) { + for pending in records { + let _ = pending.tx.send(result.clone()); + } + } + /// Serialize a slice of records into a byte buffer. pub fn serialize_batch(records: &[PendingRecord]) -> Vec { let mut buf = Vec::new(); diff --git a/src/client/producer/mod.rs b/src/client/producer/mod.rs index 21604acc..471979e2 100644 --- a/src/client/producer/mod.rs +++ b/src/client/producer/mod.rs @@ -1,12 +1,16 @@ pub(crate) mod buffers; mod config; pub(crate) mod record; +mod session; use crate::client::error::ClientError; +use crate::client::routing::TopicRouting; use crate::client::{Client, CompressionCodec}; use crate::control_plane::metadata::{EntryId, RangeId}; +use crate::data_plane::ProduceError; use crate::impl_new_struct_wrapper; use buffers::{PendingRecord, ProducerBuffers, PushResult}; pub use config::{BufferConfig, ProducerConfig}; +use session::{ClientProducerSession, ClientProducerSessionManager}; use std::collections::BTreeMap; use std::ops::Deref; @@ -27,9 +31,8 @@ pub struct Inner { topic: String, buffers: ProducerBuffers, codec: CompressionCodec, - // Idempotency seam: globally unique session ID and monotonic sequence counter - producer_id: Uuid, - next_sequence: std::sync::atomic::AtomicU32, + session_manager: ClientProducerSessionManager, + next_record_order: std::sync::atomic::AtomicU64, } impl Producer { @@ -46,38 +49,33 @@ impl Producer { config: ProducerConfig, producer_id: Uuid, ) -> Self { - // Generate a globally unique UUID for the producer session ID (idempotency seam) - Self(Arc::new(Inner { client, topic, buffers: ProducerBuffers::new(config.buffer.clone()), codec: config.codec, - producer_id, - next_sequence: std::sync::atomic::AtomicU32::new(0), + session_manager: ClientProducerSessionManager::new(producer_id), + next_record_order: std::sync::atomic::AtomicU64::new(0), })) } /// Produce a single record. Returns the committed entry ID once the batch flushes. pub async fn send(&self, key: &[u8], value: Vec) -> Result { + let order = self + .next_record_order + .fetch_add(1, std::sync::atomic::Ordering::Relaxed); let topic = &self.topic; let client = &self.client; let routing = client.resolve_topic_if_missing(topic).await?; let range_id = routing.range_id(key).ok_or(ClientError::TopicNotFound)?; - // Idempotency seam: allocate sequence number - let sequence_number = self - .next_sequence - .fetch_add(1, std::sync::atomic::Ordering::Relaxed); - let (tx, rx) = oneshot::channel(); let pending = PendingRecord { + order, key: key.to_vec(), value, - producer_id: self.producer_id, - sequence_number, tx, }; @@ -121,115 +119,154 @@ impl Producer { /// Serialize, compress, and publish a batch of records. async fn flush_records(&self, mut records_to_publish: Vec) { + records_to_publish.sort_unstable_by_key(|record| record.order); let deadline = Instant::now() + self.client.retry.deadline; let mut backoff = self.client.retry.initial_backoff; - let timeout_error = ClientError::Timeout { - waited: self.client.retry.deadline, - last_error: Some("topic routing remained stale".to_string()), - }; loop { if records_to_publish.is_empty() { return; } - let resolve_remaining = deadline.saturating_duration_since(Instant::now()); - if resolve_remaining.is_zero() { - for pending in records_to_publish { - let _ = pending.tx.send(Err(timeout_error.clone())); - } - return; - } - - let resolve_result = tokio::time::timeout( - resolve_remaining, - self.client.resolve_topic_if_missing(&self.topic), - ) - .await; - let routing = match resolve_result { - Ok(Ok(routing)) => routing, - Ok(Err(error)) => { - for pending in records_to_publish { - let _ = pending.tx.send(Err(error.clone())); - } + let routing = match self.resolve_routing(deadline).await { + Ok(routing) => routing, + Err(error) => { + PendingRecord::complete_all(records_to_publish, Err(error)); return; } - Err(_) => { - for pending in records_to_publish { - let _ = pending.tx.send(Err(timeout_error.clone())); - } + }; + + let session = match self.session_manager.ensure(&self.client, &self.topic).await { + Ok(session) => session, + Err(error) => { + PendingRecord::complete_all(records_to_publish, Err(error)); return; } }; - // Group in range-ID order while preserving record order within each range. - let mut recs_per_range: BTreeMap> = BTreeMap::new(); - - for pending in records_to_publish { - let Some(range_id) = routing.range_id(&pending.key) else { - let _ = pending.tx.send(Err(ClientError::UnexpectedResponse)); - continue; - }; - recs_per_range.entry(range_id).or_default().push(pending); - } + self.session_manager + .observe_topology(session, routing.active_range_ids()); let mut next_retry_records = Vec::new(); - - for (range_id, range_recs) in recs_per_range { - let routing_key = &range_recs[0].key; - - // Encode and compress the payload - let Ok(payload) = self.codec.encode_payload(&range_recs) else { - tracing::error!("Compression error while flushing records"); - for pending in range_recs { - let _ = pending.tx.send(Err(ClientError::UnexpectedResponse)); - } - continue; - }; - - let produce_result = match tokio::time::timeout( - deadline.saturating_duration_since(Instant::now()), - self.client.produce_to_range( - &self.topic, - range_id, - routing_key, - payload, - range_recs.len() as u32, - ), - ) - .await + for (range_id, records) in Self::group_by_range(records_to_publish, &routing) { + if let Some(records) = self + .publish_range_attempt(session, range_id, records, deadline) + .await { - Ok(result) => result, - Err(_) => Err(timeout_error.clone()), - }; - - match produce_result { - Ok(entry_id) => { - // Success! Deliver to all awaiting senders of this sub-batch - for pending in range_recs { - let _ = pending.tx.send(Ok(entry_id)); - } - } - Err(ClientError::StaleRange) => { - // Invalidate routing cache and collect records to retry in the next loop iteration - self.client.cache.invalidate(&self.topic); - next_retry_records.extend(range_recs); - } - Err(err) => { - // Terminal error. Fail this sub-batch. - for pending in range_recs { - let _ = pending.tx.send(Err(err.clone())); - } - } + next_retry_records.extend(records); } } - if !next_retry_records.is_empty() { - let backoff_remaining = deadline.saturating_duration_since(Instant::now()); - tokio::time::sleep(backoff.min(backoff_remaining)).await; - backoff = (backoff * 2).min(self.client.retry.max_backoff); + if next_retry_records.is_empty() { + return; } + + let backoff_remaining = deadline.saturating_duration_since(Instant::now()); + tokio::time::sleep(backoff.min(backoff_remaining)).await; + backoff = (backoff * 2).min(self.client.retry.max_backoff); records_to_publish = next_retry_records; } } + + async fn resolve_routing(&self, deadline: Instant) -> Result, ClientError> { + let remaining = deadline.saturating_duration_since(Instant::now()); + if remaining.is_zero() { + return Err(self.flush_timeout()); + } + + match tokio::time::timeout(remaining, self.client.resolve_topic_if_missing(&self.topic)) + .await + { + Ok(result) => result, + Err(_) => Err(self.flush_timeout()), + } + } + + fn group_by_range( + records: Vec, + routing: &TopicRouting, + ) -> BTreeMap> { + let mut grouped: BTreeMap> = BTreeMap::new(); + for pending in records { + match routing.range_id(&pending.key) { + Some(range_id) => grouped.entry(range_id).or_default().push(pending), + None => { + let _ = pending.tx.send(Err(ClientError::UnexpectedResponse)); + } + } + } + grouped + } + + /// Own one range batch for one network attempt. + /// + /// Returns the records only when the caller must retry them. Every other + /// outcome completes their senders before returning `None`. + async fn publish_range_attempt( + &self, + session: ClientProducerSession, + range_id: RangeId, + records: Vec, + deadline: Instant, + ) -> Option> { + let sequence = self.session_manager.sequence_for(session, range_id); + // Contiguous per-range sequences require serialization through durability. + let mut sequence = sequence.lock().await; + + let payload = match self.codec.encode_payload(&records) { + Ok(payload) => payload, + Err(error) => { + tracing::error!(?error, "failed to encode producer batch"); + PendingRecord::complete_all(records, Err(ClientError::UnexpectedResponse)); + return None; + } + }; + let digest = crc32fast::hash(&payload); + let result = match tokio::time::timeout( + deadline.saturating_duration_since(Instant::now()), + self.client.produce_to_range( + &self.topic, + range_id, + &records[0].key, + payload, + records.len() as u32, + Some(session.append_identity(*sequence, digest)), + ), + ) + .await + { + Ok(result) => result, + Err(_) => Err(self.flush_timeout()), + }; + + match result { + Ok(entry_id) => { + *sequence += 1; + PendingRecord::complete_all(records, Ok(entry_id)); + None + } + Err(ClientError::StaleRange) => { + self.client.cache.invalidate(&self.topic); + Some(records) + } + Err(ClientError::ProduceRejected( + ProduceError::SessionNotInstalled | ProduceError::RequestInFlight, + )) => Some(records), + Err(ClientError::ProduceRejected(ProduceError::SessionExpired)) => { + self.session_manager.mark_expired(session).await; + Some(records) + } + Err(error) => { + PendingRecord::complete_all(records, Err(error)); + None + } + } + } + + fn flush_timeout(&self) -> ClientError { + ClientError::Timeout { + waited: self.client.retry.deadline, + last_error: Some("producer flush deadline elapsed".to_string()), + } + } } diff --git a/src/client/producer/session.rs b/src/client/producer/session.rs new file mode 100644 index 00000000..ac121574 --- /dev/null +++ b/src/client/producer/session.rs @@ -0,0 +1,319 @@ +use crate::client::Client; +use crate::client::error::ClientError; +use crate::connections::protocol::{OpenProducerSessionRequest, ProducerSessionOpened}; +use crate::control_plane::metadata::RangeId; +use crate::data_plane::ProducerAppendIdentity; + +use std::collections::{HashMap, HashSet}; +use std::sync::Arc; +use tokio::sync::Mutex; +use uuid::Uuid; + +const SESSION_RENEWAL_MARGIN_MS: u64 = 1_000; + +pub(super) struct ClientProducerSessionManager { + session: Mutex, + range_sequences: std::sync::Mutex, +} + +#[derive(Clone, Copy)] +pub(super) struct ClientProducerSession { + producer_id: Uuid, + nonce: Uuid, + generation: u64, + opened: Option, +} +impl ClientProducerSession { + fn is_valid(&self, now: u64) -> bool { + if let Some(opened) = self.opened + && opened.expires_at > now.saturating_add(SESSION_RENEWAL_MARGIN_MS) + { + return true; + } + false + } + + fn expired(&self, now: u64) -> bool { + self.opened.is_some_and(|opened| opened.expires_at <= now) + } + + fn provision_open_session_req(&mut self, now: u64, topic: &str) -> OpenProducerSessionRequest { + if self.expired(now) { + self.producer_id = Uuid::new_v4(); + self.nonce = Uuid::new_v4(); + self.generation = self.generation.wrapping_add(1); + self.opened = None; + } + + let (producer_id, nonce) = (self.producer_id, self.nonce); + OpenProducerSessionRequest { + topic_name: topic.to_string(), + producer_id, + session_nonce: nonce, + } + } +} + +#[derive(Default)] +struct RangeSequences { + sequences: HashMap<(ProducerSessionKey, RangeId), Arc>>, + topologies: HashMap, + /// Highest session generation observed so far (prevents stale generation updates). + latest_generation: u64, +} + +/// Key uniquely identifying a producer session generation. +/// +/// Generation increments on session rotation (e.g. lease expiration), +/// ensuring new session sequence counters start at 0 without colliding with older sessions. +#[derive(Clone, Copy, PartialEq, Eq, Hash)] +struct ProducerSessionKey { + producer_id: Uuid, + generation: u64, +} + +/// Active range topology for a producer session. +struct ProducerTopology { + /// Since range IDs increase monotonically as ranges split over time, + /// a higher range ID implies a newer topology snapshot, allowing out-of-order + /// metadata updates to be ignored. + highest_range_id: RangeId, + /// Currently active ranges accepting appends for the topic. + active_ranges: HashSet, +} + +impl ClientProducerSessionManager { + pub(super) fn new(producer_id: Uuid) -> Self { + Self { + session: Mutex::new(ClientProducerSession { + producer_id, + nonce: Uuid::new_v4(), + generation: 0, + opened: None, + }), + range_sequences: std::sync::Mutex::new(RangeSequences::default()), + } + } + + pub(super) async fn ensure( + &self, + client: &Client, + topic: &str, + ) -> Result { + let mut current = self.session.lock().await; + let now = crate::now_ms(); + + if current.is_valid(now) { + return Ok(*current); + } + let req = current.provision_open_session_req(now, topic); + let opened = client.open_producer_session(req).await?; + + current.opened = Some(opened); + Ok(*current) + } + + pub(super) fn observe_topology( + &self, + session: ClientProducerSession, + active_ranges: impl Iterator, + ) { + self.range_sequences + .lock() + .expect("producer sequence registry poisoned") + .observe_topology(session.key(), active_ranges.collect()); + } + + pub(super) fn sequence_for( + &self, + session: ClientProducerSession, + range_id: RangeId, + ) -> Arc> { + self.range_sequences + .lock() + .expect("producer sequence registry poisoned") + .sequence_for(session.key(), range_id) + } + + pub(super) async fn mark_expired(&self, session: ClientProducerSession) { + let mut current = self.session.lock().await; + if current.key() == session.key() + && let Some(opened) = current.opened.as_mut() + { + opened.expires_at = 0; + } + } +} + +impl ClientProducerSession { + fn key(self) -> ProducerSessionKey { + ProducerSessionKey { + producer_id: self.producer_id, + generation: self.generation, + } + } + + pub(super) fn append_identity(self, sequence: u64, digest: u32) -> ProducerAppendIdentity { + let opened = self + .opened + .expect("ensured producer session must be opened"); + ProducerAppendIdentity { + producer_id: self.producer_id, + incarnation: opened.incarnation, + expires_at: opened.expires_at, + sequence, + digest, + } + } +} + +impl RangeSequences { + /// Update active range topology for a producer session generation and garbage collect retired sequences. + fn observe_topology(&mut self, session: ProducerSessionKey, active_ranges: HashSet) { + // Ignore stale topology updates from an older session generation + if session.generation < self.latest_generation { + return; + } + self.latest_generation = session.generation; + + let Some(highest_range_id) = active_ranges.iter().max().copied() else { + return; + }; + + // Monotonically advance topology: only update if high_watermark >= current high_watermark + match self.topologies.get_mut(&session) { + Some(topology) if highest_range_id >= topology.highest_range_id => { + topology.highest_range_id = highest_range_id; + topology.active_ranges = active_ranges; + } + Some(_) => {} // Ignore out-of-order stale topology observation + None => { + self.topologies.insert( + session, + ProducerTopology { + highest_range_id, + active_ranges, + }, + ); + } + } + + // Garbage collection: Keep sequence counters if: + // 1. The range is active in the current session's topology, OR + // 2. An in-flight flush task still holds a reference handle to it (Arc::strong_count > 1). + self.sequences.retain(|(key, range_id), sequence| { + (*key == session && self.topologies[&session].active_ranges.contains(range_id)) + || Arc::strong_count(sequence) > 1 + }); + + // Prune old topologies that no longer have active sequences and aren't the current session + self.topologies.retain(|key, _| { + *key == session + || self + .sequences + .keys() + .any(|(sequence_key, _)| sequence_key == key) + }); + } + + /// Retrieve or create the sequence counter handle for a (session, range_id) pair. + fn sequence_for(&mut self, session: ProducerSessionKey, range_id: RangeId) -> Arc> { + self.sequences + .entry((session, range_id)) + .or_insert_with(|| Arc::new(Mutex::new(0))) + .clone() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn ranges(ids: &[u64]) -> HashSet { + ids.iter().copied().map(RangeId).collect() + } + + fn session(generation: u64) -> ProducerSessionKey { + ProducerSessionKey { + producer_id: Uuid::new_v4(), + generation, + } + } + + #[test] + fn stale_topology_does_not_prune_newer_range_sequences() { + let mut sequences = RangeSequences::default(); + let session = session(0); + sequences.observe_topology(session, ranges(&[0])); + sequences.observe_topology(session, ranges(&[1, 2])); + sequences.sequence_for(session, RangeId(1)); + + sequences.observe_topology(session, ranges(&[0])); + + assert!(sequences.sequences.contains_key(&(session, RangeId(1)))); + assert_eq!(sequences.topologies[&session].highest_range_id, RangeId(2)); + } + + #[test] + fn retired_range_sequence_is_kept_while_a_flush_owns_it() { + let mut sequences = RangeSequences::default(); + let session = session(0); + sequences.observe_topology(session, ranges(&[0])); + let in_flight = sequences.sequence_for(session, RangeId(0)); + + sequences.observe_topology(session, ranges(&[1, 2])); + assert!(sequences.sequences.contains_key(&(session, RangeId(0)))); + + drop(in_flight); + sequences.observe_topology(session, ranges(&[1, 2])); + assert!(!sequences.sequences.contains_key(&(session, RangeId(0)))); + } + + #[test] + fn rotated_identity_never_reuses_an_old_in_flight_counter() { + let mut sequences = RangeSequences::default(); + let old_session = session(0); + let new_session = session(1); + sequences.observe_topology(old_session, ranges(&[0])); + let old_in_flight = sequences.sequence_for(old_session, RangeId(0)); + + sequences.observe_topology(new_session, ranges(&[0])); + let new_sequence = sequences.sequence_for(new_session, RangeId(0)); + + assert!(!Arc::ptr_eq(&old_in_flight, &new_sequence)); + } + + #[test] + fn old_session_cleanup_cannot_prune_new_session_sequences() { + let mut sequences = RangeSequences::default(); + let old_session = session(0); + let new_session = session(1); + sequences.observe_topology(new_session, ranges(&[1])); + sequences.sequence_for(new_session, RangeId(1)); + + sequences.observe_topology(old_session, ranges(&[0])); + + assert!(sequences.sequences.contains_key(&(new_session, RangeId(1)))); + } + + #[test] + fn failed_expired_session_open_retries_the_same_identity() { + let mut session = ClientProducerSession { + producer_id: Uuid::new_v4(), + nonce: Uuid::new_v4(), + generation: 0, + opened: Some(ProducerSessionOpened { + incarnation: 0, + expires_at: 10, + }), + }; + + let first = session.provision_open_session_req(11, "topic"); + let retry = session.provision_open_session_req(12, "topic"); + + assert_eq!(first.producer_id, retry.producer_id); + assert_eq!(first.session_nonce, retry.session_nonce); + assert_eq!(session.generation, 1); + assert!(session.opened.is_none()); + } +} diff --git a/src/client/routing.rs b/src/client/routing.rs index 79331fce..44ef88ac 100644 --- a/src/client/routing.rs +++ b/src/client/routing.rs @@ -86,6 +86,10 @@ impl TopicRouting { .find(|range| range.range_id == range_id) .and_then(|range| range.replica_addrs.first().cloned()) } + + pub(crate) fn active_range_ids(&self) -> impl Iterator + '_ { + self.ranges.iter().map(|r| r.range_id) + } } /// Thread-safe, lock-free per-topic cache. Entries are replaced on refresh and diff --git a/src/control_plane/metadata/topic.rs b/src/control_plane/metadata/topic.rs index 87f6f82f..7b30dce2 100644 --- a/src/control_plane/metadata/topic.rs +++ b/src/control_plane/metadata/topic.rs @@ -537,6 +537,10 @@ impl TopicMeta { q: ProducerAppendIdentity, received_at_ms: u64, ) -> Result { + if q.expires_at < received_at_ms { + return Err(ProduceError::SessionExpired); + } + let Some(session) = self.producer_sessions.get(&q.producer_id) else { // Data replica placement is independent from metadata-Raft placement. // A restarted data leader can therefore serve before its local metadata @@ -557,7 +561,7 @@ impl TopicMeta { return Ok(AuthorizedProducerIdentity::MetadataVerified(q)); } - Err(ProduceError::SessionExpired) + Err(ProduceError::SessionNotInstalled) } } diff --git a/src/data_plane/auxiliary_states/producer/state.rs b/src/data_plane/auxiliary_states/producer/state.rs index c41b45f0..1ee3422b 100644 --- a/src/data_plane/auxiliary_states/producer/state.rs +++ b/src/data_plane/auxiliary_states/producer/state.rs @@ -54,7 +54,7 @@ impl ProducerSession { } if identity.incarnation > self.incarnation { if !allow_session_creation { - return Err(ProduceError::SessionExpired); + return Err(ProduceError::SessionNotInstalled); } // * frontier and dedup_entries must be reset rather than reserved because: // * dedup_entries stores entries indexed by sequence number and if old dedup were kept, sending seq=0 in new incarnation could match previous one, @@ -159,7 +159,7 @@ impl ProducerTracker { let producer_key = append_key.producer_key(); if !allow_session_creation && !self.sessions.contains_key(&producer_key) { - return Err(ProduceError::SessionExpired); + return Err(ProduceError::SessionNotInstalled); } let session = self.sessions.entry(producer_key).or_insert_with(|| { @@ -353,7 +353,7 @@ mod tests { let mut coordination = ProducerTracker::default(); assert_eq!( verify(&mut coordination, existing(request(id, 0, 0, 7))), - Err(ProduceError::SessionExpired) + Err(ProduceError::SessionNotInstalled) ); assert_eq!( verify(&mut coordination, verified(request(id, 0, 0, 7))), diff --git a/src/data_plane/auxiliary_states/producer/types.rs b/src/data_plane/auxiliary_states/producer/types.rs index f399cdd1..55812162 100644 --- a/src/data_plane/auxiliary_states/producer/types.rs +++ b/src/data_plane/auxiliary_states/producer/types.rs @@ -28,6 +28,8 @@ pub enum ProduceError { ProducerFenced, #[error("producer session expired or is unknown")] SessionExpired, + #[error("producer session is not installed on this replica")] + SessionNotInstalled, #[error("sequence was reused with a different payload")] RequestIdentityConflict, #[error("duplicate position is outside the retained result window")] diff --git a/src/it/e2e/client_sdk.rs b/src/it/e2e/client_sdk.rs index ed2e3725..87b581df 100644 --- a/src/it/e2e/client_sdk.rs +++ b/src/it/e2e/client_sdk.rs @@ -966,6 +966,63 @@ fn consumer_basic_consume_earliest() -> turmoil::Result { sim.run() } +/// Producer Test: Session lease renewal, expiration, and resumption. +/// Renews inside the safety margin without resetting the sequence stream, then +/// idles past the renewed lease and verifies that a fresh producer resumes at zero. +#[test] +#[serial_test::serial] +fn producer_resumes_cleanly_after_session_lease_expiration() -> turmoil::Result { + let mut sim = build_sim(160); + host_cluster(&mut sim, &NODES, |env| { + sim_cluster(env); + }); + + sim.client("test-client", async { + let client = Arc::new(Client::connect(client_seeds()).expect("client connects")); + client + .create_topic("session-expire-test", policy()) + .await + .expect("create topic"); + + let producer = Producer::new( + client.clone(), + "session-expire-test".into(), + ProducerConfig::default(), + ); + + // 1. Send initial record + let entry1 = producer + .send(b"k1", b"v1".to_vec()) + .await + .expect("first send"); + assert_eq!(entry1, EntryId(0)); + + // Enter the one-second renewal margin of the 60-second lease. Renewal + // must reuse the nonce and preserve the existing range sequence stream. + tokio::time::sleep(Duration::from_millis(59_500)).await; + + let entry2 = producer + .send(b"k2", b"v2".to_vec()) + .await + .expect("send after proactive renewal"); + assert_eq!(entry2, EntryId(1)); + + // Idle past the renewed lease. Reopening now uses a fresh producer + // identity and starts its range sequence stream at zero. + tokio::time::sleep(Duration::from_secs(61)).await; + + let entry3 = producer + .send(b"k3", b"v3".to_vec()) + .await + .expect("send after expiration"); + assert_eq!(entry3, EntryId(2)); + + Ok(()) + }); + + sim.run() +} + /// E2E Consumer Test 1.5: Latest Skips History (Active Segment) /// StartPolicy::Latest should skip existing records in an active segment without needing a roll. #[test] diff --git a/src/it/e2e/producer_deduplication.rs b/src/it/e2e/producer_deduplication.rs index 19412cff..b1f3ca5f 100644 --- a/src/it/e2e/producer_deduplication.rs +++ b/src/it/e2e/producer_deduplication.rs @@ -124,7 +124,9 @@ fn duplicate_unknown_and_fenced_sessions_are_end_to_end_visible() -> turmoil::Re for (host, port, _) in NODES { if matches!( send_request(host, port, unknown_wire.clone()).await, - ClientResponse::Err(ServerError::ProduceRejected(ProduceError::SessionExpired)) + ClientResponse::Err(ServerError::ProduceRejected( + ProduceError::SessionNotInstalled + )) ) { rejected = true; break; From 21f46846487d530a21c343ca6c15924a4dccd9be Mon Sep 17 00:00:00 2001 From: Migorithm Date: Fri, 24 Jul 2026 01:31:59 +0400 Subject: [PATCH 39/41] routing logic optimized --- src/client/routing.rs | 177 ++++++++++++++++++++++++++---------------- 1 file changed, 111 insertions(+), 66 deletions(-) diff --git a/src/client/routing.rs b/src/client/routing.rs index 44ef88ac..074ac852 100644 --- a/src/client/routing.rs +++ b/src/client/routing.rs @@ -4,91 +4,111 @@ //! cache only ever costs an extra hop, never a wrong target. use crate::connections::protocol::{RangeDetail, TopicDetail}; use crate::control_plane::NodeAddressInfo; -use crate::control_plane::metadata::{RangeId, TopicId}; +use crate::control_plane::metadata::{RangeId, RangeState, TopicId}; use arc_swap::ArcSwap; use std::collections::HashMap; use std::net::SocketAddr; use std::sync::Arc; -/// Per-topic routing snapshot. Built once from a `DescribeTopic` response and +/// Per-topic routing cache. Built once from a `DescribeTopic` response and /// replaced wholesale on refresh. pub(crate) struct TopicRouting { pub topic_id: TopicId, - ranges: Box<[RangeRoute]>, + range_placements: Box<[RangePlacement]>, + active_ranges: Box<[ActiveRangeRoute]>, } -/// One active range's placement. Sealed/deleting ranges have no write leader. -struct RangeRoute { +/// Replica placement for active and historical ranges used by readers. +struct RangePlacement { range_id: RangeId, - keyspace_start: Vec, - /// `replica_set[0]` of the active segment, the produce target. `None` when the - /// range has no active segment (sealed/deleting) or its replica set is empty. - write_leader: Option, replica_addrs: Box<[NodeAddressInfo]>, } - -// Active segment only — correct for produce (the write head) and a tailing consumer. -// Historical/cold reads live in *sealed* segments whose placement can differ, so they -// need per-segment replicas; that's blocked on the DescribeTopic sealed-segment -// extension (`RangeDetail` exposes only `active_segment` today) and is a C3 concern. -// See c1_routing_and_connections.md / d4. -impl From<&RangeDetail> for RangeRoute { - fn from(range: &RangeDetail) -> Self { - let placement = range +impl RangePlacement { + fn from(r: &RangeDetail) -> Self { + let placement = r .active_segment .as_ref() - .or_else(|| range.sealed_segments.last()); - let replicas: Box<[NodeAddressInfo]> = placement - .map(|seg| seg.replica_set.iter().cloned().collect()) - .unwrap_or_default(); - RangeRoute { - range_id: range.range_id, - keyspace_start: range.keyspace_start.clone(), - write_leader: range - .active_segment - .as_ref() - .and_then(|_| replicas.first().cloned()), - replica_addrs: replicas, + .or_else(|| r.sealed_segments.last()); + RangePlacement { + range_id: r.range_id, + replica_addrs: placement + .map(|segment| segment.replica_set.iter().cloned().collect()) + .unwrap_or_default(), } } } +/// Active keyspace owner. Leader resolution is optional and independent of ownership. +struct ActiveRangeRoute { + range_id: RangeId, + keyspace_start: Vec, + write_leader: Option, +} + +impl ActiveRangeRoute { + fn from(t: &TopicDetail) -> Vec { + t.ranges + .iter() + .filter(|range| range.state == RangeState::Active) + .map(|range| ActiveRangeRoute { + range_id: range.range_id, + keyspace_start: range.keyspace_start.clone(), + write_leader: range + .active_segment + .as_ref() + .and_then(|segment| segment.replica_set.first().cloned()), + }) + .collect() + } +} + impl TopicRouting { + fn from_detail(detail: &TopicDetail) -> Self { + let range_placements = detail.ranges.iter().map(RangePlacement::from).collect(); + let mut active_ranges = ActiveRangeRoute::from(detail); + active_ranges.sort_unstable_by(|a, b| a.keyspace_start.cmp(&b.keyspace_start)); + + Self { + topic_id: detail.topic_id, + range_placements, + active_ranges: active_ranges.into_boxed_slice(), + } + } + + fn active_range(&self, routing_key: &[u8]) -> Option<&ActiveRangeRoute> { + let insertion = self + .active_ranges + .partition_point(|range| range.keyspace_start.as_slice() <= routing_key); + self.active_ranges.get(insertion.checked_sub(1)?) + } + /// The active range owning `routing_key` — mirrors the server's /// `route_active_range`: the active range with the greatest `keyspace_start ≤ key`. - /// Returns its write leader. + /// Returns its write leader when that address is currently resolved. pub(crate) fn write_leader(&self, routing_key: &[u8]) -> Option { - self.ranges - .iter() - .rev() - .find(|r| r.write_leader.is_some() && r.keyspace_start.as_slice() <= routing_key) - .and_then(|r| r.write_leader.clone()) + self.active_range(routing_key)?.write_leader.clone() } pub(crate) fn range_id(&self, routing_key: &[u8]) -> Option { - self.ranges - .iter() - .rev() - .find(|r| r.write_leader.is_some() && r.keyspace_start.as_slice() <= routing_key) - .map(|r| r.range_id) + Some(self.active_range(routing_key)?.range_id) } pub(crate) fn replicas(&self, range_id: RangeId) -> Option<&[NodeAddressInfo]> { - self.ranges + self.range_placements .iter() - .find(|r| r.range_id == range_id) - .map(|r| r.replica_addrs.as_ref()) + .find(|range| range.range_id == range_id) + .map(|range| range.replica_addrs.as_ref()) } pub(crate) fn write_leader_for_range(&self, range_id: RangeId) -> Option { - self.ranges + self.range_placements .iter() .find(|range| range.range_id == range_id) .and_then(|range| range.replica_addrs.first().cloned()) } pub(crate) fn active_range_ids(&self) -> impl Iterator + '_ { - self.ranges.iter().map(|r| r.range_id) + self.active_ranges.iter().map(|range| range.range_id) } } @@ -114,16 +134,7 @@ impl RoutingCache { } pub(crate) fn insert(&self, detail: &TopicDetail) { - let mut ranges: Vec = detail.ranges.iter().map(RangeRoute::from).collect(); - // Sort ranges by keyspace_start so that reverse-scan and binary search work - ranges.sort_unstable_by(|a, b| a.keyspace_start.cmp(&b.keyspace_start)); - - let topic_routing = TopicRouting { - topic_id: detail.topic_id, - ranges: ranges.into_boxed_slice(), - }; - - let routing = Arc::new(topic_routing); + let routing = Arc::new(TopicRouting::from_detail(detail)); self.topics.rcu(|current_topics| { let mut next_topics = (**current_topics).clone(); next_topics.insert(detail.name.clone(), routing.clone()); @@ -150,7 +161,7 @@ impl RoutingCache { } /// Cached *active-segment* replica addresses for a range (C3 tailing fetch). - /// Historical reads need the sealed-segment extension — see `RangeRoute`. + /// Historical reads need the sealed-segment extension — see `RangePlacement`. pub(crate) fn replicas( &self, topic: &str, @@ -171,18 +182,16 @@ impl RoutingCache { return current_topics.clone(); }; - let ranges = routing - .ranges + let range_placements = routing + .range_placements .iter() - .map(|r| RangeRoute { - range_id: r.range_id, - keyspace_start: r.keyspace_start.clone(), - write_leader: r.write_leader.as_ref().map(|_| wrong.clone()), - replica_addrs: if r.replica_addrs.is_empty() { + .map(|range| RangePlacement { + range_id: range.range_id, + replica_addrs: if range.replica_addrs.is_empty() { Box::new([]) } else { std::iter::once(wrong.clone()) - .chain(r.replica_addrs.iter().skip(1).cloned()) + .chain(range.replica_addrs.iter().skip(1).cloned()) .collect() }, }) @@ -190,7 +199,16 @@ impl RoutingCache { let poisoned = Arc::new(TopicRouting { topic_id: routing.topic_id, - ranges, + range_placements, + active_ranges: routing + .active_ranges + .iter() + .map(|range| ActiveRangeRoute { + range_id: range.range_id, + keyspace_start: range.keyspace_start.clone(), + write_leader: range.write_leader.as_ref().map(|_| wrong.clone()), + }) + .collect(), }); let mut next_topics = (**current_topics).clone(); @@ -199,3 +217,30 @@ impl RoutingCache { }); } } + +#[cfg(test)] +mod tests { + use super::*; + use crate::control_plane::{NodeAddress, NodeId}; + use std::net::{Ipv4Addr, SocketAddr, SocketAddrV4}; + + fn leader(port: u16) -> NodeAddressInfo { + let addr = SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::LOCALHOST, port)); + NodeAddressInfo::new( + NodeId::new(format!("node-{port}")), + NodeAddress::test(addr, addr), + ) + } + + fn active_range( + range_id: u64, + keyspace_start: &[u8], + leader_resolved: bool, + ) -> ActiveRangeRoute { + ActiveRangeRoute { + range_id: RangeId(range_id), + keyspace_start: keyspace_start.to_vec(), + write_leader: leader_resolved.then(|| leader(range_id as u16 + 1)), + } + } +} From a77f5dd600d07bb5e255c8242b261e325dd6f2cc Mon Sep 17 00:00:00 2001 From: Migorithm Date: Fri, 24 Jul 2026 01:38:50 +0400 Subject: [PATCH 40/41] structural property for routing cache --- src/client/routing.rs | 345 ++++++++++++++++++++++++++++++++++++++++-- 1 file changed, 336 insertions(+), 9 deletions(-) diff --git a/src/client/routing.rs b/src/client/routing.rs index 074ac852..21df9aa4 100644 --- a/src/client/routing.rs +++ b/src/client/routing.rs @@ -7,9 +7,14 @@ use crate::control_plane::NodeAddressInfo; use crate::control_plane::metadata::{RangeId, RangeState, TopicId}; use arc_swap::ArcSwap; use std::collections::HashMap; +#[cfg(any(test, debug_assertions))] +use std::collections::HashSet; use std::net::SocketAddr; use std::sync::Arc; +#[cfg(any(test, debug_assertions))] +use crate::test_traits::TAssertInvariant; + /// Per-topic routing cache. Built once from a `DescribeTopic` response and /// replaced wholesale on refresh. pub(crate) struct TopicRouting { @@ -68,11 +73,16 @@ impl TopicRouting { let mut active_ranges = ActiveRangeRoute::from(detail); active_ranges.sort_unstable_by(|a, b| a.keyspace_start.cmp(&b.keyspace_start)); - Self { + let routing = Self { topic_id: detail.topic_id, range_placements, active_ranges: active_ranges.into_boxed_slice(), - } + }; + + #[cfg(any(test, debug_assertions))] + routing.assert_invariants(); + + routing } fn active_range(&self, routing_key: &[u8]) -> Option<&ActiveRangeRoute> { @@ -140,6 +150,8 @@ impl RoutingCache { next_topics.insert(detail.name.clone(), routing.clone()); Arc::new(next_topics) }); + #[cfg(any(test, debug_assertions))] + self.assert_invariants(); } pub(crate) fn invalidate(&self, topic: &str) { @@ -152,6 +164,8 @@ impl RoutingCache { next_topics.remove(topic); Arc::new(next_topics) }); + #[cfg(any(test, debug_assertions))] + self.assert_invariants(); } /// Cached write leader for `routing_key` in `topic`, if the topic is cached and @@ -215,12 +229,108 @@ impl RoutingCache { next_topics.insert(topic.to_string(), poisoned); Arc::new(next_topics) }); + #[cfg(any(test, debug_assertions))] + self.assert_invariants(); + } +} + +#[cfg(any(test, debug_assertions))] +impl TAssertInvariant for TopicRouting { + fn assert_invariants(&self) { + // No two RangePlacement entries in range_placements share the same range_id. + let mut seen_placements = HashSet::new(); + for placement in self.range_placements.iter() { + assert!( + seen_placements.insert(placement.range_id), + "duplicate range_id {:?} in range_placements for topic {:?}", + placement.range_id, + self.topic_id + ); + } + + // Unique Active Range IDs: + let mut seen_active_ranges = HashSet::new(); + for range in self.active_ranges.iter() { + assert!( + seen_active_ranges.insert(range.range_id), + "duplicate active range_id {:?} in active_ranges for topic {:?}", + range.range_id, + self.topic_id + ); + } + + // active_ranges must be strictly sorted by keyspace_start so binary search lookup + // (partition_point) behaves deterministically and correctly. + for window in self.active_ranges.windows(2) { + assert!( + window[0].keyspace_start < window[1].keyspace_start, + "active_ranges in topic {:?} are not strictly sorted by keyspace_start: {:?} >= {:?}", + self.topic_id, + window[0].keyspace_start, + window[1].keyspace_start + ); + } + + // Placement Coverage (Subset Property): + // Every range_id in active_ranges must exist in range_placements. + for active in self.active_ranges.iter() { + assert!( + seen_placements.contains(&active.range_id), + "active range_id {:?} in topic {:?} missing from range_placements", + active.range_id, + self.topic_id + ); + } + + // Write Leader Alignment with Replica Set: + // If an active range specifies a write_leader, the corresponding RangePlacement must + // have replica_addrs non-empty and its first element equal to write_leader. + for active in self.active_ranges.iter() { + if let Some(leader) = &active.write_leader { + let placement = self + .range_placements + .iter() + .find(|p| p.range_id == active.range_id) + .unwrap_or_else(|| { + panic!( + "missing placement for active range_id {:?}", + active.range_id + ) + }); + assert!( + !placement.replica_addrs.is_empty(), + "active range {:?} has write_leader {:?} but empty replica_addrs in placement", + active.range_id, + leader + ); + assert_eq!( + placement.replica_addrs.first(), + Some(leader), + "active range {:?} write_leader {:?} mismatch with first replica address {:?}", + active.range_id, + leader, + placement.replica_addrs.first() + ); + } + } + } +} + +#[cfg(any(test, debug_assertions))] +impl TAssertInvariant for RoutingCache { + fn assert_invariants(&self) { + let topics = self.topics.load(); + for routing in topics.values() { + routing.assert_invariants(); + } } } #[cfg(test)] mod tests { use super::*; + use crate::connections::protocol::SegmentDetail; + use crate::control_plane::metadata::TopicState; use crate::control_plane::{NodeAddress, NodeId}; use std::net::{Ipv4Addr, SocketAddr, SocketAddrV4}; @@ -232,15 +342,232 @@ mod tests { ) } - fn active_range( + fn make_segment_detail(segment_id: u64, replicas: Vec) -> SegmentDetail { + SegmentDetail { + segment_id: crate::control_plane::metadata::SegmentId(segment_id), + start_entry_id: crate::control_plane::metadata::EntryId(0), + end_entry_id: None, + replica_set: replicas, + } + } + + fn make_range_detail( range_id: u64, - keyspace_start: &[u8], - leader_resolved: bool, - ) -> ActiveRangeRoute { - ActiveRangeRoute { + start_key: &[u8], + state: RangeState, + leader_addr: Option, + ) -> RangeDetail { + let active_segment = leader_addr.map(|addr| make_segment_detail(range_id * 10, vec![addr])); + RangeDetail { range_id: RangeId(range_id), - keyspace_start: keyspace_start.to_vec(), - write_leader: leader_resolved.then(|| leader(range_id as u16 + 1)), + keyspace_start: start_key.to_vec(), + keyspace_end: vec![], + state, + active_segment, + sealed_segments: Box::new([]), + split_into: None, + merged_into: None, + merged_from: None, + } + } + + fn make_topic_detail(name: &str, ranges: Vec) -> TopicDetail { + TopicDetail { + topic_id: TopicId(1), + name: name.to_string(), + state: TopicState::Active, + ranges: ranges.into_boxed_slice(), } } + + #[test] + fn test_topic_routing_from_detail_and_lookups() { + let node1 = leader(8001); + let node2 = leader(8002); + + // Pass active ranges out of order to ensure from_detail sorts them + let r2 = make_range_detail(2, b"m", RangeState::Active, Some(node2.clone())); + let r1 = make_range_detail(1, b"", RangeState::Active, Some(node1.clone())); + let r_sealed = make_range_detail(3, b"z", RangeState::Sealed, None); + + let detail = make_topic_detail("orders", vec![r2, r1, r_sealed]); + let routing = TopicRouting::from_detail(&detail); + + assert_eq!( + routing.active_range_ids().collect::>(), + vec![RangeId(1), RangeId(2)] + ); + assert_eq!(routing.write_leader(b"a"), Some(node1.clone())); + assert_eq!(routing.write_leader(b"m"), Some(node2.clone())); + assert_eq!(routing.range_id(b"abc"), Some(RangeId(1))); + assert_eq!(routing.range_id(b"xyz"), Some(RangeId(2))); + + let replicas_r1 = routing.replicas(RangeId(1)).unwrap(); + assert_eq!(replicas_r1.len(), 1); + assert_eq!(replicas_r1[0], node1); + + assert_eq!(routing.write_leader_for_range(RangeId(1)), Some(node1)); + assert_eq!(routing.write_leader_for_range(RangeId(2)), Some(node2)); + assert_eq!(routing.write_leader_for_range(RangeId(3)), None); + } + + #[test] + fn test_routing_cache_lifecycle() { + let cache = RoutingCache::new(); + assert!(cache.get("test-topic").is_none()); + + let node1 = leader(8001); + let r1 = make_range_detail(1, b"", RangeState::Active, Some(node1.clone())); + let detail = make_topic_detail("test-topic", vec![r1]); + + cache.insert(&detail); + assert!(cache.get("test-topic").is_some()); + assert_eq!(cache.range_id("test-topic", b"key1"), Some(RangeId(1))); + assert_eq!( + cache.write_leader("test-topic", b"key1"), + Some(node1.clone()) + ); + assert_eq!( + cache.replicas("test-topic", RangeId(1)).unwrap().as_ref(), + &[node1] + ); + + cache.invalidate("test-topic"); + assert!(cache.get("test-topic").is_none()); + assert_eq!(cache.write_leader("test-topic", b"key1"), None); + } + + #[test] + fn test_routing_cache_poison_write_leaders() { + let cache = RoutingCache::new(); + let orig_node = leader(8001); + let r1 = make_range_detail(1, b"", RangeState::Active, Some(orig_node)); + let detail = make_topic_detail("test-topic", vec![r1]); + + cache.insert(&detail); + + let wrong_node = leader(9999); + cache.poison_write_leaders("test-topic", wrong_node.clone()); + + assert_eq!(cache.write_leader("test-topic", b"key1"), Some(wrong_node)); + } + + #[test] + fn test_keyspace_routing_boundaries() { + let n1 = leader(8001); + let n2 = leader(8002); + let n3 = leader(8003); + + let r1 = make_range_detail(1, b"", RangeState::Active, Some(n1.clone())); + let r2 = make_range_detail(2, b"f", RangeState::Active, Some(n2.clone())); + let r3 = make_range_detail(3, b"p", RangeState::Active, Some(n3.clone())); + + let detail = make_topic_detail("boundary-topic", vec![r1, r2, r3]); + let routing = TopicRouting::from_detail(&detail); + + // Boundary tests for partition_point routing + assert_eq!(routing.range_id(b""), Some(RangeId(1))); + assert_eq!(routing.range_id(b"a"), Some(RangeId(1))); + assert_eq!(routing.range_id(b"e"), Some(RangeId(1))); + assert_eq!(routing.range_id(b"f"), Some(RangeId(2))); + assert_eq!(routing.range_id(b"g"), Some(RangeId(2))); + assert_eq!(routing.range_id(b"o"), Some(RangeId(2))); + assert_eq!(routing.range_id(b"p"), Some(RangeId(3))); + assert_eq!(routing.range_id(b"z"), Some(RangeId(3))); + } + + #[test] + fn test_invariants_pass_on_valid_routing() { + let node = leader(8001); + let r1 = make_range_detail(1, b"", RangeState::Active, Some(node)); + let detail = make_topic_detail("valid-topic", vec![r1]); + + let cache = RoutingCache::new(); + cache.insert(&detail); + cache.assert_invariants(); + } + + #[test] + fn test_invariants_fail_on_unsorted_active_ranges() { + let invalid_routing = TopicRouting { + topic_id: TopicId(1), + range_placements: Box::new([ + RangePlacement { + range_id: RangeId(1), + replica_addrs: Box::new([]), + }, + RangePlacement { + range_id: RangeId(2), + replica_addrs: Box::new([]), + }, + ]), + active_ranges: Box::new([ + ActiveRangeRoute { + range_id: RangeId(2), + keyspace_start: b"z".to_vec(), + write_leader: None, + }, + ActiveRangeRoute { + range_id: RangeId(1), + keyspace_start: b"a".to_vec(), + write_leader: None, + }, + ]), + }; + + assert!(std::panic::catch_unwind(|| invalid_routing.assert_invariants()).is_err()); + } + + #[test] + fn test_invariants_fail_on_missing_placement() { + let invalid_routing = TopicRouting { + topic_id: TopicId(1), + range_placements: Box::new([]), + active_ranges: Box::new([ActiveRangeRoute { + range_id: RangeId(1), + keyspace_start: b"a".to_vec(), + write_leader: None, + }]), + }; + + assert!(std::panic::catch_unwind(|| invalid_routing.assert_invariants()).is_err()); + } + + #[test] + fn test_invariants_fail_on_duplicate_placement() { + let invalid_routing = TopicRouting { + topic_id: TopicId(1), + range_placements: Box::new([ + RangePlacement { + range_id: RangeId(1), + replica_addrs: Box::new([]), + }, + RangePlacement { + range_id: RangeId(1), + replica_addrs: Box::new([]), + }, + ]), + active_ranges: Box::new([]), + }; + + assert!(std::panic::catch_unwind(|| invalid_routing.assert_invariants()).is_err()); + } + + #[test] + fn test_invariants_fail_on_write_leader_mismatch() { + let invalid_routing = TopicRouting { + topic_id: TopicId(1), + range_placements: Box::new([RangePlacement { + range_id: RangeId(1), + replica_addrs: Box::new([leader(8001)]), + }]), + active_ranges: Box::new([ActiveRangeRoute { + range_id: RangeId(1), + keyspace_start: b"a".to_vec(), + write_leader: Some(leader(9999)), + }]), + }; + + assert!(std::panic::catch_unwind(|| invalid_routing.assert_invariants()).is_err()); + } } From e9fa3180a98e31986cec96ea07d88a1f1304810c Mon Sep 17 00:00:00 2001 From: Migorithm Date: Fri, 24 Jul 2026 12:16:22 +0400 Subject: [PATCH 41/41] fix: timer schedule on segment roll MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit During a FlappingPartition fault, a data write leader (node-1) experiences follower replication timeouts and requests a segment roll (RequestSegmentRoll). When the metadata state machine commits the roll, DataPlane::handle_segment_roll_committed (state.rs) replays uncommitted producer entries into the new successor segment and sets self.needs_flush = true. However, handle_segment_roll_committed was not scheduling a BatchFlushTimer with the ticker system. Without a scheduled timer (or a subsequent incoming client request), flush_batch() was never triggered for the replayed records. As a result, the replayed entries remained in dirty_segments, and the client's produce request was left stranded on reply_rx.await until timing out after 40 seconds. #### Fixes Applied 1. Timer Scheduling on Segment Roll: • In state.rs, scheduled BatchFlushTimer::deadline() whenever handle_segment_roll_committed sets self.needs_flush = true. This ensures flush_batch() is promptly executed to write the WAL, trigger replication, and complete pending producer ACK channels for successor segments. 2. Scenario Recovery Validation: • In scenario.rs, updated is_valid() so FlappingPartition and PartitionLeader scenarios require at least 15s of post-heal recovery budget before simulation_secs ends, ensuring full SWIM reconnect, Raft consensus catch-up, and produce ACK verification before simulation termination. --- src/data_plane/state.rs | 9 +++++++++ src/it/sim/scenario.rs | 4 ++-- 2 files changed, 11 insertions(+), 2 deletions(-) diff --git a/src/data_plane/state.rs b/src/data_plane/state.rs index a3fb7ce9..c3f7e3c3 100644 --- a/src/data_plane/state.rs +++ b/src/data_plane/state.rs @@ -1301,6 +1301,15 @@ impl DataPlane { self.dirty_segments .insert(cmd.old_segment_key.with_segment_id(cmd.new_segment_id)); self.needs_flush = true; + + // ! Uncommitted entries are replayed into the successor segment, and self.needs_flush = true is flagged. + // ! If no BatchFlushTimer was scheduled for the ticker system, flush_batch() won't be triggered( when no subsequent client produce request arrived), + // ! leaving replayed records dirty and stranding the waiting producer client on reply_rx.await until timing out after 40s. + self.out + .store_batch_produce_timer(TimerCommand::SetSchedule { + seq: self.wal.next_lsn(), + timer: BatchFlushTimer::deadline(), + }); } fn handle_segment_meta_sealed(&mut self, cmd: SegmentMetaSealed) { diff --git a/src/it/sim/scenario.rs b/src/it/sim/scenario.rs index 4a3adf09..3584ac50 100644 --- a/src/it/sim/scenario.rs +++ b/src/it/sim/scenario.rs @@ -276,7 +276,7 @@ impl SimScenario { ) }); if !topic_exists - || event.at_secs + fault.heal_after_secs >= self.simulation_secs + || event.at_secs + fault.heal_after_secs + 15 >= self.simulation_secs { return false; } @@ -289,7 +289,7 @@ impl SimScenario { ) }); if !topic_exists - || event.at_secs + fault.heal_after_secs >= self.simulation_secs + || event.at_secs + fault.heal_after_secs + 15 >= self.simulation_secs { return false; }