diff --git a/CHANGELOG.md b/CHANGELOG.md index e4441724..1698c2aa 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -33,6 +33,7 @@ All notable changes to this project will be documented in this file. - **[PR #81]** Add data models for state machine ### Data Plane & Storage +- Graduate consumer-offset replicas through durable range snapshots before they can serve reads, while allowing joining replicas to receive live commits without delaying client acknowledgement. - **[PR #136]** feat: 135 - **[PR #132]** feat: d4 - cold path - **[PR #126]** feat: d3 - node membership reconcilation diff --git a/docs/data-plane/d9_offset_placement_graduation.md b/docs/data-plane/d9_offset_placement_graduation.md new file mode 100644 index 00000000..168105e8 --- /dev/null +++ b/docs/data-plane/d9_offset_placement_graduation.md @@ -0,0 +1,328 @@ +# D9: Offset Placement Graduation + +**Goal:** Ensure that when a range's ownership moves to new replicas, the new replicas do not +act as authoritative offset sources until they have safely copied the current consumer +checkpoints. This prevents checkpoints from temporarily appearing absent during the move. + +**Depends on:** [D8: Consumer Offset Management](d8_consumer_offset_management.md), +[D2: Segment Replication](d2_segment_replication.md), and +[C4: Consumer Offset Commits](../clients/c4_consumer_offset_commits.md). + +--- + +## The Gap (The Problem) + +Consumer offsets (checkpoints of where a consumer group is reading) are attached to a "range", not a specific segment. When a segment gets full and we "roll" to a new segment, we might decide to store that new segment on different replica nodes. However, the offsets themselves stay with the range. + +Imagine a range's data is placed like this: + +```text +old segment: A B C +new segment (roll): A B D +``` + +Before the roll, a checkpoint was saved safely on A, B, and C. During the roll, we tell A, B, and D that they are the new replica set. But we *don't* automatically copy the old checkpoint over to D. + +Immediately after this change: + +| Node | Status | Does it know it's a replica? | Does it have the old checkpoint? | +|---|---|---|---| +| A | Kept | Yes | Yes | +| B | Kept | Yes | Yes | +| C | Removed | Yes (but stale) | Yes | +| D | Added | Yes | **No (Missing)** | + +If a consumer makes a brand new commit, D will get it and be up to date. But if the consumer +is inactive, D never gets the old checkpoint. If A or B later fail, D might be asked to serve +an authoritative offset read, but it is missing the checkpoint. We need a way to bring D up +to date before we trust it. + +--- + +## Placement Is Not Readiness + +Just because a node (like D) is told "you are in the replica set", doesn't mean it has all the historical data it needs. Being *placed* in the set is different from being *ready* to serve. + +We define explicit stages for a new replica: + +```text + Placement Learned (Told it's a replica) + │ + ▼ + Joining + (Accepts incoming historical data, + but is NOT trusted yet) + │ + Durable Bootstrap Complete (Finished copying to disk) + │ + ▼ + Ready + (Trusted to serve offset data) +``` + +This strict graduation process is only for offset authority. Standard segment data continues to use its normal replication rules. + +--- + +## Range-Offset Bootstrap (How we copy the data) + +When a roll adds a new replica, an existing, fully "ready" replica takes charge. It transfers +the complete current ledger slice for that range to the new replica: the latest generation +and committed position for every known consumer group, including groups that have not been +active recently. The ledger does not retain the history of every earlier commit. + +```text +Control Plane A (Ready Source) D (Joining) + │ │ │ + │── 1. Tells A the new set ─►│ │ + │── 1. Tells D the new set ────────────────────────────────►│ + │ │ │ + │ │── 2. Sends current offsets ─►│ + │ │ │ 3. D saves to disk + │ │ │ + │ │◄── 4. D says "I am done" ────│ + │ │ │ + │ │ D becomes ready +``` + +The transfer (bootstrap) includes: + +- **Topic/Range Identity:** So D knows what data this is. +- **New Segment Identity:** So we ensure this data matches the current placement change. +- **Source/Destination:** To ignore wrong or bounced messages. +- **Group Generation & Position:** The actual checkpoint data. + +--- + +## Durable and Monotonic Import (Safe Conflict Resolution) + +Copying historical data (bootstrap) is a trusted internal transfer, not a brand new action +from a client. Therefore, we don't apply the normal "is this client generation up to date?" +checks that we do for normal commits. + +When the joining replica receives this historical data, it merges it **monotonically** (safely moving forward): + +- It only updates its stored generation if the incoming one is *greater*. +- It only updates its stored position if the incoming one is *greater*. +- If it receives the same transfer twice by mistake, nothing changes (harmless). +- If a live consumer makes a brand new commit at the same time, the new commit has a greater position, so it won't be overwritten by the older historical data. + +To survive crashes, the imported data and a special "bootstrap-completion marker" must both be written safely to the disk log (WAL). The joining replica only graduates to "ready" *after* this disk save is complete. If the node crashes halfway through, it wakes up, sees the data but no completion marker, and knows it needs to resume the import rather than blindly trusting incomplete data. + +The completion acknowledgement does not immediately make the replica ready. The leader also +waits for every ordinary offset commit already sent to that replica while it was joining to +acknowledge its WAL write. This closes the race where the snapshot becomes durable, a newer +commit is still in flight, and the replica crashes after being promoted but before storing +that commit. + +```text +bootstrap durable + │ + ▼ +bootstrap acknowledged + │ + ▼ +all earlier joining-period commits acknowledged + │ + ▼ +replica becomes ready +``` + +--- + +## Coordination and Retries + +Two leaders participate in different parts of graduation: + +| Role | Responsibility | +|---|---| +| Metadata Raft leader | Commits the successor placement and re-drives incomplete graduation | +| Range data leader | Serves offset reads and commits, and coordinates transfer from a ready source | + +The range data leader knows the active placement used by offset traffic. It only selects +fully ready replicas as bootstrap sources. The metadata leader treats bootstrap as +declarative reconciliation: it can announce the committed placement again after message +loss or leadership change without creating a second transition. + +The transfer process is self-healing: + +1. The control plane announces the new placement. New replicas become "joining". +2. A "ready" survivor sends the current range-ledger state to the joining replicas. +3. The joining replica saves it and sends back an acknowledgement. +4. If the acknowledgement is lost, the transfer is announced again and the source re-sends + the data. +5. If either leader changes, the successor reconstructs the work from committed placement + and durable readiness state rather than relying on the previous leader's memory. + +If all ready survivors are unavailable, offset reads fail closed rather than returning +incorrect empty checkpoints. Service can resume when a ready replica restarts or another +trusted copy is recovered and graduated. + +--- + +## Relationship to Segment Recovery + +Offset graduation reuses the control pattern already used to repair sealed segment data: + +| Segment repair | Offset graduation | +|---|---| +| Announce required placement | Announce required offset placement | +| Select a surviving source | Select a ready surviving source | +| Transfer missing segment bytes | Transfer the current range-ledger slice | +| Verify and durably register the segment | Monotonically merge through the shared WAL | +| Acknowledge durable catch-up | Acknowledge durable graduation | +| Re-drive until confirmed | Re-drive until confirmed | + +The payload and completion check remain separate. Segment repair copies immutable bytes and +verifies a committed end; offset graduation copies a mutable map of group generations and +positions while ordinary commits may still arrive. Sharing the reconciliation shape avoids +a second recovery lifecycle without pretending the two kinds of state have the same merge +rules. + +--- + +## Leadership Changes + +Bootstrap is fenced by the successor segment identity, which acts as the placement token for +the transition. Every transfer and durable completion marker carries this token. A delayed +message may still contribute a newer checkpoint through monotonic merge, but it cannot mark +a replica ready for a different placement. + +### Metadata leader change + +A metadata leadership change does not alter the committed range placement. The new leader +re-drives the same placement token. A joining replica resumes its transfer; a replica that +already stored the completion marker answers again without rewriting its state. + +```text +old metadata leader fails + │ +new leader reads committed placement + │ +re-announces the same graduation token + │ +joining replica resumes, or ready replica re-acknowledges +``` + +This keeps completion recoverable without placing transient transfer progress in the Raft +log. + +### Range data leader change + +Offset leadership follows the first member of the active data replica set. When that leader +fails, recovery rolls the active segment again and places a suitable surviving replica first. +That second roll creates a new successor identity and therefore supersedes the earlier +bootstrap attempt. + +```text +S2 placement: A (ready leader), B (ready), D (joining) +A fails +S3 placement: B (ready leader), D (joining), E (joining) +``` + +B continues serving offsets and bootstraps the joining replicas for S3. A delayed completion +for S2 cannot graduate D for S3. The conservative first implementation bootstraps D again +under the S3 token; a later optimization may reuse D's durable state after proving it covers +the required boundary. + +A joining replica must never become authoritative merely because it appears first in a new +placement. Recovery should choose a ready survivor first whenever one exists. If none exists, +offset reads remain unavailable until a replica is safely graduated. + +--- + +## Reads and Commits During Graduation + +While the new replicas are catching up, the system keeps working: + +- The existing "ready" leader continues serving read requests. +- New offset commits are still accepted and sent to every placed replica. +- Client success waits for every ready replica, but not for replicas that are still joining. +- Joining-replica acknowledgements are tracked so graduation waits for their in-flight commits + to drain. + +**A joining replica CAN:** + +- Receive new commits. +- Receive historical bootstrap data and save it. +- Acknowledge that it saved data. + +**A joining replica CANNOT:** + +- Tell a client that a checkpoint doesn't exist (because it might just not have downloaded it yet). +- Become the leader for reads. +- Act as a source to send data to someone else. +- Graduate to "ready" before the completion marker is safely on disk. + +If a client accidentally asks a joining replica for data, the replica will return a "not ready" error, and the client will simply retry until it finds the true leader. + +--- + +## Failure Handling + +| What went wrong | How we fix it | +|---|---| +| Transfer message is lost | The source resends it. | +| The "I am done" message is lost | The source resends the data. The joining replica safely ignores the duplicate data. | +| Joining replica crashes before saving to disk | When it restarts, it sees no completion marker and stays in the "joining" state. | +| Joining replica crashes after saving to disk | When it restarts, it sees the completion marker and resumes as "ready". | +| A replica restarts in the same placement after missing a group-epoch seal | A replicated commit from the newer epoch stays parked and prompts the replica to request the current range snapshot from the leader. After the snapshot and completion marker are durable, the parked commit is applied and acknowledged. | +| The source replica dies during transfer | Another ready replica takes over and resends. | +| Metadata leader changes | The new leader re-drives the same committed placement token. | +| Range data leader changes | A new roll supersedes the old token and a ready survivor resumes coordination. | +| An old transfer arrives after another roll | Its data may merge monotonically, but it cannot complete the new token. | +| A new commit arrives at the exact same time | The monotonic merge ensures the newer (greater) commit is kept. | +| Bootstrap finishes while a commit is in flight | The replica stays joining until that commit acknowledges durability. | +| All ready replicas are unavailable | Reads fail closed until a ready replica returns or trusted state is recovered. | +| A consumer group is completely idle | The bootstrap still copies its old checkpoint safely. | + +--- + +## Scope Boundaries + +This process applies whenever a segment roll preserves the range identity but changes the +active replica set. The roll may be triggered normally by sealing or by node-failure +recovery. + +It does **not** handle copying offsets when ranges are split into two, or merged together. Those create brand new range identities and require clients to handle the transition. It also doesn't apply to background data repair. + +--- + +## Target Guarantees + +These are the strict goals we are aiming for: + +1. **No trusting incomplete replicas:** A replacement cannot serve authoritative offset + reads until its import is safely on disk. +2. **Idle groups are protected:** We copy checkpoints for all groups, not just active ones. +3. **Safe from network weirdness:** Delayed or out-of-order messages won't accidentally roll a checkpoint backward. +4. **Disk durability:** Graduation only happens after the data is physically on the disk. +5. **Only copy from the good ones:** We never use an incomplete replica as a source. +6. **When in doubt, fail closed:** A joining replica must never return an empty checkpoint if it isn't sure. +7. **Joining replicas do not delay client commits:** Commits are sent to them, but only ready + replicas are required for client acknowledgement until graduation. + +--- + +## Implementation and Verification Status + +The graduation path is implemented. Range-ledger snapshots, placement completion, and ordinary +offset commits share the data-plane WAL, so one successful flush durably orders the imported +state before its acknowledgement. Placement announcements are declarative and re-driven until +all members graduate, while segment identity fences acknowledgements from older rolls. + +Current automated coverage verifies: + +- placement ordering, conflicting announcements, and durable readiness recovery; +- source eligibility, including rejection after an intervening placement; +- snapshot durability before acknowledgement and monotonic import; +- commits concurrent with bootstrap, including the acknowledgement drain before graduation; +- fail-closed behavior for an unready leader; +- recovery when a restarted replica missed a consumer-group epoch while it was offline; and +- deterministic end-to-end consumer-group operation across that restart. + +The remaining coverage gap is a deterministic full-system test that forces a roll onto a changed +replica set and then reads an idle group's previously committed checkpoint from the replacement. +The manager-level tests exercise that placement transition and preserve the idle ledger entry, but +the end-to-end suite currently covers the related restart/missed-epoch path rather than forcing the +complete replacement topology. diff --git a/src/client/routing.rs b/src/client/routing.rs index 42e3690b..6805fa69 100644 --- a/src/client/routing.rs +++ b/src/client/routing.rs @@ -24,7 +24,7 @@ struct RangeRoute { /// `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, - replicas: Box<[NodeAddressInfo]>, + replica_addrs: Box<[NodeAddressInfo]>, } // Active segment only — correct for produce (the write head) and a tailing consumer. @@ -48,7 +48,7 @@ impl From<&RangeDetail> for RangeRoute { .active_segment .as_ref() .and_then(|_| replicas.first().cloned()), - replicas, + replica_addrs: replicas, } } } @@ -77,14 +77,14 @@ impl TopicRouting { self.ranges .iter() .find(|r| r.range_id == range_id) - .map(|r| r.replicas.as_ref()) + .map(|r| r.replica_addrs.as_ref()) } pub(crate) fn write_leader_for_range(&self, range_id: RangeId) -> Option { self.ranges .iter() .find(|range| range.range_id == range_id) - .and_then(|range| range.replicas.first().cloned()) + .and_then(|range| range.replica_addrs.first().cloned()) } } @@ -174,7 +174,7 @@ impl RoutingCache { range_id: r.range_id, keyspace_start: r.keyspace_start.clone(), write_leader: r.write_leader.as_ref().map(|_| wrong.clone()), - replicas: r.replicas.clone(), + replica_addrs: r.replica_addrs.clone(), }) .collect(); diff --git a/src/connections/controller.rs b/src/connections/controller.rs index 2b518440..b3a801ea 100644 --- a/src/connections/controller.rs +++ b/src/connections/controller.rs @@ -232,7 +232,7 @@ impl ClientController { let cmd = CreateTopic { storage_policy, name, - replica_set: group.members, + replica_set: group.replicas, created_at, }; Ok(self @@ -598,7 +598,7 @@ impl ClientController { shard_group_id: group.id, leader_node_id: leader.as_ref().map(|e| e.leader.node_id.to_string()), leader_addr: leader.map(|e| e.leader.client_addr()), - member_node_ids: group.members.iter().map(|n| n.to_string()).collect(), + member_node_ids: group.replicas.iter().map(|n| n.to_string()).collect(), }); Ok(ClientResponse::Admin(AdminResponse::ShardInfo { detail })) } @@ -664,7 +664,7 @@ mod tests { use crate::control_plane::metadata::TopicStats as MetadataTopicStats; use crate::control_plane::metadata::strategy::{PartitionStrategy, StoragePolicy}; use crate::control_plane::metadata::{RangeId, TopicId, TopicMeta}; - use crate::control_plane::{NodeAddress, NodeId, SwimNode, SwimNodeState}; + use crate::control_plane::{NodeAddress, NodeId, Replicas, SwimNode, SwimNodeState}; use crate::data_plane::actor::DataPlaneSender; use crate::data_plane::messages::DataPlaneMessage; use crate::data_plane::messages::command::{DataPlaneCommand, ProduceAck}; @@ -681,7 +681,7 @@ mod tests { fn test_shard_group() -> ShardGroup { ShardGroup { id: ShardGroupId(42), - members: vec![node_id("node-1")], + replicas: Replicas::new(vec![node_id("node-1")]), } } @@ -745,7 +745,7 @@ mod tests { TopicMeta::new( "t1".into(), TopicId(1), - vec![node_id(leader)], + Replicas::new(vec![node_id(leader)]), 0, StoragePolicy { retention_ms: Some(3_600_000), @@ -785,7 +785,7 @@ mod tests { let owner_addr = NodeAddress::test(addr(18003), addr(8083)); let group = ShardGroup { id: ShardGroupId(42), - members: vec![owner.clone()], + replicas: Replicas::new(vec![owner.clone()]), }; let swim = swim_sender_with(move |cmd| match cmd { SwimActorCommand::Query(QueryCommand::ResolveShardGroup { reply, .. }) => { @@ -815,7 +815,7 @@ mod tests { let leader_addr = NodeAddress::test(addr(18004), addr(8084)); let group = ShardGroup { id: ShardGroupId(1), - members: vec![me.clone(), leader.clone()], + replicas: Replicas::new(vec![me.clone(), leader.clone()]), }; let swim = swim_sender_with(move |cmd| match cmd { SwimActorCommand::Query(QueryCommand::ResolveShardGroup { reply, .. }) => { @@ -849,7 +849,7 @@ mod tests { let me = node_id("self"); let group = ShardGroup { id: ShardGroupId(1), - members: vec![me.clone()], + replicas: Replicas::new(vec![me.clone()]), }; let swim = swim_sender_with(move |cmd| { if let SwimActorCommand::Query(QueryCommand::ResolveShardGroup { reply, .. }) = cmd { @@ -878,7 +878,7 @@ mod tests { let owner_addr = NodeAddress::test(addr(18005), addr(8085)); let group = ShardGroup { id: ShardGroupId(9), - members: vec![owner.clone()], + replicas: Replicas::new(vec![owner.clone()]), }; let swim = swim_sender_with(move |cmd| match cmd { SwimActorCommand::Query(QueryCommand::ResolveShardGroup { reply, .. }) => { @@ -1030,7 +1030,7 @@ mod tests { let owner_addr = NodeAddress::test(addr(18002), addr(8082)); let group = ShardGroup { id: ShardGroupId(42), - members: vec![owner.clone()], + replicas: Replicas::new(vec![owner.clone()]), }; let swim = { let group = group.clone(); @@ -1071,7 +1071,7 @@ mod tests { let me = node_id("self"); let group = ShardGroup { id: ShardGroupId(7), - members: vec![me.clone()], + replicas: Replicas::new(vec![me.clone()]), }; let swim = swim_sender_with(move |cmd| { if let SwimActorCommand::Query(QueryCommand::ResolveShardGroup { reply, .. }) = cmd { diff --git a/src/control_plane/consensus/messages/actor.rs b/src/control_plane/consensus/messages/actor.rs index 75447d1d..7124a5d3 100644 --- a/src/control_plane/consensus/messages/actor.rs +++ b/src/control_plane/consensus/messages/actor.rs @@ -5,10 +5,12 @@ use crate::control_plane::NodeId; use crate::control_plane::consensus::raft::errors::ProposalError; use crate::control_plane::membership::ShardGroupId; use crate::control_plane::metadata::{ConsumerGroupAssignment, TopicMeta, TopicStats}; -use crate::data_plane::messages::command::{CatchUpAck, SealBoundaryReport, SegmentAssignmentAck}; +use crate::data_plane::messages::command::{ + DurableSegmentEndReported, SegmentCaughtUp, SegmentPlaced, +}; use super::command::{ - CoordinatorSealRequest, EnsureGroup, InboundRaftRpc, MetadataProposal, RaftProtocolMessage, + EnsureGroup, InboundRaftRpc, MetadataProposal, ProposeSegmentRoll, RaftProtocolMessage, RemoveGroup, }; use super::timer::RaftTimeoutCallback; @@ -51,18 +53,18 @@ pub enum MultiRaftActorCommand { reply: oneshot::Sender>, }, GetConsumerGroupAssignment(GetConsumerGroupAssignment), - /// Data plane SealRequest forwarded to coordinator for Raft proposal. - Coordinator(CoordinatorSealRequest), - /// Data-leader confirmation that it received a `SegmentAssignment`. Marks the + /// Data-plane request forwarded to the metadata coordinator for proposal. + ProposeSegmentRoll(ProposeSegmentRoll), + /// Data-leader confirmation that it received a `PlaceSegment`. Marks the /// segment confirmed so the leader's heartbeat sweep stops re-driving it. - AssignmentAck(SegmentAssignmentAck), - /// A survivor's reply to a leader-crash `SealBoundaryQuery` — its durable extent + AssignmentAck(SegmentPlaced), + /// A survivor's reply to a leader-crash `RequestDurableSegmentEnd` — its durable extent /// for the segment, gathered to recover the committed seal end. - SealBoundaryReport(SealBoundaryReport), + DurableSegmentEndReported(DurableSegmentEndReported), /// A replica's confirmation that it holds a reassigned sealed segment through /// `sealed_end`. Clears the member from the coordinator's catch-up re-drive so /// the heartbeat sweep stops re-announcing the assignment. - CatchUpAck(CatchUpAck), + SegmentCaughtUp(SegmentCaughtUp), } pub struct GetConsumerGroupAssignment { diff --git a/src/control_plane/consensus/messages/command.rs b/src/control_plane/consensus/messages/command.rs index 620df58e..0429d343 100644 --- a/src/control_plane/consensus/messages/command.rs +++ b/src/control_plane/consensus/messages/command.rs @@ -3,7 +3,7 @@ use crate::control_plane::consensus::messages::rpc::{OutboundRaftPacket, RaftRpc use crate::control_plane::consensus::messages::timer::RaftTimeoutCallback; use crate::control_plane::membership::{ShardGroup, ShardGroupId}; use crate::control_plane::metadata::command::MetadataCommand; -use crate::data_plane::messages::command::SealRequest; +use crate::data_plane::messages::command::RequestSegmentRoll; use crate::impl_from_variant; pub struct InboundRaftRpc { @@ -50,6 +50,6 @@ pub enum RaftTransportCommand { } #[derive(Debug)] -pub struct CoordinatorSealRequest { - pub request: SealRequest, +pub struct ProposeSegmentRoll { + pub request: RequestSegmentRoll, } diff --git a/src/control_plane/consensus/messages/event.rs b/src/control_plane/consensus/messages/event.rs index 12114f57..8ae59bc2 100644 --- a/src/control_plane/consensus/messages/event.rs +++ b/src/control_plane/consensus/messages/event.rs @@ -48,11 +48,11 @@ pub enum RaftEvent { DisconnectPeer(NodeId), MetadataCommitted(MetadataCommitted), /// Idempotent re-delivery of assignment messages each heartbeat, so a lost - /// fire-and-forget send self-heals: active-segment `SegmentAssignment`s and - /// sealed-segment `CatchUpAssignment`s. The actor forwards them to the data + /// fire-and-forget send self-heals: active-segment `PlaceSegment`s and + /// sealed-segment `AssignSegmentCatchUp`s. The actor forwards them to the data /// transport. RedriveAssignments(Vec), - /// Leader-crash `SealBoundaryQuery` fan-out to a segment's survivors. + /// Leader-crash `RequestDurableSegmentEnd` fan-out to a segment's survivors. /// A coordinator-initiated data-plane send not tied to a committed entry. /// The actor just forwards these to the data transport. SealBoundaryQueries(Vec), diff --git a/src/control_plane/consensus/multi_raft.rs b/src/control_plane/consensus/multi_raft.rs index d7558abc..c4676325 100644 --- a/src/control_plane/consensus/multi_raft.rs +++ b/src/control_plane/consensus/multi_raft.rs @@ -1,8 +1,7 @@ use crate::control_plane::NodeId; use crate::control_plane::consensus::messages::{ - CoordinatorSealRequest, DeferredConsumerGroupAssignment, DeferredReply, InboundRaftRpc, - LogMutation, MetadataProposal, MultiRaftActorCommand, RaftEvent, RaftProtocolMessage, - RaftTimeoutCallback, + DeferredConsumerGroupAssignment, DeferredReply, InboundRaftRpc, LogMutation, MetadataProposal, + MultiRaftActorCommand, ProposeSegmentRoll, RaftEvent, RaftProtocolMessage, RaftTimeoutCallback, }; use crate::control_plane::consensus::raft::errors::ProposalError; use crate::control_plane::consensus::raft::state::{LeaderlessSegments, Raft, TimerSeqs}; @@ -16,7 +15,7 @@ use crate::control_plane::metadata::{ }; use crate::data_plane::SegmentKey; use crate::data_plane::messages::command::{ - CatchUpAck, SealBoundaryQuery, SealBoundaryReport, SegmentAssignmentAck, + DurableSegmentEndReported, RequestDurableSegmentEnd, SegmentCaughtUp, SegmentPlaced, }; use crate::data_plane::transport::command::DataTransportCommand; use std::collections::{BTreeMap, BTreeSet, HashMap, HashSet}; @@ -104,7 +103,7 @@ impl PendingSealTracker { /// Drop every pending seal context for `group_id`. Called when a leader /// steps down — any not-yet-committed proposals will either be truncated /// by the new leader or commit there, but this replica will no longer be - /// the one to dispatch the SealResponse. + /// the one to dispatch `SegmentRollCommitted`. fn drop_group(&mut self, group_id: ShardGroupId) { if let Some(stale_group) = self.by_group.remove(&group_id) { for seal in stale_group.into_values() { @@ -147,7 +146,7 @@ impl MultiRaft { pub(crate) fn reconcile_on_leadership_change(&mut self, shard_group_id: ShardGroupId) { let target_members = self.topology.group_ring_members(shard_group_id); - self.reconcile_membership(shard_group_id, target_members); + self.reconcile_membership(shard_group_id, target_members.map(|e| e.0)); // Takeover backstop: the catch-up tracker is leader-volatile, so a repair // in flight when leadership changed left it empty here. Re-seed it; the @@ -265,16 +264,16 @@ impl MultiRaft { }, )); } - MultiRaftActorCommand::Coordinator(req) => { - self.handle_seal_request(req); + MultiRaftActorCommand::ProposeSegmentRoll(cmd) => { + self.propose_segment_roll(cmd); } MultiRaftActorCommand::AssignmentAck(ack) => { - self.handle_assignment_ack(ack); + self.handle_segment_placed(ack); } - MultiRaftActorCommand::SealBoundaryReport(report) => { + MultiRaftActorCommand::DurableSegmentEndReported(report) => { self.handle_seal_boundary_report(report); } - MultiRaftActorCommand::CatchUpAck(ack) => { + MultiRaftActorCommand::SegmentCaughtUp(ack) => { self.handle_catch_up_ack(ack); } } @@ -323,12 +322,12 @@ impl MultiRaft { if self.groups.contains_key(&group.id) { return; } - if !group.members.contains(&self.node_id) { + if !group.replicas.contains(&self.node_id) { return; } let peers: HashSet = group - .members + .replicas .iter() .filter(|id| *id != &self.node_id) .cloned() @@ -428,11 +427,11 @@ impl MultiRaft { } } - fn handle_assignment_ack(&mut self, ack: SegmentAssignmentAck) { + fn handle_segment_placed(&mut self, ack: SegmentPlaced) { let Some(raft) = self.groups.get_mut(&ack.shard_group_id) else { return; }; - raft.handle_assignment_ack(ack); + raft.handle_segment_placed(ack); } fn get_leader(&self, group_id: ShardGroupId) -> Option { @@ -512,7 +511,7 @@ impl MultiRaft { } } - fn handle_seal_request(&mut self, req: CoordinatorSealRequest) { + fn propose_segment_roll(&mut self, req: ProposeSegmentRoll) { let seal = &req.request; if self.pending_seals.contains(&seal.segment_key) { return; @@ -520,7 +519,7 @@ impl MultiRaft { let Some(group_id) = self.find_group_for_topic(&seal.segment_key.topic_id) else { tracing::warn!( - "SealRequest for unknown topic {:?}", + "segment roll requested for unknown topic {:?}", seal.segment_key.topic_id ); return; @@ -530,7 +529,10 @@ impl MultiRaft { }; let Some(old_replica_set) = raft.get_replica_set(&seal.segment_key) else { - tracing::warn!("SealRequest for unknown segment {:?}", seal.segment_key); + tracing::warn!( + "segment roll requested for unknown segment {:?}", + seal.segment_key + ); return; }; @@ -559,7 +561,7 @@ impl MultiRaft { self.dirty.insert(group_id); } Err(e) => { - tracing::debug!("SealRequest proposal rejected: {:?}", e); + tracing::debug!("segment roll proposal rejected: {:?}", e); } } } @@ -574,7 +576,7 @@ impl MultiRaft { /// Feed a survivor's durable extent into the subsystem; if it completes a /// gather, execute the resulting seal. - fn handle_seal_boundary_report(&mut self, report: SealBoundaryReport) { + fn handle_seal_boundary_report(&mut self, report: DurableSegmentEndReported) { match self .seal_recovery .record(report.segment_key, report.from, report.durable_end) @@ -590,8 +592,8 @@ impl MultiRaft { } /// Route a catch-up confirmation to the owning group's `Raft`, which clears the - /// member and prunes the repair once all confirm. Mirrors `handle_assignment_ack`. - fn handle_catch_up_ack(&mut self, ack: CatchUpAck) { + /// member and prunes the repair once all confirm. Mirrors `handle_segment_placed`. + fn handle_catch_up_ack(&mut self, ack: SegmentCaughtUp) { let Some(raft) = self.groups.get_mut(&ack.shard_group_id) else { return; }; @@ -644,8 +646,7 @@ impl MultiRaft { if let Some(leader) = recency_leader && let Some(pos) = new_replica_set.iter().position(|n| *n == leader) { - let node = new_replica_set.remove(pos); - new_replica_set.insert(0, node); + new_replica_set.change_leader(pos); } let cmd = RollSegment { @@ -668,7 +669,7 @@ impl MultiRaft { } let cmd = DataTransportCommand::send_to_targets( targets.to_vec(), - SealBoundaryQuery { + RequestDurableSegmentEnd { segment_key, coordinator: self.node_id.clone(), }, @@ -806,7 +807,7 @@ impl MultiRaft { // that would otherwise prune is leader-gated and stops firing). A new // leader re-derives: re-proposes the seal / re-polls the survivors. if !raft.is_leader() { - // pending_seals: its SealResponse won't come from here. + // pending_seals: its committed roll notification won't come from here. self.pending_seals.drop_group(*id); // seal-end gathers: we can't propose their recovery roll. self.seal_recovery.drop_group(*id); @@ -833,7 +834,7 @@ impl MultiRaft { } if let RaftEvent::MetadataCommitted(committed) = &mut event { - // For leader to send a SealResponse back to the right client. + // For the leader to send the committed roll back to the requester. committed.seal_context = self.pending_seals.pop_seal_context(id, committed.log_index); } @@ -931,6 +932,7 @@ impl MultiRaft { mod tests { use super::*; + use crate::control_plane::Replicas; use crate::control_plane::consensus::raft::storage::RaftPersistentState; use crate::control_plane::consensus::seal_recovery; use crate::impls::metadata_storage::MetadataStorage; @@ -944,7 +946,7 @@ mod tests { fn shard(id: u64, members: Vec) -> ShardGroup { ShardGroup { id: ShardGroupId(id), - members, + replicas: Replicas::new(members), } } @@ -1386,7 +1388,7 @@ mod tests { replication_factor: 3, partition_strategy: PartitionStrategy::AutoSplit, }, - replica_set: vec![node("n1"), node("n2"), node("n3")], + replica_set: Replicas::new(vec![node("n1"), node("n2"), node("n3")]), created_at: 1000, }); @@ -1458,7 +1460,7 @@ mod tests { replication_factor: 3, partition_strategy: PartitionStrategy::AutoSplit, }, - replica_set: vec![node("n1"), node("n2"), node("n3")], + replica_set: Replicas::new(vec![node("n1"), node("n2"), node("n3")]), created_at: 1000, }) }; @@ -1716,7 +1718,7 @@ mod tests { .find(|id| { new_topology .group(*id) - .is_some_and(|g| g.members.contains(&node("n4"))) + .is_some_and(|g| g.replicas.contains(&node("n4"))) }) .expect("some group of n1 must gain n4 after the join"); @@ -1775,7 +1777,7 @@ mod tests { .topology .shard_groups_for_node(&n1) .iter() - .find(|g| !g.members.contains(&node("n9"))) + .find(|g| !g.replicas.contains(&node("n9"))) .map(|g| g.id) .expect("some group of n1 must exclude n9"); @@ -1849,7 +1851,7 @@ mod tests { replication_factor: rf, partition_strategy: PartitionStrategy::AutoSplit, }, - replica_set, + replica_set: Replicas::new(replica_set), created_at: 1000, }), }) @@ -1857,9 +1859,9 @@ mod tests { store.flush(); // commit + persist + apply (single-node) } - /// Targets a `SealBoundaryQuery` for `seg` was fanned out to (sorted). + /// Targets a `RequestDurableSegmentEnd` for `seg` was fanned out to (sorted). fn seal_boundary_query_targets(events: &[RaftEvent], seg: SegmentKey) -> Vec { - use crate::data_plane::messages::command::DataPlaneInterNodeCommand; + use crate::data_plane::messages::command::DataPlanePeerMessage; let mut out = Vec::new(); for e in events { let RaftEvent::SealBoundaryQueries(cmds) = e else { @@ -1867,7 +1869,7 @@ mod tests { }; for cmd in cmds { if let DataTransportCommand::SendToTargets(s) = cmd - && let DataPlaneInterNodeCommand::SealBoundaryQuery(q) = &s.message + && let DataPlanePeerMessage::RequestDurableSegmentEnd(q) = &s.message && q.segment_key == seg { out.extend(s.targets.iter().cloned()); @@ -1922,12 +1924,12 @@ mod tests { let before = proposals_after_become_leader(&store).len(); // Survivors report durable extents; the second completes the gather. - store.handle_seal_boundary_report(SealBoundaryReport { + store.handle_seal_boundary_report(DurableSegmentEndReported { segment_key: seg0, from: node("y"), durable_end: Some(EntryId(50)), }); - store.handle_seal_boundary_report(SealBoundaryReport { + store.handle_seal_boundary_report(DurableSegmentEndReported { segment_key: seg0, from: node("z"), durable_end: Some(EntryId(40)), diff --git a/src/control_plane/consensus/raft/catch_up.rs b/src/control_plane/consensus/raft/catch_up.rs index d95d9b84..04cfee7f 100644 --- a/src/control_plane/consensus/raft/catch_up.rs +++ b/src/control_plane/consensus/raft/catch_up.rs @@ -1,18 +1,18 @@ //! Sealed-segment catch-up re-drive — a `Raft`-owned tracker. //! -//! A committed `ReassignSegment` dispatches its `CatchUpAssignment` once +//! A committed `ReassignSegment` dispatches its `AssignSegmentCatchUp` once //! (fire-and-forget). Seeded at apply and driven by the leader's heartbeat, this //! tracker re-drives the assignment to each member until it acks, so a lost //! message can't strand the repair. See `.claude/rules/raft-actor.md` #9. use std::collections::{HashMap, HashSet}; -use crate::control_plane::NodeId; use crate::control_plane::membership::ShardGroupId; use crate::control_plane::metadata::EntryId; use crate::control_plane::metadata::event::SegmentReassigned; +use crate::control_plane::{NodeId, Replicas}; use crate::data_plane::SegmentKey; -use crate::data_plane::messages::command::CatchUpAssignment; +use crate::data_plane::messages::command::AssignSegmentCatchUp; use crate::data_plane::transport::command::DataTransportCommand; /// One in-flight repair: re-drive each member in `pending` until it acks. @@ -21,7 +21,7 @@ use crate::data_plane::transport::command::DataTransportCommand; struct CatchUpRepair { start_entry_id: EntryId, sealed_end: EntryId, - replica_set: Vec, + replica_set: Replicas, pending: HashSet, } @@ -54,7 +54,7 @@ impl CatchUpRepairs { segment_key: SegmentKey, start_entry_id: EntryId, sealed_end: EntryId, - replica_set: Vec, + replica_set: Replicas, ) { let pending = replica_set.iter().cloned().collect(); self.repairs.insert( @@ -86,7 +86,7 @@ impl CatchUpRepairs { for member in &repair.pending { cmds.push(DataTransportCommand::send_to_targets( vec![member.clone()], - CatchUpAssignment { + AssignSegmentCatchUp { segment_key: *segment_key, shard_group_id, start_entry_id: repair.start_entry_id, @@ -113,8 +113,9 @@ impl CatchUpRepairs { #[cfg(test)] mod tests { use super::*; + use crate::control_plane::Replicas; use crate::control_plane::metadata::{RangeId, SegmentId, TopicId}; - use crate::data_plane::messages::command::DataPlaneInterNodeCommand; + use crate::data_plane::messages::command::DataPlanePeerMessage; fn node(id: &str) -> NodeId { NodeId::new(id) @@ -127,7 +128,7 @@ mod tests { segment_key: seg(), start_entry_id: 0.into(), sealed_end: sealed_end.map(EntryId), - new_replica_set: members.iter().map(|n| node(n)).collect(), + new_replica_set: Replicas::new(members.iter().map(|n| node(n)).collect()), } } /// Sorted targets of the re-drive commands. @@ -164,8 +165,8 @@ mod tests { let DataTransportCommand::SendToTargets(s) = &cmds[0] else { panic!("expected SendToTargets"); }; - let DataPlaneInterNodeCommand::CatchUpAssignment(a) = &s.message else { - panic!("expected CatchUpAssignment"); + let DataPlanePeerMessage::AssignSegmentCatchUp(a) = &s.message else { + panic!("expected AssignSegmentCatchUp"); }; assert_eq!(a.shard_group_id, ShardGroupId(7)); assert_eq!(a.sealed_end_entry_id, EntryId(100)); diff --git a/src/control_plane/consensus/raft/mod.rs b/src/control_plane/consensus/raft/mod.rs index 1286c791..81677f48 100644 --- a/src/control_plane/consensus/raft/mod.rs +++ b/src/control_plane/consensus/raft/mod.rs @@ -1,4 +1,4 @@ -use crate::control_plane::NodeId; +use crate::control_plane::{NodeId, Replicas}; pub(crate) mod catch_up; pub(crate) mod command; @@ -18,7 +18,7 @@ pub(crate) fn compute_replacement_replica_set( old: &[NodeId], dead_nodes: &[NodeId], live_nodes: &[NodeId], -) -> Vec { +) -> Replicas { let replication_factor = old.len(); let mut new_set: Vec = old .iter() @@ -33,8 +33,6 @@ pub(crate) fn compute_replacement_replica_set( new_set.push(candidate.clone()); } } - if new_set.is_empty() { - tracing::error!("All replicas dead for segment — empty replica set"); - } - new_set + + Replicas::new(new_set) } diff --git a/src/control_plane/consensus/raft/state.rs b/src/control_plane/consensus/raft/state.rs index 38567ba1..8e74df43 100644 --- a/src/control_plane/consensus/raft/state.rs +++ b/src/control_plane/consensus/raft/state.rs @@ -1,6 +1,5 @@ #![allow(dead_code)] -use crate::control_plane::NodeId; use crate::control_plane::consensus::messages::*; use crate::control_plane::consensus::raft::catch_up::CatchUpRepairs; use crate::control_plane::consensus::raft::command::RaftCommand; @@ -14,10 +13,11 @@ use crate::control_plane::metadata::event::ApplyResult; use crate::control_plane::metadata::state_machine::MetadataStateMachine; use crate::control_plane::metadata::{ ConsumerGroupAssignment, ConsumerMemberId, EntryId, MetadataCommand, RangeId, ReassignSegment, - ReplicaSet, RollSegment, SegmentId, TopicId, TopicMeta, TopicStats, + RollSegment, SegmentId, TopicId, TopicMeta, TopicStats, }; +use crate::control_plane::{NodeId, Replicas}; use crate::data_plane::SegmentKey; -use crate::data_plane::messages::command::{CatchUpAck, SegmentAssignment, SegmentAssignmentAck}; +use crate::data_plane::messages::command::{PlaceSegment, SegmentCaughtUp, SegmentPlaced}; use crate::data_plane::transport::command::DataTransportCommand; use crate::schedulers::ticker_message::TimerCommand; #[cfg(any(test, debug_assertions))] @@ -126,9 +126,9 @@ pub struct Raft { election_epoch: u64, election_jitter: ElectionJitter, /// Segments whose data-leader has acked its - /// `SegmentAssignment`, mapped to the acking node. The heartbeat sweep skips + /// `PlaceSegment`, mapped to the acking node. The heartbeat sweep skips /// re-driving a segment whose confirmed node still matches `replica_set[0]` - confirmed_assignment: HashMap, + confirmed_placement: HashMap, /// In-flight sealed-segment repairs. Seeded at `ReassignSegment` apply; the /// heartbeat sweep re-drives until acked. Leader-volatile — cleared on step-down. catch_up: CatchUpRepairs, @@ -187,7 +187,7 @@ impl Raft { election_jitter: ElectionJitter::new(election_jitter_seed), timer_seqs, election_epoch: 0, - confirmed_assignment: HashMap::new(), + confirmed_placement: HashMap::new(), catch_up: CatchUpRepairs::default(), ring_observation_streak: None, }; @@ -224,14 +224,14 @@ impl Raft { pub(crate) fn active_segments_for_node( &self, node_id: &NodeId, - ) -> Box<[(SegmentKey, ReplicaSet)]> { + ) -> Box<[(SegmentKey, Replicas)]> { self.state_machine.active_segments_for_node(node_id) } /// Every active segment's assignment tuple `(key, replica_set, start_offset)`. /// The leader's confirmation-gated assignment sweep (`MultiRaft::build_redrive_cmds`) - /// turns these into `SegmentAssignment` re-drives for unconfirmed segments. - pub(crate) fn active_segment_assignments(&self) -> Box<[(SegmentKey, ReplicaSet, EntryId)]> { + /// turns these into `PlaceSegment` re-drives for unconfirmed segments. + pub(crate) fn active_segment_assignments(&self) -> Box<[(SegmentKey, Replicas, EntryId)]> { self.state_machine.active_segment_assignments() } @@ -502,7 +502,7 @@ impl Raft { /// /// Runs per-death (`handle_node_death`) and as the takeover backfill sweep. pub(crate) fn reconcile_segments(&mut self, live_set: &HashSet) -> bool { - let mut to_roll: Vec<(SegmentKey, ReplicaSet)> = Vec::new(); + let mut to_roll: Vec<(SegmentKey, Replicas)> = Vec::new(); for (key, rs) in self.active_segments_with_dead_members(live_set).into_vec() { if Self::only_replica_dead(&rs, live_set) { let survivors = rs @@ -561,12 +561,12 @@ impl Raft { fn repair_segments( &mut self, - segments: Box<[(SegmentKey, ReplicaSet)]>, + segments: Box<[(SegmentKey, Replicas)]>, live_set: &HashSet, build_cmd: F, ) -> bool where - F: Fn(SegmentKey, Vec) -> MetadataCommand, + F: Fn(SegmentKey, Replicas) -> MetadataCommand, { let live_nodes: Vec = live_set.iter().cloned().collect(); let mut changed = false; @@ -593,7 +593,7 @@ impl Raft { /// Re-fill known-end sealed segments left under-replicated by an earlier death fn refill_under_replicated_segments(&mut self, topology: &TopologyReader) -> bool { - let targets: Box<[(SegmentKey, Vec)]> = self + let targets: Box<[(SegmentKey, Replicas)]> = self .state_machine .topics .values() @@ -609,7 +609,8 @@ impl Raft { continue; // ring can't grow it yet — retried on the next ring check } - replica_set.extend(additions.iter().cloned()); + replica_set.extend(&additions); + let cmd = ReassignSegment { segment_key, replica_set, @@ -711,7 +712,7 @@ impl Raft { pub(crate) fn active_segments_with_dead_members( &self, live: &HashSet, - ) -> Box<[(SegmentKey, ReplicaSet)]> { + ) -> Box<[(SegmentKey, Replicas)]> { self.state_machine .topics .values() @@ -724,7 +725,7 @@ impl Raft { pub(crate) fn sealed_segments_with_dead_members( &self, live: &HashSet, - ) -> Box<[(SegmentKey, ReplicaSet)]> { + ) -> Box<[(SegmentKey, Replicas)]> { self.state_machine .topics .values() @@ -734,7 +735,7 @@ impl Raft { pub(crate) fn boundary_unknown_segments( &self, - ) -> impl Iterator { + ) -> impl Iterator)> { self.state_machine .topics .values() @@ -751,7 +752,7 @@ impl Raft { self.state_machine.get_topic(topic_id).is_some() } - pub(crate) fn get_replica_set(&self, key: &SegmentKey) -> Option { + pub(crate) fn get_replica_set(&self, key: &SegmentKey) -> Option { let topic = self.state_machine.get_topic(&key.topic_id)?; let range = topic.ranges.get(&key.range_id)?; let seg = range.segments.get(&key.segment_id)?; @@ -1168,7 +1169,7 @@ impl Raft { self.current_leader = None; self.peer_states.clear(); self.learner_states.clear(); - self.confirmed_assignment.clear(); + self.confirmed_placement.clear(); self.catch_up.clear(); } @@ -1196,7 +1197,7 @@ impl Raft { // rederive Metadata <> Datanode segment assignment let active = self.active_segment_assignments(); let active_keys: HashSet = active.iter().map(|(k, _, _)| *k).collect(); - self.confirmed_assignment + self.confirmed_placement .retain(|k, _| active_keys.contains(k)); let mut redrives = Vec::new(); @@ -1205,14 +1206,14 @@ impl Raft { continue; }; - // * If data leader acks assignment, it would have been added to confirmed_assignment through Raft::handle_assignment_ack - if self.confirmed_assignment.get(&segment_key) == Some(target) { + // * If data leader acks assignment, it would have been added to confirmed_placement through Raft::handle_segment_placed + if self.confirmed_placement.get(&segment_key) == Some(target) { continue; } redrives.push(DataTransportCommand::send_to_targets( vec![target.clone()], - SegmentAssignment { + PlaceSegment { segment_key, shard_group_id: self.shard_group_id, replica_set, @@ -1225,8 +1226,8 @@ impl Raft { } } - pub(crate) fn handle_assignment_ack(&mut self, ack: SegmentAssignmentAck) { - self.confirmed_assignment.insert(ack.segment_key, ack.from); + pub(crate) fn handle_segment_placed(&mut self, ack: SegmentPlaced) { + self.confirmed_placement.insert(ack.segment_key, ack.from); } /// Re-drive the catch-up sweep — the sealed-segment analogue of @@ -1240,7 +1241,7 @@ impl Raft { /// A member confirmed it holds a reassigned sealed segment; routed here from /// `MultiRaft`. - pub(crate) fn handle_catch_up_ack(&mut self, ack: CatchUpAck) { + pub(crate) fn handle_catch_up_ack(&mut self, ack: SegmentCaughtUp) { self.catch_up.confirm(ack.segment_key, ack.from); } @@ -1933,7 +1934,7 @@ impl crate::test_traits::TAssertInvariant for Raft { #[cfg(test)] mod tests { use super::*; - use crate::data_plane::messages::command::CatchUpAssignment; + use crate::data_plane::messages::command::AssignSegmentCatchUp; impl Raft { pub(crate) fn propose_noop(&mut self) -> Result { @@ -2922,7 +2923,7 @@ mod tests { replication_factor: 3, partition_strategy: PartitionStrategy::AutoSplit, }, - replica_set: vec![node("node-1"), node("node-2"), node("node-3")], + replica_set: Replicas::new(vec![node("node-1"), node("node-2"), node("node-3")]), created_at: 1000, }); @@ -3052,7 +3053,7 @@ mod tests { replication_factor: 3, partition_strategy: PartitionStrategy::AutoSplit, }, - replica_set: vec![node("n1"), node("n2"), node("n3")], + replica_set: Replicas::new(vec![node("n1"), node("n2"), node("n3")]), created_at: 1000, }) } @@ -3414,7 +3415,7 @@ mod tests { replication_factor: rf as u64, partition_strategy: PartitionStrategy::AutoSplit, }, - replica_set, + replica_set: Replicas::new(replica_set), created_at: 1000, }) .into(); @@ -3509,7 +3510,7 @@ mod tests { /// Leader over a real ring group plus the named extra "stale" voters. /// Ring peers have granted the election; nothing is acked yet, so /// commit_index is 0 until the test drives `ack_to_last`. - fn ring_raft_with_stale(stale_names: &[&str]) -> (Raft, TopologyReader, Vec) { + fn ring_raft_with_stale(stale_names: &[&str]) -> (Raft, TopologyReader, Replicas) { let reader = topology_reader_with(&["node-1", "node-2", "node-3"]); let group = reader .shard_groups_for_node(&node("node-1")) @@ -3518,7 +3519,7 @@ mod tests { .clone(); let mut peers: HashSet = group - .members + .replicas .iter() .filter(|m| **m != node("node-1")) .cloned() @@ -3540,7 +3541,7 @@ mod tests { epoch: u64::MAX, }); let term = raft.current_term(); - for peer in group.members.iter().filter(|m| **m != node("node-1")) { + for peer in group.replicas.iter().filter(|m| **m != node("node-1")) { raft.handle_rpc( peer.clone(), RaftRpc::RequestVoteResponse(RequestVoteResponse { @@ -3553,7 +3554,7 @@ mod tests { drain(&mut raft); assert_eq!(raft.role, Role::Leader, "must be leader after election"); - let members = group.members.clone(); + let members = group.replicas.clone(); (raft, reader, members) } @@ -3780,7 +3781,7 @@ mod tests { ); match &after[before] { RaftCommand::Metadata(MetadataCommand::RollSegment(roll)) => { - for member in &roll.new_replica_set { + for member in &roll.new_replica_set.0 { assert!( live.contains(member), "new replica_set must contain only live nodes, found {:?}", @@ -3836,7 +3837,7 @@ mod tests { MetadataCommand::RollSegment(RollSegment { segment_key: seg0, sealed_at: 2000, - new_replica_set: vec![node("node-1")], + new_replica_set: Replicas::new(vec![node("node-1")]), end_entry_id: Some(100.into()), }) .into(), @@ -3863,7 +3864,7 @@ mod tests { &after[before..] ); assert_eq!(reassigns[0].segment_key, seg0); - for member in &reassigns[0].replica_set { + for member in &reassigns[0].replica_set.0 { assert!( live.contains(member), "reassigned replica_set must be all-live, found {:?}", @@ -3898,7 +3899,11 @@ mod tests { MetadataCommand::RollSegment(RollSegment { segment_key: seg0, sealed_at: 2000, - new_replica_set: vec![node("node-1"), survivor.clone(), spare.clone()], + new_replica_set: Replicas::new(vec![ + node("node-1"), + survivor.clone(), + spare.clone(), + ]), end_entry_id: Some(100.into()), }) .into(), @@ -3966,7 +3971,7 @@ mod tests { MetadataCommand::RollSegment(RollSegment { segment_key: seg0, sealed_at: 2000, - new_replica_set: sealed_set, + new_replica_set: Replicas::new(sealed_set), end_entry_id: Some(100.into()), }) .into(), @@ -3976,7 +3981,7 @@ mod tests { seg0 } - fn reassigns_in(proposals: &[RaftCommand]) -> Vec<(SegmentKey, ReplicaSet)> { + fn reassigns_in(proposals: &[RaftCommand]) -> Vec<(SegmentKey, Replicas)> { proposals .iter() .filter_map(|c| match c { @@ -4095,7 +4100,7 @@ mod tests { MetadataCommand::RollSegment(RollSegment { segment_key: seg0, sealed_at: 2000, - new_replica_set: vec![node("node-1")], + new_replica_set: Replicas::new(vec![node("node-1")]), end_entry_id: Some(100.into()), }) .into(), @@ -4105,7 +4110,7 @@ mod tests { raft.propose( MetadataCommand::ReassignSegment(ReassignSegment { segment_key: seg0, - replica_set: new_set, + replica_set: Replicas::new(new_set), }) .into(), ) @@ -4115,8 +4120,8 @@ mod tests { } /// Run the catch-up re-drive sweep and collect (target, assignment) pairs. - fn drain_catch_up_redrives(raft: &mut Raft) -> Vec<(NodeId, CatchUpAssignment)> { - use crate::data_plane::messages::command::DataPlaneInterNodeCommand; + fn drain_catch_up_redrives(raft: &mut Raft) -> Vec<(NodeId, AssignSegmentCatchUp)> { + use crate::data_plane::messages::command::DataPlanePeerMessage; raft.maybe_redrive_catch_ups(); let mut out = Vec::new(); for event in raft.take_events() { @@ -4126,7 +4131,7 @@ mod tests { for cmd in cmds { if let DataTransportCommand::SendToTargets(s) = cmd { let target = s.targets[0].clone(); - if let DataPlaneInterNodeCommand::CatchUpAssignment(a) = s.message { + if let DataPlanePeerMessage::AssignSegmentCatchUp(a) = s.message { out.push((target, a)); } } @@ -4154,7 +4159,7 @@ mod tests { assert_eq!(a.shard_group_id, TEST_SHARD); assert_eq!(a.start_entry_id, 0.into()); assert_eq!(a.sealed_end_entry_id, 100.into()); - assert_eq!(a.replica_set, members); + assert_eq!(a.replica_set.0, members.clone().into_boxed_slice()); } } @@ -4207,7 +4212,7 @@ mod tests { MetadataCommand::RollSegment(RollSegment { segment_key: seg0, sealed_at: 2000, - new_replica_set: vec![node("y"), node("z"), node("node-1")], + new_replica_set: Replicas::new(vec![node("y"), node("z"), node("node-1")]), end_entry_id: None, }) .into(), @@ -4216,7 +4221,7 @@ mod tests { raft.simulate_flush_and_apply(); assert_eq!( raft.get_replica_set(&seg0.with_segment_id(SegmentId(1))), - Some(vec![node("y"), node("z"), node("node-1")]) + Some(Replicas::new(vec![node("y"), node("z"), node("node-1")])) ); let before = proposals_after_become_leader(&raft).len(); diff --git a/src/control_plane/consensus/seal_recovery.rs b/src/control_plane/consensus/seal_recovery.rs index d2e72834..d6009e33 100644 --- a/src/control_plane/consensus/seal_recovery.rs +++ b/src/control_plane/consensus/seal_recovery.rs @@ -17,7 +17,7 @@ use crate::control_plane::metadata::EntryId; use crate::data_plane::SegmentKey; #[derive(Debug, Clone, PartialEq, Eq)] -pub(crate) enum SealBoundaryReportError { +pub(crate) enum DurableSegmentEndReportedError { UnknownRecovery, UnexpectedReporter(NodeId), ConflictingReport { @@ -84,16 +84,16 @@ impl SealEndGather { segment_key: SegmentKey, from: NodeId, durable_end: Option, - ) -> Result, SealBoundaryReportError> { + ) -> Result, DurableSegmentEndReportedError> { if self.proposed { return Ok(None); } if !self.awaits(&from) { - return Err(SealBoundaryReportError::UnexpectedReporter(from)); + return Err(DurableSegmentEndReportedError::UnexpectedReporter(from)); } if let Some(previous) = self.reports.get(&from) { if *previous != durable_end { - return Err(SealBoundaryReportError::ConflictingReport { + return Err(DurableSegmentEndReportedError::ConflictingReport { reporter: from, previous: *previous, incoming: durable_end, @@ -247,10 +247,10 @@ impl SealEndRecovery { segment_key: SegmentKey, from: NodeId, durable_end: Option, - ) -> Result, SealBoundaryReportError> { + ) -> Result, DurableSegmentEndReportedError> { self.gathers .get_mut(&segment_key) - .ok_or(SealBoundaryReportError::UnknownRecovery)? + .ok_or(DurableSegmentEndReportedError::UnknownRecovery)? .record(segment_key, from, durable_end) } @@ -427,7 +427,7 @@ mod tests { assert!(matches!( recovery.record(seg, node("x"), Some(EntryId(10))), - Err(SealBoundaryReportError::UnexpectedReporter(reporter)) if reporter == node("x") + Err(DurableSegmentEndReportedError::UnexpectedReporter(reporter)) if reporter == node("x") )); assert!(matches!( recovery.record(seg, node("y"), Some(EntryId(50))), @@ -435,7 +435,7 @@ mod tests { )); assert!(matches!( recovery.record(seg, node("y"), Some(EntryId(49))), - Err(SealBoundaryReportError::ConflictingReport { + Err(DurableSegmentEndReportedError::ConflictingReport { reporter, previous: Some(EntryId(50)), incoming: Some(EntryId(49)), diff --git a/src/control_plane/membership/actor.rs b/src/control_plane/membership/actor.rs index 7822b8e2..0da40ee0 100644 --- a/src/control_plane/membership/actor.rs +++ b/src/control_plane/membership/actor.rs @@ -226,10 +226,10 @@ impl SwimSender { let Some(group) = self.resolve_shard_group(key).await? else { return Ok(ShardRouting::Redirect(None)); }; - if group.members.contains(node_id) { + if group.replicas.contains(node_id) { return Ok(ShardRouting::Local(group)); } - let member = self.resolve_any(&group.members).await?; + let member = self.resolve_any(&group.replicas).await?; Ok(ShardRouting::Redirect(member)) } diff --git a/src/control_plane/membership/topology.rs b/src/control_plane/membership/topology.rs index 7d918b9e..1ccba559 100644 --- a/src/control_plane/membership/topology.rs +++ b/src/control_plane/membership/topology.rs @@ -1,5 +1,5 @@ use crate::control_plane::membership::messages::dissemination_buffer::ShardLeaderInfo; -use crate::control_plane::{NodeAddressInfo, NodeId, SwimNodeState}; +use crate::control_plane::{NodeAddressInfo, NodeId, Replicas, SwimNodeState}; use crate::impl_new_struct_wrapper; #[cfg(any(test, debug_assertions))] use crate::test_traits::TAssertInvariant; @@ -31,7 +31,7 @@ impl_new_struct_wrapper!(ShardGroupId, u64); #[derive(Debug, Clone, PartialEq, Eq)] pub struct ShardGroup { pub id: ShardGroupId, - pub members: Vec, + pub replicas: Replicas, } /// Position on the consistent hash ring. Sorted by `(hash, pnode_id, replica_index)`. @@ -109,11 +109,8 @@ impl TopologyReader { /// Ring membership of `group_id` in the current snapshot, regardless of /// whether the local node is among the members. `None` when the group is /// not in the snapshot at all. - pub(crate) fn group_ring_members(&self, group_id: ShardGroupId) -> Option> { - self.0 - .load() - .group(group_id) - .map(|g| g.members.clone().into_boxed_slice()) + pub(crate) fn group_ring_members(&self, group_id: ShardGroupId) -> Option { + self.0.load().group(group_id).map(|g| g.replicas.clone()) } /// The cluster's configured replication factor — the target replica-set size a @@ -245,7 +242,7 @@ impl Topology { .chain(self.vnodes.range(..&start)) } - fn token_owners_at(&self, hash: u32, n: usize) -> Vec<&NodeId> { + fn token_owners_at(&self, hash: u32, n: usize) -> Vec { if self.vnodes.is_empty() || n == 0 { return Vec::new(); } @@ -257,16 +254,16 @@ impl Topology { hash: u32, n: usize, excluded: &HashSet, - ) -> Vec<&NodeId> { - let mut result: Vec<&NodeId> = Vec::with_capacity(n); + ) -> Vec { + let mut result: Vec = Vec::with_capacity(n); for token in self.walk_clockwise_from(hash) { if excluded.contains(&token.pnode_id) { continue; } - if result.iter().any(|o| **o == token.pnode_id) { + if result.contains(&token.pnode_id) { continue; } - result.push(&token.pnode_id); + result.push(token.pnode_id.clone()); if result.len() == n { break; } @@ -304,20 +301,21 @@ impl Topology { if !seen.insert(id) { continue; } - let owners = self.token_owners_at(token.hash, self.config.replication_factor); - let members: Vec = owners.into_iter().cloned().collect(); - for member in &members { + let replicas = + Replicas::new(self.token_owners_at(token.hash, self.config.replication_factor)); + + for member in &replicas.0 { self.node_group_ids .entry(member.clone()) .or_default() .push(id); } - self.groups.insert(id, ShardGroup { id, members }); + self.groups.insert(id, ShardGroup { id, replicas }); } } #[cfg(test)] - fn token_owners_for(&self, key: &[u8], n: usize) -> Vec<&NodeId> { + fn token_owners_for(&self, key: &[u8], n: usize) -> Vec { let hash = hash_stable(key); self.token_owners_at(hash, n) } @@ -376,7 +374,6 @@ impl Topology { let anchor = group_id.0 as u32; self.collect_distinct_owners_excluding(anchor, count, excluded) .into_iter() - .cloned() .collect() } @@ -451,7 +448,7 @@ pub mod props { .get(gid) .expect("node_group_ids references missing group"); assert!( - group.members.contains(node_id), + group.replicas.contains(node_id), "node {:?} in reverse index for group {:?} but not in group members", node_id, gid @@ -463,7 +460,7 @@ pub mod props { fn assert_groups_consistent(&self) { use std::collections::HashSet; for (gid, group) in &self.groups { - for member in &group.members { + for member in &group.replicas.0 { let ids = self .node_group_ids .get(member) @@ -475,10 +472,10 @@ pub mod props { member ); } - let unique: HashSet<&NodeId> = group.members.iter().collect(); + let unique: HashSet<&NodeId> = group.replicas.iter().collect(); assert_eq!( unique.len(), - group.members.len(), + group.replicas.len(), "group {:?} has duplicate members", gid ); @@ -634,7 +631,7 @@ mod tests { let multiple_owner = topology.token_owners_for(b"hello", 3); assert_eq!(multiple_owner.len(), 3); - let physical_node_ids: HashSet<&NodeId> = multiple_owner.into_iter().collect(); + let physical_node_ids: HashSet = multiple_owner.into_iter().collect(); assert_eq!(physical_node_ids.len(), 3); } @@ -672,7 +669,7 @@ mod tests { let new_owners = topology.token_owners_for(key, 1); assert_eq!(new_owners.len(), 1); - assert_eq!(*new_owners[0], successor_id); + assert_eq!(new_owners[0], successor_id); } // --- ShardGroup tests --- @@ -688,9 +685,9 @@ mod tests { ); let group = topology.shard_group_for(b"topic-blue").unwrap(); - assert_eq!(group.members.len(), 3); + assert_eq!(group.replicas.len(), 3); - let unique: HashSet<_> = group.members.iter().collect(); + let unique: HashSet<_> = group.replicas.iter().collect(); assert_eq!(unique.len(), 3); } @@ -737,7 +734,7 @@ mod tests { ); let group = topology.shard_group_for(b"any-key").unwrap(); - assert_eq!(group.members.len(), 2); + assert_eq!(group.replicas.len(), 2); } #[test] @@ -752,14 +749,14 @@ mod tests { let key = b"test-key"; let before = topology.shard_group_for(key).unwrap().clone(); - assert_eq!(before.members.len(), 3); + assert_eq!(before.replicas.len(), 3); - let removed = before.members[0].clone(); + let removed = before.replicas[0].clone(); topology.remove_node(&removed); let after = topology.shard_group_for(key).unwrap(); - assert_eq!(after.members.len(), 2); - assert!(!after.members.contains(&removed)); + assert_eq!(after.replicas.len(), 2); + assert!(!after.replicas.contains(&removed)); } #[test] @@ -790,8 +787,8 @@ mod tests { let groups = topology.shard_groups_for_node(&NodeId::new("node-0")); assert!(!groups.is_empty()); for group in &groups { - assert!(group.members.contains(&NodeId::new("node-0"))); - assert_eq!(group.members.len(), 3); + assert!(group.replicas.contains(&NodeId::new("node-0"))); + assert_eq!(group.replicas.len(), 3); } } @@ -913,7 +910,7 @@ mod tests { fn ring_replacements_returns_node_not_in_excluded() { let topology = five_node_topology(); let group = topology.shard_group_for(b"topic-blue").unwrap().clone(); - let current: HashSet = group.members.iter().cloned().collect(); + let current: HashSet = group.replicas.iter().cloned().collect(); let picks = topology.ring_replacements_for(group.id, ¤t, 1); @@ -931,7 +928,7 @@ mod tests { let topology = five_node_topology(); let group = topology.shard_group_for(b"topic-blue").unwrap().clone(); // Exclude four of five nodes: every member plus one extra outsider. - let mut excluded: HashSet = group.members.iter().cloned().collect(); + let mut excluded: HashSet = group.replicas.iter().cloned().collect(); let outsider = topology .live_nodes() .into_iter() @@ -959,7 +956,7 @@ mod tests { fn ring_replacements_returns_fewer_than_count_when_pool_depleted() { let topology = five_node_topology(); let group = topology.shard_group_for(b"topic-blue").unwrap().clone(); - let current: HashSet = group.members.iter().cloned().collect(); + let current: HashSet = group.replicas.iter().cloned().collect(); // 5 total, 3 excluded → 2 eligible. Ask for 5; expect 2. let picks = topology.ring_replacements_for(group.id, ¤t, 5); @@ -1002,7 +999,7 @@ mod tests { fn ring_replacements_deterministic_across_calls() { let topology = five_node_topology(); let group = topology.shard_group_for(b"topic-blue").unwrap().clone(); - let current: HashSet = group.members.iter().cloned().collect(); + let current: HashSet = group.replicas.iter().cloned().collect(); let a = topology.ring_replacements_for(group.id, ¤t, 2); let b = topology.ring_replacements_for(group.id, ¤t, 2); @@ -1031,7 +1028,7 @@ mod tests { let g2 = t2.shard_group_for(b"topic-blue").unwrap(); assert_eq!(g1.id, g2.id, "group id must be deterministic"); - let excluded: HashSet = g1.members.iter().cloned().collect(); + let excluded: HashSet = g1.replicas.iter().cloned().collect(); let p1 = t1.ring_replacements_for(g1.id, &excluded, 2); let p2 = t2.ring_replacements_for(g2.id, &excluded, 2); assert_eq!(p1, p2); @@ -1044,19 +1041,16 @@ mod tests { // Verifies we're picking the ring's natural successors, not arbitrary nodes. let topology = five_node_topology(); let group = topology.shard_group_for(b"topic-blue").unwrap().clone(); - assert_eq!(group.members.len(), 3); + assert_eq!(group.replicas.len(), 3); // Owners at RF=5 are: the 3 current members, then the 2 ring successors. - let owners_5: Vec = topology - .token_owners_at(group.id.0 as u32, 5) - .into_iter() - .cloned() - .collect(); + let owners_5 = topology.token_owners_at(group.id.0 as u32, 5); + assert_eq!(owners_5.len(), 5); - assert_eq!(&owners_5[..3], group.members.as_slice()); + assert_eq!(&owners_5[..3], &group.replicas[..3]); let expected_successors = &owners_5[3..]; - let current: HashSet = group.members.iter().cloned().collect(); + let current: HashSet = group.replicas.iter().cloned().collect(); let picks = topology.ring_replacements_for(group.id, ¤t, 2); assert_eq!(&*picks, expected_successors); } @@ -1291,6 +1285,6 @@ mod tests { let g1 = t1.shard_group_for(b"topic-blue").unwrap(); let g2 = t2.shard_group_for(b"topic-blue").unwrap(); assert_eq!(g1.id, g2.id); - assert_eq!(g1.members, g2.members); + assert_eq!(g1.replicas, g2.replicas); } } diff --git a/src/control_plane/metadata/command.rs b/src/control_plane/metadata/command.rs index 18d8cb62..83ca9e03 100644 --- a/src/control_plane/metadata/command.rs +++ b/src/control_plane/metadata/command.rs @@ -6,7 +6,7 @@ use uuid::Uuid; use crate::{ connections::protocol::ConsumerGroupSyncAction, control_plane::{ - NodeId, + Replicas, metadata::{EntryId, RangeId, SegmentId, TopicId, strategy::StoragePolicy}, }, data_plane::SegmentKey, @@ -17,7 +17,7 @@ use crate::{ pub struct CreateTopic { pub name: String, pub storage_policy: StoragePolicy, - pub replica_set: Vec, + pub replica_set: Replicas, pub created_at: u64, } @@ -25,7 +25,7 @@ pub struct CreateTopic { pub struct RollSegment { pub segment_key: SegmentKey, pub sealed_at: u64, - pub new_replica_set: Vec, + pub new_replica_set: Replicas, /// None for SWIM-death-triggered seals — the coordinator doesn't know /// the actual committed offset. Corrected later via `correct_end_offset` /// or D5 sealed segment repair. @@ -38,8 +38,8 @@ pub struct SplitRange { pub range_id: RangeId, pub split_point: Vec, pub created_at: u64, - pub left_replica_set: Vec, - pub right_replica_set: Vec, + pub left_replica_set: Replicas, + pub right_replica_set: Replicas, } #[derive(Debug, Clone, PartialEq, Eq, BorshSerialize, BorshDeserialize)] @@ -48,7 +48,7 @@ pub struct MergeRange { pub range_id_1: RangeId, pub range_id_2: RangeId, pub created_at: u64, - pub merged_replica_set: Vec, + pub merged_replica_set: Replicas, } #[derive(Debug, Clone, PartialEq, Eq, BorshSerialize, BorshDeserialize)] @@ -59,7 +59,7 @@ pub struct DeleteTopic { #[derive(Debug, Clone, PartialEq, Eq, BorshSerialize, BorshDeserialize)] pub struct ReassignSegment { pub segment_key: SegmentKey, - pub replica_set: Vec, + pub replica_set: Replicas, } /// Retention: mark an oldest-first **prefix** of one range's sealed segments diff --git a/src/control_plane/metadata/event.rs b/src/control_plane/metadata/event.rs index b3dac419..6bc7e472 100644 --- a/src/control_plane/metadata/event.rs +++ b/src/control_plane/metadata/event.rs @@ -1,4 +1,4 @@ -use crate::control_plane::NodeId; +use crate::control_plane::Replicas; use crate::control_plane::consensus::multi_raft::SealContext; use crate::control_plane::membership::ShardGroupId; use crate::control_plane::metadata::consumer_group::GenerationId; @@ -6,7 +6,7 @@ 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::messages::command::{ - CatchUpAssignment, DataPlaneInterNodeCommand, DeleteSegments, SealResponse, SegmentAssignment, + AssignSegmentCatchUp, DataPlanePeerMessage, DeleteSegments, PlaceSegment, SegmentRollCommitted, SegmentSealed, }; use crate::data_plane::transport::command::DataTransportCommand; @@ -15,13 +15,13 @@ use crate::impl_from_variant; #[derive(Debug, Clone, PartialEq, Eq)] pub struct TopicCreated { pub segment_key: SegmentKey, - pub replica_set: Vec, + pub replica_set: Replicas, } impl TopicCreated { pub fn into_command(self, shard_group_id: ShardGroupId) -> DataTransportCommand { DataTransportCommand::send_to_targets( vec![self.replica_set[0].clone()], - SegmentAssignment { + PlaceSegment { segment_key: self.segment_key, shard_group_id, replica_set: self.replica_set, @@ -34,7 +34,7 @@ impl TopicCreated { #[derive(Debug, Clone, PartialEq, Eq)] pub struct SegmentRolled { pub new_segment_key: SegmentKey, - pub new_replica_set: Vec, + pub new_replica_set: Replicas, pub end_entry_id: Option, pub consumer_group_epochs: Box<[ConsumerGroupEpochSnapshot]>, } @@ -48,7 +48,7 @@ impl SegmentRolled { let start = self.end_entry_id.map_or(EntryId::MIN, |id| id + 1); let mut v = vec![DataTransportCommand::send_to_targets( vec![self.new_replica_set[0].clone()], - SegmentAssignment { + PlaceSegment { segment_key: self.new_segment_key, shard_group_id, replica_set: self.new_replica_set.clone(), @@ -59,7 +59,7 @@ impl SegmentRolled { if let Some(ctx) = ctx { v.push(DataTransportCommand::send_to_targets( vec![ctx.requester], - SealResponse { + SegmentRollCommitted { old_segment_key: ctx.segment_key, new_segment_id: self.new_segment_key.segment_id, new_replica_set: self.new_replica_set, @@ -81,7 +81,7 @@ pub struct SegmentReassigned { /// was never established; those aren't selected for reassignment, so the /// dispatch just emits nothing. pub sealed_end: Option, - pub new_replica_set: Vec, + pub new_replica_set: Replicas, } impl SegmentReassigned { @@ -94,7 +94,7 @@ impl SegmentReassigned { .map(|member| { DataTransportCommand::send_to_targets( vec![member.clone()], - CatchUpAssignment { + AssignSegmentCatchUp { segment_key: self.segment_key, shard_group_id, start_entry_id: self.start_entry_id, @@ -110,8 +110,8 @@ impl SegmentReassigned { #[derive(Debug, Clone, PartialEq, Eq)] pub struct RangeSplit { pub topic_id: TopicId, - pub children: [(RangeId, SegmentId, Vec); 2], - pub parent_active_segment: Option<(SegmentKey, Vec)>, + pub children: [(RangeId, SegmentId, Replicas); 2], + pub parent_active_segment: Option<(SegmentKey, Replicas)>, pub consumer_group_epochs: Box<[ConsumerGroupEpochSnapshot]>, } @@ -124,7 +124,7 @@ impl RangeSplit { let target = replica_set[0].clone(); DataTransportCommand::send_to_targets( vec![target], - SegmentAssignment { + PlaceSegment { segment_key: SegmentKey::new(self.topic_id, range_id, segment_id), shard_group_id, replica_set, @@ -139,7 +139,7 @@ impl RangeSplit { && !parent_replica_set.is_empty() { cmds.push(DataTransportCommand::send_to_targets( - parent_replica_set, + parent_replica_set.0, SegmentSealed { segment_key: parent_key, committed_entry_id: None, @@ -158,7 +158,7 @@ impl RangeSplit { #[derive(Debug, Clone, PartialEq, Eq)] pub struct RangeMerged { pub segment_key: SegmentKey, - pub replica_set: Vec, + pub replica_set: Replicas, pub consumer_group_epochs: Box<[ConsumerGroupEpochSnapshot]>, } impl RangeMerged { @@ -166,7 +166,7 @@ impl RangeMerged { let target = self.replica_set[0].clone(); let mut commands = vec![DataTransportCommand::send_to_targets( vec![target], - SegmentAssignment { + PlaceSegment { segment_key: self.segment_key, shard_group_id, replica_set: self.replica_set, @@ -185,7 +185,7 @@ pub struct ConsumerGroupEpochSnapshot { pub topic_id: TopicId, pub group_id: String, pub generation: GenerationId, - pub ranges: Box<[(RangeId, Vec)]>, + pub ranges: Box<[(RangeId, Replicas)]>, } #[derive(Debug, Clone, PartialEq, Eq)] @@ -210,8 +210,8 @@ impl ConsumerGroupEpochSnapshot { .filter(|(_, replicas)| !replicas.is_empty()) .map(|(range_id, replicas)| { DataTransportCommand::send_to_targets( - replicas, - DataPlaneInterNodeCommand::ConsumerGroupEpochSeal(EpochSeal { + replicas.0, + DataPlanePeerMessage::ConsumerGroupEpochSealed(EpochSeal { key: ConsumerOffsetKey { topic_id: self.topic_id, range_id, @@ -231,7 +231,7 @@ impl ConsumerGroupEpochSnapshot { /// different replica sets, so there may be several groups. #[derive(Debug, Clone, PartialEq, Eq)] pub struct SegmentsDeleted { - pub groups: Vec<(Vec, Vec)>, + pub groups: Vec<(Replicas, Vec)>, } impl SegmentsDeleted { @@ -240,7 +240,7 @@ impl SegmentsDeleted { .into_iter() .map(|(replica_set, keys)| { DataTransportCommand::send_to_targets( - replica_set, + replica_set.0, DeleteSegments { segment_keys: keys.into_boxed_slice(), }, @@ -253,7 +253,7 @@ impl SegmentsDeleted { #[derive(Debug, Clone, PartialEq, Eq)] pub struct SegmentSealCorrected { pub segment_key: SegmentKey, - pub replica_set: Vec, + pub replica_set: Replicas, pub committed_entry_id: Option, } @@ -263,7 +263,7 @@ impl SegmentSealCorrected { vec![] } else { vec![DataTransportCommand::send_to_targets( - self.replica_set, + self.replica_set.0, SegmentSealed { segment_key: self.segment_key, committed_entry_id: self.committed_entry_id, diff --git a/src/control_plane/metadata/range.rs b/src/control_plane/metadata/range.rs index d4d55792..1b11bfa2 100644 --- a/src/control_plane/metadata/range.rs +++ b/src/control_plane/metadata/range.rs @@ -2,9 +2,12 @@ use super::constants::*; use borsh::{BorshDeserialize, BorshSerialize}; use std::collections::{HashMap, VecDeque}; -use crate::control_plane::metadata::{ - EntryId, RangeId, ReplicaSet, SegmentId, SegmentMeta, SegmentMetaState, SplitRange, - command::RollSegment, error::MetadataError, +use crate::control_plane::{ + Replicas, + metadata::{ + EntryId, RangeId, SegmentId, SegmentMeta, SegmentMetaState, SplitRange, + command::RollSegment, error::MetadataError, + }, }; #[derive(Debug, Clone, Copy, PartialEq, Eq, BorshSerialize, BorshDeserialize)] @@ -35,7 +38,7 @@ impl RangeMeta { range_id: RangeId, keyspace_start: Vec, keyspace_end: Vec, - replica_set: ReplicaSet, + replica_set: Replicas, created_at: u64, ) -> Self { let segment_id = SegmentId(0); @@ -113,7 +116,7 @@ impl RangeMeta { &mut self, segment_id: SegmentId, end_entry_id: EntryId, - ) -> Option { + ) -> Option { let seg = self.segments.get_mut(&segment_id)?; if seg.state != SegmentMetaState::Sealed || seg.end_entry_id.is_some() { @@ -186,7 +189,7 @@ impl RangeMeta { &mut self, other: &mut RangeMeta, merged_id: RangeId, - replica_set: ReplicaSet, + replica_set: Replicas, requested_at: u64, ) -> Result { self.validate_mergeable(other)?; @@ -220,7 +223,7 @@ impl RangeMeta { &self, other: &RangeMeta, merged_id: RangeId, - replica_set: ReplicaSet, + replica_set: Replicas, requested_at: u64, ) -> RangeMeta { let (merged_start, merged_end) = self.merged_keyspace(other); diff --git a/src/control_plane/metadata/segment.rs b/src/control_plane/metadata/segment.rs index c99ae0d4..200849c5 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::{ - NodeId, + Replicas, metadata::{EntryId, SegmentId, error::MetadataError}, }; @@ -12,12 +12,11 @@ pub enum SegmentMetaState { Deleting, } -pub(crate) type ReplicaSet = Vec; #[derive(Debug, Clone, PartialEq, Eq, BorshSerialize, BorshDeserialize)] pub struct SegmentMeta { pub segment_id: SegmentId, pub state: SegmentMetaState, - pub replica_set: ReplicaSet, + pub replica_set: Replicas, pub size_bytes: u64, pub start_entry_id: EntryId, pub end_entry_id: Option, @@ -28,7 +27,7 @@ pub struct SegmentMeta { impl SegmentMeta { pub(crate) fn new( segment_id: SegmentId, - replica_set: ReplicaSet, + replica_set: Replicas, start_entry_id: EntryId, created_at: u64, ) -> Self { @@ -64,7 +63,7 @@ impl SegmentMeta { /// the state stays `Sealed`. Returns whether the set actually changed: /// an identical set is an idempotent no-op, tolerating duplicate death detection /// and no-leader re-proposals (cf. `RollSegment` idempotency). - pub(crate) fn reassign(&mut self, new_replica_set: ReplicaSet) -> Result { + pub(crate) fn reassign(&mut self, new_replica_set: Replicas) -> Result { if self.state != SegmentMetaState::Sealed { return Err(MetadataError::SegmentNotSealed); } diff --git a/src/control_plane/metadata/state_machine.rs b/src/control_plane/metadata/state_machine.rs index a4f706bd..b261f345 100644 --- a/src/control_plane/metadata/state_machine.rs +++ b/src/control_plane/metadata/state_machine.rs @@ -1,8 +1,9 @@ use super::command::*; use super::event::*; -use super::segment::*; + use super::topic::{TopicMeta, TopicState, TopicStats}; use crate::control_plane::NodeId; +use crate::control_plane::Replicas; use crate::control_plane::membership::ShardGroupId; use crate::control_plane::metadata::ConsumerGroupAssignment; use crate::control_plane::metadata::{EntryId, RangeId, SegmentId, TopicId, error::MetadataError}; @@ -74,7 +75,7 @@ impl MetadataStateMachine { pub(crate) fn active_segments_for_node( &self, node_id: &NodeId, - ) -> Box<[(SegmentKey, ReplicaSet)]> { + ) -> Box<[(SegmentKey, Replicas)]> { self.topics .values() .flat_map(|t| t.active_segments_for_node(node_id)) @@ -83,7 +84,7 @@ impl MetadataStateMachine { /// Every active segment across all topics with its replica set and start /// offset, for the leader's periodic assignment re-drive. - pub(crate) fn active_segment_assignments(&self) -> Box<[(SegmentKey, ReplicaSet, EntryId)]> { + pub(crate) fn active_segment_assignments(&self) -> Box<[(SegmentKey, Replicas, EntryId)]> { self.topics .values() .flat_map(|t| t.active_segment_assignments()) @@ -235,7 +236,7 @@ impl MetadataStateMachine { return Ok(ApplyResult::Noop); } // Group the deleted segments by replica_set here. - let mut groups: Vec<(Vec, Vec)> = Vec::new(); + let mut groups: Vec<(Replicas, Vec)> = Vec::new(); for sid in &deleted_ids { let Some(seg) = range.segments.get(sid) else { continue; @@ -400,6 +401,7 @@ impl crate::test_traits::TAssertInvariant for MetadataStateMachine { mod tests { use super::super::constants::*; use super::super::range::*; + use super::super::segment::*; use super::*; use crate::connections::protocol::ConsumerGroupSyncAction; use crate::control_plane::membership::ShardGroupId; @@ -428,12 +430,12 @@ mod tests { } } - fn replica_set() -> Vec { - vec![ + fn replica_set() -> Replicas { + Replicas::new(vec![ NodeId::new("node-1"), NodeId::new("node-2"), NodeId::new("node-3"), - ] + ]) } fn create_topic(sm: &mut MetadataStateMachine, name: &str) -> TopicId { @@ -705,12 +707,12 @@ mod tests { /// A surviving subset plus a fresh replacement — what the coordinator picks /// when a replica of a sealed segment dies (node-3 → node-4 here). - fn replacement_set() -> Vec { - vec![ + fn replacement_set() -> Replicas { + Replicas::new(vec![ NodeId::new("node-1"), NodeId::new("node-2"), NodeId::new("node-4"), - ] + ]) } // --- ReassignSegment --- diff --git a/src/control_plane/metadata/topic.rs b/src/control_plane/metadata/topic.rs index f28c00f5..83e36403 100644 --- a/src/control_plane/metadata/topic.rs +++ b/src/control_plane/metadata/topic.rs @@ -4,7 +4,7 @@ use super::constants::*; use super::*; use crate::{ control_plane::{ - NodeId, + NodeId, Replicas, metadata::{ error::MetadataError, event::ConsumerGroupEpochSnapshot, @@ -37,7 +37,7 @@ impl TopicMeta { pub(crate) fn new( name: String, id: TopicId, - replica_set: ReplicaSet, + replica_set: Replicas, created_at: u64, storage_policy: StoragePolicy, ) -> Self { @@ -103,7 +103,7 @@ impl TopicMeta { pub(crate) fn active_segments_for_node( &self, node_id: &NodeId, - ) -> Box<[(SegmentKey, ReplicaSet)]> { + ) -> Box<[(SegmentKey, Replicas)]> { if self.state != TopicState::Active { return Box::new([]); } @@ -123,9 +123,9 @@ impl TopicMeta { /// Every active segment with its replica set and current start offset, for /// the leader's periodic assignment re-drive. Returns `(key, replica_set, - /// start_offset)` so a re-driven `SegmentAssignment` can recreate a segment + /// start_offset)` so a re-driven `PlaceSegment` can recreate a segment /// that missed its one-shot delivery with the correct starting offset. - pub(crate) fn active_segment_assignments(&self) -> Box<[(SegmentKey, ReplicaSet, EntryId)]> { + pub(crate) fn active_segment_assignments(&self) -> Box<[(SegmentKey, Replicas, EntryId)]> { if self.state != TopicState::Active { return Box::new([]); } @@ -147,7 +147,7 @@ impl TopicMeta { pub(crate) fn active_segments_with_dead_members( &self, live: &std::collections::HashSet, - ) -> Box<[(SegmentKey, ReplicaSet)]> { + ) -> Box<[(SegmentKey, Replicas)]> { if self.state != TopicState::Active { return Box::new([]); } @@ -176,7 +176,7 @@ impl TopicMeta { pub(crate) fn sealed_segments_with_dead_members( &self, live: &HashSet, - ) -> Box<[(SegmentKey, ReplicaSet)]> { + ) -> Box<[(SegmentKey, Replicas)]> { let mut out = Vec::new(); for range in self.ranges.values() { for seg in range.segments.values() { @@ -197,7 +197,7 @@ impl TopicMeta { /// Boundary-unknown sealed segments. These remain seal-end recovery /// candidates until a recovered boundary is committed to metadata, even if /// the crashed replica has rejoined by the next reconciliation pass. - pub(crate) fn boundary_unknown_segments(&self) -> Box<[(SegmentKey, ReplicaSet)]> { + pub(crate) fn boundary_unknown_segments(&self) -> Box<[(SegmentKey, Vec)]> { let mut out = Vec::new(); for range in self.ranges.values() { let active_replicas = range @@ -230,7 +230,7 @@ impl TopicMeta { pub(crate) fn under_replicated_sealed_segments( &self, replication_factor: usize, - ) -> Box<[(SegmentKey, ReplicaSet)]> { + ) -> Box<[(SegmentKey, Replicas)]> { self.ranges .values() .flat_map(|range| { @@ -248,12 +248,12 @@ impl TopicMeta { .collect() } - /// Every known-end sealed segment, with the data a `CatchUpAssignment` needs. + /// Every known-end sealed segment, with the data a `AssignSegmentCatchUp` needs. /// The takeover reseed (raft-actor.md #9) tracks these so a repair in flight /// when leadership changed isn't stranded. pub(crate) fn known_end_sealed_segments( &self, - ) -> Vec<(SegmentKey, EntryId, EntryId, ReplicaSet)> { + ) -> Vec<(SegmentKey, EntryId, EntryId, Replicas)> { let mut out = Vec::new(); for range in self.ranges.values() { for seg in range.segments.values() { @@ -356,7 +356,7 @@ impl TopicMeta { .active_segment .and_then(|sid| r1.segments.get(&sid)) .map(|s| s.replica_set.clone()) - .unwrap_or_default(); + .unwrap(); Some(MetadataCommand::MergeRange(MergeRange { topic_id: self.id, range_id_1: r1.range_id, @@ -635,7 +635,13 @@ mod routing_tests { } fn topic() -> TopicMeta { - TopicMeta::new("t".into(), TopicId(1), vec![NodeId::new("n")], 0, policy()) + TopicMeta::new( + "t".into(), + TopicId(1), + Replicas::new(vec![NodeId::new("n")]), + 0, + policy(), + ) } #[test] @@ -658,8 +664,8 @@ mod routing_tests { range_id: RangeId(0), split_point: vec![0x80], created_at: 1, - left_replica_set: vec![NodeId::new("n")], - right_replica_set: vec![NodeId::new("n")], + left_replica_set: Replicas::new(vec![NodeId::new("n")]), + right_replica_set: Replicas::new(vec![NodeId::new("n")]), }) .unwrap(); // left = [MIN, 0x80), right = [0x80, MAX]. diff --git a/src/control_plane/types/node.rs b/src/control_plane/types/node.rs index c7d8b2e7..d010172f 100644 --- a/src/control_plane/types/node.rs +++ b/src/control_plane/types/node.rs @@ -139,13 +139,13 @@ impl std::borrow::Borrow for NodeId { } #[derive(Clone, Debug, PartialEq, Eq, Hash, BorshSerialize, BorshDeserialize)] -pub struct Replicas(Vec); -crate::smart_pointer!(Replicas, Vec); +pub struct Replicas(pub Box<[NodeId]>); +crate::impl_new_struct_wrapper!(Replicas, Box<[NodeId]>); impl Replicas { pub fn new(nodes: Vec) -> Self { assert!(!nodes.is_empty(), "Replica set cannot be empty"); - Self(nodes) + Self(nodes.into_boxed_slice()) } pub fn leader(&self) -> Option<&NodeId> { @@ -155,4 +155,19 @@ impl Replicas { pub fn followers(&self) -> impl Iterator { self.iter().skip(1) } + + pub fn extend(&mut self, nodes: &[NodeId]) { + let new_nodes = [&*self.0, nodes].concat(); + self.0 = new_nodes.into_boxed_slice(); + } + + pub fn change_leader(&mut self, pos: usize) { + // Take until pos and rotate to right so + // [A,B,C] -> [C,A,B] + self.0[0..=pos].rotate_right(1); + } + + pub fn except_for(&self, node: &NodeId) -> Vec { + self.0.iter().filter(|n| *n != node).cloned().collect() + } } diff --git a/src/data_plane/consumer_offset_management/ledger.rs b/src/data_plane/consumer_offset_management/ledger.rs index dc4e8b7a..eb4bebd8 100644 --- a/src/data_plane/consumer_offset_management/ledger.rs +++ b/src/data_plane/consumer_offset_management/ledger.rs @@ -9,6 +9,7 @@ use borsh::{BorshDeserialize, BorshSerialize}; use crate::client::RangeId; use crate::control_plane::metadata::consumer_group::GenerationId; use crate::control_plane::metadata::{EntryId, TopicId}; +use crate::data_plane::SegmentKey; const SNAPSHOT_FILE: &str = "consumer-offsets.snapshot"; const SNAPSHOT_TEMP_FILE: &str = "consumer-offsets.snapshot.tmp"; @@ -20,6 +21,12 @@ pub(crate) struct ConsumerOffsetKey { pub(crate) group_id: String, } +impl ConsumerOffsetKey { + pub(crate) fn placement_key(&self) -> (TopicId, RangeId) { + (self.topic_id, self.range_id) + } +} + #[derive(Default, Debug, Clone, Copy, PartialEq, Eq, BorshSerialize, BorshDeserialize)] pub struct ConsumerOffsetPosition { pub entry_id: EntryId, @@ -54,19 +61,30 @@ pub struct ConsumerOffsetUpdate { pub generation: GenerationId, pub position: ConsumerOffsetPosition, } + +#[derive(Debug, Clone, PartialEq, Eq, BorshSerialize, BorshDeserialize)] +pub struct ConsumerOffsetEntry { + pub key: ConsumerOffsetKey, + pub generation: GenerationId, + pub position: Option, +} + #[derive(Debug, Clone, BorshSerialize, BorshDeserialize)] pub(crate) enum OffsetRecord { EpochSeal(EpochSeal), OffsetCommit(ConsumerOffsetUpdate), + BootstrapEntry(ConsumerOffsetEntry), + PlacementInstalled(SegmentKey), } -/// Durable consumer-group state. Live mutations are persisted by the shared +/// "Durable" consumer-group state. Live mutations are persisted by the shared /// 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 { epochs: HashMap, offsets: HashMap, + installed_placements: HashMap<(TopicId, RangeId), SegmentKey>, } impl OffsetLedger { @@ -123,18 +141,75 @@ impl OffsetLedger { }) .or_insert(position); } + OffsetRecord::BootstrapEntry(snapshot) => { + self.epochs + .entry(snapshot.key.clone()) + .and_modify(|generation| *generation = (*generation).max(snapshot.generation)) + .or_insert(snapshot.generation); + if let Some(position) = snapshot.position { + self.offsets + .entry(snapshot.key) + .and_modify(|current| *current = (*current).max(position)) + .or_insert(position); + } + } + OffsetRecord::PlacementInstalled(segment_key) => { + self.installed_placements + .entry((segment_key.topic_id, segment_key.range_id)) + .and_modify(|current| { + if segment_key.segment_id > current.segment_id { + *current = segment_key; + } + }) + .or_insert(segment_key); + } } } pub(crate) fn offset(&self, key: &ConsumerOffsetKey) -> Option { self.offsets.get(key).copied() } + + pub(crate) fn has_installed_placement(&self, segment_key: &SegmentKey) -> bool { + self.installed_placements.get(&segment_key.placement_key()) == Some(segment_key) + } + + pub(crate) fn can_graduate_to(&self, segment_key: &SegmentKey) -> bool { + let Some(installed) = self.installed_placements.get(&segment_key.placement_key()) else { + return false; + }; + + installed.segment_id.checked_add(1) == Some(segment_key.segment_id.0) + } + + pub(crate) fn can_source_snapshot_for(&self, segment_key: &SegmentKey) -> bool { + self.has_installed_placement(segment_key) || self.can_graduate_to(segment_key) + } + + pub(crate) fn snapshot_range( + &self, + topic_id: TopicId, + range_id: RangeId, + ) -> Box<[ConsumerOffsetEntry]> { + self.epochs + .iter() + .filter(|(key, _)| key.topic_id == topic_id && key.range_id == range_id) + .map(|(key, generation)| ConsumerOffsetEntry { + key: key.clone(), + generation: *generation, + position: self.offsets.get(key).copied(), + }) + .collect() + } } #[cfg(test)] mod tests { use super::*; - use crate::{client::RangeId, control_plane::metadata::TopicId}; + use crate::{ + client::RangeId, + control_plane::metadata::{SegmentId, TopicId}, + }; fn key() -> ConsumerOffsetKey { ConsumerOffsetKey { @@ -144,6 +219,18 @@ 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(); + ledger.apply(OffsetRecord::PlacementInstalled(current)); + + assert!(ledger.has_installed_placement(¤t)); + assert!(!ledger.can_graduate_to(¤t)); + assert!(ledger.can_graduate_to(¤t.with_segment_id(SegmentId(5)))); + assert!(!ledger.can_graduate_to(¤t.with_segment_id(SegmentId(6)))); + } + #[test] fn snapshot_survives_restart() { let dir = tempfile::tempdir().unwrap(); diff --git a/src/data_plane/consumer_offset_management/manager.rs b/src/data_plane/consumer_offset_management/manager.rs index 320614b2..194b7ce0 100644 --- a/src/data_plane/consumer_offset_management/manager.rs +++ b/src/data_plane/consumer_offset_management/manager.rs @@ -1,5 +1,6 @@ +use crate::control_plane::membership::ShardGroupId; use crate::control_plane::metadata::consumer_group::GenerationId; -use crate::control_plane::metadata::{RangeId, TopicId}; +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; @@ -9,27 +10,109 @@ use crate::data_plane::consumer_offset_management::ledger::{ use crate::data_plane::consumer_offset_management::replication::OffsetReplicationState; use crate::data_plane::messages::command::{ - CommitConsumerOffset, ConsumerOffsetCommitAck, ReplicaOffsetAck, ReplicaOffsetAckResult, + CommitConsumerOffset, ConsumerOffsetCommitAck, ConsumerOffsetReplicated, + ConsumerOffsetReplicationResult, ConsumerOffsetSnapshot, ConsumerOffsetSnapshotInstalled, }; +use crate::data_plane::messages::{RequestConsumerOffsetSnapshot, SegmentPlaced}; use crate::data_plane::transport::command::DataTransportCommand; -use std::collections::{HashMap, VecDeque}; +use std::collections::{HashMap, HashSet, VecDeque}; use std::path::PathBuf; use super::types::*; #[derive(Default)] pub(crate) struct ConsumerOffsetManager { - offset_ledger: OffsetLedger, - offset_placements: HashMap<(TopicId, RangeId), OffsetPlacement>, + 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(Default)] +struct ConsumerOffsetCoordination { + placements: HashMap<(TopicId, RangeId), OffsetPlacement>, pending_offset_mutations: Vec, future_offset_commits: HashMap>, - pub(crate) uncheckpointed_offset_lsns: VecDeque, - pub(crate) offset_checkpoint_lsn: u64, - pub(crate) offset_checkpoint_in_flight: bool, offset_replication: OffsetReplicationState, } +struct OffsetPlacement { + segment_key: SegmentKey, + shard_group_id: ShardGroupId, + // Desired data replicas for this placement. This also preserves their + // ordering, including which replica is the leader. + replicas: Replicas, + // Consumer-offset readiness for every member of `replicas`. + replica_states: HashMap, + placement_ack_sent: bool, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum OffsetReplicaState { + Joining, + SnapshotInstalled, + Ready, +} + +impl OffsetPlacement { + fn leader(&self) -> &NodeId { + self.replicas + .leader() + .expect("an offset placement must have a leader") + } + + fn can_bootstrap_replicas(&self, segment_key: &SegmentKey) -> bool { + self.segment_key == *segment_key && self.is_replica_ready(self.leader()) + } + + fn is_replica_ready(&self, node: &NodeId) -> bool { + self.replica_states.get(node) == Some(&OffsetReplicaState::Ready) + } + + fn all_replicas_ready(&self) -> bool { + self.replica_states + .values() + .all(|state| *state == OffsetReplicaState::Ready) + } + + fn compare(&self, other_key: &SegmentKey, other: &Replicas) -> PlacementObservation { + if self.segment_key == *other_key { + if &self.replicas != other { + tracing::warn!( + ?other_key, + current_replicas = ?self.replicas, + observed_replicas = ?other, + "ignored conflicting replica set for an existing segment" + ); + } + return if &self.replicas == other { + PlacementObservation::Unchanged + } else { + PlacementObservation::Conflict + }; + } + if self.segment_key.segment_id >= other_key.segment_id { + return PlacementObservation::Stale; + } + PlacementObservation::Accepted + } + + fn unready_followers(&self) -> Box<[NodeId]> { + self.replicas + .followers() + .filter(|node| !self.is_replica_ready(node)) + .cloned() + .collect() + } +} + +#[derive(Default)] +struct OffsetCheckpointState { + uncheckpointed_lsns: VecDeque, + checkpoint_lsn: u64, + in_flight: bool, +} + impl ConsumerOffsetManager { pub(crate) fn new(offset_ledger: OffsetLedger) -> Self { Self { @@ -38,15 +121,18 @@ impl ConsumerOffsetManager { } } - fn future_entry(&mut self, key: ConsumerOffsetKey) -> &mut Vec { - self.future_offset_commits.entry(key).or_default() + pub(crate) fn future_entry(&mut self, key: ConsumerOffsetKey) -> &mut Vec { + self.coordination + .future_offset_commits + .entry(key) + .or_default() } pub(crate) fn offset(&self, key: &ConsumerOffsetKey) -> Option { self.offset_ledger.offset(key) } - pub(crate) fn handle_consumer_offset_commit( + pub(crate) fn commit_consumer_offset( &mut self, cmd: CommitConsumerOffset, node_id: &NodeId, @@ -60,6 +146,7 @@ impl ConsumerOffsetManager { return false; } }; + let required_followers = self.get_placement_followers(&cmd); let sealed_generation = self.latest_generation(&cmd.update.key); if cmd.update.generation < sealed_generation { let _ = cmd @@ -76,53 +163,56 @@ impl ConsumerOffsetManager { return false; } - self.pending_offset_mutations - .push(PendingOffsetMutation::new( - OffsetRecord::OffsetCommit(cmd.update), - LeaderOffsetCommitApplied { - replica_set, - reply: cmd.reply, - }, - )); + self.push_pending_offset_mutation(PendingOffsetMutation::new( + OffsetRecord::OffsetCommit(cmd.update), + LeaderOffsetCommitApplied { + replica_set, + required_followers, + reply: cmd.reply, + }, + )); + true } pub(crate) fn handle_offset_checkpoint_complete(&mut self, checkpointed_lsn: u64) { - self.offset_checkpoint_lsn = self.offset_checkpoint_lsn.max(checkpointed_lsn); + self.checkpoint.checkpoint_lsn = self.checkpoint.checkpoint_lsn.max(checkpointed_lsn); // Clear now-safely-persisted LSNs from the uncheckpointed tracking queue while self - .uncheckpointed_offset_lsns + .checkpoint + .uncheckpointed_lsns .front() .is_some_and(|lsn| *lsn <= checkpointed_lsn) { - self.uncheckpointed_offset_lsns.pop_front(); + self.checkpoint.uncheckpointed_lsns.pop_front(); } - self.offset_checkpoint_in_flight = false; + self.checkpoint.in_flight = false; } // 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. pub(crate) fn reclamation_watermark(&self, reclaimable_lsn: u64) -> u64 { - self.uncheckpointed_offset_lsns + self.checkpoint + .uncheckpointed_lsns .front() .map(|lsn| lsn.saturating_sub(1)) .unwrap_or(reclaimable_lsn) } pub(crate) fn oldest_uncheckpointed_lsn(&self) -> Option<&u64> { - self.uncheckpointed_offset_lsns.front() + self.checkpoint.uncheckpointed_lsns.front() } - pub(crate) fn latest_uncheckpointed_lsn(&self) -> Option<&u64> { - self.uncheckpointed_offset_lsns.back() + fn latest_uncheckpointed_lsn(&self) -> Option<&u64> { + self.checkpoint.uncheckpointed_lsns.back() } pub(crate) fn raise_offset_checkpoint_job( &mut self, data_dir: PathBuf, ) -> Option { - self.offset_checkpoint_in_flight = true; + self.checkpoint.in_flight = true; Some(OffsetCheckpointJob { checkpointed_lsn: *self.latest_uncheckpointed_lsn()?, offsets: self.offset_ledger.clone(), @@ -130,17 +220,312 @@ impl ConsumerOffsetManager { }) } - pub(crate) fn add_placement(&mut self, segment_key: &SegmentKey, replica_set: Vec) { - if replica_set.is_empty() { - return; + pub(crate) fn offset_checkpoint_in_flight(&self) -> bool { + self.checkpoint.in_flight + } + + pub(crate) fn observe_placement( + &mut self, + segment_key: SegmentKey, + shard_group_id: ShardGroupId, + replicas: &Replicas, + node_id: &NodeId, + ) -> PlacementObservation { + let observation = self.compare_placement(&segment_key, replicas); + if observation != PlacementObservation::Accepted { + return observation; } - let replicas = Replicas::new(replica_set); - let leader = replicas.leader().cloned().unwrap(); - self.offset_placements.insert( - (segment_key.topic_id, segment_key.range_id), - OffsetPlacement { leader, replicas }, + // need to place! + let local_placement_installed = self.offset_ledger.has_installed_placement(&segment_key); + let replica_states = replicas + .iter() + .cloned() + .map(|node| { + let state = if node == *node_id && local_placement_installed { + OffsetReplicaState::Ready + } else { + OffsetReplicaState::Joining + }; + (node, state) + }) + .collect(); + self.coordination.placements.insert( + segment_key.placement_key(), + OffsetPlacement { + segment_key, + shard_group_id, + replicas: replicas.clone(), + replica_states, + placement_ack_sent: false, + }, ); + PlacementObservation::Accepted + } + + fn compare_placement( + &self, + segment_key: &SegmentKey, + replicas: &Replicas, + ) -> PlacementObservation { + self.coordination + .placements + .get(&(segment_key.topic_id, segment_key.range_id)) + .map(|placement| placement.compare(segment_key, replicas)) + .unwrap_or(PlacementObservation::Accepted) + } + + /// 1. records placement to [`ConsumerOffsetCoordination`](Self::coordination) + /// 2. Takes a State Snapshot + /// 3. Initiates Follower Bootstrapping + pub(crate) fn install_leader_placement( + &mut self, + segment_key: SegmentKey, + shard_group_id: ShardGroupId, + replicas: Replicas, + ) -> Option { + let leader = replicas.leader().cloned()?; + + let placement_key = segment_key.placement_key(); + let replica_states = + self.derive_offset_commit_replicas(segment_key, &replicas, placement_key); + + match self.observe_placement(segment_key, shard_group_id, &replicas, &leader) { + PlacementObservation::Stale | PlacementObservation::Conflict => return None, + PlacementObservation::Accepted => { + let placement = self.get_placement_mut(&placement_key).unwrap(); + placement.replica_states = replica_states; + } + PlacementObservation::Unchanged => { + let placement = self.get_placement_mut(&placement_key).unwrap(); + placement.placement_ack_sent = false; + } + } + + let placement = self + .get_placement(&placement_key) + .expect("accepted or unchanged placement must exist"); + + if !placement.is_replica_ready(&leader) { + let targets: Box<[_]> = placement.replicas.followers().cloned().collect(); + if targets.is_empty() { + return None; + } + return Some(DataTransportCommand::send_to_targets( + targets, + RequestConsumerOffsetSnapshot { + segment_key, + replicas: placement.replicas.clone(), + requester: leader, + }, + )); + } + + let targets = placement.unready_followers(); + if targets.is_empty() { + return None; + } + Some(DataTransportCommand::send_to_targets( + targets, + self.raise_snapshot(segment_key, placement.replicas.clone()), + )) + } + + fn derive_offset_commit_replicas( + &self, + segment_key: SegmentKey, + replicas: &Replicas, + placement_key: (TopicId, RangeId), + ) -> HashMap { + let leader = replicas.leader().unwrap(); + let current = self.get_placement(&placement_key); + let mut replica_states: HashMap<_, _> = replicas + .iter() + .cloned() + .map(|replica| { + let state = if current.is_some_and(|placement| placement.is_replica_ready(&replica)) + { + OffsetReplicaState::Ready + } else { + OffsetReplicaState::Joining + }; + (replica, state) + }) + .collect(); + + if segment_key.segment_id == SegmentId(0) + || self.offset_ledger.has_installed_placement(&segment_key) + { + replica_states.insert(leader.clone(), OffsetReplicaState::Ready); + } + replica_states + } + + fn get_placement(&self, placement_key: &(TopicId, RangeId)) -> Option<&OffsetPlacement> { + self.coordination.placements.get(placement_key) + } + + fn get_placement_followers(&self, cmd: &CommitConsumerOffset) -> HashSet { + self.get_placement(&cmd.update.key.placement_key()) + .map(|placement| { + placement + .replica_states + .keys() + .filter(|n| placement.is_replica_ready(n) && *n != placement.leader()) + .cloned() + .collect() + }) + .unwrap_or_default() + } + + fn get_placement_mut( + &mut self, + placement_key: &(TopicId, RangeId), + ) -> Option<&mut OffsetPlacement> { + self.coordination.placements.get_mut(placement_key) + } + + pub(crate) fn install_consumer_offsets( + &mut self, + cmd: ConsumerOffsetSnapshot, + node_id: &NodeId, + ) -> Vec { + let leader = cmd.replica_set.leader().cloned().unwrap(); + 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, + }, + )); + bootstrapped_keys + .into_iter() + .flat_map(|key| self.take_future_commits(key)) + .collect() + } + + pub(crate) fn create_offset_snapshot( + &self, + request: RequestConsumerOffsetSnapshot, + source_node_id: &NodeId, + ) -> Option { + let live_ready = self + .get_placement(&request.segment_key.placement_key()) + .is_some_and(|placement| { + placement.segment_key == request.segment_key + && placement.is_replica_ready(source_node_id) + }); + + // 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) + { + return None; + } + Some(DataTransportCommand::send_to_targets( + vec![request.requester.clone()], + self.raise_snapshot(request.segment_key, request.replicas), + )) + } + + fn raise_snapshot( + &self, + segment_key: SegmentKey, + replica_set: Replicas, + ) -> ConsumerOffsetSnapshot { + ConsumerOffsetSnapshot { + entries: self + .offset_ledger + .snapshot_range(segment_key.topic_id, segment_key.range_id), + segment_key, + replica_set, + } + } + + pub(crate) fn create_offset_snapshot_for_unready_replicas( + &self, + segment_key: SegmentKey, + ) -> Option { + let placement = self.get_placement(&segment_key.placement_key())?; + if !placement.can_bootstrap_replicas(&segment_key) { + return None; + } + + let targets = placement.unready_followers(); + + if targets.is_empty() { + return None; + } + + Some(DataTransportCommand::send_to_targets( + targets, + self.raise_snapshot(segment_key, placement.replicas.clone()), + )) + } + + pub(crate) fn handle_snapshot_installed_ack(&mut self, ack: &ConsumerOffsetSnapshotInstalled) { + let Some(placement) = self.get_placement_mut(&ack.segment_key.placement_key()) else { + return; + }; + if placement.segment_key != ack.segment_key || !placement.replicas.contains(&ack.from) { + return; + } + placement + .replica_states + .insert(ack.from.clone(), OffsetReplicaState::SnapshotInstalled); + self.mark_offset_replica_ready_if_caught_up(&ack.from); + } + + fn mark_offset_replica_ready_if_caught_up(&mut self, node: &NodeId) -> bool { + if self.coordination.offset_replication.has_pending_for(node) { + return false; + } + let mut readiness_changed = false; + for placement in self.coordination.placements.values_mut() { + if placement.replica_states.get(node) == Some(&OffsetReplicaState::SnapshotInstalled) { + placement + .replica_states + .insert(node.clone(), OffsetReplicaState::Ready); + readiness_changed = true; + } + } + readiness_changed + } + + pub(crate) fn drain_ready_placements(&mut self, from: &NodeId) -> Vec { + self.coordination + .placements + .values_mut() + .filter_map(|placement| { + if !placement.placement_ack_sent + && placement.leader() == from + && placement.all_replicas_ready() + { + placement.placement_ack_sent = true; + return Some(SegmentPlaced { + segment_key: placement.segment_key, + shard_group_id: placement.shard_group_id, + from: from.clone(), + }); + } + None + }) + .collect() } pub(crate) fn get_replica_set_if_leader( @@ -148,78 +533,59 @@ impl ConsumerOffsetManager { key: &ConsumerOffsetKey, node_id: &NodeId, ) -> Result> { - let Some(placement) = self.offset_placements.get(&(key.topic_id, key.range_id)) else { + let Some(placement) = self.get_placement(&key.placement_key()) else { return Err(None); }; - if placement.leader != *node_id { - return Err(Some(placement.leader.clone())); + if placement.leader() != node_id { + return Err(Some(placement.leader().clone())); + } + if !placement.is_replica_ready(node_id) { + return Err(None); } Ok(placement.replicas.clone()) } pub(crate) fn has_pending_offsets(&self) -> bool { - !self.pending_offset_mutations.is_empty() + !self.coordination.pending_offset_mutations.is_empty() } - pub(crate) fn handle_replica_offset_commit( - &mut self, - cmd: ReplicaOffsetCommit, - sealed_generation: GenerationId, - ) -> bool { - if cmd.update.generation > sealed_generation { - self.future_entry(cmd.update.key.clone()) - .push(FutureOffsetCommit::Replica(cmd)); - return false; - } - self.pending_offset_mutations - .push(PendingOffsetMutation::new( - OffsetRecord::OffsetCommit(cmd.update.clone()), - cmd, - )); - true + pub(crate) fn push_pending_offset_mutation(&mut self, cmd: impl Into) { + self.coordination.pending_offset_mutations.push(cmd.into()); } - pub(crate) fn handle_replica_offset_ack(&mut self, ack: ReplicaOffsetAck) { - self.offset_replication.process_ack(ack); + 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.mark_offset_replica_ready_if_caught_up(&from) } pub(crate) fn handle_consumer_group_epoch_seal( &mut self, cmd: EpochSeal, ) -> Option> { + let key = cmd.key.clone(); // Idempotency - if cmd.generation <= self.latest_generation(&cmd.key) { + if cmd.generation <= self.latest_generation(&key) { return None; } - - self.pending_offset_mutations - .push(PendingOffsetMutation::new( - OffsetRecord::EpochSeal(cmd.clone()), - OffsetMutationCompletion::EpochSeal, - )); - - Some( - self.future_offset_commits - .remove(&cmd.key) - .unwrap_or_default(), - ) + self.push_pending_offset_mutation(cmd); + Some(self.take_future_commits(key)) } pub(crate) fn latest_generation(&self, key: &ConsumerOffsetKey) -> GenerationId { - self.pending_offset_mutations + self.coordination + .pending_offset_mutations .iter() .rev() // Look backwards through the pending queue (newest first) - .find_map(|pending| { - // Find the most recent EpochSeal for this specific key - if let OffsetRecord::EpochSeal(EpochSeal { + .find_map(|pending| match &pending.record { + OffsetRecord::EpochSeal(EpochSeal { key: pending_key, generation, - }) = &pending.record - && pending_key == key - { - return Some(*generation); + }) if pending_key == key => Some(*generation), + OffsetRecord::BootstrapEntry(snapshot) if &snapshot.key == key => { + Some(snapshot.generation) } - None + _ => None, }) // No pending EpochSeals, fall back to the actual ledger's generation .unwrap_or_else(|| self.offset_ledger.generation(key)) @@ -227,6 +593,7 @@ impl ConsumerOffsetManager { 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)) @@ -235,7 +602,7 @@ impl ConsumerOffsetManager { 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) { + for pending in std::mem::take(&mut self.coordination.pending_offset_mutations) { if let OffsetMutationCompletion::LeaderCommit(commit) = pending.completion { let _ = commit .reply @@ -248,11 +615,11 @@ impl ConsumerOffsetManager { } pub(crate) fn take_pending(&mut self) -> Vec { - std::mem::take(&mut self.pending_offset_mutations) + std::mem::take(&mut self.coordination.pending_offset_mutations) } pub(crate) fn fail_all(&mut self, error: &str) { - self.offset_replication.fail_all(error); + self.coordination.offset_replication.fail_all(error); for pending in self.take_pending() { if let OffsetMutationCompletion::LeaderCommit(commit) = pending.completion { @@ -263,6 +630,7 @@ impl ConsumerOffsetManager { } for parked in self + .coordination .future_offset_commits .drain() .flat_map(|(_, commits)| commits) @@ -302,15 +670,20 @@ impl ConsumerOffsetManager { continue; } - let seq = self.offset_replication.begin( + let seq = self.coordination.offset_replication.begin( followers.clone().into_iter().collect(), + commit.required_followers, offset_commit.clone(), commit.reply, ); transports.push(DataTransportCommand::send_to_targets( followers, - ReplicaOffsetCommit { + ReplicateConsumerOffset { seq, + segment_key: self + .get_placement(&offset_commit.key.placement_key()) + .map(|placement| placement.segment_key) + .unwrap(), replica_set: commit.replica_set, update: offset_commit, }, @@ -322,21 +695,45 @@ impl ConsumerOffsetManager { }; transports.push(DataTransportCommand::send_to_targets( vec![leader.clone()], - ReplicaOffsetAck { + ConsumerOffsetReplicated { seq: commit.seq, update: commit.update, from: node_id.clone(), - result: ReplicaOffsetAckResult::Committed, + result: ConsumerOffsetReplicationResult::Committed, }, )); } + OffsetMutationCompletion::Bootstrap(ack) => { + transports.push(DataTransportCommand::send_to_targets( + vec![ack.leader.clone()], + ack, + )); + } } } - self.uncheckpointed_offset_lsns + self.checkpoint + .uncheckpointed_lsns .push_back(uncheckpointed_lsn); transports } + + fn take_future_commits(&mut self, key: ConsumerOffsetKey) -> Vec { + self.coordination + .future_offset_commits + .remove(&key) + .unwrap_or_default() + } + + #[cfg(test)] + pub(crate) fn uncheckpointed_offset_lsns(&self) -> &VecDeque { + &self.checkpoint.uncheckpointed_lsns + } + + #[cfg(test)] + pub(crate) fn offset_checkpoint_lsn(&self) -> u64 { + self.checkpoint.checkpoint_lsn + } } #[cfg(test)] @@ -348,6 +745,7 @@ mod tests { use crate::data_plane::consumer_offset_management::ledger::{ ConsumerOffsetPosition, ConsumerOffsetUpdate, }; + use crate::data_plane::messages::DataPlanePeerMessage; use tokio::sync::oneshot; fn key() -> ConsumerOffsetKey { @@ -366,11 +764,338 @@ mod tests { } } + fn segment(id: u64) -> SegmentKey { + SegmentKey::new(TopicId(1), RangeId(2), id.into()) + } + + #[test] + fn observed_placement_restores_local_durable_readiness() { + let local = NodeId::new("b"); + let current = segment(2); + let mut ledger = OffsetLedger::default(); + ledger.apply(OffsetRecord::PlacementInstalled(current)); + let mut manager = ConsumerOffsetManager::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 + .placements + .get(&(current.topic_id, current.range_id)) + .unwrap(); + assert!(placement.is_replica_ready(&local)); + } + + #[test] + fn repeated_follower_placement_preserves_live_coordination_state() { + let local = NodeId::new("b"); + let leader = NodeId::new("a"); + let replicas = Replicas::new(vec![leader, local.clone()]); + let current = segment(2); + let mut manager = ConsumerOffsetManager::new(OffsetLedger::default()); + manager.observe_placement(current, ShardGroupId(0), &replicas, &local); + let live_placement = manager + .coordination + .placements + .get_mut(&(current.topic_id, current.range_id)) + .unwrap(); + live_placement + .replica_states + .insert(local.clone(), OffsetReplicaState::SnapshotInstalled); + live_placement.placement_ack_sent = true; + + manager.observe_placement(current, ShardGroupId(0), &replicas, &local); + + let placement = manager + .coordination + .placements + .get(&(current.topic_id, current.range_id)) + .unwrap(); + assert_eq!( + placement.replica_states.get(&local), + Some(&OffsetReplicaState::SnapshotInstalled) + ); + assert!(placement.placement_ack_sent); + } + + #[test] + fn stale_follower_batch_cannot_replace_newer_placement() { + let local = NodeId::new("b"); + let current_replicas = Replicas::new(vec![NodeId::new("a"), local.clone()]); + 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()); + manager.observe_placement(current, ShardGroupId(0), ¤t_replicas, &local); + + manager.observe_placement(stale, ShardGroupId(0), &stale_replicas, &local); + + let placement = manager + .coordination + .placements + .get(&(current.topic_id, current.range_id)) + .unwrap(); + assert_eq!(placement.segment_key, current); + assert_eq!(placement.replicas, current_replicas); + } + + #[test] + fn leader_placement_rejects_stale_and_conflicting_assignments() { + let leader = NodeId::new("a"); + 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()); + + manager.install_leader_placement(current, ShardGroupId(7), current_replicas.clone()); + + assert!( + manager + .install_leader_placement( + segment(1), + ShardGroupId(7), + Replicas::new(vec![leader.clone(), NodeId::new("old")]), + ) + .is_none() + ); + assert!( + manager + .install_leader_placement( + current, + ShardGroupId(7), + Replicas::new(vec![leader, NodeId::new("conflict")]), + ) + .is_none() + ); + + let placement = manager + .coordination + .placements + .get(&(current.topic_id, current.range_id)) + .unwrap(); + assert_eq!(placement.segment_key, current); + assert_eq!(placement.replicas, current_replicas); + } + + #[test] + fn a_new_unready_leader_cannot_commit_or_publish_snapshots() { + let old_leader = NodeId::new("a"); + let retained = NodeId::new("b"); + let new_leader = NodeId::new("d"); + let mut manager = ConsumerOffsetManager::new(OffsetLedger::default()); + manager.install_leader_placement( + segment(1), + ShardGroupId(7), + Replicas::new(vec![old_leader, retained.clone()]), + ); + manager + .get_placement_mut(&segment(1).placement_key()) + .unwrap() + .replica_states + .insert(retained.clone(), OffsetReplicaState::Ready); + + let snapshots = manager.install_leader_placement( + segment(2), + ShardGroupId(7), + Replicas::new(vec![new_leader.clone(), retained.clone()]), + ); + + assert!(snapshots.iter().any(|command| matches!( + command, + DataTransportCommand::SendToTargets(send) + if send.targets.contains(&retained) + && matches!( + send.message, + DataPlanePeerMessage::RequestConsumerOffsetSnapshot(_) + ) + ))); + let placement = manager.get_placement(&segment(2).placement_key()).unwrap(); + assert!(placement.is_replica_ready(&retained)); + assert!(!placement.is_replica_ready(&new_leader)); + + let (reply, mut response) = oneshot::channel(); + assert!(!manager.commit_consumer_offset( + CommitConsumerOffset { + update: ConsumerOffsetUpdate { + key: key(), + generation: GenerationId(0), + position: position(), + }, + reply, + }, + &new_leader, + )); + assert_eq!( + response.try_recv().unwrap(), + ConsumerOffsetCommitAck::NotWriteLeader(None) + ); + } + + #[test] + fn only_a_durably_ready_replica_can_bootstrap_a_new_leader() { + let source = NodeId::new("b"); + let new_leader = NodeId::new("d"); + let old_segment = segment(1); + let new_segment = segment(2); + let replica_set = Replicas::new(vec![new_leader.clone(), source]); + let request = RequestConsumerOffsetSnapshot { + segment_key: new_segment, + replicas: replica_set.clone(), + requester: new_leader.clone(), + }; + + let empty = ConsumerOffsetManager::new(OffsetLedger::default()); + assert!( + empty + .create_offset_snapshot(request.clone(), &NodeId::new("empty")) + .is_none() + ); + + let mut ledger = OffsetLedger::default(); + ledger.apply(OffsetRecord::PlacementInstalled(old_segment)); + ledger.apply(OffsetRecord::EpochSeal(EpochSeal { + key: key(), + generation: GenerationId(1), + })); + ledger.apply(OffsetRecord::OffsetCommit(ConsumerOffsetUpdate { + key: key(), + generation: GenerationId(1), + position: position(), + })); + let source_manager = ConsumerOffsetManager::new(ledger); + + let skipped_request = RequestConsumerOffsetSnapshot { + segment_key: segment(3), + replicas: replica_set.clone(), + requester: new_leader.clone(), + }; + assert!( + source_manager + .create_offset_snapshot(skipped_request, &NodeId::new("b")) + .is_none(), + "a replica removed for an intervening placement is stale" + ); + + let DataTransportCommand::SendToTargets(send) = source_manager + .create_offset_snapshot(request, &NodeId::new("b")) + .unwrap() + else { + panic!("snapshot must be sent directly to the new leader"); + }; + assert_eq!(send.targets.as_ref(), &[new_leader]); + let DataPlanePeerMessage::InstallConsumerOffsetSnapshot(snapshot) = send.message else { + panic!("expected an offset snapshot"); + }; + assert_eq!(snapshot.segment_key, new_segment); + assert_eq!(snapshot.replica_set, replica_set); + assert_eq!(snapshot.entries.len(), 1); + } + + #[test] + fn joining_replica_is_not_required_until_bootstrap_and_commit_drain() { + let leader = NodeId::new("a"); + let retained = NodeId::new("b"); + let removed = NodeId::new("c"); + let joining = NodeId::new("d"); + let mut ledger = OffsetLedger::default(); + ledger.apply(OffsetRecord::EpochSeal(EpochSeal { + key: key(), + generation: GenerationId(1), + })); + ledger.apply(OffsetRecord::OffsetCommit(ConsumerOffsetUpdate { + key: key(), + generation: GenerationId(1), + position: position(), + })); + let mut manager = ConsumerOffsetManager::new(ledger); + manager.coordination.placements.insert( + (TopicId(1), RangeId(2)), + OffsetPlacement { + segment_key: segment(1), + shard_group_id: ShardGroupId(7), + replicas: Replicas::new(vec![leader.clone(), retained.clone(), removed.clone()]), + replica_states: HashMap::from([ + (leader.clone(), OffsetReplicaState::Ready), + (retained.clone(), OffsetReplicaState::Ready), + (removed, OffsetReplicaState::Ready), + ]), + placement_ack_sent: true, + }, + ); + + let bootstraps = manager.install_leader_placement( + segment(2), + ShardGroupId(7), + Replicas::new(vec![leader.clone(), retained.clone(), joining.clone()]), + ); + assert!(bootstraps.is_some(), "only the joining replica bootstraps"); + + let advanced = ConsumerOffsetPosition { + entry_id: EntryId(4), + batch_offset: 0, + absolute_offset: 5, + }; + let update = ConsumerOffsetUpdate { + key: key(), + generation: GenerationId(1), + position: advanced, + }; + let (reply, mut response) = oneshot::channel(); + assert!(manager.commit_consumer_offset( + CommitConsumerOffset { + update: update.clone(), + reply, + }, + &leader, + )); + let transports = manager.flush_batch(&leader, 1); + let replicated = transports + .into_iter() + .find_map(|transport| match transport { + DataTransportCommand::SendToTargets(send) => match send.message { + DataPlanePeerMessage::ReplicateConsumerOffset(commit) => Some(commit), + _ => None, + }, + _ => None, + }) + .unwrap(); + assert!(response.try_recv().is_err()); + + manager.handle_snapshot_installed_ack(&ConsumerOffsetSnapshotInstalled { + segment_key: segment(2), + from: joining.clone(), + leader: leader.clone(), + }); + assert!(manager.drain_ready_placements(&leader).is_empty()); + + manager.handle_replica_offset_ack(ConsumerOffsetReplicated { + seq: replicated.seq, + update: update.clone(), + from: retained, + result: ConsumerOffsetReplicationResult::Committed, + }); + assert_eq!( + response.try_recv().unwrap(), + ConsumerOffsetCommitAck::Committed + ); + assert!(manager.drain_ready_placements(&leader).is_empty()); + + manager.handle_replica_offset_ack(ConsumerOffsetReplicated { + seq: replicated.seq, + update, + from: joining, + result: ConsumerOffsetReplicationResult::Committed, + }); + assert_eq!(manager.drain_ready_placements(&leader).len(), 1); + } + #[test] fn fail_all_resolves_pending_in_flight_and_future_clients() { let mut manager = ConsumerOffsetManager::new(OffsetLedger::default()); let (pending_reply, mut pending_response) = oneshot::channel(); manager + .coordination .pending_offset_mutations .push(PendingOffsetMutation::new( OffsetRecord::OffsetCommit(ConsumerOffsetUpdate { @@ -380,12 +1105,13 @@ mod tests { }), LeaderOffsetCommitApplied { replica_set: Replicas::new(vec![NodeId::new("leader")]), + required_followers: HashSet::new(), reply: pending_reply, }, )); let (future_reply, mut future_response) = oneshot::channel(); - manager.future_offset_commits.insert( + manager.coordination.future_offset_commits.insert( key(), vec![FutureOffsetCommit::Client(CommitConsumerOffset { update: ConsumerOffsetUpdate { @@ -398,7 +1124,8 @@ mod tests { ); let (replication_reply, mut replication_response) = oneshot::channel(); - manager.offset_replication.begin( + manager.coordination.offset_replication.begin( + HashSet::from_iter([NodeId::new("follower")]), HashSet::from_iter([NodeId::new("follower")]), ConsumerOffsetUpdate { key: key(), @@ -414,7 +1141,7 @@ mod tests { assert_eq!(pending_response.try_recv().unwrap(), expected); assert_eq!(future_response.try_recv().unwrap(), expected); assert_eq!(replication_response.try_recv().unwrap(), expected); - assert!(manager.pending_offset_mutations.is_empty()); - assert!(manager.future_offset_commits.is_empty()); + assert!(manager.coordination.pending_offset_mutations.is_empty()); + assert!(manager.coordination.future_offset_commits.is_empty()); } } diff --git a/src/data_plane/consumer_offset_management/replication.rs b/src/data_plane/consumer_offset_management/replication.rs index a4dde597..5be7a558 100644 --- a/src/data_plane/consumer_offset_management/replication.rs +++ b/src/data_plane/consumer_offset_management/replication.rs @@ -1,7 +1,7 @@ use crate::control_plane::NodeId; use crate::data_plane::consumer_offset_management::ledger::ConsumerOffsetUpdate; use crate::data_plane::messages::command::{ - ConsumerOffsetCommitAck, ReplicaOffsetAck, ReplicaOffsetAckResult, + ConsumerOffsetCommitAck, ConsumerOffsetReplicated, ConsumerOffsetReplicationResult, }; use std::collections::{HashMap, HashSet}; use tokio::sync::oneshot; @@ -15,40 +15,57 @@ pub(crate) struct OffsetReplicationState { struct PendingOffsetReplication { update: ConsumerOffsetUpdate, pending_acks: HashSet, - reply: oneshot::Sender, + required_acks: HashSet, + reply: Option>, } impl PendingOffsetReplication { fn need_more_acks(&self) -> bool { - !self.pending_acks.is_empty() + !self.required_acks.is_empty() } - fn is_valid_ack(&self, c: &ReplicaOffsetAck) -> bool { + fn is_valid_ack(&self, c: &ConsumerOffsetReplicated) -> bool { self.update == c.update && self.pending_acks.contains(&c.from) } + + fn is_consumer_waiting(&self) -> bool { + self.reply + .as_ref() + .is_some_and(|pending_reply| !pending_reply.is_closed()) + } } impl OffsetReplicationState { pub(crate) fn begin( &mut self, followers: HashSet, + required: HashSet, update: ConsumerOffsetUpdate, reply: oneshot::Sender, ) -> u64 { - self.pending.retain(|_, pending| !pending.reply.is_closed()); + self.pending + .retain(|_, pending| pending.need_more_acks() || pending.is_consumer_waiting()); + self.seq = self.seq.wrapping_add(1); + let reply = if required.is_empty() { + let _ = reply.send(ConsumerOffsetCommitAck::Committed); + None + } else { + Some(reply) + }; self.pending.insert( self.seq, PendingOffsetReplication { update: update.clone(), pending_acks: followers, + required_acks: required, reply, }, ); self.seq } - pub(crate) fn process_ack(&mut self, ack: ReplicaOffsetAck) { + pub(crate) fn process_ack(&mut self, ack: ConsumerOffsetReplicated) { let Some(pending) = self.pending.get_mut(&ack.seq) else { return; }; @@ -59,27 +76,39 @@ impl OffsetReplicationState { return; } - if let ReplicaOffsetAckResult::StaleEpoch(stale) = ack.result { - let rejected = self.pending.remove(&ack.seq).unwrap(); - let _ = rejected - .reply - .send(ConsumerOffsetCommitAck::StaleEpoch(stale)); - return; + if let ConsumerOffsetReplicationResult::StaleEpoch(stale) = ack.result { + pending.pending_acks.remove(&ack.from); + if pending.required_acks.remove(&ack.from) + && let Some(reply) = pending.reply.take() + { + let _ = reply.send(ConsumerOffsetCommitAck::StaleEpoch(stale)); + } + } else { + pending.pending_acks.remove(&ack.from); + pending.required_acks.remove(&ack.from); } - pending.pending_acks.remove(&ack.from); - if pending.need_more_acks() { - return; + if !pending.need_more_acks() + && let Some(reply) = pending.reply.take() + { + let _ = reply.send(ConsumerOffsetCommitAck::Committed); + } + if pending.pending_acks.is_empty() { + self.pending.remove(&ack.seq); } - let committed = self.pending.remove(&ack.seq).unwrap(); - let _ = committed.reply.send(ConsumerOffsetCommitAck::Committed); + } + + pub(crate) fn has_pending_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) { - for (_, pending) in self.pending.drain() { - let _ = pending - .reply - .send(ConsumerOffsetCommitAck::InternalError(error.to_string())); + 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/types.rs b/src/data_plane/consumer_offset_management/types.rs index 46a9515c..d673b75e 100644 --- a/src/data_plane/consumer_offset_management/types.rs +++ b/src/data_plane/consumer_offset_management/types.rs @@ -1,6 +1,13 @@ +use std::collections::HashSet; + use crate::control_plane::{NodeId, Replicas}; -use crate::data_plane::consumer_offset_management::ledger::{ConsumerOffsetUpdate, OffsetRecord}; -use crate::data_plane::messages::command::{CommitConsumerOffset, ConsumerOffsetCommitAck}; +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, +}; use crate::impl_from_variant; use borsh::{BorshDeserialize, BorshSerialize}; @@ -22,36 +29,57 @@ impl PendingOffsetMutation { } } +impl From for PendingOffsetMutation { + fn from(cmd: ReplicateConsumerOffset) -> Self { + Self::new(OffsetRecord::OffsetCommit(cmd.update.clone()), cmd) + } +} +impl From for PendingOffsetMutation { + fn from(cmd: EpochSeal) -> Self { + PendingOffsetMutation::new( + OffsetRecord::EpochSeal(cmd.clone()), + OffsetMutationCompletion::EpochSeal, + ) + } +} + pub(crate) struct LeaderOffsetCommitApplied { pub(crate) replica_set: Replicas, + pub(crate) required_followers: HashSet, pub(crate) reply: oneshot::Sender, } pub(crate) enum OffsetMutationCompletion { EpochSeal, LeaderCommit(LeaderOffsetCommitApplied), - ReplicaCommit(ReplicaOffsetCommit), + ReplicaCommit(ReplicateConsumerOffset), + Bootstrap(ConsumerOffsetSnapshotInstalled), } impl_from_variant!( OffsetMutationCompletion, LeaderCommit(LeaderOffsetCommitApplied), - ReplicaCommit(ReplicaOffsetCommit) + ReplicaCommit(ReplicateConsumerOffset), + Bootstrap(ConsumerOffsetSnapshotInstalled) ); #[derive(Debug, Clone, BorshSerialize, BorshDeserialize)] -pub struct ReplicaOffsetCommit { +pub struct ReplicateConsumerOffset { pub seq: u64, + pub segment_key: SegmentKey, pub replica_set: Replicas, pub update: ConsumerOffsetUpdate, } pub(crate) enum FutureOffsetCommit { Client(CommitConsumerOffset), - Replica(ReplicaOffsetCommit), + Replica(ReplicateConsumerOffset), } -pub(crate) struct OffsetPlacement { - pub(crate) leader: NodeId, - pub(crate) replicas: Replicas, +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum PlacementObservation { + Accepted, + Unchanged, + Stale, + Conflict, } diff --git a/src/data_plane/messages/command.rs b/src/data_plane/messages/command.rs index fdab4575..d02f96a6 100644 --- a/src/data_plane/messages/command.rs +++ b/src/data_plane/messages/command.rs @@ -1,7 +1,9 @@ +use crate::control_plane::Replicas; +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::ReplicaOffsetCommit; +use crate::data_plane::consumer_offset_management::types::ReplicateConsumerOffset; use crate::impl_from_variant; use crate::impl_from_variant_via; use borsh::{BorshDeserialize, BorshSerialize}; @@ -22,14 +24,20 @@ pub enum DataPlaneCommand { SegmentCheckpointComplete(SegmentCheckpointComplete), OffsetCheckpointComplete(OffsetCheckpointComplete), DataPlaneTimeoutCallback(DataPlaneTimeoutCallback), - DataPlaneInterNodeCommand(DataPlaneInterNodeCommand), + ReceivePeerMessage(ReceivePeerMessage), /// Internal (not a wire message): the cold-read pool's reply for a catch-up - /// source read. The worker turns it into `CatchUpChunk`s on the transport. + /// source read. The worker turns it into `CatchUpEntries`s on the transport. CatchUpReadComplete(CatchUpReadComplete), OrphanGcCheck(OrphanGcCheck), CommitConsumerOffset(CommitConsumerOffset), } +#[derive(Debug, Clone, BorshSerialize, BorshDeserialize)] +pub struct ReceivePeerMessage { + pub from: NodeId, + pub message: Box, +} + /// What the orphan-GC handler replies to the ticker: keep ticking while strays remain, or /// stop once `recovered` is drained. pub enum OrphanGcSignal { @@ -64,60 +72,81 @@ pub enum ConsumerOffsetCommitAck { } #[derive(Debug, Clone, BorshSerialize, BorshDeserialize)] -pub struct SegmentAssignment { +pub struct PlaceSegment { pub segment_key: SegmentKey, pub shard_group_id: ShardGroupId, - pub replica_set: Vec, + pub replica_set: Replicas, pub start_entry_id: EntryId, } #[derive(Debug, Clone, BorshSerialize, BorshDeserialize)] -pub struct SegmentAssignmentAck { +pub struct SegmentPlaced { pub segment_key: SegmentKey, pub shard_group_id: ShardGroupId, pub from: NodeId, } #[derive(Debug, Clone, BorshSerialize, BorshDeserialize)] -pub struct ReplicaAppend { +pub struct ReplicateSegmentEntries { pub segment_key: SegmentKey, - pub replica_set: Vec, + pub replicas: Replicas, pub data: EntryPayload, pub record_count: u32, pub entry_id: EntryId, } #[derive(Debug, Clone, BorshSerialize, BorshDeserialize)] -pub struct ReplicaAck { +pub struct ReplicaEntriesAppended { pub segment_key: SegmentKey, pub entry_id: EntryId, pub from: NodeId, } #[derive(Debug, Clone, PartialEq, Eq, BorshSerialize, BorshDeserialize)] -pub enum ReplicaOffsetAckResult { +pub enum ConsumerOffsetReplicationResult { Committed, StaleEpoch(StaleEpoch), } -// impl_from_variant!(ReplicaOffsetAckResult, StaleEpoch); +// impl_from_variant!(ConsumerOffsetReplicationResult, StaleEpoch); #[derive(Debug, Clone, BorshSerialize, BorshDeserialize)] -pub struct ReplicaOffsetAck { +pub struct ConsumerOffsetReplicated { pub seq: u64, pub update: ConsumerOffsetUpdate, pub from: NodeId, - pub result: ReplicaOffsetAckResult, + pub result: ConsumerOffsetReplicationResult, +} + +#[derive(Debug, Clone, BorshSerialize, BorshDeserialize)] +pub struct ConsumerOffsetSnapshot { + pub segment_key: SegmentKey, + pub replica_set: Replicas, + pub entries: Box<[ConsumerOffsetEntry]>, +} + +#[derive(Debug, Clone, BorshSerialize, BorshDeserialize)] +pub struct RequestConsumerOffsetSnapshot { + pub segment_key: SegmentKey, + pub replicas: Replicas, + pub requester: NodeId, +} + +#[derive(Debug, Clone, BorshSerialize, BorshDeserialize)] +pub struct ConsumerOffsetSnapshotInstalled { + pub segment_key: SegmentKey, + pub from: NodeId, + pub leader: NodeId, } #[derive(Debug, Clone, BorshSerialize, BorshDeserialize)] -pub struct CommitAdvance { +pub struct AdvanceReplicaCommit { pub segment_key: SegmentKey, pub committed_entry_id: EntryId, } #[derive(Debug, Clone, BorshSerialize, BorshDeserialize)] -pub struct SealRequest { +pub struct RequestSegmentRoll { pub from: NodeId, pub segment_key: SegmentKey, pub failed_nodes: Vec, @@ -125,10 +154,10 @@ pub struct SealRequest { } #[derive(Debug, Clone, BorshSerialize, BorshDeserialize)] -pub struct SealResponse { +pub struct SegmentRollCommitted { pub old_segment_key: SegmentKey, pub new_segment_id: SegmentId, - pub new_replica_set: Vec, + pub new_replica_set: Replicas, } #[derive(Debug, Clone, BorshSerialize, BorshDeserialize)] @@ -140,31 +169,31 @@ pub struct SegmentSealed { // Catch-up: re-replicate a sealed segment to a newly assigned replica. // // A node assigned a sealed segment it doesn't fully hold fetches the missing -// suffix from a healthy peer. Messages on the `DataPlaneInterNodeCommand` wire: +// suffix from a healthy peer. Messages on the `DataPlanePeerMessage` wire: // -// coordinator ─CatchUpAssignment─▶ each replica "own `key` [start, sealed_end]" -// replacement ─CatchUpRequest────▶ a peer "I have through `local_end`; send the rest" -// source ─CatchUpChunk(s)───▶ replacement batches of entries -// source ─CatchUpStreamEnd──▶ replacement end of stream -// replacement ─CatchUpAck────────▶ coordinator "have it through `sealed_end`" +// coordinator ─AssignSegmentCatchUp───▶ each replica "own `key` [start, sealed_end]" +// replacement ─RequestCatchUpEntries──▶ source replica "send after `local_end`" +// source ─CatchUpEntries─────────▶ replacement batches of entries +// source ─CatchUpEntriesSent─────▶ replacement end of stream +// replacement ─SegmentCaughtUp────────▶ coordinator "have it through `sealed_end`" // -// `CatchUpAssignment`/`CatchUpAck` are the coordinator↔replica pair; the rest is +// `AssignSegmentCatchUp`/`SegmentCaughtUp` are the coordinator↔replica pair; the rest is // the replacement↔source transfer. The ack lets the coordinator stop re-driving. #[derive(Debug, Clone, BorshSerialize, BorshDeserialize)] -pub struct CatchUpAssignment { +pub struct AssignSegmentCatchUp { pub segment_key: SegmentKey, - /// Echoed into the `CatchUpAck` so the reply reaches this group's coordinator. + /// Echoed into the `SegmentCaughtUp` so the reply reaches this group's coordinator. pub shard_group_id: ShardGroupId, pub start_entry_id: EntryId, pub sealed_end_entry_id: EntryId, - pub replica_set: Vec, + pub replica_set: Replicas, } #[derive(Debug, Clone, BorshSerialize, BorshDeserialize)] -pub struct CatchUpRequest { +pub struct RequestCatchUpEntries { pub segment_key: SegmentKey, - /// The requesting node — the source streams its `CatchUpChunk`s back here. + /// The requesting node — the source streams `CatchUpEntries` back here. pub from: NodeId, /// Highest entry id the requester already holds locally; the source streams /// `(local_end, sealed_end]`. `None` means nothing is held. @@ -172,7 +201,7 @@ pub struct CatchUpRequest { } #[derive(Debug, Clone, BorshSerialize, BorshDeserialize)] -pub struct CatchUpChunk { +pub struct CatchUpEntries { pub segment_key: SegmentKey, /// Contiguous entries in ascending `entry_id` order. pub entries: Box<[CatchUpEntry]>, @@ -195,7 +224,7 @@ impl CatchUpEntry { } #[derive(Debug, Clone, BorshSerialize, BorshDeserialize)] -pub struct CatchUpStreamEnd { +pub struct CatchUpEntriesSent { pub segment_key: SegmentKey, } @@ -203,7 +232,7 @@ pub struct CatchUpStreamEnd { /// coordinator stop re-driving. Sent after a transfer verifies, and on a /// zero-transfer full match so a re-drive re-confirms. #[derive(Debug, Clone, BorshSerialize, BorshDeserialize)] -pub struct CatchUpAck { +pub struct SegmentCaughtUp { pub segment_key: SegmentKey, pub shard_group_id: ShardGroupId, pub from: NodeId, @@ -215,22 +244,22 @@ pub struct CatchUpAck { // committed end (the leader was the one tracking it). The coordinator gathers // each survivor's durable (fsync'd) extent and seals at their `min` — the // highest offset present on every survivor, hence committed (commit requires -// all-replica ack). Two messages on the `DataPlaneInterNodeCommand` wire: +// all-replica ack). Two messages on the `DataPlanePeerMessage` wire: // -// coordinator ─SealBoundaryQuery──▶ each survivor "what's your durable extent for `key`?" -// survivor ─SealBoundaryReport─▶ coordinator durable end (or `None` if it holds nothing) +// coordinator ─RequestDurableSegmentEnd──▶ each survivor "what's your durable extent for `key`?" +// survivor ─DurableSegmentEndReported─▶ coordinator durable end (or `None` if it holds nothing) // // The originating coordinator rides the query so every report returns to the // same gather even while shard-leader caches disagree during an election. // See `docs/data-plane/leader_crash_seal_boundary.md`. #[derive(Debug, Clone, BorshSerialize, BorshDeserialize)] -pub struct SealBoundaryQuery { +pub struct RequestDurableSegmentEnd { pub segment_key: SegmentKey, pub coordinator: NodeId, } #[derive(Debug, Clone, BorshSerialize, BorshDeserialize)] -pub struct SealBoundaryReport { +pub struct DurableSegmentEndReported { pub segment_key: SegmentKey, pub from: NodeId, pub durable_end: Option, @@ -248,49 +277,55 @@ pub struct DeleteSegments { } #[derive(Debug, Clone, BorshSerialize, BorshDeserialize)] -pub enum DataPlaneInterNodeCommand { - SegmentAssignment(SegmentAssignment), - SegmentAssignmentAck(SegmentAssignmentAck), - ReplicaAppend(ReplicaAppend), - ReplicaAck(ReplicaAck), - ReplicaOffsetCommit(ReplicaOffsetCommit), - ReplicaOffsetAck(ReplicaOffsetAck), - CommitAdvance(CommitAdvance), - SealRequest(SealRequest), - SealResponse(SealResponse), +pub enum DataPlanePeerMessage { + PlaceSegment(PlaceSegment), + SegmentPlaced(SegmentPlaced), + ReplicateSegmentEntries(ReplicateSegmentEntries), + ReplicaEntriesAppended(ReplicaEntriesAppended), + ReplicateConsumerOffset(ReplicateConsumerOffset), + ConsumerOffsetReplicated(ConsumerOffsetReplicated), + InstallConsumerOffsetSnapshot(ConsumerOffsetSnapshot), + RequestConsumerOffsetSnapshot(RequestConsumerOffsetSnapshot), + ConsumerOffsetSnapshotInstalled(ConsumerOffsetSnapshotInstalled), + AdvanceReplicaCommit(AdvanceReplicaCommit), + RequestSegmentRoll(RequestSegmentRoll), + SegmentRollCommitted(SegmentRollCommitted), SegmentSealed(SegmentSealed), - CatchUpAssignment(CatchUpAssignment), - CatchUpRequest(CatchUpRequest), - CatchUpChunk(CatchUpChunk), - CatchUpStreamEnd(CatchUpStreamEnd), - CatchUpAck(CatchUpAck), - SealBoundaryQuery(SealBoundaryQuery), - SealBoundaryReport(SealBoundaryReport), + AssignSegmentCatchUp(AssignSegmentCatchUp), + RequestCatchUpEntries(RequestCatchUpEntries), + CatchUpEntries(CatchUpEntries), + CatchUpEntriesSent(CatchUpEntriesSent), + SegmentCaughtUp(SegmentCaughtUp), + RequestDurableSegmentEnd(RequestDurableSegmentEnd), + DurableSegmentEndReported(DurableSegmentEndReported), DeleteSegments(DeleteSegments), - ConsumerGroupEpochSeal(EpochSeal), + ConsumerGroupEpochSealed(EpochSeal), } impl_from_variant!( - DataPlaneInterNodeCommand, - SegmentAssignment, - SegmentAssignmentAck, - ReplicaAppend, - ReplicaAck, - ReplicaOffsetCommit, - ReplicaOffsetAck, - CommitAdvance, - SealRequest, - SealResponse, + DataPlanePeerMessage, + PlaceSegment, + SegmentPlaced, + ReplicateSegmentEntries, + ReplicaEntriesAppended, + ReplicateConsumerOffset, + ConsumerOffsetReplicated, + InstallConsumerOffsetSnapshot(ConsumerOffsetSnapshot), + RequestConsumerOffsetSnapshot, + ConsumerOffsetSnapshotInstalled, + AdvanceReplicaCommit, + RequestSegmentRoll, + SegmentRollCommitted, SegmentSealed, - CatchUpAssignment, - CatchUpRequest, - CatchUpChunk, - CatchUpStreamEnd, - CatchUpAck, - SealBoundaryQuery, - SealBoundaryReport, + AssignSegmentCatchUp, + RequestCatchUpEntries, + CatchUpEntries, + CatchUpEntriesSent, + SegmentCaughtUp, + RequestDurableSegmentEnd, + DurableSegmentEndReported, DeleteSegments, - ConsumerGroupEpochSeal(EpochSeal), + ConsumerGroupEpochSealed(EpochSeal), ); #[derive(Debug)] @@ -334,7 +369,7 @@ impl_from_variant!( OrphanGcCheck, CommitConsumerOffset, DataPlaneTimeoutCallback(DataPlaneTimeoutCallback), - DataPlaneInterNodeCommand(DataPlaneInterNodeCommand), + ReceivePeerMessage, ); use crate::data_plane::timer::{BatchFlushCallback, ReplicationCallback, SegmentAgeCallback}; diff --git a/src/data_plane/messages/mod.rs b/src/data_plane/messages/mod.rs index 92b1dc80..4dc7f210 100644 --- a/src/data_plane/messages/mod.rs +++ b/src/data_plane/messages/mod.rs @@ -31,7 +31,7 @@ impl_from_variant_via!( Produce, SegmentCheckpointComplete, DataPlaneTimeoutCallback, - DataPlaneInterNodeCommand, + ReceivePeerMessage, CatchUpReadComplete, CommitConsumerOffset ); diff --git a/src/data_plane/mod.rs b/src/data_plane/mod.rs index 54e9c7ba..39556be9 100644 --- a/src/data_plane/mod.rs +++ b/src/data_plane/mod.rs @@ -48,6 +48,10 @@ impl SegmentKey { } } + pub(crate) fn placement_key(&self) -> (TopicId, RangeId) { + (self.topic_id, self.range_id) + } + /// Path to this segment's file. The filename encodes the segment's /// `start_entry` (its first entry id) so the file is self-describing for /// crash recovery — discovery derives the base entry id from the name diff --git a/src/data_plane/state.rs b/src/data_plane/state.rs index 394eb45c..b5e4171c 100644 --- a/src/data_plane/state.rs +++ b/src/data_plane/state.rs @@ -2,7 +2,7 @@ use super::SegmentKey; use super::actor::DataPlaneSender; use super::cold_read::{CatchUpReadReply, ColdReadReply, ColdReadRequest}; use super::messages::DataPlaneMessage; -use super::messages::command::DataPlaneInterNodeCommand; +use super::messages::command::DataPlanePeerMessage; use super::messages::command::*; use super::messages::pending::DataPlaneOutputs; use super::messages::query::{ @@ -15,20 +15,21 @@ use super::segment_writer::SegmentAppender; use super::states::replication::PendingReplicationBatch; use super::states::replication::ReplicationState; -use super::states::seal_request::PendingSealRequests; +use super::states::seal_request::PendingSegmentRollRequests; use super::states::segment_store::SegmentStore; use super::timer::DataPlaneTimeoutCallback; use super::transport::command::DataTransportCommand; use super::wal::WalRecord; use super::wal::WalStorage; use crate::config::DataNodeConfig; -use crate::control_plane::NodeId; -use crate::control_plane::consensus::messages::{CoordinatorSealRequest, MultiRaftActorCommand}; +use crate::control_plane::consensus::messages::{MultiRaftActorCommand, ProposeSegmentRoll}; use crate::control_plane::membership::ShardGroupId; use crate::control_plane::metadata::EntryId; +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::consumer_offset_management::types::*; use crate::data_plane::messages::query::ReadConsumerOffset; use crate::data_plane::states::segment::tracker::{SegmentRole, SegmentTracker}; @@ -41,7 +42,7 @@ use crate::schedulers::ticker_message::TimerCommand; use crate::test_traits::TAssertInvariant; use std::collections::{BTreeSet, HashMap}; -/// Same rational as the Raft transport's 4MiB cap, Per-`CatchUpChunk` read cap. +/// Same rational as the Raft transport's 4MiB cap, Per-`CatchUpEntries` read cap. /// A large segment streams as several chunks via the read-complete re-arm loop. const CATCH_UP_CHUNK_MAX_BYTES: u64 = 4 * 1024 * 1024; @@ -71,7 +72,7 @@ pub struct DataPlane { replication: ReplicationState, - pending_seal_requests: PendingSealRequests, + pending_seal_requests: PendingSegmentRollRequests, /// In-progress catch-up receives (replacement side), keyed by segment. pending_catch_ups: HashMap, @@ -116,7 +117,7 @@ impl DataPlane { buffer_byte_count: 0, needs_flush: false, replication: ReplicationState::default(), - pending_seal_requests: PendingSealRequests::default(), + pending_seal_requests: PendingSegmentRollRequests::default(), pending_catch_ups: HashMap::new(), cold_read_handoff_sender, self_tx, @@ -148,14 +149,14 @@ impl DataPlane { DataPlaneCommand::DataPlaneTimeoutCallback(callback) => { self.handle_timeout(callback); } - DataPlaneCommand::DataPlaneInterNodeCommand(inter) => { - self.process_inter_node(inter); + DataPlaneCommand::ReceivePeerMessage(received) => { + self.receive_peer_message(received); } DataPlaneCommand::CatchUpReadComplete(cmd) => { self.handle_catch_up_read_complete(cmd); } DataPlaneCommand::OrphanGcCheck(check) => self.handle_orphan_gc_check(check), - DataPlaneCommand::CommitConsumerOffset(cmd) => self.handle_commit_consumer_offset(cmd), + DataPlaneCommand::CommitConsumerOffset(cmd) => self.commit_consumer_offset(cmd), } #[cfg(any(test, debug_assertions))] @@ -199,7 +200,8 @@ impl DataPlane { }; self.buffer_byte_count += cmd.data.len(); - tracker.stage_entry(cmd.segment_key, cmd.data, cmd.record_count); + 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); self.out @@ -391,7 +393,7 @@ impl DataPlane { self.wal.delete_below(watermark); } - if self.consumer_offsets.offset_checkpoint_in_flight { + if self.consumer_offsets.offset_checkpoint_in_flight() { return; } let Some(reclaimable_lsn) = self.wal.reclaimable_lsn() else { @@ -464,57 +466,68 @@ impl DataPlane { let _ = check.reply.try_send(signal); } - fn process_inter_node(&mut self, cmd: DataPlaneInterNodeCommand) { - use DataPlaneInterNodeCommand as C; - match cmd { - C::SegmentAssignment(cmd) => self.handle_segment_assignment(cmd), - C::ReplicaAppend(cmd) => self.process_replica_append(cmd), - C::ReplicaAck(cmd) => self.handle_replica_ack(cmd), - C::ReplicaOffsetCommit(cmd) => self.handle_replica_offset_commit(cmd), - C::ReplicaOffsetAck(cmd) => self.consumer_offsets.handle_replica_offset_ack(cmd), - C::CommitAdvance(cmd) => self.handle_commit_advance(cmd), - C::SealResponse(cmd) => self.handle_seal_response(cmd), + fn receive_peer_message(&mut self, received: ReceivePeerMessage) { + use DataPlanePeerMessage as C; + let ReceivePeerMessage { from, message } = received; + + match *message { + C::PlaceSegment(cmd) => self.place_segment(cmd), + C::ReplicateSegmentEntries(cmd) => self.replicate_segment_entries(&from, cmd), + C::ReplicaEntriesAppended(cmd) => self.handle_replica_ack(cmd), + C::ReplicateConsumerOffset(cmd) => self.replicate_consumer_offset(cmd), + C::ConsumerOffsetReplicated(cmd) => self.handle_replica_offset_ack(cmd), + C::InstallConsumerOffsetSnapshot(cmd) => { + self.install_consumer_offset_snapshot(&from, cmd) + } + C::RequestConsumerOffsetSnapshot(cmd) => { + self.request_consumer_offset_snapshot(&from, cmd) + } + C::ConsumerOffsetSnapshotInstalled(cmd) => { + self.handle_consumer_offset_snapshot_installed(cmd) + } + C::AdvanceReplicaCommit(cmd) => self.handle_commit_advance(cmd), + C::SegmentRollCommitted(message) => self.handle_segment_roll_committed(message), C::SegmentSealed(cmd) => self.handle_segment_sealed(cmd), - // Pass-through to MultiRaftActor — SealRequest is a control plane + // Pass-through to MultiRaftActor — RequestSegmentRoll is a control plane // message that shares the data transport wire format. The // coordinator loads its own live-nodes view from the shared // topology snapshot, so the data plane just forwards the request. - C::SealRequest(cmd) => { + C::RequestSegmentRoll(cmd) => { self.out - .store_coordinator_cmd(MultiRaftActorCommand::Coordinator( - CoordinatorSealRequest { request: cmd }, + .store_coordinator_cmd(MultiRaftActorCommand::ProposeSegmentRoll( + ProposeSegmentRoll { request: cmd }, )); } // Assignment confirmation from a data-leader — forward to the local // MultiRaftActor (we are this shard's coordinator), which marks the // segment confirmed so the heartbeat sweep stops re-driving it. - C::SegmentAssignmentAck(cmd) => { + C::SegmentPlaced(cmd) => { self.out .store_coordinator_cmd(MultiRaftActorCommand::AssignmentAck(cmd)); } // Catch-ups: re-replicate a sealed segment to a newly assigned replica. - C::CatchUpAssignment(cmd) => self.handle_catch_up_assignment(cmd), - C::CatchUpRequest(cmd) => self.handle_catch_up_request(cmd), - C::CatchUpChunk(cmd) => self.handle_catch_up_chunk(cmd), - C::CatchUpStreamEnd(cmd) => self.handle_catch_up_stream_end(cmd), + C::AssignSegmentCatchUp(cmd) => self.handle_catch_up_assignment(cmd), + C::RequestCatchUpEntries(cmd) => self.handle_catch_up_request(cmd), + C::CatchUpEntries(cmd) => self.handle_catch_up_chunk(cmd), + C::CatchUpEntriesSent(cmd) => self.handle_catch_up_stream_end(cmd), // Pass-through to the local MultiRaftActor (we coordinate this group): // a replica's confirmation that it finished catching up. - C::CatchUpAck(cmd) => { + C::SegmentCaughtUp(cmd) => { self.out - .store_coordinator_cmd(MultiRaftActorCommand::CatchUpAck(cmd)); + .store_coordinator_cmd(MultiRaftActorCommand::SegmentCaughtUp(cmd)); } - C::SealBoundaryQuery(cmd) => self.handle_seal_boundary_query(cmd), + C::RequestDurableSegmentEnd(cmd) => self.handle_seal_boundary_query(cmd), // Pass-through to the local MultiRaftActor - C::SealBoundaryReport(cmd) => { + C::DurableSegmentEndReported(cmd) => { self.out - .store_coordinator_cmd(MultiRaftActorCommand::SealBoundaryReport(cmd)); + .store_coordinator_cmd(MultiRaftActorCommand::DurableSegmentEndReported(cmd)); } C::DeleteSegments(cmd) => self.handle_delete_segments(cmd), - C::ConsumerGroupEpochSeal(cmd) => self.handle_consumer_group_epoch_seal(cmd), + C::ConsumerGroupEpochSealed(cmd) => self.handle_consumer_group_epoch_seal(cmd), } } @@ -543,22 +556,28 @@ impl DataPlane { }; self.needs_flush = true; + self.resume_future_offset_commits(parked_commits); + } + + fn resume_future_offset_commits(&mut self, parked_commits: Vec) { for parked in parked_commits { match parked { FutureOffsetCommit::Client(pending) => { - self.handle_commit_consumer_offset(pending); + self.commit_consumer_offset(pending); + } + FutureOffsetCommit::Replica(pending) => { + self.replicate_consumer_offset(pending); } - FutureOffsetCommit::Replica(pending) => self.handle_replica_offset_commit(pending), } } } - fn handle_seal_boundary_query(&mut self, cmd: SealBoundaryQuery) { + fn handle_seal_boundary_query(&mut self, cmd: RequestDurableSegmentEnd) { let durable_end = self.durable_end(&cmd.segment_key); self.out .store_transport_cmd(DataTransportCommand::send_to_targets( vec![cmd.coordinator], - SealBoundaryReport { + DurableSegmentEndReported { segment_key: cmd.segment_key, from: self.node_id.clone(), durable_end, @@ -573,7 +592,7 @@ impl DataPlane { /// - already hold it through `sealed_end` → register if needed, transfer nothing; /// - already reconciling → no-op while progressing, resume if stalled; /// - else fetch `(local, sealed_end]` from a peer — a delta or a full copy. - fn handle_catch_up_assignment(&mut self, cmd: CatchUpAssignment) { + fn handle_catch_up_assignment(&mut self, cmd: AssignSegmentCatchUp) { if let Some(pending) = self.pending_catch_ups.get_mut(&cmd.segment_key) { // A receive is already in flight: no-op while it appends; once stalled // (no append across `CATCH_UP_IDLE_REDRIVES` re-drives) re-request the @@ -590,7 +609,7 @@ impl DataPlane { self.out .store_transport_cmd(DataTransportCommand::send_to_targets( vec![source], - CatchUpRequest { + RequestCatchUpEntries { segment_key: cmd.segment_key, from: self.node_id.clone(), local_end, @@ -646,7 +665,7 @@ impl DataPlane { self.out .store_transport_cmd(DataTransportCommand::send_to_targets( vec![source], - CatchUpRequest { + RequestCatchUpEntries { segment_key: cmd.segment_key, from: self.node_id.clone(), local_end: local, @@ -666,7 +685,7 @@ impl DataPlane { /// Highest durable (fsync'd) entry id this node holds for `key`, across every /// place a copy can live: a live tracker (the active segment a survivor is /// still following), the sealed store, or the recovered inventory. `None` if - /// we hold nothing. Reported in a `SealBoundaryReport` so the coordinator can seal + /// we hold nothing. Reported in a `DurableSegmentEndReported` so the coordinator can seal /// a leader-crashed segment at the `min` of survivors' durable extents. fn durable_end(&self, key: &SegmentKey) -> Option { let live = self @@ -711,7 +730,7 @@ impl DataPlane { /// Append a streamed batch to the in-progress receive, building sparse-index /// anchors as we go (same predicate the checkpoint/recovery paths use). /// A write failure aborts the receive — the source/coordinator re-drive. - fn handle_catch_up_chunk(&mut self, cmd: CatchUpChunk) { + fn handle_catch_up_chunk(&mut self, cmd: CatchUpEntries) { let Some(pending) = self.pending_catch_ups.get_mut(&cmd.segment_key) else { tracing::debug!( "catch-up chunk for an unknown receive {:?}; dropping", @@ -756,7 +775,7 @@ impl DataPlane { /// End of stream. Fsync, re-scan to verify it reaches `sealed_end`, then on /// success index the anchors, register the segment, and ack the coordinator. - fn handle_catch_up_stream_end(&mut self, cmd: CatchUpStreamEnd) { + fn handle_catch_up_stream_end(&mut self, cmd: CatchUpEntriesSent) { let Some(mut pending) = self.pending_catch_ups.remove(&cmd.segment_key) else { tracing::debug!( "catch-up done for an unknown receive {:?}; dropping", @@ -807,12 +826,12 @@ impl DataPlane { } /// Tell the coordinator we hold this segment, so its sweep stops re-announcing. - /// Routed via `SendToCoordinator`, mirroring `SegmentAssignmentAck`. + /// Routed via `SendToCoordinator`, mirroring `SegmentPlaced`. fn send_catch_up_ack(&mut self, shard_group_id: ShardGroupId, segment_key: SegmentKey) { self.out .store_transport_cmd(DataTransportSendToCoordinator { shard_group_id, - message: CatchUpAck { + message: SegmentCaughtUp { segment_key, shard_group_id, from: self.node_id.clone(), @@ -822,7 +841,7 @@ impl DataPlane { } /// Source side: a peer wants `cmd.segment_key` brought up to its local end. - fn handle_catch_up_request(&mut self, cmd: CatchUpRequest) { + fn handle_catch_up_request(&mut self, cmd: RequestCatchUpEntries) { let Some((start_offset, sealed_end)) = self.segments.sealed_bounds(&cmd.segment_key) else { tracing::warn!( "catch-up source has no sealed segment {:?}; dropping request", @@ -870,8 +889,8 @@ impl DataPlane { } /// The cold-read pool finished one batch for a catch-up source read. Emit it - /// as a `CatchUpChunk` to the requester, then re-arm the next read until the - /// sealed end is reached, at which point send `CatchUpStreamEnd`. + /// as a `CatchUpEntries` to the requester, then re-arm the next read until the + /// sealed end is reached, at which point send `CatchUpEntriesSent`. fn handle_catch_up_read_complete(&mut self, cmd: CatchUpReadComplete) { let has_entries = !cmd.entries.is_empty(); if has_entries { @@ -883,7 +902,7 @@ impl DataPlane { self.out .store_transport_cmd(DataTransportCommand::send_to_targets( vec![cmd.requester.clone()], - CatchUpChunk { + CatchUpEntries { segment_key: cmd.segment_key, entries, }, @@ -909,11 +928,11 @@ impl DataPlane { self.out .store_transport_cmd(DataTransportCommand::send_to_targets( vec![requester], - CatchUpStreamEnd { segment_key }, + CatchUpEntriesSent { segment_key }, )); } - fn handle_replica_ack(&mut self, cmd: ReplicaAck) { + fn handle_replica_ack(&mut self, cmd: ReplicaEntriesAppended) { let Some(committed) = self.replication.process_ack(&cmd.segment_key, &cmd.from) else { return; }; @@ -935,47 +954,152 @@ impl DataPlane { self.check_pending_seal(cmd.segment_key); } - fn handle_commit_consumer_offset(&mut self, cmd: CommitConsumerOffset) { + fn commit_consumer_offset(&mut self, cmd: CommitConsumerOffset) { self.needs_flush |= self .consumer_offsets - .handle_consumer_offset_commit(cmd, &self.node_id); + .commit_consumer_offset(cmd, &self.node_id); } - fn handle_replica_offset_commit(&mut self, cmd: ReplicaOffsetCommit) { - if !cmd.replica_set.contains(&self.node_id) { + + fn replicate_consumer_offset(&mut self, cmd: ReplicateConsumerOffset) { + if !self.is_follower_of(&cmd.replica_set) { return; } let Some(leader) = cmd.replica_set.leader().cloned() else { return; }; - if leader == self.node_id { - tracing::debug!("leader received replica offset commit message"); - return; - } let sealed_generation = self.consumer_offsets.latest_generation(&cmd.update.key); if cmd.update.generation < sealed_generation { let transport = DataTransportCommand::send_to_targets( vec![leader], - ReplicaOffsetAck { + ConsumerOffsetReplicated { seq: cmd.seq, update: cmd.update, from: self.node_id.clone(), - result: ReplicaOffsetAckResult::StaleEpoch(StaleEpoch(sealed_generation)), + result: ConsumerOffsetReplicationResult::StaleEpoch(StaleEpoch( + sealed_generation, + )), }, ); self.out.store_transport_cmd(transport); return; } - self.needs_flush |= self + if cmd.update.generation > sealed_generation { + let targets = cmd.replica_set.except_for(&self.node_id); + if !targets.is_empty() { + let transport = DataTransportCommand::send_to_targets( + targets, + RequestConsumerOffsetSnapshot { + segment_key: cmd.segment_key, + replicas: cmd.replica_set.clone(), + requester: self.node_id.clone(), + }, + ); + self.out.store_transport_cmd(transport) + } + + self.consumer_offsets + .future_entry(cmd.update.key.clone()) + .push(FutureOffsetCommit::Replica(cmd)); + return; + } + + self.consumer_offsets.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) { + self.ack_ready_offset_placements(); + } + } + + fn install_consumer_offset_snapshot(&mut self, from: &NodeId, cmd: ConsumerOffsetSnapshot) { + if !cmd.replica_set.contains(&self.node_id) || !cmd.replica_set.contains(from) { + return; + } + let Some(leader) = cmd.replica_set.leader() else { + return; + }; + if leader != from && leader != &self.node_id { + return; + }; + + let parked = self .consumer_offsets - .handle_replica_offset_commit(cmd, sealed_generation); + .install_consumer_offsets(cmd, &self.node_id); + self.resume_future_offset_commits(parked); + + self.needs_flush = true; } - fn handle_segment_assignment(&mut self, cmd: SegmentAssignment) { - self.consumer_offsets - .add_placement(&cmd.segment_key, cmd.replica_set.clone()); + fn request_consumer_offset_snapshot( + &mut self, + from: &NodeId, + cmd: RequestConsumerOffsetSnapshot, + ) { + if &cmd.requester != from || !cmd.replicas.contains(&self.node_id) { + return; + } + if !cmd.replicas.contains(&cmd.requester) { + return; + } + + if let Some(bootstrap) = self + .consumer_offsets + .create_offset_snapshot(cmd, &self.node_id) + { + self.out.store_transport_cmd(bootstrap); + } + } + + fn handle_consumer_offset_snapshot_installed(&mut self, cmd: ConsumerOffsetSnapshotInstalled) { + if self.node_id != cmd.leader { + return; + } + + self.consumer_offsets.handle_snapshot_installed_ack(&cmd); + + if let Some(bootstrap) = self + .consumer_offsets + .create_offset_snapshot_for_unready_replicas(cmd.segment_key) + { + self.out.store_transport_cmd(bootstrap); + } + self.ack_ready_offset_placements(); + } + + fn handle_commit_advance(&mut self, cmd: AdvanceReplicaCommit) { + let Some(tracker) = self.segments.get_mut(&cmd.segment_key) else { + return; + }; + + if tracker.role() != SegmentRole::Follower { + tracing::warn!( + "AdvanceReplicaCommit received by non-follower: {:?}", + cmd.segment_key + ); + return; + } + tracker.commit_entry(cmd.committed_entry_id); + } + + fn place_segment(&mut self, cmd: PlaceSegment) { + if !self.is_leader_of(&cmd.replica_set) { + return; + } + + if let Some(transport) = self.consumer_offsets.install_leader_placement( + cmd.segment_key, + cmd.shard_group_id, + cmd.replica_set.clone(), + ) { + self.out.store_transport_cmd(transport); + } if !self.segments.contains_key(&cmd.segment_key) { let tracker = SegmentTracker::new_with_start_entry_id( @@ -989,83 +1113,69 @@ impl DataPlane { self.segments.insert_active(cmd.segment_key, tracker); } - self.out - .store_transport_cmd(DataTransportSendToCoordinator { - shard_group_id: cmd.shard_group_id, - message: SegmentAssignmentAck { - segment_key: cmd.segment_key, - shard_group_id: cmd.shard_group_id, - from: self.node_id.clone(), - } - .into(), - }); + self.ack_ready_offset_placements(); + } + + /// Checks if there are any consumer offset placements that have reached full readiness (all replicas in the set + /// 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) { + self.out + .store_transport_cmd(DataTransportSendToCoordinator { + shard_group_id: ack.shard_group_id, + message: ack.into(), + }); + } } - fn process_replica_append(&mut self, cmd: ReplicaAppend) { - if cmd.replica_set.is_empty() { + fn replicate_segment_entries(&mut self, from: &NodeId, cmd: ReplicateSegmentEntries) { + if cmd.replicas.leader() != Some(from) { + return; + } + if !self.is_follower_of(&cmd.replicas) { + tracing::warn!("{} is not follower for {:?}", self.node_id, cmd.segment_key); return; } - let replica_set = cmd.replica_set.clone(); - if !self.segments.contains_key(&cmd.segment_key) { - if !cmd.replica_set.contains(&self.node_id) { - tracing::warn!( - "ReplicaAppend for segment I'm not in: {:?}", - cmd.segment_key - ); - return; - } - if cmd.replica_set.first() == Some(&self.node_id) { - tracing::warn!("ReplicaAppend from self as leader: {:?}", cmd.segment_key); - return; - } + if matches!( + self.consumer_offsets.observe_placement( + cmd.segment_key, + ShardGroupId(0), + &cmd.replicas, + &self.node_id + ), + PlacementObservation::Stale | PlacementObservation::Conflict + ) { + return; + }; + // If it sees this segment for the first time + if !self.segments.contains_key(&cmd.segment_key) { self.segments.insert_active( cmd.segment_key, SegmentTracker::new_with_start_entry_id( cmd.segment_key .file_path(&self.config.data_dir, cmd.entry_id), SegmentRole::Follower, - cmd.replica_set, + cmd.replicas, ShardGroupId(0), cmd.entry_id, ), ); } - - self.consumer_offsets - .add_placement(&cmd.segment_key, replica_set); - let Some(tracker) = self.segments.get_mut(&cmd.segment_key) else { return; }; - tracker.stage_entry_from_replica(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); self.dirty_segments.insert(cmd.segment_key); self.needs_flush = true; } - fn handle_commit_advance(&mut self, cmd: CommitAdvance) { - let Some(tracker) = self.segments.get_mut(&cmd.segment_key) else { - return; - }; - - if tracker.role() != SegmentRole::Follower { - tracing::warn!( - "CommitAdvance received by non-follower: {:?}", - cmd.segment_key - ); - return; - } - tracker.commit_entry(cmd.committed_entry_id); - } - - fn handle_seal_response(&mut self, cmd: SealResponse) { - if cmd.new_replica_set.is_empty() { - return; - } - + // TODO refactor + fn handle_segment_roll_committed(&mut self, cmd: SegmentRollCommitted) { self.pending_seal_requests.clear(&cmd.old_segment_key); let Some(old_tracker) = self.segments.get(&cmd.old_segment_key) else { return; @@ -1093,9 +1203,6 @@ impl DataPlane { )); } - self.consumer_offsets - .add_placement(&cmd.old_segment_key, cmd.new_replica_set.clone()); - self.replication.segment_handoff( cmd.old_segment_key, cmd.old_segment_key.with_segment_id(cmd.new_segment_id), @@ -1103,7 +1210,7 @@ impl DataPlane { self.retire_old_segment(cmd.old_segment_key); // The successor may already exist: on a roll the coordinator co-dispatches a - // `SegmentAssignment` for the new active segment alongside this `SealResponse`, + // `PlaceSegment` for the new active segment alongside this committed result, // and it can win the race (creating the successor empty). Carry the tail onto // the existing tracker rather than a fresh one that `insert_active` would // silently refuse — which would drop the records (data loss) and strand the @@ -1133,10 +1240,12 @@ impl DataPlane { .get_mut(&cmd.old_segment_key.with_segment_id(cmd.new_segment_id)) { for (data, record_count) 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, ); } } @@ -1165,7 +1274,7 @@ impl DataPlane { self.out .store_transport_cmd(DataTransportSendToCoordinator { shard_group_id: tracker.shard_group_id(), - message: SealRequest { + message: RequestSegmentRoll { from: self.node_id.clone(), segment_key: cmd.segment_key, failed_nodes: vec![], @@ -1214,7 +1323,7 @@ impl DataPlane { self.out .store_transport_cmd(DataTransportSendToCoordinator { shard_group_id: tracker.shard_group_id(), - message: SealRequest { + message: RequestSegmentRoll { from: self.node_id.clone(), segment_key, failed_nodes: vec![], @@ -1254,7 +1363,7 @@ impl DataPlane { self.out .store_transport_cmd(DataTransportCommand::send_to_targets( followers, - CommitAdvance { + AdvanceReplicaCommit { segment_key, committed_entry_id: entry_id, }, @@ -1297,7 +1406,7 @@ impl DataPlane { self.out .store_transport_cmd(DataTransportSendToCoordinator { shard_group_id: tracker.shard_group_id(), - message: SealRequest { + message: RequestSegmentRoll { from: self.node_id.clone(), segment_key, failed_nodes: failed_nodes.clone(), @@ -1322,7 +1431,7 @@ impl DataPlane { self.out .store_transport_cmd(DataTransportSendToCoordinator { shard_group_id: tracker.shard_group_id(), - message: SealRequest { + message: RequestSegmentRoll { from: self.node_id.clone(), segment_key, failed_nodes: vec![], @@ -1411,7 +1520,7 @@ impl DataPlane { self.out .store_transport_cmd(DataTransportCommand::send_to_targets( vec![tracker.leader_node()], - ReplicaAck { + ReplicaEntriesAppended { segment_key: key, entry_id: entry.entry_id, from: self.node_id.clone(), @@ -1523,7 +1632,7 @@ impl DataPlane { self.out .store_transport_cmd(DataTransportSendToCoordinator { shard_group_id: tracker.shard_group_id(), - message: SealRequest { + message: RequestSegmentRoll { from: self.node_id.clone(), segment_key, failed_nodes, @@ -1533,6 +1642,14 @@ impl DataPlane { }); } } + + fn is_follower_of(&self, replicas: &Replicas) -> bool { + replicas.contains(&self.node_id) && replicas.leader() != Some(&self.node_id) + } + + fn is_leader_of(&self, replicas: &Replicas) -> bool { + replicas.contains(&self.node_id) && replicas.leader() == Some(&self.node_id) + } } #[cfg(any(test, debug_assertions))] @@ -1593,12 +1710,63 @@ mod tests { fn has_buffered_data(&self) -> bool { self.buffer_byte_count > 0 } + + fn process_peer(&mut self, message: impl Into) { + self.handle_command(receive_peer_message(message.into())); + } } fn test_key() -> SegmentKey { SegmentKey::new(TopicId(1), RangeId(0), SegmentId(0)) } + fn receive_peer_message(message: DataPlanePeerMessage) -> DataPlaneCommand { + let from = match &message { + DataPlanePeerMessage::ReplicateSegmentEntries(command) => { + command.replicas.leader().cloned().unwrap() + } + DataPlanePeerMessage::SegmentPlaced(event) => event.from.clone(), + DataPlanePeerMessage::ReplicaEntriesAppended(event) => event.from.clone(), + DataPlanePeerMessage::ReplicateConsumerOffset(command) => { + command.replica_set.leader().cloned().unwrap() + } + DataPlanePeerMessage::ConsumerOffsetReplicated(event) => event.from.clone(), + DataPlanePeerMessage::InstallConsumerOffsetSnapshot(command) => { + command.replica_set.leader().cloned().unwrap() + } + DataPlanePeerMessage::RequestConsumerOffsetSnapshot(request) => { + request.requester.clone() + } + DataPlanePeerMessage::ConsumerOffsetSnapshotInstalled(event) => event.from.clone(), + DataPlanePeerMessage::RequestSegmentRoll(request) => request.from.clone(), + DataPlanePeerMessage::RequestCatchUpEntries(request) => request.from.clone(), + DataPlanePeerMessage::SegmentCaughtUp(event) => event.from.clone(), + DataPlanePeerMessage::RequestDurableSegmentEnd(request) => request.coordinator.clone(), + DataPlanePeerMessage::DurableSegmentEndReported(event) => event.from.clone(), + DataPlanePeerMessage::PlaceSegment(_) + | DataPlanePeerMessage::AdvanceReplicaCommit(_) + | DataPlanePeerMessage::SegmentRollCommitted(_) + | DataPlanePeerMessage::SegmentSealed(_) + | DataPlanePeerMessage::AssignSegmentCatchUp(_) + | DataPlanePeerMessage::CatchUpEntries(_) + | DataPlanePeerMessage::CatchUpEntriesSent(_) + | DataPlanePeerMessage::DeleteSegments(_) + | DataPlanePeerMessage::ConsumerGroupEpochSealed(_) => NodeId::new("test-peer"), + }; + receive_peer_message_from(from, message) + } + + fn receive_peer_message_from( + from: NodeId, + message: impl Into, + ) -> DataPlaneCommand { + ReceivePeerMessage { + from, + message: Box::new(message.into()), + } + .into() + } + fn consumer_offset_key() -> ConsumerOffsetKey { ConsumerOffsetKey { topic_id: TopicId(1), @@ -1631,16 +1799,14 @@ mod tests { Err(tokio::sync::oneshot::error::TryRecvError::Empty) )); - dp.process(DataPlaneInterNodeCommand::ConsumerGroupEpochSeal( - EpochSeal { - generation: 2.into(), - key: ConsumerOffsetKey { - topic_id: TopicId(1), - range_id: RangeId(0), - group_id: "group".into(), - }, + dp.process_peer(DataPlanePeerMessage::ConsumerGroupEpochSealed(EpochSeal { + generation: 2.into(), + key: ConsumerOffsetKey { + topic_id: TopicId(1), + range_id: RangeId(0), + group_id: "group".into(), }, - )); + })); assert!(matches!( response.try_recv(), Err(tokio::sync::oneshot::error::TryRecvError::Empty) @@ -1656,7 +1822,7 @@ mod tests { ); assert_eq!( dp.consumer_offsets - .uncheckpointed_offset_lsns + .uncheckpointed_offset_lsns() .iter() .copied() .collect::>(), @@ -1711,8 +1877,8 @@ mod tests { dp.handle_command(OffsetCheckpointComplete { checkpointed_lsn: 2, }); - assert!(dp.consumer_offsets.uncheckpointed_offset_lsns.is_empty()); - assert_eq!(dp.consumer_offsets.offset_checkpoint_lsn, 2); + assert!(dp.consumer_offsets.uncheckpointed_offset_lsns().is_empty()); + assert_eq!(dp.consumer_offsets.offset_checkpoint_lsn(), 2); assert_eq!(dp.wal.reclaimable_lsn(), None); } @@ -1721,16 +1887,14 @@ mod tests { let dir = tempfile::tempdir().unwrap(); let mut dp = make_data_plane(&dir); dp.process(assign_segment(test_key(), vec![test_node_id()])); - dp.process(DataPlaneInterNodeCommand::ConsumerGroupEpochSeal( - EpochSeal { - generation: 4.into(), - key: ConsumerOffsetKey { - topic_id: TopicId(1), - range_id: RangeId(0), - group_id: "group".into(), - }, + dp.process_peer(DataPlanePeerMessage::ConsumerGroupEpochSealed(EpochSeal { + generation: 4.into(), + key: ConsumerOffsetKey { + topic_id: TopicId(1), + range_id: RangeId(0), + group_id: "group".into(), }, - )); + })); let (reply, mut response) = oneshot::channel(); dp.process(CommitConsumerOffset { update: ConsumerOffsetUpdate { @@ -1755,7 +1919,7 @@ mod tests { let dir = tempfile::tempdir().unwrap(); let mut dp = make_data_plane(&dir); dp.process(assign_segment(test_key(), vec![test_node_id()])); - dp.process(DataPlaneInterNodeCommand::SegmentSealed(SegmentSealed { + dp.process_peer(DataPlanePeerMessage::SegmentSealed(SegmentSealed { segment_key: test_key(), committed_entry_id: None, })); @@ -1797,12 +1961,19 @@ mod tests { test_key(), vec![test_node_id(), follower_1.clone(), follower_2.clone()], )); - dp.process(DataPlaneInterNodeCommand::ConsumerGroupEpochSeal( - EpochSeal { - generation: 1.into(), - key: consumer_offset_key(), - }, - )); + for follower in [&follower_1, &follower_2] { + dp.process_peer(DataPlanePeerMessage::ConsumerOffsetSnapshotInstalled( + ConsumerOffsetSnapshotInstalled { + segment_key: test_key(), + from: follower.clone(), + leader: test_node_id(), + }, + )); + } + dp.process_peer(DataPlanePeerMessage::ConsumerGroupEpochSealed(EpochSeal { + generation: 1.into(), + key: consumer_offset_key(), + })); dp.flush_batch(); dp.out.transport_cmds.clear(); @@ -1832,25 +2003,28 @@ mod tests { .iter() .find_map(|command| match command { DataTransportCommand::SendToTargets(send) => match &send.message { - DataPlaneInterNodeCommand::ReplicaOffsetCommit(commit) => Some(commit.clone()), - DataPlaneInterNodeCommand::SegmentAssignment(_) - | DataPlaneInterNodeCommand::SegmentAssignmentAck(_) - | DataPlaneInterNodeCommand::ReplicaAppend(_) - | DataPlaneInterNodeCommand::ReplicaAck(_) - | DataPlaneInterNodeCommand::ReplicaOffsetAck(_) - | DataPlaneInterNodeCommand::CommitAdvance(_) - | DataPlaneInterNodeCommand::SealRequest(_) - | DataPlaneInterNodeCommand::SealResponse(_) - | DataPlaneInterNodeCommand::SegmentSealed(_) - | DataPlaneInterNodeCommand::CatchUpAssignment(_) - | DataPlaneInterNodeCommand::CatchUpRequest(_) - | DataPlaneInterNodeCommand::CatchUpChunk(_) - | DataPlaneInterNodeCommand::CatchUpStreamEnd(_) - | DataPlaneInterNodeCommand::CatchUpAck(_) - | DataPlaneInterNodeCommand::SealBoundaryQuery(_) - | DataPlaneInterNodeCommand::SealBoundaryReport(_) - | DataPlaneInterNodeCommand::DeleteSegments(_) - | DataPlaneInterNodeCommand::ConsumerGroupEpochSeal(_) => None, + DataPlanePeerMessage::ReplicateConsumerOffset(commit) => Some(commit.clone()), + DataPlanePeerMessage::PlaceSegment(_) + | DataPlanePeerMessage::SegmentPlaced(_) + | DataPlanePeerMessage::ReplicateSegmentEntries(_) + | DataPlanePeerMessage::ReplicaEntriesAppended(_) + | DataPlanePeerMessage::ConsumerOffsetReplicated(_) + | DataPlanePeerMessage::InstallConsumerOffsetSnapshot(_) + | DataPlanePeerMessage::RequestConsumerOffsetSnapshot(_) + | DataPlanePeerMessage::ConsumerOffsetSnapshotInstalled(_) + | DataPlanePeerMessage::AdvanceReplicaCommit(_) + | DataPlanePeerMessage::RequestSegmentRoll(_) + | DataPlanePeerMessage::SegmentRollCommitted(_) + | DataPlanePeerMessage::SegmentSealed(_) + | DataPlanePeerMessage::AssignSegmentCatchUp(_) + | DataPlanePeerMessage::RequestCatchUpEntries(_) + | DataPlanePeerMessage::CatchUpEntries(_) + | DataPlanePeerMessage::CatchUpEntriesSent(_) + | DataPlanePeerMessage::SegmentCaughtUp(_) + | DataPlanePeerMessage::RequestDurableSegmentEnd(_) + | DataPlanePeerMessage::DurableSegmentEndReported(_) + | DataPlanePeerMessage::DeleteSegments(_) + | DataPlanePeerMessage::ConsumerGroupEpochSealed(_) => None, }, DataTransportCommand::SendToCoordinator(_) | DataTransportCommand::DisconnectPeer(_) => None, @@ -1858,12 +2032,12 @@ mod tests { .expect("leader must replicate the offset commit"); for follower in [follower_1, follower_2] { - dp.process(DataPlaneInterNodeCommand::ReplicaOffsetAck( - ReplicaOffsetAck { + dp.process_peer(DataPlanePeerMessage::ConsumerOffsetReplicated( + ConsumerOffsetReplicated { seq: replicated.seq, update: replicated.update.clone(), from: follower, - result: ReplicaOffsetAckResult::Committed, + result: ConsumerOffsetReplicationResult::Committed, }, )); } @@ -1883,9 +2057,10 @@ mod tests { batch_offset: 1, absolute_offset: 6, }; - dp.process(DataPlaneInterNodeCommand::ReplicaOffsetCommit( - ReplicaOffsetCommit { + dp.process_peer(DataPlanePeerMessage::ReplicateConsumerOffset( + ReplicateConsumerOffset { seq: 9, + segment_key: test_key(), replica_set: Replicas::new(vec![leader, test_node_id()]), update: ConsumerOffsetUpdate { key: consumer_offset_key(), @@ -1906,15 +2081,123 @@ mod tests { DataTransportCommand::SendToTargets(send) if matches!( &send.message, - DataPlaneInterNodeCommand::ReplicaOffsetAck(ReplicaOffsetAck { + DataPlanePeerMessage::ConsumerOffsetReplicated(ConsumerOffsetReplicated { seq: 9, - result: ReplicaOffsetAckResult::Committed, + result: ConsumerOffsetReplicationResult::Committed, .. }) ) ))); } + #[test] + fn consumer_offset_bootstrap_acks_only_after_wal_flush() { + let dir = tempfile::tempdir().unwrap(); + let mut dp = make_data_plane(&dir); + let leader = NodeId::new("leader"); + let position = ConsumerOffsetPosition { + entry_id: EntryId(8), + batch_offset: 2, + absolute_offset: 11, + }; + + dp.process_peer(DataPlanePeerMessage::InstallConsumerOffsetSnapshot( + 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), + }, + ] + .into_boxed_slice(), + }, + )); + + assert!(dp.out.transport_cmds.iter().all(|cmd| !matches!( + cmd, + DataTransportCommand::SendToTargets(send) + if matches!(send.message, DataPlanePeerMessage::ConsumerOffsetSnapshotInstalled(_)) + ))); + + dp.flush_batch(); + + assert_eq!( + dp.consumer_offsets.offset(&consumer_offset_key()), + Some(position) + ); + assert!(dp.out.transport_cmds.iter().any(|cmd| matches!( + cmd, + DataTransportCommand::SendToTargets(send) + if send.targets.as_ref() == [leader.clone()] + && matches!(send.message, DataPlanePeerMessage::ConsumerOffsetSnapshotInstalled(_)) + ))); + } + + #[test] + fn consumer_offset_bootstrap_releases_a_parked_newer_epoch_commit() { + let dir = tempfile::tempdir().unwrap(); + let mut dp = make_data_plane(&dir); + let leader = NodeId::new("leader"); + let position = ConsumerOffsetPosition { + entry_id: EntryId(9), + batch_offset: 0, + absolute_offset: 12, + }; + let update = ConsumerOffsetUpdate { + key: consumer_offset_key(), + generation: 2.into(), + position, + }; + dp.process_peer(DataPlanePeerMessage::ReplicateConsumerOffset( + ReplicateConsumerOffset { + seq: 10, + segment_key: test_key(), + replica_set: Replicas::new(vec![leader.clone(), test_node_id()]), + update: update.clone(), + }, + )); + assert!(!dp.consumer_offsets.has_pending_offsets()); + assert!(dp.out.transport_cmds.iter().any(|cmd| matches!( + cmd, + DataTransportCommand::SendToTargets(send) + if send.targets.as_ref() == [leader.clone()] + && matches!(send.message, DataPlanePeerMessage::RequestConsumerOffsetSnapshot(_)) + ))); + + dp.process_peer(DataPlanePeerMessage::InstallConsumerOffsetSnapshot( + 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, + }, + ] + .into_boxed_slice(), + }, + )); + dp.flush_batch(); + + assert_eq!( + dp.consumer_offsets.offset(&consumer_offset_key()), + Some(position) + ); + assert!(dp.out.transport_cmds.iter().any(|cmd| matches!( + cmd, + DataTransportCommand::SendToTargets(send) + if matches!( + &send.message, + DataPlanePeerMessage::ConsumerOffsetReplicated(ack) + if ack.seq == 10 && ack.update == update + ) + ))); + } + fn test_node_id() -> NodeId { NodeId::new("test-node") } @@ -2005,11 +2288,11 @@ mod tests { } fn assign_segment(key: SegmentKey, replica_set: Vec) -> DataPlaneCommand { - DataPlaneCommand::DataPlaneInterNodeCommand( - SegmentAssignment { + receive_peer_message( + PlaceSegment { segment_key: key, shard_group_id: ShardGroupId(1), - replica_set, + replica_set: Replicas::new(replica_set), start_entry_id: EntryId(0), } .into(), @@ -2032,11 +2315,11 @@ mod tests { new_segment_id: SegmentId, new_replica_set: Vec, ) -> DataPlaneCommand { - DataPlaneCommand::DataPlaneInterNodeCommand( - SealResponse { + receive_peer_message( + SegmentRollCommitted { old_segment_key: old, new_segment_id, - new_replica_set, + new_replica_set: Replicas::new(new_replica_set), } .into(), ) @@ -2093,7 +2376,7 @@ mod tests { let dir = tempfile::tempdir().unwrap(); let mut dp = make_data_plane(&dir); - dp.handle_command(assign_segment(test_key(), vec![])); + dp.handle_command(assign_segment(test_key(), vec![test_node_id()])); let (tx, _) = oneshot::channel(); dp.handle_command(Produce { @@ -2112,7 +2395,7 @@ mod tests { fn segment_assignment_creates_tracker() { let dir = tempfile::tempdir().unwrap(); let mut dp = make_data_plane(&dir); - dp.handle_command(assign_segment(test_key(), vec![])); + dp.handle_command(assign_segment(test_key(), vec![test_node_id()])); assert!(dp.segments.contains_key(&test_key())); } @@ -2121,7 +2404,7 @@ mod tests { let dir = tempfile::tempdir().unwrap(); let mut dp = make_data_plane(&dir); - dp.handle_command(assign_segment(test_key(), vec![])); + dp.handle_command(assign_segment(test_key(), vec![test_node_id()])); let (cmd, _) = produce(test_key()); dp.handle_command(cmd); @@ -2146,7 +2429,7 @@ mod tests { let dir = tempfile::tempdir().unwrap(); let mut dp = make_data_plane(&dir); - dp.handle_command(assign_segment(test_key(), vec![])); + dp.handle_command(assign_segment(test_key(), vec![test_node_id()])); let (cmd, _) = produce(test_key()); dp.handle_command(cmd); @@ -2164,7 +2447,7 @@ mod tests { let dir = tempfile::tempdir().unwrap(); let mut dp = make_data_plane(&dir); - dp.handle_command(assign_segment(test_key(), vec![])); + dp.handle_command(assign_segment(test_key(), vec![test_node_id()])); let (cmd, _) = produce(test_key()); dp.handle_command(cmd); @@ -2182,7 +2465,7 @@ mod tests { let dir = tempfile::tempdir().unwrap(); let mut dp = make_data_plane(&dir); - dp.handle_command(assign_segment(test_key(), vec![])); + dp.handle_command(assign_segment(test_key(), vec![dp.node_id.clone()])); let (cmd, _) = produce(test_key()); dp.handle_command(cmd); @@ -2198,7 +2481,7 @@ mod tests { fn volume_trigger_flushes_immediately() { let dir = tempfile::tempdir().unwrap(); let mut dp = make_data_plane(&dir); - dp.handle_command(assign_segment(test_key(), vec![])); + dp.handle_command(assign_segment(test_key(), vec![dp.node_id.clone()])); dp.handle_command(Produce { segment_key: test_key(), @@ -2216,7 +2499,7 @@ mod tests { fn volume_trigger_flush_leaves_stale_timer_harmless() { let dir = tempfile::tempdir().unwrap(); let mut dp = make_data_plane(&dir); - dp.handle_command(assign_segment(test_key(), vec![])); + dp.handle_command(assign_segment(test_key(), vec![dp.node_id.clone()])); dp.handle_command(Produce { segment_key: test_key(), @@ -2237,7 +2520,7 @@ mod tests { fn cache_pressure_checkpoint_disabled_when_budget_zero() { let dir = tempfile::tempdir().unwrap(); let mut dp = make_data_plane(&dir); - dp.handle_command(assign_segment(test_key(), vec![])); + dp.handle_command(assign_segment(test_key(), vec![dp.node_id.clone()])); dp.handle_command(Produce { segment_key: test_key(), @@ -2262,8 +2545,10 @@ mod tests { let small = SegmentKey::new(TopicId(1), RangeId(0), SegmentId(0)); let large = SegmentKey::new(TopicId(2), RangeId(0), SegmentId(0)); - dp.handle_command(assign_segment(small, vec![])); - dp.handle_command(assign_segment(large, vec![])); + let replica_set = vec![NodeId::new(dp.node_id.to_string())]; + + dp.handle_command(assign_segment(small, replica_set.clone())); + dp.handle_command(assign_segment(large, replica_set)); dp.handle_command(Produce { segment_key: small, @@ -2320,7 +2605,10 @@ mod tests { let mut dp = make_data_plane(&dir); dp.config.hot_cache_budget_bytes = 10; dp.config.hot_cache_pressure_watermark = 0.5; - dp.handle_command(assign_segment(test_key(), vec![])); + dp.handle_command(assign_segment( + test_key(), + vec![NodeId::new(dp.node_id.to_string())], + )); dp.handle_command(Produce { segment_key: test_key(), @@ -2377,8 +2665,8 @@ mod tests { assert!(dp.has_buffered_data()); } - /// On a roll the coordinator co-dispatches a `SegmentAssignment` for the new - /// active segment alongside the `SealResponse`. If the assignment wins the race + /// On a roll the coordinator co-dispatches a `PlaceSegment` for the new + /// active segment alongside `SegmentRollCommitted`. If the assignment wins the race /// (creating the successor empty first), the seal handoff must still carry the /// old segment's staged tail onto that existing successor — not build a fresh /// tracker that `insert_active` refuses, dropping the records (data loss) and @@ -2394,11 +2682,11 @@ mod tests { dp.handle_command(cmd); assert!(dp.has_buffered_data()); - // The successor's SegmentAssignment wins the race: seg1 is created empty. + // The successor's PlaceSegment wins the race: seg1 is created empty. let successor = test_key().with_segment_id(SegmentId(1)); dp.handle_command(assign_segment(successor, vec![test_node_id()])); - // The SealResponse handoff arrives second. Pre-fix `insert_active` refuses + // The committed roll handoff arrives second. Pre-fix `insert_active` refuses // (seg1 already active) → the staged produce is dropped and `buffer_byte_count` // strands, tripping `assert_invariants` after the command. dp.handle_command(seal_response( @@ -2419,8 +2707,8 @@ mod tests { let key1 = SegmentKey::new(TopicId(1), RangeId(0), SegmentId(0)); let key2 = SegmentKey::new(TopicId(2), RangeId(0), SegmentId(0)); - dp.handle_command(assign_segment(key1, vec![])); - dp.handle_command(assign_segment(key2, vec![])); + dp.handle_command(assign_segment(key1, vec![test_node_id()])); + dp.handle_command(assign_segment(key2, vec![test_node_id()])); let wal_file_count = || -> usize { std::fs::read_dir(dir.path().join("wal")).unwrap().count() }; @@ -2451,8 +2739,8 @@ mod tests { let dir = tempfile::tempdir().unwrap(); let mut dp = make_data_plane(&dir); - dp.handle_command(assign_segment(test_key(), vec![])); - dp.handle_command(assign_segment(test_key(), vec![])); + dp.handle_command(assign_segment(test_key(), vec![test_node_id()])); + dp.handle_command(assign_segment(test_key(), vec![test_node_id()])); assert_eq!(dp.segments.len(), 1); } @@ -2462,7 +2750,7 @@ mod tests { let dir = tempfile::tempdir().unwrap(); let mut dp = make_data_plane(&dir); - dp.handle_command(assign_segment(test_key(), vec![])); + dp.handle_command(assign_segment(test_key(), vec![test_node_id()])); let (cmd1, _rx1) = produce(test_key()); let (cmd2, _rx2) = produce(test_key()); @@ -2519,7 +2807,7 @@ mod tests { let dir = tempfile::tempdir().unwrap(); let mut dp = make_data_plane(&dir); - dp.handle_command(assign_segment(test_key(), vec![])); + dp.handle_command(assign_segment(test_key(), vec![test_node_id()])); let (cmd, _rx) = produce(test_key()); dp.handle_command(cmd); @@ -2592,10 +2880,10 @@ mod tests { let mut dp = make_data_plane(&dir); let leader = NodeId::new("leader-node"); - dp.handle_command(DataPlaneCommand::DataPlaneInterNodeCommand( - ReplicaAppend { + dp.handle_command(receive_peer_message( + ReplicateSegmentEntries { segment_key: test_key(), - replica_set: vec![leader, test_node_id()], + replicas: Replicas::new(vec![leader, test_node_id()]), data: b"setup".to_vec().into(), record_count: 1, entry_id: EntryId(0), @@ -2617,10 +2905,10 @@ mod tests { let mut dp = make_data_plane(&dir); let leader = NodeId::new("leader-node"); - dp.handle_command(DataPlaneCommand::DataPlaneInterNodeCommand( - ReplicaAppend { + dp.handle_command(receive_peer_message( + ReplicateSegmentEntries { segment_key: test_key(), - replica_set: vec![leader.clone(), test_node_id()], + replicas: Replicas::new(vec![leader.clone(), test_node_id()]), data: b"data".to_vec().into(), record_count: 1, entry_id: EntryId(0), @@ -2635,15 +2923,39 @@ mod tests { ); } + #[test] + fn segment_replication_rejects_a_non_leader_transport_peer() { + let dir = tempfile::tempdir().unwrap(); + let mut dp = make_data_plane(&dir); + let leader = NodeId::new("leader-node"); + + dp.handle_command(receive_peer_message_from( + NodeId::new("other-peer"), + ReplicateSegmentEntries { + segment_key: test_key(), + replicas: Replicas::new(vec![leader, test_node_id()]), + data: b"data".to_vec().into(), + record_count: 1, + entry_id: EntryId(0), + }, + )); + + assert!(!dp.segments.contains_key(&test_key())); + assert!(!dp.needs_flush); + } + #[test] fn replica_append_rejected_if_not_in_replica_set() { let dir = tempfile::tempdir().unwrap(); let mut dp = make_data_plane(&dir); - dp.handle_command(DataPlaneCommand::DataPlaneInterNodeCommand( - ReplicaAppend { + dp.handle_command(receive_peer_message( + ReplicateSegmentEntries { segment_key: test_key(), - replica_set: vec![NodeId::new("other-leader"), NodeId::new("other-follower")], + replicas: Replicas::new(vec![ + NodeId::new("other-leader"), + NodeId::new("other-follower"), + ]), data: b"data".to_vec().into(), record_count: 1, entry_id: EntryId(0), @@ -2660,10 +2972,10 @@ mod tests { let mut dp = make_data_plane(&dir); let leader = NodeId::new("leader-node"); - dp.handle_command(DataPlaneCommand::DataPlaneInterNodeCommand( - ReplicaAppend { + dp.handle_command(receive_peer_message( + ReplicateSegmentEntries { segment_key: test_key(), - replica_set: vec![leader.clone(), test_node_id()], + replicas: Replicas::new(vec![leader.clone(), test_node_id()]), data: b"data".to_vec().into(), record_count: 1, entry_id: EntryId(0), @@ -2672,8 +2984,8 @@ mod tests { )); dp.flush_batch(); - dp.handle_command(DataPlaneCommand::DataPlaneInterNodeCommand( - CommitAdvance { + dp.handle_command(receive_peer_message( + AdvanceReplicaCommit { segment_key: test_key(), committed_entry_id: EntryId(0), } @@ -2779,10 +3091,10 @@ mod tests { assert!(!dp.needs_flush); - dp.handle_command(DataPlaneCommand::DataPlaneInterNodeCommand( - ReplicaAppend { + dp.handle_command(receive_peer_message( + ReplicateSegmentEntries { segment_key: test_key(), - replica_set: vec![leader, test_node_id()], + replicas: Replicas::new(vec![leader, test_node_id()]), data: b"data".to_vec().into(), record_count: 1, entry_id: EntryId(0), @@ -2794,7 +3106,7 @@ mod tests { } #[test] - fn seal_response_no_uncommitted_skips_wal_write() { + fn segment_roll_committed_without_uncommitted_data_skips_wal_write() { let dir = tempfile::tempdir().unwrap(); let mut dp = make_data_plane(&dir); let follower = NodeId::new("follower"); @@ -2812,11 +3124,11 @@ mod tests { .unwrap() .commit_entry(EntryId(0)); - dp.handle_command(DataPlaneCommand::DataPlaneInterNodeCommand( - SealResponse { + dp.handle_command(receive_peer_message( + SegmentRollCommitted { old_segment_key: test_key(), new_segment_id: SegmentId(1), - new_replica_set: vec![test_node_id(), NodeId::new("new-follower")], + new_replica_set: Replicas::new(vec![test_node_id(), NodeId::new("new-follower")]), } .into(), )); @@ -2857,8 +3169,8 @@ mod tests { dp.enqueue_seal_request(test_key()); dp.enqueue_seal_request(test_key()); - // Count only seal requests — `handle_segment_assignment` also emits a - // SegmentAssignmentAck via SendToCoordinator, which is not a seal. + // Count only seal requests — `handle_place_segment` also emits a + // SegmentPlaced via SendToCoordinator, which is not a seal. assert_eq!( dp.out .transport_cmds @@ -2866,7 +3178,7 @@ mod tests { .filter(|c| matches!( c, DataTransportCommand::SendToCoordinator(s) - if matches!(s.message, DataPlaneInterNodeCommand::SealRequest(_)) + if matches!(s.message, DataPlanePeerMessage::RequestSegmentRoll(_)) )) .count(), 1 @@ -2879,10 +3191,10 @@ mod tests { let mut dp = make_data_plane(&dir); let leader = NodeId::new("leader-node"); - dp.handle_command(DataPlaneCommand::DataPlaneInterNodeCommand( - ReplicaAppend { + dp.handle_command(receive_peer_message( + ReplicateSegmentEntries { segment_key: test_key(), - replica_set: vec![leader, test_node_id()], + replicas: Replicas::new(vec![leader, test_node_id()]), data: b"data".to_vec().into(), record_count: 1, entry_id: EntryId(0), @@ -2909,7 +3221,7 @@ mod tests { // ── D4 integration tests ────────────────────────────────────────── // // These exercise the resolver end-to-end via `DataPlane` commands — - // SegmentAssignment routes through `handle_segment_assignment`, seal + // PlaceSegment routes through `handle_place_segment`, seal // routes through `handle_seal_response`, etc. The unit-level tests for // the index itself live in `states/segment_store.rs`. @@ -2919,7 +3231,7 @@ mod tests { fn segment_assignment_indexes_for_resolver() { let dir = tempfile::tempdir().unwrap(); let mut dp = make_data_plane(&dir); - dp.handle_command(assign_segment(test_key(), vec![])); + dp.handle_command(assign_segment(test_key(), vec![test_node_id()])); let resolved = dp .segments @@ -2928,19 +3240,19 @@ mod tests { assert!(matches!(resolved, SegmentReadState::Active(k) if k == test_key())); } - /// Verifies the leader-side seal path: after `SealResponse`, the old + /// Verifies the leader-side seal path: after `SegmentRollCommitted`, the old /// segment migrates to the sealed table with `end_entry_id` set from /// the tracker's `committed_entry_id`, and the new segment becomes the /// active index entry for the range. #[test] - fn seal_response_routes_old_offsets_through_sealed_table() { + fn segment_roll_committed_routes_old_offsets_through_sealed_table() { let dir = tempfile::tempdir().unwrap(); let mut dp = make_data_plane(&dir); // Single-replica produce auto-commits on flush, so entry 0 is // committed at the time we issue the seal. After seal: end=0 on the // sealed entry, new segment starts at offset 1. - dp.handle_command(assign_segment(test_key(), vec![])); + dp.handle_command(assign_segment(test_key(), vec![test_node_id()])); let (cmd, _rx) = produce(test_key()); dp.handle_command(cmd); process_and_flush( @@ -2950,11 +3262,11 @@ mod tests { ), ); - dp.handle_command(DataPlaneCommand::DataPlaneInterNodeCommand( - SealResponse { + dp.handle_command(receive_peer_message( + SegmentRollCommitted { old_segment_key: test_key(), new_segment_id: SegmentId(1), - new_replica_set: vec![test_node_id()], + new_replica_set: Replicas::new(vec![test_node_id()]), } .into(), )); @@ -3014,11 +3326,11 @@ mod tests { ); // Replication-timeout failover seals the never-committed segment. - dp.handle_command(DataPlaneCommand::DataPlaneInterNodeCommand( - SealResponse { + dp.handle_command(receive_peer_message( + SegmentRollCommitted { old_segment_key: test_key(), new_segment_id: SegmentId(1), - new_replica_set: vec![test_node_id()], + new_replica_set: Replicas::new(vec![test_node_id()]), } .into(), )); @@ -3088,7 +3400,7 @@ mod tests { ); // Assign + produce 3 records (single replica → commit inline on flush). - dp.handle_command(assign_segment(test_key(), vec![])); + dp.handle_command(assign_segment(test_key(), vec![test_node_id()])); let records: [(&[u8], u32); 3] = [(b"alpha", 2), (b"bravo", 5), (b"charlie", 1)]; for (payload, record_count) in records { let (tx, _rx) = oneshot::channel(); @@ -3127,7 +3439,7 @@ mod tests { .unwrap(); sparse.put_batch(index_entries).unwrap(); - dp.handle_command(DataPlaneCommand::DataPlaneInterNodeCommand( + dp.handle_command(receive_peer_message( SegmentSealed { segment_key: test_key(), committed_entry_id: None, @@ -3201,7 +3513,7 @@ mod tests { let mut tracker = SegmentTracker::new_with_start_entry_id( dir.path().to_path_buf(), SegmentRole::Leader, - vec![], + Replicas::new(vec![test_node_id()]), ShardGroupId(1), EntryId(0), ); @@ -3220,8 +3532,8 @@ mod tests { let (mut dp, cold_read_rx) = source_with_sealed_segment(&dir); let requester = NodeId::new("replacement"); - dp.handle_command(DataPlaneCommand::DataPlaneInterNodeCommand( - CatchUpRequest { + dp.handle_command(receive_peer_message( + RequestCatchUpEntries { segment_key: test_key(), from: requester.clone(), local_end: Some(EntryId(4)), @@ -3239,7 +3551,7 @@ mod tests { } /// A requester already at (or past) the source's committed end gets an - /// immediate `CatchUpStreamEnd` and no read is dispatched — the source short- + /// immediate `CatchUpEntriesSent` and no read is dispatched — the source short- /// circuits on its OWN `sealed_end`, not on anything the requester relayed. #[test] fn catch_up_request_already_caught_up_sends_done_without_reading() { @@ -3247,8 +3559,8 @@ mod tests { let (mut dp, cold_read_rx) = source_with_sealed_segment(&dir); let requester = NodeId::new("replacement"); - dp.handle_command(DataPlaneCommand::DataPlaneInterNodeCommand( - CatchUpRequest { + dp.handle_command(receive_peer_message( + RequestCatchUpEntries { segment_key: test_key(), from: requester.clone(), local_end: Some(EntryId(10)), // == the source's sealed end @@ -3258,7 +3570,7 @@ mod tests { // Nothing handed to the pool... assert!(cold_read_rx.try_recv().is_err()); - // ...and a CatchUpStreamEnd went straight back to the requester. + // ...and a CatchUpEntriesSent went straight back to the requester. assert_eq!(dp.out.transport_cmds.len(), 1); let DataTransportCommand::SendToTargets(s) = &dp.out.transport_cmds[0] else { panic!("expected SendToTargets"); @@ -3266,12 +3578,12 @@ mod tests { assert_eq!(s.targets[0], requester); assert!(matches!( &s.message, - DataPlaneInterNodeCommand::CatchUpStreamEnd(d) if d.segment_key == test_key() + DataPlanePeerMessage::CatchUpEntriesSent(d) if d.segment_key == test_key() )); } - /// A batch whose `next_offset` is past the sealed end emits one `CatchUpChunk` - /// (entries intact, with `record_count`) followed by `CatchUpStreamEnd` — no re-arm. + /// A batch whose `next_offset` is past the sealed end emits one `CatchUpEntries` + /// (entries intact, with `record_count`) followed by `CatchUpEntriesSent` — no re-arm. #[test] fn catch_up_read_complete_emits_chunk_then_done() { use crate::data_plane::states::segment::cache::CachedEntry; @@ -3316,8 +3628,8 @@ mod tests { panic!("expected chunk SendToTargets"); }; assert_eq!(chunk.targets[0], requester); - let DataPlaneInterNodeCommand::CatchUpChunk(c) = &chunk.message else { - panic!("expected CatchUpChunk"); + let DataPlanePeerMessage::CatchUpEntries(c) = &chunk.message else { + panic!("expected CatchUpEntries"); }; assert_eq!(c.segment_key, test_key()); let got: Vec<(u64, u32, Vec)> = c @@ -3339,12 +3651,12 @@ mod tests { }; assert!(matches!( &done.message, - DataPlaneInterNodeCommand::CatchUpStreamEnd(_) + DataPlanePeerMessage::CatchUpEntriesSent(_) )); } /// A batch that ends below the sealed end emits its chunk and re-arms the next - /// read against the pool, resuming at `next_offset` — no `CatchUpStreamEnd` yet. + /// read against the pool, resuming at `next_offset` — no `CatchUpEntriesSent` yet. #[test] fn catch_up_read_complete_re_arms_until_sealed_end() { use crate::data_plane::states::segment::cache::CachedEntry; @@ -3391,7 +3703,7 @@ mod tests { }; assert!(matches!( &chunk.message, - DataPlaneInterNodeCommand::CatchUpChunk(_) + DataPlanePeerMessage::CatchUpEntries(_) )); // The next read was re-armed against the pool, resuming at next_offset. @@ -3418,13 +3730,13 @@ mod tests { let source = NodeId::new("source"); // Coordinator assigns this node the sealed segment [0, 2]. - dp.handle_command(DataPlaneCommand::DataPlaneInterNodeCommand( - CatchUpAssignment { + dp.handle_command(receive_peer_message( + AssignSegmentCatchUp { segment_key: test_key(), shard_group_id: ShardGroupId(1), start_entry_id: EntryId(0), sealed_end_entry_id: EntryId(2), - replica_set: vec![test_node_id(), source.clone()], + replica_set: Replicas::new(vec![test_node_id(), source.clone()]), } .into(), )); @@ -3437,13 +3749,13 @@ mod tests { assert_eq!(s.targets[0], source); assert!(matches!( &s.message, - DataPlaneInterNodeCommand::CatchUpRequest(r) + DataPlanePeerMessage::RequestCatchUpEntries(r) if r.segment_key == test_key() && r.local_end.is_none() )); // Source streams the three entries, then signals end of stream. - dp.handle_command(DataPlaneCommand::DataPlaneInterNodeCommand( - CatchUpChunk { + dp.handle_command(receive_peer_message( + CatchUpEntries { segment_key: test_key(), entries: vec![ CatchUpEntry { @@ -3466,8 +3778,8 @@ mod tests { } .into(), )); - dp.handle_command(DataPlaneCommand::DataPlaneInterNodeCommand( - CatchUpStreamEnd { + dp.handle_command(receive_peer_message( + CatchUpEntriesSent { segment_key: test_key(), } .into(), @@ -3488,11 +3800,11 @@ mod tests { matches!(t, CheckpointTask::PutAnchors(anchors) if !anchors.is_empty()) }) ); - // ...and a CatchUpAck confirmed completion back to the coordinator. + // ...and a SegmentCaughtUp confirmed completion back to the coordinator. assert!(dp.out.transport_cmds.iter().any(|c| matches!( c, DataTransportCommand::SendToCoordinator(coord) - if matches!(&coord.message, DataPlaneInterNodeCommand::CatchUpAck(a) if a.segment_key == test_key()) + if matches!(&coord.message, DataPlanePeerMessage::SegmentCaughtUp(a) if a.segment_key == test_key()) ))); } @@ -3505,19 +3817,19 @@ mod tests { let dir = tempfile::tempdir().unwrap(); let mut dp = make_data_plane(&dir); - dp.handle_command(DataPlaneCommand::DataPlaneInterNodeCommand( - CatchUpAssignment { + dp.handle_command(receive_peer_message( + AssignSegmentCatchUp { segment_key: test_key(), shard_group_id: ShardGroupId(1), start_entry_id: EntryId(0), sealed_end_entry_id: EntryId(5), // need through 5... - replica_set: vec![test_node_id(), NodeId::new("source")], + replica_set: Replicas::new(vec![test_node_id(), NodeId::new("source")]), } .into(), )); // ...but only 0..=2 arrive. - dp.handle_command(DataPlaneCommand::DataPlaneInterNodeCommand( - CatchUpChunk { + dp.handle_command(receive_peer_message( + CatchUpEntries { segment_key: test_key(), entries: vec![ CatchUpEntry { @@ -3540,8 +3852,8 @@ mod tests { } .into(), )); - dp.handle_command(DataPlaneCommand::DataPlaneInterNodeCommand( - CatchUpStreamEnd { + dp.handle_command(receive_peer_message( + CatchUpEntriesSent { segment_key: test_key(), } .into(), @@ -3555,11 +3867,11 @@ mod tests { .iter() .any(|t| matches!(t, CheckpointTask::PutAnchors(_))) ); - // ...and NO CatchUpAck — an unverified receive must not confirm completion. + // ...and NO SegmentCaughtUp — an unverified receive must not confirm completion. assert!(!dp.out.transport_cmds.iter().any(|c| matches!( c, DataTransportCommand::SendToCoordinator(s) - if matches!(&s.message, DataPlaneInterNodeCommand::CatchUpAck(_)) + if matches!(&s.message, DataPlanePeerMessage::SegmentCaughtUp(_)) ))); } @@ -3572,18 +3884,18 @@ mod tests { // Recovered from disk through entry 5 — past the sealed end of 2. let mut dp = make_data_plane_with(&dir, inventory_with(test_key(), 5)); - dp.handle_command(DataPlaneCommand::DataPlaneInterNodeCommand( - CatchUpAssignment { + dp.handle_command(receive_peer_message( + AssignSegmentCatchUp { segment_key: test_key(), shard_group_id: ShardGroupId(1), start_entry_id: EntryId(0), sealed_end_entry_id: EntryId(2), - replica_set: vec![test_node_id(), NodeId::new("peer")], + replica_set: Replicas::new(vec![test_node_id(), NodeId::new("peer")]), } .into(), )); - // Zero transfer — no CatchUpRequest dispatched, only a CatchUpAck back to + // Zero transfer — no RequestCatchUpEntries dispatched, only a SegmentCaughtUp back to // the coordinator so its re-drive stops re-announcing the assignment. assert_eq!(dp.out.transport_cmds.len(), 1); let DataTransportCommand::SendToCoordinator(s) = &dp.out.transport_cmds[0] else { @@ -3591,7 +3903,7 @@ mod tests { }; assert!(matches!( &s.message, - DataPlaneInterNodeCommand::CatchUpAck(a) if a.segment_key == test_key() + DataPlanePeerMessage::SegmentCaughtUp(a) if a.segment_key == test_key() )); // ...and the segment is now registered cold-readable at the sealed bounds. assert_eq!( @@ -3662,13 +3974,13 @@ mod tests { write_seg_file(&path, &[(b"a", 1), (b"b", 1)]); let mut dp = make_data_plane_with(&dir, inventory_with(test_key(), 1)); - dp.handle_command(DataPlaneCommand::DataPlaneInterNodeCommand( - CatchUpAssignment { + dp.handle_command(receive_peer_message( + AssignSegmentCatchUp { segment_key: test_key(), shard_group_id: ShardGroupId(1), start_entry_id: EntryId(0), sealed_end_entry_id: EntryId(4), - replica_set: vec![test_node_id(), NodeId::new("source")], + replica_set: Replicas::new(vec![test_node_id(), NodeId::new("source")]), } .into(), )); @@ -3713,7 +4025,7 @@ mod tests { .insert_sealed_from_catch_up(key, EntryId(0), EntryId(0)); assert!(dp.segments.sealed_bounds(&key).is_some()); - dp.handle_command(DataPlaneCommand::DataPlaneInterNodeCommand( + dp.handle_command(receive_peer_message( DeleteSegments { segment_keys: Box::new([key]), } @@ -3742,7 +4054,7 @@ mod tests { let dir = tempfile::tempdir().unwrap(); let mut dp = make_data_plane(&dir); - dp.handle_command(DataPlaneCommand::DataPlaneInterNodeCommand( + dp.handle_command(receive_peer_message( DeleteSegments { segment_keys: Box::new([test_key()]), } @@ -3770,13 +4082,13 @@ mod tests { ); let mut dp = make_data_plane_with(&dir, inventory_with(test_key(), 1)); - dp.handle_command(DataPlaneCommand::DataPlaneInterNodeCommand( - CatchUpAssignment { + dp.handle_command(receive_peer_message( + AssignSegmentCatchUp { segment_key: test_key(), shard_group_id: ShardGroupId(1), start_entry_id: EntryId(0), sealed_end_entry_id: EntryId(4), - replica_set: vec![test_node_id(), NodeId::new("source")], + replica_set: Replicas::new(vec![test_node_id(), NodeId::new("source")]), } .into(), )); @@ -3788,7 +4100,7 @@ mod tests { assert!(matches!( &s.message, - DataPlaneInterNodeCommand::CatchUpRequest(r) if r.local_end == Some(EntryId(1)) + DataPlanePeerMessage::RequestCatchUpEntries(r) if r.local_end == Some(EntryId(1)) )); } @@ -3801,14 +4113,14 @@ mod tests { let leader = NodeId::new("leader"); let follower = NodeId::new("follower"); - dp.handle_command(DataPlaneCommand::DataPlaneInterNodeCommand( - CatchUpAssignment { + dp.handle_command(receive_peer_message( + AssignSegmentCatchUp { segment_key: test_key(), shard_group_id: ShardGroupId(1), start_entry_id: EntryId(0), sealed_end_entry_id: EntryId(2), // self (test-node) is a follower here; `leader` is replica_set[0]. - replica_set: vec![leader.clone(), test_node_id(), follower.clone()], + replica_set: Replicas::new(vec![leader.clone(), test_node_id(), follower.clone()]), } .into(), )); @@ -3825,21 +4137,21 @@ mod tests { // ── Re-drive rescues a stalled in-flight receive ─────────────────────── fn catch_up_assignment_to(source: &NodeId, sealed_end: u64) -> DataPlaneCommand { - DataPlaneCommand::DataPlaneInterNodeCommand( - CatchUpAssignment { + receive_peer_message( + AssignSegmentCatchUp { segment_key: test_key(), shard_group_id: ShardGroupId(1), start_entry_id: EntryId(0), sealed_end_entry_id: EntryId(sealed_end), - replica_set: vec![test_node_id(), source.clone()], + replica_set: Replicas::new(vec![test_node_id(), source.clone()]), } .into(), ) } fn catch_up_chunk_of(ids: &[u64]) -> DataPlaneCommand { - DataPlaneCommand::DataPlaneInterNodeCommand( - CatchUpChunk { + receive_peer_message( + CatchUpEntries { segment_key: test_key(), entries: ids .iter() @@ -3883,7 +4195,7 @@ mod tests { assert_eq!(s.targets[0], source); assert!(matches!( &s.message, - DataPlaneInterNodeCommand::CatchUpRequest(r) + DataPlanePeerMessage::RequestCatchUpEntries(r) if r.segment_key == test_key() && r.local_end.is_none() )); } @@ -3926,8 +4238,8 @@ mod tests { dp.handle_command(catch_up_chunk_of(&[0, 1, 2])); // Overlap: 1,2 are duplicates (skipped), 3,4 are new (appended). dp.handle_command(catch_up_chunk_of(&[1, 2, 3, 4])); - dp.handle_command(DataPlaneCommand::DataPlaneInterNodeCommand( - CatchUpStreamEnd { + dp.handle_command(receive_peer_message( + CatchUpEntriesSent { segment_key: test_key(), } .into(), @@ -3945,7 +4257,7 @@ mod tests { // --- Leader-crash boundary recovery: durable extent + query handler ----- /// A follower publishes (fsyncs) entries on flush but isn't told of the - /// commit until a `CommitAdvance` — so `durable_end` (the fsync'd extent) + /// commit until a `AdvanceReplicaCommit` — so `durable_end` (the fsync'd extent) /// runs ahead of `committed_entry_id`. Boundary recovery reads the former. #[test] fn durable_end_reports_published_extent_not_commit_cursor() { @@ -3954,10 +4266,10 @@ mod tests { let leader = NodeId::new("leader-node"); for entry_id in 0..3 { - dp.handle_command(DataPlaneCommand::DataPlaneInterNodeCommand( - ReplicaAppend { + dp.handle_command(receive_peer_message( + ReplicateSegmentEntries { segment_key: test_key(), - replica_set: vec![leader.clone(), test_node_id()], + replicas: Replicas::new(vec![leader.clone(), test_node_id()]), data: b"x".to_vec().into(), record_count: 1, entry_id: EntryId(entry_id), @@ -3965,7 +4277,7 @@ mod tests { .into(), )); } - dp.flush_batch(); // publishes 0,1,2 past the WAL fsync; no CommitAdvance yet + dp.flush_batch(); // publishes 0,1,2 past the WAL fsync; no AdvanceReplicaCommit yet assert_eq!( dp.segments.get(&test_key()).unwrap().committed_entry_id(), @@ -3999,15 +4311,15 @@ mod tests { assert_eq!(empty_dp.durable_end(&test_key()), None); } - /// A `SealBoundaryQuery` is answered with our durable extent, addressed back to + /// A `RequestDurableSegmentEnd` is answered with our durable extent, addressed back to /// the coordinator that owns this gather. #[test] fn seal_boundary_query_replies_with_durable_end() { let dir = tempfile::tempdir().unwrap(); let mut dp = make_data_plane_with(&dir, inventory_with(test_key(), 4)); - dp.handle_command(DataPlaneCommand::DataPlaneInterNodeCommand( - SealBoundaryQuery { + dp.handle_command(receive_peer_message( + RequestDurableSegmentEnd { segment_key: test_key(), coordinator: NodeId::new("coordinator"), } @@ -4019,8 +4331,8 @@ mod tests { panic!("boundary report must route to the coordinator"); }; assert_eq!(s.targets.as_ref(), &[NodeId::new("coordinator")]); - let DataPlaneInterNodeCommand::SealBoundaryReport(report) = &s.message else { - panic!("expected a SealBoundaryReport"); + let DataPlanePeerMessage::DurableSegmentEndReported(report) = &s.message else { + panic!("expected a DurableSegmentEndReported"); }; assert_eq!(report.segment_key, test_key()); assert_eq!(report.from, test_node_id()); diff --git a/src/data_plane/states/replication.rs b/src/data_plane/states/replication.rs index 7b8be1f0..9087fdbe 100644 --- a/src/data_plane/states/replication.rs +++ b/src/data_plane/states/replication.rs @@ -3,10 +3,10 @@ use std::sync::Arc; use tokio::sync::oneshot; -use crate::control_plane::NodeId; use crate::control_plane::metadata::EntryId; +use crate::control_plane::{NodeId, Replicas}; use crate::data_plane::SegmentKey; -use crate::data_plane::messages::command::{ProduceAck, ReplicaAppend}; +use crate::data_plane::messages::command::{ProduceAck, ReplicateSegmentEntries}; use crate::data_plane::states::segment::cache::CachedEntry; #[derive(Default)] @@ -32,16 +32,16 @@ pub(crate) struct AckCommitted { pub(crate) struct PendingReplicationBatch { pub segment_key: SegmentKey, pub entry: Arc, - pub replica_set: Vec, + pub replica_set: Replicas, pub followers: Vec, } impl PendingReplicationBatch { - pub(crate) fn into_replica_append(self) -> (Vec, ReplicaAppend) { + pub(crate) fn into_replica_append(self) -> (Vec, ReplicateSegmentEntries) { let targets = self.followers; - let message = ReplicaAppend { + let message = ReplicateSegmentEntries { segment_key: self.segment_key, - replica_set: self.replica_set, + replicas: self.replica_set, data: self.entry.data.clone(), record_count: self.entry.record_count, entry_id: self.entry.entry_id, @@ -241,7 +241,7 @@ mod tests { PendingReplicationBatch { segment_key, entry: cached_entry(entry_id), - replica_set: vec![node("leader")], + replica_set: Replicas::new(vec![node("leader")]), followers, } } diff --git a/src/data_plane/states/seal_request.rs b/src/data_plane/states/seal_request.rs index 8d3a8a50..b3de0141 100644 --- a/src/data_plane/states/seal_request.rs +++ b/src/data_plane/states/seal_request.rs @@ -1,7 +1,7 @@ -//! Pending seal requests — the leader→coordinator `SealRequest` retry tracker. +//! Pending segment-roll requests from the write leader to the coordinator. //! -//! A leader that can't replicate sends a `SealRequest` and tracks it here until a -//! `SealResponse` clears it; the data plane re-sends ones that time out. Owns the +//! A leader that can't replicate sends `RequestSegmentRoll` and tracks it until +//! `SegmentRollCommitted` clears it; the data plane re-sends timed-out requests. Owns the //! dedup + timing bookkeeping; building/sending the request stays in the data //! plane (it reads the live tracker's committed end). @@ -16,18 +16,18 @@ use tokio::time::Instant; use crate::control_plane::NodeId; use crate::data_plane::SegmentKey; -struct PendingSealRequest { +struct PendingSegmentRollRequest { sent_at: Instant, failed_nodes: Vec, } -/// Seal requests awaiting a `SealResponse`, keyed by segment. +/// Segment-roll requests awaiting `SegmentRollCommitted`, keyed by segment. #[derive(Default)] -pub(crate) struct PendingSealRequests { - requests: HashMap, +pub(crate) struct PendingSegmentRollRequests { + requests: HashMap, } -impl PendingSealRequests { +impl PendingSegmentRollRequests { /// Whether a request for `key` is already in flight (the dedup guard). pub(crate) fn is_tracked(&self, key: &SegmentKey) -> bool { self.requests.contains_key(key) @@ -37,14 +37,14 @@ impl PendingSealRequests { pub(crate) fn track(&mut self, key: SegmentKey, failed_nodes: Vec, now: Instant) { self.requests.insert( key, - PendingSealRequest { + PendingSegmentRollRequest { sent_at: now, failed_nodes, }, ); } - /// Drop a request once its `SealResponse` lands. + /// Drop a request once its committed result lands. pub(crate) fn clear(&mut self, key: &SegmentKey) { self.requests.remove(key); } @@ -83,7 +83,7 @@ mod tests { #[test] fn track_then_clear() { - let mut p = PendingSealRequests::default(); + let mut p = PendingSegmentRollRequests::default(); assert!(!p.is_tracked(&seg())); p.track(seg(), vec![], Instant::now()); assert!(p.is_tracked(&seg())); @@ -93,7 +93,7 @@ mod tests { #[test] fn take_due_returns_only_timed_out_then_refreshes() { - let mut p = PendingSealRequests::default(); + let mut p = PendingSegmentRollRequests::default(); let t0 = Instant::now(); let timeout = Duration::from_secs(5); p.track(seg(), vec![NodeId::new("d")], t0); diff --git a/src/data_plane/states/segment/tracker.rs b/src/data_plane/states/segment/tracker.rs index 690c9b65..15c82d25 100644 --- a/src/data_plane/states/segment/tracker.rs +++ b/src/data_plane/states/segment/tracker.rs @@ -1,8 +1,8 @@ use std::{path::PathBuf, sync::Arc, time::Duration}; -use crate::control_plane::NodeId; 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 super::cache::CachedEntry; @@ -23,7 +23,7 @@ pub(crate) struct SegmentTracker { checkpoint_lsn: u64, segment_file_path: PathBuf, role: SegmentRole, - replica_set: Vec, + replicas: Replicas, shard_group_id: ShardGroupId, committed_entry_id: EntryId, next_entry_id: EntryId, @@ -39,7 +39,7 @@ impl SegmentTracker { pub(crate) fn new( path: PathBuf, role: SegmentRole, - replica_set: Vec, + replicas: Replicas, shard_group_id: ShardGroupId, ) -> Self { Self { @@ -48,7 +48,7 @@ impl SegmentTracker { checkpoint_lsn: 0, segment_file_path: path, role, - replica_set, + replicas, shard_group_id, committed_entry_id: EntryId::MIN, next_entry_id: EntryId::MIN, @@ -61,7 +61,7 @@ impl SegmentTracker { pub(crate) fn new_with_start_entry_id( path: PathBuf, role: SegmentRole, - replica_set: Vec, + replica_set: Replicas, shard_group_id: ShardGroupId, start_entry_id: EntryId, ) -> Self { @@ -75,12 +75,12 @@ impl SegmentTracker { self.start_entry_id } - pub(crate) fn replica_set(&self) -> Vec { - self.replica_set.clone() + pub(crate) fn replica_set(&self) -> Replicas { + self.replicas.clone() } pub(crate) fn leader_node(&self) -> NodeId { - self.replica_set[0].clone() + self.replicas[0].clone() } pub(crate) fn role(&self) -> SegmentRole { @@ -88,8 +88,8 @@ impl SegmentTracker { } pub(crate) fn followers(&self) -> &[NodeId] { - if self.replica_set.len() > 1 { - &self.replica_set[1..] + if self.replicas.len() > 1 { + &self.replicas[1..] } else { &[] } @@ -262,20 +262,9 @@ impl SegmentTracker { segment_key: SegmentKey, data: EntryPayload, record_count: u32, - ) { - self.size_bytes += data.len() as u64; - self.staged_entries - .push(StagedEntry::new(data, record_count, segment_key)); - } - - pub(crate) fn stage_entry_from_replica( - &mut self, - segment_key: SegmentKey, - data: EntryPayload, - record_count: u32, entry_id: EntryId, ) { - let expected = self.next_entry_id + self.staged_entries.len() as u64; + let expected = self.next_staged_entry_id(); if entry_id < expected { return; } @@ -288,6 +277,10 @@ impl SegmentTracker { .push(StagedEntry::new(data, record_count, segment_key)); } + pub(crate) fn next_staged_entry_id(&self) -> EntryId { + self.next_entry_id + self.staged_entries.len() as u64 + } + pub(crate) fn shard_group_id(&self) -> ShardGroupId { self.shard_group_id } @@ -381,7 +374,7 @@ pub mod tests { SegmentTracker::new( PathBuf::from("/tmp/test.seg"), role, - vec![NodeId::new("leader"), NodeId::new("follower")], + Replicas::new(vec![NodeId::new("leader"), NodeId::new("follower")]), ShardGroupId(1), ) } @@ -389,7 +382,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); + t.stage_entry(test_key(), Bytes::from("abcde").into(), 2, EntryId(0)); assert_eq!(t.size_bytes, 5); assert!(t.has_staged()); @@ -401,11 +394,11 @@ pub mod tests { let mut t = SegmentTracker::new_with_start_entry_id( PathBuf::from("/tmp/test.seg"), SegmentRole::Follower, - vec![NodeId::new("leader"), NodeId::new("follower")], + Replicas::new(vec![NodeId::new("leader"), NodeId::new("follower")]), ShardGroupId(1), EntryId(5), ); - t.stage_entry_from_replica(test_key(), Bytes::from("data").into(), 1, EntryId(5)); + t.stage_entry(test_key(), Bytes::from("data").into(), 1, EntryId(5)); assert!(t.has_staged()); // Publish to advance next_entry_id @@ -414,7 +407,7 @@ pub mod tests { t.publish_staged(1); // Duplicate entry_id (5) should be skipped since next is now 6 - t.stage_entry_from_replica(test_key(), Bytes::from("dup").into(), 1, EntryId(5)); + t.stage_entry(test_key(), Bytes::from("dup").into(), 1, EntryId(5)); assert!(!t.has_staged()); assert_eq!(t.next_entry_id, EntryId(6)); t.assert_invariants(); @@ -426,7 +419,12 @@ pub mod tests { let mut wal_buf = Vec::new(); for i in 0..3u64 { - t.stage_entry(test_key(), Bytes::from(format!("entry-{i}")).into(), 1); + t.stage_entry( + test_key(), + Bytes::from(format!("entry-{i}")).into(), + 1, + EntryId(i), + ); t.stage_to_wal(&mut wal_buf); t.publish_staged(i + 1); } @@ -451,11 +449,11 @@ pub mod tests { let mut t = SegmentTracker::new_with_start_entry_id( PathBuf::new(), SegmentRole::Follower, - vec![NodeId::new("leader"), NodeId::new("follower")], + Replicas::new(vec![NodeId::new("leader"), NodeId::new("follower")]), ShardGroupId(1), EntryId(2), ); - t.stage_entry_from_replica(test_key(), Bytes::from("entry-2").into(), 1, EntryId(2)); + t.stage_entry(test_key(), Bytes::from("entry-2").into(), 1, EntryId(2)); t.publish_staged(1); t.commit_entry(EntryId(1)); @@ -468,7 +466,7 @@ 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); + t.stage_entry(test_key(), Bytes::from("payload").into(), 3, EntryId(0)); t.stage_to_wal(&mut wal_buf); assert!(!wal_buf.is_empty()); @@ -492,7 +490,7 @@ pub mod tests { let single = SegmentTracker::new( PathBuf::from("/tmp/t.seg"), SegmentRole::Leader, - vec![NodeId::new("solo")], + Replicas::new(vec![NodeId::new("solo")]), ShardGroupId(1), ); assert!(single.followers().is_empty()); diff --git a/src/data_plane/states/segment_store.rs b/src/data_plane/states/segment_store.rs index 0437d16b..68f9a542 100644 --- a/src/data_plane/states/segment_store.rs +++ b/src/data_plane/states/segment_store.rs @@ -107,7 +107,7 @@ impl SegmentStore { /// Insert an active tracker AND its `active_by_range` index entry. /// Idempotent — if the key already exists the new tracker is dropped, so - /// callers can blindly retry a SegmentAssignment without checking first. + /// callers can blindly retry a PlaceSegment without checking first. /// The active index keys by `tracker.start_entry_id()` (invariant 2). /// /// Refuses to resurrect an already-sealed segment. A sealed segment is @@ -331,6 +331,7 @@ impl SegmentStore { mod tests { use super::*; use crate::control_plane::membership::ShardGroupId; + use crate::control_plane::{NodeId, Replicas}; use crate::data_plane::states::segment::tracker::SegmentRole; use std::path::PathBuf; @@ -342,7 +343,7 @@ mod tests { SegmentTracker::new_with_start_entry_id( PathBuf::from("/tmp"), SegmentRole::Leader, - vec![], + Replicas::new(vec![NodeId::new("tracker")]), ShardGroupId(1), EntryId(start), ) diff --git a/src/data_plane/transport/command.rs b/src/data_plane/transport/command.rs index 8e0f69bd..ec8b09b3 100644 --- a/src/data_plane/transport/command.rs +++ b/src/data_plane/transport/command.rs @@ -1,18 +1,18 @@ use crate::control_plane::NodeId; use crate::control_plane::membership::ShardGroupId; -use crate::data_plane::messages::command::DataPlaneInterNodeCommand; +use crate::data_plane::messages::command::DataPlanePeerMessage; use crate::impl_from_variant; #[derive(Debug)] pub(crate) struct DataTransportSendToTargets { pub targets: Box<[NodeId]>, - pub message: DataPlaneInterNodeCommand, + pub message: DataPlanePeerMessage, } #[derive(Debug)] pub(crate) struct DataTransportSendToCoordinator { pub shard_group_id: ShardGroupId, - pub message: DataPlaneInterNodeCommand, + pub message: DataPlanePeerMessage, } #[derive(Debug)] @@ -31,11 +31,11 @@ impl_from_variant!( impl DataTransportCommand { pub(crate) fn send_to_targets( - targets: Vec, - message: impl Into, + targets: impl Into>, + message: impl Into, ) -> Self { Self::SendToTargets(DataTransportSendToTargets { - targets: targets.into_boxed_slice(), + targets: targets.into(), message: message.into(), }) } diff --git a/src/data_plane/transport/reader.rs b/src/data_plane/transport/reader.rs index 4755980d..b6a62a58 100644 --- a/src/data_plane/transport/reader.rs +++ b/src/data_plane/transport/reader.rs @@ -3,7 +3,7 @@ use tokio::sync::mpsc; use crate::control_plane::NodeId; use crate::data_plane::actor::DataPlaneSender; -use crate::data_plane::messages::command::{DataPlaneCommand, DataPlaneInterNodeCommand}; +use crate::data_plane::messages::command::{DataPlaneCommand, ReceivePeerMessage}; use crate::net::OwnedReadHalf; const NODE_ID_FRAME_MAX: usize = 1024; @@ -32,12 +32,17 @@ impl DataReader { disconnect_tx: mpsc::Sender, ) { loop { - match self - .read_frame::(DATA_FRAME_MAX) - .await - { - Ok(msg) => { - let _ = data_plane_tx.send(DataPlaneCommand::DataPlaneInterNodeCommand(msg)); + match self.read_frame::(DATA_FRAME_MAX).await { + Ok(message) => { + if message.from != peer { + tracing::warn!( + transport_peer = ?peer, + claimed_sender = ?message.from, + "rejected peer message whose sender differs from the connection peer" + ); + return; + } + let _ = data_plane_tx.send(DataPlaneCommand::ReceivePeerMessage(message)); } Err(e) => { tracing::debug!("DataReader connection closed: {e}"); diff --git a/src/data_plane/transport/writers.rs b/src/data_plane/transport/writers.rs index f7a7ea99..22725ae1 100644 --- a/src/data_plane/transport/writers.rs +++ b/src/data_plane/transport/writers.rs @@ -8,7 +8,9 @@ use tokio::time::Instant; use crate::control_plane::NodeId; use crate::control_plane::membership::actor::SwimSender; use crate::data_plane::actor::DataPlaneSender; -use crate::data_plane::messages::command::{DataPlaneCommand, DataPlaneInterNodeCommand}; +use crate::data_plane::messages::command::{ + DataPlaneCommand, DataPlanePeerMessage, ReceivePeerMessage, +}; use crate::net::{OwnedWriteHalf, TcpStream}; use super::reader::DataReader; @@ -63,17 +65,20 @@ impl TransportState { pub async fn send( &mut self, targets: &[NodeId], - msg: &DataPlaneInterNodeCommand, + msg: &DataPlanePeerMessage, swim_tx: &SwimSender, data_plane_tx: &DataPlaneSender, disconnect_tx: &mpsc::Sender, ) { for target in targets { - // Self-delivery: a node can be its own target (e.g. a SegmentAssignment + // Self-delivery: a node can be its own target (e.g. a PlaceSegment // to `replica_set[0]` if *target == self.node_id { let _ = - data_plane_tx.send(DataPlaneCommand::DataPlaneInterNodeCommand(msg.clone())); + data_plane_tx.send(DataPlaneCommand::ReceivePeerMessage(ReceivePeerMessage { + from: self.node_id.clone(), + message: Box::new(msg.clone()), + })); continue; } @@ -117,7 +122,7 @@ impl TransportState { async fn connect_and_send( &mut self, target_id: NodeId, - msg: &DataPlaneInterNodeCommand, + msg: &DataPlanePeerMessage, swim_tx: &SwimSender, ) -> anyhow::Result { let node_addr = swim_tx @@ -174,13 +179,16 @@ impl TransportState { async fn write_message( &mut self, target: &NodeId, - msg: &DataPlaneInterNodeCommand, + msg: &DataPlanePeerMessage, ) -> anyhow::Result<()> { let writer = self .writers .get_mut(target) .context("no writer for target")?; - let bytes = borsh::to_vec(msg)?; + let bytes = borsh::to_vec(&ReceivePeerMessage { + from: self.node_id.clone(), + message: Box::new(msg.clone()), + })?; let len = bytes.len() as u32; let mut buf = Vec::with_capacity(4 + bytes.len()); buf.extend_from_slice(&len.to_be_bytes()); diff --git a/src/it/e2e/client_protocol.rs b/src/it/e2e/client_protocol.rs index 69af0f22..b9b6b790 100644 --- a/src/it/e2e/client_protocol.rs +++ b/src/it/e2e/client_protocol.rs @@ -553,7 +553,7 @@ fn age_seal_rolls_the_active_segment() -> turmoil::Result { } // The segment ages past `max_segment_age_secs`; the leader enqueues a - // SealRequest, the coordinator commits a RollSegment, and the active + // RequestSegmentRoll, the coordinator commits a RollSegment, and the active // segment rolls. Poll DescribeTopic until the successor's start_offset // advances past 0 (the known-end seal landed). let mut rolled_start = None; @@ -585,7 +585,7 @@ fn age_seal_rolls_the_active_segment() -> turmoil::Result { /// End-to-end sealed-segment repair. Produce + age-seal a segment, kill one of its /// replicas, and confirm the coordinator reassigns the segment to the spare node, /// which catches up and serves the old data on a cold fetch. Exercises the whole -/// loop: SWIM death → reconcile → `ReassignSegment` → `CatchUpAssignment` → +/// loop: SWIM death → reconcile → `ReassignSegment` → `AssignSegmentCatchUp` → /// catch-up → cold-serve, across the control/data-plane seam. #[test] #[serial_test::serial] @@ -1137,7 +1137,7 @@ fn sealed_repair_survives_coordinator_crash() -> turmoil::Result { /// End-to-end catch-up redrive under message loss. Seal a segment, then **partition /// the spare from the coordinator** before killing one of the segment's replicas (a /// *non-coordinator* one, so the coordinator survives to drive the repair). With the -/// link down, the coordinator's `CatchUpAssignment` to the spare — and every heartbeat +/// link down, the coordinator's `AssignSegmentCatchUp` to the spare — and every heartbeat /// redrive of it — are dropped. After a window we heal the link; the next redrive /// delivers, the spare pulls the segment from a surviving source (never partitioned), /// and serves the old data. Guards the redrive hardening (`raft-actor.md` #9): a lost @@ -1533,7 +1533,7 @@ async fn metadata_leader(topic: &str, nodes: &[(&str, u16)]) -> Option<(ShardGro // ── produce/fetch e2e helpers ────────────────────────────────────────────── /// Wait until every node sees `expected` alive members. Topic creation emits the -/// initial `SegmentAssignment` exactly once (fire-and-forget, no retry); if it's +/// initial `PlaceSegment` exactly once (fire-and-forget, no retry); if it's /// sent before SWIM has converged, the data-leader's address isn't yet resolvable /// on the metadata leader and the assignment is dropped permanently. Gating /// creation on convergence keeps the produce path off that drop window. @@ -1622,7 +1622,7 @@ async fn produce_until_acked<'a>( nodes: &'a [(&'a str, u16)], ) -> Option<(&'a str, u16)> { // Generous budget: the data plane, replication, and checkpoint run on real - // OS threads (outside turmoil's deterministic runtime), so SegmentAssignment + // OS threads (outside turmoil's deterministic runtime), so PlaceSegment // delivery + first commit can take a while on a slow interleaving. for _ in 0..80 { for &(host, port) in nodes { diff --git a/src/it/e2e/consumer_group_test.rs b/src/it/e2e/consumer_group_test.rs index 5814bf8d..60b88bf2 100644 --- a/src/it/e2e/consumer_group_test.rs +++ b/src/it/e2e/consumer_group_test.rs @@ -740,6 +740,119 @@ fn consumer_group_rebalance_survives_broker_restart_without_offset_replay() { sim.run().unwrap(); } +#[test] +#[serial_test::serial] +fn restarted_offset_replica_bootstraps_missed_epoch_before_commit_ack() { + let mut sim = Builder::new() + .tick_duration(Duration::from_millis(10)) + .simulation_duration(Duration::from_secs(75)) + .rng_seed(187) + .build(); + + host_cluster(&mut sim, NODES, |env| { + env.node_id_suffix = Some("sim-187".to_string()); + env.vnodes_per_node = 16; + }); + + // 0 = setup, 1 = crash replica, 2 = replica down, 3 = replica restarted. + let phase = Arc::new(AtomicU8::new(0)); + let client_phase = phase.clone(); + sim.client("client", async move { + let client_addr: std::net::SocketAddr = + format!("{}:9091", turmoil::lookup("n1")).parse().unwrap(); + let admin = Arc::new(Client::connect([client_addr]).unwrap()); + let producer_client = Arc::new(Client::connect([client_addr]).unwrap()); + let first_client = Arc::new(Client::connect([client_addr]).unwrap()); + let second_client = Arc::new(Client::connect([client_addr]).unwrap()); + let topic = "rebalance_restart_offsets"; + create_group_test_topics(&admin, topic).await; + let producer = Producer::new( + producer_client, + topic.to_string(), + ProducerConfig::default(), + ); + let group_id = "offset-replica-epoch-bootstrap-group".to_string(); + let first = Consumer::new( + first_client, + topic.to_string(), + KeyInterest::AllKeys, + group_config(group_id.clone()), + ) + .await + .unwrap(); + + producer.send(b"before", b"before".to_vec()).await.unwrap(); + let before = tokio::time::timeout(Duration::from_secs(8), first.next_record()) + .await + .expect("initial consume timed out") + .expect("initial consume failed") + .expect("initial consumer closed"); + first.ack(&before).unwrap(); + first.commit().await.unwrap(); + + client_phase.store(1, Ordering::SeqCst); + while client_phase.load(Ordering::SeqCst) < 2 { + tokio::time::sleep(Duration::from_millis(10)).await; + } + + // This membership change advances the group generation while n2 is down, + // so n2 misses the epoch seal but remains in the unchanged data placement. + let second = Consumer::new( + second_client, + topic.to_string(), + KeyInterest::AllKeys, + group_config(group_id), + ) + .await + .unwrap(); + + while client_phase.load(Ordering::SeqCst) < 3 { + tokio::time::sleep(Duration::from_millis(10)).await; + } + + producer.send(b"after", b"after".to_vec()).await.unwrap(); + tokio::select! { + result = first.next_record() => { + let record = result.unwrap().expect("first consumer closed"); + assert_eq!(record.value, b"after"); + first.ack(&record).unwrap(); + first.commit().await.unwrap(); + } + result = second.next_record() => { + let record = result.unwrap().expect("second consumer closed"); + assert_eq!(record.value, b"after"); + second.ack(&record).unwrap(); + second.commit().await.unwrap(); + } + _ = tokio::time::sleep(Duration::from_secs(12)) => { + panic!("post-restart record was not delivered"); + } + } + Ok(()) + }); + + while phase.load(Ordering::SeqCst) < 1 { + assert!( + sim.elapsed() < Duration::from_secs(35), + "client never reached the replica-crash phase" + ); + sim.step().unwrap(); + } + sim.crash("n2"); + phase.store(2, Ordering::SeqCst); + let restart_at = sim.elapsed() + Duration::from_secs(6); + while sim.elapsed() < restart_at { + sim.step().unwrap(); + } + sim.bounce("n2"); + let recovered_at = sim.elapsed() + Duration::from_secs(6); + while sim.elapsed() < recovered_at { + sim.step().unwrap(); + } + phase.store(3, Ordering::SeqCst); + sim.run().unwrap(); +} + #[test] fn consumer_group_split_rebalance() { use crate::control_plane::metadata::RangeId; diff --git a/src/it/raft/election.rs b/src/it/raft/election.rs index 19003b2e..f9a51890 100644 --- a/src/it/raft/election.rs +++ b/src/it/raft/election.rs @@ -7,12 +7,12 @@ use tokio::io::{AsyncReadExt, AsyncWriteExt}; use tokio::sync::{mpsc, oneshot}; use turmoil::Builder; -use crate::control_plane::NodeId; use crate::control_plane::consensus::actor::MultiRaftActor; use crate::control_plane::consensus::messages::{EnsureGroup, MultiRaftActorCommand, RaftTimer}; use crate::control_plane::consensus::transport::RaftTransportActor; use crate::control_plane::membership::actor::SwimActor; use crate::control_plane::membership::{ShardGroup, ShardGroupId}; +use crate::control_plane::{NodeId, Replicas}; use crate::impls::metadata_storage::MetadataStorage; use crate::net::{TcpListener, TcpStream}; @@ -171,14 +171,14 @@ fn three_node_raft_elects_leader() -> turmoil::Result { .build(); let group_id = ShardGroupId(42); - let members: Vec = vec![ + let members = Replicas::new(vec![ NodeId::new("node-1"), NodeId::new("node-2"), NodeId::new("node-3"), - ]; + ]); let group = ShardGroup { id: group_id, - members, + replicas: members, }; for (name, port, peers) in [ diff --git a/src/it/raft/leader_event.rs b/src/it/raft/leader_event.rs index 42943c0b..69985030 100644 --- a/src/it/raft/leader_event.rs +++ b/src/it/raft/leader_event.rs @@ -14,7 +14,7 @@ use crate::control_plane::membership::actor::SwimActor; use crate::control_plane::membership::{ QueryCommand, ShardGroup, ShardGroupId, SwimActorCommand, SwimCommand, }; -use crate::control_plane::{NodeAddress, NodeId}; +use crate::control_plane::{NodeAddress, NodeId, Replicas}; use crate::impls::metadata_storage::MetadataStorage; use crate::net::{TcpListener, TcpStream}; use crate::schedulers::actor::spawn_scheduling_actor; @@ -60,11 +60,11 @@ fn leader_election_emits_leader_change_event() -> turmoil::Result { let group_id = ShardGroupId(77); let group = ShardGroup { id: group_id, - members: vec![ + replicas: Replicas::new(vec![ NodeId::new("node-1"), NodeId::new("node-2"), NodeId::new("node-3"), - ], + ]), }; for (name, port, peers) in [ diff --git a/src/it/raft/membership_change.rs b/src/it/raft/membership_change.rs index 7877c03b..e786101b 100644 --- a/src/it/raft/membership_change.rs +++ b/src/it/raft/membership_change.rs @@ -14,7 +14,7 @@ use crate::control_plane::consensus::messages::*; use crate::control_plane::consensus::transport::RaftTransportActor; use crate::control_plane::membership::actor::SwimActor; use crate::control_plane::membership::{ShardGroup, ShardGroupId, Topology}; -use crate::control_plane::{NodeId, SwimNodeState}; +use crate::control_plane::{NodeId, Replicas, SwimNodeState}; use crate::impls::metadata_storage::MetadataStorage; use crate::net::{TcpListener, TcpStream}; use crate::schedulers::actor::spawn_scheduling_actor; @@ -196,11 +196,11 @@ fn node_death_triggers_remove_peer() -> turmoil::Result { let group_id = ShardGroupId(99); let group = ShardGroup { id: group_id, - members: vec![ + replicas: Replicas::new(vec![ NodeId::new("node-1"), NodeId::new("node-2"), NodeId::new("node-3"), - ], + ]), }; for (name, port, peers) in [ @@ -275,11 +275,11 @@ fn node_death_removes_dead_voter_unreachable_spare_stays_non_voting() -> turmoil let group_id = ShardGroupId(101); let initial_group = ShardGroup { id: group_id, - members: vec![ + replicas: Replicas::new(vec![ NodeId::new("node-1"), NodeId::new("node-2"), NodeId::new("node-3"), - ], + ]), }; // peer_names threaded into start_raft_node seed the stub topology — all