krabka_remote_storage/diskless/capture/
projection.rs

1//! Reduction of committed index events into a deterministic capture boundary.
2
3use {
4    super::{
5        super::{
6            WalIndexEntry,
7            index::{
8                WalDeleteFloorKey, WalDeleteFloorRecord, WalFlushRecord, WalIndexKey,
9                wal_index_key_version,
10            },
11        },
12        CapturedWalRange, DisklessPartitionCapture, DisklessWalCapture,
13    },
14    std::{
15        collections::{BTreeMap, HashMap},
16        hash::BuildHasher,
17    },
18    uuid::Uuid,
19};
20
21/// A partition in a capture, ordered by topic name, then topic id, then partition.
22#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
23struct GroupKey {
24    topic: String,
25    topic_id: Uuid,
26    partition: i32,
27}
28
29/// Reducer for committed keyed index events.
30#[derive(Default)]
31pub struct WalCaptureProjection {
32    ranges: HashMap<WalIndexKey, CapturedWalRange>,
33    floors: HashMap<WalDeleteFloorKey, i64>,
34}
35
36impl WalCaptureProjection {
37    /// Apply one committed keyed value or tombstone.
38    ///
39    /// # Errors
40    /// Returns an error for malformed keys, values, or inconsistent records.
41    pub fn apply(&mut self, key: Option<&[u8]>, value: Option<&[u8]>) -> Result<(), String> {
42        // Every index record is keyed: the flusher publishes one keyed record
43        // per range, and compaction keeps only the latest per key. An
44        // unkeyed record has no range it could be the latest for.
45        let key = key.ok_or_else(|| "diskless WAL index record has no key".to_owned())?;
46        if let Some(floor_key) = WalDeleteFloorKey::from_bytes(key) {
47            if let Some(value) = value {
48                let record = WalDeleteFloorRecord::from_bytes(value)?;
49                if record.topic_id != floor_key.topic_id
50                    || record.partition != floor_key.partition
51                    || record.floor < 0
52                {
53                    return Err("diskless WAL delete-floor key/value mismatch".to_owned());
54                }
55                let floor = self.floors.entry(floor_key).or_default();
56                *floor = (*floor).max(record.floor);
57            } else {
58                self.floors.remove(&floor_key);
59            }
60            return Ok(());
61        }
62        let range_key =
63            WalIndexKey::from_bytes(key).ok_or_else(|| match wal_index_key_version(key) {
64                Some(version)
65                    if version != WalIndexKey::KEY_VERSION
66                        && version != WalDeleteFloorKey::KEY_VERSION =>
67                {
68                    format!("unknown diskless WAL index key version {version}")
69                }
70                _ => "invalid diskless WAL index key".to_owned(),
71            })?;
72        let Some(value) = value else {
73            self.ranges.remove(&range_key);
74            return Ok(());
75        };
76        let record = WalFlushRecord::from_bytes(value)?;
77        let entry = record
78            .entries
79            .into_iter()
80            .find(|entry| WalIndexKey::from(entry) == range_key)
81            .ok_or_else(|| "diskless WAL index key/value mismatch".to_owned())?;
82        Self::validate_range(&entry)?;
83        self.store_range(range_key, record.object_key, entry);
84        Ok(())
85    }
86
87    fn validate_range(entry: &WalIndexEntry) -> Result<(), String> {
88        if entry.first_offset < 0 || entry.last_offset < entry.first_offset || entry.byte_len == 0 {
89            return Err("invalid diskless WAL index range".to_owned());
90        }
91        Ok(())
92    }
93
94    fn store_range(&mut self, key: WalIndexKey, object_key: String, entry: WalIndexEntry) {
95        self.ranges
96            .insert(key, CapturedWalRange { object_key, entry });
97    }
98
99    /// Freeze the projection into a deterministic portable capture.
100    ///
101    /// # Errors
102    /// Returns an error for unknown topic ids or inconsistent live ranges.
103    pub fn capture<S: BuildHasher>(
104        &self,
105        topics: &HashMap<Uuid, (String, i32), S>,
106        source_cutoffs: Vec<i64>,
107        captured_at_ms: u64,
108    ) -> Result<DisklessWalCapture, String> {
109        let mut grouped: BTreeMap<GroupKey, Vec<CapturedWalRange>> = BTreeMap::new();
110        for (topic_id, (topic, partition_count)) in topics {
111            if *partition_count < 0 {
112                return Err(format!(
113                    "negative partition count for diskless topic {topic}"
114                ));
115            }
116            for partition in 0..*partition_count {
117                grouped
118                    .entry(GroupKey {
119                        topic: topic.clone(),
120                        topic_id: *topic_id,
121                        partition,
122                    })
123                    .or_default();
124            }
125        }
126        for range in self.ranges.values() {
127            let Some((topic, _)) = topics.get(&range.entry.topic_id) else {
128                continue;
129            };
130            grouped
131                .entry(GroupKey {
132                    topic: topic.clone(),
133                    topic_id: range.entry.topic_id,
134                    partition: range.entry.partition,
135                })
136                .or_default()
137                .push(range.clone());
138        }
139        for WalDeleteFloorKey {
140            topic_id,
141            partition,
142        } in self.floors.keys()
143        {
144            let Some((topic, _)) = topics.get(topic_id) else {
145                continue;
146            };
147            grouped
148                .entry(GroupKey {
149                    topic: topic.clone(),
150                    topic_id: *topic_id,
151                    partition: *partition,
152                })
153                .or_default();
154        }
155        let mut partitions = Vec::with_capacity(grouped.len());
156        for (
157            GroupKey {
158                topic,
159                topic_id,
160                partition,
161            },
162            mut ranges,
163        ) in grouped
164        {
165            ranges.sort_by_key(|range| range.entry.first_offset);
166            if ranges
167                .iter()
168                .zip(ranges.iter().skip(1))
169                .any(|(earlier, later)| earlier.entry.last_offset >= later.entry.first_offset)
170            {
171                return Err(format!(
172                    "overlapping diskless WAL ranges for {topic}-{partition}"
173                ));
174            }
175            let mut runs: Vec<CapturedWalRange> = Vec::with_capacity(ranges.len());
176            for range in ranges {
177                let joins_previous = runs.last().is_some_and(|previous| {
178                    previous.object_key == range.object_key
179                        && previous.entry.last_offset.checked_add(1)
180                            == Some(range.entry.first_offset)
181                        && previous
182                            .entry
183                            .byte_start
184                            .checked_add(u64::from(previous.entry.byte_len))
185                            == Some(range.entry.byte_start)
186                });
187                if joins_previous {
188                    if let Some(previous) = runs.last_mut() {
189                        previous.entry.last_offset = range.entry.last_offset;
190                        previous.entry.byte_len = previous
191                            .entry
192                            .byte_len
193                            .checked_add(range.entry.byte_len)
194                            .ok_or_else(|| {
195                                format!("diskless WAL run exceeds 4 GiB for {topic}-{partition}")
196                            })?;
197                        previous.entry.max_timestamp_ms = previous
198                            .entry
199                            .max_timestamp_ms
200                            .max(range.entry.max_timestamp_ms);
201                    }
202                } else {
203                    runs.push(range);
204                }
205            }
206            let ranges = runs;
207            let delete_floor = self
208                .floors
209                .get(&WalDeleteFloorKey {
210                    topic_id,
211                    partition,
212                })
213                .copied()
214                .unwrap_or(0);
215            let recovery_cutoff = ranges
216                .last()
217                .map_or(delete_floor, |range| {
218                    range.entry.last_offset.saturating_add(1)
219                })
220                .max(delete_floor);
221            partitions.push(DisklessPartitionCapture {
222                topic,
223                topic_id,
224                partition,
225                delete_floor,
226                recovery_cutoff,
227                ranges,
228            });
229        }
230        Ok(DisklessWalCapture {
231            format_version: DisklessWalCapture::FORMAT_VERSION,
232            captured_at_ms,
233            source_cutoffs,
234            partitions,
235            metadata_snapshot_sha256: None,
236            rlmm_snapshot_sha256: None,
237            group_offsets_sha256: None,
238            authentication: None,
239        })
240    }
241}
krabka_remote_storage/diskless/capture/projection.rs

krabka_remote_storage/diskless/capture/
projection.rs
