Skip to main content

root/
mount.rs

1use dynamic::{MsgPack, MsgUnpack};
2use rand::random_range;
3use scc::HashMap;
4use smol_str::SmolStr;
5
6use anyhow::{Result, anyhow};
7
8use super::sync_await;
9use crate::directory;
10use crate::node::Node;
11use fjall::{KeyspaceCreateOptions, OptimisticTxDatabase, OptimisticTxKeyspace};
12use redis::AsyncCommands;
13use redis::Commands;
14use std::sync::Arc;
15
16use rslock::LockManager;
17
18pub enum Mount<T> {
19    Memory(Arc<HashMap<SmolStr, Node<T>>>),
20    Redis { client: redis::Client, rl: LockManager },
21    Fjall { values: OptimisticTxKeyspace, write_lock: Arc<std::sync::Mutex<()>> },
22}
23
24impl<T: std::fmt::Debug + MsgPack + MsgUnpack + Default + Send> Mount<T> {
25    pub fn memory() -> Self {
26        Self::Memory(Arc::new(HashMap::new()))
27    }
28
29    pub fn redis(url: &str) -> Result<Self> {
30        let client = redis::Client::open(url)?;
31        let mut conn = client.get_connection()?;
32        directory::rebuild_once(&mut conn)?;
33        let rl = LockManager::new(vec![url]);
34        Ok(Self::Redis { client, rl })
35    }
36
37    pub fn fjall(data_dir: &str) -> Result<Self> {
38        let db = OptimisticTxDatabase::builder(data_dir).open()?;
39        let values = db.keyspace("root", KeyspaceCreateOptions::default)?;
40        Ok(Self::Fjall { values, write_lock: Arc::new(std::sync::Mutex::new(())) })
41    }
42
43    pub fn add(&self, name: &str, value: T) -> bool {
44        match self {
45            Self::Memory(m) => {
46                m.upsert_sync(name.into(), Node::Object(value));
47                true
48            }
49            Self::Redis { client, rl: _ } => {
50                let mut buf = Vec::new();
51                value.encode(&mut buf);
52                let Ok(mut conn) = client.get_connection() else {
53                    return false;
54                };
55                conn.set::<&str, Vec<u8>, ()>(name, buf).is_ok() && directory::add_path(&mut conn, name).is_ok()
56            }
57            Self::Fjall { values, write_lock } => {
58                let mut buf = Vec::new();
59                value.encode(&mut buf);
60                let Ok(_guard) = write_lock.lock() else {
61                    return false;
62                };
63                fjall_clear_node(values, name).is_ok() && values.insert(fjall_object_key(name), buf).is_ok()
64            }
65        }
66    }
67
68    pub fn contains(&self, name: &str) -> bool {
69        match self {
70            Self::Memory(m) => m.contains_sync(name),
71            Self::Redis { client, rl: _ } => client.get_connection().and_then(|mut conn| conn.exists::<&str, bool>(name)).unwrap_or(false),
72            Self::Fjall { values, .. } => values.contains_key(fjall_object_key(name)).unwrap_or(false) || values.contains_key(fjall_type_key(name)).unwrap_or(false),
73        }
74    }
75
76    pub fn get_mut<R: Send + 'static, F: FnMut(&mut T) -> R>(&self, name: &str, mut f: F) -> Result<R>
77    where
78        F: Send + 'static,
79    {
80        match self {
81            Self::Memory(m) => m
82                .update_sync(name, |_, v| match v {
83                    Node::Object(v) => Some(f(v)),
84                    _ => None,
85                })
86                .flatten()
87                .ok_or(anyhow!("{} 不存在", name)),
88            Self::Redis { client, rl } => {
89                let name = String::from(name);
90                let rl = rl.clone();
91                let client = client.clone();
92                sync_await!(async move {
93                    loop {
94                        let time_out = random_range(0..1000); //等待随机时间
95                        if let Ok(lock) = rl.lock(name.as_str(), std::time::Duration::from_millis(time_out)).await {
96                            let mut conn = client.get_multiplexed_async_connection().await?;
97                            let mut buf: Vec<u8> = conn.get(name.as_str()).await?;
98                            let (mut v, _) = T::decode(buf.as_slice())?;
99                            let r = f(&mut v);
100                            buf.clear();
101                            v.encode(&mut buf);
102                            conn.set::<&str, Vec<u8>, ()>(name.as_str(), buf).await?;
103                            rl.unlock(&lock).await;
104                            break Ok(r);
105                        }
106                    }
107                })
108            }
109            Self::Fjall { values, write_lock } => {
110                let _guard = write_lock.lock().map_err(|e| anyhow!("无法获取 fjall 写锁: {}", e))?;
111                let mut v = fjall_get_object::<T>(values, name)?;
112                let r = f(&mut v);
113                fjall_insert_object(values, name, &v)?;
114                Ok(r)
115            }
116        }
117    }
118
119    pub fn get<R, F: FnOnce(&T) -> R>(&self, name: &str, f: F) -> Result<R> {
120        match self {
121            Self::Memory(m) => m
122                .read_sync(name, |_, v| match v {
123                    Node::Object(v) => Some(f(v)),
124                    _ => None,
125                })
126                .flatten()
127                .ok_or(anyhow!("{} 不存在", name)),
128            Self::Redis { client, rl: _ } => {
129                let mut conn = client.get_connection()?;
130                let buf: Vec<u8> = conn.get(name)?;
131                let (v, _) = T::decode(buf.as_slice())?;
132                Ok(f(&v))
133            }
134            Self::Fjall { values, .. } => fjall_get_object(values, name).map(|v| f(&v)),
135        }
136    }
137
138    pub fn get_key_mut<'a, R: Send + 'static, F: FnOnce(&mut T) -> R>(&'a self, name: &'a str, key: &'a str, f: F) -> Result<R>
139    where
140        F: Send + 'static,
141    {
142        match self {
143            Self::Memory(m) => m
144                .update_sync(name, |_, v| match v {
145                    Node::Map(m) => m.update_sync(key, |_, v| f(v)),
146                    _ => None,
147                })
148                .flatten()
149                .ok_or(anyhow!("{} 不存在", name)),
150            Self::Redis { client, rl } => {
151                let name = String::from(name);
152                let key = String::from(key);
153                let rl = rl.clone();
154                let client = client.clone();
155                sync_await!(async move {
156                    loop {
157                        let time_out = random_range(0..1000); //等待随机时间
158                        let lock_name = format!("{}::{}", name, key); //为这个 name 里面的 key 单独上锁
159                        if let Ok(lock) = rl.lock(lock_name.as_str(), std::time::Duration::from_millis(time_out)).await {
160                            let mut conn = client.get_multiplexed_async_connection().await?;
161                            let mut buf: Vec<u8> = conn.hget(name.as_str(), key.as_str()).await?;
162                            let (mut v, _) = T::decode(buf.as_slice())?;
163                            let r = f(&mut v);
164                            buf.clear();
165                            v.encode(&mut buf);
166                            conn.hset::<&str, &str, Vec<u8>, ()>(name.as_str(), key.as_str(), buf).await?;
167                            rl.unlock(&lock).await;
168                            break Ok(r);
169                        }
170                    }
171                })
172            }
173            Self::Fjall { values, write_lock } => {
174                let _guard = write_lock.lock().map_err(|e| anyhow!("无法获取 fjall 写锁: {}", e))?;
175                let mut v = fjall_get_map_item::<T>(values, name, key)?;
176                let r = f(&mut v);
177                fjall_insert_map_item(values, name, key, &v)?;
178                Ok(r)
179            }
180        }
181    }
182
183    pub fn dir(&self, name: &str) -> Result<Vec<SmolStr>> {
184        let prefix = if name.is_empty() || name.ends_with('/') { name.to_string() } else { format!("{name}/") };
185        if let Self::Redis { client, rl: _ } = self {
186            let mut conn = client.get_connection()?;
187            return Ok(Self::dir_entries_from_children(&prefix, directory::children(&mut conn, name)?));
188        }
189
190        let raw = self.dir_raw(&prefix)?;
191        Ok(Self::dir_entries_from_raw(&prefix, raw))
192    }
193
194    fn dir_entries_from_children(prefix: &str, children: Vec<SmolStr>) -> Vec<SmolStr> {
195        children.into_iter().map(|child| if prefix.is_empty() { child } else { format!("{prefix}{child}").into() }).collect()
196    }
197
198    fn dir_entries_from_raw(prefix: &str, raw: Vec<SmolStr>) -> Vec<SmolStr> {
199        let mut seen = std::collections::HashSet::new();
200        let mut names = Vec::new();
201        for key in raw {
202            let Some(rest) = key.strip_prefix(&prefix) else {
203                continue;
204            };
205            let end = rest.find('/').unwrap_or(rest.len());
206            let entry = &rest[..end];
207            if !entry.is_empty() && seen.insert(entry.to_string()) {
208                names.push(format!("{prefix}{entry}").into());
209            }
210        }
211        names
212    }
213
214    pub fn dir_raw(&self, name: &str) -> Result<Vec<SmolStr>> {
215        let mut names = Vec::new();
216        match self {
217            Self::Memory(m) => {
218                m.iter_sync(|key, _| {
219                    if key.starts_with(name) {
220                        names.push(key.clone())
221                    }
222                    true
223                });
224            }
225            Self::Redis { client, rl: _ } => {
226                let mut conn = client.get_connection()?;
227                let prefix = if name.is_empty() || name.ends_with('/') { name.to_string() } else { format!("{name}/") };
228                names.append(&mut Self::dir_entries_from_children(&prefix, directory::children(&mut conn, name)?));
229            }
230            Self::Fjall { values, .. } => {
231                names.append(&mut fjall_paths_with_prefix(values, name)?);
232            }
233        }
234        Ok(names)
235    }
236
237    pub fn len(&self, name: &str) -> Result<usize> {
238        match self {
239            Self::Memory(m) => m.read_sync(name, |_, v| v.len()).ok_or(anyhow!("{} 不是列表", name)),
240            Self::Redis { client, rl: _ } => {
241                let mut conn = client.get_connection()?;
242                let ty: String = conn.key_type(name)?;
243                match ty.as_str() {
244                    "list" => Ok(conn.llen(name)?),
245                    "hash" => Ok(conn.hlen(name)?),
246                    _ => Ok(1),
247                }
248            }
249            Self::Fjall { values, .. } => match fjall_node_type(values, name)? {
250                Some(FjallNodeType::List) => Ok(fjall_count_prefix(values, fjall_list_prefix(name))?),
251                Some(FjallNodeType::Map) => Ok(fjall_count_prefix(values, fjall_map_prefix(name))?),
252                None if values.contains_key(fjall_object_key(name))? => Ok(1),
253                None => Err(anyhow!("{} 不存在", name)),
254            },
255        }
256    }
257
258    pub fn remove(&self, name: &str) -> Result<T> {
259        match self {
260            Self::Memory(m) => m.remove_sync(name).and_then(|(_, v)| v.into_object()).ok_or(anyhow!("{} 不存在", name)),
261            Self::Redis { client, rl: _ } => {
262                let mut conn = client.get_connection()?;
263                let buf: Vec<u8> = conn.get(name)?;
264                let (v, _) = T::decode(buf.as_slice())?;
265                let removed: usize = conn.del(name)?;
266                if removed != 0 {
267                    directory::remove_path(&mut conn, name)?;
268                }
269                Ok(v)
270            }
271            Self::Fjall { values, .. } => {
272                let v = fjall_get_object(values, name)?;
273                values.remove(fjall_object_key(name))?;
274                Ok(v)
275            }
276        }
277    }
278
279    pub fn add_list(&self, name: &str) {
280        match self {
281            Self::Memory(m) => {
282                m.upsert_sync(name.into(), Node::<T>::list());
283            } // 强制插入 肯定成功
284            Self::Redis { client: _, rl: _ } => {}
285            Self::Fjall { values, write_lock } => {
286                if let Ok(_guard) = write_lock.lock() {
287                    let _ = fjall_clear_node(values, name).and_then(|_| fjall_set_node_type(values, name, FjallNodeType::List));
288                }
289            }
290        }
291    }
292
293    pub fn add_map(&self, name: &str) {
294        match self {
295            Self::Memory(m) => {
296                m.upsert_sync(name.into(), Node::<T>::map());
297            } // 强制插入 肯定成功
298            Self::Redis { client: _, rl: _ } => {}
299            Self::Fjall { values, write_lock } => {
300                if let Ok(_guard) = write_lock.lock() {
301                    let _ = fjall_clear_node(values, name).and_then(|_| fjall_set_node_type(values, name, FjallNodeType::Map));
302                }
303            }
304        }
305    }
306
307    pub fn push(&self, name: &str, value: T) -> Result<usize> {
308        match self {
309            Self::Memory(m) => m.update_sync(name, |_, v| v.push(value)).flatten().ok_or(anyhow!("push {} 失败", name)),
310            Self::Redis { client, rl: _ } => {
311                let mut conn = client.get_connection()?;
312                let mut buf = Vec::new();
313                value.encode(&mut buf);
314                let len = conn.rpush(name, buf)?;
315                directory::add_path(&mut conn, name)?;
316                Ok(len)
317            }
318            Self::Fjall { values, write_lock } => {
319                let _guard = write_lock.lock().map_err(|e| anyhow!("无法获取 fjall 写锁: {}", e))?;
320                match fjall_node_type(values, name)? {
321                    Some(FjallNodeType::List) => {}
322                    Some(FjallNodeType::Map) => return Err(anyhow!("push {} 失败", name)),
323                    None => fjall_set_node_type(values, name, FjallNodeType::List)?,
324                }
325                let idx = fjall_next_list_idx(values, name)?;
326                fjall_insert_list_item(values, name, idx, &value)?;
327                Ok(idx)
328            }
329        }
330    }
331
332    pub fn get_idx<R, F: FnOnce(&T) -> R>(&self, name: &str, idx: usize, f: F) -> Result<R> {
333        match self {
334            Self::Memory(m) => m.read_sync(name, |_, v| v.get_idx(idx, f)).flatten().ok_or(anyhow!("get_idx {} 失败", name)),
335            Self::Redis { client, rl: _ } => {
336                let mut conn = client.get_connection()?;
337                let buf: Vec<u8> = conn.lindex(name, idx as isize)?;
338                let (v, _) = T::decode(buf.as_slice())?;
339                Ok(f(&v))
340            }
341            Self::Fjall { values, .. } => fjall_get_list_item(values, name, idx).map(|v| f(&v)),
342        }
343    }
344
345    pub fn get_idx_mut<R, F: FnMut(&mut T) -> R>(&self, name: &str, idx: usize, mut f: F) -> Result<R> {
346        match self {
347            Self::Memory(m) => m.update_sync(name, |_, v| v.get_idx_mut(idx, f)).flatten().ok_or(anyhow!("get_idx {} 失败", name)),
348            Self::Redis { client, rl: _ } => {
349                let mut conn = client.get_connection()?;
350                let buf: Vec<u8> = conn.lindex(name, idx as isize)?;
351                let (mut v, _) = T::decode(buf.as_slice())?;
352                Ok(f(&mut v))
353            }
354            Self::Fjall { values, write_lock } => {
355                let _guard = write_lock.lock().map_err(|e| anyhow!("无法获取 fjall 写锁: {}", e))?;
356                let mut v = fjall_get_list_item::<T>(values, name, idx)?;
357                let r = f(&mut v);
358                fjall_insert_list_item(values, name, idx, &v)?;
359                Ok(r)
360            }
361        }
362    }
363
364    pub fn remove_idx(&self, name: &str, idx: usize) -> Result<T> {
365        match self {
366            Self::Memory(m) => m.update_sync(name, |_, v| v.remove_idx(idx)).flatten().ok_or(anyhow!("remove_idx {} 失败", name)),
367            Self::Redis { client, rl: _ } => {
368                let mut conn = client.get_connection()?;
369                let buf: Vec<u8> = conn.lindex(name, idx as isize)?;
370                let v = T::decode(buf.as_slice()).map(|(v, _)| v).unwrap_or(T::default());
371                let _: () = conn.lset(name, idx as isize, Vec::new())?;
372                Ok(v)
373            }
374            Self::Fjall { values, .. } => {
375                let key = fjall_list_item_key(name, idx);
376                let buf = values.get(&key)?.ok_or(anyhow!("remove_idx {} 失败", name))?;
377                let (v, _) = T::decode(buf.as_ref())?;
378                values.remove(key)?;
379                Ok(v)
380            }
381        }
382    }
383
384    pub fn insert(&self, name: &str, key: &str, value: T) -> Option<T> {
385        match self {
386            Self::Memory(m) => m.update_sync(name, |_, v| v.insert(key.into(), value)).flatten(),
387            Self::Redis { client, rl: _ } => {
388                if let Ok(mut conn) = client.get_connection() {
389                    let mut buf = Vec::new();
390                    value.encode(&mut buf);
391                    if conn.hset::<&str, &str, Vec<u8>, ()>(name, key, buf).is_ok() {
392                        let _ = directory::add_path(&mut conn, name);
393                    }
394                }
395                None
396            }
397            Self::Fjall { values, write_lock } => {
398                let _guard = write_lock.lock().ok()?;
399                match fjall_node_type(values, name).ok()? {
400                    Some(FjallNodeType::Map) => {}
401                    Some(FjallNodeType::List) => return None,
402                    None => fjall_set_node_type(values, name, FjallNodeType::Map).ok()?,
403                }
404                let old = fjall_get_map_item(values, name, key).ok();
405                fjall_insert_map_item(values, name, key, &value).ok()?;
406                old
407            }
408        }
409    }
410
411    pub fn get_key<R, F: FnOnce(&T) -> R>(&self, name: &str, key: &str, f: F) -> Result<R> {
412        match self {
413            Self::Memory(m) => m.read_sync(name, |_, v| v.get_key(key, f)).flatten().ok_or(anyhow!("get_key {} 失败", name)),
414            Self::Redis { client, rl: _ } => {
415                let mut conn = client.get_connection()?;
416                let buf: Vec<u8> = conn.hget(name, key)?;
417                let (v, _) = T::decode(buf.as_slice())?;
418                Ok(f(&v))
419            }
420            Self::Fjall { values, .. } => fjall_get_map_item(values, name, key).map(|v| f(&v)),
421        }
422    }
423
424    pub fn keys(&self, name: &str) -> Result<Vec<SmolStr>> {
425        match self {
426            Self::Memory(m) => m.read_sync(name, |_, v| v.keys()).flatten().ok_or(anyhow!("get_key {} 失败", name)),
427            Self::Redis { client, rl: _ } => {
428                let mut conn = client.get_connection()?;
429                let keys: Vec<String> = conn.hkeys(name).unwrap_or(Vec::new());
430                Ok(keys.into_iter().map(|k| k.into()).collect())
431            }
432            Self::Fjall { values, .. } => fjall_map_keys(values, name),
433        }
434    }
435
436    pub fn remove_key(&self, name: &str, key: &str) -> Result<T> {
437        match self {
438            Self::Memory(m) => m.update_sync(name, |_, v| v.remove_key(key)).flatten().ok_or(anyhow!("remove_key {} 失败", name)),
439            Self::Redis { client, rl: _ } => {
440                let mut conn = client.get_connection()?;
441                let buf: Vec<u8> = conn.hget(name, key)?;
442                let v = T::decode(buf.as_slice())?.0;
443                let _: usize = conn.hdel(name, key)?;
444                if !conn.exists::<_, bool>(name)? {
445                    directory::remove_path(&mut conn, name)?;
446                }
447                Ok(v)
448            }
449            Self::Fjall { values, .. } => {
450                let item_key = fjall_map_item_key(name, key);
451                let buf = values.get(&item_key)?.ok_or(anyhow!("remove_key {} 失败", name))?;
452                let (v, _) = T::decode(buf.as_ref())?;
453                values.remove(item_key)?;
454                Ok(v)
455            }
456        }
457    }
458}
459
460#[derive(Clone, Copy)]
461enum FjallNodeType {
462    List,
463    Map,
464}
465
466const FJALL_OBJECT_PREFIX: u8 = b'o';
467const FJALL_TYPE_PREFIX: u8 = b't';
468const FJALL_LIST_PREFIX: u8 = b'l';
469const FJALL_LIST_COUNTER_PREFIX: u8 = b'c';
470const FJALL_MAP_PREFIX: u8 = b'm';
471const FJALL_SEPARATOR: u8 = 0;
472
473fn fjall_key(prefix: u8, name: &str) -> Vec<u8> {
474    let mut key = Vec::with_capacity(2 + name.len());
475    key.push(prefix);
476    key.push(FJALL_SEPARATOR);
477    key.extend_from_slice(name.as_bytes());
478    key
479}
480
481fn fjall_object_key(name: &str) -> Vec<u8> {
482    fjall_key(FJALL_OBJECT_PREFIX, name)
483}
484
485fn fjall_type_key(name: &str) -> Vec<u8> {
486    fjall_key(FJALL_TYPE_PREFIX, name)
487}
488
489fn fjall_list_counter_key(name: &str) -> Vec<u8> {
490    fjall_key(FJALL_LIST_COUNTER_PREFIX, name)
491}
492
493fn fjall_item_prefix(prefix: u8, name: &str) -> Vec<u8> {
494    let mut key = fjall_key(prefix, name);
495    key.push(FJALL_SEPARATOR);
496    key
497}
498
499fn fjall_list_prefix(name: &str) -> Vec<u8> {
500    fjall_item_prefix(FJALL_LIST_PREFIX, name)
501}
502
503fn fjall_map_prefix(name: &str) -> Vec<u8> {
504    fjall_item_prefix(FJALL_MAP_PREFIX, name)
505}
506
507fn fjall_list_item_key(name: &str, idx: usize) -> Vec<u8> {
508    let mut key = fjall_list_prefix(name);
509    key.extend_from_slice(&(idx as u64).to_be_bytes());
510    key
511}
512
513fn fjall_map_item_key(name: &str, key: &str) -> Vec<u8> {
514    let mut item_key = fjall_map_prefix(name);
515    item_key.extend_from_slice(key.as_bytes());
516    item_key
517}
518
519fn fjall_decode_value<T: MsgUnpack>(buf: &[u8]) -> Result<T> {
520    T::decode(buf).map(|(value, _)| value)
521}
522
523fn fjall_encode_value<T: MsgPack>(value: &T) -> Vec<u8> {
524    let mut buf = Vec::new();
525    value.encode(&mut buf);
526    buf
527}
528
529fn fjall_get_object<T: MsgUnpack>(values: &OptimisticTxKeyspace, name: &str) -> Result<T> {
530    let buf = values.get(fjall_object_key(name))?.ok_or(anyhow!("{} 不存在", name))?;
531    fjall_decode_value(buf.as_ref())
532}
533
534fn fjall_insert_object<T: MsgPack>(values: &OptimisticTxKeyspace, name: &str, value: &T) -> Result<()> {
535    Ok(values.insert(fjall_object_key(name), fjall_encode_value(value))?)
536}
537
538fn fjall_get_list_item<T: MsgUnpack>(values: &OptimisticTxKeyspace, name: &str, idx: usize) -> Result<T> {
539    let buf = values.get(fjall_list_item_key(name, idx))?.ok_or(anyhow!("get_idx {} 失败", name))?;
540    fjall_decode_value(buf.as_ref())
541}
542
543fn fjall_insert_list_item<T: MsgPack>(values: &OptimisticTxKeyspace, name: &str, idx: usize, value: &T) -> Result<()> {
544    Ok(values.insert(fjall_list_item_key(name, idx), fjall_encode_value(value))?)
545}
546
547fn fjall_get_map_item<T: MsgUnpack>(values: &OptimisticTxKeyspace, name: &str, key: &str) -> Result<T> {
548    let buf = values.get(fjall_map_item_key(name, key))?.ok_or(anyhow!("get_key {} 失败", name))?;
549    fjall_decode_value(buf.as_ref())
550}
551
552fn fjall_insert_map_item<T: MsgPack>(values: &OptimisticTxKeyspace, name: &str, key: &str, value: &T) -> Result<()> {
553    Ok(values.insert(fjall_map_item_key(name, key), fjall_encode_value(value))?)
554}
555
556fn fjall_node_type(values: &OptimisticTxKeyspace, name: &str) -> Result<Option<FjallNodeType>> {
557    Ok(values.get(fjall_type_key(name))?.and_then(|value| match value.as_ref() {
558        b"L" => Some(FjallNodeType::List),
559        b"M" => Some(FjallNodeType::Map),
560        _ => None,
561    }))
562}
563
564fn fjall_set_node_type(values: &OptimisticTxKeyspace, name: &str, node_type: FjallNodeType) -> Result<()> {
565    let value = match node_type {
566        FjallNodeType::List => b"L".as_slice(),
567        FjallNodeType::Map => b"M".as_slice(),
568    };
569    Ok(values.insert(fjall_type_key(name), value)?)
570}
571
572fn fjall_clear_node(values: &OptimisticTxKeyspace, name: &str) -> Result<()> {
573    values.remove(fjall_object_key(name))?;
574    values.remove(fjall_type_key(name))?;
575    values.remove(fjall_list_counter_key(name))?;
576    fjall_remove_prefix(values, fjall_list_prefix(name))?;
577    fjall_remove_prefix(values, fjall_map_prefix(name))?;
578    Ok(())
579}
580
581fn fjall_remove_prefix(values: &OptimisticTxKeyspace, prefix: Vec<u8>) -> Result<()> {
582    let mut keys = Vec::new();
583    for item in values.inner().prefix(prefix) {
584        keys.push(item.key()?);
585    }
586    for key in keys {
587        values.remove(key)?;
588    }
589    Ok(())
590}
591
592fn fjall_next_list_idx(values: &OptimisticTxKeyspace, name: &str) -> Result<usize> {
593    let key = fjall_list_counter_key(name);
594    let current = values.get(&key)?.and_then(|buf| buf.as_ref().try_into().ok().map(u64::from_be_bytes)).unwrap_or(0);
595    values.insert(key, (current + 1).to_be_bytes())?;
596    Ok(current as usize)
597}
598
599fn fjall_count_prefix(values: &OptimisticTxKeyspace, prefix: Vec<u8>) -> Result<usize> {
600    let mut count = 0;
601    for item in values.inner().prefix(prefix) {
602        item.key()?;
603        count += 1;
604    }
605    Ok(count)
606}
607
608fn fjall_paths_with_prefix(values: &OptimisticTxKeyspace, prefix: &str) -> Result<Vec<SmolStr>> {
609    let mut names = Vec::new();
610    names.extend(fjall_paths_for_key_prefix(values, FJALL_OBJECT_PREFIX, prefix)?);
611    names.extend(fjall_paths_for_key_prefix(values, FJALL_TYPE_PREFIX, prefix)?);
612    names.sort();
613    names.dedup();
614    Ok(names)
615}
616
617fn fjall_paths_for_key_prefix(values: &OptimisticTxKeyspace, key_prefix: u8, path_prefix: &str) -> Result<Vec<SmolStr>> {
618    let mut scan_prefix = fjall_key(key_prefix, path_prefix);
619    if path_prefix.is_empty() {
620        scan_prefix.truncate(2);
621    }
622    let mut names = Vec::new();
623    for item in values.inner().prefix(scan_prefix) {
624        let key = item.key()?;
625        let Some(path) = key.get(2..) else {
626            continue;
627        };
628        if let Ok(path) = std::str::from_utf8(path) {
629            names.push(path.into());
630        }
631    }
632    Ok(names)
633}
634
635fn fjall_map_keys(values: &OptimisticTxKeyspace, name: &str) -> Result<Vec<SmolStr>> {
636    let prefix = fjall_map_prefix(name);
637    let mut keys = Vec::new();
638    for item in values.inner().prefix(&prefix) {
639        let key = item.key()?;
640        let Some(map_key) = key.get(prefix.len()..) else {
641            continue;
642        };
643        if let Ok(map_key) = std::str::from_utf8(map_key) {
644            keys.push(map_key.into());
645        }
646    }
647    Ok(keys)
648}
649
650use std::sync::RwLock;
651
652#[derive(Debug)]
653pub struct Root<T> {
654    mounts: Arc<RwLock<Vec<(SmolStr, Mount<T>)>>>,
655}
656
657impl<T: std::fmt::Debug> std::fmt::Debug for Mount<T> {
658    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
659        match self {
660            Self::Memory(_) => write!(f, "Mount::Memory"),
661            Self::Redis { client: _, rl: _ } => write!(f, "Mount::Redis"),
662            Self::Fjall { .. } => write!(f, "Mount::Fjall"),
663        }
664    }
665}
666
667impl<T> Clone for Mount<T> {
668    fn clone(&self) -> Self {
669        match self {
670            Self::Memory(m) => Self::Memory(m.clone()),
671            Self::Redis { client, rl } => Self::Redis { client: client.clone(), rl: rl.clone() },
672            Self::Fjall { values, write_lock } => Self::Fjall { values: values.clone(), write_lock: write_lock.clone() },
673        }
674    }
675}
676
677impl<T: std::fmt::Debug + MsgPack + MsgUnpack + Default + Send> Root<T> {
678    pub fn new() -> Self {
679        Self { mounts: Arc::new(RwLock::new(vec![("local".into(), Mount::<T>::memory())])) }
680    }
681
682    pub fn mount_memory(&self, name: &str) -> bool {
683        let mounts = self.mounts.write();
684        match mounts {
685            Ok(mut mounts) => {
686                if mounts.iter().any(|(n, _)| n == name) {
687                    return false;
688                }
689                mounts.push((name.into(), Mount::<T>::memory()));
690                true
691            }
692            Err(_) => false,
693        }
694    }
695
696    pub fn mount_redis(&self, name: &str, url: &str) -> Result<bool> {
697        let mounts = self.mounts.write();
698        match mounts {
699            Ok(mut mounts) => {
700                if mounts.iter().any(|(n, _)| n == name) {
701                    return Ok(false);
702                }
703                mounts.push((name.into(), Mount::<T>::redis(url)?));
704                Ok(true)
705            }
706            Err(e) => Err(anyhow!("无法获取写锁: {}", e)),
707        }
708    }
709
710    pub fn mount_fjall(&self, name: &str, data_dir: &str) -> Result<bool> {
711        let mounts = self.mounts.write();
712        match mounts {
713            Ok(mut mounts) => {
714                if mounts.iter().any(|(n, _)| n == name) {
715                    return Ok(false);
716                }
717                mounts.push((name.into(), Mount::<T>::fjall(data_dir)?));
718                Ok(true)
719            }
720            Err(e) => Err(anyhow!("无法获取写锁: {}", e)),
721        }
722    }
723
724    pub fn get_mount<'a>(&self, name: &'a str) -> Result<(Mount<T>, &'a str)> {
725        let (mount, name) = name.split_once('/').ok_or(anyhow!("{} 没有 root 路径", name))?;
726        let mounts = self.mounts.read().map_err(|e| anyhow!("无法获取读锁: {}", e))?;
727        let m = mounts.iter().find_map(|m| if m.0 == mount { Some(m.1.clone()) } else { None }).ok_or(anyhow!("没有找到 {}", name))?;
728        Ok((m, name))
729    }
730}
731
732#[cfg(test)]
733mod tests {
734    use super::*;
735    use dynamic::Dynamic;
736
737    #[test]
738    fn dir_returns_only_immediate_children() {
739        let root = Root::<Dynamic>::new();
740        let (mount, name) = root.get_mount("local/test/dir/a").unwrap();
741        assert!(mount.add(name, 1.into()));
742        let (mount, name) = root.get_mount("local/test/dir/sub/item").unwrap();
743        assert!(mount.add(name, 2.into()));
744        let (mount, name) = root.get_mount("local/test/dir2/x").unwrap();
745        assert!(mount.add(name, 3.into()));
746
747        let (mount, name) = root.get_mount("local/test/dir").unwrap();
748        let mut entries = mount.dir(name).unwrap();
749        entries.sort();
750
751        assert_eq!(entries, vec![SmolStr::new("test/dir/a"), SmolStr::new("test/dir/sub")]);
752    }
753
754    #[test]
755    fn dir_accepts_trailing_slash() {
756        let root = Root::<Dynamic>::new();
757        let (mount, name) = root.get_mount("local/test/slash/a").unwrap();
758        assert!(mount.add(name, 1.into()));
759
760        let (mount, name) = root.get_mount("local/test/slash/").unwrap();
761        assert_eq!(mount.dir(name).unwrap(), vec![SmolStr::new("test/slash/a")]);
762    }
763
764    #[test]
765    fn fjall_mount_persists_values_and_dirs() {
766        let data_dir = std::env::temp_dir().join(format!("zust-root-fjall-{}", uuid::Uuid::new_v4()));
767        let data_dir_str = data_dir.to_str().unwrap();
768
769        {
770            let root = Root::<Dynamic>::new();
771            assert!(root.mount_fjall("fjall", data_dir_str).unwrap());
772            let (mount, name) = root.get_mount("fjall/test/kv/a").unwrap();
773            assert!(mount.add(name, 42.into()));
774            let (mount, name) = root.get_mount("fjall/test/kv/b").unwrap();
775            assert!(mount.add(name, "persisted".into()));
776            let (mount, name) = root.get_mount("fjall/test/list").unwrap();
777            mount.add_list(name);
778            assert_eq!(mount.push(name, 7.into()).unwrap(), 0);
779            let (mount, name) = root.get_mount("fjall/test/map").unwrap();
780            mount.add_map(name);
781            mount.insert(name, "answer", 42.into());
782        }
783
784        {
785            let root = Root::<Dynamic>::new();
786            assert!(root.mount_fjall("fjall", data_dir_str).unwrap());
787            let (mount, name) = root.get_mount("fjall/test/kv/a").unwrap();
788            assert_eq!(mount.get(name, |v| v.as_int()).unwrap(), Some(42));
789            let (mount, name) = root.get_mount("fjall/test").unwrap();
790            let mut entries = mount.dir(name).unwrap();
791            entries.sort();
792            assert_eq!(entries, vec![SmolStr::new("test/kv"), SmolStr::new("test/list"), SmolStr::new("test/map")]);
793            let (mount, name) = root.get_mount("fjall/test/list").unwrap();
794            assert_eq!(mount.get_idx(name, 0, |v| v.as_int()).unwrap(), Some(7));
795            let (mount, name) = root.get_mount("fjall/test/map").unwrap();
796            assert_eq!(mount.get_key(name, "answer", |v| v.as_int()).unwrap(), Some(42));
797        }
798
799        let _ = std::fs::remove_dir_all(data_dir);
800    }
801}