1use 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
35pub struct WalBuffer {
43 records: RwLock<VecDeque<(u64, Arc<[u8]>)>>,
45 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 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 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 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 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 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#[derive(Debug, Default)]
152struct LogicalWalSpoolState {
153 current_lsn: u64,
154 seek_index: Vec<(u64, u64)>,
158 write_offset: u64,
162 record_count: u64,
165}
166
167impl LogicalWalSpoolState {
168 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 if self.seek_index.last().map(|(l, _)| *l) != Some(lsn) {
177 self.seek_index.push((lsn, offset));
178 }
179 }
180 }
181
182 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
198pub 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 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 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 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 let frame =
294 reddb_file::encode_logical_wal_v3(term, ownership_epoch, lsn, timestamp_ms, data)?;
295
296 file.write_all(&frame)?;
297 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 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 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 #[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 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 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#[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 pub region: Option<String>,
424 pub rebootstrapping: bool,
431}
432
433#[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
459pub 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 pub commit_waiter: Arc<crate::replication::commit_waiter::CommitWaiter>,
476 topology_epoch: std::sync::atomic::AtomicU64,
483 partial_resync_count: std::sync::atomic::AtomicU64,
489 full_resync_count: std::sync::atomic::AtomicU64,
496}
497
498#[derive(Debug, Clone, Copy, PartialEq, Eq)]
500pub enum ResumeMode {
501 PartialResync { resume_lsn: u64 },
505 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 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 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 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 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 pub fn topology_epoch(&self) -> u64 {
776 self.topology_epoch
777 .load(std::sync::atomic::Ordering::Relaxed)
778 }
779
780 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 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(replicas);
835 self.commit_waiter.record_replica_ack(id, durable_lsn);
836 }
837
838 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 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 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 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 pub fn partial_resync_count(&self) -> u64 {
1067 self.partial_resync_count
1068 .load(std::sync::atomic::Ordering::Relaxed)
1069 }
1070
1071 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;