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 read_since_shared_bounded(
96 &self,
97 since_lsn: u64,
98 max_count: usize,
99 max_payload_bytes: usize,
100 ) -> Vec<(u64, Arc<[u8]>)> {
101 let records = self.records.read().unwrap_or_else(|e| e.into_inner());
102 let mut payload_bytes = 0usize;
103 records
104 .iter()
105 .filter(|(lsn, _)| *lsn > since_lsn)
106 .take(max_count)
107 .take_while(|(_, data)| {
108 let next_payload_bytes = payload_bytes.saturating_add(data.len());
109 let include = payload_bytes == 0 || next_payload_bytes <= max_payload_bytes;
110 if include {
111 payload_bytes = next_payload_bytes;
112 }
113 include
114 })
115 .cloned()
116 .collect()
117 }
118
119 pub fn current_lsn(&self) -> u64 {
121 *self.current_lsn.read().unwrap_or_else(|e| e.into_inner())
122 }
123
124 pub fn set_current_lsn(&self, lsn: u64) {
125 let mut current = self.current_lsn.write().unwrap_or_else(|e| e.into_inner());
126 *current = (*current).max(lsn);
127 }
128
129 pub fn prune_through(&self, upto_lsn: u64) {
130 let mut records = self.records.write().unwrap_or_else(|e| e.into_inner());
131 while records
132 .front()
133 .map(|(lsn, _)| *lsn <= upto_lsn)
134 .unwrap_or(false)
135 {
136 records.pop_front();
137 }
138 }
139
140 pub fn oldest_lsn(&self) -> Option<u64> {
142 let records = self.records.read().unwrap_or_else(|e| e.into_inner());
143 records.front().map(|(lsn, _)| *lsn)
144 }
145}
146
147fn logical_wal_entry_term(entry: &reddb_file::LogicalWalEntry) -> u64 {
148 if entry.term == 0 {
149 term_from_payload(&entry.data)
150 } else {
151 entry.term
152 }
153}
154
155fn logical_wal_data_with_framing_term(entry: &reddb_file::LogicalWalEntry) -> Vec<u8> {
156 let term = logical_wal_entry_term(entry);
157 match crate::replication::cdc::ChangeRecord::decode(&entry.data) {
158 Ok(mut record)
159 if record.term != term || record.ownership_epoch != entry.ownership_epoch =>
160 {
161 record.term = term;
162 if entry.ownership_epoch.is_some() {
163 record.ownership_epoch = entry.ownership_epoch;
164 }
165 record.encode()
166 }
167 _ => entry.data.clone(),
168 }
169}
170
171#[derive(Debug, Default)]
179struct LogicalWalSpoolState {
180 current_lsn: u64,
181 seek_index: Vec<(u64, u64)>,
185 write_offset: u64,
189 record_count: u64,
192}
193
194impl LogicalWalSpoolState {
195 fn note_record(&mut self, ordinal: u64, lsn: u64, offset: u64) {
199 if ordinal.is_multiple_of(reddb_file::LOGICAL_WAL_SEEK_INDEX_INTERVAL) {
200 if self.seek_index.last().map(|(l, _)| *l) != Some(lsn) {
204 self.seek_index.push((lsn, offset));
205 }
206 }
207 }
208
209 fn seek_floor_offset(&self, since_lsn: u64) -> u64 {
214 match self
215 .seek_index
216 .binary_search_by(|(lsn, _)| lsn.cmp(&since_lsn))
217 {
218 Ok(idx) => self.seek_index[idx].1,
219 Err(0) => 0,
220 Err(idx) => self.seek_index[idx - 1].1,
221 }
222 }
223}
224
225pub struct LogicalWalSpool {
229 path: PathBuf,
230 state: Mutex<LogicalWalSpoolState>,
231}
232
233impl LogicalWalSpool {
234 pub fn path_for(data_path: &Path) -> PathBuf {
235 reddb_file::layout::logical_wal_path(data_path)
236 }
237
238 pub fn open(data_path: &Path) -> io::Result<Self> {
239 let path = Self::path_for(data_path);
240 if let Some(parent) = path.parent() {
241 fs::create_dir_all(parent)?;
242 }
243 if !path.exists() {
244 File::create(&path)?;
245 }
246 let entries = reddb_file::read_and_repair_logical_wal_entries(&path)?;
251 let current_lsn = entries.last().map(|entry| entry.lsn).unwrap_or(0);
252 let (seek_index, write_offset, record_count) =
255 reddb_file::build_logical_wal_seek_index(&path)?;
256 Ok(Self {
257 path,
258 state: Mutex::new(LogicalWalSpoolState {
259 current_lsn,
260 seek_index,
261 write_offset,
262 record_count,
263 }),
264 })
265 }
266
267 pub fn append(&self, lsn: u64, data: &[u8]) -> io::Result<()> {
268 let timestamp_ms = SystemTime::now()
269 .duration_since(UNIX_EPOCH)
270 .map(|d| d.as_millis() as u64)
271 .unwrap_or(0);
272 self.append_with_timestamp(lsn, timestamp_ms, data)
273 }
274
275 pub fn append_with_timestamp(
279 &self,
280 lsn: u64,
281 timestamp_ms: u64,
282 data: &[u8],
283 ) -> io::Result<()> {
284 let (term, ownership_epoch) = authority_from_payload(data);
285 self.append_with_term_epoch_and_timestamp(term, ownership_epoch, lsn, timestamp_ms, data)
286 }
287
288 pub fn append_with_term_and_timestamp(
289 &self,
290 term: u64,
291 lsn: u64,
292 timestamp_ms: u64,
293 data: &[u8],
294 ) -> io::Result<()> {
295 let (_, ownership_epoch) = authority_from_payload(data);
296 self.append_with_term_epoch_and_timestamp(term, ownership_epoch, lsn, timestamp_ms, data)
297 }
298
299 pub fn append_with_term_epoch_and_timestamp(
300 &self,
301 term: u64,
302 ownership_epoch: Option<u64>,
303 lsn: u64,
304 timestamp_ms: u64,
305 data: &[u8],
306 ) -> io::Result<()> {
307 let mut file = OpenOptions::new()
308 .create(true)
309 .append(true)
310 .open(&self.path)?;
311 let frame =
321 reddb_file::encode_logical_wal_v3(term, ownership_epoch, lsn, timestamp_ms, data)?;
322
323 file.write_all(&frame)?;
324 file.sync_all()?;
329
330 let mut state = self.state.lock().unwrap_or_else(|e| e.into_inner());
331 state.current_lsn = state.current_lsn.max(lsn);
332 let record_start = state.write_offset;
336 let ordinal = state.record_count;
337 state.note_record(ordinal, lsn, record_start);
338 state.write_offset = record_start + frame.len() as u64;
339 state.record_count = ordinal + 1;
340 Ok(())
341 }
342
343 pub fn read_since(&self, since_lsn: u64, max_count: usize) -> io::Result<Vec<(u64, Vec<u8>)>> {
344 let start_offset = {
350 let state = self.state.lock().unwrap_or_else(|e| e.into_inner());
351 state.seek_floor_offset(since_lsn)
352 };
353 let entries = reddb_file::read_logical_wal_entries_from(&self.path, start_offset)?;
354 Ok(entries
355 .into_iter()
356 .filter(|entry| entry.lsn > since_lsn)
357 .take(max_count)
358 .map(|entry| (entry.lsn, logical_wal_data_with_framing_term(&entry)))
359 .collect())
360 }
361
362 pub fn read_since_bounded(
363 &self,
364 since_lsn: u64,
365 max_count: usize,
366 max_payload_bytes: usize,
367 ) -> io::Result<Vec<(u64, Vec<u8>)>> {
368 let start_offset = {
369 let state = self.state.lock().unwrap_or_else(|e| e.into_inner());
370 state.seek_floor_offset(since_lsn)
371 };
372 Ok(reddb_file::read_logical_wal_entries_from_bounded(
373 &self.path,
374 start_offset,
375 since_lsn,
376 max_count,
377 max_payload_bytes,
378 )?
379 .into_iter()
380 .map(|entry| (entry.lsn, logical_wal_data_with_framing_term(&entry)))
381 .collect())
382 }
383
384 #[cfg(test)]
388 fn seek_floor_offset(&self, since_lsn: u64) -> u64 {
389 self.state
390 .lock()
391 .unwrap_or_else(|e| e.into_inner())
392 .seek_floor_offset(since_lsn)
393 }
394
395 pub fn current_lsn(&self) -> u64 {
396 self.state
397 .lock()
398 .unwrap_or_else(|e| e.into_inner())
399 .current_lsn
400 }
401
402 pub fn oldest_lsn(&self) -> io::Result<Option<u64>> {
403 Ok(reddb_file::read_and_repair_logical_wal_entries(&self.path)?
404 .into_iter()
405 .next()
406 .map(|entry| entry.lsn))
407 }
408
409 pub fn prune_through(&self, upto_lsn: u64) -> io::Result<()> {
410 let previous_lsn = self.current_lsn();
411 let mut retained: Vec<_> = reddb_file::read_and_repair_logical_wal_entries(&self.path)?
412 .into_iter()
413 .filter(|entry| entry.lsn > upto_lsn)
414 .collect();
415 for entry in &mut retained {
416 entry.term = logical_wal_entry_term(entry);
417 }
418 let temp_path = reddb_file::layout::logical_wal_temp_path(&self.path);
419 for entry in &mut retained {
420 let timestamp_ms = if entry.timestamp_ms > 0 {
428 entry.timestamp_ms
429 } else {
430 SystemTime::now()
431 .duration_since(UNIX_EPOCH)
432 .map(|d| d.as_millis() as u64)
433 .unwrap_or(0)
434 };
435 entry.timestamp_ms = timestamp_ms;
436 }
437 let current_lsn =
438 reddb_file::rewrite_logical_wal_entries(&self.path, &temp_path, &retained)?;
439
440 let (seek_index, write_offset, record_count) =
443 reddb_file::build_logical_wal_seek_index(&self.path)?;
444 let mut state = self.state.lock().unwrap_or_else(|e| e.into_inner());
445 state.current_lsn = previous_lsn.max(current_lsn).max(upto_lsn);
446 state.seek_index = seek_index;
447 state.write_offset = write_offset;
448 state.record_count = record_count;
449 Ok(())
450 }
451}
452
453#[derive(Debug, Clone)]
459pub struct ReplicaState {
460 pub id: String,
461 pub last_acked_lsn: u64,
462 pub last_sent_lsn: u64,
463 pub last_durable_lsn: u64,
464 pub apply_error_count: u64,
465 pub divergence_count: u64,
466 pub connected_at_unix_ms: u128,
467 pub last_seen_at_unix_ms: u128,
468 pub region: Option<String>,
473 pub rebootstrapping: bool,
480}
481
482#[derive(Debug, Clone, Copy, PartialEq, Eq)]
484pub struct ReplicationProgress {
485 pub lag_lsn: u64,
486 pub safe_replay_lsn: u64,
487}
488
489impl ReplicationProgress {
490 pub fn from_replicas(replicas: &[ReplicaState]) -> Option<Self> {
491 let max_sent_lsn = replicas.iter().map(|replica| replica.last_sent_lsn).max()?;
492 let min_acked_lsn = replicas
493 .iter()
494 .map(|replica| replica.last_acked_lsn)
495 .min()?;
496 let safe_replay_lsn = replicas
497 .iter()
498 .map(|replica| replica.last_durable_lsn)
499 .min()?;
500
501 Some(Self {
502 lag_lsn: max_sent_lsn.saturating_sub(min_acked_lsn),
503 safe_replay_lsn,
504 })
505 }
506}
507
508pub struct PrimaryReplication {
510 pub wal_buffer: Arc<WalBuffer>,
511 pub logical_wal_spool: Option<Arc<LogicalWalSpool>>,
512 pub replicas: RwLock<Vec<ReplicaState>>,
513 pub(crate) routing_endpoints: RwLock<BTreeMap<String, reddb_wire::routing::RoutingEndpoints>>,
514 wal_appended: (Mutex<u64>, Condvar),
515 metadata_store: Option<Arc<crate::storage::unified::UnifiedStore>>,
516 metadata_error: Mutex<Option<String>>,
517 slot_path: Option<PathBuf>,
518 slot_catalog_path: Option<PathBuf>,
519 primary_replica_file_plan: Option<reddb_file::PrimaryReplicaFilePlan>,
520 primary_replica_wal_lock: Mutex<()>,
521 slots: RwLock<BTreeMap<String, ReplicationSlot>>,
522 slot_retention_max_lag_lsn: u64,
523 slot_idle_timeout_ms: u64,
524 pub commit_waiter: Arc<crate::replication::commit_waiter::CommitWaiter>,
528 topology_epoch: std::sync::atomic::AtomicU64,
535 partial_resync_count: std::sync::atomic::AtomicU64,
541 full_resync_count: std::sync::atomic::AtomicU64,
548}
549
550#[derive(Debug, Clone, Copy, PartialEq, Eq)]
552pub enum ResumeMode {
553 PartialResync { resume_lsn: u64 },
557 FullRebootstrap {
560 cause: ReplicationSlotInvalidationCause,
561 },
562}
563
564impl PrimaryReplication {
565 pub fn slot_path_for(data_path: &Path) -> PathBuf {
566 reddb_file::layout::legacy_logical_slots_path(data_path)
567 }
568
569 pub fn primary_replica_root_for(data_path: &Path) -> PathBuf {
570 reddb_file::layout::primary_replica_root(data_path)
571 }
572
573 pub fn slot_catalog_path_for(data_path: &Path) -> PathBuf {
574 Self::primary_replica_file_plan_for(data_path).slots_path()
575 }
576
577 fn primary_replica_file_plan_for(data_path: &Path) -> reddb_file::PrimaryReplicaFilePlan {
578 let root = Self::primary_replica_root_for(data_path);
579 let timeline =
580 Self::primary_replica_current_timeline_for_root(&root).unwrap_or_else(|err| {
581 warn!(
582 target: "reddb::replication::primary",
583 error = %err,
584 "failed to read primary-replica timeline history; using initial timeline"
585 );
586 reddb_file::TimelineId::initial()
587 });
588 reddb_file::PrimaryReplicaFilePlan::new(root, timeline)
589 }
590
591 fn primary_replica_current_timeline_for_root(
592 root: &Path,
593 ) -> Result<reddb_file::TimelineId, reddb_file::RdbFileError> {
594 let path = reddb_file::PrimaryReplicaFilePlan::new(root, reddb_file::TimelineId::initial())
595 .timeline_history_path();
596 match reddb_file::TimelineHistory::read_from_path(&path) {
597 Ok(history) => Ok(history
598 .current()
599 .unwrap_or_else(reddb_file::TimelineId::initial)),
600 Err(reddb_file::RdbFileError::Io(err))
601 if err.kind() == std::io::ErrorKind::NotFound =>
602 {
603 Ok(reddb_file::TimelineId::initial())
604 }
605 Err(err) => Err(err),
606 }
607 }
608
609 pub fn new(data_path: Option<&Path>) -> Self {
610 Self::new_with_config(data_path, &crate::replication::ReplicationConfig::primary())
611 }
612
613 pub fn new_with_config(
614 data_path: Option<&Path>,
615 config: &crate::replication::ReplicationConfig,
616 ) -> Self {
617 let now_ms = crate::utils::now_unix_millis() as u128;
618 let slot_path = data_path.map(Self::slot_path_for);
619 let slot_catalog_path = data_path.map(Self::slot_catalog_path_for);
620 let primary_replica_file_plan = data_path.map(Self::primary_replica_file_plan_for);
621 let mut slots = load_replication_slot_catalog(slot_catalog_path.as_deref(), now_ms);
622 slots.extend(load_replication_slots(slot_path.as_deref(), now_ms));
623 let logical_wal_spool = data_path
624 .and_then(|path| LogicalWalSpool::open(path).ok())
625 .map(Arc::new);
626 let current_lsn = logical_wal_spool
627 .as_ref()
628 .map(|spool| spool.current_lsn())
629 .unwrap_or(0);
630 Self {
631 wal_buffer: Arc::new(WalBuffer::new(100_000)),
632 logical_wal_spool,
633 replicas: RwLock::new(Vec::new()),
634 routing_endpoints: RwLock::new(BTreeMap::new()),
635 wal_appended: (Mutex::new(current_lsn), Condvar::new()),
636 metadata_store: None,
637 metadata_error: Mutex::new(None),
638 slot_path,
639 slot_catalog_path,
640 primary_replica_file_plan,
641 primary_replica_wal_lock: Mutex::new(()),
642 slots: RwLock::new(slots),
643 slot_retention_max_lag_lsn: config.slot_retention_max_lag_lsn,
644 slot_idle_timeout_ms: config.slot_idle_timeout_ms,
645 commit_waiter: Arc::new(crate::replication::commit_waiter::CommitWaiter::new()),
646 topology_epoch: std::sync::atomic::AtomicU64::new(0),
647 partial_resync_count: std::sync::atomic::AtomicU64::new(0),
648 full_resync_count: std::sync::atomic::AtomicU64::new(0),
649 }
650 }
651
652 pub(crate) fn new_with_store(
653 data_path: Option<&Path>,
654 config: &crate::replication::ReplicationConfig,
655 store: Arc<crate::storage::unified::UnifiedStore>,
656 ) -> io::Result<Self> {
657 let mut primary = Self::new_with_config(data_path, config);
658 if let Some(path) = data_path.filter(|_| primary.logical_wal_spool.is_none()) {
659 primary.logical_wal_spool = Some(Arc::new(LogicalWalSpool::open(path)?));
662 }
663 if let Some(bytes) = super::metadata::read(&store, super::metadata::SLOTS)? {
664 let catalog =
665 reddb_file::ReplicationSlotCatalog::decode(&bytes).map_err(io::Error::other)?;
666 primary.slots = RwLock::new(
667 catalog
668 .slots
669 .into_iter()
670 .map(|slot| (slot.replica_id.clone(), slot))
671 .collect(),
672 );
673 } else {
674 let mut valid_source = false;
677 let mut import_error = None;
678 for path in [&primary.slot_catalog_path, &primary.slot_path]
679 .into_iter()
680 .flatten()
681 {
682 if path.try_exists()? {
683 let result = if primary.slot_catalog_path.as_ref() == Some(path) {
684 reddb_file::ReplicationSlotCatalog::read_from_path(path)
685 } else {
686 reddb_file::ReplicationSlotCatalog::read_legacy_json_from_path(
687 path,
688 crate::utils::now_unix_millis() as u128,
689 )
690 };
691 match result {
692 Ok(_) => valid_source = true,
693 Err(error) => import_error = Some(error),
694 }
695 }
696 }
697 if !valid_source {
698 if let Some(error) = import_error {
699 return Err(io::Error::other(error));
700 }
701 }
702 let slots = primary
703 .slots
704 .read()
705 .unwrap_or_else(|error| error.into_inner());
706 let bytes = slots::slot_catalog_from_map(&slots)
707 .and_then(|catalog| catalog.encode())
708 .map_err(io::Error::other)?;
709 super::metadata::write(&store, super::metadata::SLOTS, &bytes)?;
710 }
711 for path in [&primary.slot_catalog_path, &primary.slot_path]
712 .into_iter()
713 .flatten()
714 {
715 super::metadata::retire_legacy(path)?;
716 }
717 primary.metadata_store = Some(store);
718 Ok(primary)
719 }
720
721 fn current_timeline(
722 &self,
723 root: &Path,
724 ) -> Result<reddb_file::TimelineId, reddb_file::RdbFileError> {
725 if let Some(store) = &self.metadata_store {
726 let path =
727 reddb_file::PrimaryReplicaFilePlan::new(root, reddb_file::TimelineId::initial())
728 .timeline_history_path();
729 return super::metadata::timeline(store, &path)
730 .map(|history| {
731 history
732 .current()
733 .unwrap_or_else(reddb_file::TimelineId::initial)
734 })
735 .map_err(reddb_file::RdbFileError::Io);
736 }
737 Self::primary_replica_current_timeline_for_root(root)
738 }
739
740 pub fn append_logical_record(&self, lsn: u64, encoded: Vec<u8>) {
741 if let Err(error) = self.try_append_logical_record(lsn, encoded) {
742 warn!(%error, lsn, "failed to publish logical replication record");
743 }
744 }
745
746 pub(crate) fn try_append_logical_record(&self, lsn: u64, encoded: Vec<u8>) -> io::Result<()> {
750 let _guard = self
751 .primary_replica_wal_lock
752 .lock()
753 .unwrap_or_else(|error| error.into_inner());
754 if let Some(plan) = &self.primary_replica_file_plan {
755 let timeline = self
756 .current_timeline(&plan.root)
757 .map_err(io::Error::other)?;
758 let plan = reddb_file::PrimaryReplicaFilePlan::new(plan.root.clone(), timeline);
759 plan.append_wal_record(lsn, &encoded)
760 .map_err(io::Error::other)?;
761 }
762 if let Some(spool) = &self.logical_wal_spool {
763 if spool.current_lsn() < lsn {
764 spool.append(lsn, &encoded)?;
765 }
766 }
767 if self.wal_buffer.current_lsn() < lsn {
768 self.wal_buffer.append(lsn, encoded);
769 }
770 let (lock, cvar) = &self.wal_appended;
771 let mut latest = lock.lock().unwrap_or_else(|error| error.into_inner());
772 *latest = (*latest).max(lsn);
773 cvar.notify_all();
774 Ok(())
775 }
776
777 pub fn wait_for_logical_lsn_after(&self, since_lsn: u64, timeout: Duration) -> bool {
778 if self.current_logical_lsn() > since_lsn {
779 return true;
780 }
781 let deadline = Instant::now() + timeout;
782 let (lock, cvar) = &self.wal_appended;
783 let mut latest = lock.lock().unwrap_or_else(|e| e.into_inner());
784 while *latest <= since_lsn {
785 let now = Instant::now();
786 if now >= deadline {
787 return false;
788 }
789 let remaining = deadline.saturating_duration_since(now);
790 let (guard, result) = cvar
791 .wait_timeout(latest, remaining)
792 .unwrap_or_else(|e| e.into_inner());
793 latest = guard;
794 if result.timed_out() && *latest <= since_lsn {
795 return false;
796 }
797 }
798 true
799 }
800
801 pub fn register_replica(&self, id: String) -> u64 {
802 self.register_replica_with_region(id, None)
803 }
804
805 pub fn register_replica_with_region(&self, id: String, region: Option<String>) -> u64 {
822 self.register_replica_with_region_at(id, region, self.current_logical_lsn())
823 }
824
825 fn register_replica_with_region_at(
826 &self,
827 id: String,
828 region: Option<String>,
829 initial_lsn: u64,
830 ) -> u64 {
831 let now_ms = crate::utils::now_unix_millis() as u128;
832 let resume_lsn = self.ensure_slot(&id, initial_lsn);
833 let mut replicas = self.replicas.write().unwrap_or_else(|e| e.into_inner());
834 if let Some(existing) = replicas.iter_mut().find(|r| r.id == id) {
835 existing.last_seen_at_unix_ms = now_ms;
836 if region.is_some() {
837 existing.region = region;
838 }
839 return resume_lsn;
840 }
841 replicas.push(ReplicaState {
842 id,
843 last_acked_lsn: resume_lsn,
844 last_sent_lsn: resume_lsn,
845 last_durable_lsn: resume_lsn,
846 apply_error_count: 0,
847 divergence_count: 0,
848 connected_at_unix_ms: now_ms,
849 last_seen_at_unix_ms: now_ms,
850 region,
851 rebootstrapping: false,
852 });
853 drop(replicas);
854 self.bump_topology_epoch();
855 resume_lsn
856 }
857
858 pub fn set_replica_rebootstrapping(&self, id: &str, rebootstrapping: bool) -> bool {
873 let mut replicas = self.replicas.write().unwrap_or_else(|e| e.into_inner());
874 let Some(state) = replicas.iter_mut().find(|r| r.id == id) else {
875 return false;
876 };
877 if state.rebootstrapping == rebootstrapping {
878 return true;
879 }
880 state.rebootstrapping = rebootstrapping;
881 drop(replicas);
882 self.bump_topology_epoch();
883 true
884 }
885
886 pub fn ensure_replica_registered(&self, id: &str, requested_since_lsn: u64) -> bool {
897 let already = self
898 .replicas
899 .read()
900 .unwrap_or_else(|e| e.into_inner())
901 .iter()
902 .any(|r| r.id == id);
903 if already {
904 return false;
905 }
906 let initial_lsn = requested_since_lsn.min(self.current_logical_lsn());
907 self.register_replica_with_region_at(id.to_string(), None, initial_lsn);
908 true
909 }
910
911 pub fn unregister_replica(&self, id: &str) -> bool {
915 let mut replicas = self.replicas.write().unwrap_or_else(|e| e.into_inner());
916 let before = replicas.len();
917 replicas.retain(|r| r.id != id);
918 let removed = replicas.len() != before;
919 drop(replicas);
920 if removed {
921 self.routing_endpoints
922 .write()
923 .expect("replica routing endpoint map lock must not be poisoned")
924 .remove(id);
925 self.commit_waiter.drop_replica(id);
926 self.bump_topology_epoch();
927 }
928 removed
929 }
930
931 pub fn topology_epoch(&self) -> u64 {
934 self.topology_epoch
935 .load(std::sync::atomic::Ordering::Relaxed)
936 }
937
938 pub fn bump_topology_epoch(&self) {
946 self.topology_epoch
947 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
948 }
949
950 pub fn ack_replica(&self, id: &str, lsn: u64) {
951 let now_ms = crate::utils::now_unix_millis() as u128;
952 let mut replicas = self.replicas.write().unwrap_or_else(|e| e.into_inner());
953 if let Some(r) = replicas.iter_mut().find(|r| r.id == id) {
954 r.last_acked_lsn = r.last_acked_lsn.max(lsn);
955 r.last_durable_lsn = r.last_durable_lsn.max(lsn);
956 r.last_seen_at_unix_ms = now_ms;
957 }
958 drop(replicas);
959 self.commit_waiter.record_replica_ack(id, lsn);
960 }
961
962 pub fn ack_replica_lsn(&self, id: &str, applied_lsn: u64, durable_lsn: u64) {
968 self.ack_replica_lsn_with_observability(id, applied_lsn, durable_lsn, 0, 0);
969 }
970
971 pub fn ack_replica_lsn_with_observability(
972 &self,
973 id: &str,
974 applied_lsn: u64,
975 durable_lsn: u64,
976 apply_error_count: u64,
977 divergence_count: u64,
978 ) {
979 let now_ms = crate::utils::now_unix_millis() as u128;
980 self.advance_slot(id, applied_lsn, durable_lsn, now_ms);
981 let mut replicas = self.replicas.write().unwrap_or_else(|e| e.into_inner());
982 if let Some(r) = replicas.iter_mut().find(|r| r.id == id) {
983 r.last_acked_lsn = r.last_acked_lsn.max(applied_lsn);
984 r.last_durable_lsn = r.last_durable_lsn.max(durable_lsn);
985 r.apply_error_count = r.apply_error_count.max(apply_error_count);
986 r.divergence_count = r.divergence_count.max(divergence_count);
987 r.last_seen_at_unix_ms = now_ms;
988 }
989 drop(replicas);
993 self.commit_waiter.record_replica_ack(id, durable_lsn);
994 }
995
996 pub fn note_replica_pull(&self, id: &str, last_sent_lsn: u64) {
1001 let now_ms = crate::utils::now_unix_millis() as u128;
1002 self.touch_slot(id, now_ms);
1003 let mut replicas = self.replicas.write().unwrap_or_else(|e| e.into_inner());
1004 if let Some(r) = replicas.iter_mut().find(|r| r.id == id) {
1005 r.last_sent_lsn = r.last_sent_lsn.max(last_sent_lsn);
1006 r.last_seen_at_unix_ms = now_ms;
1007 }
1008 }
1009
1010 pub fn replica_snapshots(&self) -> Vec<ReplicaState> {
1014 self.replicas
1015 .read()
1016 .unwrap_or_else(|e| e.into_inner())
1017 .clone()
1018 }
1019
1020 pub fn replication_progress(&self) -> Option<ReplicationProgress> {
1021 let replicas = self.replicas.read().unwrap_or_else(|e| e.into_inner());
1022 ReplicationProgress::from_replicas(&replicas)
1023 }
1024
1025 pub fn slot_snapshots(&self) -> Vec<ReplicationSlot> {
1026 self.slots
1027 .read()
1028 .unwrap_or_else(|e| e.into_inner())
1029 .values()
1030 .cloned()
1031 .collect()
1032 }
1033
1034 pub fn retention_floor_lsn(&self) -> Option<u64> {
1035 let slots = self.slots.read().unwrap_or_else(|e| e.into_inner());
1038 if self.ensure_slot_metadata_durable().is_err() {
1039 return Some(0);
1040 }
1041 slots
1042 .values()
1043 .filter(|slot| slot.invalidation_reason.is_none())
1044 .map(|slot| slot.restart_lsn)
1045 .min()
1046 }
1047
1048 pub fn prune_retained_wal_through(&self, archived_lsn: u64) -> io::Result<u64> {
1049 self.enforce_retention_limits(crate::utils::now_unix_millis() as u128);
1050 let prune_lsn = self
1051 .retention_floor_lsn()
1052 .map(|floor| floor.min(archived_lsn))
1053 .unwrap_or(archived_lsn);
1054 if prune_lsn > 0 {
1055 if let Some(spool) = &self.logical_wal_spool {
1056 spool.prune_through(prune_lsn)?;
1057 }
1058 self.wal_buffer.prune_through(prune_lsn);
1059 }
1060 Ok(prune_lsn)
1061 }
1062
1063 pub fn replica_count(&self) -> usize {
1064 self.replicas
1065 .read()
1066 .unwrap_or_else(|e| e.into_inner())
1067 .len()
1068 }
1069
1070 pub fn current_logical_lsn(&self) -> u64 {
1074 self.logical_wal_spool
1075 .as_ref()
1076 .map(|spool| spool.current_lsn())
1077 .unwrap_or_else(|| self.wal_buffer.current_lsn())
1078 }
1079
1080 fn ensure_slot(&self, id: &str, initial_lsn: u64) -> u64 {
1081 let now_ms = crate::utils::now_unix_millis() as u128;
1082 let _publication = self
1083 .metadata_store
1084 .as_ref()
1085 .map(|store| store.vault_publication_lock.lock());
1086 let mut slots = self.slots.write().unwrap_or_else(|e| e.into_inner());
1087 if let Some(slot) = slots.get_mut(id) {
1088 slot.last_seen_at_unix_ms = now_ms;
1089 let restart_lsn = slot.restart_lsn;
1090 self.persist_slots_locked(&slots);
1091 return restart_lsn;
1092 }
1093 let mut slot = ReplicationSlot::new(
1094 id.to_string(),
1095 reddb_file::TimelineId::initial(),
1096 initial_lsn,
1097 );
1098 slot.last_seen_at_unix_ms = now_ms;
1099 slots.insert(id.to_string(), slot);
1100 let restart_lsn = initial_lsn;
1101 self.persist_slots_locked(&slots);
1102 restart_lsn
1103 }
1104
1105 fn advance_slot(&self, id: &str, confirmed_lsn: u64, restart_lsn: u64, now_ms: u128) {
1106 let _publication = self
1107 .metadata_store
1108 .as_ref()
1109 .map(|store| store.vault_publication_lock.lock());
1110 let mut slots = self.slots.write().unwrap_or_else(|e| e.into_inner());
1111 let slot = slots.entry(id.to_string()).or_insert_with(|| {
1112 let mut slot =
1113 ReplicationSlot::new(id.to_string(), reddb_file::TimelineId::initial(), 0);
1114 slot.last_seen_at_unix_ms = now_ms;
1115 slot
1116 });
1117 if slot.invalidation_reason.is_some() {
1118 return;
1119 }
1120 slot.confirmed_write_lsn = slot.confirmed_lsn().max(confirmed_lsn).max(restart_lsn);
1121 slot.restart_lsn = slot.restart_lsn.max(restart_lsn);
1122 slot.confirmed_flush_lsn = slot.confirmed_flush_lsn.max(slot.restart_lsn);
1123 slot.confirmed_apply_lsn = slot.confirmed_apply_lsn.max(slot.restart_lsn);
1124 slot.last_seen_at_unix_ms = now_ms;
1125 self.persist_slots_locked(&slots);
1126 }
1127
1128 pub fn touch_slot(&self, id: &str, now_ms: u128) {
1129 let _publication = self
1130 .metadata_store
1131 .as_ref()
1132 .map(|store| store.vault_publication_lock.lock());
1133 let mut slots = self.slots.write().unwrap_or_else(|e| e.into_inner());
1134 let mut changed = false;
1135 if let Some(slot) = slots.get_mut(id) {
1136 if slot.invalidation_reason.is_none() {
1137 slot.last_seen_at_unix_ms = now_ms;
1138 changed = true;
1139 }
1140 }
1141 if changed {
1142 self.persist_slots_locked(&slots);
1143 }
1144 }
1145
1146 pub fn enforce_retention_limits(
1147 &self,
1148 now_ms: u128,
1149 ) -> Vec<(String, ReplicationSlotInvalidationCause)> {
1150 let current_lsn = self.current_logical_lsn();
1151 let mut invalidated = Vec::new();
1152 let _publication = self
1153 .metadata_store
1154 .as_ref()
1155 .map(|store| store.vault_publication_lock.lock());
1156 let mut slots = self.slots.write().unwrap_or_else(|e| e.into_inner());
1157 for slot in slots.values_mut() {
1158 if slot.invalidation_reason.is_some() {
1159 continue;
1160 }
1161 let reason = if self.slot_retention_max_lag_lsn > 0
1162 && current_lsn.saturating_sub(slot.restart_lsn) > self.slot_retention_max_lag_lsn
1163 {
1164 Some(ReplicationSlotInvalidationCause::Horizon)
1165 } else if self.slot_idle_timeout_ms > 0
1166 && now_ms.saturating_sub(slot.last_seen_at_unix_ms)
1167 > u128::from(self.slot_idle_timeout_ms)
1168 {
1169 Some(ReplicationSlotInvalidationCause::IdleTimeout)
1170 } else {
1171 None
1172 };
1173 if let Some(reason) = reason {
1174 slot.invalidation_reason = Some(reason);
1175 slot.invalidated_at_unix_ms = Some(now_ms);
1176 invalidated.push((slot.replica_id.clone(), reason));
1177 }
1178 }
1179 if !invalidated.is_empty() {
1180 self.persist_slots_locked(&slots);
1181 }
1182 invalidated
1183 }
1184
1185 pub fn slot_rebootstrap_reason(
1186 &self,
1187 id: &str,
1188 requested_since_lsn: u64,
1189 oldest_available_lsn: Option<u64>,
1190 ) -> Option<ReplicationSlotInvalidationCause> {
1191 let now_ms = crate::utils::now_unix_millis() as u128;
1192 let _publication = self
1193 .metadata_store
1194 .as_ref()
1195 .map(|store| store.vault_publication_lock.lock());
1196 let mut slots = self.slots.write().unwrap_or_else(|e| e.into_inner());
1197 let slot = slots.get_mut(id)?;
1198 if let Some(reason) = slot.invalidation_reason {
1199 return Some(reason);
1200 }
1201 let slot_floor = slot.restart_lsn.max(requested_since_lsn);
1202 if oldest_available_lsn
1203 .map(|oldest| oldest > slot_floor.saturating_add(1))
1204 .unwrap_or(false)
1205 {
1206 slot.invalidation_reason = Some(ReplicationSlotInvalidationCause::WalRemoved);
1207 slot.invalidated_at_unix_ms = Some(now_ms);
1208 self.persist_slots_locked(&slots);
1209 return Some(ReplicationSlotInvalidationCause::WalRemoved);
1210 }
1211 None
1212 }
1213
1214 pub fn plan_replica_resume(
1223 &self,
1224 id: &str,
1225 requested_since_lsn: u64,
1226 oldest_available_lsn: Option<u64>,
1227 ) -> ResumeMode {
1228 if let Some(cause) =
1229 self.slot_rebootstrap_reason(id, requested_since_lsn, oldest_available_lsn)
1230 {
1231 self.full_resync_count
1232 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
1233 return ResumeMode::FullRebootstrap { cause };
1234 }
1235 let resume_lsn = self
1236 .slot_snapshots()
1237 .into_iter()
1238 .find(|slot| slot.replica_id == id)
1239 .map(|slot| requested_since_lsn.max(slot.restart_lsn))
1240 .unwrap_or(requested_since_lsn);
1241 self.partial_resync_count
1242 .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
1243 ResumeMode::PartialResync { resume_lsn }
1244 }
1245
1246 pub fn partial_resync_count(&self) -> u64 {
1249 self.partial_resync_count
1250 .load(std::sync::atomic::Ordering::Relaxed)
1251 }
1252
1253 pub fn full_resync_count(&self) -> u64 {
1257 self.full_resync_count
1258 .load(std::sync::atomic::Ordering::Relaxed)
1259 }
1260
1261 fn persist_slots_locked(&self, slots: &BTreeMap<String, ReplicationSlot>) {
1262 if let Some(store) = &self.metadata_store {
1263 let result = slots::slot_catalog_from_map(slots)
1264 .and_then(|catalog| catalog.encode())
1265 .map_err(io::Error::other)
1266 .and_then(|bytes| super::metadata::write(store, super::metadata::SLOTS, &bytes));
1267 let mut last_error = self
1268 .metadata_error
1269 .lock()
1270 .unwrap_or_else(|error| error.into_inner());
1271 if let Err(error) = result {
1272 *last_error = Some(error.to_string());
1273 warn!(target: "reddb::replication::slots", %error,
1274 "failed to persist replication slots inside database");
1275 } else {
1276 *last_error = None;
1277 }
1278 return;
1279 }
1280 if let Err(err) = persist_replication_slots(self.slot_path.as_deref(), slots) {
1281 warn!(
1282 target: "reddb::replication::slots",
1283 error = %err,
1284 "failed to persist replication slots"
1285 );
1286 }
1287 if let Err(err) = persist_replication_slot_catalog(self.slot_catalog_path.as_deref(), slots)
1288 {
1289 warn!(
1290 target: "reddb::replication::slots",
1291 error = %err,
1292 "failed to persist binary replication slot catalog"
1293 );
1294 }
1295 }
1296
1297 pub(crate) fn ensure_slot_metadata_durable(&self) -> io::Result<()> {
1298 match self
1299 .metadata_error
1300 .lock()
1301 .unwrap_or_else(|error| error.into_inner())
1302 .as_ref()
1303 {
1304 Some(error) => Err(io::Error::other(error.clone())),
1305 None => Ok(()),
1306 }
1307 }
1308}
1309
1310#[cfg(test)]
1311mod tests;