1use std::collections::{BTreeMap, HashMap, VecDeque};
2use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
3use std::sync::{Arc, Mutex, OnceLock};
4use std::time::Duration;
5
6use bytes::Bytes;
7use futures::future::join_all;
8use futures::stream::{FuturesUnordered, StreamExt};
9use sha2::{Digest, Sha256};
10use tokio::sync::watch;
11
12use crate::grpc::pack_append;
13use crate::metrics::Metrics;
14use crate::record::{RecordError, RecordFrame};
15use crate::transport::{
16 AppendToken, LaneDurableChange, PackedAppend, Replica, ReplicaSnapshot, TransportCode,
17 TransportError, FORMAT_VERSION, META_FORMAT,
18};
19
20pub(crate) const SUPPORTED_REPLICA_COUNTS: [usize; 3] = [1, 3, 5];
23
24pub(crate) const DEFAULT_LANE_STALL_TIMEOUT: Duration = Duration::from_secs(5);
29
30pub(crate) fn majority(replica_count: usize) -> usize {
32 replica_count / 2 + 1
33}
34
35fn select_recovery_size(sizes: &mut [i64], replica_count: usize) -> Option<i64> {
36 let quorum = majority(replica_count);
37 if sizes.len() < quorum || sizes.len() > replica_count {
38 return None;
39 }
40 sizes.sort_unstable();
41 let unavailable = replica_count - sizes.len();
42 let required_available_support = quorum.checked_sub(unavailable)?;
43 (required_available_support > 0).then(|| sizes[sizes.len() - required_available_support])
44}
45
46pub(crate) type AttemptedBytes = Arc<dyn Fn(u64) + Send + Sync>;
47
48#[derive(Clone, Debug)]
49pub struct ClientConfig {
54 pub max_retries: usize,
56 pub retry_base: Duration,
59}
60
61impl Default for ClientConfig {
62 fn default() -> Self {
63 Self {
64 max_retries: 5,
65 retry_base: Duration::from_millis(20),
66 }
67 }
68}
69
70pub(crate) struct QuorumVolume {
71 replicas: Vec<Arc<dyn Replica>>,
72 config: ClientConfig,
73 metadata: HashMap<String, String>,
74 metrics: Arc<Metrics>,
75}
76
77struct LaneHandle {
80 work: tokio::sync::mpsc::UnboundedSender<LaneBatch>,
81 done: tokio::task::JoinHandle<Option<AppendToken>>,
82 budget: Arc<LaneBudget>,
83 stall_timeout: Arc<LaneStallTimeout>,
84}
85
86#[derive(Debug)]
87struct LaneBudget {
88 outstanding: AtomicUsize,
89 limit: AtomicUsize,
90}
91
92impl LaneBudget {
93 fn new() -> Arc<Self> {
94 Arc::new(Self {
95 outstanding: AtomicUsize::new(0),
96 limit: AtomicUsize::new(usize::MAX),
97 })
98 }
99
100 fn set_limit(&self, limit: usize) {
101 self.limit.store(limit, Ordering::Relaxed);
102 }
103
104 fn try_reserve(self: &Arc<Self>, bytes: usize) -> Option<Arc<LaneReservation>> {
105 let mut current = self.outstanding.load(Ordering::Relaxed);
106 loop {
107 let next = current.checked_add(bytes)?;
108 if next > self.limit.load(Ordering::Relaxed) {
109 return None;
110 }
111 match self.outstanding.compare_exchange_weak(
112 current,
113 next,
114 Ordering::Relaxed,
115 Ordering::Relaxed,
116 ) {
117 Ok(_) => {
118 return Some(Arc::new(LaneReservation {
119 budget: Arc::clone(self),
120 bytes,
121 }));
122 }
123 Err(observed) => current = observed,
124 }
125 }
126 }
127}
128
129#[derive(Debug)]
130struct LaneStallTimeout {
131 nanos: AtomicU64,
132}
133
134impl LaneStallTimeout {
135 fn new() -> Arc<Self> {
136 Arc::new(Self {
137 nanos: AtomicU64::new(Self::encode(DEFAULT_LANE_STALL_TIMEOUT)),
138 })
139 }
140
141 fn set(&self, timeout: Duration) {
142 self.nanos.store(Self::encode(timeout), Ordering::Relaxed);
143 }
144
145 fn get(&self) -> Duration {
146 Duration::from_nanos(self.nanos.load(Ordering::Relaxed))
147 }
148
149 fn encode(timeout: Duration) -> u64 {
150 u64::try_from(timeout.as_nanos()).unwrap_or(u64::MAX).max(1)
151 }
152}
153
154#[derive(Debug)]
155struct LaneReservation {
156 budget: Arc<LaneBudget>,
157 bytes: usize,
158}
159
160impl Drop for LaneReservation {
161 fn drop(&mut self) {
162 self.budget
163 .outstanding
164 .fetch_sub(self.bytes, Ordering::Relaxed);
165 }
166}
167
168pub(crate) struct Writer {
169 replicas: Vec<Arc<dyn Replica>>,
170 config: ClientConfig,
171 lanes: Vec<Option<LaneHandle>>,
172 admitted: AdmittedPrefix,
173 commits: Arc<CommitTracker>,
174 sealed: bool,
175 metadata: HashMap<String, String>,
176 metrics: Arc<Metrics>,
177}
178
179struct AdmittedPrefix {
188 records: usize,
189 bytes: usize,
190 digest: DigestState,
191 crc32c: u32,
192}
193
194enum DigestState {
195 Ready(Sha256),
196 Pending {
197 sender: tokio::sync::mpsc::UnboundedSender<Arc<[Bytes]>>,
198 task: tokio::task::JoinHandle<Sha256>,
199 },
200}
201
202#[derive(Clone, Debug, Default, Eq, PartialEq)]
203pub(crate) struct SealReport {
204 finalized: Vec<Option<ReplicaSnapshot>>,
205}
206
207impl SealReport {
208 pub fn all_replicas_finalized(&self) -> bool {
209 !self.finalized.is_empty()
212 && self.finalized.iter().flatten().count() == self.finalized.len()
213 }
214}
215
216#[derive(Clone, Debug, Default)]
217pub(crate) struct CanonicalPrefix {
218 bytes: Vec<u8>,
219 records: Vec<RecordFrame>,
220 record_ends: Vec<usize>,
221}
222
223pub(crate) struct RecoveredTail {
228 canonical: CanonicalPrefix,
229 had_discarded_suffix: bool,
230}
231
232pub(crate) enum RecoveryCandidate {
234 Absent,
236 Empty {
241 reusable_writer: Option<Box<Writer>>,
242 },
243 NonEmpty(RecoveredTail),
245}
246
247impl RecoveredTail {
248 pub fn len(&self) -> usize {
249 self.canonical.len()
250 }
251
252 pub fn digest(&self) -> String {
255 digest_bytes(&self.canonical.bytes)
256 }
257
258 pub fn crc32c(&self) -> u32 {
260 crc32c::crc32c(&self.canonical.bytes)
261 }
262
263 pub fn canonical(&self) -> &CanonicalPrefix {
264 &self.canonical
265 }
266
267 pub fn had_discarded_suffix(&self) -> bool {
271 self.had_discarded_suffix
272 }
273}
274
275#[derive(Clone, Debug)]
276enum CommitFailure {
277 Poisoned,
278 Fenced(String),
279 Transport(TransportError),
280}
281
282impl CommitFailure {
283 fn protocol_error(&self) -> ProtocolError {
284 match self {
285 Self::Poisoned => ProtocolError::Poisoned,
286 Self::Fenced(error) => ProtocolError::Fenced(error.clone()),
287 Self::Transport(error) => ProtocolError::Transport(error.clone()),
288 }
289 }
290}
291
292#[derive(Clone, Debug, Default)]
293struct CommitSnapshot {
294 committed: usize,
295 failure: Option<CommitFailure>,
296}
297
298#[derive(Debug)]
299struct LaneCommitState {
300 durable: i64,
301 represented_end: i64,
302 finished: bool,
303 error: Option<TransportError>,
304}
305
306impl Default for LaneCommitState {
307 fn default() -> Self {
308 Self {
309 durable: 0,
310 represented_end: 0,
311 finished: true,
312 error: None,
313 }
314 }
315}
316
317#[derive(Debug)]
318struct CommitState {
319 boundaries: VecDeque<i64>,
320 admitted: usize,
321 committed: usize,
322 committed_bytes: usize,
323 failure: Option<CommitFailure>,
324 lanes: Vec<LaneCommitState>,
325}
326
327struct CommitTracker {
331 quorum: usize,
332 state: Mutex<CommitState>,
333 updates: watch::Sender<CommitSnapshot>,
334 metrics: Arc<Metrics>,
335}
336
337pub(crate) struct CommitRange {
338 first_offset: usize,
339 end_offset: usize,
340 updates: watch::Receiver<CommitSnapshot>,
341}
342
343#[cfg(test)]
344pub(crate) struct PendingCommit {
345 pub logical_offset: u64,
346 updates: watch::Receiver<CommitSnapshot>,
347}
348
349#[cfg(test)]
350impl PendingCommit {
351 pub async fn wait(mut self) -> Result<u64, ProtocolError> {
352 loop {
353 let snapshot = self.updates.borrow_and_update().clone();
354 if snapshot.committed > self.logical_offset as usize {
355 return Ok(self.logical_offset);
356 }
357 if let Some(failure) = snapshot.failure {
358 return Err(failure.protocol_error());
359 }
360 self.updates
361 .changed()
362 .await
363 .map_err(|_| ProtocolError::PipelineClosed)?;
364 }
365 }
366}
367
368#[derive(Debug, thiserror::Error)]
369pub(crate) enum ProtocolError {
370 #[error("the replica count must be 1, 3, or 5")]
371 ReplicaCount,
372 #[error("operation did not reach a replica quorum")]
373 NoQuorum,
374 #[error("writer is poisoned by an indeterminate record and must be recovered")]
375 Poisoned,
376 #[error("writer was fenced: {0}")]
377 Fenced(String),
378 #[error("recovery witnesses contain different bytes at record {record_index}")]
379 ConflictingPrefix { record_index: usize },
380 #[error("recovery prefix has {actual} records, expected at least {expected}")]
381 RecoveryPrefixTooShort { expected: usize, actual: usize },
382 #[error("recovered seal digest {actual} does not match committed digest {expected}")]
383 SealDigestMismatch { expected: String, actual: String },
384 #[error("recovered seal CRC32C {actual:08x} does not match committed CRC32C {expected:08x}")]
385 SealCrc32cMismatch { expected: u32, actual: u32 },
386 #[error("manifest register is invalid: {0}")]
387 InvalidManifest(String),
388 #[error(
389 "the manifest segment directory is full: truncate the WAL to free \
390 retained sealed segments before sealing again"
391 )]
392 SegmentDirectoryFull,
393 #[error(transparent)]
394 ManifestStore(#[from] crate::manifest_store::ManifestStoreError),
395 #[error("manifest register is unavailable")]
396 ManifestUnavailable,
397 #[error("commit pipeline closed before reporting a result")]
398 PipelineClosed,
399 #[error("segment writer is sealed")]
400 Finalized,
401 #[error(transparent)]
402 Record(#[from] RecordError),
403 #[error("transport error: {0}")]
404 Transport(#[from] TransportError),
405}
406
407impl CanonicalPrefix {
408 pub fn len(&self) -> usize {
409 self.records.len()
410 }
411
412 pub fn into_records(self) -> Vec<RecordFrame> {
413 self.records
414 }
415
416 fn truncate(&mut self, records: usize) {
417 self.records.truncate(records);
418 self.record_ends.truncate(records);
419 self.bytes.truncate(self.committed_bytes_len(records));
420 }
421
422 fn committed_bytes_len(&self, records: usize) -> usize {
423 records
424 .checked_sub(1)
425 .and_then(|index| self.record_ends.get(index).copied())
426 .unwrap_or(0)
427 }
428
429 fn record_bytes(&self, index: usize) -> &[u8] {
430 let start = index
431 .checked_sub(1)
432 .and_then(|previous| self.record_ends.get(previous).copied())
433 .unwrap_or(0);
434 &self.bytes[start..self.record_ends[index]]
435 }
436
437 fn from_snapshot(snapshot: &ReplicaSnapshot) -> Self {
438 let (records, consumed) = RecordFrame::decode_complete_prefix(&snapshot.bytes);
439 let mut record_ends = Vec::with_capacity(records.len());
440 let mut end = 0usize;
441 for record in &records {
442 end += record.encode().expect("a decoded record must encode").len();
443 record_ends.push(end);
444 }
445 Self {
446 bytes: snapshot.bytes[..consumed].to_vec(),
447 records,
448 record_ends,
449 }
450 }
451}
452
453impl Default for AdmittedPrefix {
454 fn default() -> Self {
455 Self {
456 records: 0,
457 bytes: 0,
458 digest: DigestState::Ready(Sha256::new()),
459 crc32c: 0,
460 }
461 }
462}
463
464impl AdmittedPrefix {
465 fn len(&self) -> usize {
466 self.records
467 }
468
469 fn is_empty(&self) -> bool {
470 self.records == 0
471 }
472
473 fn bytes_len(&self) -> usize {
474 self.bytes
475 }
476
477 fn extend_metadata(&mut self, chunks: &[Bytes]) {
478 for chunk in chunks {
479 self.crc32c = crc32c::crc32c_append(self.crc32c, chunk);
480 self.bytes += chunk.len();
481 self.records += 1;
482 }
483 }
484
485 fn queue_digest(&mut self, chunks: Arc<[Bytes]>) {
486 let previous = std::mem::replace(&mut self.digest, DigestState::Ready(Sha256::new()));
487 let pending = match previous {
488 DigestState::Ready(mut hasher) => {
489 let (sender, mut receiver) = tokio::sync::mpsc::unbounded_channel::<Arc<[Bytes]>>();
490 let task = tokio::spawn(async move {
491 while let Some(chunks) = receiver.recv().await {
492 hasher = tokio::task::spawn_blocking(move || {
493 for chunk in chunks.iter() {
494 hasher.update(chunk);
495 }
496 hasher
497 })
498 .await
499 .expect("ordered digest task failed");
500 }
501 hasher
502 });
503 sender.send(chunks).expect("digest worker failed");
504 DigestState::Pending { sender, task }
505 }
506 DigestState::Pending { sender, task } => {
507 sender.send(chunks).expect("digest worker failed");
508 DigestState::Pending { sender, task }
509 }
510 };
511 self.digest = pending;
512 }
513
514 async fn digest(&mut self) -> String {
515 let pending = std::mem::replace(&mut self.digest, DigestState::Ready(Sha256::new()));
516 let hasher = match pending {
517 DigestState::Ready(hasher) => hasher,
518 DigestState::Pending { sender, task } => {
519 drop(sender);
520 task.await.expect("ordered digest task failed")
521 }
522 };
523 let digest = digest_hex(hasher.clone().finalize());
524 self.digest = DigestState::Ready(hasher);
525 digest
526 }
527
528 async fn shutdown_digest(&mut self) {
529 let digest = std::mem::replace(&mut self.digest, DigestState::Ready(Sha256::new()));
530 if let DigestState::Pending { sender, task } = digest {
531 drop(sender);
532 let _ = task.await;
533 }
534 }
535
536 fn crc32c(&self) -> u32 {
537 self.crc32c
538 }
539}
540
541impl Drop for AdmittedPrefix {
542 fn drop(&mut self) {
543 let digest = std::mem::replace(&mut self.digest, DigestState::Ready(Sha256::new()));
544 if let DigestState::Pending { sender, task } = digest {
545 drop(sender);
546 task.abort();
547 }
548 }
549}
550
551impl CommitTracker {
552 fn new(lanes: usize, quorum: usize, metrics: Arc<Metrics>) -> Arc<Self> {
553 let snapshot = CommitSnapshot::default();
554 let (updates, _) = watch::channel(snapshot);
555 Arc::new(Self {
556 quorum,
557 state: Mutex::new(CommitState {
558 boundaries: VecDeque::new(),
559 admitted: 0,
560 committed: 0,
561 committed_bytes: 0,
562 failure: None,
563 lanes: (0..lanes).map(|_| LaneCommitState::default()).collect(),
564 }),
565 updates,
566 metrics,
567 })
568 }
569
570 fn activate_lane(&self, zone: usize, durable: i64) {
571 let mut state = self
572 .state
573 .lock()
574 .unwrap_or_else(|poisoned| poisoned.into_inner());
575 let lane = &mut state.lanes[zone];
576 let previous_lag = lane_durable_lag(lane);
577 lane.durable = durable.max(0);
578 lane.finished = false;
579 self.metrics
580 .adjust_zone_durable_lag(zone, lane_durable_lag(lane) - previous_lag);
581 }
582
583 fn admitted_len(&self) -> usize {
584 self.state
585 .lock()
586 .unwrap_or_else(|poisoned| poisoned.into_inner())
587 .admitted
588 }
589
590 fn committed_len(&self) -> usize {
591 self.state
592 .lock()
593 .unwrap_or_else(|poisoned| poisoned.into_inner())
594 .committed
595 }
596
597 fn committed_bytes(&self) -> usize {
598 self.state
599 .lock()
600 .unwrap_or_else(|poisoned| poisoned.into_inner())
601 .committed_bytes
602 }
603
604 fn is_poisoned(&self) -> bool {
605 self.state
606 .lock()
607 .unwrap_or_else(|poisoned| poisoned.into_inner())
608 .failure
609 .is_some()
610 }
611
612 fn subscribe(&self) -> watch::Receiver<CommitSnapshot> {
613 self.updates.subscribe()
614 }
615
616 fn admit_window(&self, boundaries: &[i64], represented_zones: &[usize]) -> CommitRange {
617 let mut state = self
618 .state
619 .lock()
620 .unwrap_or_else(|poisoned| poisoned.into_inner());
621 let first_offset = state.admitted;
622 let end = boundaries
623 .last()
624 .copied()
625 .expect("an admitted window is non-empty");
626 state.admitted += boundaries.len();
627 state.boundaries.extend(boundaries.iter().copied());
628 for &zone in represented_zones {
629 let lane = &mut state.lanes[zone];
630 let previous_lag = lane_durable_lag(lane);
631 lane.represented_end = lane.represented_end.max(end);
632 self.metrics
633 .adjust_zone_durable_lag(zone, lane_durable_lag(lane) - previous_lag);
634 }
635 self.recompute_locked(&mut state);
636 CommitRange {
637 first_offset,
638 end_offset: state.admitted,
639 updates: self.subscribe(),
640 }
641 }
642
643 fn publish_durable(&self, zone: usize, durable: i64) {
644 let mut state = self
645 .state
646 .lock()
647 .unwrap_or_else(|poisoned| poisoned.into_inner());
648 let lane = &mut state.lanes[zone];
649 if durable <= lane.durable {
650 return;
651 }
652 let previous_lag = lane_durable_lag(lane);
653 lane.durable = durable;
654 self.metrics
655 .adjust_zone_durable_lag(zone, lane_durable_lag(lane) - previous_lag);
656 self.recompute_locked(&mut state);
657 }
658
659 fn finish_lane(&self, zone: usize, error: Option<TransportError>) {
660 let mut state = self
661 .state
662 .lock()
663 .unwrap_or_else(|poisoned| poisoned.into_inner());
664 let fence = error
665 .as_ref()
666 .filter(|error| error.code.fences_writer())
667 .map(ToString::to_string);
668 let lane = &mut state.lanes[zone];
669 let previous_lag = lane_durable_lag(lane);
670 lane.finished = true;
671 if error.is_some() {
672 lane.error = error;
673 }
674 self.metrics
675 .adjust_zone_durable_lag(zone, lane_durable_lag(lane) - previous_lag);
676 if let Some(fence) = fence {
677 if state.failure.is_none() {
680 state.failure = Some(CommitFailure::Fenced(fence));
681 self.publish_locked(&state);
682 }
683 return;
684 }
685 self.recompute_locked(&mut state);
686 }
687
688 fn poison(&self) {
689 let mut state = self
690 .state
691 .lock()
692 .unwrap_or_else(|poisoned| poisoned.into_inner());
693 if state.failure.is_none() {
694 state.failure = Some(CommitFailure::Poisoned);
695 self.publish_locked(&state);
696 }
697 }
698
699 fn recompute_locked(&self, state: &mut CommitState) {
700 if state.failure.is_some() {
701 return;
702 }
703 let quorum_watermark = quorum_durable_watermark(&state.lanes, self.quorum);
704 let mut changed = false;
705 while state
706 .boundaries
707 .front()
708 .is_some_and(|boundary| *boundary <= quorum_watermark)
709 {
710 let boundary = state
711 .boundaries
712 .pop_front()
713 .expect("front boundary was present");
714 state.committed += 1;
715 state.committed_bytes =
716 usize::try_from(boundary).expect("record boundaries are nonnegative");
717 changed = true;
718 }
719
720 if let Some(&oldest) = state.boundaries.front() {
721 let possible = state
722 .lanes
723 .iter()
724 .filter(|lane| {
725 lane.durable >= oldest || (!lane.finished && lane.represented_end >= oldest)
726 })
727 .count();
728 if possible < self.quorum {
729 state.failure = Some(select_commit_failure(&state.lanes, oldest));
730 changed = true;
731 }
732 }
733 if changed {
734 self.publish_locked(state);
735 }
736 }
737
738 fn publish_locked(&self, state: &CommitState) {
739 self.updates.send_replace(CommitSnapshot {
740 committed: state.committed,
741 failure: state.failure.clone(),
742 });
743 }
744}
745
746impl Drop for CommitTracker {
747 fn drop(&mut self) {
748 let state = self
749 .state
750 .get_mut()
751 .unwrap_or_else(|poisoned| poisoned.into_inner());
752 for (zone, lane) in state.lanes.iter_mut().enumerate() {
753 let lag = lane_durable_lag(lane);
754 if lag != 0 {
755 self.metrics.adjust_zone_durable_lag(zone, -lag);
756 lane.finished = true;
757 }
758 }
759 }
760}
761
762fn lane_durable_lag(lane: &LaneCommitState) -> i64 {
763 if lane.finished {
764 0
765 } else {
766 lane.represented_end.saturating_sub(lane.durable).max(0)
767 }
768}
769
770fn quorum_durable_watermark(lanes: &[LaneCommitState], quorum: usize) -> i64 {
771 let mut durables: Vec<_> = lanes.iter().map(|lane| lane.durable).collect();
772 durables.sort_unstable();
773 durables[durables.len() - quorum]
774}
775
776fn select_commit_failure(lanes: &[LaneCommitState], boundary: i64) -> CommitFailure {
777 let represented = lanes.iter().filter(|lane| lane.represented_end >= boundary);
778 let mut fenced = None;
779 let mut terminal = None;
780 for error in represented.filter_map(|lane| lane.error.clone()) {
781 if error.code.fences_writer() {
782 prefer_lower_zone(&mut fenced, error);
783 } else if !error.code.transient() {
784 prefer_lower_zone(&mut terminal, error);
785 }
786 }
787 match (fenced, terminal) {
788 (Some(error), _) => CommitFailure::Fenced(error.to_string()),
789 (None, Some(error)) => CommitFailure::Transport(error),
790 (None, None) => CommitFailure::Poisoned,
791 }
792}
793
794impl CommitRange {
795 pub(crate) fn first_offset(&self) -> usize {
796 self.first_offset
797 }
798
799 pub(crate) fn end_offset(&self) -> usize {
800 self.end_offset
801 }
802
803 #[cfg(test)]
804 pub(crate) fn into_pending(self) -> Vec<PendingCommit> {
805 (self.first_offset..self.end_offset)
806 .map(|logical_offset| PendingCommit {
807 logical_offset: logical_offset as u64,
808 updates: self.updates.clone(),
809 })
810 .collect()
811 }
812
813 pub(crate) fn progress(&mut self) -> (usize, Option<ProtocolError>) {
814 let snapshot = self.updates.borrow_and_update().clone();
815 (
816 snapshot.committed,
817 snapshot.failure.map(|failure| failure.protocol_error()),
818 )
819 }
820
821 pub(crate) async fn changed(&mut self) -> Result<(), ProtocolError> {
822 self.updates
823 .changed()
824 .await
825 .map_err(|_| ProtocolError::PipelineClosed)
826 }
827}
828
829pub(crate) fn digest_bytes(bytes: &[u8]) -> String {
830 digest_hex(Sha256::digest(bytes))
831}
832
833fn digest_hex(digest: impl AsRef<[u8]>) -> String {
834 use std::fmt::Write;
835
836 let digest = digest.as_ref();
837 let mut encoded = String::with_capacity(digest.len() * 2);
838 for byte in digest {
839 write!(&mut encoded, "{byte:02x}").expect("writing to a String cannot fail");
840 }
841 encoded
842}
843
844impl QuorumVolume {
845 pub fn with_metadata(
851 replicas: Vec<Arc<dyn Replica>>,
852 config: ClientConfig,
853 metadata: HashMap<String, String>,
854 metrics: Arc<Metrics>,
855 ) -> Result<Self, ProtocolError> {
856 if !SUPPORTED_REPLICA_COUNTS.contains(&replicas.len()) {
857 return Err(ProtocolError::ReplicaCount);
858 }
859 Ok(Self {
860 replicas,
861 config,
862 metadata,
863 metrics,
864 })
865 }
866
867 fn quorum(&self) -> usize {
868 majority(self.replicas.len())
869 }
870
871 pub async fn create_writer(&self) -> Result<Writer, ProtocolError> {
876 let metadata = self.metadata.clone();
877 let creates = join_all(
878 self.replicas
879 .iter()
880 .map(|replica| create_session_with_retry(replica, metadata.clone(), &self.config)),
881 )
882 .await;
883 let tokens: Vec<_> = creates.into_iter().flatten().collect();
884 if tokens.len() < self.quorum() {
885 return Err(ProtocolError::NoQuorum);
886 }
887 Ok(Writer::new(
888 self.replicas.clone(),
889 self.config.clone(),
890 self.metadata.clone(),
891 tokens,
892 Arc::clone(&self.metrics),
893 ))
894 }
895
896 pub async fn recover_for_seal(
901 &self,
902 expected_records: Option<usize>,
903 ) -> Result<RecoveredTail, ProtocolError> {
904 match self.recover_candidate(expected_records).await? {
905 RecoveryCandidate::NonEmpty(tail) => Ok(tail),
906 RecoveryCandidate::Empty { .. } | RecoveryCandidate::Absent => {
907 Err(ProtocolError::RecoveryPrefixTooShort {
908 expected: expected_records.unwrap_or(1),
909 actual: 0,
910 })
911 }
912 }
913 }
914
915 pub(crate) async fn recover_candidate(
924 &self,
925 expected_records: Option<usize>,
926 ) -> Result<RecoveryCandidate, ProtocolError> {
927 enum Observation {
928 Live(AppendToken),
929 Finalized(ReplicaSnapshot),
930 Missing(usize),
931 }
932
933 let attempts = join_all(self.replicas.iter().map(|replica| {
934 let replica = Arc::clone(replica);
935 let config = self.config.clone();
936 async move {
937 match takeover_current_with_retry(&replica, &config).await {
938 Ok(token) => Ok::<_, ProtocolError>(Observation::Live(token)),
939 Err(error) if error.code == TransportCode::FailedPrecondition => {
940 let snapshot = snapshot_with_retry(&replica, &config).await?;
941 if valid_format(&snapshot.metadata) {
942 Ok(Observation::Finalized(snapshot))
943 } else {
944 Err(ProtocolError::NoQuorum)
945 }
946 }
947 Err(error) if error.code == TransportCode::NotFound => {
948 Ok(Observation::Missing(error.zone))
949 }
950 Err(error)
951 if error.code.transient()
952 || matches!(
953 error.code,
954 TransportCode::DataLoss | TransportCode::Ambiguous
955 ) =>
956 {
957 Err(ProtocolError::NoQuorum)
958 }
959 Err(error) => Err(error.into()),
960 }
961 }
962 }))
963 .await;
964
965 let mut observations = Vec::new();
966 for attempt in attempts {
967 match attempt {
968 Ok(observation) => observations.push(observation),
969 Err(ProtocolError::NoQuorum) => {}
970 Err(error) => return Err(error),
971 }
972 }
973 if observations.len() < self.quorum() {
974 return Err(ProtocolError::NoQuorum);
975 }
976 if observations
977 .iter()
978 .all(|observation| matches!(observation, Observation::Missing(_)))
979 {
980 return Ok(RecoveryCandidate::Absent);
981 }
982
983 let missing_zones: Vec<_> = observations
984 .iter()
985 .filter_map(|observation| match observation {
986 Observation::Missing(zone) => Some(*zone),
987 _ => None,
988 })
989 .collect();
990
991 let mut sizes = observations
992 .iter()
993 .map(|observation| match observation {
994 Observation::Live(token) => token.persisted_size,
995 Observation::Finalized(snapshot) => snapshot.persisted_size,
996 Observation::Missing(_) => 0,
997 })
998 .collect::<Vec<_>>();
999 if sizes.len() < self.quorum() {
1000 return Err(ProtocolError::NoQuorum);
1001 }
1002 let max_observed_size = *sizes
1003 .iter()
1004 .max()
1005 .expect("a quorum supplied at least one size");
1006 let committed_size = select_recovery_size(&mut sizes, self.replicas.len())
1015 .expect("a reachable majority intersects every write quorum");
1016
1017 if committed_size == 0 {
1018 let mut empty_witnesses = observations
1019 .iter()
1020 .filter(|observation| {
1021 matches!(
1022 observation,
1023 Observation::Live(token) if token.persisted_size == 0
1024 ) || matches!(
1025 observation,
1026 Observation::Finalized(snapshot) if snapshot.persisted_size == 0
1027 )
1028 })
1029 .count();
1030 let mut tokens = observations
1031 .into_iter()
1032 .filter_map(|observation| match observation {
1033 Observation::Live(token) if token.persisted_size == 0 => Some(token),
1034 _ => None,
1035 })
1036 .collect::<Vec<_>>();
1037 if tokens.len() < self.quorum() {
1038 let creates = join_all(missing_zones.into_iter().map(|zone| {
1039 create_session_with_retry(
1040 &self.replicas[zone],
1041 self.metadata.clone(),
1042 &self.config,
1043 )
1044 }))
1045 .await;
1046 for token in creates.into_iter().flatten() {
1047 empty_witnesses += 1;
1048 tokens.push(token);
1049 }
1050 }
1051 if empty_witnesses < self.quorum() {
1052 return Err(ProtocolError::NoQuorum);
1053 }
1054 let reusable_writer = (tokens.len() == self.replicas.len()).then(|| {
1055 Box::new(Writer::new(
1056 self.replicas.clone(),
1057 self.config.clone(),
1058 self.metadata.clone(),
1059 tokens,
1060 Arc::clone(&self.metrics),
1061 ))
1062 });
1063 return Ok(RecoveryCandidate::Empty { reusable_writer });
1064 }
1065
1066 let mut live_tokens = Vec::new();
1067 let mut recovery_snapshots = Vec::new();
1068 for observation in observations {
1069 match observation {
1070 Observation::Live(token) => live_tokens.push(token),
1071 Observation::Finalized(snapshot) => recovery_snapshots.push(snapshot),
1072 Observation::Missing(_) => {}
1073 }
1074 }
1075 let frozen = join_all(live_tokens.into_iter().map(|mut token| {
1076 let replica = Arc::clone(&self.replicas[token.zone]);
1077 let config = self.config.clone();
1078 async move {
1079 let persisted_size = token.persisted_size;
1080 finalize_with_retry(&replica, &mut token, persisted_size, &config).await
1081 }
1082 }))
1083 .await;
1084 let frozen_zones = frozen
1085 .into_iter()
1086 .flatten()
1087 .map(|snapshot| snapshot.zone)
1088 .collect::<Vec<_>>();
1089 if frozen_zones.len() + recovery_snapshots.len() + missing_zones.len() < self.quorum() {
1090 return Err(ProtocolError::NoQuorum);
1091 }
1092
1093 let committed_size = usize::try_from(committed_size).map_err(|_| {
1094 ProtocolError::InvalidManifest("persisted segment size does not fit in usize".into())
1095 })?;
1096 let max_observed_size = usize::try_from(max_observed_size).map_err(|_| {
1097 ProtocolError::InvalidManifest("persisted segment size does not fit in usize".into())
1098 })?;
1099 let reads = join_all(
1100 frozen_zones
1101 .into_iter()
1102 .map(|zone| snapshot_with_retry(&self.replicas[zone], &self.config)),
1103 )
1104 .await;
1105 for snapshot in reads.into_iter().flatten() {
1106 if valid_format(&snapshot.metadata) {
1107 recovery_snapshots.push(snapshot);
1108 }
1109 }
1110 for snapshot in &mut recovery_snapshots {
1111 snapshot
1112 .bytes
1113 .truncate(snapshot.bytes.len().min(committed_size));
1114 }
1115 recovery_snapshots.extend(missing_zones.into_iter().map(|zone| ReplicaSnapshot {
1116 zone,
1117 generation: 0,
1118 metageneration: 0,
1119 persisted_size: 0,
1120 finalized: true,
1121 crc32c: Some(crc32c::crc32c(&[])),
1122 metadata: self.metadata.clone(),
1123 bytes: Vec::new(),
1124 }));
1125 let (mut canonical, _) = select_canonical_quorum(&recovery_snapshots, self.quorum())?;
1126 let had_discarded_suffix = max_observed_size > canonical.bytes.len();
1127 if let Some(expected) = expected_records {
1128 if canonical.len() < expected {
1129 return Err(ProtocolError::RecoveryPrefixTooShort {
1130 expected,
1131 actual: canonical.len(),
1132 });
1133 }
1134 canonical.truncate(expected);
1135 }
1136 if canonical.len() == 0 {
1137 return Ok(RecoveryCandidate::Empty {
1138 reusable_writer: None,
1139 });
1140 }
1141 Ok(RecoveryCandidate::NonEmpty(RecoveredTail {
1142 canonical,
1143 had_discarded_suffix,
1144 }))
1145 }
1146
1147 pub async fn enforce_seal(&self, canonical: &CanonicalPrefix) -> Result<(), ProtocolError> {
1158 let data = Bytes::from(canonical.bytes.clone());
1159 let reads = join_all(
1160 self.replicas
1161 .iter()
1162 .map(|replica| snapshot_with_retry(replica, &self.config)),
1163 )
1164 .await;
1165 let mut witnesses: Vec<ReplicaSnapshot> = Vec::new();
1166 let mut missing: Vec<usize> = Vec::new();
1167 for (zone, read) in reads.into_iter().enumerate() {
1168 match read {
1169 Ok(snapshot) if valid_format(&snapshot.metadata) => witnesses.push(snapshot),
1170 Ok(_) => {}
1171 Err(error) if error.code == TransportCode::NotFound => missing.push(zone),
1172 Err(error) if error.code == TransportCode::DataLoss => {
1173 match stat_with_retry(&self.replicas[zone], &self.config).await {
1174 Ok(snapshot) if valid_format(&snapshot.metadata) => {
1175 witnesses.push(snapshot);
1179 }
1180 Ok(_) => {}
1181 Err(stat_error) if stat_error.code == TransportCode::NotFound => {
1182 missing.push(zone);
1183 }
1184 Err(stat_error) if stat_error.code.transient() => {}
1185 Err(stat_error) => return Err(stat_error.into()),
1186 }
1187 }
1188 Err(error) if error.code.transient() => {}
1189 Err(error) => return Err(error.into()),
1190 }
1191 }
1192 witnesses.sort_by_key(|snapshot| {
1194 let shared = snapshot
1195 .bytes
1196 .iter()
1197 .zip(&canonical.bytes)
1198 .take_while(|(a, b)| a == b)
1199 .count();
1200 (shared == canonical.bytes.len() && snapshot.bytes.len() == canonical.bytes.len())
1201 as usize
1202 * canonical.bytes.len()
1203 + shared
1204 });
1205 let mut finalized = 0usize;
1206 for snapshot in witnesses {
1207 match self.enforce_witness(snapshot, &data).await {
1208 Ok(true) => finalized += 1,
1209 Ok(false) => {}
1210 Err(error) => return Err(error),
1211 }
1212 }
1213 for zone in missing {
1214 let replica = Arc::clone(&self.replicas[zone]);
1215 let Ok(created) =
1216 create_with_retry(&replica, self.metadata.clone(), &self.config).await
1217 else {
1218 continue;
1219 };
1220 match self.enforce_witness(created, &data).await {
1221 Ok(true) => finalized += 1,
1222 Ok(false) => {}
1223 Err(error) => return Err(error),
1224 }
1225 }
1226 if finalized >= self.quorum() {
1227 Ok(())
1228 } else {
1229 Err(ProtocolError::NoQuorum)
1230 }
1231 }
1232
1233 async fn enforce_witness(
1234 &self,
1235 mut snapshot: ReplicaSnapshot,
1236 data: &Bytes,
1237 ) -> Result<bool, ProtocolError> {
1238 let zone = snapshot.zone;
1239 let replica = Arc::clone(&self.replicas[zone]);
1240 for _ in 0..=self.config.max_retries {
1241 if snapshot.finalized {
1242 if snapshot.bytes == data[..] {
1243 return Ok(true);
1244 }
1245 } else if snapshot.bytes == data[..] {
1247 let Ok(mut token) = takeover_with_retry(&replica, &snapshot, &self.config).await
1248 else {
1249 snapshot = snapshot_with_retry(&replica, &self.config).await?;
1250 continue;
1251 };
1252 match finalize_with_retry(&replica, &mut token, data.len() as i64, &self.config)
1253 .await
1254 {
1255 Ok(_) => return Ok(true),
1256 Err(error) if error.code == TransportCode::FailedPrecondition => {
1257 snapshot = snapshot_with_retry(&replica, &self.config).await?;
1258 continue;
1259 }
1260 Err(error) if error.code.transient() => return Ok(false),
1261 Err(error) => return Err(error.into()),
1262 }
1263 }
1264 let mut token = match replace_with_retry(
1265 &replica,
1266 snapshot.clone(),
1267 data.clone(),
1268 self.metadata.clone(),
1269 &self.config,
1270 )
1271 .await
1272 {
1273 Ok(token) => token,
1274 Err(error) if error.code == TransportCode::FailedPrecondition => {
1275 snapshot = snapshot_with_retry(&replica, &self.config).await?;
1276 continue;
1277 }
1278 Err(error) if error.code.transient() => return Ok(false),
1279 Err(error) => return Err(error.into()),
1280 };
1281 match finalize_with_retry(&replica, &mut token, data.len() as i64, &self.config).await {
1282 Ok(_) => return Ok(true),
1283 Err(error) if error.code == TransportCode::FailedPrecondition => {
1284 snapshot = snapshot_with_retry(&replica, &self.config).await?;
1285 }
1286 Err(error) if error.code.transient() => return Ok(false),
1287 Err(error) => return Err(error.into()),
1288 }
1289 }
1290 Ok(false)
1291 }
1292}
1293
1294impl Drop for Writer {
1295 fn drop(&mut self) {
1296 for lane in self.lanes.iter_mut().filter_map(Option::take) {
1297 lane.done.abort();
1298 }
1299 }
1300}
1301
1302impl Writer {
1303 fn new(
1304 replicas: Vec<Arc<dyn Replica>>,
1305 config: ClientConfig,
1306 metadata: HashMap<String, String>,
1307 tokens: Vec<AppendToken>,
1308 metrics: Arc<Metrics>,
1309 ) -> Self {
1310 let mut by_zone: HashMap<usize, AppendToken> = tokens
1311 .into_iter()
1312 .map(|token| (token.zone, token))
1313 .collect();
1314 let commits = CommitTracker::new(
1315 replicas.len(),
1316 majority(replicas.len()),
1317 Arc::clone(&metrics),
1318 );
1319 let lanes = (0..replicas.len())
1320 .map(|zone| {
1321 by_zone.remove(&zone).map(|token| {
1322 commits.activate_lane(zone, token.persisted_size);
1323 let (work, rx) = tokio::sync::mpsc::unbounded_channel();
1324 let replica = Arc::clone(&replicas[zone]);
1325 let config = config.clone();
1326 let metrics = Arc::clone(&metrics);
1327 let budget = LaneBudget::new();
1328 let stall_timeout = LaneStallTimeout::new();
1329 let lane_commits = Arc::clone(&commits);
1330 LaneHandle {
1331 work,
1332 done: tokio::spawn(run_lane(
1333 replica,
1334 token,
1335 config,
1336 metrics,
1337 rx,
1338 lane_commits,
1339 Arc::clone(&stall_timeout),
1340 )),
1341 budget,
1342 stall_timeout,
1343 }
1344 })
1345 })
1346 .collect();
1347 Self {
1348 replicas,
1349 config,
1350 metadata,
1351 lanes,
1352 admitted: AdmittedPrefix::default(),
1353 commits,
1354 sealed: false,
1355 metrics,
1356 }
1357 }
1358
1359 fn quorum(&self) -> usize {
1360 majority(self.replicas.len())
1361 }
1362
1363 pub fn admitted_len(&self) -> usize {
1364 self.commits.admitted_len()
1365 }
1366
1367 pub(crate) fn set_max_replica_lag_bytes(&mut self, limit: usize) {
1368 for lane in self.lanes.iter().flatten() {
1369 lane.budget.set_limit(limit);
1370 }
1371 }
1372
1373 pub(crate) fn set_lane_stall_timeout(&mut self, timeout: Duration) {
1374 for lane in self.lanes.iter().flatten() {
1375 lane.stall_timeout.set(timeout);
1376 }
1377 }
1378
1379 pub(crate) async fn shutdown_background_tasks(&mut self) {
1380 let lanes = std::mem::take(&mut self.lanes);
1381 let mut tasks = Vec::new();
1382 for lane in lanes.into_iter().flatten() {
1383 lane.done.abort();
1384 tasks.push(lane.done);
1385 }
1386 let _ = join_all(tasks).await;
1387 self.admitted.shutdown_digest().await;
1388 join_all(self.replicas.iter().map(|replica| replica.shutdown())).await;
1389 }
1390
1391 pub fn committed_len(&self) -> usize {
1392 self.commits.committed_len()
1393 }
1394
1395 pub fn is_poisoned(&self) -> bool {
1396 self.commits.is_poisoned()
1397 }
1398
1399 pub fn physical_size(&self) -> usize {
1400 self.commits.committed_bytes()
1401 }
1402
1403 pub async fn seal_digest(&mut self) -> String {
1408 self.admitted.digest().await
1409 }
1410
1411 pub fn seal_crc32c(&self) -> u32 {
1415 self.admitted.crc32c()
1416 }
1417
1418 pub async fn enqueue_data_window(
1419 &mut self,
1420 records: Vec<RecordFrame>,
1421 on_attempted: AttemptedBytes,
1422 ) -> Result<CommitRange, ProtocolError> {
1423 if records.is_empty() {
1424 return Ok(CommitRange {
1425 first_offset: self.admitted_len(),
1426 end_offset: self.admitted_len(),
1427 updates: self.commits.subscribe(),
1428 });
1429 }
1430 if self.sealed {
1431 return Err(ProtocolError::Finalized);
1432 }
1433 if self.is_poisoned() {
1434 return Err(ProtocolError::Poisoned);
1435 }
1436 let chunks: Result<Vec<_>, _> = records.iter().map(RecordFrame::encode).collect();
1437 let chunks: Arc<[Bytes]> = chunks?.into();
1438 drop(records);
1439 let quorum = self.quorum();
1440 let active_lanes = self.lanes.iter().filter(|lane| lane.is_some()).count();
1441 if active_lanes < quorum {
1442 tracing::warn!(active_lanes, required_lanes = quorum, "WAL writer poisoned");
1443 self.commits.poison();
1444 return Err(ProtocolError::NoQuorum);
1445 }
1446 self.metrics.batches_sent.increment();
1447 let batch_bytes = chunks.iter().map(Bytes::len).sum::<usize>();
1448 let mut reservations = Vec::with_capacity(self.lanes.len());
1449 let mut retired_lanes = Vec::new();
1450 for zone in 0..self.lanes.len() {
1451 let reservation = self.lanes[zone]
1452 .as_ref()
1453 .and_then(|lane| lane.budget.try_reserve(batch_bytes));
1454 if reservation.is_none() && self.lanes[zone].is_some() {
1455 let lane = self.lanes[zone].take().expect("lane checked present");
1456 lane.done.abort();
1457 retired_lanes.push(lane.done);
1458 self.commits.finish_lane(zone, None);
1459 self.metrics.lane_capacity_drops.increment();
1460 tracing::warn!(
1461 zone,
1462 batch_bytes,
1463 "replica lane exceeded its retained-byte budget"
1464 );
1465 }
1466 reservations.push(reservation);
1467 }
1468 let _ = join_all(retired_lanes).await;
1469 if reservations.iter().flatten().count() < quorum {
1470 self.commits.poison();
1471 return Err(ProtocolError::NoQuorum);
1472 }
1473 let start = self.admitted.bytes_len() as i64;
1474 let mut boundaries = Vec::with_capacity(chunks.len());
1475 let mut acc = start;
1476 for chunk in chunks.iter() {
1477 acc += chunk.len() as i64;
1478 boundaries.push(acc);
1479 }
1480 let boundaries: Arc<[i64]> = boundaries.into();
1481 self.admitted.extend_metadata(&chunks);
1482 let batch = Arc::new(BatchDescriptor {
1483 start,
1484 chunks: Arc::clone(&chunks),
1485 boundaries,
1486 on_attempted: Arc::clone(&on_attempted),
1487 bytes: batch_bytes,
1488 pending_lanes: AtomicUsize::new(reservations.iter().flatten().count()),
1489 packed_groups: std::sync::Mutex::new(BTreeMap::new()),
1490 });
1491 let mut represented_zones = Vec::with_capacity(self.lanes.len());
1492 for (zone, reservation) in reservations.iter_mut().enumerate() {
1493 if let (Some(lane), Some(reservation)) = (&self.lanes[zone], reservation.take()) {
1494 let lane_batch = LaneBatch::new(Arc::clone(&batch), reservation);
1495 if let Err(error) = lane.work.send(lane_batch) {
1496 drop(error.0);
1497 let lane = self.lanes[zone]
1501 .take()
1502 .expect("failed lane send still owns its task handle");
1503 let _ = lane.done.await;
1504 self.commits.finish_lane(zone, None);
1505 } else {
1506 represented_zones.push(zone);
1507 }
1508 }
1509 }
1510 self.admitted.queue_digest(chunks);
1514 Ok(self
1515 .commits
1516 .admit_window(&batch.boundaries, &represented_zones))
1517 }
1518
1519 pub async fn seal(&mut self) -> Result<SealReport, ProtocolError> {
1533 self.sealed = true;
1534 let total = self.admitted.len();
1535 if self.admitted.is_empty() || self.is_poisoned() {
1536 self.shutdown_background_tasks().await;
1537 return Err(ProtocolError::Poisoned);
1538 }
1539 let expected_digest = self.seal_digest().await;
1540 let expected_crc32c = self.seal_crc32c();
1541 let quorum = self.quorum();
1542 let lanes = std::mem::take(&mut self.lanes);
1547 let aborts = AbortLanesOnDrop(
1548 lanes
1549 .iter()
1550 .flatten()
1551 .map(|lane| lane.done.abort_handle())
1552 .collect(),
1553 );
1554 let mut draining: FuturesUnordered<_> = lanes
1555 .into_iter()
1556 .enumerate()
1557 .filter_map(|(zone, lane)| {
1558 lane.map(|lane| {
1559 drop(lane.work);
1560 let done = lane.done;
1561 async move { done.await.ok().flatten().map(|token| (zone, token)) }
1562 })
1563 })
1564 .collect();
1565 let result = async {
1566 let mut settling = self.commits.subscribe();
1570 loop {
1571 let snapshot = settling.borrow_and_update().clone();
1572 if snapshot.failure.is_some() {
1573 return Err(ProtocolError::Poisoned);
1574 }
1575 if snapshot.committed >= total {
1576 break;
1577 }
1578 settling
1579 .changed()
1580 .await
1581 .map_err(|_| ProtocolError::PipelineClosed)?;
1582 }
1583 let mut drained = Vec::new();
1594 while drained.len() < quorum {
1595 match draining.next().await {
1596 Some(Some((zone, token))) => drained.push((zone, token)),
1597 Some(None) => {}
1598 None => break,
1599 }
1600 }
1601 if !draining.is_empty() {
1602 let grace = tokio::time::sleep(LANE_DRAIN_GRACE);
1603 tokio::pin!(grace);
1604 loop {
1605 tokio::select! {
1606 biased;
1607 next = draining.next() => match next {
1608 Some(Some((zone, token))) => drained.push((zone, token)),
1609 Some(None) => {}
1610 None => break,
1611 },
1612 () = &mut grace => break,
1613 }
1614 }
1615 }
1616 let write_offset = self.physical_size() as i64;
1617 let finalizations = join_all(drained.into_iter().map(|(zone, mut token)| {
1618 let replica = Arc::clone(&self.replicas[zone]);
1619 let config = self.config.clone();
1620 async move {
1621 finalize_with_retry(&replica, &mut token, write_offset, &config)
1622 .await
1623 .map(|snapshot| (zone, snapshot))
1624 }
1625 }))
1626 .await;
1627 let mut finalized = vec![None; self.replicas.len()];
1628 for result in finalizations {
1629 match result {
1630 Ok((zone, snapshot)) => {
1631 finalized[zone] = Some(snapshot);
1632 }
1633 Err(error) => {
1634 tracing::warn!(
1635 zone = error.zone,
1636 code = ?error.code,
1637 %error,
1638 write_offset,
1639 "live segment finalization failed"
1640 );
1641 }
1642 }
1643 }
1644 if finalized.iter().flatten().count() >= self.quorum() {
1645 return Ok(SealReport { finalized });
1646 }
1647 let volume = QuorumVolume::with_metadata(
1648 self.replicas.clone(),
1649 self.config.clone(),
1650 self.metadata.clone(),
1651 Arc::clone(&self.metrics),
1652 )?;
1653 let recovered = volume.recover_for_seal(Some(total)).await?;
1654 let actual_digest = recovered.digest();
1655 if actual_digest != expected_digest {
1656 return Err(ProtocolError::SealDigestMismatch {
1657 expected: expected_digest,
1658 actual: actual_digest,
1659 });
1660 }
1661 let actual_crc32c = recovered.crc32c();
1662 if actual_crc32c != expected_crc32c {
1663 return Err(ProtocolError::SealCrc32cMismatch {
1664 expected: expected_crc32c,
1665 actual: actual_crc32c,
1666 });
1667 }
1668 volume.enforce_seal(recovered.canonical()).await?;
1669 Ok(SealReport::default())
1672 }
1673 .await;
1674 aborts.abort();
1675 while draining.next().await.is_some() {}
1676 self.shutdown_background_tasks().await;
1677 result
1678 }
1679
1680 #[cfg(test)]
1682 pub fn poison(&self) {
1683 self.commits.poison();
1684 }
1685}
1686
1687fn prefer_lower_zone(current: &mut Option<TransportError>, candidate: TransportError) {
1688 if current
1689 .as_ref()
1690 .is_none_or(|existing| candidate.zone < existing.zone)
1691 {
1692 *current = Some(candidate);
1693 }
1694}
1695
1696async fn snapshot_with_retry(
1697 replica: &Arc<dyn Replica>,
1698 config: &ClientConfig,
1699) -> Result<ReplicaSnapshot, TransportError> {
1700 let mut attempt = 0usize;
1701 loop {
1702 match replica.snapshot().await {
1703 Ok(snapshot) => return Ok(snapshot),
1704 Err(error) if error.code.transient() && attempt < config.max_retries => {
1705 retry_sleep(config, attempt).await;
1706 attempt += 1;
1707 }
1708 Err(error) => return Err(error),
1709 }
1710 }
1711}
1712
1713async fn stat_with_retry(
1714 replica: &Arc<dyn Replica>,
1715 config: &ClientConfig,
1716) -> Result<ReplicaSnapshot, TransportError> {
1717 let mut attempt = 0usize;
1718 loop {
1719 match replica.stat().await {
1720 Ok(snapshot) => return Ok(snapshot),
1721 Err(error) if error.code.transient() && attempt < config.max_retries => {
1722 retry_sleep(config, attempt).await;
1723 attempt += 1;
1724 }
1725 Err(error) => return Err(error),
1726 }
1727 }
1728}
1729
1730async fn create_with_retry(
1731 replica: &Arc<dyn Replica>,
1732 metadata: HashMap<String, String>,
1733 config: &ClientConfig,
1734) -> Result<ReplicaSnapshot, TransportError> {
1735 let mut attempt = 0usize;
1736 loop {
1737 match replica.create_appendable(metadata.clone()).await {
1738 Ok(snapshot) => return Ok(snapshot),
1739 Err(error) if error.code.transient() && attempt < config.max_retries => {
1740 retry_sleep(config, attempt).await;
1741 attempt += 1;
1742 }
1743 Err(error) => return Err(error),
1744 }
1745 }
1746}
1747
1748async fn create_session_with_retry(
1749 replica: &Arc<dyn Replica>,
1750 metadata: HashMap<String, String>,
1751 config: &ClientConfig,
1752) -> Result<AppendToken, TransportError> {
1753 let mut attempt = 0usize;
1754 loop {
1755 match replica.create_append_session(metadata.clone()).await {
1756 Ok(token) => return Ok(token),
1757 Err(error) if error.code.transient() && attempt < config.max_retries => {
1758 retry_sleep(config, attempt).await;
1759 attempt += 1;
1760 }
1761 Err(error) => return Err(error),
1762 }
1763 }
1764}
1765
1766async fn takeover_with_retry(
1767 replica: &Arc<dyn Replica>,
1768 observed: &ReplicaSnapshot,
1769 config: &ClientConfig,
1770) -> Result<AppendToken, TransportError> {
1771 let mut attempt = 0usize;
1772 loop {
1773 match replica.takeover(observed).await {
1774 Ok(token) => return Ok(token),
1775 Err(error) if error.code.transient() && attempt < config.max_retries => {
1776 retry_sleep(config, attempt).await;
1777 attempt += 1;
1778 }
1779 Err(error) => return Err(error),
1780 }
1781 }
1782}
1783
1784async fn takeover_current_with_retry(
1785 replica: &Arc<dyn Replica>,
1786 config: &ClientConfig,
1787) -> Result<AppendToken, TransportError> {
1788 let mut attempt = 0usize;
1789 loop {
1790 match replica.takeover_current().await {
1791 Ok(token) => return Ok(token),
1792 Err(error) if error.code.transient() && attempt < config.max_retries => {
1793 retry_sleep(config, attempt).await;
1794 attempt += 1;
1795 }
1796 Err(error) => return Err(error),
1797 }
1798 }
1799}
1800
1801async fn replace_with_retry(
1802 replica: &Arc<dyn Replica>,
1803 mut observed: ReplicaSnapshot,
1804 data: Bytes,
1805 metadata: HashMap<String, String>,
1806 config: &ClientConfig,
1807) -> Result<AppendToken, TransportError> {
1808 let mut attempt = 0usize;
1809 loop {
1810 match replica
1811 .replace_appendable(&observed, data.clone(), metadata.clone())
1812 .await
1813 {
1814 Ok(token) => return Ok(token),
1815 Err(error) if error.code.transient() && attempt < config.max_retries => {
1816 retry_sleep(config, attempt).await;
1817 attempt += 1;
1818 observed = snapshot_with_retry(replica, config).await?;
1819 if observed.bytes == data[..]
1820 && observed.metadata == metadata
1821 && !observed.finalized
1822 {
1823 return Ok(AppendToken {
1824 zone: observed.zone,
1825 generation: Some(observed.generation),
1826 metageneration: Some(observed.metageneration),
1827 persisted_size: data.len() as i64,
1828 write_handle: None,
1829 });
1830 }
1831 }
1832 Err(error) => return Err(error),
1833 }
1834 }
1835}
1836
1837struct BatchDescriptor {
1839 start: i64,
1840 chunks: Arc<[Bytes]>,
1841 boundaries: Arc<[i64]>,
1842 on_attempted: AttemptedBytes,
1843 bytes: usize,
1844 pending_lanes: AtomicUsize,
1845 packed_groups: Mutex<BTreeMap<i64, Arc<OnceLock<Arc<PackedAppend>>>>>,
1846}
1847
1848impl BatchDescriptor {
1849 fn end(&self) -> i64 {
1850 self.boundaries
1851 .last()
1852 .copied()
1853 .expect("an admitted batch is non-empty")
1854 }
1855
1856 fn lane_staged(&self) {
1857 let previous = self.pending_lanes.fetch_sub(1, Ordering::AcqRel);
1858 debug_assert!(previous > 0, "batch lane count underflow");
1859 if previous == 1 {
1860 self.packed_groups
1861 .lock()
1862 .unwrap_or_else(|poisoned| poisoned.into_inner())
1863 .clear();
1864 }
1865 }
1866}
1867
1868struct LaneBatch {
1870 batch: Arc<BatchDescriptor>,
1871 reservation: Option<Arc<LaneReservation>>,
1872 staged: bool,
1873}
1874
1875impl LaneBatch {
1876 fn new(batch: Arc<BatchDescriptor>, reservation: Arc<LaneReservation>) -> Self {
1877 Self {
1878 batch,
1879 reservation: Some(reservation),
1880 staged: false,
1881 }
1882 }
1883
1884 fn into_retained(mut self) -> RetainedBatch {
1885 let retained = RetainedBatch {
1886 batch: Arc::clone(&self.batch),
1887 next_chunk: 0,
1888 _reservation: self
1889 .reservation
1890 .take()
1891 .expect("an unstaged lane batch owns its reservation"),
1892 };
1893 self.staged = true;
1894 self.batch.lane_staged();
1895 retained
1896 }
1897}
1898
1899impl Drop for LaneBatch {
1900 fn drop(&mut self) {
1901 if !self.staged {
1902 self.batch.lane_staged();
1903 }
1904 }
1905}
1906
1907struct RetainedBatch {
1908 batch: Arc<BatchDescriptor>,
1909 next_chunk: usize,
1910 _reservation: Arc<LaneReservation>,
1911}
1912
1913fn packed_group(batches: &[LaneBatch]) -> Arc<PackedAppend> {
1914 let first = &batches.first().expect("coalesced group is non-empty").batch;
1915 let end = batches
1916 .last()
1917 .expect("coalesced group is non-empty")
1918 .batch
1919 .end();
1920 let cell = {
1921 let mut groups = first
1922 .packed_groups
1923 .lock()
1924 .unwrap_or_else(|poisoned| poisoned.into_inner());
1925 Arc::clone(
1926 groups
1927 .entry(end)
1928 .or_insert_with(|| Arc::new(OnceLock::new())),
1929 )
1930 };
1931 Arc::clone(cell.get_or_init(|| {
1932 let chunks = batches
1933 .iter()
1934 .flat_map(|batch| batch.batch.chunks.iter().cloned())
1935 .collect();
1936 Arc::new(pack_append(chunks))
1937 }))
1938}
1939
1940const LANE_DRAIN_GRACE: std::time::Duration = std::time::Duration::from_millis(100);
1945
1946struct AbortLanesOnDrop(Vec<tokio::task::AbortHandle>);
1950
1951impl AbortLanesOnDrop {
1952 fn abort(&self) {
1953 for lane in &self.0 {
1954 lane.abort();
1955 }
1956 }
1957}
1958
1959impl Drop for AbortLanesOnDrop {
1960 fn drop(&mut self) {
1961 self.abort();
1962 }
1963}
1964
1965#[derive(Debug)]
1976enum LaneDeath {
1977 Stalled,
1978 Transport(TransportError),
1979}
1980
1981impl LaneDeath {
1982 fn stalled(metrics: &Metrics) -> Self {
1983 metrics.lane_timeouts.increment();
1984 Self::Stalled
1985 }
1986}
1987
1988struct LaneRuntime {
1989 replica: Arc<dyn Replica>,
1990 token: AppendToken,
1991 config: ClientConfig,
1992 metrics: Arc<Metrics>,
1993 commits: Arc<CommitTracker>,
1994 stall_timeout: Arc<LaneStallTimeout>,
1995 durable: i64,
1996 retained: VecDeque<RetainedBatch>,
1997 attempted: Option<AttemptedBytes>,
1998 monitor_session: bool,
1999 last_progress: tokio::time::Instant,
2000}
2001
2002impl LaneRuntime {
2003 fn new(
2004 replica: Arc<dyn Replica>,
2005 token: AppendToken,
2006 config: ClientConfig,
2007 metrics: Arc<Metrics>,
2008 commits: Arc<CommitTracker>,
2009 stall_timeout: Arc<LaneStallTimeout>,
2010 ) -> Self {
2011 let durable = token.persisted_size;
2012 Self {
2013 replica,
2014 token,
2015 config,
2016 metrics,
2017 commits,
2018 stall_timeout,
2019 durable,
2020 retained: VecDeque::new(),
2021 attempted: None,
2022 monitor_session: true,
2026 last_progress: tokio::time::Instant::now(),
2027 }
2028 }
2029
2030 fn zone(&self) -> usize {
2031 self.token.zone
2032 }
2033
2034 fn stall_deadline(&self) -> tokio::time::Instant {
2035 self.last_progress + self.stall_timeout.get()
2036 }
2037
2038 fn publish_advance(&mut self, change: LaneDurableChange) -> Result<bool, LaneDeath> {
2039 publish_lane_advance(
2040 change,
2041 self.zone(),
2042 &mut self.durable,
2043 &mut self.last_progress,
2044 &mut self.retained,
2045 &self.commits,
2046 )
2047 }
2048
2049 async fn confirm_timeout(&mut self) -> Result<bool, LaneDeath> {
2055 let change = confirm_lane_stall(
2056 &self.replica,
2057 self.durable,
2058 self.stall_deadline(),
2059 &self.metrics,
2060 )
2061 .await?;
2062 let stream_failed = self.publish_advance(change)?;
2063 Ok(stream_failed || !self.retained.is_empty())
2064 }
2065
2066 async fn stage(&mut self, batches: Vec<LaneBatch>) -> Result<bool, LaneDeath> {
2067 match tokio::time::timeout_at(
2068 self.stall_deadline(),
2069 stage_group(
2070 &self.replica,
2071 batches,
2072 &mut self.attempted,
2073 &mut self.retained,
2074 ),
2075 )
2076 .await
2077 {
2078 Ok(failed) => Ok(failed),
2079 Err(_) => self.confirm_timeout().await,
2080 }
2081 }
2082}
2083
2084async fn run_lane(
2085 replica: Arc<dyn Replica>,
2086 token: AppendToken,
2087 config: ClientConfig,
2088 metrics: Arc<Metrics>,
2089 mut work: tokio::sync::mpsc::UnboundedReceiver<LaneBatch>,
2090 commits: Arc<CommitTracker>,
2091 stall_timeout: Arc<LaneStallTimeout>,
2092) -> Option<AppendToken> {
2093 let mut lane = LaneRuntime::new(replica, token, config, metrics, commits, stall_timeout);
2094 let mut closed = false;
2095 let death: Option<LaneDeath> = loop {
2096 if closed && lane.retained.is_empty() {
2097 break None;
2098 }
2099 tokio::select! {
2100 biased;
2101 changed = lane_progress(
2109 &lane.replica,
2110 lane.durable,
2111 !lane.retained.is_empty(),
2112 lane.stall_deadline(),
2113 &lane.metrics,
2114 ), if lane.monitor_session || !lane.retained.is_empty() => {
2115 match changed {
2116 LaneProgress::Advanced(change) => {
2117 match lane.publish_advance(change) {
2118 Ok(false) => lane.monitor_session = true,
2119 Ok(true) => {
2120 if let Err(error) = lane.recover().await {
2121 break Some(error);
2122 }
2123 lane.monitor_session = true;
2124 }
2125 Err(error) => break Some(error),
2126 }
2127 }
2128 LaneProgress::Stalled => {
2129 break Some(LaneDeath::Stalled);
2130 }
2131 LaneProgress::Failed(error) if !error.code.transient() => {
2132 break Some(LaneDeath::Transport(error));
2133 }
2134 LaneProgress::Failed(_) => {
2135 lane.monitor_session = false;
2136 if lane.retained.is_empty() {
2137 continue;
2138 }
2139 if let Err(error) = lane.recover().await {
2140 break Some(error);
2141 }
2142 lane.monitor_session = true;
2143 }
2144 }
2145 }
2146 batch = work.recv(), if !closed => match batch {
2147 Some(batch) => {
2148 let mut batches = vec![batch];
2153 let queued = work.len();
2154 for _ in 0..queued {
2155 match work.try_recv() {
2156 Ok(batch) => batches.push(batch),
2157 Err(tokio::sync::mpsc::error::TryRecvError::Empty) => break,
2158 Err(tokio::sync::mpsc::error::TryRecvError::Disconnected) => {
2159 closed = true;
2160 break;
2161 }
2162 }
2163 }
2164 if lane.retained.is_empty() {
2165 lane.last_progress = tokio::time::Instant::now();
2166 }
2167 let failed = match lane.stage(batches).await {
2168 Ok(failed) => failed,
2169 Err(error) => break Some(error),
2170 };
2171 if failed {
2172 if let Err(error) = lane.recover().await {
2173 break Some(error);
2174 }
2175 }
2176 lane.monitor_session = true;
2177 }
2178 None => closed = true,
2179 },
2180 }
2181 };
2182 match death {
2183 None => {
2184 lane.commits.finish_lane(lane.zone(), None);
2185 lane.token.persisted_size = lane.durable;
2186 Some(lane.token)
2187 }
2188 Some(LaneDeath::Stalled) => {
2189 tracing::warn!(
2190 zone = lane.zone(),
2191 durable_offset = lane.durable,
2192 retained_chunks = retained_chunk_count(&lane.retained),
2193 stall_timeout_ms = lane.stall_timeout.get().as_millis(),
2194 "append lane shed after making no durable progress"
2195 );
2196 lane.commits.finish_lane(lane.zone(), None);
2197 lane.retained.clear();
2198 work.close();
2199 while work.try_recv().is_ok() {}
2200 None
2201 }
2202 Some(LaneDeath::Transport(error)) => {
2203 tracing::warn!(
2204 zone = error.zone,
2205 code = ?error.code,
2206 error = %error,
2207 durable_offset = lane.durable,
2208 retained_chunks = retained_chunk_count(&lane.retained),
2209 "append lane died"
2210 );
2211 tracing::debug!(
2212 zone = error.zone,
2213 code = ?error.code,
2214 message = error.message.as_str(),
2215 durable_offset = lane.durable,
2216 retained_chunks = retained_chunk_count(&lane.retained),
2217 "append lane dropped"
2218 );
2219 lane.commits.finish_lane(lane.zone(), Some(error));
2220 lane.retained.clear();
2221 work.close();
2222 while work.try_recv().is_ok() {}
2223 None
2224 }
2225 }
2226}
2227
2228async fn stage_group(
2233 replica: &Arc<dyn Replica>,
2234 batches: Vec<LaneBatch>,
2235 attempted: &mut Option<AttemptedBytes>,
2236 retained: &mut VecDeque<RetainedBatch>,
2237) -> bool {
2238 let group_start = batches
2239 .first()
2240 .expect("coalesced group is non-empty")
2241 .batch
2242 .start;
2243 let packed = packed_group(&batches);
2244 for batch in batches {
2245 (batch.batch.on_attempted)(batch.batch.bytes as u64);
2246 *attempted = Some(Arc::clone(&batch.batch.on_attempted));
2247 retained.push_back(batch.into_retained());
2248 }
2249 replica
2250 .lane_send_packed(group_start, &packed)
2251 .await
2252 .is_err()
2253}
2254
2255fn ack_through(durable: i64, retained: &mut VecDeque<RetainedBatch>) {
2259 while let Some(batch) = retained.front_mut() {
2260 batch.next_chunk = batch
2261 .batch
2262 .boundaries
2263 .partition_point(|boundary| *boundary <= durable)
2264 .max(batch.next_chunk);
2265 if batch.next_chunk == batch.batch.chunks.len() {
2266 retained.pop_front();
2267 } else {
2268 break;
2269 }
2270 }
2271}
2272
2273fn retained_chunk_count(retained: &VecDeque<RetainedBatch>) -> usize {
2274 retained
2275 .iter()
2276 .map(|batch| batch.batch.chunks.len() - batch.next_chunk)
2277 .sum()
2278}
2279
2280enum LaneProgress {
2284 Advanced(LaneDurableChange),
2285 Stalled,
2286 Failed(TransportError),
2287}
2288
2289async fn lane_progress(
2290 replica: &Arc<dyn Replica>,
2291 seen: i64,
2292 active: bool,
2293 stall_deadline: tokio::time::Instant,
2294 metrics: &Metrics,
2295) -> LaneProgress {
2296 loop {
2297 if !active {
2298 match replica.lane_durable_change(seen).await {
2299 Ok(change) if change.persisted_size > seen => {
2300 return LaneProgress::Advanced(change);
2301 }
2302 Ok(_) => tokio::task::yield_now().await,
2303 Err(error) => return LaneProgress::Failed(error),
2304 }
2305 continue;
2306 }
2307 tokio::select! {
2308 biased;
2309 result = replica.lane_durable_change(seen) => match result {
2310 Ok(change) if change.persisted_size > seen => {
2311 return LaneProgress::Advanced(change);
2312 }
2313 Ok(_) if tokio::time::Instant::now() < stall_deadline => {
2314 tokio::task::yield_now().await;
2315 }
2316 Ok(_) => {
2317 let _ = LaneDeath::stalled(metrics);
2318 return LaneProgress::Stalled;
2319 }
2320 Err(error) => return LaneProgress::Failed(error),
2321 },
2322 _ = tokio::time::sleep_until(stall_deadline) => {
2323 let _ = LaneDeath::stalled(metrics);
2324 return LaneProgress::Stalled;
2325 }
2326 }
2327 }
2328}
2329
2330fn publish_lane_progress(
2331 tail: i64,
2332 zone: usize,
2333 durable: &mut i64,
2334 last_progress: &mut tokio::time::Instant,
2335 retained: &mut VecDeque<RetainedBatch>,
2336 commits: &CommitTracker,
2337) {
2338 let previous = *durable;
2339 *durable = (*durable).max(tail);
2340 if *durable > previous {
2341 *last_progress = tokio::time::Instant::now();
2342 }
2343 ack_through(*durable, retained);
2344 commits.publish_durable(zone, *durable);
2345}
2346
2347fn publish_lane_advance(
2350 change: LaneDurableChange,
2351 zone: usize,
2352 durable: &mut i64,
2353 last_progress: &mut tokio::time::Instant,
2354 retained: &mut VecDeque<RetainedBatch>,
2355 commits: &CommitTracker,
2356) -> Result<bool, LaneDeath> {
2357 publish_lane_progress(
2358 change.persisted_size,
2359 zone,
2360 durable,
2361 last_progress,
2362 retained,
2363 commits,
2364 );
2365 match change.error {
2366 None => Ok(false),
2367 Some(error) if error.code.transient() => Ok(true),
2368 Some(error) => Err(LaneDeath::Transport(error)),
2369 }
2370}
2371
2372async fn confirm_lane_stall(
2373 replica: &Arc<dyn Replica>,
2374 seen: i64,
2375 stall_deadline: tokio::time::Instant,
2376 metrics: &Metrics,
2377) -> Result<LaneDurableChange, LaneDeath> {
2378 match lane_progress(replica, seen, true, stall_deadline, metrics).await {
2379 LaneProgress::Advanced(change) => Ok(change),
2380 LaneProgress::Stalled => Err(LaneDeath::Stalled),
2381 LaneProgress::Failed(error) if error.code.transient() => Err(LaneDeath::stalled(metrics)),
2382 LaneProgress::Failed(error) => Err(LaneDeath::Transport(error)),
2383 }
2384}
2385
2386impl LaneRuntime {
2387 async fn recover(&mut self) -> Result<(), LaneDeath> {
2393 let mut attempt = 0usize;
2394 loop {
2395 self.metrics.lane_retries.increment();
2396 let deadline = self.stall_deadline();
2397 let resumed =
2398 tokio::time::timeout_at(deadline, self.replica.resume_tail(&mut self.token)).await;
2399 let resumed = match resumed {
2400 Ok(resumed) => resumed,
2401 Err(_) => {
2402 if self.confirm_timeout().await? {
2403 continue;
2404 }
2405 return Ok(());
2406 }
2407 };
2408 match resumed {
2409 Ok(tail) => {
2410 publish_lane_progress(
2411 tail,
2412 self.zone(),
2413 &mut self.durable,
2414 &mut self.last_progress,
2415 &mut self.retained,
2416 &self.commits,
2417 );
2418 let mut suffix = Vec::with_capacity(retained_chunk_count(&self.retained));
2419 for retained_batch in &self.retained {
2420 for index in retained_batch.next_chunk..retained_batch.batch.chunks.len() {
2421 let chunk = &retained_batch.batch.chunks[index];
2422 if suffix.is_empty() {
2423 let offset = index
2424 .checked_sub(1)
2425 .and_then(|previous| {
2426 retained_batch.batch.boundaries.get(previous).copied()
2427 })
2428 .unwrap_or(retained_batch.batch.start);
2429 let skip = usize::try_from((self.durable - offset).max(0))
2430 .unwrap_or(chunk.len());
2431 if skip >= chunk.len() {
2432 continue;
2433 }
2434 suffix.push(chunk.slice(skip..));
2435 } else {
2436 suffix.push(chunk.clone());
2437 }
2438 }
2439 }
2440 if suffix.is_empty() {
2441 return Ok(());
2442 }
2443 let packed = pack_append(suffix);
2444 let resend_bytes = packed.len();
2445 tracing::debug!(
2446 zone = self.zone(),
2447 durable_offset = self.durable,
2448 chunks = packed.chunks().len(),
2449 bytes = resend_bytes,
2450 attempt,
2451 "resending append lane batch after recovery"
2452 );
2453 if let Some(attempted) = &self.attempted {
2454 attempted(resend_bytes as u64);
2455 }
2456 let sent = tokio::time::timeout_at(
2457 self.stall_deadline(),
2458 self.replica.lane_send_packed(self.durable, &packed),
2459 )
2460 .await;
2461 let sent = match sent {
2462 Ok(sent) => sent,
2463 Err(_) => {
2464 if self.confirm_timeout().await? {
2465 continue;
2466 }
2467 return Ok(());
2468 }
2469 };
2470 match sent {
2471 Ok(()) => return Ok(()),
2472 Err(error)
2473 if error.code.transient() && attempt < self.config.max_retries =>
2474 {
2475 if tokio::time::timeout_at(
2476 self.stall_deadline(),
2477 retry_sleep(&self.config, attempt),
2478 )
2479 .await
2480 .is_err()
2481 {
2482 if self.confirm_timeout().await? {
2483 continue;
2484 }
2485 return Ok(());
2486 }
2487 attempt += 1;
2488 }
2489 Err(error) => return Err(LaneDeath::Transport(error)),
2490 }
2491 }
2492 Err(error) if error.code.transient() && attempt < self.config.max_retries => {
2493 if tokio::time::timeout_at(
2494 self.stall_deadline(),
2495 retry_sleep(&self.config, attempt),
2496 )
2497 .await
2498 .is_err()
2499 {
2500 if self.confirm_timeout().await? {
2501 continue;
2502 }
2503 return Ok(());
2504 }
2505 attempt += 1;
2506 }
2507 Err(error) => return Err(LaneDeath::Transport(error)),
2508 }
2509 }
2510 }
2511}
2512
2513async fn finalize_with_retry(
2514 replica: &Arc<dyn Replica>,
2515 token: &mut AppendToken,
2516 write_offset: i64,
2517 config: &ClientConfig,
2518) -> Result<ReplicaSnapshot, TransportError> {
2519 let mut attempt = 0usize;
2520 loop {
2521 match replica.finalize(token, write_offset).await {
2522 Ok(snapshot) => return Ok(snapshot),
2523 Err(error)
2524 if error.code == TransportCode::FailedPrecondition || error.code.transient() =>
2525 {
2526 let Some(generation) = token.generation else {
2530 return Err(error);
2531 };
2532 if error.code.transient() {
2533 if attempt >= config.max_retries {
2534 return Err(error);
2535 }
2536 retry_sleep(config, attempt).await;
2537 attempt += 1;
2538 }
2539 let snapshot = stat_with_retry(replica, config).await?;
2544 if snapshot.finalized
2545 && snapshot.generation == generation
2546 && snapshot.persisted_size == write_offset
2547 {
2548 return Ok(snapshot);
2549 }
2550 if error.code == TransportCode::FailedPrecondition {
2551 return Err(error);
2552 }
2553 }
2554 Err(error) => return Err(error),
2555 }
2556 }
2557}
2558
2559pub(crate) fn retry_delay(config: &ClientConfig, attempt: usize) -> Duration {
2560 let multiplier = 1u32.checked_shl(attempt.min(16) as u32).unwrap_or(u32::MAX);
2561 config.retry_base.saturating_mul(multiplier)
2562}
2563
2564pub(crate) async fn retry_sleep(config: &ClientConfig, attempt: usize) {
2565 tokio::time::sleep(retry_delay(config, attempt)).await;
2566}
2567
2568pub(crate) fn canonical_prefix(
2583 snapshots: &[ReplicaSnapshot],
2584 quorum: usize,
2585) -> Result<Vec<RecordFrame>, ProtocolError> {
2586 select_canonical_quorum(snapshots, quorum).map(|(prefix, _)| prefix.into_records())
2587}
2588
2589fn quorum_subsets(count: usize, quorum: usize) -> Vec<Vec<usize>> {
2592 fn extend(
2593 start: usize,
2594 count: usize,
2595 quorum: usize,
2596 current: &mut Vec<usize>,
2597 subsets: &mut Vec<Vec<usize>>,
2598 ) {
2599 if current.len() == quorum {
2600 subsets.push(current.clone());
2601 return;
2602 }
2603 for index in start..count {
2604 current.push(index);
2605 extend(index + 1, count, quorum, current, subsets);
2606 current.pop();
2607 }
2608 }
2609 let mut subsets = Vec::new();
2610 extend(
2611 0,
2612 count,
2613 quorum,
2614 &mut Vec::with_capacity(quorum),
2615 &mut subsets,
2616 );
2617 subsets
2618}
2619
2620fn select_canonical_quorum(
2621 snapshots: &[ReplicaSnapshot],
2622 quorum: usize,
2623) -> Result<(CanonicalPrefix, Vec<ReplicaSnapshot>), ProtocolError> {
2624 if snapshots.len() < quorum {
2625 return Err(ProtocolError::NoQuorum);
2626 }
2627 let decoded: Vec<_> = snapshots
2628 .iter()
2629 .map(CanonicalPrefix::from_snapshot)
2630 .collect();
2631 let mut conflicts = vec![false; decoded.len() * decoded.len()];
2632 let mut first_conflict = None;
2633 for left in 0..decoded.len() {
2634 for right in left + 1..decoded.len() {
2635 let overlap = decoded[left].len().min(decoded[right].len());
2636 let conflict = (0..overlap).find(|index| {
2637 decoded[left].record_bytes(*index) != decoded[right].record_bytes(*index)
2638 });
2639 if let Some(index) = conflict {
2640 first_conflict.get_or_insert(index);
2641 conflicts[left * decoded.len() + right] = true;
2642 }
2643 }
2644 }
2645 let mut candidates = Vec::new();
2646 for subset in quorum_subsets(decoded.len(), quorum) {
2647 let consistent = subset.iter().enumerate().all(|(position, &left)| {
2648 subset[position + 1..]
2649 .iter()
2650 .all(|&right| !conflicts[left * decoded.len() + right])
2651 });
2652 if !consistent {
2653 continue;
2654 }
2655 let mut longest_member = subset[0];
2656 for &member in &subset[1..] {
2657 if decoded[member].len() > decoded[longest_member].len() {
2658 longest_member = member;
2659 }
2660 }
2661 candidates.push((
2662 decoded[longest_member].clone(),
2663 subset
2664 .iter()
2665 .map(|&member| snapshots[member].clone())
2666 .collect::<Vec<_>>(),
2667 ));
2668 }
2669 let longest = candidates
2670 .iter()
2671 .map(|(candidate, _)| candidate.len())
2672 .max()
2673 .ok_or(ProtocolError::ConflictingPrefix {
2674 record_index: first_conflict.unwrap_or(0),
2675 })?;
2676 let mut longest_candidates = candidates
2677 .into_iter()
2678 .filter(|(candidate, _)| candidate.len() == longest);
2679 let first = longest_candidates
2680 .next()
2681 .expect("longest length came from a candidate");
2682 let mut equivalent = vec![first];
2683 for candidate in longest_candidates {
2684 if candidate.0.bytes != equivalent[0].0.bytes {
2685 let record_index = (0..candidate.0.len())
2686 .find(|index| {
2687 candidate.0.record_bytes(*index) != equivalent[0].0.record_bytes(*index)
2688 })
2689 .unwrap_or(0);
2690 return Err(ProtocolError::ConflictingPrefix { record_index });
2691 }
2692 equivalent.push(candidate);
2693 }
2694 Ok(equivalent
2695 .into_iter()
2696 .max_by(|(_, left_witnesses), (_, right_witnesses)| {
2697 let left_zones: Vec<_> = left_witnesses.iter().map(|copy| copy.zone).collect();
2698 let right_zones: Vec<_> = right_witnesses.iter().map(|copy| copy.zone).collect();
2699 right_zones.cmp(&left_zones)
2700 })
2701 .expect("at least one equivalent candidate remains"))
2702}
2703
2704pub(crate) fn protocol_metadata() -> HashMap<String, String> {
2705 HashMap::from([(META_FORMAT.to_string(), FORMAT_VERSION.to_string())])
2706}
2707
2708pub(crate) fn valid_format(metadata: &HashMap<String, String>) -> bool {
2709 metadata.get(META_FORMAT).map(String::as_str) == Some(FORMAT_VERSION)
2710}
2711
2712#[cfg(test)]
2713fn encode_records(records: &[RecordFrame]) -> Result<Vec<u8>, RecordError> {
2714 let encoded: Result<Vec<_>, _> = records.iter().map(RecordFrame::encode).collect();
2715 Ok(encoded?.into_iter().flatten().collect())
2716}
2717
2718#[cfg(test)]
2719mod tests {
2720 use super::*;
2721 use std::collections::{HashMap, VecDeque};
2722 use std::sync::Arc;
2723
2724 use async_trait::async_trait;
2725 use tokio::sync::{mpsc, oneshot, watch, Mutex};
2726
2727 use crate::metrics::{test_support::TestMetricsRecorder, Metrics};
2728
2729 fn snapshot(zone: usize, records: &[RecordFrame]) -> ReplicaSnapshot {
2730 let bytes = encode_records(records).unwrap();
2731 ReplicaSnapshot {
2732 zone,
2733 generation: 1,
2734 metageneration: 1,
2735 persisted_size: 0,
2736 finalized: false,
2737 crc32c: Some(crc32c::crc32c(&bytes)),
2738 metadata: protocol_metadata(),
2739 bytes,
2740 }
2741 }
2742
2743 fn record(value: &[u8]) -> RecordFrame {
2744 RecordFrame {
2745 payload: Bytes::copy_from_slice(value),
2746 }
2747 }
2748
2749 fn batch_descriptor(
2750 start: i64,
2751 chunks: Vec<Bytes>,
2752 pending_lanes: usize,
2753 ) -> Arc<BatchDescriptor> {
2754 let mut end = start;
2755 let mut boundaries = Vec::with_capacity(chunks.len());
2756 for chunk in &chunks {
2757 end += chunk.len() as i64;
2758 boundaries.push(end);
2759 }
2760 let bytes = chunks.iter().map(Bytes::len).sum();
2761 Arc::new(BatchDescriptor {
2762 start,
2763 chunks: chunks.into(),
2764 boundaries: boundaries.into(),
2765 on_attempted: Arc::new(|_| {}),
2766 bytes,
2767 pending_lanes: AtomicUsize::new(pending_lanes),
2768 packed_groups: std::sync::Mutex::new(BTreeMap::new()),
2769 })
2770 }
2771
2772 struct ScriptedLaneReplica {
2773 zone: usize,
2774 durable: watch::Sender<i64>,
2775 send_releases: Mutex<VecDeque<oneshot::Receiver<()>>>,
2776 sends: mpsc::UnboundedSender<i64>,
2777 }
2778
2779 impl ScriptedLaneReplica {
2780 fn new(
2781 zone: usize,
2782 send_releases: VecDeque<oneshot::Receiver<()>>,
2783 sends: mpsc::UnboundedSender<i64>,
2784 ) -> Self {
2785 let (durable, _) = watch::channel(0);
2786 Self {
2787 zone,
2788 durable,
2789 send_releases: Mutex::new(send_releases),
2790 sends,
2791 }
2792 }
2793
2794 fn error(&self, code: TransportCode, message: &str) -> TransportError {
2795 TransportError {
2796 zone: self.zone,
2797 code,
2798 message: message.into(),
2799 }
2800 }
2801 }
2802
2803 struct ReaderTerminalReplica {
2804 zone: usize,
2805 reader_failed: watch::Sender<bool>,
2806 resume_calls: AtomicUsize,
2807 }
2808
2809 struct StalledReplica {
2810 zone: usize,
2811 }
2812
2813 impl ReaderTerminalReplica {
2814 fn new(zone: usize) -> Self {
2815 let (reader_failed, _) = watch::channel(false);
2816 Self {
2817 zone,
2818 reader_failed,
2819 resume_calls: AtomicUsize::new(0),
2820 }
2821 }
2822
2823 fn error(&self) -> TransportError {
2824 TransportError {
2825 zone: self.zone,
2826 code: TransportCode::PermissionDenied,
2827 message: "async reader rejected append".into(),
2828 }
2829 }
2830
2831 fn resume_calls(&self) -> usize {
2832 self.resume_calls.load(Ordering::SeqCst)
2833 }
2834 }
2835
2836 struct SingleReplicaFactory {
2837 replica: Arc<dyn Replica>,
2838 }
2839
2840 impl SingleReplicaFactory {
2841 fn new(replica: Arc<dyn Replica>) -> Self {
2842 Self { replica }
2843 }
2844 }
2845
2846 #[async_trait]
2847 impl Replica for ScriptedLaneReplica {
2848 async fn snapshot(&self) -> Result<ReplicaSnapshot, TransportError> {
2849 panic!("snapshot is not used in this test")
2850 }
2851
2852 async fn stat(&self) -> Result<ReplicaSnapshot, TransportError> {
2853 panic!("stat is not used in this test")
2854 }
2855
2856 async fn create_appendable(
2857 &self,
2858 _metadata: HashMap<String, String>,
2859 ) -> Result<ReplicaSnapshot, TransportError> {
2860 panic!("create_appendable is not used in this test")
2861 }
2862
2863 async fn create_append_session(
2864 &self,
2865 _metadata: HashMap<String, String>,
2866 ) -> Result<AppendToken, TransportError> {
2867 panic!("create_append_session is not used in this test")
2868 }
2869
2870 async fn create_register(
2871 &self,
2872 _metadata: HashMap<String, String>,
2873 ) -> Result<ReplicaSnapshot, TransportError> {
2874 panic!("create_register is not used in this test")
2875 }
2876
2877 async fn update_register(
2878 &self,
2879 _metageneration: i64,
2880 _metadata: HashMap<String, String>,
2881 ) -> Result<ReplicaSnapshot, TransportError> {
2882 panic!("update_register is not used in this test")
2883 }
2884
2885 async fn resume_tail(&self, _token: &mut AppendToken) -> Result<i64, TransportError> {
2886 Err(self.error(
2887 TransportCode::Internal,
2888 "resume_tail should not run in this test",
2889 ))
2890 }
2891
2892 async fn takeover(
2893 &self,
2894 _observed: &ReplicaSnapshot,
2895 ) -> Result<AppendToken, TransportError> {
2896 panic!("takeover is not used in this test")
2897 }
2898
2899 async fn replace_appendable(
2900 &self,
2901 _observed: &ReplicaSnapshot,
2902 _data: Bytes,
2903 _metadata: HashMap<String, String>,
2904 ) -> Result<AppendToken, TransportError> {
2905 panic!("replace_appendable is not used in this test")
2906 }
2907
2908 async fn append(
2909 &self,
2910 _token: &AppendToken,
2911 _write_offset: i64,
2912 _data: Vec<u8>,
2913 ) -> Result<i64, TransportError> {
2914 panic!("append is not used in this test")
2915 }
2916
2917 async fn lane_send(
2918 &self,
2919 write_offset: i64,
2920 chunks: &[Bytes],
2921 ) -> Result<(), TransportError> {
2922 let end = write_offset + chunks.iter().map(|chunk| chunk.len() as i64).sum::<i64>();
2923 self.durable.send_replace(end);
2924 let _ = self.sends.send(end);
2925 let release = self.send_releases.lock().await.pop_front();
2926 if let Some(release) = release {
2927 let _ = release.await;
2928 }
2929 Ok(())
2930 }
2931
2932 async fn lane_durable_change(
2933 &self,
2934 seen: i64,
2935 ) -> Result<LaneDurableChange, TransportError> {
2936 let mut durable = self.durable.subscribe();
2937 loop {
2938 let current = *durable.borrow_and_update();
2939 if current > seen {
2940 return Ok(LaneDurableChange {
2941 persisted_size: current,
2942 error: None,
2943 });
2944 }
2945 durable
2946 .changed()
2947 .await
2948 .map_err(|_| self.error(TransportCode::Unavailable, "durable watch closed"))?;
2949 }
2950 }
2951
2952 async fn delete(&self, _generation: i64) -> Result<(), TransportError> {
2953 panic!("delete is not used in this test")
2954 }
2955
2956 async fn finalize(
2957 &self,
2958 _token: &mut AppendToken,
2959 _write_offset: i64,
2960 ) -> Result<ReplicaSnapshot, TransportError> {
2961 panic!("finalize is not used in this test")
2962 }
2963 }
2964
2965 #[async_trait]
2966 impl Replica for ReaderTerminalReplica {
2967 async fn snapshot(&self) -> Result<ReplicaSnapshot, TransportError> {
2968 panic!("snapshot is not used in this test")
2969 }
2970
2971 async fn stat(&self) -> Result<ReplicaSnapshot, TransportError> {
2972 panic!("stat is not used in this test")
2973 }
2974
2975 async fn create_appendable(
2976 &self,
2977 _metadata: HashMap<String, String>,
2978 ) -> Result<ReplicaSnapshot, TransportError> {
2979 panic!("create_appendable is not used in this test")
2980 }
2981
2982 async fn create_append_session(
2983 &self,
2984 _metadata: HashMap<String, String>,
2985 ) -> Result<AppendToken, TransportError> {
2986 Ok(AppendToken {
2987 zone: self.zone,
2988 generation: Some(1),
2989 metageneration: Some(1),
2990 persisted_size: 0,
2991 write_handle: None,
2992 })
2993 }
2994
2995 async fn create_register(
2996 &self,
2997 _metadata: HashMap<String, String>,
2998 ) -> Result<ReplicaSnapshot, TransportError> {
2999 panic!("create_register is not used in this test")
3000 }
3001
3002 async fn update_register(
3003 &self,
3004 _metageneration: i64,
3005 _metadata: HashMap<String, String>,
3006 ) -> Result<ReplicaSnapshot, TransportError> {
3007 panic!("update_register is not used in this test")
3008 }
3009
3010 async fn resume_tail(&self, _token: &mut AppendToken) -> Result<i64, TransportError> {
3011 self.resume_calls.fetch_add(1, Ordering::SeqCst);
3012 Err(TransportError {
3013 zone: self.zone,
3014 code: TransportCode::Unavailable,
3015 message: "recovery must not run after a terminal reader error".into(),
3016 })
3017 }
3018
3019 async fn takeover(
3020 &self,
3021 _observed: &ReplicaSnapshot,
3022 ) -> Result<AppendToken, TransportError> {
3023 panic!("takeover is not used in this test")
3024 }
3025
3026 async fn replace_appendable(
3027 &self,
3028 _observed: &ReplicaSnapshot,
3029 _data: Bytes,
3030 _metadata: HashMap<String, String>,
3031 ) -> Result<AppendToken, TransportError> {
3032 panic!("replace_appendable is not used in this test")
3033 }
3034
3035 async fn append(
3036 &self,
3037 _token: &AppendToken,
3038 _write_offset: i64,
3039 _data: Vec<u8>,
3040 ) -> Result<i64, TransportError> {
3041 panic!("append is not used in this test")
3042 }
3043
3044 async fn lane_send(
3045 &self,
3046 _write_offset: i64,
3047 _chunks: &[Bytes],
3048 ) -> Result<(), TransportError> {
3049 self.reader_failed.send_replace(true);
3050 Ok(())
3051 }
3052
3053 async fn lane_durable_change(
3054 &self,
3055 _seen: i64,
3056 ) -> Result<LaneDurableChange, TransportError> {
3057 let mut reader_failed = self.reader_failed.subscribe();
3058 loop {
3059 if *reader_failed.borrow_and_update() {
3060 return Err(self.error());
3061 }
3062 reader_failed.changed().await.map_err(|_| TransportError {
3063 zone: self.zone,
3064 code: TransportCode::Unavailable,
3065 message: "reader failure watch closed".into(),
3066 })?;
3067 }
3068 }
3069
3070 async fn delete(&self, _generation: i64) -> Result<(), TransportError> {
3071 panic!("delete is not used in this test")
3072 }
3073
3074 async fn finalize(
3075 &self,
3076 _token: &mut AppendToken,
3077 _write_offset: i64,
3078 ) -> Result<ReplicaSnapshot, TransportError> {
3079 panic!("finalize is not used in this test")
3080 }
3081 }
3082
3083 #[async_trait]
3084 impl Replica for StalledReplica {
3085 async fn snapshot(&self) -> Result<ReplicaSnapshot, TransportError> {
3086 panic!("snapshot is not used in this test")
3087 }
3088
3089 async fn stat(&self) -> Result<ReplicaSnapshot, TransportError> {
3090 panic!("stat is not used in this test")
3091 }
3092
3093 async fn create_appendable(
3094 &self,
3095 _metadata: HashMap<String, String>,
3096 ) -> Result<ReplicaSnapshot, TransportError> {
3097 panic!("create_appendable is not used in this test")
3098 }
3099
3100 async fn create_append_session(
3101 &self,
3102 _metadata: HashMap<String, String>,
3103 ) -> Result<AppendToken, TransportError> {
3104 Ok(AppendToken {
3105 zone: self.zone,
3106 generation: Some(1),
3107 metageneration: Some(1),
3108 persisted_size: 0,
3109 write_handle: None,
3110 })
3111 }
3112
3113 async fn create_register(
3114 &self,
3115 _metadata: HashMap<String, String>,
3116 ) -> Result<ReplicaSnapshot, TransportError> {
3117 panic!("create_register is not used in this test")
3118 }
3119
3120 async fn update_register(
3121 &self,
3122 _metageneration: i64,
3123 _metadata: HashMap<String, String>,
3124 ) -> Result<ReplicaSnapshot, TransportError> {
3125 panic!("update_register is not used in this test")
3126 }
3127
3128 async fn resume_tail(&self, _token: &mut AppendToken) -> Result<i64, TransportError> {
3129 panic!("a no-progress timeout must shed instead of recovering the lane")
3130 }
3131
3132 async fn takeover(
3133 &self,
3134 _observed: &ReplicaSnapshot,
3135 ) -> Result<AppendToken, TransportError> {
3136 panic!("takeover is not used in this test")
3137 }
3138
3139 async fn replace_appendable(
3140 &self,
3141 _observed: &ReplicaSnapshot,
3142 _data: Bytes,
3143 _metadata: HashMap<String, String>,
3144 ) -> Result<AppendToken, TransportError> {
3145 panic!("replace_appendable is not used in this test")
3146 }
3147
3148 async fn append(
3149 &self,
3150 _token: &AppendToken,
3151 _write_offset: i64,
3152 _data: Vec<u8>,
3153 ) -> Result<i64, TransportError> {
3154 panic!("append is not used in this test")
3155 }
3156
3157 async fn lane_send(
3158 &self,
3159 _write_offset: i64,
3160 _chunks: &[Bytes],
3161 ) -> Result<(), TransportError> {
3162 Ok(())
3163 }
3164
3165 async fn lane_durable_change(
3166 &self,
3167 _seen: i64,
3168 ) -> Result<LaneDurableChange, TransportError> {
3169 std::future::pending().await
3170 }
3171
3172 async fn delete(&self, _generation: i64) -> Result<(), TransportError> {
3173 panic!("delete is not used in this test")
3174 }
3175
3176 async fn finalize(
3177 &self,
3178 _token: &mut AppendToken,
3179 _write_offset: i64,
3180 ) -> Result<ReplicaSnapshot, TransportError> {
3181 panic!("finalize is not used in this test")
3182 }
3183 }
3184
3185 #[async_trait]
3186 impl crate::transport::ReplicaFactory for SingleReplicaFactory {
3187 fn bucket_name(&self) -> &str {
3188 "single-replica"
3189 }
3190
3191 fn replica(&self, _object: &str) -> Arc<dyn Replica> {
3192 self.replica.clone()
3193 }
3194
3195 async fn list(
3196 &self,
3197 _prefix: &str,
3198 ) -> Result<Vec<crate::transport::ListedObject>, TransportError> {
3199 Ok(Vec::new())
3200 }
3201 }
3202
3203 #[test]
3204 fn majority_matches_supported_widths() {
3205 assert_eq!(majority(1), 1);
3206 assert_eq!(majority(3), 2);
3207 assert_eq!(majority(5), 3);
3208 let lanes = |durables: &[i64]| {
3209 durables
3210 .iter()
3211 .map(|durable| LaneCommitState {
3212 durable: *durable,
3213 ..LaneCommitState::default()
3214 })
3215 .collect::<Vec<_>>()
3216 };
3217 assert_eq!(quorum_durable_watermark(&lanes(&[7]), majority(1)), 7);
3218 assert_eq!(quorum_durable_watermark(&lanes(&[1, 9, 5]), majority(3)), 5);
3219 assert_eq!(
3220 quorum_durable_watermark(&lanes(&[1, 9, 5, 7, 3]), majority(5)),
3221 5
3222 );
3223 }
3224
3225 #[tokio::test]
3226 async fn admitted_prefix_hashes_framed_bytes_in_order_off_path() {
3227 let first = record(b"first").encode().unwrap();
3228 let second = record(b"second").encode().unwrap();
3229 let third = record(b"third").encode().unwrap();
3230 let expected = [first.as_ref(), second.as_ref(), third.as_ref()].concat();
3231
3232 let mut admitted = AdmittedPrefix::default();
3233 let first_batch: Arc<[Bytes]> = vec![first].into();
3234 admitted.extend_metadata(&first_batch);
3235 admitted.queue_digest(first_batch);
3236 let second_batch: Arc<[Bytes]> = vec![second, third].into();
3237 admitted.extend_metadata(&second_batch);
3238 admitted.queue_digest(second_batch);
3239
3240 assert_eq!(admitted.len(), 3);
3241 assert_eq!(admitted.bytes_len(), expected.len());
3242 assert_eq!(admitted.digest().await, digest_bytes(&expected));
3243 assert_eq!(admitted.crc32c(), crc32c::crc32c(&expected));
3244 }
3245
3246 #[test]
3247 fn matching_lane_groups_share_one_packed_wire_payload() {
3248 let first = batch_descriptor(0, vec![Bytes::from_static(b"first")], 2);
3249 let second = batch_descriptor(first.end(), vec![Bytes::from_static(b"second")], 2);
3250 let first_budget = LaneBudget::new();
3251 let second_budget = LaneBudget::new();
3252 let group_one = vec![
3253 LaneBatch::new(
3254 Arc::clone(&first),
3255 first_budget.try_reserve(first.bytes).unwrap(),
3256 ),
3257 LaneBatch::new(
3258 Arc::clone(&second),
3259 first_budget.try_reserve(second.bytes).unwrap(),
3260 ),
3261 ];
3262 let group_two = vec![
3263 LaneBatch::new(
3264 Arc::clone(&first),
3265 second_budget.try_reserve(first.bytes).unwrap(),
3266 ),
3267 LaneBatch::new(
3268 Arc::clone(&second),
3269 second_budget.try_reserve(second.bytes).unwrap(),
3270 ),
3271 ];
3272
3273 let packed_one = packed_group(&group_one);
3274 let packed_two = packed_group(&group_two);
3275
3276 assert!(Arc::ptr_eq(&packed_one, &packed_two));
3277 assert_eq!(
3278 packed_one.chunks(),
3279 &[Bytes::from_static(b"first"), Bytes::from_static(b"second")]
3280 );
3281 drop(group_one);
3282 drop(group_two);
3283 assert!(first
3284 .packed_groups
3285 .lock()
3286 .unwrap_or_else(|poisoned| poisoned.into_inner())
3287 .is_empty());
3288 }
3289
3290 #[test]
3291 fn packed_wire_payload_preserves_offsets_bytes_and_checksums() {
3292 let first = Bytes::from(vec![7; 262_143]);
3293 let second = Bytes::from_static(b"xy");
3294 let packed = pack_append(vec![first.clone(), second.clone()]);
3295 let messages = packed.messages();
3296
3297 assert_eq!(messages.len(), 2);
3298 assert_eq!(messages[0].relative_offset, 0);
3299 assert_eq!(messages[0].content, first);
3300 assert_eq!(messages[0].crc32c, crc32c::crc32c(&messages[0].content));
3301 assert_eq!(messages[1].relative_offset, 262_143);
3302 assert_eq!(messages[1].content, second);
3303 assert_eq!(messages[1].crc32c, crc32c::crc32c(&messages[1].content));
3304 }
3305
3306 #[test]
3307 fn quorum_subsets_enumerate_lexicographically() {
3308 assert_eq!(quorum_subsets(1, 1), vec![vec![0]]);
3309 assert_eq!(
3310 quorum_subsets(3, 2),
3311 vec![vec![0, 1], vec![0, 2], vec![1, 2]]
3312 );
3313 assert_eq!(quorum_subsets(5, 3).len(), 10);
3314 assert_eq!(quorum_subsets(5, 3)[0], vec![0, 1, 2]);
3315 }
3316
3317 #[test]
3318 fn canonical_promotes_a_tail_visible_on_one_recovery_witness() {
3319 let first = record(b"first");
3320 let second = record(b"second");
3321 let snapshots = vec![
3322 snapshot(0, &[first.clone(), second.clone()]),
3323 snapshot(1, std::slice::from_ref(&first)),
3324 ];
3325 assert_eq!(
3326 canonical_prefix(&snapshots, 2).unwrap(),
3327 vec![first, second]
3328 );
3329 }
3330
3331 #[test]
3332 fn canonical_rejects_conflicting_bytes_at_the_same_record() {
3333 let snapshots = vec![snapshot(0, &[record(b"a")]), snapshot(1, &[record(b"b")])];
3334 assert!(matches!(
3335 canonical_prefix(&snapshots, 2),
3336 Err(ProtocolError::ConflictingPrefix { record_index: 0 })
3337 ));
3338 }
3339
3340 #[test]
3341 fn canonical_ignores_one_conflicting_lane_when_two_exact_copies_agree() {
3342 let good = record(b"good");
3343 let snapshots = vec![
3344 snapshot(0, &[record(b"bad")]),
3345 snapshot(1, std::slice::from_ref(&good)),
3346 snapshot(2, std::slice::from_ref(&good)),
3347 ];
3348 assert_eq!(canonical_prefix(&snapshots, 2).unwrap(), vec![good]);
3349 }
3350
3351 #[test]
3352 fn canonical_rejects_equal_length_candidates_without_a_quorum_choice() {
3353 let snapshots = vec![
3354 snapshot(0, &[]),
3355 snapshot(1, &[record(b"left")]),
3356 snapshot(2, &[record(b"right")]),
3357 ];
3358 assert!(matches!(
3359 canonical_prefix(&snapshots, 2),
3360 Err(ProtocolError::ConflictingPrefix { record_index: 0 })
3361 ));
3362 }
3363
3364 #[test]
3365 fn canonical_stops_at_a_partial_tail() {
3366 let first = record(b"first");
3367 let second = record(b"second");
3368 let mut damaged = encode_records(&[first.clone(), second]).unwrap();
3369 damaged.truncate(damaged.len() - 2);
3370 let snapshots = vec![
3371 ReplicaSnapshot {
3372 bytes: damaged,
3373 ..snapshot(0, &[])
3374 },
3375 snapshot(1, std::slice::from_ref(&first)),
3376 ];
3377 assert_eq!(canonical_prefix(&snapshots, 2).unwrap(), vec![first]);
3378 }
3379
3380 #[test]
3381 fn canonical_accepts_a_single_replica_witness() {
3382 let first = record(b"first");
3383 let snapshots = vec![snapshot(0, std::slice::from_ref(&first))];
3384 assert_eq!(canonical_prefix(&snapshots, 1).unwrap(), vec![first]);
3385 }
3386
3387 #[test]
3388 fn canonical_five_zone_quorum_requires_three_consistent_witnesses() {
3389 let good = record(b"good");
3390 let consistent = vec![
3391 snapshot(0, std::slice::from_ref(&good)),
3392 snapshot(1, std::slice::from_ref(&good)),
3393 snapshot(2, &[record(b"divergent")]),
3394 snapshot(3, std::slice::from_ref(&good)),
3395 ];
3396 assert_eq!(
3397 canonical_prefix(&consistent, 3).unwrap(),
3398 vec![good.clone()]
3399 );
3400
3401 let insufficient = vec![
3405 snapshot(0, std::slice::from_ref(&good)),
3406 snapshot(1, std::slice::from_ref(&good)),
3407 snapshot(2, &[record(b"divergent")]),
3408 ];
3409 assert!(matches!(
3410 canonical_prefix(&insufficient, 3),
3411 Err(ProtocolError::ConflictingPrefix { .. })
3412 ));
3413 }
3414
3415 #[test]
3416 fn canonical_five_zone_promotes_the_longest_member_of_the_quorum() {
3417 let first = record(b"first");
3418 let second = record(b"second");
3419 let snapshots = vec![
3420 snapshot(0, std::slice::from_ref(&first)),
3421 snapshot(1, &[first.clone(), second.clone()]),
3422 snapshot(2, &[]),
3423 ];
3424 assert_eq!(
3425 canonical_prefix(&snapshots, 3).unwrap(),
3426 vec![first, second]
3427 );
3428 }
3429
3430 #[test]
3431 fn recovery_size_uses_the_available_quorum_intersection_rank() {
3432 let mut all_three = [10, 20, 20];
3433 assert_eq!(select_recovery_size(&mut all_three, 3), Some(20));
3434
3435 let mut two_of_three = [10, 20];
3436 assert_eq!(select_recovery_size(&mut two_of_three, 3), Some(20));
3437
3438 let mut four_of_five = [10, 20, 30, 40];
3439 assert_eq!(select_recovery_size(&mut four_of_five, 5), Some(30));
3440
3441 let mut all_five = [10, 20, 30, 40, 50];
3442 assert_eq!(select_recovery_size(&mut all_five, 5), Some(30));
3443 }
3444
3445 fn active_commit_tracker(lanes: usize) -> Arc<CommitTracker> {
3446 let recorder = TestMetricsRecorder::default();
3447 let metrics = Arc::new(Metrics::new(&recorder, lanes));
3448 let tracker = CommitTracker::new(lanes, majority(lanes), metrics);
3449 for zone in 0..lanes {
3450 tracker.activate_lane(zone, 0);
3451 }
3452 tracker
3453 }
3454
3455 #[test]
3456 fn quorum_byte_watermark_evicts_resolved_record_boundaries() {
3457 let tracker = active_commit_tracker(3);
3458 let mut range = tracker.admit_window(&[10, 20, 30], &[0, 1, 2]);
3459
3460 tracker.publish_durable(0, 30);
3461 assert_eq!(range.progress().0, 0);
3462 tracker.publish_durable(1, 20);
3463 assert_eq!(range.progress().0, 2);
3464 {
3465 let state = tracker
3466 .state
3467 .lock()
3468 .unwrap_or_else(|poisoned| poisoned.into_inner());
3469 assert_eq!(state.committed_bytes, 20);
3470 assert_eq!(state.boundaries, VecDeque::from([30]));
3471 }
3472
3473 tracker.publish_durable(2, 30);
3474 assert_eq!(range.progress().0, 3);
3475 assert!(tracker
3476 .state
3477 .lock()
3478 .unwrap_or_else(|poisoned| poisoned.into_inner())
3479 .boundaries
3480 .is_empty());
3481 }
3482
3483 #[test]
3484 fn per_zone_durable_lag_tracks_admitted_and_persisted_bytes() {
3485 let recorder = Arc::new(TestMetricsRecorder::default());
3486 let metrics = Arc::new(Metrics::new(recorder.as_ref(), 1));
3487 let tracker = CommitTracker::new(1, 1, metrics);
3488 tracker.activate_lane(0, 0);
3489 let _range = tracker.admit_window(&[10, 20], &[0]);
3490 assert_eq!(
3491 recorder.labeled_gauge("chorus.wal.replica.durable_lag_bytes", &[("zone", "0")]),
3492 20
3493 );
3494
3495 tracker.publish_durable(0, 10);
3496 assert_eq!(
3497 recorder.labeled_gauge("chorus.wal.replica.durable_lag_bytes", &[("zone", "0")]),
3498 10
3499 );
3500 tracker.finish_lane(0, None);
3501 assert_eq!(
3502 recorder.labeled_gauge("chorus.wal.replica.durable_lag_bytes", &[("zone", "0")]),
3503 0
3504 );
3505 }
3506
3507 #[test]
3508 fn retired_lanes_cannot_support_future_admissions() {
3509 let tracker = active_commit_tracker(3);
3510 tracker.finish_lane(0, None);
3511 let mut range = tracker.admit_window(&[10], &[1, 2]);
3512
3513 tracker.finish_lane(1, None);
3514
3515 let (committed, failure) = range.progress();
3516 assert_eq!(committed, 0);
3517 assert!(matches!(failure, Some(ProtocolError::Poisoned)));
3518 }
3519
3520 #[tokio::test]
3521 async fn durable_progress_resets_the_lane_stall_timeout() {
3522 let recorder = Arc::new(TestMetricsRecorder::default());
3523 let metrics = Arc::new(Metrics::new(recorder.as_ref(), 1));
3524 let (sends, _) = mpsc::unbounded_channel();
3525 let replica = Arc::new(ScriptedLaneReplica::new(0, VecDeque::new(), sends));
3526 let replica_for_progress: Arc<dyn Replica> = replica.clone();
3527 let updater = replica.clone();
3528 let updates = tokio::spawn(async move {
3529 for durable in [10, 20, 30] {
3530 tokio::time::sleep(Duration::from_millis(5)).await;
3531 updater.durable.send_replace(durable);
3532 }
3533 });
3534 let timeout = Duration::from_millis(100);
3535 let mut seen = 0;
3536
3537 for expected in [10, 20, 30] {
3538 match lane_progress(
3539 &replica_for_progress,
3540 seen,
3541 true,
3542 tokio::time::Instant::now() + timeout,
3543 &metrics,
3544 )
3545 .await
3546 {
3547 LaneProgress::Advanced(change) => {
3548 assert_eq!(change.persisted_size, expected);
3549 assert!(change.error.is_none());
3550 seen = change.persisted_size;
3551 }
3552 LaneProgress::Stalled => panic!("advancing lane was falsely shed"),
3553 LaneProgress::Failed(error) => panic!("advancing lane failed: {error}"),
3554 }
3555 }
3556 updates.await.unwrap();
3557 assert_eq!(recorder.counter("chorus.wal.lane.timeouts"), 0);
3558 }
3559
3560 #[tokio::test]
3561 async fn ready_progress_wins_an_expired_operation_deadline() {
3562 let recorder = Arc::new(TestMetricsRecorder::default());
3563 let metrics = Arc::new(Metrics::new(recorder.as_ref(), 1));
3564 let tracker = CommitTracker::new(1, 1, Arc::clone(&metrics));
3565 tracker.activate_lane(0, 0);
3566 let (sends, _) = mpsc::unbounded_channel();
3567 let replica = Arc::new(ScriptedLaneReplica::new(0, VecDeque::new(), sends));
3568 replica.durable.send_replace(10);
3569 let replica: Arc<dyn Replica> = replica;
3570 let mut durable = 0;
3571 let mut last_progress = tokio::time::Instant::now() - Duration::from_secs(1);
3572 let mut retained = VecDeque::new();
3573
3574 let change = confirm_lane_stall(&replica, durable, tokio::time::Instant::now(), &metrics)
3575 .await
3576 .expect("ready durable progress must win the deadline tie");
3577 publish_lane_advance(
3578 change,
3579 0,
3580 &mut durable,
3581 &mut last_progress,
3582 &mut retained,
3583 &tracker,
3584 )
3585 .expect("ready progress has no stream failure");
3586
3587 assert_eq!(durable, 10);
3588 assert_eq!(recorder.counter("chorus.wal.lane.timeouts"), 0);
3589 }
3590
3591 #[tokio::test]
3592 async fn fencing_lane_failure_stops_the_writer_after_publishing_progress() {
3593 let tracker = active_commit_tracker(1);
3594 let range = tracker.admit_window(&[10], &[0]);
3595
3596 tracker.publish_durable(0, 10);
3597 tracker.finish_lane(
3598 0,
3599 Some(TransportError {
3600 zone: 0,
3601 code: TransportCode::FailedPrecondition,
3602 message: "newer writer took over".into(),
3603 }),
3604 );
3605
3606 assert_eq!(range.into_pending().remove(0).wait().await.unwrap(), 0);
3607 assert!(tracker.is_poisoned());
3608 let mut updates = tracker.subscribe();
3609 let snapshot = updates.borrow_and_update().clone();
3610 assert!(matches!(snapshot.failure, Some(CommitFailure::Fenced(_))));
3611 }
3612
3613 #[tokio::test]
3614 async fn pending_commits_follow_the_prefix_watermark_and_gap_poison() {
3615 let tracker = active_commit_tracker(3);
3616 let range = tracker.admit_window(&[10, 20], &[0, 1, 2]);
3617 let mut pending = range.into_pending();
3618 let second = pending.pop().expect("second pending commit");
3619 let first = pending.pop().expect("first pending commit");
3620
3621 tracker.publish_durable(0, 20);
3622 tracker.publish_durable(1, 10);
3623
3624 tracker.finish_lane(
3625 1,
3626 Some(TransportError {
3627 zone: 1,
3628 code: TransportCode::PermissionDenied,
3629 message: "terminal".into(),
3630 }),
3631 );
3632 tracker.finish_lane(
3633 2,
3634 Some(TransportError {
3635 zone: 2,
3636 code: TransportCode::Unavailable,
3637 message: "transient".into(),
3638 }),
3639 );
3640 tracker.publish_durable(2, 20);
3641
3642 assert_eq!(first.wait().await.unwrap(), 0);
3643 assert!(matches!(
3644 second.wait().await,
3645 Err(ProtocolError::Transport(TransportError {
3646 zone: 1,
3647 code: TransportCode::PermissionDenied,
3648 ..
3649 }))
3650 ));
3651 assert_eq!(tracker.committed_len(), 1);
3652 }
3653
3654 #[tokio::test]
3655 async fn ready_progress_releases_lane_budget_before_more_work_is_staged() {
3656 let recorder = Arc::new(TestMetricsRecorder::default());
3657 let metrics = Arc::new(Metrics::new(recorder.as_ref(), 1));
3658 let (first_send_release_tx, first_send_release_rx) = oneshot::channel();
3659 let (second_send_release_tx, second_send_release_rx) = oneshot::channel();
3660 let (sends_tx, mut sends_rx) = mpsc::unbounded_channel();
3661 let replica: Arc<dyn Replica> = Arc::new(ScriptedLaneReplica::new(
3662 0,
3663 VecDeque::from([first_send_release_rx, second_send_release_rx]),
3664 sends_tx,
3665 ));
3666 let mut writer = Writer::new(
3667 vec![replica],
3668 ClientConfig::default(),
3669 protocol_metadata(),
3670 vec![AppendToken {
3671 zone: 0,
3672 generation: Some(1),
3673 metageneration: Some(1),
3674 persisted_size: 0,
3675 write_handle: None,
3676 }],
3677 metrics,
3678 );
3679 let encoded = record(b"first").encode().unwrap().len();
3680 writer.set_max_replica_lag_bytes(encoded * 2);
3681 let attempted: AttemptedBytes = Arc::new(|_| {});
3682
3683 let first = writer
3684 .enqueue_data_window(vec![record(b"first")], attempted.clone())
3685 .await
3686 .unwrap()
3687 .into_pending()
3688 .remove(0);
3689 assert_eq!(sends_rx.recv().await, Some(encoded as i64));
3690
3691 let second = writer
3692 .enqueue_data_window(vec![record(b"other")], attempted.clone())
3693 .await
3694 .unwrap()
3695 .into_pending()
3696 .remove(0);
3697 first_send_release_tx
3698 .send(())
3699 .expect("the first staged group should still be blocked");
3700 assert_eq!(sends_rx.recv().await, Some((encoded * 2) as i64));
3701
3702 let third = writer
3703 .enqueue_data_window(vec![record(b"third")], attempted)
3704 .await
3705 .unwrap()
3706 .into_pending()
3707 .remove(0);
3708 second_send_release_tx
3709 .send(())
3710 .expect("the second staged group should still be blocked");
3711
3712 assert_eq!(first.wait().await.unwrap(), 0);
3713 assert_eq!(second.wait().await.unwrap(), 1);
3714 assert_eq!(third.wait().await.unwrap(), 2);
3715 assert_eq!(writer.committed_len(), 3);
3716 assert_eq!(recorder.counter("chorus.wal.lane.capacity_drops"), 0);
3717 assert!(!writer.is_poisoned());
3718 }
3719
3720 #[tokio::test]
3721 async fn terminal_lane_failures_preserve_transport_errors_in_completions() {
3722 let replica = Arc::new(ReaderTerminalReplica::new(0));
3723 let factory: Arc<dyn crate::transport::ReplicaFactory> =
3724 Arc::new(SingleReplicaFactory::new(replica.clone()));
3725 let manifest_store =
3726 Arc::new(crate::manifest_store::test_support::InMemoryManifestStore::default());
3727 let volume = crate::segment::SegmentedVolume::new_with_factories_and_manifest_store(
3728 vec![factory],
3729 manifest_store,
3730 "terminal-completion",
3731 ClientConfig {
3732 max_retries: 0,
3733 retry_base: Duration::ZERO,
3734 },
3735 )
3736 .unwrap();
3737 let writer = volume.recover_writer().await.unwrap();
3738 let mut handle = crate::engine::WalEngine::start(
3739 writer,
3740 crate::WalEngineConfig {
3741 repair_interval: None,
3742 ..Default::default()
3743 },
3744 )
3745 .unwrap();
3746 let completion = handle
3747 .enqueue_append(
3748 crate::segment::WalSeqNo::ZERO,
3749 Bytes::from_static(b"terminal"),
3750 )
3751 .await
3752 .unwrap();
3753 let error = completion.await.unwrap_err();
3754 assert!(matches!(
3755 error,
3756 crate::Error::Transport {
3757 code: TransportCode::PermissionDenied,
3758 ..
3759 }
3760 ));
3761 assert_eq!(replica.resume_calls(), 0);
3762 let _ = tokio::time::timeout(Duration::from_secs(1), handle.shutdown())
3763 .await
3764 .expect("engine shutdown timed out");
3765 }
3766
3767 #[tokio::test]
3768 async fn stalled_lane_poison_releases_blocked_admission() {
3769 let replica: Arc<dyn Replica> = Arc::new(StalledReplica { zone: 0 });
3770 let factory: Arc<dyn crate::transport::ReplicaFactory> =
3771 Arc::new(SingleReplicaFactory::new(replica));
3772 let manifest_store =
3773 Arc::new(crate::manifest_store::test_support::InMemoryManifestStore::default());
3774 let volume = crate::segment::SegmentedVolume::new_with_factories_and_manifest_store(
3775 vec![factory],
3776 manifest_store,
3777 "stalled-admission",
3778 ClientConfig {
3779 max_retries: 0,
3780 retry_base: Duration::ZERO,
3781 },
3782 )
3783 .unwrap();
3784 let writer = volume.recover_writer().await.unwrap();
3785 let payload = Bytes::from_static(b"stalled");
3786 let encoded_bytes = payload.len() + 4;
3787 let stall_timeout = Duration::from_millis(20);
3788 let mut handle = crate::engine::WalEngine::start(
3789 writer,
3790 crate::WalEngineConfig {
3791 queue_capacity: 2,
3792 max_record_bytes: payload.len(),
3793 pipeline_window_records: 1,
3794 max_inflight_bytes: encoded_bytes,
3795 max_replica_lag_bytes: encoded_bytes,
3796 lane_stall_timeout: stall_timeout,
3797 repair_interval: None,
3798 ..Default::default()
3799 },
3800 )
3801 .unwrap();
3802 let first = handle
3803 .enqueue_append(crate::segment::WalSeqNo::ZERO, payload.clone())
3804 .await
3805 .unwrap();
3806
3807 let second = tokio::time::timeout(
3808 stall_timeout.saturating_mul(10),
3809 handle.enqueue_append(crate::segment::WalSeqNo::record(1), payload),
3810 )
3811 .await
3812 .expect("blocked admission did not wake after the writer poisoned");
3813 assert!(matches!(second, Err(crate::Error::Closed)));
3814
3815 let first = tokio::time::timeout(stall_timeout.saturating_mul(10), first)
3816 .await
3817 .expect("admitted append did not receive terminal poison");
3818 assert!(matches!(first, Err(crate::Error::Poisoned)));
3819 let _ = handle.shutdown().await;
3820 }
3821}