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