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;
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 = node_id.parse()?;
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); 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); let lock_name = format!("{}::{}", name, key); 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 } 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 } 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}