laser_wire/
snapshot.rs

1use crate::agent::ConversationId;
2use crate::error::InvalidError;
3use crate::validate::Validate;
4use serde::{Deserialize, Deserializer, Serialize, Serializer};
5
6/// One source partition and its last folded offset.
7#[derive(Clone, Copy, Debug, PartialEq, Eq)]
8pub struct SnapshotOffset {
9    pub topic_id: u32,
10    pub topic_created_at_micros: u64,
11    pub partition_id: u32,
12    pub offset: u64,
13}
14
15impl SnapshotOffset {
16    pub const fn new(
17        topic_id: u32,
18        topic_created_at_micros: u64,
19        partition_id: u32,
20        offset: u64,
21    ) -> Self {
22        Self {
23            topic_id,
24            topic_created_at_micros,
25            partition_id,
26            offset,
27        }
28    }
29
30    fn key(self) -> (u32, u64, u32) {
31        (
32            self.topic_id,
33            self.topic_created_at_micros,
34            self.partition_id,
35        )
36    }
37}
38
39impl Serialize for SnapshotOffset {
40    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
41    where
42        S: Serializer,
43    {
44        (
45            self.topic_id,
46            self.topic_created_at_micros,
47            self.partition_id,
48            self.offset,
49        )
50            .serialize(serializer)
51    }
52}
53
54impl<'de> Deserialize<'de> for SnapshotOffset {
55    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
56    where
57        D: Deserializer<'de>,
58    {
59        let (topic_id, topic_created_at_micros, partition_id, offset) =
60            <(u32, u64, u32, u64)>::deserialize(deserializer)?;
61        Ok(Self::new(
62            topic_id,
63            topic_created_at_micros,
64            partition_id,
65            offset,
66        ))
67    }
68}
69
70/// A fold state and the exact source positions included in it.
71#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
72pub struct FoldSnapshot {
73    pub stream: String,
74    pub stream_id: u32,
75    pub stream_created_at_micros: u64,
76    pub conversation: ConversationId,
77    pub fold: String,
78    /// Sorted by topic ID, topic creation time, and partition ID.
79    pub as_of: Vec<SnapshotOffset>,
80    /// The application chooses the codec for these bytes.
81    #[serde(with = "crate::encoding::bin_bytes")]
82    pub state: Vec<u8>,
83}
84
85impl FoldSnapshot {
86    /// Return the next offset for one exact source partition.
87    pub fn resume_offset(
88        &self,
89        topic_id: u32,
90        topic_created_at_micros: u64,
91        partition_id: u32,
92    ) -> u64 {
93        let key = (topic_id, topic_created_at_micros, partition_id);
94        self.as_of
95            .binary_search_by_key(&key, |entry| entry.key())
96            .ok()
97            .map_or(0, |index| self.as_of[index].offset.saturating_add(1))
98    }
99}
100
101impl Validate for FoldSnapshot {
102    fn validate(&self) -> Result<(), InvalidError> {
103        if self.stream.is_empty() || self.fold.is_empty() {
104            return Err(InvalidError::new(
105                "snapshot stream and fold must be non-empty",
106            ));
107        }
108        if self
109            .as_of
110            .windows(2)
111            .any(|pair| pair[0].key() >= pair[1].key())
112        {
113            return Err(InvalidError::new(
114                "snapshot offsets must be sorted with no duplicate source",
115            ));
116        }
117        Ok(())
118    }
119}
120
121#[cfg(test)]
122mod tests {
123    use super::*;
124
125    fn snapshot() -> FoldSnapshot {
126        FoldSnapshot {
127            stream: "agents".to_owned(),
128            stream_id: 0,
129            stream_created_at_micros: 100,
130            conversation: ConversationId::from_u128(1),
131            fold: "planner".to_owned(),
132            as_of: vec![
133                SnapshotOffset::new(2, 20, 0, 41),
134                SnapshotOffset::new(2, 20, 1, 9),
135            ],
136            state: vec![1, 2, 3],
137        }
138    }
139
140    #[test]
141    fn given_a_snapshot_when_resumed_then_should_match_the_exact_topic_generation() {
142        let snapshot = snapshot();
143        snapshot.validate().expect("valid snapshot");
144        assert_eq!(snapshot.resume_offset(2, 20, 0), 42);
145        assert_eq!(snapshot.resume_offset(2, 20, 1), 10);
146        assert_eq!(snapshot.resume_offset(2, 21, 0), 0);
147        assert_eq!(snapshot.resume_offset(3, 20, 0), 0);
148    }
149
150    #[test]
151    fn given_duplicate_or_unsorted_sources_when_validated_then_should_reject_them() {
152        let mut snapshot = snapshot();
153        snapshot.as_of.push(SnapshotOffset::new(2, 20, 1, 10));
154        assert!(snapshot.validate().is_err());
155        snapshot.as_of[2] = SnapshotOffset::new(1, 20, 0, 10);
156        assert!(snapshot.validate().is_err());
157    }
158}
laser_wire/snapshot.rs

laser_wire/
snapshot.rs
