hyperloglog/storage/
file.rs

1use crate::{HyperLogLog, Result, HllError};
2use crate::storage::Storage;
3use async_trait::async_trait;
4use std::path::{Path, PathBuf};
5use tokio::fs;
6use tokio::io::{AsyncReadExt, AsyncWriteExt};
7
8/// File-based storage backend for HyperLogLog structures
9#[derive(Debug, Clone)]
10pub struct FileStorage {
11    base_path: PathBuf,
12}
13
14impl FileStorage {
15    /// Create a new FileStorage with the given base directory
16    pub async fn new(base_path: impl AsRef<Path>) -> Result<Self> {
17        let base_path = base_path.as_ref().to_path_buf();
18        fs::create_dir_all(&base_path).await?;
19
20        Ok(Self { base_path })
21    }
22
23    fn key_to_path(&self, key: &str) -> PathBuf {
24        self.base_path.join(format!("{}.hll", key))
25    }
26}
27
28#[async_trait]
29impl Storage for FileStorage {
30    async fn store(&self, key: &str, hll: &HyperLogLog) -> Result<()> {
31        let path = self.key_to_path(key);
32        let serialized = serde_json::to_vec(hll)?;
33
34        let mut file = fs::File::create(&path).await?;
35        file.write_all(&serialized).await?;
36        file.flush().await?;
37
38        Ok(())
39    }
40
41    async fn load(&self, key: &str) -> Result<HyperLogLog> {
42        let path = self.key_to_path(key);
43
44        if !path.exists() {
45            return Err(HllError::NotFound(key.to_string()));
46        }
47
48        let mut file = fs::File::open(&path).await?;
49        let mut contents = Vec::new();
50        file.read_to_end(&mut contents).await?;
51
52        let hll = serde_json::from_slice(&contents)?;
53        Ok(hll)
54    }
55
56    async fn delete(&self, key: &str) -> Result<()> {
57        let path = self.key_to_path(key);
58
59        if path.exists() {
60            fs::remove_file(&path).await?;
61        }
62
63        Ok(())
64    }
65
66    async fn exists(&self, key: &str) -> Result<bool> {
67        let path = self.key_to_path(key);
68        Ok(path.exists())
69    }
70
71    async fn list_keys(&self) -> Result<Vec<String>> {
72        let mut keys = Vec::new();
73        let mut entries = fs::read_dir(&self.base_path).await?;
74
75        while let Some(entry) = entries.next_entry().await? {
76            let path = entry.path();
77            if let Some(ext) = path.extension() {
78                if ext == "hll" {
79                    if let Some(stem) = path.file_stem() {
80                        if let Some(key) = stem.to_str() {
81                            keys.push(key.to_string());
82                        }
83                    }
84                }
85            }
86        }
87
88        Ok(keys)
89    }
90}
91
92#[cfg(test)]
93mod tests {
94    use super::*;
95
96    #[tokio::test]
97    async fn test_file_storage() {
98        let temp_dir = std::env::temp_dir().join("hll_test");
99        let storage = FileStorage::new(&temp_dir).await.unwrap();
100
101        let mut hll = HyperLogLog::new(10).unwrap();
102        hll.add_str("test1");
103        hll.add_str("test2");
104
105        storage.store("test_key", &hll).await.unwrap();
106        assert!(storage.exists("test_key").await.unwrap());
107
108        let loaded = storage.load("test_key").await.unwrap();
109        assert_eq!(loaded.precision(), hll.precision());
110
111        storage.delete("test_key").await.unwrap();
112        assert!(!storage.exists("test_key").await.unwrap());
113
114        let _ = fs::remove_dir_all(&temp_dir).await;
115    }
116}
hyperloglog/storage/file.rs

hyperloglog/storage/
file.rs
