krabka_replicator/
offset_sync_store.rs

1//! In-memory source-to-target offset translation built from MM2 offset-sync
2//! records.
3//! Translation is "at-or-before": find the latest sync with upstream <= committed,
4//! then downstream + (committed - upstream). Never maps past un-replicated data.
5
6use std::collections::BTreeMap;
7
8use crate::{
9    ids::{CommittedOffset, DownstreamOffset, PartitionIndex, UpstreamOffset},
10    mm2::OffsetSync,
11};
12
13/// Accumulates [`OffsetSync`] records and translates committed source offsets to
14/// their corresponding target offsets.
15///
16/// The store deliberately retains every distinct upstream anchor for its
17/// process lifetime. This conservative policy lets consumer groups that have
18/// been inactive for a long time still translate from an old anchor. A caller
19/// must clear the store when it detects that the offset-sync topic was truncated
20/// or recreated; otherwise, mappings are never expired merely due to age or the
21/// arrival of newer syncs.
22#[derive(Default)]
23pub struct OffsetSyncStore {
24    /// Outer key: (topic, partition). Inner key: upstream offset → downstream offset.
25    syncs: BTreeMap<(String, PartitionIndex), BTreeMap<UpstreamOffset, DownstreamOffset>>,
26}
27
28impl OffsetSyncStore {
29    /// Ingest one offset-sync record. It replaces any prior entry for the same
30    /// `(topic, partition, upstream)` triple.
31    pub fn ingest(&mut self, s: OffsetSync) {
32        self.syncs
33            .entry((s.topic, s.partition))
34            .or_default()
35            .insert(s.upstream, s.downstream);
36    }
37
38    /// Remove all translation anchors.
39    ///
40    /// Incremental readers must call this before rebuilding when any input
41    /// partition cursor is no longer valid. Offset-sync topic partitions do not
42    /// correspond one-to-one with the source topic-partitions stored here, so a
43    /// partial reset could retain stale translations.
44    pub fn clear(&mut self) {
45        self.syncs.clear();
46    }
47
48    /// Translate a committed source offset to its target offset.
49    ///
50    /// Returns `None` when:
51    /// - the `(topic, partition)` pair has no syncs, or
52    /// - `committed` is below every known upstream offset (un-replicated data).
53    #[must_use]
54    pub fn translate(
55        &self,
56        topic: &str,
57        partition: PartitionIndex,
58        committed: CommittedOffset,
59    ) -> Option<DownstreamOffset> {
60        let m = self.syncs.get(&(topic.to_string(), partition))?;
61        // The inner map is keyed by `UpstreamOffset`; `committed` lives in the
62        // same (source) address space, so bound the range at the matching
63        // upstream offset.
64        let (&up, &down) = m.range(..=UpstreamOffset::from(committed)).next_back()?;
65        // `committed - up` is a source-space delta; adding it to the paired
66        // downstream offset yields the target offset. The arithmetic crosses
67        // newtype boundaries, so unwrap via `.0` rather than deriving `Add`/`Sub`.
68        Some(DownstreamOffset(down.0 + (committed.0 - up.0)))
69    }
70}
71
72#[cfg(test)]
73mod tests {
74    use super::*;
75    use crate::mm2::OffsetSync;
76
77    #[test]
78    fn translates_at_or_before_nearest_sync() {
79        let mut s = OffsetSyncStore::default();
80        s.ingest(OffsetSync {
81            topic: "orders".into(),
82            partition: PartitionIndex(0),
83            upstream: UpstreamOffset(100),
84            downstream: DownstreamOffset(70),
85        });
86        s.ingest(OffsetSync {
87            topic: "orders".into(),
88            partition: PartitionIndex(0),
89            upstream: UpstreamOffset(200),
90            downstream: DownstreamOffset(165),
91        });
92        for (partition, upstream, want) in [
93            (0, 250, Some(215)), // 165 + (250-200)
94            (0, 150, Some(120)), // 70 + (150-100)
95            (0, 50, None),       // below first sync
96            (9, 100, None),      // unknown partition
97        ] {
98            assert2::assert!(
99                s.translate(
100                    "orders",
101                    PartitionIndex(partition),
102                    CommittedOffset(upstream)
103                ) == want.map(DownstreamOffset)
104            );
105        }
106    }
107
108    #[test]
109    fn retains_old_anchors_after_newer_syncs_arrive() {
110        let mut store = OffsetSyncStore::default();
111        for (upstream, downstream) in [(10, 100), (20, 200), (30, 300)] {
112            store.ingest(OffsetSync {
113                topic: "orders".into(),
114                partition: PartitionIndex(0),
115                upstream: UpstreamOffset(upstream),
116                downstream: DownstreamOffset(downstream),
117            });
118        }
119
120        assert2::assert!(
121            store.translate("orders", PartitionIndex(0), CommittedOffset(15))
122                == Some(DownstreamOffset(105))
123        );
124        assert2::assert!(
125            store.translate("orders", PartitionIndex(0), CommittedOffset(35))
126                == Some(DownstreamOffset(305))
127        );
128    }
129
130    #[test]
131    fn clear_removes_every_mapping() {
132        let mut store = OffsetSyncStore::default();
133        for (topic, partition, downstream) in
134            [("orders", 0, 100), ("orders", 1, 200), ("users", 0, 300)]
135        {
136            store.ingest(OffsetSync {
137                topic: topic.into(),
138                partition: PartitionIndex(partition),
139                upstream: UpstreamOffset(10),
140                downstream: DownstreamOffset(downstream),
141            });
142        }
143
144        store.clear();
145
146        assert2::assert!(
147            store
148                .translate("orders", PartitionIndex(0), CommittedOffset(10))
149                .is_none()
150        );
151        assert2::assert!(
152            store
153                .translate("orders", PartitionIndex(1), CommittedOffset(10))
154                .is_none()
155        );
156        assert2::assert!(
157            store
158                .translate("users", PartitionIndex(0), CommittedOffset(10))
159                .is_none()
160        );
161    }
162}
krabka_replicator/offset_sync_store.rs

krabka_replicator/
offset_sync_store.rs
