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