external_buffered_stream/buffer/
sled.rs

1use std::sync::atomic::{AtomicU64, Ordering};
2
3use crate::{Error, ExternalBufferSerde};
4
5use super::ExternalBuffer;
6
7/// Sled as the persistent buffer with FIFO queue order
8pub struct ExternalBufferSled {
9    db: sled::Db,
10    head_counter: AtomicU64,
11    tail_counter: AtomicU64,
12}
13
14impl ExternalBufferSled {
15    pub fn new<P: AsRef<std::path::Path>>(path: P) -> Result<Self, Error> {
16        let db = sled::open(path)?;
17
18        // Initialize counters by scanning existing keys
19        let (head, tail) = Self::initialize_counters(&db)?;
20
21        Ok(Self {
22            db,
23            head_counter: AtomicU64::new(head),
24            tail_counter: AtomicU64::new(tail),
25        })
26    }
27
28    fn initialize_counters(db: &sled::Db) -> Result<(u64, u64), Error> {
29        let mut min_key = u64::MAX;
30        let mut max_key = 0u64;
31        let mut has_keys = false;
32
33        for result in db.iter() {
34            let (key, _) = result?;
35            if key.len() == 8 {
36                let key_u64 = u64::from_be_bytes(
37                    key.as_ref()
38                        .try_into()
39                        .map_err(|_| Error::InvalidSledKeyFormat)?,
40                );
41                min_key = min_key.min(key_u64);
42                max_key = max_key.max(key_u64);
43                has_keys = true;
44            }
45        }
46
47        if has_keys {
48            Ok((min_key, max_key + 1))
49        } else {
50            Ok((0, 0))
51        }
52    }
53
54    fn key_from_u64(value: u64) -> [u8; 8] {
55        value.to_be_bytes()
56    }
57}
58
59#[async_trait::async_trait]
60impl<T: ExternalBufferSerde + Send + 'static> ExternalBuffer<T> for ExternalBufferSled {
61    async fn push(&self, item: T) -> Result<(), Error> {
62        let serialized = item.into_external_buffer()?;
63        let key = self.tail_counter.fetch_add(1, Ordering::SeqCst);
64        let key_bytes = Self::key_from_u64(key);
65
66        self.db.insert(&key_bytes, serialized)?;
67        Ok(())
68    }
69
70    async fn shift(&self) -> Result<Option<T>, Error> {
71        loop {
72            let current_head = self.head_counter.load(Ordering::SeqCst);
73            let current_tail = self.tail_counter.load(Ordering::SeqCst);
74
75            // Check if buffer is empty
76            if current_head >= current_tail {
77                return Ok(None);
78            }
79
80            let key_bytes = Self::key_from_u64(current_head);
81
82            // Try to remove the item atomically
83            match self.db.remove(&key_bytes)? {
84                Some(data) => {
85                    // Successfully removed, update head counter
86                    self.head_counter.fetch_add(1, Ordering::SeqCst);
87
88                    // Deserialize and return the item
89                    let item = T::from_external_buffer(&data)?;
90                    return Ok(Some(item));
91                }
92                None => {
93                    // Item was already removed by another thread, try next
94                    self.head_counter.fetch_add(1, Ordering::SeqCst);
95                    continue;
96                }
97            }
98        }
99    }
100}
101
102#[cfg(test)]
103mod tests {
104    use super::*;
105    use bincode::{Decode, Encode};
106    use tempfile::TempDir;
107
108    #[derive(Debug, Clone, PartialEq, Encode, Decode)]
109    struct TestItem {
110        id: u32,
111        name: String,
112    }
113
114    #[tokio::test]
115    async fn test_push_and_shift() {
116        let temp_dir = TempDir::new().unwrap();
117        let buffer = ExternalBufferSled::new(temp_dir.path().join("test_db")).unwrap();
118
119        let item1 = TestItem {
120            id: 1,
121            name: "first".to_string(),
122        };
123        let item2 = TestItem {
124            id: 2,
125            name: "second".to_string(),
126        };
127
128        // Push items
129        buffer.push(item1.clone()).await.unwrap();
130        buffer.push(item2.clone()).await.unwrap();
131
132        // Shift items (should come out in FIFO order)
133        let shifted1 = buffer.shift().await.unwrap();
134        assert_eq!(shifted1, Some(item1));
135
136        let shifted2 = buffer.shift().await.unwrap();
137        assert_eq!(shifted2, Some(item2));
138
139        // Buffer should be empty now
140        let shifted3: Option<TestItem> = buffer.shift().await.unwrap();
141        assert_eq!(shifted3, None);
142    }
143
144    #[tokio::test]
145    async fn test_empty_buffer() {
146        let temp_dir = TempDir::new().unwrap();
147        let buffer = ExternalBufferSled::new(temp_dir.path().join("empty_db")).unwrap();
148
149        // Empty buffer should return None
150        let result: Option<TestItem> = buffer.shift().await.unwrap();
151        assert_eq!(result, None);
152    }
153
154    #[tokio::test]
155    async fn test_persistence() {
156        let temp_dir = TempDir::new().unwrap();
157        let db_path = temp_dir.path().join("persistent_db");
158
159        let item = TestItem {
160            id: 42,
161            name: "persistent".to_string(),
162        };
163
164        // Create buffer, push item, and drop it
165        {
166            let buffer = ExternalBufferSled::new(&db_path).unwrap();
167            buffer.push(item.clone()).await.unwrap();
168        }
169
170        // Create new buffer with same path and verify item is still there
171        {
172            let buffer = ExternalBufferSled::new(&db_path).unwrap();
173            let retrieved = buffer.shift().await.unwrap();
174            assert_eq!(retrieved, Some(item));
175        }
176    }
177
178    #[tokio::test]
179    async fn test_multiple_pushes_and_shifts() {
180        let temp_dir = TempDir::new().unwrap();
181        let buffer = ExternalBufferSled::new(temp_dir.path().join("multi_db")).unwrap();
182
183        let items: Vec<TestItem> = (0..10)
184            .map(|i| TestItem {
185                id: i,
186                name: format!("item_{}", i),
187            })
188            .collect();
189
190        // Push all items
191        for item in &items {
192            buffer.push(item.clone()).await.unwrap();
193        }
194
195        // Shift all items and verify order
196        for expected_item in &items {
197            let shifted = buffer.shift().await.unwrap();
198            assert_eq!(shifted, Some(expected_item.clone()));
199        }
200
201        // Buffer should be empty
202        let result: Option<TestItem> = buffer.shift().await.unwrap();
203        assert_eq!(result, None);
204    }
205
206    #[tokio::test]
207    async fn test_interleaved_push_and_shift() {
208        let temp_dir = TempDir::new().unwrap();
209        let buffer = ExternalBufferSled::new(temp_dir.path().join("interleaved_db")).unwrap();
210
211        let item1 = TestItem {
212            id: 1,
213            name: "first".to_string(),
214        };
215        let item2 = TestItem {
216            id: 2,
217            name: "second".to_string(),
218        };
219        let item3 = TestItem {
220            id: 3,
221            name: "third".to_string(),
222        };
223
224        // Push one, shift one
225        buffer.push(item1.clone()).await.unwrap();
226        let shifted1 = buffer.shift().await.unwrap();
227        assert_eq!(shifted1, Some(item1));
228
229        // Push two, shift two
230        buffer.push(item2.clone()).await.unwrap();
231        buffer.push(item3.clone()).await.unwrap();
232
233        let shifted2 = buffer.shift().await.unwrap();
234        assert_eq!(shifted2, Some(item2));
235
236        let shifted3 = buffer.shift().await.unwrap();
237        assert_eq!(shifted3, Some(item3));
238
239        // Should be empty
240        let result: Option<TestItem> = buffer.shift().await.unwrap();
241        assert_eq!(result, None);
242    }
243}
external_buffered_stream/buffer/sled.rs

external_buffered_stream/buffer/
sled.rs
