Skip to main content

faucet_source_kinesis/
state.rs

1//! Per-shard bookmark format: `{ "shards": { "<shard-id>": "<sequence>" } }`.
2//!
3//! Every emitted page carries the **cumulative** map, so any page's bookmark
4//! is a valid resume point. On resume, a bookmarked shard restarts at
5//! `AFTER_SEQUENCE_NUMBER(<persisted>)`; shards absent from the map use the
6//! configured start position.
7
8use serde::{Deserialize, Serialize};
9use serde_json::Value;
10use std::collections::BTreeMap;
11
12/// The persisted bookmark value.
13#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
14pub struct ShardBookmarks {
15    /// Shard id → last durably-emitted sequence number.
16    #[serde(default)]
17    pub shards: BTreeMap<String, String>,
18}
19
20impl ShardBookmarks {
21    /// Record a shard's newest sequence number.
22    pub fn advance(&mut self, shard_id: &str, sequence: &str) {
23        self.shards
24            .insert(shard_id.to_string(), sequence.to_string());
25    }
26
27    /// The persisted sequence for a shard, if any.
28    pub fn get(&self, shard_id: &str) -> Option<&str> {
29        self.shards.get(shard_id).map(String::as_str)
30    }
31
32    pub fn to_value(&self) -> Value {
33        serde_json::to_value(self).unwrap_or(Value::Null)
34    }
35
36    /// Parse a bookmark `Value`; a malformed value is treated as absent
37    /// (never fails a run — worst case the shard re-reads from the start
38    /// position, which is at-least-once-safe).
39    pub fn from_value(v: &Value) -> Self {
40        serde_json::from_value(v.clone()).unwrap_or_default()
41    }
42}
43
44/// The source's stable state key.
45pub fn state_key(stream_name: &str) -> String {
46    format!("kinesis:{stream_name}")
47}
48
49#[cfg(test)]
50mod tests {
51    use super::*;
52    use serde_json::json;
53
54    #[test]
55    fn bookmark_round_trips_and_advances() {
56        let mut b = ShardBookmarks::default();
57        b.advance("shardId-000000000000", "491");
58        b.advance("shardId-000000000001", "530");
59        b.advance("shardId-000000000000", "495"); // newer overwrites
60        let v = b.to_value();
61        assert_eq!(v["shards"]["shardId-000000000000"], "495");
62        let back = ShardBookmarks::from_value(&v);
63        assert_eq!(back, b);
64        assert_eq!(back.get("shardId-000000000001"), Some("530"));
65        assert_eq!(back.get("missing"), None);
66    }
67
68    #[test]
69    fn malformed_bookmark_is_treated_as_fresh() {
70        assert_eq!(
71            ShardBookmarks::from_value(&json!("not-a-map")),
72            ShardBookmarks::default()
73        );
74        assert_eq!(
75            ShardBookmarks::from_value(&json!(null)),
76            ShardBookmarks::default()
77        );
78    }
79
80    #[test]
81    fn state_key_shape_is_valid() {
82        assert_eq!(state_key("events"), "kinesis:events");
83        faucet_core::state::validate_state_key(&state_key("events")).unwrap();
84    }
85}