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