Skip to main content

reddb_server/replication/
primary.rs

1//! Primary-side replication: WAL record production and snapshot serving.
2//!
3//! The logical WAL spool byte format is a `reddb-file` contract. This module
4//! owns runtime policy: appending after writes, syncing acknowledged records,
5//! serving replica pulls, and pruning once slots make records removable.
6
7use std::collections::{BTreeMap, VecDeque};
8use std::fs::{self, File, OpenOptions};
9use std::io::{self, Write};
10use std::path::{Path, PathBuf};
11use std::sync::{Arc, Condvar, Mutex, RwLock};
12use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};
13
14use reddb_file::{ReplicationSlot, ReplicationSlotInvalidationCause};
15use tracing::warn;
16
17mod slots;
18use slots::{
19    load_replication_slot_catalog, load_replication_slots, persist_replication_slot_catalog,
20    persist_replication_slots,
21};
22
23fn term_from_payload(payload: &[u8]) -> u64 {
24    crate::replication::cdc::ChangeRecord::decode(payload)
25        .map(|record| record.term)
26        .unwrap_or(crate::replication::DEFAULT_REPLICATION_TERM)
27}
28
29fn authority_from_payload(payload: &[u8]) -> (u64, Option<u64>) {
30    crate::replication::cdc::ChangeRecord::decode(payload)
31        .map(|record| (record.term, record.ownership_epoch))
32        .unwrap_or((crate::replication::DEFAULT_REPLICATION_TERM, None))
33}
34
35/// In-memory WAL buffer for replication.
36/// Primary appends records here; replicas consume from it.
37///
38/// Each record payload is stored behind an `Arc<[u8]>` so fan-out to
39/// multiple replicas shares a single heap allocation per record
40/// (issue #832): a pull clones the `Arc` handle, never the bytes, so
41/// adding replicas does not multiply the primary's send-buffer memory.
42pub struct WalBuffer {
43    /// Circular buffer of (lsn, ref-counted serialized record) pairs.
44    records: RwLock<VecDeque<(u64, Arc<[u8]>)>>,
45    /// Current write LSN.
46    current_lsn: RwLock<u64>,
47}
48
49impl WalBuffer {
50    pub fn new(max_size: usize) -> Self {
51        Self {
52            records: RwLock::new(VecDeque::with_capacity(max_size)),
53            current_lsn: RwLock::new(0),
54        }
55    }
56
57    /// Append a WAL record. Called by the storage engine after each write.
58    pub fn append(&self, lsn: u64, data: Vec<u8>) {
59        let mut records = self.records.write().unwrap_or_else(|e| e.into_inner());
60        records.push_back((lsn, Arc::from(data.into_boxed_slice())));
61
62        let mut current = self.current_lsn.write().unwrap_or_else(|e| e.into_inner());
63        *current = (*current).max(lsn);
64    }
65
66    /// Read records since the given LSN (exclusive), copying each
67    /// payload into an owned `Vec<u8>`. Kept for callers (WAL
68    /// archiving, retention bookkeeping) that need owned bytes; the
69    /// per-replica fan-out path should prefer [`Self::read_since_shared`]
70    /// to avoid copying.
71    pub fn read_since(&self, since_lsn: u64, max_count: usize) -> Vec<(u64, Vec<u8>)> {
72        self.read_since_shared(since_lsn, max_count)
73            .into_iter()
74            .map(|(lsn, data)| (lsn, data.to_vec()))
75            .collect()
76    }
77
78    /// Read records since the given LSN (exclusive) sharing the stored
79    /// `Arc<[u8]>` payloads. Fan-out to N replicas clones only the
80    /// reference-counted handles, so the buffer's bytes are never
81    /// duplicated per replica (issue #832).
82    pub fn read_since_shared(&self, since_lsn: u64, max_count: usize) -> Vec<(u64, Arc<[u8]>)> {
83        let records = self.records.read().unwrap_or_else(|e| e.into_inner());
84        records
85            .iter()
86            .filter(|(lsn, _)| *lsn > since_lsn)
87            .take(max_count)
88            .cloned()
89            .collect()
90    }
91
92    /// Current LSN.
93    pub fn current_lsn(&self) -> u64 {
94        *self.current_lsn.read().unwrap_or_else(|e| e.into_inner())
95    }
96
97    pub fn set_current_lsn(&self, lsn: u64) {
98        let mut current = self.current_lsn.write().unwrap_or_else(|e| e.into_inner());
99        *current = (*current).max(lsn);
100    }
101
102    pub fn prune_through(&self, upto_lsn: u64) {
103        let mut records = self.records.write().unwrap_or_else(|e| e.into_inner());
104        while records
105            .front()
106            .map(|(lsn, _)| *lsn <= upto_lsn)
107            .unwrap_or(false)
108        {
109            records.pop_front();
110        }
111    }
112
113    /// Oldest available LSN (for gap detection).
114    pub fn oldest_lsn(&self) -> Option<u64> {
115        let records = self.records.read().unwrap_or_else(|e| e.into_inner());
116        records.front().map(|(lsn, _)| *lsn)
117    }
118}
119
120fn logical_wal_entry_term(entry: &reddb_file::LogicalWalEntry) -> u64 {
121    if entry.term == 0 {
122        term_from_payload(&entry.data)
123    } else {
124        entry.term
125    }
126}
127
128fn logical_wal_data_with_framing_term(entry: &reddb_file::LogicalWalEntry) -> Vec<u8> {
129    let term = logical_wal_entry_term(entry);
130    match crate::replication::cdc::ChangeRecord::decode(&entry.data) {
131        Ok(mut record)
132            if record.term != term || record.ownership_epoch != entry.ownership_epoch =>
133        {
134            record.term = term;
135            if entry.ownership_epoch.is_some() {
136                record.ownership_epoch = entry.ownership_epoch;
137            }
138            record.encode()
139        }
140        _ => entry.data.clone(),
141    }
142}
143
144/// One in every `SEEK_INDEX_INTERVAL` records is checkpointed into the
145/// spool's in-memory seek index. A briefly-disconnected replica
146/// resuming from its slot LSN binary-searches this sparse index and
147/// seeks straight to the nearest preceding checkpoint, then scans
148/// forward at most `SEEK_INDEX_INTERVAL` records — turning resume from
149/// an O(n) full-file scan into a sub-linear seek (issue #832). The
150/// index is rebuilt on `open` and extended on every `append`.
151#[derive(Debug, Default)]
152struct LogicalWalSpoolState {
153    current_lsn: u64,
154    /// Sparse, strictly LSN-ascending `(lsn, byte_offset)` checkpoints
155    /// into the spool file. `byte_offset` is the start of the record
156    /// whose LSN is `lsn`.
157    seek_index: Vec<(u64, u64)>,
158    /// Byte length of the spool file (offset at which the next append
159    /// lands). Tracked so `append` can record a checkpoint's offset
160    /// without an extra `stat`.
161    write_offset: u64,
162    /// Total records appended/recovered, used to space checkpoints
163    /// `SEEK_INDEX_INTERVAL` records apart.
164    record_count: u64,
165}
166
167impl LogicalWalSpoolState {
168    /// Push a checkpoint for the record at `offset` if it falls on a
169    /// `SEEK_INDEX_INTERVAL` boundary. `ordinal` is the record's
170    /// zero-based position in the spool.
171    fn note_record(&mut self, ordinal: u64, lsn: u64, offset: u64) {
172        if ordinal.is_multiple_of(reddb_file::LOGICAL_WAL_SEEK_INDEX_INTERVAL) {
173            // Keep the index strictly ascending even if LSNs repeat
174            // (they should not, but a defensive guard keeps the binary
175            // search total).
176            if self.seek_index.last().map(|(l, _)| *l) != Some(lsn) {
177                self.seek_index.push((lsn, offset));
178            }
179        }
180    }
181
182    /// Byte offset to start a forward scan from when resuming at
183    /// `since_lsn` (exclusive). Returns the offset of the latest
184    /// checkpoint whose LSN is `<= since_lsn`, or `0` when no such
185    /// checkpoint exists.
186    fn seek_floor_offset(&self, since_lsn: u64) -> u64 {
187        match self
188            .seek_index
189            .binary_search_by(|(lsn, _)| lsn.cmp(&since_lsn))
190        {
191            Ok(idx) => self.seek_index[idx].1,
192            Err(0) => 0,
193            Err(idx) => self.seek_index[idx - 1].1,
194        }
195    }
196}
197
198/// Durable append-only logical WAL spool kept beside the main `.rdb` file.
199///
200/// This is not the storage-engine WAL; it is a structured replication/PITR log.
201pub struct LogicalWalSpool {
202    path: PathBuf,
203    state: Mutex<LogicalWalSpoolState>,
204}
205
206impl LogicalWalSpool {
207    pub fn path_for(data_path: &Path) -> PathBuf {
208        reddb_file::layout::logical_wal_path(data_path)
209    }
210
211    pub fn open(data_path: &Path) -> io::Result<Self> {
212        let path = Self::path_for(data_path);
213        if let Some(parent) = path.parent() {
214            fs::create_dir_all(parent)?;
215        }
216        if !path.exists() {
217            File::create(&path)?;
218        }
219        // Recover-or-truncate to the longest valid prefix. A torn tail
220        // from the previous process exit (power loss, OOM kill, ENOSPC
221        // mid-write) is silently dropped; the warning surfaces to the
222        // operator log but the spool stays open.
223        let entries = reddb_file::read_and_repair_logical_wal_entries(&path)?;
224        let current_lsn = entries.last().map(|entry| entry.lsn).unwrap_or(0);
225        // Rebuild the sparse seek index from the (now repaired) file so
226        // a post-restart resume is sub-linear from the first pull.
227        let (seek_index, write_offset, record_count) =
228            reddb_file::build_logical_wal_seek_index(&path)?;
229        Ok(Self {
230            path,
231            state: Mutex::new(LogicalWalSpoolState {
232                current_lsn,
233                seek_index,
234                write_offset,
235                record_count,
236            }),
237        })
238    }
239
240    pub fn append(&self, lsn: u64, data: &[u8]) -> io::Result<()> {
241        let timestamp_ms = SystemTime::now()
242            .duration_since(UNIX_EPOCH)
243            .map(|d| d.as_millis() as u64)
244            .unwrap_or(0);
245        self.append_with_timestamp(lsn, timestamp_ms, data)
246    }
247
248    /// Append a record with an explicit framing timestamp. Used in
249    /// tests to produce deterministic timestamps; production callers
250    /// should use `append`.
251    pub fn append_with_timestamp(
252        &self,
253        lsn: u64,
254        timestamp_ms: u64,
255        data: &[u8],
256    ) -> io::Result<()> {
257        let (term, ownership_epoch) = authority_from_payload(data);
258        self.append_with_term_epoch_and_timestamp(term, ownership_epoch, lsn, timestamp_ms, data)
259    }
260
261    pub fn append_with_term_and_timestamp(
262        &self,
263        term: u64,
264        lsn: u64,
265        timestamp_ms: u64,
266        data: &[u8],
267    ) -> io::Result<()> {
268        let (_, ownership_epoch) = authority_from_payload(data);
269        self.append_with_term_epoch_and_timestamp(term, ownership_epoch, lsn, timestamp_ms, data)
270    }
271
272    pub fn append_with_term_epoch_and_timestamp(
273        &self,
274        term: u64,
275        ownership_epoch: Option<u64>,
276        lsn: u64,
277        timestamp_ms: u64,
278        data: &[u8],
279    ) -> io::Result<()> {
280        let mut file = OpenOptions::new()
281            .create(true)
282            .append(true)
283            .open(&self.path)?;
284        // Pre-build the record in memory so a single write_all keeps
285        // the on-disk record contiguous. Two side-effects:
286        //   (a) crash recovery sees either a complete record or a torn
287        //       header, never an interleaved partial frame from two
288        //       writers (the spool is not multi-writer today, but the
289        //       single-write semantics make that future-safe);
290        //   (b) crc32 is computed exactly once over the same bytes the
291        //       reader will checksum, with zero risk of header/payload
292        //       drift from a partial flush.
293        let frame =
294            reddb_file::encode_logical_wal_v3(term, ownership_epoch, lsn, timestamp_ms, data)?;
295
296        file.write_all(&frame)?;
297        // PLAN.md Phase 2 mandates `sync_all` for logical WAL durability.
298        // `flush()` only drains the std::io userspace buffer; without
299        // `sync_all` the kernel page cache may still be dirty when an
300        // acknowledged write supposedly committed.
301        file.sync_all()?;
302
303        let mut state = self.state.lock().unwrap_or_else(|e| e.into_inner());
304        state.current_lsn = state.current_lsn.max(lsn);
305        // The record we just wrote starts at the prior end-of-file.
306        // Checkpoint it into the seek index if it lands on an interval
307        // boundary, then advance the tracked write offset.
308        let record_start = state.write_offset;
309        let ordinal = state.record_count;
310        state.note_record(ordinal, lsn, record_start);
311        state.write_offset = record_start + frame.len() as u64;
312        state.record_count = ordinal + 1;
313        Ok(())
314    }
315
316    pub fn read_since(&self, since_lsn: u64, max_count: usize) -> io::Result<Vec<(u64, Vec<u8>)>> {
317        // Seek straight to the nearest indexed checkpoint at or before
318        // `since_lsn` instead of scanning the whole spool from offset 0
319        // (issue #832). The file was already repaired on `open`, so the
320        // forward scan from the checkpoint is non-repairing and simply
321        // stops at the first torn tail (left for the next `open` to fix).
322        let start_offset = {
323            let state = self.state.lock().unwrap_or_else(|e| e.into_inner());
324            state.seek_floor_offset(since_lsn)
325        };
326        let entries = reddb_file::read_logical_wal_entries_from(&self.path, start_offset)?;
327        Ok(entries
328            .into_iter()
329            .filter(|entry| entry.lsn > since_lsn)
330            .take(max_count)
331            .map(|entry| (entry.lsn, logical_wal_data_with_framing_term(&entry)))
332            .collect())
333    }
334
335    /// Byte offset a resume at `since_lsn` would seek to before
336    /// forward-scanning. Exposed for tests asserting the resume is
337    /// sub-linear (starts past offset 0 for a mid-spool LSN).
338    #[cfg(test)]
339    fn seek_floor_offset(&self, since_lsn: u64) -> u64 {
340        self.state
341            .lock()
342            .unwrap_or_else(|e| e.into_inner())
343            .seek_floor_offset(since_lsn)
344    }
345
346    pub fn current_lsn(&self) -> u64 {
347        self.state
348            .lock()
349            .unwrap_or_else(|e| e.into_inner())
350            .current_lsn
351    }
352
353    pub fn oldest_lsn(&self) -> io::Result<Option<u64>> {
354        Ok(reddb_file::read_and_repair_logical_wal_entries(&self.path)?
355            .into_iter()
356            .next()
357            .map(|entry| entry.lsn))
358    }
359
360    pub fn prune_through(&self, upto_lsn: u64) -> io::Result<()> {
361        let previous_lsn = self.current_lsn();
362        let mut retained: Vec<_> = reddb_file::read_and_repair_logical_wal_entries(&self.path)?
363            .into_iter()
364            .filter(|entry| entry.lsn > upto_lsn)
365            .collect();
366        for entry in &mut retained {
367            entry.term = logical_wal_entry_term(entry);
368        }
369        let temp_path = reddb_file::layout::logical_wal_temp_path(&self.path);
370        for entry in &mut retained {
371            // Re-frame as v3 so the spool only ever contains current records
372            // after a prune. Legacy v1 records are upgraded by carrying
373            // their original LSN and default term forward; the framing timestamp is
374            // re-stamped to wall-clock-now because the original v1
375            // record didn't carry one — downstream consumers that need
376            // the operation's logical timestamp continue to use the
377            // payload's own ChangeRecord::timestamp field.
378            let timestamp_ms = if entry.timestamp_ms > 0 {
379                entry.timestamp_ms
380            } else {
381                SystemTime::now()
382                    .duration_since(UNIX_EPOCH)
383                    .map(|d| d.as_millis() as u64)
384                    .unwrap_or(0)
385            };
386            entry.timestamp_ms = timestamp_ms;
387        }
388        let current_lsn =
389            reddb_file::rewrite_logical_wal_entries(&self.path, &temp_path, &retained)?;
390
391        // The rewrite shifted every record's byte offset, so the old
392        // seek index is stale — rebuild it from the compacted file.
393        let (seek_index, write_offset, record_count) =
394            reddb_file::build_logical_wal_seek_index(&self.path)?;
395        let mut state = self.state.lock().unwrap_or_else(|e| e.into_inner());
396        state.current_lsn = previous_lsn.max(current_lsn).max(upto_lsn);
397        state.seek_index = seek_index;
398        state.write_offset = write_offset;
399        state.record_count = record_count;
400        Ok(())
401    }
402}
403
404/// State of a connected replica. PLAN.md Phase 11.4 fields:
405/// `last_seen_at_unix_ms` updates on every interaction (pull or ack);
406/// `last_sent_lsn` updates when the primary serves a `pull_wal_records`
407/// batch; `last_durable_lsn` updates when the replica reports its WAL
408/// is durably written via `ack_replica_lsn`.
409#[derive(Debug, Clone)]
410pub struct ReplicaState {
411    pub id: String,
412    pub last_acked_lsn: u64,
413    pub last_sent_lsn: u64,
414    pub last_durable_lsn: u64,
415    pub apply_error_count: u64,
416    pub divergence_count: u64,
417    pub connected_at_unix_ms: u128,
418    pub last_seen_at_unix_ms: u128,
419    /// Region identifier declared by the replica at handshake time
420    /// (Phase 2.6 multi-region PG parity). `None` until the replica
421    /// handshake extension lands in 2.6.2; the quorum coordinator's
422    /// region-binding map covers the in-process case meanwhile.
423    pub region: Option<String>,
424    /// `true` while this replica is re-bootstrapping — loading a fresh
425    /// snapshot to replace its current dataset (issue #837). It keeps
426    /// serving non-causal reads from the old data, but the advertiser
427    /// surfaces this flag so a causal reader routes bookmark reads
428    /// elsewhere: the replica's `last_acked_lsn` describes data it is
429    /// about to discard. Cleared atomically when the swap completes.
430    pub rebootstrapping: bool,
431}
432
433/// Primary-side replication progress derived from the replica registry.
434#[derive(Debug, Clone, Copy, PartialEq, Eq)]
435pub struct ReplicationProgress {
436    pub lag_lsn: u64,
437    pub safe_replay_lsn: u64,
438}
439
440impl ReplicationProgress {
441    pub fn from_replicas(replicas: &[ReplicaState]) -> Option<Self> {
442        let max_sent_lsn = replicas.iter().map(|replica| replica.last_sent_lsn).max()?;
443        let min_acked_lsn = replicas
444            .iter()
445            .map(|replica| replica.last_acked_lsn)
446            .min()?;
447        let safe_replay_lsn = replicas
448            .iter()
449            .map(|replica| replica.last_durable_lsn)
450            .min()?;
451
452        Some(Self {
453            lag_lsn: max_sent_lsn.saturating_sub(min_acked_lsn),
454            safe_replay_lsn,
455        })
456    }
457}
458
459/// Primary replication manager.
460pub struct PrimaryReplication {
461    pub wal_buffer: Arc<WalBuffer>,
462    pub logical_wal_spool: Option<Arc<LogicalWalSpool>>,
463    pub replicas: RwLock<Vec<ReplicaState>>,
464    wal_appended: (Mutex<u64>, Condvar),
465    slot_path: Option<PathBuf>,
466    slot_catalog_path: Option<PathBuf>,
467    primary_replica_file_plan: Option<reddb_file::PrimaryReplicaFilePlan>,
468    primary_replica_wal_lock: Mutex<()>,
469    slots: RwLock<BTreeMap<String, ReplicationSlot>>,
470    slot_retention_max_lag_lsn: u64,
471    slot_idle_timeout_ms: u64,
472    /// PLAN.md Phase 11.4 — ack-driven commit synchronization. Always
473    /// allocated so the policy enum can flip from `Local` to
474    /// `AckN`/`Quorum` without touching this struct's shape.
475    pub commit_waiter: Arc<crate::replication::commit_waiter::CommitWaiter>,
476    /// Monotonic registry-change counter consumed by the
477    /// `TopologyAdvertiser` (issue #167). Bumps on register,
478    /// unregister, and the periodic health sweep when a replica
479    /// flips between healthy/unhealthy. Clients use the epoch to
480    /// detect stale advertisements without comparing the full
481    /// replica list element-wise.
482    topology_epoch: std::sync::atomic::AtomicU64,
483    /// Count of pulls served as a partial resync — a replica resuming
484    /// incrementally from its retained slot position rather than
485    /// triggering a full re-bootstrap (issue #832). Surfaced as a
486    /// replication metric so a brief disconnect that recovers via
487    /// partial resync is observable.
488    partial_resync_count: std::sync::atomic::AtomicU64,
489    /// Count of pulls that forced a full re-bootstrap — the replica's
490    /// retained WAL no longer covers its requested position, so it must
491    /// discard its dataset and reload a fresh snapshot (issue #839).
492    /// This is the primary alert signal: a healthy cluster re-bootstraps
493    /// rarely, so any sustained rise means slots are being invalidated
494    /// faster than replicas can keep up.
495    full_resync_count: std::sync::atomic::AtomicU64,
496}
497
498/// How a replica's pull should be served, decided from its slot state.
499#[derive(Debug, Clone, Copy, PartialEq, Eq)]
500pub enum ResumeMode {
501    /// Resume incrementally from `resume_lsn` (the replica's slot
502    /// position, never behind it). The retained WAL still covers the
503    /// gap, so a brief disconnect costs only a partial resync.
504    PartialResync { resume_lsn: u64 },
505    /// The slot is past the retention cap (or otherwise invalidated);
506    /// the replica must discard and re-bootstrap from a fresh snapshot.
507    FullRebootstrap {
508        cause: ReplicationSlotInvalidationCause,
509    },
510}
511
512impl PrimaryReplication {
513    pub fn slot_path_for(data_path: &Path) -> PathBuf {
514        reddb_file::layout::legacy_logical_slots_path(data_path)
515    }
516
517    pub fn primary_replica_root_for(data_path: &Path) -> PathBuf {
518        reddb_file::layout::primary_replica_root(data_path)
519    }
520
521    pub fn slot_catalog_path_for(data_path: &Path) -> PathBuf {
522        Self::primary_replica_file_plan_for(data_path).slots_path()
523    }
524
525    fn primary_replica_file_plan_for(data_path: &Path) -> reddb_file::PrimaryReplicaFilePlan {
526        let root = Self::primary_replica_root_for(data_path);
527        let timeline =
528            Self::primary_replica_current_timeline_for_root(&root).unwrap_or_else(|err| {
529                warn!(
530                    target: "reddb::replication::primary",
531                    error = %err,
532                    "failed to read primary-replica timeline history; using initial timeline"
533                );
534                reddb_file::TimelineId::initial()
535            });
536        reddb_file::PrimaryReplicaFilePlan::new(root, timeline)
537    }
538
539    fn primary_replica_current_timeline_for_root(
540        root: &Path,
541    ) -> Result<reddb_file::TimelineId, reddb_file::RdbFileError> {
542        let path = reddb_file::PrimaryReplicaFilePlan::new(root, reddb_file::TimelineId::initial())
543            .timeline_history_path();
544        match reddb_file::TimelineHistory::read_from_path(&path) {
545            Ok(history) => Ok(history
546                .current()
547                .unwrap_or_else(reddb_file::TimelineId::initial)),
548            Err(reddb_file::RdbFileError::Io(err))
549                if err.kind() == std::io::ErrorKind::NotFound =>
550            {
551                Ok(reddb_file::TimelineId::initial())
552            }
553            Err(err) => Err(err),
554        }
555    }
556
557    pub fn new(data_path: Option<&Path>) -> Self {
558        Self::new_with_config(data_path, &crate::replication::ReplicationConfig::primary())
559    }
560
561    pub fn new_with_config(
562        data_path: Option<&Path>,
563        config: &crate::replication::ReplicationConfig,
564    ) -> Self {
565        let now_ms = crate::utils::now_unix_millis() as u128;
566        let slot_path = data_path.map(Self::slot_path_for);
567        let slot_catalog_path = data_path.map(Self::slot_catalog_path_for);
568        let primary_replica_file_plan = data_path.map(Self::primary_replica_file_plan_for);
569        let mut slots = load_replication_slot_catalog(slot_catalog_path.as_deref(), now_ms);
570        slots.extend(load_replication_slots(slot_path.as_deref(), now_ms));
571        let logical_wal_spool = data_path
572            .and_then(|path| LogicalWalSpool::open(path).ok())
573            .map(Arc::new);
574        let current_lsn = logical_wal_spool
575            .as_ref()
576            .map(|spool| spool.current_lsn())
577            .unwrap_or(0);
578        Self {
579            wal_buffer: Arc::new(WalBuffer::new(100_000)),
580            logical_wal_spool,
581            replicas: RwLock::new(Vec::new()),
582            wal_appended: (Mutex::new(current_lsn), Condvar::new()),
583            slot_path,
584            slot_catalog_path,
585            primary_replica_file_plan,
586            primary_replica_wal_lock: Mutex::new(()),
587            slots: RwLock::new(slots),
588            slot_retention_max_lag_lsn: config.slot_retention_max_lag_lsn,
589            slot_idle_timeout_ms: config.slot_idle_timeout_ms,
590            commit_waiter: Arc::new(crate::replication::commit_waiter::CommitWaiter::new()),
591            topology_epoch: std::sync::atomic::AtomicU64::new(0),
592            partial_resync_count: std::sync::atomic::AtomicU64::new(0),
593            full_resync_count: std::sync::atomic::AtomicU64::new(0),
594        }
595    }
596
597    pub fn append_logical_record(&self, lsn: u64, encoded: Vec<u8>) {
598        self.wal_buffer.append(lsn, encoded.clone());
599        if let Some(spool) = &self.logical_wal_spool {
600            let _ = spool.append(lsn, &encoded);
601        }
602        if let Some(plan) = &self.primary_replica_file_plan {
603            let _guard = self
604                .primary_replica_wal_lock
605                .lock()
606                .unwrap_or_else(|err| err.into_inner());
607            match Self::primary_replica_current_timeline_for_root(&plan.root) {
608                Ok(timeline) => {
609                    let plan = reddb_file::PrimaryReplicaFilePlan::new(plan.root.clone(), timeline);
610                    if let Err(err) = plan.append_wal_record(lsn, &encoded) {
611                        warn!(
612                            target: "reddb::replication::primary",
613                            lsn,
614                            error = %err,
615                            "failed to append primary-replica WAL segment"
616                        );
617                    }
618                }
619                Err(err) => {
620                    warn!(
621                        target: "reddb::replication::primary",
622                        lsn,
623                        error = %err,
624                        "failed to read primary-replica timeline history; skipping WAL append"
625                    );
626                }
627            }
628        }
629        let (lock, cvar) = &self.wal_appended;
630        let mut latest = lock.lock().unwrap_or_else(|e| e.into_inner());
631        *latest = (*latest).max(lsn);
632        cvar.notify_all();
633    }
634
635    pub fn wait_for_logical_lsn_after(&self, since_lsn: u64, timeout: Duration) -> bool {
636        if self.current_logical_lsn() > since_lsn {
637            return true;
638        }
639        let deadline = Instant::now() + timeout;
640        let (lock, cvar) = &self.wal_appended;
641        let mut latest = lock.lock().unwrap_or_else(|e| e.into_inner());
642        while *latest <= since_lsn {
643            let now = Instant::now();
644            if now >= deadline {
645                return false;
646            }
647            let remaining = deadline.saturating_duration_since(now);
648            let (guard, result) = cvar
649                .wait_timeout(latest, remaining)
650                .unwrap_or_else(|e| e.into_inner());
651            latest = guard;
652            if result.timed_out() && *latest <= since_lsn {
653                return false;
654            }
655        }
656        true
657    }
658
659    pub fn register_replica(&self, id: String) -> u64 {
660        self.register_replica_with_region(id, None)
661    }
662
663    /// Register a replica with an explicit region tag (Phase 2.6 multi-region).
664    ///
665    /// Preferred when the replica handshake declares a region — the quorum
666    /// coordinator uses this field to decide whether the replica counts
667    /// toward a `QuorumMode::Regions` commit.
668    ///
669    /// Idempotent on reconnect (issue #812): if a replica with `id` is
670    /// already registered, the existing entry is *updated in place* rather
671    /// than duplicated — progress LSNs (`last_acked_lsn`, `last_sent_lsn`,
672    /// `last_durable_lsn`) are preserved so a reconnecting replica is not
673    /// rewound, only `last_seen_at_unix_ms` is refreshed (and `region` when
674    /// a non-`None` value is supplied). A re-registration is not a
675    /// registry-shape change, so it does **not** bump the topology epoch.
676    /// Returns the slot `restart_lsn` the replica should resume streaming from:
677    /// the current WAL LSN for a fresh registration, or the durable slot
678    /// restart point for a reconnect.
679    pub fn register_replica_with_region(&self, id: String, region: Option<String>) -> u64 {
680        self.register_replica_with_region_at(id, region, self.current_logical_lsn())
681    }
682
683    fn register_replica_with_region_at(
684        &self,
685        id: String,
686        region: Option<String>,
687        initial_lsn: u64,
688    ) -> u64 {
689        let now_ms = crate::utils::now_unix_millis() as u128;
690        let resume_lsn = self.ensure_slot(&id, initial_lsn);
691        let mut replicas = self.replicas.write().unwrap_or_else(|e| e.into_inner());
692        if let Some(existing) = replicas.iter_mut().find(|r| r.id == id) {
693            existing.last_seen_at_unix_ms = now_ms;
694            if region.is_some() {
695                existing.region = region;
696            }
697            return resume_lsn;
698        }
699        replicas.push(ReplicaState {
700            id,
701            last_acked_lsn: resume_lsn,
702            last_sent_lsn: resume_lsn,
703            last_durable_lsn: resume_lsn,
704            apply_error_count: 0,
705            divergence_count: 0,
706            connected_at_unix_ms: now_ms,
707            last_seen_at_unix_ms: now_ms,
708            region,
709            rebootstrapping: false,
710        });
711        drop(replicas);
712        self.bump_topology_epoch();
713        resume_lsn
714    }
715
716    /// Mark (or clear) a replica's re-bootstrap state (issue #837).
717    ///
718    /// While `rebootstrapping` is `true` the replica keeps serving
719    /// non-causal reads from its existing data, but the advertiser
720    /// surfaces the flag so causal (bookmark) reads route to a
721    /// caught-up peer instead — the rebuilding replica's applied
722    /// frontier describes data it is about to discard. The primary
723    /// flips this back to `false` when the replica reports its atomic
724    /// snapshot swap complete.
725    ///
726    /// A change to the flag is a registry-shape change for routing
727    /// purposes, so it bumps the topology epoch to force consumers to
728    /// re-read the advertisement. Returns `true` when a replica with
729    /// `id` was present and updated.
730    pub fn set_replica_rebootstrapping(&self, id: &str, rebootstrapping: bool) -> bool {
731        let mut replicas = self.replicas.write().unwrap_or_else(|e| e.into_inner());
732        let Some(state) = replicas.iter_mut().find(|r| r.id == id) else {
733            return false;
734        };
735        if state.rebootstrapping == rebootstrapping {
736            return true;
737        }
738        state.rebootstrapping = rebootstrapping;
739        drop(replicas);
740        self.bump_topology_epoch();
741        true
742    }
743
744    /// Ensure a replica identifying itself with `id` is present in the
745    /// registry (issue #812). This is the production self-registration hook
746    /// used by the `pull_wal_records` path: the first time a replica sends
747    /// its `replica_id` on a pull, the primary registers it so it is no
748    /// longer blind to that replica's existence; subsequent pulls are
749    /// idempotent no-ops. Returns `true` when a new registration was
750    /// created. A fresh slot starts at the replica's requested cursor rather
751    /// than the primary's current LSN, so an empty replica does not skip WAL
752    /// that was written before its first pull. Reconnects preserve existing
753    /// progress and do not bump the topology epoch.
754    pub fn ensure_replica_registered(&self, id: &str, requested_since_lsn: u64) -> bool {
755        let already = self
756            .replicas
757            .read()
758            .unwrap_or_else(|e| e.into_inner())
759            .iter()
760            .any(|r| r.id == id);
761        if already {
762            return false;
763        }
764        self.register_replica_with_region_at(id.to_string(), None, requested_since_lsn);
765        true
766    }
767
768    /// Unregister a replica by id. Returns `true` when the replica
769    /// was present (and removed). Bumps the topology epoch so a
770    /// pending advertisement reflects the new fleet size.
771    pub fn unregister_replica(&self, id: &str) -> bool {
772        let mut replicas = self.replicas.write().unwrap_or_else(|e| e.into_inner());
773        let before = replicas.len();
774        replicas.retain(|r| r.id != id);
775        let removed = replicas.len() != before;
776        drop(replicas);
777        if removed {
778            self.commit_waiter.drop_replica(id);
779            self.bump_topology_epoch();
780        }
781        removed
782    }
783
784    /// Current topology epoch. Strictly monotonic, bumps on every
785    /// registry-shape change consumed by `TopologyAdvertiser`.
786    pub fn topology_epoch(&self) -> u64 {
787        self.topology_epoch
788            .load(std::sync::atomic::Ordering::Relaxed)
789    }
790
791    /// Advance the topology epoch. Call sites: register, unregister,
792    /// and the health-sweep tick that flips a replica between
793    /// healthy/unhealthy. Wrapping is not a concern in practice
794    /// (`u64::MAX` events would take centuries at any realistic ack
795    /// rate) but `fetch_add` saturates implicitly via wrap-around;
796    /// the consumer treats epoch as opaque so a wrap is still
797    /// strictly "different" from the previous value.
798    pub fn bump_topology_epoch(&self) {
799        self.topology_epoch
800            .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
801    }
802
803    pub fn ack_replica(&self, id: &str, lsn: u64) {
804        let now_ms = crate::utils::now_unix_millis() as u128;
805        let mut replicas = self.replicas.write().unwrap_or_else(|e| e.into_inner());
806        if let Some(r) = replicas.iter_mut().find(|r| r.id == id) {
807            r.last_acked_lsn = r.last_acked_lsn.max(lsn);
808            r.last_durable_lsn = r.last_durable_lsn.max(lsn);
809            r.last_seen_at_unix_ms = now_ms;
810        }
811        drop(replicas);
812        self.commit_waiter.record_replica_ack(id, lsn);
813    }
814
815    /// PLAN.md Phase 11.4 — replica reports applied + durable LSN
816    /// after persisting a batch. Idempotent: only advances LSNs
817    /// monotonically. `last_seen_at_unix_ms` always refreshes.
818    /// Also signals `commit_waiter` so any thread blocked on
819    /// `ack_n` / `quorum` can wake and re-check its threshold.
820    pub fn ack_replica_lsn(&self, id: &str, applied_lsn: u64, durable_lsn: u64) {
821        self.ack_replica_lsn_with_observability(id, applied_lsn, durable_lsn, 0, 0);
822    }
823
824    pub fn ack_replica_lsn_with_observability(
825        &self,
826        id: &str,
827        applied_lsn: u64,
828        durable_lsn: u64,
829        apply_error_count: u64,
830        divergence_count: u64,
831    ) {
832        let now_ms = crate::utils::now_unix_millis() as u128;
833        self.advance_slot(id, applied_lsn, durable_lsn, now_ms);
834        let mut replicas = self.replicas.write().unwrap_or_else(|e| e.into_inner());
835        if let Some(r) = replicas.iter_mut().find(|r| r.id == id) {
836            r.last_acked_lsn = r.last_acked_lsn.max(applied_lsn);
837            r.last_durable_lsn = r.last_durable_lsn.max(durable_lsn);
838            r.apply_error_count = r.apply_error_count.max(apply_error_count);
839            r.divergence_count = r.divergence_count.max(divergence_count);
840            r.last_seen_at_unix_ms = now_ms;
841        }
842        // Drop the write lock before signaling so a waiter that
843        // wakes immediately can read replica state without
844        // contending against us.
845        drop(replicas);
846        self.commit_waiter.record_replica_ack(id, durable_lsn);
847    }
848
849    /// PLAN.md Phase 11.4 — primary records the LSN it last sent to a
850    /// replica via pull_wal_records. Helpful for `lag_records =
851    /// last_sent_lsn - last_acked_lsn` to distinguish pull-side delay
852    /// from apply-side delay.
853    pub fn note_replica_pull(&self, id: &str, last_sent_lsn: u64) {
854        let now_ms = crate::utils::now_unix_millis() as u128;
855        self.touch_slot(id, now_ms);
856        let mut replicas = self.replicas.write().unwrap_or_else(|e| e.into_inner());
857        if let Some(r) = replicas.iter_mut().find(|r| r.id == id) {
858            r.last_sent_lsn = r.last_sent_lsn.max(last_sent_lsn);
859            r.last_seen_at_unix_ms = now_ms;
860        }
861    }
862
863    /// Snapshot of all currently registered replicas, for /metrics +
864    /// /admin/status. Returns owned clones so callers don't hold the
865    /// lock during serialization.
866    pub fn replica_snapshots(&self) -> Vec<ReplicaState> {
867        self.replicas
868            .read()
869            .unwrap_or_else(|e| e.into_inner())
870            .clone()
871    }
872
873    pub fn replication_progress(&self) -> Option<ReplicationProgress> {
874        let replicas = self.replicas.read().unwrap_or_else(|e| e.into_inner());
875        ReplicationProgress::from_replicas(&replicas)
876    }
877
878    pub fn slot_snapshots(&self) -> Vec<ReplicationSlot> {
879        self.slots
880            .read()
881            .unwrap_or_else(|e| e.into_inner())
882            .values()
883            .cloned()
884            .collect()
885    }
886
887    pub fn retention_floor_lsn(&self) -> Option<u64> {
888        self.slots
889            .read()
890            .unwrap_or_else(|e| e.into_inner())
891            .values()
892            .filter(|slot| slot.invalidation_reason.is_none())
893            .map(|slot| slot.restart_lsn)
894            .min()
895    }
896
897    pub fn prune_retained_wal_through(&self, archived_lsn: u64) -> io::Result<u64> {
898        self.enforce_retention_limits(crate::utils::now_unix_millis() as u128);
899        let prune_lsn = self
900            .retention_floor_lsn()
901            .map(|floor| floor.min(archived_lsn))
902            .unwrap_or(archived_lsn);
903        if prune_lsn > 0 {
904            if let Some(spool) = &self.logical_wal_spool {
905                spool.prune_through(prune_lsn)?;
906            }
907            self.wal_buffer.prune_through(prune_lsn);
908        }
909        Ok(prune_lsn)
910    }
911
912    pub fn replica_count(&self) -> usize {
913        self.replicas
914            .read()
915            .unwrap_or_else(|e| e.into_inner())
916            .len()
917    }
918
919    /// Current primary write position (logical WAL LSN, falling back to
920    /// the in-memory WAL buffer). Used as the reference point for
921    /// per-replica lag — including issue #826 flow control.
922    pub fn current_logical_lsn(&self) -> u64 {
923        self.logical_wal_spool
924            .as_ref()
925            .map(|spool| spool.current_lsn())
926            .unwrap_or_else(|| self.wal_buffer.current_lsn())
927    }
928
929    fn ensure_slot(&self, id: &str, initial_lsn: u64) -> u64 {
930        let now_ms = crate::utils::now_unix_millis() as u128;
931        let mut slots = self.slots.write().unwrap_or_else(|e| e.into_inner());
932        if let Some(slot) = slots.get_mut(id) {
933            slot.last_seen_at_unix_ms = now_ms;
934            let restart_lsn = slot.restart_lsn;
935            self.persist_slots_locked(&slots);
936            return restart_lsn;
937        }
938        let mut slot = ReplicationSlot::new(
939            id.to_string(),
940            reddb_file::TimelineId::initial(),
941            initial_lsn,
942        );
943        slot.last_seen_at_unix_ms = now_ms;
944        slots.insert(id.to_string(), slot);
945        let restart_lsn = initial_lsn;
946        self.persist_slots_locked(&slots);
947        restart_lsn
948    }
949
950    fn advance_slot(&self, id: &str, confirmed_lsn: u64, restart_lsn: u64, now_ms: u128) {
951        let mut slots = self.slots.write().unwrap_or_else(|e| e.into_inner());
952        let slot = slots.entry(id.to_string()).or_insert_with(|| {
953            let mut slot =
954                ReplicationSlot::new(id.to_string(), reddb_file::TimelineId::initial(), 0);
955            slot.last_seen_at_unix_ms = now_ms;
956            slot
957        });
958        if slot.invalidation_reason.is_some() {
959            return;
960        }
961        slot.confirmed_write_lsn = slot.confirmed_lsn().max(confirmed_lsn).max(restart_lsn);
962        slot.restart_lsn = slot.restart_lsn.max(restart_lsn);
963        slot.confirmed_flush_lsn = slot.confirmed_flush_lsn.max(slot.restart_lsn);
964        slot.confirmed_apply_lsn = slot.confirmed_apply_lsn.max(slot.restart_lsn);
965        slot.last_seen_at_unix_ms = now_ms;
966        self.persist_slots_locked(&slots);
967    }
968
969    pub fn touch_slot(&self, id: &str, now_ms: u128) {
970        let mut slots = self.slots.write().unwrap_or_else(|e| e.into_inner());
971        let mut changed = false;
972        if let Some(slot) = slots.get_mut(id) {
973            if slot.invalidation_reason.is_none() {
974                slot.last_seen_at_unix_ms = now_ms;
975                changed = true;
976            }
977        }
978        if changed {
979            self.persist_slots_locked(&slots);
980        }
981    }
982
983    pub fn enforce_retention_limits(
984        &self,
985        now_ms: u128,
986    ) -> Vec<(String, ReplicationSlotInvalidationCause)> {
987        let current_lsn = self.current_logical_lsn();
988        let mut invalidated = Vec::new();
989        let mut slots = self.slots.write().unwrap_or_else(|e| e.into_inner());
990        for slot in slots.values_mut() {
991            if slot.invalidation_reason.is_some() {
992                continue;
993            }
994            let reason = if self.slot_retention_max_lag_lsn > 0
995                && current_lsn.saturating_sub(slot.restart_lsn) > self.slot_retention_max_lag_lsn
996            {
997                Some(ReplicationSlotInvalidationCause::Horizon)
998            } else if self.slot_idle_timeout_ms > 0
999                && now_ms.saturating_sub(slot.last_seen_at_unix_ms)
1000                    > u128::from(self.slot_idle_timeout_ms)
1001            {
1002                Some(ReplicationSlotInvalidationCause::IdleTimeout)
1003            } else {
1004                None
1005            };
1006            if let Some(reason) = reason {
1007                slot.invalidation_reason = Some(reason);
1008                slot.invalidated_at_unix_ms = Some(now_ms);
1009                invalidated.push((slot.replica_id.clone(), reason));
1010            }
1011        }
1012        if !invalidated.is_empty() {
1013            self.persist_slots_locked(&slots);
1014        }
1015        invalidated
1016    }
1017
1018    pub fn slot_rebootstrap_reason(
1019        &self,
1020        id: &str,
1021        requested_since_lsn: u64,
1022        oldest_available_lsn: Option<u64>,
1023    ) -> Option<ReplicationSlotInvalidationCause> {
1024        let now_ms = crate::utils::now_unix_millis() as u128;
1025        let mut slots = self.slots.write().unwrap_or_else(|e| e.into_inner());
1026        let slot = slots.get_mut(id)?;
1027        if let Some(reason) = slot.invalidation_reason {
1028            return Some(reason);
1029        }
1030        let slot_floor = slot.restart_lsn.max(requested_since_lsn);
1031        if oldest_available_lsn
1032            .map(|oldest| oldest > slot_floor.saturating_add(1))
1033            .unwrap_or(false)
1034        {
1035            slot.invalidation_reason = Some(ReplicationSlotInvalidationCause::WalRemoved);
1036            slot.invalidated_at_unix_ms = Some(now_ms);
1037            self.persist_slots_locked(&slots);
1038            return Some(ReplicationSlotInvalidationCause::WalRemoved);
1039        }
1040        None
1041    }
1042
1043    /// Decide how a reconnecting replica's pull should be served
1044    /// (issue #832). If the slot is invalidated or the requested
1045    /// position has fallen behind the retained WAL floor, the replica
1046    /// must re-bootstrap; otherwise it resumes via a partial resync
1047    /// from its slot position (never rewound behind it). Every
1048    /// partial-resync decision bumps the `partial_resync_count` metric
1049    /// so a brief disconnect that recovers without a full re-bootstrap
1050    /// is observable.
1051    pub fn plan_replica_resume(
1052        &self,
1053        id: &str,
1054        requested_since_lsn: u64,
1055        oldest_available_lsn: Option<u64>,
1056    ) -> ResumeMode {
1057        if let Some(cause) =
1058            self.slot_rebootstrap_reason(id, requested_since_lsn, oldest_available_lsn)
1059        {
1060            self.full_resync_count
1061                .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
1062            return ResumeMode::FullRebootstrap { cause };
1063        }
1064        let resume_lsn = self
1065            .slot_snapshots()
1066            .into_iter()
1067            .find(|slot| slot.replica_id == id)
1068            .map(|slot| requested_since_lsn.max(slot.restart_lsn))
1069            .unwrap_or(requested_since_lsn);
1070        self.partial_resync_count
1071            .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
1072        ResumeMode::PartialResync { resume_lsn }
1073    }
1074
1075    /// Number of pulls served as a partial resync since process start.
1076    /// Surfaced in the replication metrics/status payload (issue #832).
1077    pub fn partial_resync_count(&self) -> u64 {
1078        self.partial_resync_count
1079            .load(std::sync::atomic::Ordering::Relaxed)
1080    }
1081
1082    /// Number of pulls that forced a full re-bootstrap since process
1083    /// start (issue #839). Surfaced as `reddb_replication_full_resync_total`
1084    /// and in `/replication/status` — the primary operator alert signal.
1085    pub fn full_resync_count(&self) -> u64 {
1086        self.full_resync_count
1087            .load(std::sync::atomic::Ordering::Relaxed)
1088    }
1089
1090    fn persist_slots_locked(&self, slots: &BTreeMap<String, ReplicationSlot>) {
1091        if let Err(err) = persist_replication_slots(self.slot_path.as_deref(), slots) {
1092            warn!(
1093                target: "reddb::replication::slots",
1094                error = %err,
1095                "failed to persist replication slots"
1096            );
1097        }
1098        if let Err(err) = persist_replication_slot_catalog(self.slot_catalog_path.as_deref(), slots)
1099        {
1100            warn!(
1101                target: "reddb::replication::slots",
1102                error = %err,
1103                "failed to persist binary replication slot catalog"
1104            );
1105        }
1106    }
1107}
1108
1109#[cfg(test)]
1110mod tests;