1use std::collections::VecDeque;
2use std::future::Future;
3use std::pin::Pin;
4use std::sync::Arc;
5use std::task::{Context, Poll};
6
7use bytes::Bytes;
8use futures::future::join_all;
9use tokio::sync::{mpsc, oneshot, watch, OwnedSemaphorePermit, Semaphore};
10use tokio::time::{Duration, Sleep};
11
12use crate::error::Error;
13use crate::manifest::ManifestUpdate;
14use crate::metrics::Metrics;
15use crate::record::RecordFrame;
16use crate::segment::{
17 PendingSwap, RecordCommitRange, RegisteredSpare, SegmentedWriter, TruncationReport, WalSeqNo,
18};
19use crate::transport::Replica;
20
21const DEFAULT_PIPELINE_WINDOW_RECORDS: usize = 32;
22const DEFAULT_QUEUE_CAPACITY: usize = DEFAULT_PIPELINE_WINDOW_RECORDS * 8;
23const DEFAULT_SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(300);
24
25#[derive(Clone, Debug)]
26pub struct WalEngineConfig {
28 pub queue_capacity: usize,
34 pub max_record_bytes: usize,
36 pub pipeline_window_records: usize,
39 pub max_inflight_bytes: usize,
43 pub max_replica_lag_bytes: usize,
48 pub lane_stall_timeout: Duration,
57 pub max_segment_bytes: usize,
69 pub max_active_segment_bytes: usize,
80 pub repair_interval: Option<Duration>,
88 pub shutdown_timeout: Duration,
94}
95
96impl Default for WalEngineConfig {
97 fn default() -> Self {
98 Self {
99 queue_capacity: DEFAULT_QUEUE_CAPACITY,
100 max_record_bytes: 1024 * 1024,
101 pipeline_window_records: DEFAULT_PIPELINE_WINDOW_RECORDS,
102 max_inflight_bytes: 64 * 1024 * 1024,
103 max_replica_lag_bytes: 64 * 1024 * 1024,
104 lane_stall_timeout: crate::protocol::DEFAULT_LANE_STALL_TIMEOUT,
105 max_segment_bytes: 256 * 1024 * 1024,
106 max_active_segment_bytes: usize::try_from(4_u64 * 1024 * 1024 * 1024)
107 .unwrap_or(usize::MAX),
108 repair_interval: Some(Duration::from_secs(300)),
109 shutdown_timeout: DEFAULT_SHUTDOWN_TIMEOUT,
110 }
111 }
112}
113
114impl WalEngineConfig {
115 fn validate(&self) -> Result<(), Error> {
116 if self.queue_capacity == 0 {
117 return Err(Error::InvalidConfig("queue_capacity must be nonzero"));
118 }
119 if self.max_record_bytes == 0 {
120 return Err(Error::InvalidConfig("max_record_bytes must be nonzero"));
121 }
122 if self.max_record_bytes > RecordFrame::MAX_PAYLOAD_BYTES {
123 return Err(Error::InvalidConfig(
124 "max_record_bytes exceeds the record format limit",
125 ));
126 }
127 if self.pipeline_window_records == 0 {
128 return Err(Error::InvalidConfig(
129 "pipeline_window_records must be nonzero",
130 ));
131 }
132 if self.max_inflight_bytes == 0 {
133 return Err(Error::InvalidConfig("max_inflight_bytes must be nonzero"));
134 }
135 if self.max_replica_lag_bytes == 0 {
136 return Err(Error::InvalidConfig(
137 "max_replica_lag_bytes must be nonzero",
138 ));
139 }
140 let max_encoded_record = self
141 .max_record_bytes
142 .checked_add(4)
143 .ok_or(Error::InvalidConfig("max_record_bytes is too large"))?;
144 if max_encoded_record > self.max_inflight_bytes {
145 return Err(Error::InvalidConfig(
146 "max_inflight_bytes must fit one maximum-size encoded record",
147 ));
148 }
149 if self.max_replica_lag_bytes < self.max_inflight_bytes {
150 return Err(Error::InvalidConfig(
151 "max_replica_lag_bytes must be at least max_inflight_bytes",
152 ));
153 }
154 if self.lane_stall_timeout == Duration::ZERO {
155 return Err(Error::InvalidConfig("lane_stall_timeout must be nonzero"));
156 }
157 if self.max_inflight_bytes > Semaphore::MAX_PERMITS {
158 return Err(Error::InvalidConfig(
159 "max_inflight_bytes exceeds the semaphore limit",
160 ));
161 }
162 if self.max_segment_bytes == 0 {
163 return Err(Error::InvalidConfig("max_segment_bytes must be nonzero"));
164 }
165 if self.max_active_segment_bytes == 0 {
166 return Err(Error::InvalidConfig(
167 "max_active_segment_bytes must be nonzero",
168 ));
169 }
170 if self.max_active_segment_bytes < max_encoded_record {
171 return Err(Error::InvalidConfig(
172 "max_active_segment_bytes must fit one maximum-size encoded record",
173 ));
174 }
175 if self.max_active_segment_bytes < self.max_segment_bytes {
176 return Err(Error::InvalidConfig(
177 "max_active_segment_bytes must be at least max_segment_bytes",
178 ));
179 }
180 if self.repair_interval == Some(Duration::ZERO) {
181 return Err(Error::InvalidConfig("repair_interval must be nonzero"));
182 }
183 if self.shutdown_timeout == Duration::ZERO {
184 return Err(Error::InvalidConfig("shutdown_timeout must be nonzero"));
185 }
186 Ok(())
187 }
188}
189
190#[derive(Clone, Debug, Eq, PartialEq)]
191pub struct AppendReceipt {
196 pub seqno: WalSeqNo,
198}
199
200impl AppendReceipt {
201 pub fn next_seqno(&self) -> WalSeqNo {
203 WalSeqNo::record(self.seqno.record_index + 1)
204 }
205}
206
207pub struct AppendCompletion {
215 receiver: oneshot::Receiver<Result<AppendReceipt, Error>>,
216}
217
218impl Future for AppendCompletion {
219 type Output = Result<AppendReceipt, Error>;
220
221 fn poll(mut self: Pin<&mut Self>, context: &mut Context<'_>) -> Poll<Self::Output> {
222 match Pin::new(&mut self.receiver).poll(context) {
223 Poll::Ready(Ok(result)) => Poll::Ready(result),
224 Poll::Ready(Err(_)) => Poll::Ready(Err(Error::Closed)),
225 Poll::Pending => Poll::Pending,
226 }
227 }
228}
229
230pub struct WalEngine;
232
233pub struct WalHandle {
241 sender: mpsc::Sender<Command>,
242 metrics: Arc<Metrics>,
243 engine_task: tokio::task::JoinHandle<()>,
244 provisioner_task: tokio::task::JoinHandle<()>,
245 provisioner_shutdown: watch::Sender<bool>,
246 maintenance: crate::maintenance::MaintenanceHandle,
247 maintenance_task: tokio::task::JoinHandle<()>,
248 rotation_recheck: mpsc::Sender<()>,
249 next_seqno: u64,
250 max_record_bytes: usize,
251 max_inflight_bytes: usize,
252 max_active_segment_bytes: usize,
253 total_admitted_bytes: u128,
254 active_capacity: watch::Receiver<ActiveSegmentCapacity>,
255 inflight_bytes: Arc<Semaphore>,
256 queue_slots: Arc<Semaphore>,
257 shutdown_timeout: Duration,
258}
259
260struct AdmittedAppend {
261 seqno: WalSeqNo,
262 record: RecordFrame,
263 payload_bytes: usize,
264 encoded_bytes: usize,
265 admission_end_bytes: u128,
270 admitted_at: tokio::time::Instant,
273 _inflight_bytes: OwnedSemaphorePermit,
274 queue_slot: Option<OwnedSemaphorePermit>,
275 completion: oneshot::Sender<Result<AppendReceipt, Error>>,
276}
277
278enum Command {
279 Append(AdmittedAppend),
280 Shutdown { response: oneshot::Sender<()> },
281}
282
283struct CompletionBatch {
284 commits: RecordCommitRange,
285 appends: VecDeque<AdmittedAppend>,
286}
287
288struct EngineState {
289 queue: VecDeque<Command>,
290 completion_batches: VecDeque<CompletionBatch>,
291 pending_records: usize,
292 next_completion: u64,
293 next_admission: u64,
294 last_dispatched_admission_bytes: u128,
295 input_closed: bool,
296 shutdown_response: Option<oneshot::Sender<()>>,
297}
298
299impl EngineState {
300 fn new(next_record_index: u64) -> Self {
301 Self {
302 queue: VecDeque::new(),
303 completion_batches: VecDeque::new(),
304 pending_records: 0,
305 next_completion: next_record_index,
306 next_admission: next_record_index,
307 last_dispatched_admission_bytes: 0,
308 input_closed: false,
309 shutdown_response: None,
310 }
311 }
312
313 fn drain_available(&mut self, receiver: &mut mpsc::Receiver<Command>) {
314 loop {
315 match receiver.try_recv() {
316 Ok(command) => {
317 admit_command(command, &mut self.queue, &mut self.next_admission);
318 }
319 Err(mpsc::error::TryRecvError::Empty) => break,
320 Err(mpsc::error::TryRecvError::Disconnected) => {
321 self.input_closed = true;
322 break;
323 }
324 }
325 }
326 }
327
328 fn stop_admission(&mut self, receiver: &mut mpsc::Receiver<Command>) {
329 receiver.close();
330 while let Ok(command) = receiver.try_recv() {
331 self.queue.push_back(command);
332 }
333 }
334
335 fn fail_queued(
336 &mut self,
337 receiver: &mut mpsc::Receiver<Command>,
338 error: Error,
339 metrics: &Metrics,
340 ) {
341 self.stop_admission(receiver);
342 fail_all(&mut self.queue, error, metrics);
343 metrics.operation_failures.increment();
344 }
345
346 fn poison(&mut self, receiver: &mut mpsc::Receiver<Command>, metrics: &Metrics) {
347 self.stop_admission(receiver);
348 fail_completion_batches(&mut self.completion_batches, Error::Poisoned, metrics);
349 fail_all(&mut self.queue, Error::Poisoned, metrics);
350 metrics.operation_failures.increment();
351 }
352}
353
354enum CompletionFailure {
355 Append(AdmittedAppend, Error),
356 Invariant(&'static str),
357}
358
359#[derive(Clone, Copy, Debug)]
360struct ActiveSegmentCapacity {
361 admission_start_bytes: u128,
363 existing_bytes: usize,
369 seal_room: bool,
372}
373
374struct EngineControl {
375 catalog: watch::Sender<Vec<crate::segment::SegmentDescriptor>>,
376 maintenance: crate::maintenance::MaintenanceHandle,
377 active_capacity: watch::Sender<ActiveSegmentCapacity>,
378 queue_slots: Arc<Semaphore>,
379 rotation_rechecks: mpsc::Receiver<()>,
380 provision_requests: mpsc::Sender<ProvisionAttempt>,
381 provision_results: mpsc::Receiver<Result<ProvisionOutcome, Error>>,
382}
383
384impl WalEngine {
385 pub fn start(mut writer: SegmentedWriter, config: WalEngineConfig) -> Result<WalHandle, Error> {
394 config.validate()?;
395 let (sender, receiver) = mpsc::channel(config.queue_capacity);
396 let metrics = writer.metrics();
397 let next_seqno = writer.committed_record_end();
398 let active_segment_bytes = writer.active_segment_bytes();
399 let active_segment_seal_room = writer.active_segment_has_seal_room();
400 if active_segment_bytes != 0 && !active_segment_seal_room {
401 return Err(Error::SegmentDirectoryFull);
402 }
403 let max_record_bytes = config.max_record_bytes;
404 let max_inflight_bytes = config.max_inflight_bytes;
405 let max_active_segment_bytes = config.max_active_segment_bytes;
406 let shutdown_timeout = config.shutdown_timeout;
407 let inflight_bytes = Arc::new(Semaphore::new(max_inflight_bytes));
408 let queue_slots = Arc::new(Semaphore::new(config.queue_capacity));
409 writer.set_max_replica_lag_bytes(config.max_replica_lag_bytes);
410 writer.set_lane_stall_timeout(config.lane_stall_timeout);
411 let (catalog_tx, catalog_rx) = watch::channel(writer.sealed_segments_snapshot());
417 let (maintenance, maintenance_task) = crate::maintenance::start(
418 writer.maintenance_config(config.repair_interval),
419 catalog_rx,
420 Arc::clone(&metrics),
421 );
422 let (active_capacity_tx, active_capacity) = watch::channel(ActiveSegmentCapacity {
423 admission_start_bytes: 0,
424 existing_bytes: active_segment_bytes,
425 seal_room: active_segment_seal_room,
426 });
427 let (rotation_recheck, rotation_rechecks) = mpsc::channel(1);
428 let (provision_requests, provision_attempts) = mpsc::channel::<ProvisionAttempt>(1);
429 let (provision_results_tx, provision_results) = mpsc::channel(1);
430 let (provisioner_shutdown, provisioner_shutdown_rx) = watch::channel(false);
431 let provisioner_task = tokio::spawn(run_provisioner(
432 provision_attempts,
433 provision_results_tx,
434 provisioner_shutdown_rx,
435 ));
436 tracing::info!(
437 next_record_index = next_seqno,
438 queue_capacity = config.queue_capacity,
439 pipeline_window_records = config.pipeline_window_records,
440 max_inflight_bytes = config.max_inflight_bytes,
441 max_replica_lag_bytes = config.max_replica_lag_bytes,
442 lane_stall_timeout_ms = config.lane_stall_timeout.as_millis(),
443 max_segment_bytes = config.max_segment_bytes,
444 max_active_segment_bytes = config.max_active_segment_bytes,
445 "WAL engine started"
446 );
447 let engine_task = tokio::spawn(run_engine(
448 writer,
449 config,
450 receiver,
451 Arc::clone(&metrics),
452 EngineControl {
453 catalog: catalog_tx,
454 maintenance: maintenance.clone(),
455 active_capacity: active_capacity_tx,
456 queue_slots: Arc::clone(&queue_slots),
457 rotation_rechecks,
458 provision_requests,
459 provision_results,
460 },
461 ));
462 Ok(WalHandle {
463 sender,
464 metrics,
465 engine_task,
466 provisioner_task,
467 provisioner_shutdown,
468 maintenance,
469 maintenance_task,
470 rotation_recheck,
471 next_seqno,
472 max_record_bytes,
473 max_inflight_bytes,
474 max_active_segment_bytes,
475 total_admitted_bytes: 0,
476 active_capacity,
477 inflight_bytes,
478 queue_slots,
479 shutdown_timeout,
480 })
481 }
482}
483
484impl WalHandle {
485 pub async fn enqueue_append(
505 &mut self,
506 seqno: WalSeqNo,
507 record: Bytes,
508 ) -> Result<AppendCompletion, Error> {
509 if self.sender.is_closed() {
510 return Err(Error::Closed);
511 }
512 if seqno.record_index != self.next_seqno {
513 return Err(Error::OutOfOrder {
514 expected: WalSeqNo::record(self.next_seqno),
515 actual: seqno,
516 });
517 }
518 if seqno.record_index == u64::MAX {
519 return Err(Error::SequenceExhausted);
520 }
521 let payload_bytes = record.len();
522 if payload_bytes > self.max_record_bytes {
523 return Err(Error::RecordTooLarge {
524 max: self.max_record_bytes,
525 actual: payload_bytes,
526 });
527 }
528
529 let record = RecordFrame { payload: record };
530 let encoded_bytes = record.encoded_len().map_err(|_| Error::RecordTooLarge {
531 max: self.max_record_bytes.min(RecordFrame::MAX_PAYLOAD_BYTES),
532 actual: payload_bytes,
533 })?;
534 let capacity = *self.active_capacity.borrow_and_update();
535 if !capacity.seal_room {
536 return Err(Error::SegmentDirectoryFull);
537 }
538 let admitted_here = self
539 .total_admitted_bytes
540 .checked_sub(capacity.admission_start_bytes)
541 .ok_or_else(|| {
542 Error::Internal(
543 "active-segment admission boundary exceeds total admitted bytes".into(),
544 )
545 })?;
546 let current_active = (capacity.existing_bytes as u128)
547 .checked_add(admitted_here)
548 .ok_or_else(|| Error::Internal("active-segment byte accounting overflowed".into()))?;
549 let projected = current_active
550 .checked_add(encoded_bytes as u128)
551 .ok_or_else(|| Error::Internal("active-segment byte accounting overflowed".into()))?;
552 if projected > self.max_active_segment_bytes as u128 {
553 return Err(Error::ActiveSegmentFull {
554 max: self.max_active_segment_bytes,
555 current: usize::try_from(current_active).unwrap_or(usize::MAX),
556 requested: encoded_bytes,
557 });
558 }
559 let admission_end_bytes = self
564 .total_admitted_bytes
565 .checked_add(encoded_bytes as u128)
566 .ok_or_else(|| Error::Internal("cumulative admission bytes overflowed".into()))?;
567 let permits = u32::try_from(encoded_bytes)
568 .map_err(|_| Error::Internal("encoded record length exceeds u32".into()))?;
569 let inflight_bytes = Arc::clone(&self.inflight_bytes)
570 .acquire_many_owned(permits)
571 .await
572 .map_err(|_| Error::Closed)?;
573 let queue_slot = Arc::clone(&self.queue_slots)
574 .acquire_owned()
575 .await
576 .map_err(|_| Error::Closed)?;
577 self.metrics.max_inflight_bytes.update_max(
578 self.max_inflight_bytes
579 .saturating_sub(self.inflight_bytes.available_permits()) as u64,
580 );
581 let (completion, receiver) = oneshot::channel();
582 self.sender
583 .send(Command::Append(AdmittedAppend {
584 seqno,
585 record,
586 payload_bytes,
587 encoded_bytes,
588 admission_end_bytes,
589 admitted_at: tokio::time::Instant::now(),
590 _inflight_bytes: inflight_bytes,
591 queue_slot: Some(queue_slot),
592 completion,
593 }))
594 .await
595 .map_err(|_| Error::Closed)?;
596
597 self.total_admitted_bytes = admission_end_bytes;
598 self.next_seqno += 1;
599 self.metrics.append_records.increment();
600 self.metrics.append_bytes.add(payload_bytes as u64);
601 Ok(AppendCompletion { receiver })
602 }
603
604 pub async fn truncate_before(&self, floor: WalSeqNo) -> Result<TruncationReport, Error> {
620 let report = self.maintenance.truncate(floor).await?;
621 let _ = self.rotation_recheck.try_send(());
626 Ok(report)
627 }
628
629 pub async fn abort(self) {
635 let WalHandle {
636 sender,
637 engine_task,
638 provisioner_task,
639 provisioner_shutdown,
640 maintenance,
641 maintenance_task,
642 shutdown_timeout,
643 ..
644 } = self;
645 let mut engine_task = engine_task;
646 let mut provisioner_task = provisioner_task;
647 let mut maintenance_task = maintenance_task;
648 let _ = provisioner_shutdown.send(true);
649 engine_task.abort();
650 maintenance_task.abort();
651 drop(sender);
652 drop(maintenance);
653 if tokio::time::timeout(shutdown_timeout, async {
654 let _ = tokio::join!(
655 &mut engine_task,
656 &mut provisioner_task,
657 &mut maintenance_task
658 );
659 })
660 .await
661 .is_err()
662 {
663 engine_task.abort();
664 provisioner_task.abort();
665 maintenance_task.abort();
666 let _ = tokio::join!(
667 &mut engine_task,
668 &mut provisioner_task,
669 &mut maintenance_task
670 );
671 tracing::error!("WAL abort required forced cancellation after the shutdown deadline");
672 }
673 }
674
675 pub async fn shutdown(self) -> Result<(), Error> {
679 let WalHandle {
680 sender,
681 engine_task,
682 provisioner_task,
683 provisioner_shutdown,
684 maintenance,
685 maintenance_task,
686 shutdown_timeout,
687 ..
688 } = self;
689 let mut engine_task = engine_task;
690 let mut provisioner_task = provisioner_task;
691 let mut maintenance_task = maintenance_task;
692 let maintenance_shutdown = maintenance.clone();
693 let provisioner_timeout_shutdown = provisioner_shutdown.clone();
694 let graceful = async {
695 let (response, receiver) = oneshot::channel();
696 let sent = sender.send(Command::Shutdown { response }).await;
697 drop(sender);
698 let acknowledged = if sent.is_ok() {
699 receiver.await.is_ok()
700 } else {
701 false
702 };
703 let _ = provisioner_shutdown.send(true);
707 let (engine_result, provisioner_result) =
708 tokio::join!(&mut engine_task, &mut provisioner_task);
709 let engine_joined = engine_result.is_ok();
710 let provisioner_joined = provisioner_result.is_ok();
711
712 maintenance.shutdown();
716 drop(maintenance);
717 let maintenance_joined = (&mut maintenance_task).await.is_ok();
718
719 acknowledged && engine_joined && provisioner_joined && maintenance_joined
720 };
721
722 match tokio::time::timeout(shutdown_timeout, graceful).await {
723 Ok(true) => Ok(()),
724 Ok(false) => Err(Error::Closed),
725 Err(_) => {
726 tracing::error!(
727 ?shutdown_timeout,
728 "WAL graceful shutdown exceeded its deadline; aborting owned tasks"
729 );
730 maintenance_shutdown.shutdown();
731 let _ = provisioner_timeout_shutdown.send(true);
732 engine_task.abort();
733 maintenance_task.abort();
734 let cleanup = async {
735 let _ = tokio::join!(
736 &mut engine_task,
737 &mut provisioner_task,
738 &mut maintenance_task
739 );
740 };
741 if tokio::time::timeout(shutdown_timeout, cleanup)
742 .await
743 .is_err()
744 {
745 engine_task.abort();
746 provisioner_task.abort();
747 maintenance_task.abort();
748 let _ = tokio::join!(
749 &mut engine_task,
750 &mut provisioner_task,
751 &mut maintenance_task
752 );
753 tracing::error!(
754 "WAL shutdown required forced cancellation after the cleanup deadline"
755 );
756 }
757 Err(Error::ShutdownTimeout {
758 timeout: shutdown_timeout,
759 })
760 }
761 }
762 }
763}
764
765async fn run_engine(
766 mut writer: SegmentedWriter,
767 config: WalEngineConfig,
768 mut receiver: mpsc::Receiver<Command>,
769 metrics: Arc<Metrics>,
770 control: EngineControl,
771) {
772 let EngineControl {
773 catalog,
774 maintenance,
775 active_capacity,
776 queue_slots,
777 mut rotation_rechecks,
778 provision_requests,
779 mut provision_results,
780 } = control;
781 let client_config = writer.client_config();
782 let mut state = EngineState::new(writer.committed_record_end());
783 let mut rotation =
784 writer
785 .take_recovered_fold()
786 .map_or(Rotation::Idle, |swap| Rotation::Draining {
787 swap: Box::new(swap),
788 fold_attempts: 0,
789 retry: None,
790 fold_capacity_blocked: false,
791 successor_due: false,
792 });
793 let mut rotation_rechecks_open = true;
794 let mut spare_requested = false;
798 let attempted_stats = Arc::clone(&metrics);
799 let on_attempted: crate::protocol::AttemptedBytes = Arc::new(move |bytes| {
800 attempted_stats.replica_bytes_attempted.add(bytes);
801 });
802
803 'engine: loop {
804 metrics.rotation_state.set(rotation.metric_value());
805 state.drain_available(&mut receiver);
806 observe_queue_depth(&metrics, config.queue_capacity, &queue_slots);
807
808 let failed = 'completion: loop {
814 let Some(batch) = state.completion_batches.front_mut() else {
815 break None;
816 };
817 let (committed_end, failure) = batch.commits.progress();
818 while batch
819 .appends
820 .front()
821 .is_some_and(|append| append.seqno.record_index < committed_end)
822 {
823 let Some(append) = batch.appends.pop_front() else {
824 break 'completion Some(CompletionFailure::Invariant(
828 "commit watermark advanced without a queued append",
829 ));
830 };
831 if state.pending_records == 0 {
832 break 'completion Some(CompletionFailure::Invariant(
833 "completion queue exceeded the pending record count",
834 ));
835 }
836 debug_assert_eq!(append.seqno.record_index, state.next_completion);
837 complete_append(append, &metrics);
838 state.pending_records -= 1;
839 state.next_completion += 1;
840 rotation.mark_due(writer.rotation_due(config.max_segment_bytes));
841 }
842 if batch.appends.is_empty() {
843 state.completion_batches.pop_front();
844 continue;
845 }
846 if let Some(error) = failure {
847 let Some(append) = batch.appends.pop_front() else {
848 break Some(CompletionFailure::Invariant(
849 "failed commit batch had no unresolved append",
850 ));
851 };
852 if state.pending_records == 0 {
853 break Some(CompletionFailure::Invariant(
854 "failed completion exceeded the pending record count",
855 ));
856 }
857 state.pending_records -= 1;
858 break Some(CompletionFailure::Append(append, error));
859 }
860 break None;
861 };
862 if let Some(failure) = failed {
863 match failure {
864 CompletionFailure::Append(append, error) => {
865 tracing::warn!(
866 record_index = append.seqno.record_index,
867 %error,
868 "WAL engine poisoned by an indeterminate append"
869 );
870 fail_append(append, error, &metrics);
871 }
872 CompletionFailure::Invariant(message) => {
873 tracing::error!(message, "WAL engine completion invariant failed");
874 }
875 }
876 state.poison(&mut receiver, &metrics);
877 break 'engine;
878 }
879
880 if !spare_requested {
885 let attempt = if rotation.fold_ready(state.next_completion)
886 && writer.unregistered_spare_ready()
887 {
888 let fold = match rotation
889 .swap()
890 .ok_or_else(|| Error::Internal("fold readiness without a swap".into()))
891 .and_then(|swap| writer.pending_fold_request(swap))
892 {
893 Ok(fold) => fold,
894 Err(error) => {
895 state.fail_queued(&mut receiver, error, &metrics);
896 break 'engine;
897 }
898 };
899 let parts = match writer.provision_parts() {
900 Ok(parts) => parts,
901 Err(error) => {
902 state.fail_queued(&mut receiver, error, &metrics);
903 break 'engine;
904 }
905 };
906 Some(ProvisionAttempt {
907 future: Box::pin(async move {
908 crate::segment::fold_registered_pending(parts.manifest, fold)
909 .await
910 .map(ProvisionOutcome::Folded)
911 }),
912 replicas: Vec::new(),
913 })
914 } else if writer.spare_wanted() {
915 let parts = match writer.provision_parts() {
916 Ok(parts) => parts,
917 Err(error) => {
918 state.fail_queued(&mut receiver, error, &metrics);
919 break 'engine;
920 }
921 };
922 let id = writer.next_segment_id();
923 let replicas =
924 crate::segment::provision_replicas(&parts.factories, &parts.prefix, &id);
925 if rotation.has_pending_fold() {
926 let cancellation_replicas = replicas.clone();
927 Some(ProvisionAttempt {
928 future: Box::pin(async move {
929 crate::segment::provision_spare_with_replicas(
930 replicas,
931 parts.client_config,
932 parts.max_replica_lag_bytes,
933 parts.lane_stall_timeout,
934 id,
935 parts.metrics,
936 )
937 .await
938 .map(|(id, writer)| ProvisionOutcome::Unregistered(id, writer))
939 }),
940 replicas: cancellation_replicas,
941 })
942 } else {
943 let cancellation_replicas = replicas.clone();
944 Some(ProvisionAttempt {
945 future: Box::pin(async move {
946 crate::segment::provision_registered_spare_with_replicas(
947 parts, id, replicas,
948 )
949 .await
950 .map(|spare| ProvisionOutcome::Registered(Box::new(spare)))
951 }),
952 replicas: cancellation_replicas,
953 })
954 }
955 } else {
956 None
957 };
958 if let Some(attempt) = attempt {
959 if provision_requests.try_send(attempt).is_ok() {
960 spare_requested = true;
961 metrics.spare_provisioning_attempts.increment();
962 }
963 }
964 }
965
966 if rotation.swap_ready() && writer.spare_ready() && writer.swap_boundary_ready() {
967 match writer.begin_swap().await {
974 Ok(Some(swap)) => {
975 rotation = Rotation::Draining {
976 swap: Box::new(swap),
977 fold_attempts: 0,
978 retry: None,
979 fold_capacity_blocked: false,
980 successor_due: false,
981 };
982 active_capacity.send_replace(ActiveSegmentCapacity {
983 admission_start_bytes: state.last_dispatched_admission_bytes,
984 existing_bytes: writer.active_segment_bytes(),
985 seal_room: writer.active_segment_has_seal_room(),
986 });
987 continue 'engine;
991 }
992 Ok(None) => rotation = Rotation::Idle,
993 Err(error) => {
994 tracing::warn!(%error, "WAL engine rotation swap failed");
995 state.fail_queued(&mut receiver, error, &metrics);
996 break 'engine;
997 }
998 }
999 }
1000
1001 if state.pending_records == 0 && !matches!(state.queue.front(), Some(Command::Append(_))) {
1002 match state.queue.pop_front() {
1003 Some(Command::Shutdown { response }) => {
1004 fail_all(&mut state.queue, Error::Closed, &metrics);
1005 state.shutdown_response = Some(response);
1006 break;
1007 }
1008 Some(Command::Append(_)) => unreachable!(),
1009 None => {}
1010 }
1011 }
1012
1013 while !rotation.dispatch_paused()
1017 && matches!(state.queue.front(), Some(Command::Append(_)))
1018 && state.pending_records < config.pipeline_window_records
1019 {
1020 let outstanding = state.pending_records;
1021 let room = config.pipeline_window_records - outstanding;
1022 let mut appends = Vec::with_capacity(room);
1023 for _ in 0..room {
1024 if !matches!(state.queue.front(), Some(Command::Append(_))) {
1025 break;
1026 }
1027 let Some(Command::Append(append)) = state.queue.pop_front() else {
1028 unreachable!();
1029 };
1030 appends.push(append);
1031 }
1032 let records = appends.iter().map(|append| append.record.clone()).collect();
1033 let wal_record_bytes: u64 = appends
1034 .iter()
1035 .map(|append| append.encoded_bytes as u64)
1036 .sum();
1037 let batch_admission_end = appends
1038 .last()
1039 .map(|append| append.admission_end_bytes)
1040 .unwrap_or(state.last_dispatched_admission_bytes);
1041 match writer
1042 .enqueue_records_for_engine(records, Arc::clone(&on_attempted))
1043 .await
1044 {
1045 Ok(commits) => {
1046 for append in &mut appends {
1047 drop(append.queue_slot.take());
1048 }
1049 state.last_dispatched_admission_bytes = batch_admission_end;
1050 if outstanding > 0 {
1051 metrics.pipeline_refills.increment();
1052 }
1053 metrics.wal_record_bytes.add(wal_record_bytes);
1054 let expected_first = appends
1055 .first()
1056 .map(|append| append.seqno.record_index)
1057 .expect("an admitted engine batch is non-empty");
1058 let expected_end = appends
1059 .last()
1060 .map(|append| append.seqno.record_index + 1)
1061 .expect("an admitted engine batch is non-empty");
1062 if commits.first_global_record_index() != expected_first
1063 || commits.end_global_record_index() != expected_end
1064 {
1065 let error = Error::Internal(format!(
1066 "WAL assigned records {}..{}, caller admitted {expected_first}..{expected_end}",
1067 commits.first_global_record_index(),
1068 commits.end_global_record_index()
1069 ));
1070 state.stop_admission(&mut receiver);
1071 for append in appends {
1072 fail_append(append, error.clone(), &metrics);
1073 }
1074 fail_all(&mut state.queue, error, &metrics);
1075 metrics.operation_failures.increment();
1076 break 'engine;
1077 }
1078 state.pending_records += appends.len();
1079 state.completion_batches.push_back(CompletionBatch {
1080 commits,
1081 appends: appends.into(),
1082 });
1083 metrics
1084 .max_inflight_records
1085 .update_max(state.pending_records as u64);
1086 rotation.mark_due(writer.rotation_due(config.max_segment_bytes));
1087 }
1088 Err(error) => {
1089 tracing::warn!(%error, "WAL engine append pipeline failed");
1090 state.stop_admission(&mut receiver);
1091 for append in appends {
1092 fail_append(append, error.clone(), &metrics);
1093 }
1094 metrics.operation_failures.increment();
1095 fail_all(&mut state.queue, error, &metrics);
1096 break 'engine;
1097 }
1098 }
1099 }
1100 observe_queue_depth(&metrics, config.queue_capacity, &queue_slots);
1101
1102 if state.input_closed && state.queue.is_empty() && state.pending_records == 0 {
1103 break;
1104 }
1105
1106 metrics.rotation_state.set(rotation.metric_value());
1107 tokio::select! {
1117 biased;
1118 result = next_commit_update(&mut state.completion_batches), if state.pending_records != 0 => {
1119 if let Err(error) = result {
1120 state.stop_admission(&mut receiver);
1121 if let Some(append) = state.completion_batches
1122 .front_mut()
1123 .and_then(|batch| batch.appends.pop_front())
1124 {
1125 fail_append(append, error, &metrics);
1126 }
1127 fail_completion_batches(&mut state.completion_batches, Error::Poisoned, &metrics);
1128 fail_all(&mut state.queue, Error::Poisoned, &metrics);
1129 metrics.operation_failures.increment();
1130 break 'engine;
1131 }
1132 }
1133 event = async {
1137 match &mut rotation {
1138 Rotation::Sealing { enforced } => RotationEvent::Sealed(enforced.await),
1139 Rotation::Draining { retry: Some(retry), .. } => {
1140 retry.as_mut().await;
1141 RotationEvent::RetryElapsed
1142 }
1143 _ => unreachable!("rotation wait checked"),
1144 }
1145 }, if matches!(rotation, Rotation::Sealing { .. }) || rotation.retry_pending() => {
1146 match event {
1147 RotationEvent::Sealed(Ok(())) => {
1148 rotation = Rotation::Idle;
1149 rotation.mark_due(writer.rotation_due(config.max_segment_bytes));
1150 }
1151 RotationEvent::Sealed(Err(_)) => {
1152 tracing::warn!(
1153 "seal enforcement exhausted retries; rotation disabled until restart"
1154 );
1155 rotation = Rotation::Disabled {
1156 due: writer.rotation_due(config.max_segment_bytes),
1157 };
1158 metrics.operation_failures.increment();
1159 }
1160 RotationEvent::RetryElapsed => {
1161 if let Rotation::Draining { retry, .. } = &mut rotation {
1162 *retry = None;
1163 }
1164 }
1165 }
1166 }
1167 request = rotation_rechecks.recv(), if rotation_rechecks_open => {
1168 match request {
1169 Some(()) => {
1170 match writer.refresh_rotation_due(config.max_segment_bytes).await {
1171 Ok(due) => {
1172 rotation.release_fold_capacity_block();
1173 rotation.mark_due(due);
1174 active_capacity.send_modify(|capacity| {
1175 capacity.seal_room =
1176 writer.active_segment_has_seal_room();
1177 });
1178 }
1179 Err(error) => {
1180 tracing::warn!(
1181 %error,
1182 "failed to refresh rotation eligibility after truncation"
1183 );
1184 metrics.operation_failures.increment();
1185 }
1186 }
1187 }
1188 None => rotation_rechecks_open = false,
1189 }
1190 }
1191 result = provision_results.recv() => {
1195 spare_requested = false;
1196 match result {
1197 Some(Ok(ProvisionOutcome::Registered(spare))) => {
1198 writer.adopt_registered_spare(*spare);
1199 }
1200 Some(Ok(ProvisionOutcome::Unregistered(id, spare))) => {
1201 writer.adopt_unregistered_spare(id, spare);
1202 }
1203 Some(Ok(ProvisionOutcome::Folded(update))) => {
1204 let taken = std::mem::replace(&mut rotation, Rotation::Idle);
1205 let Rotation::Draining { swap, .. } = taken else {
1206 tracing::error!("pending fold completed without a rotation");
1207 state.fail_queued(&mut receiver, Error::Poisoned, &metrics);
1208 break 'engine;
1209 };
1210 if let Err(error) = writer.confirm_fold(&swap, update) {
1211 state.fail_queued(&mut receiver, error, &metrics);
1212 break 'engine;
1213 }
1214 active_capacity.send_modify(|capacity| {
1215 capacity.seal_room = writer.active_segment_has_seal_room();
1216 });
1217 let _ = catalog.send(writer.sealed_segments_snapshot());
1218 if let Some(segment) = swap.into_segment() {
1219 match maintenance.seal_segment(segment).await {
1220 Ok(enforced) => rotation = Rotation::Sealing { enforced },
1221 Err(error) => {
1222 tracing::warn!(
1223 %error,
1224 "maintenance task stopped before committed seal was queued"
1225 );
1226 rotation = Rotation::Disabled {
1227 due: writer.rotation_due(config.max_segment_bytes),
1228 };
1229 metrics.operation_failures.increment();
1230 }
1231 }
1232 } else {
1233 rotation = Rotation::Idle;
1234 rotation.mark_due(writer.rotation_due(config.max_segment_bytes));
1235 }
1236 }
1237 Some(Err(error)) => {
1238 if matches!(error, Error::SegmentDirectoryFull)
1239 && rotation.has_pending_fold()
1240 {
1241 rotation.block_fold_on_capacity();
1249 tracing::warn!(
1250 "manifest directory is full; pending fold waits for truncation"
1251 );
1252 } else if matches!(error, Error::Poisoned | Error::Fenced(_)) {
1253 metrics.spare_provisioning_failures.increment();
1254 tracing::warn!(%error, "background rotation work was fenced");
1255 state.fail_queued(&mut receiver, error, &metrics);
1256 break 'engine;
1257 } else {
1258 metrics.spare_provisioning_failures.increment();
1259 tracing::warn!(%error, "background rotation work failed; will retry");
1260 if rotation.has_pending_fold() && writer.unregistered_spare_ready() {
1261 rotation.arm_fold_retry(&client_config);
1262 }
1263 }
1264 }
1265 None => {
1266 metrics.spare_provisioning_failures.increment();
1267 tracing::warn!("spare provisioner stopped");
1268 state.fail_queued(&mut receiver, Error::Closed, &metrics);
1269 break 'engine;
1270 }
1271 }
1272 }
1273 command = receiver.recv(), if !state.input_closed => {
1274 match command {
1275 Some(command) => admit_command(
1276 command,
1277 &mut state.queue,
1278 &mut state.next_admission,
1279 ),
1280 None => state.input_closed = true,
1281 }
1282 }
1283 }
1284 }
1285 writer.shutdown_background_tasks().await;
1286 rotation.shutdown_background_tasks().await;
1287 if let Some(response) = state.shutdown_response {
1288 let _ = response.send(());
1289 tracing::info!("WAL engine shut down");
1290 } else {
1291 tracing::warn!("WAL engine stopped without graceful shutdown");
1292 }
1293 metrics.queue_depth.set(0);
1294 metrics.rotation_state.set(0);
1295}
1296
1297enum RotationEvent {
1301 Sealed(Result<(), oneshot::error::RecvError>),
1302 RetryElapsed,
1303}
1304
1305enum Rotation {
1307 Idle,
1309 Due,
1312 Draining {
1317 swap: Box<PendingSwap>,
1318 fold_attempts: usize,
1319 retry: Option<Pin<Box<Sleep>>>,
1320 fold_capacity_blocked: bool,
1324 successor_due: bool,
1328 },
1329 Sealing { enforced: oneshot::Receiver<()> },
1335 Disabled { due: bool },
1339}
1340
1341impl Rotation {
1342 fn metric_value(&self) -> i64 {
1343 match self {
1344 Self::Idle => 0,
1345 Self::Due => 1,
1346 Self::Draining { .. } => 2,
1347 Self::Sealing { .. } => 3,
1348 Self::Disabled { .. } => 4,
1349 }
1350 }
1351
1352 fn mark_due(&mut self, due: bool) {
1353 if !due {
1354 return;
1355 }
1356 match self {
1357 Rotation::Idle => *self = Rotation::Due,
1358 Rotation::Draining { successor_due, .. } => *successor_due = true,
1359 Rotation::Disabled { due: disabled_due } => *disabled_due = true,
1360 Rotation::Due | Rotation::Sealing { .. } => {}
1361 }
1362 }
1363
1364 fn swap_ready(&self) -> bool {
1365 matches!(self, Rotation::Due)
1366 }
1367
1368 fn fold_ready(&self, next_completion: u64) -> bool {
1369 matches!(
1370 self,
1371 Rotation::Draining {
1372 swap,
1373 retry: None,
1374 fold_capacity_blocked: false,
1375 ..
1376 }
1377 if next_completion > swap.end_record_index
1378 )
1379 }
1380
1381 fn swap(&self) -> Option<&PendingSwap> {
1382 match self {
1383 Rotation::Draining { swap, .. } => Some(swap),
1384 _ => None,
1385 }
1386 }
1387
1388 fn has_pending_fold(&self) -> bool {
1389 matches!(self, Rotation::Draining { .. })
1390 }
1391
1392 fn retry_pending(&self) -> bool {
1393 matches!(self, Rotation::Draining { retry: Some(_), .. })
1394 }
1395
1396 fn dispatch_paused(&self) -> bool {
1397 matches!(
1398 self,
1399 Rotation::Due
1400 | Rotation::Draining {
1401 successor_due: true,
1402 ..
1403 }
1404 | Rotation::Disabled { due: true }
1405 )
1406 }
1407
1408 fn arm_fold_retry(&mut self, config: &crate::protocol::ClientConfig) {
1409 if let Rotation::Draining {
1410 fold_attempts,
1411 retry,
1412 ..
1413 } = self
1414 {
1415 let delay = crate::protocol::retry_delay(config, *fold_attempts);
1416 *fold_attempts = fold_attempts.saturating_add(1);
1417 *retry = Some(Box::pin(tokio::time::sleep(delay)));
1418 }
1419 }
1420
1421 fn block_fold_on_capacity(&mut self) {
1422 if let Rotation::Draining {
1423 fold_capacity_blocked,
1424 ..
1425 } = self
1426 {
1427 *fold_capacity_blocked = true;
1428 }
1429 }
1430
1431 fn release_fold_capacity_block(&mut self) {
1432 if let Rotation::Draining {
1433 fold_capacity_blocked,
1434 ..
1435 } = self
1436 {
1437 *fold_capacity_blocked = false;
1438 }
1439 }
1440
1441 async fn shutdown_background_tasks(&mut self) {
1442 if let Rotation::Draining { swap, .. } = self {
1443 if let Some(writer) = &mut swap.writer {
1444 writer.shutdown_background_tasks().await;
1445 }
1446 }
1447 }
1448}
1449
1450enum ProvisionOutcome {
1451 Registered(Box<RegisteredSpare>),
1452 Unregistered(String, crate::protocol::Writer),
1453 Folded(ManifestUpdate),
1454}
1455
1456type ProvisionFuture = Pin<Box<dyn Future<Output = Result<ProvisionOutcome, Error>> + Send>>;
1457
1458struct ProvisionAttempt {
1459 future: ProvisionFuture,
1460 replicas: Vec<Arc<dyn Replica>>,
1461}
1462
1463async fn run_provisioner(
1469 mut requests: mpsc::Receiver<ProvisionAttempt>,
1470 results: mpsc::Sender<Result<ProvisionOutcome, Error>>,
1471 mut shutdown: watch::Receiver<bool>,
1472) {
1473 loop {
1474 let attempt = tokio::select! {
1475 biased;
1476 changed = shutdown.changed() => {
1477 if changed.is_err() || *shutdown.borrow_and_update() {
1478 return;
1479 }
1480 continue;
1481 }
1482 attempt = requests.recv() => {
1483 let Some(attempt) = attempt else {
1484 return;
1485 };
1486 attempt
1487 }
1488 };
1489 let ProvisionAttempt {
1490 mut future,
1491 replicas,
1492 } = attempt;
1493 let result = tokio::select! {
1494 biased;
1495 changed = shutdown.changed() => {
1496 if changed.is_err() || *shutdown.borrow_and_update() {
1497 None
1498 } else {
1499 continue;
1500 }
1501 }
1502 result = &mut future => Some(result),
1503 };
1504 let Some(result) = result else {
1505 cancel_provision_attempt(future, replicas).await;
1506 return;
1507 };
1508 tokio::select! {
1509 biased;
1510 changed = shutdown.changed() => {
1511 if changed.is_err() || *shutdown.borrow_and_update() {
1512 return;
1513 }
1514 }
1515 sent = results.send(result) => {
1516 if sent.is_err() {
1517 return;
1518 }
1519 }
1520 }
1521 }
1522}
1523
1524async fn cancel_provision_attempt(mut future: ProvisionFuture, replicas: Vec<Arc<dyn Replica>>) {
1525 if replicas.is_empty() {
1526 let _ = future.await;
1530 return;
1531 }
1532 let shutdown = join_all(replicas.iter().map(|replica| replica.shutdown()));
1533 let _ = tokio::join!(&mut future, shutdown);
1534}
1535
1536async fn next_commit_update(batches: &mut VecDeque<CompletionBatch>) -> Result<(), Error> {
1537 let batch = batches.front_mut().ok_or(Error::Poisoned)?;
1541 batch.commits.changed().await
1542}
1543
1544fn observe_queue_depth(metrics: &Metrics, queue_capacity: usize, queue_slots: &Semaphore) {
1545 metrics
1546 .queue_depth
1547 .set_usize(queue_capacity.saturating_sub(queue_slots.available_permits()));
1548}
1549
1550fn complete_append(append: AdmittedAppend, metrics: &Metrics) {
1551 metrics
1552 .append_commit_latency
1553 .record_duration(append.admitted_at.elapsed());
1554 let _ = append.completion.send(Ok(AppendReceipt {
1555 seqno: append.seqno,
1556 }));
1557 metrics.committed_records.increment();
1558 metrics.committed_bytes.add(append.payload_bytes as u64);
1559 metrics
1560 .committed_records_watermark
1561 .set_u64(append.seqno.record_index + 1);
1562}
1563
1564fn admit_command(command: Command, queue: &mut VecDeque<Command>, next_admission: &mut u64) {
1565 match command {
1566 Command::Append(append) => {
1567 debug_assert_eq!(append.seqno.record_index, *next_admission);
1568 queue.push_back(Command::Append(append));
1569 *next_admission += 1;
1570 }
1571 Command::Shutdown { response } => {
1572 queue.push_back(Command::Shutdown { response });
1573 }
1574 }
1575}
1576
1577fn fail_all(queue: &mut VecDeque<Command>, error: Error, metrics: &Metrics) {
1578 while let Some(command) = queue.pop_front() {
1579 match command {
1580 Command::Append(append) => fail_append(append, error.clone(), metrics),
1581 Command::Shutdown { response } => {
1582 let _ = response.send(());
1583 }
1584 }
1585 }
1586}
1587
1588fn fail_completion_batches(
1589 batches: &mut VecDeque<CompletionBatch>,
1590 error: Error,
1591 metrics: &Metrics,
1592) {
1593 while let Some(mut batch) = batches.pop_front() {
1594 while let Some(append) = batch.appends.pop_front() {
1595 fail_append(append, error.clone(), metrics);
1596 }
1597 }
1598}
1599
1600fn fail_append(append: AdmittedAppend, error: Error, metrics: &Metrics) {
1601 metrics.append_failures.increment();
1602 let _ = append.completion.send(Err(error));
1603}