diff --git a/docs/clients/c2_producer.md b/docs/clients/c2_producer.md index 8df7f3a2..1dd4a57d 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 -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. +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 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 +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_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/clients/client_roadmap.md b/docs/clients/client_roadmap.md index 4e0f09df..490d42e1 100644 --- a/docs/clients/client_roadmap.md +++ b/docs/clients/client_roadmap.md @@ -121,8 +121,11 @@ 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 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/d10_transparent_retry_deduplication.md b/docs/data-plane/d10_transparent_retry_deduplication.md new file mode 100644 index 00000000..bb5ab20e --- /dev/null +++ b/docs/data-plane/d10_transparent_retry_deduplication.md @@ -0,0 +1,371 @@ +# 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) +│ +├── lease timestamp expired +│ └── SessionExpired +├── unexpired session not installed on this replica +│ └── SessionNotInstalled +├── 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. + +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: + +- 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. + 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. +- 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`. +- 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. + +--- + +## 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. + +### 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 shape + +1. **Storage and WAL framing.** Extend entry identity, generalize the consumer-only + 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 + 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 SDK producer path. The low-level raw append path remains +at-least-once by design. diff --git a/docs/data-plane/d8_consumer_offset_management.md b/docs/data-plane/d8_consumer_offset_management.md index aa007b43..b4cdbcbe 100644 --- a/docs/data-plane/d8_consumer_offset_management.md +++ b/docs/data-plane/d8_consumer_offset_management.md @@ -168,12 +168,56 @@ 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. 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 +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 eaa99bcc..9240ba95 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-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`. @@ -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. diff --git a/src/client/consumer/group.rs b/src/client/consumer/group.rs index aee91e70..384f4b90 100644 --- a/src/client/consumer/group.rs +++ b/src/client/consumer/group.rs @@ -7,11 +7,11 @@ 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::consumer_offset_management::ledger::{ +use crate::data_plane::auxiliary_states::consumer_offsets::state::{ ConsumerOffsetKey, ConsumerOffsetPosition, ConsumerOffsetUpdate, }; @@ -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 33e26404..25a0f26c 100644 --- a/src/client/consumer/mod.rs +++ b/src/client/consumer/mod.rs @@ -10,11 +10,12 @@ 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::consumer_offset_management::ledger::ConsumerOffsetPosition; +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 25b01335..b2168415 100644 --- a/src/client/consumer/range_fetcher.rs +++ b/src/client/consumer/range_fetcher.rs @@ -5,12 +5,13 @@ 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::consumer_offset_management::ledger::ConsumerOffsetPosition; +use crate::data_plane::auxiliary_states::consumer_offsets::state::ConsumerOffsetPosition; +use crate::data_plane::messages::query::RangeOffsets; const EMPTY_FETCHES_BEFORE_REFRESH: u8 = 20; @@ -137,124 +138,25 @@ 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 + 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; @@ -262,8 +164,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 +173,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/consumer/topic_fetch_manager.rs b/src/client/consumer/topic_fetch_manager.rs index 2b814854..3834f174 100644 --- a/src/client/consumer/topic_fetch_manager.rs +++ b/src/client/consumer/topic_fetch_manager.rs @@ -10,7 +10,8 @@ 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::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/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/client/mod.rs b/src/client/mod.rs index ad32e4da..1a48029b 100644 --- a/src/client/mod.rs +++ b/src/client/mod.rs @@ -22,16 +22,17 @@ 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; 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, }; +use crate::data_plane::messages::query::RangeOffsets; pub use codec::CompressionCodec; pub use consumer::{ CommitMode, Consumer, ConsumerConfig, ConsumerRecord, DeliverySemantic, KeyInterest, @@ -45,9 +46,10 @@ use uuid::Uuid; use crate::connections::protocol::{ ClientDataPlaneRequest, ClientRequest, ClientResponse, CommitConsumerOffsetRequest, ConsumerGroupAssignmentResponse, ConsumerGroupSyncAction, ControlPlaneRequest, - ControlPlaneResponse, DataPlaneResponse, FetchConsumerOffsetRequest, ProduceRequest, + FetchConsumerOffsetRequest, OpenProducerSessionRequest, ProduceRequest, ProducerSessionOpened, RangeOffsetRequest, }; +use crate::data_plane::{PayloadBytes, ProduceError, ProducerAppendIdentity}; use futures::{StreamExt, stream::FuturesUnordered}; pub use redirect::RetryPolicy; @@ -120,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; @@ -155,11 +157,8 @@ impl Client { let served = self.call(addr, req).await?; match served.response { - ClientResponse::DataPlane(DataPlaneResponse::RangeOffset { - start_entry_id, - next_entry_id, - }) => return Ok((start_entry_id, next_entry_id)), - ClientResponse::DataPlane(DataPlaneResponse::SegmentNotLocal) => { + ClientResponse::Ok(ClientSuccess::RangeOffset(res)) => return Ok(res), + ClientResponse::Err(ServerError::SegmentNotLocal) => { last_error = Some("range offsets segment not local".to_string()); } _ => return Err(ClientError::UnexpectedResponse), @@ -190,8 +189,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, @@ -229,12 +228,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, @@ -260,8 +259,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), } } @@ -275,7 +274,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), } } @@ -299,7 +298,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) @@ -311,20 +310,17 @@ 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::Ok(ClientSuccess::ProducerSessionOpened(session)) => Ok(session), + _ => Err(ClientError::UnexpectedResponse), + } } /// Produce one entry under `routing_key`, returning the committed `entry_id`. @@ -335,8 +331,9 @@ impl Client { topic: &str, range_id: RangeId, routing_key: &[u8], - data: Vec, + data: impl Into, record_count: u32, + producer_identity: Option, ) -> Result { // Describe once to seed the cache (gives the first hop). let routing = self.resolve_topic_if_missing(topic).await?; @@ -352,8 +349,9 @@ impl Client { topic_name: topic.to_string(), range_id, routing_key: routing_key.to_vec(), - data, + data: data.into(), record_count, + producer_identity, }; let served = self.call(start, request).await?; @@ -362,15 +360,16 @@ impl Client { if served.redirected || matches!( &served.response, - ClientResponse::DataPlane(DataPlaneResponse::StaleRange) + ClientResponse::Err(ServerError::StaleRange) ) { self.cache.invalidate(topic); } match served.response { - ClientResponse::DataPlane(DataPlaneResponse::Produced { entry_id }) => Ok(entry_id), - ClientResponse::DataPlane(DataPlaneResponse::StaleRange) => { - Err(ClientError::StaleRange) + 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), } @@ -424,12 +423,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::DataPlane(DataPlaneResponse::ShardNotLocal { - hint_node: None, - }) => Redirect::Done, _ => Self::redirect_target(&response), } } else { @@ -474,49 +471,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::InternalError(_) => Redirect::Reresolve, - ControlPlaneResponse::TopicCreated - | ControlPlaneResponse::AlreadyExists - | ControlPlaneResponse::TopicDeleted - | ControlPlaneResponse::ConsumerGroupAssignment(_) - | ControlPlaneResponse::ConsumerGroupLeft - | ControlPlaneResponse::TopicList { .. } - | ControlPlaneResponse::TopicDetail(_) => Redirect::Done, - }, - ClientResponse::DataPlane(dp) => match dp { - DataPlaneResponse::NotWriteLeader { leader_addr } => match leader_addr { + 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, }, - DataPlaneResponse::ShardNotLocal { hint_node } => match hint_node { + ServerError::NotRaftLeader { leader_addr } + | ServerError::NotWriteLeader { leader_addr } => match leader_addr { Some(addr) => Redirect::Follow(addr.client_addr()), None => Redirect::Reresolve, }, - DataPlaneResponse::TopicNotFound => Redirect::NotFound, - DataPlaneResponse::StaleRange => Redirect::Done, - DataPlaneResponse::InternalError(_) => Redirect::Reresolve, - DataPlaneResponse::SegmentNotLocal - | DataPlaneResponse::Produced { .. } - | 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/producer/buffers.rs b/src/client/producer/buffers.rs index a5c92b21..2f7380f6 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,13 +48,18 @@ 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 { + 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/redirect.rs b/src/client/redirect.rs index 0294e033..3b219ee5 100644 --- a/src/client/redirect.rs +++ b/src/client/redirect.rs @@ -14,7 +14,7 @@ use tokio::time::Instant; use crate::client::Client; use crate::client::error::ClientError; use crate::connections::protocol::{ - ClientDataPlaneRequest, ClientRequest, ClientResponse, ControlPlaneResponse, DataPlaneResponse, + ClientDataPlaneRequest, ClientRequest, ClientResponse, ClientSuccess, ServerError, }; /// Consecutive hint-follows before forcing a backoff — breaks a redirect cycle @@ -150,35 +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::TopicCreated - )), + classify(&ClientResponse::Err(ServerError::ShardNotLocal { + hint_node: None + })), + "reresolve" + ); + assert_eq!( + classify(&ClientResponse::Ok(ClientSuccess::TopicCreated)), "done" ); } @@ -186,33 +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::ShardNotLocal { hint_node: None } - )), + classify(&ClientResponse::Err(ServerError::SegmentNotLocal)), "reresolve" ); assert_eq!( - classify(&ClientResponse::DataPlane(DataPlaneResponse::TopicNotFound)), - "notfound" - ); - assert_eq!( - classify(&ClientResponse::DataPlane( - DataPlaneResponse::SegmentNotLocal - )), - "done" - ); - assert_eq!( - classify(&ClientResponse::DataPlane(DataPlaneResponse::Produced { - entry_id: 7.into() - })), + classify(&ClientResponse::Ok(ClientSuccess::Produced(7.into()))), "done" ); } diff --git a/src/client/routing.rs b/src/client/routing.rs index 79331fce..21df9aa4 100644 --- a/src/client/routing.rs +++ b/src/client/routing.rs @@ -4,88 +4,122 @@ //! 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; +#[cfg(any(test, debug_assertions))] +use std::collections::HashSet; use std::net::SocketAddr; use std::sync::Arc; -/// Per-topic routing snapshot. Built once from a `DescribeTopic` response and +#[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 { 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)); + + 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> { + 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.active_ranges.iter().map(|range| range.range_id) + } } /// Thread-safe, lock-free per-topic cache. Entries are replaced on refresh and @@ -110,21 +144,14 @@ 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()); Arc::new(next_topics) }); + #[cfg(any(test, debug_assertions))] + self.assert_invariants(); } pub(crate) fn invalidate(&self, topic: &str) { @@ -137,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 @@ -146,7 +175,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, @@ -167,18 +196,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() }, }) @@ -186,12 +213,361 @@ 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(); 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}; + + 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 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, + 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: 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()); } } 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/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 641ec1f8..9392a83a 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)] @@ -39,55 +39,55 @@ 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, +} + +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, Leave, } -#[derive(Debug, BorshSerialize, BorshDeserialize)] -pub enum ControlPlaneResponse { - // CreateTopic - TopicCreated, - AlreadyExists, - // DeleteTopic - TopicDeleted, - TopicNotFound, - // 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, - // All control plane operations - InternalError(String), +#[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, @@ -97,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`. @@ -226,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, @@ -324,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 c07740c6..4e8efeee 100644 --- a/src/connections/protocol/data_plane.rs +++ b/src/connections/protocol/data_plane.rs @@ -9,22 +9,45 @@ use borsh::{BorshDeserialize, BorshSerialize}; use crate::{ - control_plane::{ - NodeAddressInfo, - metadata::{ - EntryId, RangeId, RangeMeta, RangeState, TopicId, TopicMeta, - consumer_group::GenerationId, - }, + client::ClientSuccess, + control_plane::metadata::{ + EntryId, RangeId, RangeMeta, RangeState, TopicId, TopicMeta, consumer_group::GenerationId, }, data_plane::{ - consumer_offset_management::ledger::{ - ConsumerOffsetKey, ConsumerOffsetPosition, ConsumerOffsetUpdate, - }, - messages::query::{FetchResult, ListOffsetsResult}, + PayloadBytes, ProducerAppendIdentity, + auxiliary_states::consumer_offsets::state::{ConsumerOffsetKey, ConsumerOffsetUpdate}, + messages::query::FetchedRecords, }, 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, +} + +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` @@ -58,8 +81,9 @@ 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, } #[derive(Debug, Clone, BorshSerialize, BorshDeserialize)] @@ -104,107 +128,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, - }, - // 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, - }, - ShardNotLocal { - hint_node: Option, - }, - StaleRange, - TopicNotFound, - SegmentNotLocal, - - ConsumerOffsetCommitted, - StaleConsumerGroupEpoch(GenerationId), - ConsumerOffset(Option), - ConsumerOffsetGenerationMismatch(ConsumerOffsetGenerationMismatch), - InternalError(String), -} - -impl DataPlaneResponse { - pub fn is_routing_error(&self) -> bool { - matches!( - self, - DataPlaneResponse::SegmentNotLocal - | DataPlaneResponse::ShardNotLocal { .. } - | 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, @@ -227,7 +155,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 { @@ -236,7 +164,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..11c225fc --- /dev/null +++ b/src/connections/protocol/error.rs @@ -0,0 +1,87 @@ +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 + ) + } +} + +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/connections/protocol/mod.rs b/src/connections/protocol/mod.rs index 7d8e08ad..9e8b2c56 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::RangeOffsets, + }, + impl_from_variant, impl_from_variant_via, }; // ── Top-level dispatch ───────────────────────────────────────────────────── @@ -42,12 +49,65 @@ 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(EntryId), + Fetched { + entries: Box<[EntryPayload]>, + next_entry_id: EntryId, + progress_signal: RangeProgressSignal, + }, + RangeOffset(RangeOffsets), + ConsumerOffsetCommitted, + ConsumerOffset(Option), + + // Admin + ClusterInfo { + nodes: Box<[NodeInfo]>, + }, + TopicStats { + topics: Box<[TopicStats]>, + }, + ShardInfo { + detail: Option, + }, + ShardLeader { + leader: Option, + }, +} + +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 +115,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/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( 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/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(); } } 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/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/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/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/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/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/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)); +} 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 55b9f893..7b30dce2 100644 --- a/src/control_plane/metadata/topic.rs +++ b/src/control_plane/metadata/topic.rs @@ -3,15 +3,20 @@ use std::collections::{HashMap, HashSet}; use super::constants::*; use super::*; use crate::{ + client::ServerError, control_plane::{ NodeId, Replicas, metadata::{ error::MetadataError, event::ConsumerGroupEpochSnapshot, + producer_sessions::ProducerSessions, strategy::{PartitionStrategy, StoragePolicy}, }, }, - data_plane::SegmentKey, + data_plane::{ + ProduceError, ProducerAppendIdentity, SegmentKey, + messages::command::AuthorizedProducerIdentity, + }, }; use borsh::{BorshDeserialize, BorshSerialize}; @@ -32,7 +37,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 +66,7 @@ impl TopicMeta { ranges: HashMap::from([(range_id, range)]), next_range_id: 1, consumer_groups: HashMap::new(), + producer_sessions: Default::default(), } } @@ -78,6 +86,28 @@ 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) + } + + /// 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, @@ -495,11 +525,44 @@ 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(); } } + + pub(crate) fn verify_producer_session( + &self, + 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 + // 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::SessionNotInstalled) + } } pub struct TopicStats { @@ -529,6 +592,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/actor.rs b/src/data_plane/actor.rs index 3c660ace..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()) + 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> { - self.0.send_async(msg.into()).await + pub async fn send_async(&self, msg: impl Into) -> Result<(), ServerError> { + self.0 + .send_async(msg.into()) + .await + .map_err(|err| ServerError::Internal(format!("Channel dropped: {}", err))) } } 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..1226a8f5 --- /dev/null +++ b/src/data_plane/auxiliary_states/consumer_offsets/mod.rs @@ -0,0 +1,209 @@ +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(Debug, 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 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 { + 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() + } +} + +#[derive(Debug)] +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 93% rename from src/data_plane/consumer_offset_management/replication.rs rename to src/data_plane/auxiliary_states/consumer_offsets/replication.rs index 5be7a558..784ab79c 100644 --- a/src/data_plane/consumer_offset_management/replication.rs +++ b/src/data_plane/auxiliary_states/consumer_offsets/replication.rs @@ -1,17 +1,18 @@ +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, }; 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, @@ -98,13 +99,13 @@ 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)) } - 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/consumer_offset_management/ledger.rs b/src/data_plane/auxiliary_states/consumer_offsets/state.rs similarity index 81% rename from src/data_plane/consumer_offset_management/ledger.rs rename to src/data_plane/auxiliary_states/consumer_offsets/state.rs index eb4bebd8..5f019576 100644 --- a/src/data_plane/consumer_offset_management/ledger.rs +++ b/src/data_plane/auxiliary_states/consumer_offsets/state.rs @@ -1,8 +1,5 @@ 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}; @@ -11,9 +8,6 @@ use crate::control_plane::metadata::consumer_group::GenerationId; use crate::control_plane::metadata::{EntryId, TopicId}; use crate::data_plane::SegmentKey; -const SNAPSHOT_FILE: &str = "consumer-offsets.snapshot"; -const SNAPSHOT_TEMP_FILE: &str = "consumer-offsets.snapshot.tmp"; - #[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash, BorshSerialize, BorshDeserialize)] pub(crate) struct ConsumerOffsetKey { pub(crate) topic_id: TopicId, @@ -81,38 +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 OffsetLedger { +pub(crate) struct ConsumerOffsets { epochs: HashMap, offsets: HashMap, installed_placements: HashMap<(TopicId, RangeId), SegmentKey>, } -impl OffsetLedger { - 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), - } - } - - 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)) } @@ -222,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 = OffsetLedger::default(); + let mut ledger = ConsumerOffsets::default(); ledger.apply(OffsetRecord::PlacementInstalled(current)); assert!(ledger.has_installed_placement(¤t)); @@ -232,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 = OffsetLedger::default(); + let mut ledger = ConsumerOffsets::default(); ledger.apply(OffsetRecord::EpochSeal(EpochSeal { key: key(), generation: GenerationId(3), @@ -249,10 +217,7 @@ mod tests { generation: GenerationId(3), position, })); - ledger.write_snapshot(dir.path()).unwrap(); - - let recovered = OffsetLedger::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/consumer_offset_management/types.rs b/src/data_plane/auxiliary_states/consumer_offsets/types.rs similarity index 94% rename from src/data_plane/consumer_offset_management/types.rs rename to src/data_plane/auxiliary_states/consumer_offsets/types.rs index d673b75e..112bf5f9 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, }; @@ -13,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, @@ -43,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), @@ -71,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/consumer_offset_management/manager.rs b/src/data_plane/auxiliary_states/manager.rs similarity index 75% rename from src/data_plane/consumer_offset_management/manager.rs rename to src/data_plane/auxiliary_states/manager.rs index 5b1f87ad..3696446c 100644 --- a/src/data_plane/consumer_offset_management/manager.rs +++ b/src/data_plane/auxiliary_states/manager.rs @@ -1,139 +1,104 @@ +use super::consumer_offsets::state::{ + ConsumerOffsetKey, ConsumerOffsetPosition, ConsumerOffsets, EpochSeal, OffsetRecord, StaleEpoch, +}; +use super::snapshot::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}; 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::auxiliary_states::producer::types::AppendKey; +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::ProduceError; +use crate::data_plane::ProducerAppendIdentity; +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; -use super::types::*; +use super::consumer_offsets::types::*; -#[derive(Default)] -pub(crate) struct ConsumerOffsetManager { - offset_ledger: OffsetLedger, // durable offsets, epochs, and local placement readiness - coordination: ConsumerOffsetCoordination, // live placements, pending mutations, parked commits, and replication tracking - checkpoint: OffsetCheckpointState, // WAL reclamation and checkpoint progress +#[derive(Debug, Default)] +pub(crate) struct AuxiliaryStateManager { + offsets: ConsumerOffsets, + consumer_coord: ConsumerOffsetCoordination, // live placements, pending mutations, parked commits, and replication tracking + producer_tracker: ProducerTracker, + checkpoint: AuxiliaryCheckpointState, // 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, +#[derive(Debug, Default)] +struct AuxiliaryCheckpointState { + uncheckpointed_lsns: VecDeque, + checkpoint_lsn: u64, + in_flight: bool, } -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) +impl AuxiliaryStateManager { + #[cfg(test)] + pub(crate) fn new(offsets: ConsumerOffsets) -> Self { + Self::from_recovery(offsets, ProducerTracker::default()) } - 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; + pub(crate) fn from_recovery( + offsets: ConsumerOffsets, + producer_tracker: ProducerTracker, + ) -> Self { + Self { + offsets, + producer_tracker, + ..Default::default() } - PlacementObservation::Accepted } - fn unready_followers(&self) -> Box<[NodeId]> { - self.replicas - .followers() - .filter(|node| !self.is_replica_ready(node)) - .cloned() - .collect() + pub(crate) fn verify_producer( + &mut self, + segment_key: SegmentKey, + producer: AuthorizedProducerIdentity, + received_at_ms: u64, + ) -> Result<(), ProduceError> { + self.producer_tracker + .verify(segment_key, producer, received_at_ms) } -} -#[derive(Default)] -struct OffsetCheckpointState { - uncheckpointed_lsns: VecDeque, - checkpoint_lsn: u64, - in_flight: bool, -} + pub(crate) fn unstage_producer(&mut self, append_key: &AppendKey) { + self.producer_tracker.unstage(append_key); + } -impl ConsumerOffsetManager { - pub(crate) fn new(offset_ledger: OffsetLedger) -> Self { - Self { - offset_ledger, - ..Default::default() - } + pub(crate) fn advance_producer( + &mut self, + segment_key: SegmentKey, + producer: ProducerAppendIdentity, + entry_id: EntryId, + lsn: u64, + ) { + self.producer_tracker + .advance(segment_key, producer, entry_id); + self.mark_auxiliary_state_dirty(lsn); } - pub(crate) fn future_entry(&mut self, key: ConsumerOffsetKey) -> &mut Vec { - self.coordination - .future_offset_commits - .entry(key) - .or_default() + pub(crate) fn push_future(&mut self, key: ConsumerOffsetKey, future: FutureOffsetCommit) { + self.consumer_coord.future_entry(key).push(future); } pub(crate) fn offset(&self, key: &ConsumerOffsetKey) -> Option { - self.offset_ledger.offset(key) + self.offsets.offset(key) } pub(crate) fn durable_generation(&self, key: &ConsumerOffsetKey) -> GenerationId { - self.offset_ledger.generation(key) + self.offsets.generation(key) } pub(crate) fn commit_consumer_offset( @@ -161,8 +126,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 +143,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 +158,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 +182,27 @@ 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(), + snapshot: AuxiliarySnapshot { + consumer_offsets: self.offsets.clone(), + producer_sessions: self.producer_tracker.sessions(), + }, data_dir, }) } - pub(crate) fn offset_checkpoint_in_flight(&self) -> bool { + #[cfg(any(test, debug_assertions))] + pub(crate) fn assert_producer_invariants(&self) { + self.producer_tracker.assert_invariants(); + } + + pub(crate) fn auxiliary_checkpoint_in_flight(&self) -> bool { self.checkpoint.in_flight } @@ -241,7 +219,7 @@ impl ConsumerOffsetManager { } // need to place! - let local_placement_installed = self.offset_ledger.has_installed_placement(&segment_key); + let local_placement_installed = self.offsets.has_installed_placement(&segment_key); let replica_states = replicas .iter() .cloned() @@ -254,7 +232,7 @@ impl ConsumerOffsetManager { (node, state) }) .collect(); - self.coordination.placements.insert( + self.consumer_coord.placements.insert( segment_key.placement_key(), OffsetPlacement { segment_key, @@ -272,7 +250,7 @@ impl ConsumerOffsetManager { 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)) @@ -358,7 +336,7 @@ impl ConsumerOffsetManager { .collect(); if segment_key.segment_id == SegmentId(0) - || self.offset_ledger.has_installed_placement(&segment_key) + || self.offsets.has_installed_placement(&segment_key) { replica_states.insert(leader.clone(), OffsetReplicaState::Ready); } @@ -366,7 +344,7 @@ impl ConsumerOffsetManager { } 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 { @@ -386,7 +364,7 @@ impl ConsumerOffsetManager { &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( @@ -398,23 +376,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.consumer_coord.push_pending(PendingOffsetMutation::new( + OffsetRecord::BootstrapEntry(entry), + OffsetMutationCompletion::EpochSeal, )); + } + self.consumer_coord.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 +408,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.offsets.can_source_snapshot_for(&request.segment_key) { return None; } Some(DataTransportCommand::send_to_targets( @@ -454,7 +424,7 @@ impl ConsumerOffsetManager { ) -> ConsumerOffsetSnapshot { ConsumerOffsetSnapshot { entries: self - .offset_ledger + .offsets .snapshot_range(segment_key.topic_id, segment_key.range_id), segment_key, replica_set, @@ -496,11 +466,11 @@ 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.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 @@ -512,7 +482,7 @@ impl ConsumerOffsetManager { } pub(crate) fn drain_ready_placements(&mut self, from: &NodeId) -> Vec { - self.coordination + self.consumer_coord .placements .values_mut() .filter_map(|placement| { @@ -550,20 +520,20 @@ impl ConsumerOffsetManager { } pub(crate) fn has_pending_offsets(&self) -> bool { - !self.coordination.pending_offset_mutations.is_empty() + self.consumer_coord.has_pending_offsets() } pub(crate) fn push_pending_offset_mutation(&mut self, cmd: impl Into) { - self.coordination.pending_offset_mutations.push(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.offset_replication.process_ack(ack); + self.consumer_coord.ack_replication(ack); 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> { @@ -577,74 +547,22 @@ 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, - }) + self.consumer_coord + .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.offsets.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.consumer_coord.encode_offsets() } pub(crate) fn take_pending(&mut self) -> Vec { - std::mem::take(&mut self.coordination.pending_offset_mutations) + self.consumer_coord.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 drop_coordination(&mut self, error: &str) { + self.consumer_coord.drop_all(error); } pub(crate) fn flush_batch( @@ -659,7 +577,7 @@ impl ConsumerOffsetManager { } for PendingOffsetMutation { record, completion } in pending { - self.offset_ledger.apply(record.clone()); + self.offsets.apply(record.clone()); match completion { OffsetMutationCompletion::EpochSeal => {} OffsetMutationCompletion::LeaderCommit(commit) => { @@ -674,7 +592,7 @@ impl ConsumerOffsetManager { continue; } - let seq = self.coordination.offset_replication.begin( + let seq = self.consumer_coord.begin_replication( followers.clone().into_iter().collect(), commit.required_followers, offset_commit.clone(), @@ -723,10 +641,7 @@ impl ConsumerOffsetManager { } fn take_future_commits(&mut self, key: ConsumerOffsetKey) -> Vec { - self.coordination - .future_offset_commits - .remove(&key) - .unwrap_or_default() + self.consumer_coord.take_future_commits(&key) } #[cfg(test)] @@ -735,7 +650,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 +661,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,15 +691,15 @@ 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 = ConsumerOffsets::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); let placement = manager - .coordination + .consumer_coord .placements .get(&(current.topic_id, current.range_id)) .unwrap(); @@ -797,10 +712,10 @@ 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(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(); @@ -812,7 +727,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(); @@ -830,13 +745,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 = ConsumerOffsetManager::new(OffsetLedger::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(); @@ -850,7 +765,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(ConsumerOffsets::default()); manager.install_leader_placement(current, ShardGroupId(7), current_replicas.clone()); @@ -874,7 +789,7 @@ mod tests { ); let placement = manager - .coordination + .consumer_coord .placements .get(&(current.topic_id, current.range_id)) .unwrap(); @@ -887,7 +802,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(ConsumerOffsets::default()); manager.install_leader_placement( segment(1), ShardGroupId(7), @@ -949,14 +864,14 @@ mod tests { requester: new_leader.clone(), }; - let empty = ConsumerOffsetManager::new(OffsetLedger::default()); + let empty = AuxiliaryStateManager::new(ConsumerOffsets::default()); assert!( empty .create_offset_snapshot(request.clone(), &NodeId::new("empty")) .is_none() ); - let mut ledger = OffsetLedger::default(); + let mut ledger = ConsumerOffsets::default(); ledger.apply(OffsetRecord::PlacementInstalled(old_segment)); ledger.apply(OffsetRecord::EpochSeal(EpochSeal { key: key(), @@ -967,7 +882,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 +917,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 = ConsumerOffsets::default(); ledger.apply(OffsetRecord::EpochSeal(EpochSeal { key: key(), generation: GenerationId(1), @@ -1012,8 +927,8 @@ mod tests { generation: GenerationId(1), position: position(), })); - let mut manager = ConsumerOffsetManager::new(ledger); - manager.coordination.placements.insert( + let mut manager = AuxiliaryStateManager::new(ledger); + manager.consumer_coord.placements.insert( (TopicId(1), RangeId(2)), OffsetPlacement { segment_key: segment(1), @@ -1096,12 +1011,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(ConsumerOffsets::default()); let (pending_reply, mut pending_response) = oneshot::channel(); manager - .coordination - .pending_offset_mutations - .push(PendingOffsetMutation::new( + .consumer_coord + .push_pending(PendingOffsetMutation::new( OffsetRecord::OffsetCommit(ConsumerOffsetUpdate { key: key(), generation: GenerationId(1), @@ -1115,20 +1029,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.consumer_coord.begin_replication( HashSet::from_iter([NodeId::new("follower")]), HashSet::from_iter([NodeId::new("follower")]), ConsumerOffsetUpdate { @@ -1139,13 +1053,12 @@ mod tests { replication_reply, ); - manager.fail_all("WAL failed"); + manager.drop_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.consumer_coord.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..68055102 --- /dev/null +++ b/src/data_plane/auxiliary_states/mod.rs @@ -0,0 +1,5 @@ +pub(crate) mod consumer_offsets; +pub(crate) mod manager; +pub(crate) mod producer; +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 new file mode 100644 index 00000000..ca0a4d5e --- /dev/null +++ b/src/data_plane/auxiliary_states/producer/mod.rs @@ -0,0 +1,5 @@ +mod state; +pub(crate) mod types; +use crate::data_plane::messages::command::AuthorizedProducerIdentity; +pub(crate) use state::{ProducerSessions, ProducerTracker}; +pub use types::{ProduceError, ProducerAppendIdentity}; 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..1ee3422b --- /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::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, + // * 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::SessionNotInstalled); + } + + 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::SessionNotInstalled) + ); + 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" + ); + } +} 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..55812162 --- /dev/null +++ b/src/data_plane/auxiliary_states/producer/types.rs @@ -0,0 +1,70 @@ +use borsh::{BorshDeserialize, BorshSerialize}; +use uuid::Uuid; + +use crate::{ + client::RangeId, + control_plane::metadata::{EntryId, 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)] +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("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")] + 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), + #[error("duplicate sequence produce: {0:?}")] + Duplicate(EntryId), +} + +#[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) + } +} 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 48a0918c..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::consumer_offset_management::ledger::OffsetLedger; +use crate::data_plane::auxiliary_states::snapshot::AuxiliarySnapshot; 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) => { - if let Err(e) = job.offsets.write_snapshot(&job.data_dir) { - tracing::error!("consumer offset checkpoint failed: {e}"); + CheckpointTask::AuxiliaryState(job) => { + if let Err(e) = job.snapshot.write(&job.data_dir) { + tracing::error!("auxiliary state 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(OffsetCheckpointJob), + 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 snapshot: AuxiliarySnapshot, pub data_dir: PathBuf, pub checkpointed_lsn: u64, } diff --git a/src/data_plane/cold_read.rs b/src/data_plane/cold_read.rs index 386aab84..5ec7335f 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 @@ -252,6 +253,7 @@ impl ColdReadPool { // lsn is a live-WAL concept; cold reads serve durable // data, so it carries no meaningful LSN. lsn: 0, + producer_identity: None, })); if bytes_read >= req.max_bytes { 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/command.rs b/src/data_plane/messages/command.rs index 675e61d0..dadb8252 100644 --- a/src/data_plane/messages/command.rs +++ b/src/data_plane/messages/command.rs @@ -1,29 +1,31 @@ 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, 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; +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 @@ -54,11 +56,39 @@ 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, 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), + } + } +} + +#[derive(Debug)] pub struct CommitConsumerOffset { pub update: ConsumerOffsetUpdate, pub reply: oneshot::Sender, @@ -91,9 +121,10 @@ 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, } #[derive(Debug, Clone, BorshSerialize, BorshDeserialize)] @@ -211,7 +242,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 { @@ -335,10 +366,8 @@ 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, - }, - Err(String), + Ok(EntryId), + Err(ProduceError), } pub struct SegmentCheckpointComplete { @@ -348,7 +377,7 @@ pub struct SegmentCheckpointComplete { pub checkpointed_bytes: u64, } -pub struct OffsetCheckpointComplete { +pub struct AuxiliaryCheckpointComplete { pub checkpointed_lsn: u64, } @@ -366,7 +395,7 @@ impl_from_variant!( DataPlaneCommand, Produce, SegmentCheckpointComplete, - OffsetCheckpointComplete, + AuxiliaryCheckpointComplete, CatchUpReadComplete, OrphanGcCheck, CommitConsumerOffset, @@ -374,8 +403,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/messages/pending.rs b/src/data_plane/messages/pending.rs index 6ba1eab4..2c331a00 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}; @@ -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); @@ -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(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..f4e672b4 100644 --- a/src/data_plane/messages/query.rs +++ b/src/data_plane/messages/query.rs @@ -7,13 +7,14 @@ use std::sync::Arc; +use borsh::{BorshDeserialize, BorshSerialize}; use tokio::sync::oneshot; use crate::connections::protocol::{ConsumerOffsetGenerationMismatch, RangeProgressSignal}; 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; @@ -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 39556be9..20478da0 100644 --- a/src/data_plane/mod.rs +++ b/src/data_plane/mod.rs @@ -6,9 +6,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 recovery; pub(crate) mod segment_writer; @@ -28,12 +28,14 @@ pub struct SegmentKey { pub segment_id: SegmentId, } +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/recovery/inventory.rs b/src/data_plane/recovery/inventory.rs index 8a3ce46d..6c1c2861 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::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 @@ -70,13 +70,18 @@ impl LocalInventory { pub(crate) struct RecoveryOutput { pub(crate) inventory: LocalInventory, pub(crate) data_dir: PathBuf, - pub(crate) offsets: OffsetLedger, + 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)) @@ -126,7 +131,10 @@ mod tests { let output = RecoveryOutput { inventory: inv, data_dir: PathBuf::from("/var/lib/eastguard/data"), - offsets: OffsetLedger::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 f80b193f..72cd6029 100644 --- a/src/data_plane/recovery/mod.rs +++ b/src/data_plane/recovery/mod.rs @@ -17,7 +17,10 @@ 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::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; use crate::data_plane::sparse_index::SparseIndex; use crate::data_plane::states::segment::record::RoutingHeader; use crate::data_plane::wal::WalRecordType; @@ -56,7 +59,8 @@ struct Replaying { data_dir: PathBuf, writer: ReplayWriter, scanner: WalScanner, - offsets: OffsetLedger, + offsets: ConsumerOffsets, + producer: ProducerTracker, } impl Replaying { @@ -64,12 +68,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 +95,13 @@ impl Replaying { let (header, entry_data) = RoutingHeader::split_wal_payload(&record.payload)?; self.writer.replay(&header, entry_data)?; + if let Some(producer_identity) = header.producer_identity { + self.producer.advance( + header.segment_key(), + producer_identity, + header.entry_id, + ); + } } WalRecordType::ConsumerOffset => { let offset = OffsetRecord::try_from_slice(&record.payload) @@ -124,6 +136,7 @@ impl Replaying { scanner: self.scanner, recovered, offsets: self.offsets, + producer: self.producer, }) } } @@ -132,7 +145,8 @@ struct Durable { data_dir: PathBuf, scanner: WalScanner, recovered: RecoveredSegments, - offsets: OffsetLedger, + offsets: ConsumerOffsets, + producer: ProducerTracker, } impl Durable { @@ -154,12 +168,16 @@ 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, + auxiliary_state: AuxiliaryStateManager::from_recovery(self.offsets, self.producer), }) } } @@ -174,11 +192,12 @@ 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::messages::AuthorizedProducerIdentity; use crate::data_plane::wal::{WalRecord, WalStorage, WalWriter}; + use crate::data_plane::{ProduceError, ProducerAppendIdentity, SegmentKey}; /// One bare-payload `(Data, BatchEnd)` batch — the segment-file framing the /// checkpoint worker writes (no routing header). @@ -199,6 +218,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()) } @@ -308,12 +338,52 @@ 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.auxiliary_state.durable_generation(&key), 4); + assert_eq!(restarted.auxiliary_state.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 = 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(); - assert_eq!(*restarted.offsets.generation(&key), 4); - assert_eq!(restarted.offsets.offset(&key), Some(position)); + let mut restarted_tracker = restarted.auxiliary_state; + + assert_eq!( + restarted_tracker.verify_producer( + key, + AuthorizedProducerIdentity::ExistingOnly(producer), + 1, + ), + Err(ProduceError::Duplicate(EntryId(7))), + "the auxiliary snapshot must recover producer sessions after WAL reclamation" + ); } #[test] diff --git a/src/data_plane/state.rs b/src/data_plane/state.rs index c3989d46..c3f7e3c3 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, ); } } @@ -1264,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) { @@ -1356,7 +1402,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 +1518,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 +1526,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 +1555,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 +1571,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 +1592,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 +1643,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 +1662,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 +1767,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 +1957,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 +2011,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 +2086,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 +2213,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 +2245,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 +2263,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 +2297,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 +2309,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 +2386,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 +2444,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 +2524,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 +2534,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 +2749,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 +2769,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 +2792,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 +2822,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 +2859,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 +2888,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 +3163,7 @@ mod tests { data: b"setup".to_vec().into(), record_count: 1, entry_id: EntryId(0), + producer_identity: None, } .into(), )); @@ -2991,6 +3189,7 @@ mod tests { data: b"data".to_vec().into(), record_count: 1, entry_id: EntryId(0), + producer_identity: None, } .into(), )); @@ -3016,6 +3215,7 @@ mod tests { data: b"data".to_vec().into(), record_count: 1, entry_id: EntryId(0), + producer_identity: None, }, )); @@ -3038,6 +3238,7 @@ mod tests { data: b"data".to_vec().into(), record_count: 1, entry_id: EntryId(0), + producer_identity: None, } .into(), )); @@ -3058,6 +3259,7 @@ mod tests { data: b"data".to_vec().into(), record_count: 1, entry_id: EntryId(0), + producer_identity: None, } .into(), )); @@ -3181,6 +3383,7 @@ mod tests { data: b"data".to_vec().into(), record_count: 1, entry_id: EntryId(0), + producer_identity: None, } .into(), )); @@ -3311,6 +3514,7 @@ mod tests { data: b"data".to_vec().into(), record_count: 1, entry_id: EntryId(0), + producer_identity: None, } .into(), )); @@ -3508,7 +3712,7 @@ mod tests { RecoveryOutput { inventory: empty_inventory(), data_dir: dir.path().to_path_buf(), - offsets: OffsetLedger::default(), + auxiliary_state: AuxiliaryStateManager::default(), }, ); @@ -3521,6 +3725,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 +3780,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 +3823,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 +3915,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 +3996,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 +4011,7 @@ mod tests { record_count: 1, entry_id: EntryId(4), lsn: 0, + producer_identity: None, })], next_offset: EntryId(5), })); @@ -4386,6 +4593,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 9087fdbe..16ced5ad 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_identity: 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_identity: 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,20 +373,18 @@ 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 .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)) )); } diff --git a/src/data_plane/states/segment/cache.rs b/src/data_plane/states/segment/cache.rs index 81a68291..62a0069f 100644 --- a/src/data_plane/states/segment/cache.rs +++ b/src/data_plane/states/segment/cache.rs @@ -7,14 +7,15 @@ use arc_swap::ArcSwapOption; use tokio::sync::Notify; use crate::control_plane::metadata::EntryId; -use crate::data_plane::EntryPayload; +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, + pub(crate) producer_identity: Option, } impl CachedEntry { @@ -321,16 +322,17 @@ 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, + producer_identity: None, }) } diff --git a/src/data_plane/states/segment/record.rs b/src/data_plane/states/segment/record.rs index 8808b295..23656e29 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::{PayloadBytes, 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,23 +146,35 @@ 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) } } 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, } impl StagedEntry { - pub(crate) fn new(data: EntryPayload, record_count: u32, segment_key: SegmentKey) -> Self { + pub(crate) fn new( + data: PayloadBytes, + 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"); + } } diff --git a/src/data_plane/states/segment/tracker.rs b/src/data_plane/states/segment/tracker.rs index 6a89f497..5aa13076 100644 --- a/src/data_plane/states/segment/tracker.rs +++ b/src/data_plane/states/segment/tracker.rs @@ -3,7 +3,8 @@ 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::{EntryPayload, SegmentKey, checkpoint::CheckpointJob, wal::WalRecord}; +use crate::data_plane::ProducerAppendIdentity; +use crate::data_plane::{PayloadBytes, 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_identity); 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_identity) + .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_identity: s.producer_identity, }); 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_identity, + ) + }) }) } @@ -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_identity)) } // 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() } @@ -261,9 +275,10 @@ impl SegmentTracker { pub(crate) fn stage_entry( &mut self, segment_key: SegmentKey, - data: EntryPayload, + data: PayloadBytes, 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_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_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_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(); @@ -425,6 +444,7 @@ pub mod tests { 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_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_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_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 778c2788..74f71f2b 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_entry(key, data, record_count, entry_id, producer_identity); + Some(entry_id) + } + pub(crate) fn values(&self) -> impl Iterator { self.by_key.values() } diff --git a/src/it/e2e/client_protocol.rs b/src/it/e2e/client_protocol.rs index fba3c8f8..637157b7 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, 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, }; @@ -1830,15 +1819,16 @@ 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, })); 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::DataPlane(DataPlaneResponse::ShardNotLocal { + ClientResponse::Err(ServerError::ShardNotLocal { hint_node: Some(addr), }) => ProduceOutcome::Redirect(addr.client_addr()), _ => ProduceOutcome::Retry, @@ -1851,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 eba16eaf..87b581df 100644 --- a/src/it/e2e/client_sdk.rs +++ b/src/it/e2e/client_sdk.rs @@ -13,16 +13,16 @@ 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; 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 @@ -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!( @@ -445,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); @@ -510,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); @@ -548,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); @@ -607,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); @@ -842,13 +848,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 +888,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 @@ -958,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] @@ -1892,7 +1957,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; @@ -1942,7 +2007,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; @@ -1956,8 +2028,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/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/e2e/producer_deduplication.rs b/src/it/e2e/producer_deduplication.rs new file mode 100644 index 00000000..b1f3ca5f --- /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, ServerError, StoragePolicy}; +use crate::connections::protocol::{ + ClientDataPlaneRequest, ClientRequest, ClientResponse, OpenProducerSessionRequest, + ProduceRequest, +}; +use crate::data_plane::{PayloadBytes, 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]) -> 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.into() +} + +#[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_identity: Some(unknown), + })); + let mut rejected = false; + for (host, port, _) in NODES { + if matches!( + send_request(host, port, unknown_wire.clone()).await, + ClientResponse::Err(ServerError::ProduceRejected( + ProduceError::SessionNotInstalled + )) + ) { + 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() +} 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 2ea4f519..3584ac50 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; @@ -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; } @@ -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(_) => {} @@ -401,13 +398,14 @@ 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, })); 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(_) => {} @@ -440,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) => @@ -554,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; } 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 +}