Skip to main content

chorus_client/
engine.rs

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)]
26/// Capacity controls for the in-process transactional WAL pipeline.
27pub struct WalEngineConfig {
28    /// Maximum records waiting to enter the dispatched pipeline.
29    ///
30    /// This is one combined bound across the handle-to-engine channel and the
31    /// engine's internal queue, not a separate allowance for each stage. The
32    /// default is 256 records, eight times the default pipeline window.
33    pub queue_capacity: usize,
34    /// Maximum application payload bytes accepted in one record.
35    pub max_record_bytes: usize,
36    /// Maximum dispatched but not yet logically committed records. Slots are
37    /// replenished individually as ordered quorum completions arrive.
38    pub pipeline_window_records: usize,
39    /// Maximum encoded WAL bytes admitted but not yet resolved through the
40    /// ordered quorum completion stream. Admission waits for this byte budget
41    /// before taking ownership of another record.
42    pub max_inflight_bytes: usize,
43    /// Maximum unacknowledged encoded bytes retained for one replica lane.
44    /// A lane that falls farther behind is dropped so a healthy quorum can
45    /// continue without retaining an unbounded retry suffix. Must be at
46    /// least [`max_inflight_bytes`](Self::max_inflight_bytes).
47    pub max_replica_lag_bytes: usize,
48    /// Maximum interval in which a lane with retained writes may make no
49    /// durable-tail progress before it is shed.
50    ///
51    /// Any increase in `persisted_size` resets the interval, so a slow but
52    /// steadily advancing lane remains live. Shedding never changes the
53    /// strict-majority quorum requirement: if the remaining lanes cannot cover
54    /// an admitted record, the writer poisons instead of advancing its commit
55    /// watermark. The default is five seconds.
56    pub lane_stall_timeout: Duration,
57    /// Target encoded size for automatic segment rotation.
58    ///
59    /// The engine checks this threshold between records, so a segment may exceed
60    /// it by one complete encoded record. This is not a hard object-size limit.
61    ///
62    /// Treat this as the minimum segment-size floor for the single-pending
63    /// design. With sustained encoded throughput `T` bytes/s and worst-case
64    /// provision-plus-fold latency `L` seconds, choose at least `T * L` times
65    /// an operational safety factor. If the active pending segment crosses
66    /// this threshold before its refill is registered, dispatch pauses
67    /// fail-closed rather than acknowledging an unregistered successor.
68    pub max_segment_bytes: usize,
69    /// Hard encoded-byte ceiling for one active segment object.
70    ///
71    /// The default is 4 GiB, deliberately far below GCS's 5 TiB object limit.
72    /// Rotation normally occurs at [`max_segment_bytes`](Self::max_segment_bytes),
73    /// but the manifest's bounded sealed-segment directory can defer that
74    /// rotation until truncation removes retained entries. Admission stops
75    /// cleanly at this ceiling with [`Error::ActiveSegmentFull`] instead of
76    /// letting the active object grow toward the provider limit or poisoning
77    /// the healthy writer. The ceiling must fit one maximum encoded record and
78    /// be at least the advisory rotation target.
79    pub max_active_segment_bytes: usize,
80    /// Interval between background maintenance passes.
81    ///
82    /// Each startup and periodic pass runs the dead-incarnation sweep deferred
83    /// by recovery, retries deletion tombstones below the committed floor, then
84    /// repairs retained sealed segments. `None` disables periodic passes.
85    /// Startup cleanup and repair, plus targeted repair after a degraded
86    /// rotation, still run without timers.
87    pub repair_interval: Option<Duration>,
88    /// Maximum time graceful shutdown may spend draining accepted work and
89    /// joining owned background tasks before aborting them.
90    ///
91    /// The default is five minutes, long enough for the default storage retry
92    /// budget while still turning a wedged backend or task into a bounded error.
93    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)]
191/// Durability evidence for one caller-numbered record.
192///
193/// Success means the record reached a strict-majority quorum and every preceding
194/// sequence number committed.
195pub struct AppendReceipt {
196    /// Caller-assigned sequence number of the committed record.
197    pub seqno: WalSeqNo,
198}
199
200impl AppendReceipt {
201    /// Exclusive replay and checkpoint boundary after this append.
202    pub fn next_seqno(&self) -> WalSeqNo {
203        WalSeqNo::record(self.seqno.record_index + 1)
204    }
205}
206
207/// Owned durability future returned after a record is admitted in sequence.
208///
209/// Dropping this value does not cancel the append. The WAL continues processing
210/// admitted work so shutdown can drain it and recovery can preserve it. A
211/// failure may preserve a representative terminal [`crate::TransportCode`] from
212/// one failed lane. Use [`Error::may_have_committed`] to determine whether
213/// recovery must resolve the record's durable outcome.
214pub 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
230/// Starts the background append pipeline and maintenance task.
231pub struct WalEngine;
232
233/// Control plane for a running WAL engine.
234///
235/// This handle is intentionally single-owner and not cloneable. Calling
236/// [`enqueue_append`](Self::enqueue_append) through `&mut self` gives the
237/// database one linear admission point for its caller-assigned sequence
238/// numbers, while returned [`AppendCompletion`] values remain independently
239/// awaitable by the database's apply pipeline.
240pub 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    /// Cumulative encoded admission boundary after this record. The engine
266    /// publishes the last boundary assigned to an old segment when it swaps,
267    /// so the handle can charge every later queued record to the successor
268    /// without inspecting or racing the command queue.
269    admission_end_bytes: u128,
270    /// Admission time for the commit-latency histogram. `tokio::time`
271    /// follows the simulator's virtual clock under DST.
272    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    /// Cumulative engine-lifetime admission bytes preceding this segment.
362    admission_start_bytes: u128,
363    /// Bytes already present when the engine adopted this active segment.
364    ///
365    /// Normal recovery starts a fresh empty object. Tests and low-level users
366    /// may start the engine around a writer that already committed records,
367    /// which still need to count against the hard ceiling.
368    existing_bytes: usize,
369    /// Whether recovery can still add this active object to the manifest
370    /// directory after it gains a committed record.
371    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    /// Start the background engine around an already recovered writer.
386    ///
387    /// Configuration limits must be nonzero and internally consistent: byte
388    /// budgets must fit a maximum encoded record, replica lag must cover the
389    /// in-flight budget, and the hard active-segment ceiling must cover the
390    /// advisory rotation target. Invalid limits return [`Error::InvalidConfig`].
391    /// This function requires a running Tokio runtime and returns immediately
392    /// after spawning the task.
393    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        // Maintenance (sealed-segment repair, floor-committed truncation)
412        // runs on its own task, serialized internally, concurrent with the
413        // engine: it shares no mutable writer state — only catalog snapshots
414        // published through the watch channel and an epoch-free manifest
415        // handle of its own.
416        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    /// Admit one caller-numbered opaque record without waiting for durability.
486    ///
487    /// Before waiting, this method verifies that `seqno` is exactly the next
488    /// contiguous sequence number and that the record fits the configured
489    /// limits. If the manifest has no slot recovery could use to seal this
490    /// active object, admission returns [`Error::SegmentDirectoryFull`]. If the
491    /// active object is at its hard ceiling, admission returns
492    /// [`Error::ActiveSegmentFull`]. Neither condition consumes `seqno`;
493    /// truncation may free directory capacity, after which the same append can
494    /// be retried. Otherwise this waits for encoded-byte admission capacity and
495    /// a bounded queue slot. On success the WAL owns the record and the
496    /// returned [`AppendCompletion`] may be moved to another task. Await that
497    /// completion before applying the transaction to the database.
498    ///
499    /// The payload should be the database's complete encoded transaction batch.
500    /// The WAL treats it as opaque bytes and never merges it with another call.
501    /// Dropping the completion does not cancel admitted work. On an eventual
502    /// completion error, use [`Error::may_have_committed`] before deciding
503    /// whether recovery must resolve the sequence number.
504    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        // Sequence numbers cap an engine lifetime at fewer than 2^64 records
560        // and every encoded record is u32-bounded, so u128 has ample headroom.
561        // Keep the check anyway: protocol-invariant failures are errors, never
562        // panics.
563        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    /// Advance the application checkpoint floor and delete eligible sealed
605    /// segments.
606    ///
607    /// `floor` is the first record the database still needs. Call this only after
608    /// the database has durably checkpointed every record before that boundary.
609    /// The WAL deletes only whole sealed segments; it never truncates the active
610    /// segment or deletes a segment containing `floor`.
611    ///
612    /// Deletion is best effort across zones. The returned report contains only
613    /// deletion statistics; the database remains the authority for the durable
614    /// checkpoint it supplied. Startup and periodic maintenance retry
615    /// tombstones whose zones were unavailable during this call. A successful
616    /// truncation also makes the engine re-read directory capacity, so an active
617    /// segment held at [`Error::ActiveSegmentFull`] can rotate and resume
618    /// admission once enough entries are removed.
619    pub async fn truncate_before(&self, floor: WalSeqNo) -> Result<TruncationReport, Error> {
620        let report = self.maintenance.truncate(floor).await?;
621        // The maintenance manifest and writer manifest are independent
622        // handles. One pending wake is enough to make the engine refresh its
623        // copy after every completed truncation: if the slot is full, its
624        // queued wake has not been received yet; otherwise this call fills it.
625        let _ = self.rotation_recheck.try_send(());
626        Ok(report)
627    }
628
629    /// Consume the admission handle, abort every owned task, and await their
630    /// termination.
631    ///
632    /// Accepted completions may resolve as closed or ambiguous. Use
633    /// [`shutdown`](Self::shutdown) when accepted work must drain gracefully.
634    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    /// Consume the admission handle, drain accepted work, finish every
676    /// committed seal, and wait for all owned tasks to exit.
677    ///
678    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            // Provisioning is speculative. Once the engine has acknowledged
704            // its drain, signal the worker, cancel its in-flight attempt through
705            // the worker's select, and join both owned tasks.
706            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            // Stop periodic work explicitly, then let maintenance drain every
713            // queued seal before exiting. The dedicated signal cannot be starved
714            // by an overdue repair tick.
715            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    // Spare provisioning runs on a dedicated worker; its result channel is a
795    // `select!` wake source, so a swap waiting on a slow spare can never park
796    // the engine. `spare_requested` keeps at most one attempt outstanding.
797    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        // Consume the segment commit watermark directly. Every batch is in
809        // global record order, and each segment watermark is already a
810        // contiguous prefix, so no second completion reorder is needed.
811        // The active pending segment was registered before the swap, so its
812        // acknowledgments do not wait for the later fold.
813        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                    // The commit tracker and completion queue are separate
825                    // state machines. If they ever disagree, poison the WAL
826                    // rather than panicking the host process.
827                    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        // The dedicated worker owns all create/open and manifest work. Before
881        // a rotation it provisions then confirms `pending`; after a swap it
882        // provisions an unregistered refill, then atomically folds the old
883        // tail and registers that refill.
884        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            // Swap rotation, first half: once every admitted old-tail record is
968            // committed, wait only for the already-dispatched digest pipeline
969            // and route admissions to the spare in memory. There is no manifest
970            // write, object creation, or stream open here. Pinning the successor
971            // base to this committed boundary prevents a pending segment from
972            // being spliced above an unresolved global sequence gap.
973            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                    // The consumed pending slot now needs an unregistered
988                    // refill. Re-enter the loop so the worker request is
989                    // queued before the engine reaches its single wait point.
990                    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        // Dispatch pauses when rotation is due without a confirmed pending
1014        // segment, including when the consumed pending segment itself crosses
1015        // the threshold before fold/refill completes.
1016        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        // The single wait point: every event that can unblock the engine is a
1108        // branch here, so nothing the loop is waiting for can fail to wake it.
1109        //
1110        // `biased` makes the poll order deterministic: an unbiased `select!`
1111        // draws its starting branch from a thread-local RNG, so two replays
1112        // of one simulation seed could service simultaneously-ready arms in
1113        // different orders and diverge. Commit movement drains first, then the
1114        // rare rotation and provisioning wakes, then admission. The shared
1115        // record permits bound the channel and internal queue together.
1116        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            // Rotation wake sources share one arm (one `&mut rotation`
1134            // borrow). Seal enforcement and fold retry backoff must wake the
1135            // engine even when no append completion is in flight.
1136            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            // provisioning results arrive the same way: with every caller
1192            // blocked on a queued append no command wakes the loop, and a due
1193            // rotation must swap the moment its spare lands
1194            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                            // The consumed pending segment is already the
1242                            // active writer, so abandoning this fold would
1243                            // lose its manifest transition. Retrying cannot
1244                            // create capacity and a zero-delay test client can
1245                            // otherwise spin forever. Park the exact fold until
1246                            // truncate_before refreshes the directory and
1247                            // sends the rotation recheck above.
1248                            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
1297/// What the rotation wake arm of the engine's `select!` observed: the two
1298/// rotation states that wait on a future share one arm so only one `&mut
1299/// rotation` borrow exists across the `select!`.
1300enum RotationEvent {
1301    Sealed(Result<(), oneshot::error::RecvError>),
1302    RetryElapsed,
1303}
1304
1305/// Rotation control state for one preregistered pending segment.
1306enum Rotation {
1307    /// Active segment within its advisory size budget.
1308    Idle,
1309    /// Budget crossed. Dispatch is fail-closed until a confirmed pending
1310    /// segment is available for the in-memory swap.
1311    Due,
1312    /// Admissions already route to the adopted spare; the swapped-out
1313    /// segment is draining toward its admitted end. The active successor was
1314    /// preregistered before the swap, so its acknowledgments are independent
1315    /// of the later background fold.
1316    Draining {
1317        swap: Box<PendingSwap>,
1318        fold_attempts: usize,
1319        retry: Option<Pin<Box<Sleep>>>,
1320        /// A fold rejected for directory capacity cannot make progress until
1321        /// application truncation removes retained entries. Keep the exact
1322        /// transition parked instead of retrying a permanent local predicate.
1323        fold_capacity_blocked: bool,
1324        /// The consumed pending segment itself crossed the rotation floor
1325        /// before fold/refill completed. With no second pending slot, dispatch
1326        /// pauses until the fold registers the refill.
1327        successor_due: bool,
1328    },
1329    /// The swapped-out segment's seal enforcement is in flight. The next
1330    /// swap would overwrite the manifest seal record that segment needs for
1331    /// recovery, so rotation holds until `enforced` resolves — at most one
1332    /// sealed segment ever lacks a finalized quorum, and it is always the
1333    /// manifest's current seal record.
1334    Sealing { enforced: oneshot::Receiver<()> },
1335    /// Live sealing and bounded idempotent reconstruction both failed: no
1336    /// further rotation until restart recovery enforces the segment named by
1337    /// the manifest's seal record.
1338    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
1463/// Dedicated worker running spare-provisioning attempts one at a time. It
1464/// exists so the engine adopts spares through a channel — a `select!` wake
1465/// source — instead of polling a task handle whose completion could otherwise
1466/// arrive while no engine branch is waiting on it. Exits when the engine drops
1467/// the request sender.
1468async 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        // A fold has no data-plane sessions to close. Drive its guarded
1527        // manifest CAS to a terminal result so shutdown does not return while
1528        // a submitted control-plane mutation is still unresolved.
1529        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    // `pending_records` gates this call, but it is maintained independently
1538    // from the batch deque. Treat disagreement as poison so a protocol bug
1539    // shuts down this WAL without aborting the process.
1540    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}