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
59impl<T: ExternalBufferSerde> ExternalBuffer<T> for ExternalBufferSled {
60    fn push(&self, item: T) -> Result<(), Error> {
61        let serialized = item.into_external_buffer()?;
62        let key = self.tail_counter.fetch_add(1, Ordering::SeqCst);
63        let key_bytes = Self::key_from_u64(key);
64
65        self.db.insert(&key_bytes, serialized)?;
66        Ok(())
67    }
68
69    fn shift(&self) -> Result<Option<T>, Error> {
70        loop {
71            let current_head = self.head_counter.load(Ordering::SeqCst);
72            let current_tail = self.tail_counter.load(Ordering::SeqCst);
73
74            // Check if buffer is empty
75            if current_head >= current_tail {
76                return Ok(None);
77            }
78
79            let key_bytes = Self::key_from_u64(current_head);
80
81            // Try to remove the item atomically
82            match self.db.remove(&key_bytes)? {
83                Some(data) => {
84                    // Successfully removed, update head counter
85                    self.head_counter.fetch_add(1, Ordering::SeqCst);
86
87                    // Deserialize and return the item
88                    let item = T::from_external_buffer(&data)?;
89                    return Ok(Some(item));
90                }
91                None => {
92                    // Item was already removed by another thread, try next
93                    self.head_counter.fetch_add(1, Ordering::SeqCst);
94                    continue;
95                }
96            }
97        }
98    }
99}
100
101#[cfg(test)]
102mod tests {
103    use super::*;
104    use bincode::{Decode, Encode};
105    use tempfile::TempDir;
106
107    #[derive(Debug, Clone, PartialEq, Encode, Decode)]
108    struct TestItem {
109        id: u32,
110        name: String,
111    }
112
113    #[test]
114    fn test_push_and_shift() {
115        let temp_dir = TempDir::new().unwrap();
116        let buffer = ExternalBufferSled::new(temp_dir.path().join("test_db")).unwrap();
117
118        let item1 = TestItem {
119            id: 1,
120            name: "first".to_string(),
121        };
122        let item2 = TestItem {
123            id: 2,
124            name: "second".to_string(),
125        };
126
127        // Push items
128        buffer.push(item1.clone()).unwrap();
129        buffer.push(item2.clone()).unwrap();
130
131        // Shift items (should come out in FIFO order)
132        let shifted1 = buffer.shift().unwrap();
133        assert_eq!(shifted1, Some(item1));
134
135        let shifted2 = buffer.shift().unwrap();
136        assert_eq!(shifted2, Some(item2));
137
138        // Buffer should be empty now
139        let shifted3: Option<TestItem> = buffer.shift().unwrap();
140        assert_eq!(shifted3, None);
141    }
142
143    #[test]
144    fn test_empty_buffer() {
145        let temp_dir = TempDir::new().unwrap();
146        let buffer = ExternalBufferSled::new(temp_dir.path().join("empty_db")).unwrap();
147
148        // Empty buffer should return None
149        let result: Option<TestItem> = buffer.shift().unwrap();
150        assert_eq!(result, None);
151    }
152
153    #[test]
154    fn test_persistence() {
155        let temp_dir = TempDir::new().unwrap();
156        let db_path = temp_dir.path().join("persistent_db");
157
158        let item = TestItem {
159            id: 42,
160            name: "persistent".to_string(),
161        };
162
163        // Create buffer, push item, and drop it
164        {
165            let buffer = ExternalBufferSled::new(&db_path).unwrap();
166            buffer.push(item.clone()).unwrap();
167        }
168
169        // Create new buffer with same path and verify item is still there
170        {
171            let buffer = ExternalBufferSled::new(&db_path).unwrap();
172            let retrieved = buffer.shift().unwrap();
173            assert_eq!(retrieved, Some(item));
174        }
175    }
176
177    #[test]
178    fn test_multiple_pushes_and_shifts() {
179        let temp_dir = TempDir::new().unwrap();
180        let buffer = ExternalBufferSled::new(temp_dir.path().join("multi_db")).unwrap();
181
182        let items: Vec<TestItem> = (0..10)
183            .map(|i| TestItem {
184                id: i,
185                name: format!("item_{}", i),
186            })
187            .collect();
188
189        // Push all items
190        for item in &items {
191            buffer.push(item.clone()).unwrap();
192        }
193
194        // Shift all items and verify order
195        for expected_item in &items {
196            let shifted = buffer.shift().unwrap();
197            assert_eq!(shifted, Some(expected_item.clone()));
198        }
199
200        // Buffer should be empty
201        let result: Option<TestItem> = buffer.shift().unwrap();
202        assert_eq!(result, None);
203    }
204
205    #[test]
206    fn test_interleaved_push_and_shift() {
207        let temp_dir = TempDir::new().unwrap();
208        let buffer = ExternalBufferSled::new(temp_dir.path().join("interleaved_db")).unwrap();
209
210        let item1 = TestItem {
211            id: 1,
212            name: "first".to_string(),
213        };
214        let item2 = TestItem {
215            id: 2,
216            name: "second".to_string(),
217        };
218        let item3 = TestItem {
219            id: 3,
220            name: "third".to_string(),
221        };
222
223        // Push one, shift one
224        buffer.push(item1.clone()).unwrap();
225        let shifted1 = buffer.shift().unwrap();
226        assert_eq!(shifted1, Some(item1));
227
228        // Push two, shift two
229        buffer.push(item2.clone()).unwrap();
230        buffer.push(item3.clone()).unwrap();
231
232        let shifted2 = buffer.shift().unwrap();
233        assert_eq!(shifted2, Some(item2));
234
235        let shifted3 = buffer.shift().unwrap();
236        assert_eq!(shifted3, Some(item3));
237
238        // Should be empty
239        let result: Option<TestItem> = buffer.shift().unwrap();
240        assert_eq!(result, None);
241    }
242}
external_buffered_stream/buffer/sled.rs

external_buffered_stream/buffer/
sled.rs
