kotoba_storage_rocksdb/
lib.rs

1//! `kotoba-storage-rocksdb`
2//!
3//! RocksDB adapter implementation for kotoba-storage port.
4//! Provides persistent key-value storage using RocksDB.
5
6use std::path::Path;
7use async_trait::async_trait;
8use rocksdb::{DB, Options, IteratorMode};
9use anyhow::Result;
10
11use kotoba_storage::KeyValueStore;
12
13/// RocksDB-based key-value store implementation
14pub struct RocksDbStore {
15    db: DB,
16}
17
18impl RocksDbStore {
19    /// Create a new RocksDB store at the specified path
20    pub fn new<P: AsRef<Path>>(path: P) -> Result<Self> {
21        let mut opts = Options::default();
22        opts.create_if_missing(true);
23
24        let db = DB::open(&opts, path)?;
25        Ok(Self { db })
26    }
27
28    /// Open an existing RocksDB store
29    pub fn open<P: AsRef<Path>>(path: P) -> Result<Self> {
30        let opts = Options::default();
31        let db = DB::open(&opts, path)?;
32        Ok(Self { db })
33    }
34}
35
36#[async_trait]
37impl KeyValueStore for RocksDbStore {
38    async fn put(&self, key: &[u8], value: &[u8]) -> Result<()> {
39        self.db.put(key, value)?;
40        Ok(())
41    }
42
43    async fn get(&self, key: &[u8]) -> Result<Option<Vec<u8>>> {
44        match self.db.get(key)? {
45            Some(value) => Ok(Some(value)),
46            None => Ok(None),
47        }
48    }
49
50    async fn delete(&self, key: &[u8]) -> Result<()> {
51        self.db.delete(key)?;
52        Ok(())
53    }
54
55    async fn scan(&self, prefix: &[u8]) -> Result<Vec<(Vec<u8>, Vec<u8>)>> {
56        let prefix_len = prefix.len();
57        let mut results = Vec::new();
58
59        let iter = self.db.iterator(IteratorMode::Start);
60        for item in iter {
61            let (key, value) = item?;
62
63            // Check if key starts with prefix
64            if key.len() >= prefix_len && &key[..prefix_len] == prefix {
65                results.push((key.to_vec(), value.to_vec()));
66            } else if key.len() >= prefix_len && &key[..prefix_len] > prefix {
67                break; // No more matching keys
68            } else if key.len() < prefix_len && key.as_ref() > prefix {
69                break; // Key is shorter than prefix but lexicographically greater
70            }
71        }
72
73        // Sort results by key for consistent ordering (RocksDB returns sorted results)
74        results.sort_by(|a, b| a.0.cmp(&b.0));
75
76        Ok(results)
77    }
78}
79
80impl Drop for RocksDbStore {
81    fn drop(&mut self) {
82        // RocksDB will be automatically closed when DB is dropped
83    }
84}
85
86#[cfg(test)]
87mod tests {
88    use super::*;
89    use std::sync::Arc;
90    use tempfile::TempDir;
91    use kotoba_storage::KeyValueStore;
92
93    fn create_temp_db() -> (RocksDbStore, TempDir) {
94        let temp_dir = TempDir::new().unwrap();
95        let db_path = temp_dir.path().join("test.db");
96        let store = RocksDbStore::new(&db_path).unwrap();
97        (store, temp_dir)
98    }
99
100    #[tokio::test]
101    async fn test_rocksdb_store_creation() {
102        let (_store, _temp_dir) = create_temp_db();
103        // Test passes if creation succeeds
104    }
105
106    #[tokio::test]
107    async fn test_rocksdb_basic_operations() {
108        let (store, _temp_dir) = create_temp_db();
109
110        // Test put and get
111        store.put(b"key1", b"value1").await.unwrap();
112        let value = store.get(b"key1").await.unwrap();
113        assert_eq!(value, Some(b"value1".to_vec()));
114
115        // Test get non-existent key
116        let value = store.get(b"nonexistent").await.unwrap();
117        assert_eq!(value, None);
118
119        // Test delete
120        store.delete(b"key1").await.unwrap();
121        let value = store.get(b"key1").await.unwrap();
122        assert_eq!(value, None);
123    }
124
125    #[tokio::test]
126    async fn test_rocksdb_multiple_keys() {
127        let (store, _temp_dir) = create_temp_db();
128
129        // Put multiple key-value pairs
130        store.put(b"key1", b"value1").await.unwrap();
131        store.put(b"key2", b"value2").await.unwrap();
132        store.put(b"key3", b"value3").await.unwrap();
133
134        // Verify all keys exist
135        assert_eq!(store.get(b"key1").await.unwrap(), Some(b"value1".to_vec()));
136        assert_eq!(store.get(b"key2").await.unwrap(), Some(b"value2".to_vec()));
137        assert_eq!(store.get(b"key3").await.unwrap(), Some(b"value3".to_vec()));
138    }
139
140    #[tokio::test]
141    async fn test_rocksdb_scan() {
142        let (store, _temp_dir) = create_temp_db();
143
144        // Put keys with common prefix
145        store.put(b"prefix_key1", b"value1").await.unwrap();
146        store.put(b"prefix_key2", b"value2").await.unwrap();
147        store.put(b"prefix_key3", b"value3").await.unwrap();
148        store.put(b"other_key", b"other_value").await.unwrap();
149
150        // Scan with prefix
151        let results = store.scan(b"prefix_").await.unwrap();
152        assert_eq!(results.len(), 3);
153
154        // Verify results are sorted (RocksDB returns sorted results)
155        assert_eq!(results[0], (b"prefix_key1".to_vec(), b"value1".to_vec()));
156        assert_eq!(results[1], (b"prefix_key2".to_vec(), b"value2".to_vec()));
157        assert_eq!(results[2], (b"prefix_key3".to_vec(), b"value3".to_vec()));
158
159        // Scan with empty prefix (should return all)
160        let all_results = store.scan(b"").await.unwrap();
161        assert_eq!(all_results.len(), 4);
162
163        // Scan with non-existent prefix
164        let no_results = store.scan(b"nonexistent").await.unwrap();
165        assert_eq!(no_results.len(), 0);
166    }
167
168    #[tokio::test]
169    async fn test_rocksdb_overwrite() {
170        let (store, _temp_dir) = create_temp_db();
171
172        // Put initial value
173        store.put(b"key", b"initial").await.unwrap();
174        assert_eq!(store.get(b"key").await.unwrap(), Some(b"initial".to_vec()));
175
176        // Overwrite with new value
177        store.put(b"key", b"updated").await.unwrap();
178        assert_eq!(store.get(b"key").await.unwrap(), Some(b"updated".to_vec()));
179    }
180
181    #[tokio::test]
182    async fn test_rocksdb_empty_keys_values() {
183        let (store, _temp_dir) = create_temp_db();
184
185        // Test empty key
186        store.put(b"", b"empty_key_value").await.unwrap();
187        assert_eq!(store.get(b"").await.unwrap(), Some(b"empty_key_value".to_vec()));
188
189        // Test empty value
190        store.put(b"empty_value_key", b"").await.unwrap();
191        assert_eq!(store.get(b"empty_value_key").await.unwrap(), Some(b"".to_vec()));
192    }
193
194    #[tokio::test]
195    async fn test_rocksdb_delete_nonexistent() {
196        let (store, _temp_dir) = create_temp_db();
197
198        // Delete non-existent key should not panic
199        store.delete(b"nonexistent").await.unwrap();
200    }
201
202    #[tokio::test]
203    async fn test_rocksdb_large_data() {
204        let (store, _temp_dir) = create_temp_db();
205
206        // Test with large key and value
207        let large_key = vec![42u8; 1024]; // 1KB key
208        let large_value = vec![255u8; 1024 * 1024]; // 1MB value
209
210        store.put(&large_key, &large_value).await.unwrap();
211        let retrieved = store.get(&large_key).await.unwrap();
212
213        assert_eq!(retrieved, Some(large_value));
214    }
215
216    #[tokio::test]
217    async fn test_rocksdb_unicode_keys_values() {
218        let (store, _temp_dir) = create_temp_db();
219
220        // Test Unicode keys and values
221        let unicode_key = "🚀 котоба 🔥".as_bytes();
222        let unicode_value = "Hello 世界 🌍".as_bytes();
223
224        store.put(unicode_key, unicode_value).await.unwrap();
225        let retrieved = store.get(unicode_key).await.unwrap();
226
227        assert_eq!(retrieved, Some(unicode_value.to_vec()));
228    }
229
230    #[tokio::test]
231    async fn test_rocksdb_persistence() {
232        let temp_dir = TempDir::new().unwrap();
233        let db_path = temp_dir.path().join("persistence_test.db");
234
235        // Create and populate database
236        {
237            let store = RocksDbStore::new(&db_path).unwrap();
238            store.put(b"persistent_key", b"persistent_value").await.unwrap();
239        } // Store goes out of scope and closes
240
241        // Reopen database and check persistence
242        {
243            let store = RocksDbStore::open(&db_path).unwrap();
244            let value = store.get(b"persistent_key").await.unwrap();
245            assert_eq!(value, Some(b"persistent_value".to_vec()));
246        }
247    }
248
249    #[tokio::test]
250    async fn test_rocksdb_concurrent_access() {
251        let temp_dir = TempDir::new().unwrap();
252        let db_path = temp_dir.path().join("concurrent_test.db");
253        let store = Arc::new(RocksDbStore::new(&db_path).unwrap());
254
255        let mut handles = vec![];
256
257        // Spawn multiple tasks to test concurrent access
258        for i in 0..10 {
259            let store_clone = Arc::clone(&store);
260            let handle = tokio::spawn(async move {
261                let key = format!("concurrent_key_{}", i).into_bytes();
262                let value = format!("concurrent_value_{}", i).into_bytes();
263
264                // Put operation
265                store_clone.put(&key, &value).await.unwrap();
266
267                // Get operation
268                let retrieved = store_clone.get(&key).await.unwrap();
269                assert_eq!(retrieved, Some(value));
270            });
271            handles.push(handle);
272        }
273
274        // Wait for all tasks to complete
275        for handle in handles {
276            handle.await.unwrap();
277        }
278
279        // Verify all keys were stored
280        for i in 0..10 {
281            let key = format!("concurrent_key_{}", i).into_bytes();
282            let expected_value = format!("concurrent_value_{}", i).into_bytes();
283            let retrieved = store.get(&key).await.unwrap();
284            assert_eq!(retrieved, Some(expected_value));
285        }
286    }
287
288    #[tokio::test]
289    async fn test_rocksdb_scan_ordering() {
290        let (store, _temp_dir) = create_temp_db();
291
292        // Insert keys in reverse order
293        store.put(b"key3", b"value3").await.unwrap();
294        store.put(b"key1", b"value1").await.unwrap();
295        store.put(b"key2", b"value2").await.unwrap();
296
297        // Scan should return results in sorted order
298        let results = store.scan(b"key").await.unwrap();
299        assert_eq!(results.len(), 3);
300        assert_eq!(results[0], (b"key1".to_vec(), b"value1".to_vec()));
301        assert_eq!(results[1], (b"key2".to_vec(), b"value2".to_vec()));
302        assert_eq!(results[2], (b"key3".to_vec(), b"value3".to_vec()));
303    }
304
305    #[tokio::test]
306    async fn test_rocksdb_scan_prefix_boundaries() {
307        let (store, _temp_dir) = create_temp_db();
308
309        // Insert keys with different prefixes
310        store.put(b"aaa_prefix_key", b"value1").await.unwrap();
311        store.put(b"aaa_other_key", b"value2").await.unwrap();
312        store.put(b"aab_prefix_key", b"value3").await.unwrap();
313        store.put(b"prefix_key", b"value4").await.unwrap();
314
315        // Scan with prefix should stop at boundary
316        let results = store.scan(b"aaa_").await.unwrap();
317        assert_eq!(results.len(), 2);
318
319        // Should include both aaa_ keys but not aab_ or prefix_
320        let keys: Vec<_> = results.iter().map(|(k, _)| k.clone()).collect();
321        assert!(keys.contains(&b"aaa_prefix_key".to_vec()));
322        assert!(keys.contains(&b"aaa_other_key".to_vec()));
323        assert!(!keys.contains(&b"aab_prefix_key".to_vec()));
324        assert!(!keys.contains(&b"prefix_key".to_vec()));
325    }
326
327    #[test]
328    fn test_rocksdb_error_handling() {
329        // Test opening non-existent database without create_if_missing
330        let temp_dir = TempDir::new().unwrap();
331        let db_path = temp_dir.path().join("nonexistent.db");
332        let result = RocksDbStore::open(&db_path);
333        // This should fail as we're trying to open a non-existent database without create_if_missing
334        assert!(result.is_err());
335
336        // Test with valid path should work
337        let temp_dir2 = TempDir::new().unwrap();
338        let db_path2 = temp_dir2.path().join("valid.db");
339        let result2 = RocksDbStore::new(&db_path2);
340        assert!(result2.is_ok());
341    }
342
343    #[tokio::test]
344    async fn test_rocksdb_binary_data() {
345        let (store, _temp_dir) = create_temp_db();
346
347        // Test with binary data (null bytes, etc.)
348        let binary_key = vec![0, 1, 2, 255, 0, 1, 2, 255];
349        let binary_value = vec![255, 254, 253, 0, 1, 2, 3, 255, 254, 253];
350
351        store.put(&binary_key, &binary_value).await.unwrap();
352        let retrieved = store.get(&binary_key).await.unwrap();
353
354        assert_eq!(retrieved, Some(binary_value));
355    }
356
357    #[tokio::test]
358    async fn test_rocksdb_performance() {
359        let (store, _temp_dir) = create_temp_db();
360
361        // Insert many key-value pairs for performance testing
362        let num_entries = 1000;
363        let start_time = std::time::Instant::now();
364
365        for i in 0..num_entries {
366            let key = format!("perf_key_{:04}", i).into_bytes();
367            let value = format!("perf_value_{}", i).into_bytes();
368            store.put(&key, &value).await.unwrap();
369        }
370
371        let insert_time = start_time.elapsed();
372
373        // Read all entries
374        let start_time = std::time::Instant::now();
375        for i in 0..num_entries {
376            let key = format!("perf_key_{:04}", i).into_bytes();
377            let expected_value = format!("perf_value_{}", i).into_bytes();
378            let value = store.get(&key).await.unwrap();
379            assert_eq!(value, Some(expected_value));
380        }
381
382        let read_time = start_time.elapsed();
383
384        println!("RocksDB Performance Test:");
385        println!("  Inserted {} entries in {:?}", num_entries, insert_time);
386        println!("  Read {} entries in {:?}", num_entries, read_time);
387
388        // Basic performance assertions (these will depend on the system)
389        assert!(insert_time.as_millis() < 5000, "Insert performance regression");
390        assert!(read_time.as_millis() < 5000, "Read performance regression");
391    }
392}
kotoba_storage_rocksdb/lib.rs

kotoba_storage_rocksdb/
lib.rs
