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        let now_ms = crate::utils::now_unix_millis() as u128;
681        let resume_lsn = self.ensure_slot(&id, self.current_logical_lsn());
682        let mut replicas = self.replicas.write().unwrap_or_else(|e| e.into_inner());
683        if let Some(existing) = replicas.iter_mut().find(|r| r.id == id) {
684            existing.last_seen_at_unix_ms = now_ms;
685            if region.is_some() {
686                existing.region = region;
687            }
688            return resume_lsn;
689        }
690        replicas.push(ReplicaState {
691            id,
692            last_acked_lsn: resume_lsn,
693            last_sent_lsn: resume_lsn,
694            last_durable_lsn: resume_lsn,
695            apply_error_count: 0,
696            divergence_count: 0,
697            connected_at_unix_ms: now_ms,
698            last_seen_at_unix_ms: now_ms,
699            region,
700            rebootstrapping: false,
701        });
702        drop(replicas);
703        self.bump_topology_epoch();
704        resume_lsn
705    }
706
707    /// Mark (or clear) a replica's re-bootstrap state (issue #837).
708    ///
709    /// While `rebootstrapping` is `true` the replica keeps serving
710    /// non-causal reads from its existing data, but the advertiser
711    /// surfaces the flag so causal (bookmark) reads route to a
712    /// caught-up peer instead — the rebuilding replica's applied
713    /// frontier describes data it is about to discard. The primary
714    /// flips this back to `false` when the replica reports its atomic
715    /// snapshot swap complete.
716    ///
717    /// A change to the flag is a registry-shape change for routing
718    /// purposes, so it bumps the topology epoch to force consumers to
719    /// re-read the advertisement. Returns `true` when a replica with
720    /// `id` was present and updated.
721    pub fn set_replica_rebootstrapping(&self, id: &str, rebootstrapping: bool) -> bool {
722        let mut replicas = self.replicas.write().unwrap_or_else(|e| e.into_inner());
723        let Some(state) = replicas.iter_mut().find(|r| r.id == id) else {
724            return false;
725        };
726        if state.rebootstrapping == rebootstrapping {
727            return true;
728        }
729        state.rebootstrapping = rebootstrapping;
730        drop(replicas);
731        self.bump_topology_epoch();
732        true
733    }
734
735    /// Ensure a replica identifying itself with `id` is present in the
736    /// registry (issue #812). This is the production self-registration hook
737    /// used by the `pull_wal_records` path: the first time a replica sends
738    /// its `replica_id` on a pull, the primary registers it so it is no
739    /// longer blind to that replica's existence; subsequent pulls are
740    /// idempotent no-ops. Returns `true` when a new registration was
741    /// created. Delegates to `register_replica_with_region`, so reconnects
742    /// preserve progress and do not bump the topology epoch.
743    pub fn ensure_replica_registered(&self, id: &str) -> bool {
744        let already = self
745            .replicas
746            .read()
747            .unwrap_or_else(|e| e.into_inner())
748            .iter()
749            .any(|r| r.id == id);
750        if already {
751            return false;
752        }
753        self.register_replica(id.to_string());
754        true
755    }
756
757    /// Unregister a replica by id. Returns `true` when the replica
758    /// was present (and removed). Bumps the topology epoch so a
759    /// pending advertisement reflects the new fleet size.
760    pub fn unregister_replica(&self, id: &str) -> bool {
761        let mut replicas = self.replicas.write().unwrap_or_else(|e| e.into_inner());
762        let before = replicas.len();
763        replicas.retain(|r| r.id != id);
764        let removed = replicas.len() != before;
765        drop(replicas);
766        if removed {
767            self.commit_waiter.drop_replica(id);
768            self.bump_topology_epoch();
769        }
770        removed
771    }
772
773    /// Current topology epoch. Strictly monotonic, bumps on every
774    /// registry-shape change consumed by `TopologyAdvertiser`.
775    pub fn topology_epoch(&self) -> u64 {
776        self.topology_epoch
777            .load(std::sync::atomic::Ordering::Relaxed)
778    }
779
780    /// Advance the topology epoch. Call sites: register, unregister,
781    /// and the health-sweep tick that flips a replica between
782    /// healthy/unhealthy. Wrapping is not a concern in practice
783    /// (`u64::MAX` events would take centuries at any realistic ack
784    /// rate) but `fetch_add` saturates implicitly via wrap-around;
785    /// the consumer treats epoch as opaque so a wrap is still
786    /// strictly "different" from the previous value.
787    pub fn bump_topology_epoch(&self) {
788        self.topology_epoch
789            .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
790    }
791
792    pub fn ack_replica(&self, id: &str, lsn: u64) {
793        let now_ms = crate::utils::now_unix_millis() as u128;
794        let mut replicas = self.replicas.write().unwrap_or_else(|e| e.into_inner());
795        if let Some(r) = replicas.iter_mut().find(|r| r.id == id) {
796            r.last_acked_lsn = r.last_acked_lsn.max(lsn);
797            r.last_durable_lsn = r.last_durable_lsn.max(lsn);
798            r.last_seen_at_unix_ms = now_ms;
799        }
800        drop(replicas);
801        self.commit_waiter.record_replica_ack(id, lsn);
802    }
803
804    /// PLAN.md Phase 11.4 — replica reports applied + durable LSN
805    /// after persisting a batch. Idempotent: only advances LSNs
806    /// monotonically. `last_seen_at_unix_ms` always refreshes.
807    /// Also signals `commit_waiter` so any thread blocked on
808    /// `ack_n` / `quorum` can wake and re-check its threshold.
809    pub fn ack_replica_lsn(&self, id: &str, applied_lsn: u64, durable_lsn: u64) {
810        self.ack_replica_lsn_with_observability(id, applied_lsn, durable_lsn, 0, 0);
811    }
812
813    pub fn ack_replica_lsn_with_observability(
814        &self,
815        id: &str,
816        applied_lsn: u64,
817        durable_lsn: u64,
818        apply_error_count: u64,
819        divergence_count: u64,
820    ) {
821        let now_ms = crate::utils::now_unix_millis() as u128;
822        self.advance_slot(id, applied_lsn, durable_lsn, now_ms);
823        let mut replicas = self.replicas.write().unwrap_or_else(|e| e.into_inner());
824        if let Some(r) = replicas.iter_mut().find(|r| r.id == id) {
825            r.last_acked_lsn = r.last_acked_lsn.max(applied_lsn);
826            r.last_durable_lsn = r.last_durable_lsn.max(durable_lsn);
827            r.apply_error_count = r.apply_error_count.max(apply_error_count);
828            r.divergence_count = r.divergence_count.max(divergence_count);
829            r.last_seen_at_unix_ms = now_ms;
830        }
831        // Drop the write lock before signaling so a waiter that
832        // wakes immediately can read replica state without
833        // contending against us.
834        drop(replicas);
835        self.commit_waiter.record_replica_ack(id, durable_lsn);
836    }
837
838    /// PLAN.md Phase 11.4 — primary records the LSN it last sent to a
839    /// replica via pull_wal_records. Helpful for `lag_records =
840    /// last_sent_lsn - last_acked_lsn` to distinguish pull-side delay
841    /// from apply-side delay.
842    pub fn note_replica_pull(&self, id: &str, last_sent_lsn: u64) {
843        let now_ms = crate::utils::now_unix_millis() as u128;
844        self.touch_slot(id, now_ms);
845        let mut replicas = self.replicas.write().unwrap_or_else(|e| e.into_inner());
846        if let Some(r) = replicas.iter_mut().find(|r| r.id == id) {
847            r.last_sent_lsn = r.last_sent_lsn.max(last_sent_lsn);
848            r.last_seen_at_unix_ms = now_ms;
849        }
850    }
851
852    /// Snapshot of all currently registered replicas, for /metrics +
853    /// /admin/status. Returns owned clones so callers don't hold the
854    /// lock during serialization.
855    pub fn replica_snapshots(&self) -> Vec<ReplicaState> {
856        self.replicas
857            .read()
858            .unwrap_or_else(|e| e.into_inner())
859            .clone()
860    }
861
862    pub fn replication_progress(&self) -> Option<ReplicationProgress> {
863        let replicas = self.replicas.read().unwrap_or_else(|e| e.into_inner());
864        ReplicationProgress::from_replicas(&replicas)
865    }
866
867    pub fn slot_snapshots(&self) -> Vec<ReplicationSlot> {
868        self.slots
869            .read()
870            .unwrap_or_else(|e| e.into_inner())
871            .values()
872            .cloned()
873            .collect()
874    }
875
876    pub fn retention_floor_lsn(&self) -> Option<u64> {
877        self.slots
878            .read()
879            .unwrap_or_else(|e| e.into_inner())
880            .values()
881            .filter(|slot| slot.invalidation_reason.is_none())
882            .map(|slot| slot.restart_lsn)
883            .min()
884    }
885
886    pub fn prune_retained_wal_through(&self, archived_lsn: u64) -> io::Result<u64> {
887        self.enforce_retention_limits(crate::utils::now_unix_millis() as u128);
888        let prune_lsn = self
889            .retention_floor_lsn()
890            .map(|floor| floor.min(archived_lsn))
891            .unwrap_or(archived_lsn);
892        if prune_lsn > 0 {
893            if let Some(spool) = &self.logical_wal_spool {
894                spool.prune_through(prune_lsn)?;
895            }
896            self.wal_buffer.prune_through(prune_lsn);
897        }
898        Ok(prune_lsn)
899    }
900
901    pub fn replica_count(&self) -> usize {
902        self.replicas
903            .read()
904            .unwrap_or_else(|e| e.into_inner())
905            .len()
906    }
907
908    /// Current primary write position (logical WAL LSN, falling back to
909    /// the in-memory WAL buffer). Used as the reference point for
910    /// per-replica lag — including issue #826 flow control.
911    pub fn current_logical_lsn(&self) -> u64 {
912        self.logical_wal_spool
913            .as_ref()
914            .map(|spool| spool.current_lsn())
915            .unwrap_or_else(|| self.wal_buffer.current_lsn())
916    }
917
918    fn ensure_slot(&self, id: &str, initial_lsn: u64) -> u64 {
919        let now_ms = crate::utils::now_unix_millis() as u128;
920        let mut slots = self.slots.write().unwrap_or_else(|e| e.into_inner());
921        if let Some(slot) = slots.get_mut(id) {
922            slot.last_seen_at_unix_ms = now_ms;
923            let restart_lsn = slot.restart_lsn;
924            self.persist_slots_locked(&slots);
925            return restart_lsn;
926        }
927        let mut slot = ReplicationSlot::new(
928            id.to_string(),
929            reddb_file::TimelineId::initial(),
930            initial_lsn,
931        );
932        slot.last_seen_at_unix_ms = now_ms;
933        slots.insert(id.to_string(), slot);
934        let restart_lsn = initial_lsn;
935        self.persist_slots_locked(&slots);
936        restart_lsn
937    }
938
939    fn advance_slot(&self, id: &str, confirmed_lsn: u64, restart_lsn: u64, now_ms: u128) {
940        let mut slots = self.slots.write().unwrap_or_else(|e| e.into_inner());
941        let slot = slots.entry(id.to_string()).or_insert_with(|| {
942            let mut slot =
943                ReplicationSlot::new(id.to_string(), reddb_file::TimelineId::initial(), 0);
944            slot.last_seen_at_unix_ms = now_ms;
945            slot
946        });
947        if slot.invalidation_reason.is_some() {
948            return;
949        }
950        slot.confirmed_write_lsn = slot.confirmed_lsn().max(confirmed_lsn).max(restart_lsn);
951        slot.restart_lsn = slot.restart_lsn.max(restart_lsn);
952        slot.confirmed_flush_lsn = slot.confirmed_flush_lsn.max(slot.restart_lsn);
953        slot.confirmed_apply_lsn = slot.confirmed_apply_lsn.max(slot.restart_lsn);
954        slot.last_seen_at_unix_ms = now_ms;
955        self.persist_slots_locked(&slots);
956    }
957
958    pub fn touch_slot(&self, id: &str, now_ms: u128) {
959        let mut slots = self.slots.write().unwrap_or_else(|e| e.into_inner());
960        let mut changed = false;
961        if let Some(slot) = slots.get_mut(id) {
962            if slot.invalidation_reason.is_none() {
963                slot.last_seen_at_unix_ms = now_ms;
964                changed = true;
965            }
966        }
967        if changed {
968            self.persist_slots_locked(&slots);
969        }
970    }
971
972    pub fn enforce_retention_limits(
973        &self,
974        now_ms: u128,
975    ) -> Vec<(String, ReplicationSlotInvalidationCause)> {
976        let current_lsn = self.current_logical_lsn();
977        let mut invalidated = Vec::new();
978        let mut slots = self.slots.write().unwrap_or_else(|e| e.into_inner());
979        for slot in slots.values_mut() {
980            if slot.invalidation_reason.is_some() {
981                continue;
982            }
983            let reason = if self.slot_retention_max_lag_lsn > 0
984                && current_lsn.saturating_sub(slot.restart_lsn) > self.slot_retention_max_lag_lsn
985            {
986                Some(ReplicationSlotInvalidationCause::Horizon)
987            } else if self.slot_idle_timeout_ms > 0
988                && now_ms.saturating_sub(slot.last_seen_at_unix_ms)
989                    > u128::from(self.slot_idle_timeout_ms)
990            {
991                Some(ReplicationSlotInvalidationCause::IdleTimeout)
992            } else {
993                None
994            };
995            if let Some(reason) = reason {
996                slot.invalidation_reason = Some(reason);
997                slot.invalidated_at_unix_ms = Some(now_ms);
998                invalidated.push((slot.replica_id.clone(), reason));
999            }
1000        }
1001        if !invalidated.is_empty() {
1002            self.persist_slots_locked(&slots);
1003        }
1004        invalidated
1005    }
1006
1007    pub fn slot_rebootstrap_reason(
1008        &self,
1009        id: &str,
1010        requested_since_lsn: u64,
1011        oldest_available_lsn: Option<u64>,
1012    ) -> Option<ReplicationSlotInvalidationCause> {
1013        let now_ms = crate::utils::now_unix_millis() as u128;
1014        let mut slots = self.slots.write().unwrap_or_else(|e| e.into_inner());
1015        let slot = slots.get_mut(id)?;
1016        if let Some(reason) = slot.invalidation_reason {
1017            return Some(reason);
1018        }
1019        let slot_floor = slot.restart_lsn.max(requested_since_lsn);
1020        if oldest_available_lsn
1021            .map(|oldest| oldest > slot_floor.saturating_add(1))
1022            .unwrap_or(false)
1023        {
1024            slot.invalidation_reason = Some(ReplicationSlotInvalidationCause::WalRemoved);
1025            slot.invalidated_at_unix_ms = Some(now_ms);
1026            self.persist_slots_locked(&slots);
1027            return Some(ReplicationSlotInvalidationCause::WalRemoved);
1028        }
1029        None
1030    }
1031
1032    /// Decide how a reconnecting replica's pull should be served
1033    /// (issue #832). If the slot is invalidated or the requested
1034    /// position has fallen behind the retained WAL floor, the replica
1035    /// must re-bootstrap; otherwise it resumes via a partial resync
1036    /// from its slot position (never rewound behind it). Every
1037    /// partial-resync decision bumps the `partial_resync_count` metric
1038    /// so a brief disconnect that recovers without a full re-bootstrap
1039    /// is observable.
1040    pub fn plan_replica_resume(
1041        &self,
1042        id: &str,
1043        requested_since_lsn: u64,
1044        oldest_available_lsn: Option<u64>,
1045    ) -> ResumeMode {
1046        if let Some(cause) =
1047            self.slot_rebootstrap_reason(id, requested_since_lsn, oldest_available_lsn)
1048        {
1049            self.full_resync_count
1050                .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
1051            return ResumeMode::FullRebootstrap { cause };
1052        }
1053        let resume_lsn = self
1054            .slot_snapshots()
1055            .into_iter()
1056            .find(|slot| slot.replica_id == id)
1057            .map(|slot| requested_since_lsn.max(slot.restart_lsn))
1058            .unwrap_or(requested_since_lsn);
1059        self.partial_resync_count
1060            .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
1061        ResumeMode::PartialResync { resume_lsn }
1062    }
1063
1064    /// Number of pulls served as a partial resync since process start.
1065    /// Surfaced in the replication metrics/status payload (issue #832).
1066    pub fn partial_resync_count(&self) -> u64 {
1067        self.partial_resync_count
1068            .load(std::sync::atomic::Ordering::Relaxed)
1069    }
1070
1071    /// Number of pulls that forced a full re-bootstrap since process
1072    /// start (issue #839). Surfaced as `reddb_replication_full_resync_total`
1073    /// and in `/replication/status` — the primary operator alert signal.
1074    pub fn full_resync_count(&self) -> u64 {
1075        self.full_resync_count
1076            .load(std::sync::atomic::Ordering::Relaxed)
1077    }
1078
1079    fn persist_slots_locked(&self, slots: &BTreeMap<String, ReplicationSlot>) {
1080        if let Err(err) = persist_replication_slots(self.slot_path.as_deref(), slots) {
1081            warn!(
1082                target: "reddb::replication::slots",
1083                error = %err,
1084                "failed to persist replication slots"
1085            );
1086        }
1087        if let Err(err) = persist_replication_slot_catalog(self.slot_catalog_path.as_deref(), slots)
1088        {
1089            warn!(
1090                target: "reddb::replication::slots",
1091                error = %err,
1092                "failed to persist binary replication slot catalog"
1093            );
1094        }
1095    }
1096}
1097
1098#[cfg(test)]
1099mod tests;