Skip to main content

firecloud_storage/
store.rs

1//! Local chunk storage using Sled (lightweight embedded DB)
2
3use crate::{StorageError, StorageResult};
4use firecloud_core::{Chunk, ChunkHash, ChunkMetadata};
5use sled::Db;
6use std::path::Path;
7use tracing::{debug, info};
8
9/// Tree names for different data types
10const TREE_CHUNKS: &str = "chunks";
11const TREE_METADATA: &str = "metadata";
12
13/// Local chunk store backed by Sled
14pub struct ChunkStore {
15    db: Db,
16}
17
18impl ChunkStore {
19    /// Open or create a chunk store at the given path
20    pub fn open<P: AsRef<Path>>(path: P) -> StorageResult<Self> {
21        let db = sled::open(path)
22            .map_err(|e| StorageError::Database(e.to_string()))?;
23        info!("Opened chunk store with Sled");
24        Ok(Self { db })
25    }
26
27    /// Store a chunk
28    pub fn put(&self, chunk: &Chunk) -> StorageResult<()> {
29        let chunks_tree = self.db.open_tree(TREE_CHUNKS)
30            .map_err(|e| StorageError::Database(e.to_string()))?;
31        let metadata_tree = self.db.open_tree(TREE_METADATA)
32            .map_err(|e| StorageError::Database(e.to_string()))?;
33
34        let key = chunk.metadata.hash.0;
35
36        // Store chunk data
37        chunks_tree.insert(key, chunk.data.as_ref())
38            .map_err(|e| StorageError::Database(e.to_string()))?;
39
40        // Store metadata
41        let metadata_bytes = bincode::serialize(&chunk.metadata)
42            .map_err(|e| StorageError::Serialization(e.to_string()))?;
43        metadata_tree.insert(key, metadata_bytes)
44            .map_err(|e| StorageError::Database(e.to_string()))?;
45
46        debug!("Stored chunk: {}", chunk.metadata.hash);
47        Ok(())
48    }
49
50    /// Retrieve a chunk by hash
51    pub fn get(&self, hash: &ChunkHash) -> StorageResult<Option<Chunk>> {
52        let chunks_tree = self.db.open_tree(TREE_CHUNKS)
53            .map_err(|e| StorageError::Database(e.to_string()))?;
54        let metadata_tree = self.db.open_tree(TREE_METADATA)
55            .map_err(|e| StorageError::Database(e.to_string()))?;
56
57        let key = hash.0;
58
59        let data = match chunks_tree.get(key)
60            .map_err(|e| StorageError::Database(e.to_string()))? {
61            Some(d) => d.to_vec(),
62            None => return Ok(None),
63        };
64
65        let metadata_bytes = metadata_tree.get(key)
66            .map_err(|e| StorageError::Database(e.to_string()))?
67            .ok_or_else(|| StorageError::ChunkNotFound(hash.to_hex()))?;
68
69        let metadata: ChunkMetadata = bincode::deserialize(&metadata_bytes)
70            .map_err(|e| StorageError::Serialization(e.to_string()))?;
71
72        Ok(Some(Chunk { data: data.into(), metadata }))
73    }
74
75    /// Check if a chunk exists
76    pub fn contains(&self, hash: &ChunkHash) -> StorageResult<bool> {
77        let chunks_tree = self.db.open_tree(TREE_CHUNKS)
78            .map_err(|e| StorageError::Database(e.to_string()))?;
79        Ok(chunks_tree.contains_key(hash.0)
80            .map_err(|e| StorageError::Database(e.to_string()))?)
81    }
82
83    /// Delete a chunk
84    pub fn delete(&self, hash: &ChunkHash) -> StorageResult<bool> {
85        let chunks_tree = self.db.open_tree(TREE_CHUNKS)
86            .map_err(|e| StorageError::Database(e.to_string()))?;
87        let metadata_tree = self.db.open_tree(TREE_METADATA)
88            .map_err(|e| StorageError::Database(e.to_string()))?;
89
90        let key = hash.0;
91        let existed = chunks_tree.remove(key)
92            .map_err(|e| StorageError::Database(e.to_string()))?
93            .is_some();
94        metadata_tree.remove(key)
95            .map_err(|e| StorageError::Database(e.to_string()))?;
96
97        if existed { debug!("Deleted chunk: {}", hash); }
98        Ok(existed)
99    }
100
101    /// Get total number of chunks stored
102    pub fn count(&self) -> StorageResult<usize> {
103        let chunks_tree = self.db.open_tree(TREE_CHUNKS)
104            .map_err(|e| StorageError::Database(e.to_string()))?;
105        Ok(chunks_tree.len())
106    }
107
108    /// Get total size of all stored chunks in bytes
109    pub fn total_size(&self) -> StorageResult<u64> {
110        let chunks_tree = self.db.open_tree(TREE_CHUNKS)
111            .map_err(|e| StorageError::Database(e.to_string()))?;
112        
113        let mut total = 0u64;
114        for result in chunks_tree.iter() {
115            let (_, value) = result.map_err(|e| StorageError::Database(e.to_string()))?;
116            total += value.len() as u64;
117        }
118        Ok(total)
119    }
120
121    /// Flush all pending writes to disk
122    pub fn flush(&self) -> StorageResult<()> {
123        self.db.flush().map_err(|e| StorageError::Database(e.to_string()))?;
124        Ok(())
125    }
126}