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