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