millipede_core/
snapshot.rs

1//! Failure-artifact capture and reload support.
2//!
3//! Artifacts are captured by crawler kinds on the `HandlerFailed` cleanup path.
4//! Execute-time failures that occur before a handler context exists produce no snapshot.
5
6use crate::{
7    request::Request,
8    storage::{KeyValueStore, StorageResult},
9};
10use std::{
11    collections::hash_map::DefaultHasher,
12    hash::{Hash, Hasher},
13    sync::Arc,
14};
15
16/// A failure-time artifact reloaded from storage.
17#[derive(Debug, Clone, PartialEq)]
18#[non_exhaustive]
19pub struct ErrorSnapshot {
20    /// MIME content type of the captured artifact.
21    pub content_type: String,
22    /// Captured artifact bytes.
23    pub bytes: bytes::Bytes,
24}
25
26/// Captures and reloads failure-time artifacts in a crawler key-value store.
27pub struct ErrorSnapshotter {
28    kvs: Arc<dyn KeyValueStore>,
29}
30
31impl ErrorSnapshotter {
32    /// Creates a snapshotter backed by `kvs`.
33    pub fn new(kvs: Arc<dyn KeyValueStore>) -> Self {
34        Self { kvs }
35    }
36
37    /// Returns the deterministic base storage key for `request`.
38    pub fn base_key(request: &Request) -> String {
39        let mut h = DefaultHasher::new();
40        request.unique_key.hash(&mut h);
41        format!("ERROR_SNAPSHOT_{:016x}", h.finish())
42    }
43
44    /// Captures an artifact under the request base key and `suffix`.
45    pub async fn capture(
46        &self,
47        request: &Request,
48        suffix: &str,
49        bytes: bytes::Bytes,
50        content_type: &str,
51    ) -> StorageResult<String> {
52        let key = format!("{}.{suffix}", Self::base_key(request));
53        self.kvs.set_bytes(&key, bytes, content_type).await?;
54        Ok(key)
55    }
56
57    /// Reloads a captured artifact by its storage key.
58    pub async fn load(&self, key: &str) -> StorageResult<Option<ErrorSnapshot>> {
59        Ok(self.kvs.get_bytes(key).await?.map(|entry| ErrorSnapshot {
60            content_type: entry.content_type,
61            bytes: entry.value,
62        }))
63    }
64}
65
66#[cfg(test)]
67mod tests {
68    use super::{ErrorSnapshot, ErrorSnapshotter};
69    use crate::storage::{KeyList, KeyValueStore, KvEntry, ListKeysOptions, StorageResult};
70    use std::{
71        collections::HashMap,
72        sync::{Arc, Mutex},
73    };
74
75    #[derive(Default)]
76    struct MapKvs(Mutex<HashMap<String, KvEntry>>);
77
78    #[async_trait::async_trait]
79    impl KeyValueStore for MapKvs {
80        async fn get_bytes(&self, key: &str) -> StorageResult<Option<KvEntry>> {
81            Ok(self.0.lock().unwrap().get(key).cloned())
82        }
83
84        async fn set_bytes(
85            &self,
86            key: &str,
87            bytes: bytes::Bytes,
88            content_type: &str,
89        ) -> StorageResult<()> {
90            self.0.lock().unwrap().insert(
91                key.into(),
92                KvEntry {
93                    key: key.into(),
94                    value: bytes,
95                    content_type: content_type.into(),
96                },
97            );
98            Ok(())
99        }
100
101        async fn delete(&self, key: &str) -> StorageResult<()> {
102            self.0.lock().unwrap().remove(key);
103            Ok(())
104        }
105
106        async fn list_keys(&self, _: ListKeysOptions) -> StorageResult<KeyList> {
107            Ok(KeyList {
108                keys: Vec::new(),
109                is_truncated: false,
110                next_exclusive_start_key: None,
111            })
112        }
113    }
114
115    fn request(url: &str) -> crate::request::Request {
116        crate::request::Request::get(url).build().unwrap()
117    }
118
119    #[test]
120    fn base_key_is_deterministic_and_request_specific() {
121        let first = request("https://example.com/a");
122        let same = request("https://example.com/a");
123        let different = request("https://example.com/b");
124
125        assert_eq!(
126            ErrorSnapshotter::base_key(&first),
127            ErrorSnapshotter::base_key(&same)
128        );
129        assert_ne!(
130            ErrorSnapshotter::base_key(&first),
131            ErrorSnapshotter::base_key(&different)
132        );
133    }
134
135    #[tokio::test]
136    async fn capture_then_load_round_trips() {
137        let snapshotter = ErrorSnapshotter::new(Arc::new(MapKvs::default()));
138        let req = request("https://example.com/a");
139        let key = snapshotter
140            .capture(
141                &req,
142                "html",
143                bytes::Bytes::from_static(b"<html></html>"),
144                "text/html",
145            )
146            .await
147            .unwrap();
148
149        assert_eq!(
150            snapshotter.load(&key).await.unwrap(),
151            Some(ErrorSnapshot {
152                content_type: "text/html".into(),
153                bytes: bytes::Bytes::from_static(b"<html></html>"),
154            })
155        );
156    }
157
158    #[tokio::test]
159    async fn repeated_capture_same_suffix_overwrites() {
160        let snapshotter = ErrorSnapshotter::new(Arc::new(MapKvs::default()));
161        let req = request("https://example.com/a");
162        let first_key = snapshotter
163            .capture(
164                &req,
165                "body",
166                bytes::Bytes::from_static(b"first"),
167                "text/plain",
168            )
169            .await
170            .unwrap();
171        let second_key = snapshotter
172            .capture(
173                &req,
174                "body",
175                bytes::Bytes::from_static(b"second"),
176                "application/octet-stream",
177            )
178            .await
179            .unwrap();
180
181        assert_eq!(first_key, second_key);
182        assert_eq!(
183            snapshotter.load(&first_key).await.unwrap(),
184            Some(ErrorSnapshot {
185                content_type: "application/octet-stream".into(),
186                bytes: bytes::Bytes::from_static(b"second"),
187            })
188        );
189    }
190
191    #[tokio::test]
192    async fn load_missing_key_returns_none() {
193        let snapshotter = ErrorSnapshotter::new(Arc::new(MapKvs::default()));
194
195        assert_eq!(snapshotter.load("missing").await.unwrap(), None);
196    }
197}
millipede_core/snapshot.rs

millipede_core/
snapshot.rs
