agentos_vfs/local/
file_block_store.rs

1use async_trait::async_trait;
2use std::collections::{HashMap, VecDeque};
3use std::fs;
4use std::path::{Path, PathBuf};
5use std::sync::{Arc, Mutex};
6use vfs::engine::block::BlockStore;
7use vfs::engine::error::{VfsError, VfsResult};
8use vfs::engine::types::BlockKey;
9
10#[derive(Debug, Clone)]
11pub struct FileBlockStore {
12    root: PathBuf,
13    cache: Arc<Mutex<BlockCache>>,
14}
15
16const DEFAULT_BLOCK_CACHE_BYTES: usize = 16 * 1024 * 1024;
17
18#[derive(Debug)]
19struct BlockCache {
20    entries: HashMap<BlockKey, Vec<u8>>,
21    insertion_order: VecDeque<BlockKey>,
22    bytes: usize,
23    max_bytes: usize,
24}
25
26impl BlockCache {
27    fn new(max_bytes: usize) -> Self {
28        Self {
29            entries: HashMap::new(),
30            insertion_order: VecDeque::new(),
31            bytes: 0,
32            max_bytes,
33        }
34    }
35
36    fn get(&self, key: &BlockKey) -> Option<Vec<u8>> {
37        self.entries.get(key).cloned()
38    }
39
40    fn contains(&self, key: &BlockKey) -> bool {
41        self.entries.contains_key(key)
42    }
43
44    fn insert(&mut self, key: BlockKey, data: &[u8]) {
45        if data.len() > self.max_bytes {
46            return;
47        }
48        if let Some(previous) = self.entries.insert(key.clone(), data.to_vec()) {
49            self.bytes = self.bytes.saturating_sub(previous.len());
50        } else {
51            self.insertion_order.push_back(key);
52        }
53        self.bytes = self.bytes.saturating_add(data.len());
54        while self.bytes > self.max_bytes {
55            let Some(oldest) = self.insertion_order.pop_front() else {
56                break;
57            };
58            if let Some(removed) = self.entries.remove(&oldest) {
59                self.bytes = self.bytes.saturating_sub(removed.len());
60            }
61        }
62    }
63
64    fn remove(&mut self, key: &BlockKey) {
65        if let Some(removed) = self.entries.remove(key) {
66            self.bytes = self.bytes.saturating_sub(removed.len());
67        }
68    }
69}
70
71impl FileBlockStore {
72    pub fn new(root: impl Into<PathBuf>) -> VfsResult<Self> {
73        let root = root.into();
74        fs::create_dir_all(&root)
75            .map_err(|err| VfsError::eio(format!("create block root {}: {err}", root.display())))?;
76        Ok(Self {
77            root,
78            cache: Arc::new(Mutex::new(BlockCache::new(DEFAULT_BLOCK_CACHE_BYTES))),
79        })
80    }
81
82    fn path_for(&self, key: &BlockKey) -> PathBuf {
83        block_path(&self.root, key)
84    }
85
86    fn ensure_safe_key(key: &BlockKey) -> VfsResult<()> {
87        if key.0.contains('/') || key.0.contains('\\') || key.0 == "." || key.0 == ".." {
88            return Err(VfsError::einval(format!("unsafe block key: {}", key.0)));
89        }
90        Ok(())
91    }
92
93    pub fn root(&self) -> &Path {
94        &self.root
95    }
96}
97
98#[async_trait]
99impl BlockStore for FileBlockStore {
100    async fn get(&self, key: &BlockKey) -> VfsResult<Vec<u8>> {
101        Self::ensure_safe_key(key)?;
102        if let Some(data) = self
103            .cache
104            .lock()
105            .expect("block cache mutex poisoned")
106            .get(key)
107        {
108            return Ok(data);
109        }
110        let path = self.path_for(key);
111        let data = fs::read(&path).map_err(|err| {
112            if err.kind() == std::io::ErrorKind::NotFound {
113                VfsError::enoent(&key.0)
114            } else {
115                VfsError::eio(format!("read block {}: {err}", path.display()))
116            }
117        })?;
118        self.cache
119            .lock()
120            .expect("block cache mutex poisoned")
121            .insert(key.clone(), &data);
122        Ok(data)
123    }
124
125    async fn get_range(&self, key: &BlockKey, off: u64, len: u64) -> VfsResult<Vec<u8>> {
126        let data = self.get(key).await?;
127        let start = usize::try_from(off)
128            .map_err(|_| VfsError::einval(format!("range offset is too large: {off}")))?;
129        let len = usize::try_from(len)
130            .map_err(|_| VfsError::einval(format!("range length is too large: {len}")))?;
131        if start >= data.len() {
132            return Ok(Vec::new());
133        }
134        let end = start.saturating_add(len).min(data.len());
135        Ok(data[start..end].to_vec())
136    }
137
138    async fn put(&self, key: &BlockKey, data: &[u8]) -> VfsResult<()> {
139        Self::ensure_safe_key(key)?;
140        let path = self.path_for(key);
141        if let Some(parent) = path.parent() {
142            fs::create_dir_all(parent)
143                .map_err(|err| VfsError::eio(format!("create block dir: {err}")))?;
144        }
145        fs::write(&path, data)
146            .map_err(|err| VfsError::eio(format!("write block {}: {err}", path.display())))?;
147        self.cache
148            .lock()
149            .expect("block cache mutex poisoned")
150            .insert(key.clone(), data);
151        Ok(())
152    }
153
154    async fn exists(&self, key: &BlockKey) -> VfsResult<bool> {
155        Self::ensure_safe_key(key)?;
156        if self
157            .cache
158            .lock()
159            .expect("block cache mutex poisoned")
160            .contains(key)
161        {
162            return Ok(true);
163        }
164        Ok(self.path_for(key).exists())
165    }
166
167    async fn delete_many(&self, keys: &[BlockKey]) -> VfsResult<()> {
168        let mut errors = Vec::new();
169        for key in keys {
170            if let Err(error) = Self::ensure_safe_key(key) {
171                errors.push(error.to_string());
172                continue;
173            }
174            self.cache
175                .lock()
176                .expect("block cache mutex poisoned")
177                .remove(key);
178            match fs::remove_file(self.path_for(key)) {
179                Ok(()) => {}
180                Err(err) if err.kind() == std::io::ErrorKind::NotFound => {}
181                Err(err) => errors.push(format!("delete block {}: {err}", key.0)),
182            }
183        }
184        if errors.is_empty() {
185            Ok(())
186        } else {
187            Err(VfsError::eio(format!(
188                "delete {} local blocks failed: {}",
189                errors.len(),
190                errors.join("; ")
191            )))
192        }
193    }
194
195    async fn copy(&self, src: &BlockKey, dst: &BlockKey) -> VfsResult<()> {
196        let data = self.get(src).await?;
197        self.put(dst, &data).await
198    }
199
200    async fn sync(&self) -> VfsResult<()> {
201        for prefix in fs::read_dir(&self.root).map_err(|err| {
202            VfsError::eio(format!("read block root {}: {err}", self.root.display()))
203        })? {
204            let prefix = prefix.map_err(|err| {
205                VfsError::eio(format!(
206                    "read block root entry {}: {err}",
207                    self.root.display()
208                ))
209            })?;
210            let file_type = prefix.file_type().map_err(|err| {
211                VfsError::eio(format!(
212                    "stat block entry {}: {err}",
213                    prefix.path().display()
214                ))
215            })?;
216            if !file_type.is_dir() {
217                continue;
218            }
219            for block in fs::read_dir(prefix.path()).map_err(|err| {
220                VfsError::eio(format!(
221                    "read block directory {}: {err}",
222                    prefix.path().display()
223                ))
224            })? {
225                let block = block
226                    .map_err(|err| VfsError::eio(format!("read block directory entry: {err}")))?;
227                if block
228                    .file_type()
229                    .map_err(|err| {
230                        VfsError::eio(format!("stat block {}: {err}", block.path().display()))
231                    })?
232                    .is_file()
233                {
234                    fs::File::open(block.path())
235                        .and_then(|file| file.sync_all())
236                        .map_err(|err| {
237                            VfsError::eio(format!("sync block {}: {err}", block.path().display()))
238                        })?;
239                }
240            }
241            fs::File::open(prefix.path())
242                .and_then(|directory| directory.sync_all())
243                .map_err(|err| {
244                    VfsError::eio(format!(
245                        "sync block directory {}: {err}",
246                        prefix.path().display()
247                    ))
248                })?;
249        }
250        fs::File::open(&self.root)
251            .and_then(|directory| directory.sync_all())
252            .map_err(|err| VfsError::eio(format!("sync block root {}: {err}", self.root.display())))
253    }
254}
255
256fn block_path(root: &Path, key: &BlockKey) -> PathBuf {
257    let (prefix, suffix) = key.0.split_at(key.0.len().min(2));
258    root.join(prefix).join(suffix)
259}
agentos_vfs/local/file_block_store.rs

agentos_vfs/local/
file_block_store.rs
