diff --git a/src/database/map.rs b/src/database/map.rs index b55a5e449..40196a4ba 100644 --- a/src/database/map.rs +++ b/src/database/map.rs @@ -47,7 +47,10 @@ use self::watch::Watch; /// /// `Get` accepts raw keys, while `Qry` serializes structured keys before /// lookup. Both yield pinned value handles through an asynchronous stream. -pub use self::{get_batch::Get, qry_batch::Qry}; +pub use self::{ + get_batch::{Get, RecursiveGetOutput}, + qry_batch::Qry, +}; use crate::{Engine, util::map_err}; /// Provides typed and raw access to one RocksDB column family. diff --git a/src/database/map/get_batch.rs b/src/database/map/get_batch.rs index 388d4bb82..2f12e252c 100644 --- a/src/database/map/get_batch.rs +++ b/src/database/map/get_batch.rs @@ -1,4 +1,4 @@ -use std::sync::Arc; +use std::{collections::HashSet, hash::Hash, sync::Arc}; use futures::{Stream, StreamExt, TryStreamExt}; use rocksdb::{DBPinnableSlice, ReadOptions}; @@ -11,7 +11,7 @@ use tuwunel_core::{ }; use super::get::{cached_handle_from, handle_from}; -use crate::Handle; +use crate::{Handle, util::map_err}; /// Extends a stream of raw keys with batched map lookup. /// @@ -135,3 +135,157 @@ where .batched_multi_get_cf_opt(&self.cf(), keys, SORTED, read_options) .into_iter() } + +/// Result container for recursive multi-get DAG traversals. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct RecursiveGetOutput { + /// Values successfully fetched and parsed during traversal. + pub values: Vec, + + /// Keys requested during traversal that were missing from the database. + pub missing: Vec, + + /// Indicates whether the traversal stopped early due to node or depth caps. + pub truncated: bool, +} + +/// Performs a recursive breadth-first traversal over database keys. +/// +/// Starting from `roots`, each batch of keys is fetched in RocksDB using +/// `batched_multi_get_cf_opt` against a point-in-time snapshot. Returned values +/// are parsed by `parse_value`, and any child keys appended to the sink buffer +/// by `extract_children` are queued for the next level of traversal. +/// +/// # Traversal Ordering +/// Results are ordered level-by-level (BFS order). Within a single level, +/// results reflect key sorting order. Note that this is **not** a topological +/// sort. +/// +/// # Bounds & Limits +/// Traversal halts early if `max_nodes` (total parsed values) or `max_depth` +/// (BFS depth iterations) is reached, marking `truncated = true` on the +/// returned output. +/// +/// # Errors +/// Fails fast on server shutdown, key parsing failure, RocksDB I/O errors, or +/// block corruption. +#[implement(super::Map)] +#[tracing::instrument(skip_all, level = "trace")] +pub async fn recursive_multi_get( + self: &Arc, + roots: I, + max_nodes: Option, + max_depth: Option, + parse_value: P, + extract_children: F, +) -> Result> +where + K: AsRef<[u8]> + Ord + Hash + Clone + Send + Sync + 'static, + V: Send + 'static, + P: Fn(&[u8]) -> Result + Send + Sync + 'static, + F: Fn(&V, &mut Vec) + Send + Sync + 'static, + I: IntoIterator + Send + 'static, +{ + let map = self.clone(); + + tokio::task::spawn_blocking(move || { + const SORTED: bool = true; + + map.engine.ctx.server.check_running()?; + + let snapshot = map.engine.db.snapshot(); + let mut read_options = super::options::read_options_default(&map.engine); + read_options.set_snapshot(&snapshot); + + let mut visited = HashSet::new(); + let mut current_batch = Vec::new(); + for root in roots { + if visited.insert(root.clone()) { + current_batch.push(root); + } + } + + let mut values = Vec::new(); + let mut missing = Vec::new(); + let mut depth: usize = 0; + let mut truncated = false; + + while !current_batch.is_empty() { + if let Some(max_d) = max_depth + && depth >= max_d + { + truncated = true; + break; + } + + // Sort keys for optimal sequential RocksDB multi-get access + current_batch.sort_unstable_by(|a, b| a.as_ref().cmp(b.as_ref())); + + if max_nodes.is_some_and(|max_n| values.len() >= max_n) { + truncated = true; + break; + } + + let db_results = map.engine.db.batched_multi_get_cf_opt( + &map.cf(), + current_batch.iter(), + SORTED, + &read_options, + ); + + let mut next_batch = Vec::with_capacity(current_batch.len().saturating_mul(2)); + + for (key, result) in current_batch.into_iter().zip(db_results) { + match result { + | Ok(Some(slice)) => + if max_nodes.is_none_or(|max_n| values.len() < max_n) { + let parsed_value = parse_value(slice.as_ref())?; + extract_children(&parsed_value, &mut next_batch); + values.push(parsed_value); + + if max_nodes.is_some_and(|max_n| values.len() >= max_n) { + truncated = true; + } + } else { + truncated = true; + }, + | Ok(None) => { + missing.push(key); + }, + | Err(e) => { + tracing::error!( + key = ?key.as_ref(), + %e, + "RocksDB multi-get failure during recursive DAG traversal" + ); + return Err(map_err(e)); + }, + } + } + + // Filter out already visited keys from next_batch while preserving order + next_batch.retain(|child| visited.insert(child.clone())); + + depth = depth.saturating_add(1); + + if truncated { + break; + } + + current_batch = next_batch; + + map.engine.ctx.server.check_running()?; + } + + Ok(RecursiveGetOutput { values, missing, truncated }) + }) + .await + .map_err(|e| { + if e.is_panic() { + tracing::error!("blocking task panicked during recursive_multi_get"); + std::io::Error::other("recursive_multi_get task panicked") + } else { + std::io::Error::other("recursive_multi_get task cancelled") + } + })? +} diff --git a/src/database/mod.rs b/src/database/mod.rs index 98e572b20..5ec7e48c9 100644 --- a/src/database/mod.rs +++ b/src/database/mod.rs @@ -40,7 +40,7 @@ pub use self::{ engine::Engine, handle::Handle, keyval::{KeyBuf, KeyVal, Slice, serialize_key, serialize_val}, - map::{Get, Map, Qry, compact}, + map::{Get, Map, Qry, RecursiveGetOutput, compact}, ser::{Cbor, Interfix, Json, SEP, Separator, serialize, serialize_to, serialize_to_vec}, txn::Txn, }; diff --git a/src/database/tests.rs b/src/database/tests.rs index b5f8d4e57..a0901ba01 100644 --- a/src/database/tests.rs +++ b/src/database/tests.rs @@ -1304,3 +1304,115 @@ async fn a_restore_is_not_repeated_on_reopen() -> Result { Ok(()) } + +#[tokio::test] +async fn recursive_multi_get_traversal() -> Result<()> { + let root = var("TMPDIR").unwrap_or_else(|_| "/tmp".into()); + let path = format!("{root}/tuwunel-database-recursive-{}", process_id()); + let raw_config = Figment::new() + .merge(("server_name", "localhost")) + .merge(("database_path", &path)) + .merge(("test", ["fresh", "cleanup"])); + + let config = Config::new(&raw_config)?; + let runtime = Handle::current(); + let logging = Logging { + subscriber: Arc::new(NoSubscriber::new()), + reload: LogLevelReloadHandles::default(), + capture: Arc::new(State::new()), + }; + + let metrics = Metrics::new(Some(&runtime)); + let server = + Arc::new(Server::new(config, Sources::default(), Some(&runtime), logging, metrics)); + + let db = Database::open(&server).await?; + let map = &db["global"]; + + // Insert DAG nodes: + // A -> B, C + // B -> A (cycle) & D (diamond convergence) + // C -> D (diamond convergence) & M (missing) + // D -> E + map.insert(b"node_A", b"node_B,node_C"); + map.insert(b"node_B", b"node_A,node_D"); + map.insert(b"node_C", b"node_D,node_M"); // node_M will be missing + map.insert(b"node_D", b"node_E"); + map.insert(b"node_E", b""); + + let parse_val = |slice: &[u8]| -> Result { + String::from_utf8(slice.to_vec()).map_err(|e| std::io::Error::other(e).into()) + }; + + let extract_children = |val: &String, sink: &mut Vec>| { + if !val.is_empty() { + for part in val.split(',') { + sink.push(part.as_bytes().to_vec()); + } + } + }; + + // Test 1: Full traversal with cycle, diamond, and missing key detection + let output = map + .recursive_multi_get( + vec![b"node_A".to_vec(), b"node_A".to_vec()], + None, + None, + parse_val, + extract_children, + ) + .await?; + + assert!(!output.truncated); + assert_eq!(output.missing, vec![b"node_M".to_vec()]); + assert_eq!(output.values, vec![ + "node_B,node_C", + "node_A,node_D", + "node_D,node_M", + "node_E", + "", + ]); + + // Test 2: Truncation via max_depth + let depth_output = map + .recursive_multi_get(vec![b"node_A".to_vec()], None, Some(1), parse_val, extract_children) + .await?; + + assert!(depth_output.truncated); + assert_eq!(depth_output.values.len(), 1); + assert_eq!(depth_output.values, vec!["node_B,node_C"]); + + // Test 3: Truncation via max_nodes + let node_output = map + .recursive_multi_get(vec![b"node_A".to_vec()], Some(2), None, parse_val, extract_children) + .await?; + + assert!(node_output.truncated); + assert_eq!(node_output.values.len(), 2); + assert_eq!(node_output.values, vec!["node_B,node_C", "node_A,node_D"]); + + // Test 4: Truncation via max_nodes = Some(0) + let zero_node_output = map + .recursive_multi_get(vec![b"node_A".to_vec()], Some(0), None, parse_val, extract_children) + .await?; + + assert!(zero_node_output.truncated); + assert!(zero_node_output.values.is_empty()); + + // Test 5: Mid-batch truncation preserves missing key recording and error checks + let mid_batch_output = map + .recursive_multi_get( + vec![b"node_C".to_vec(), b"node_M".to_vec()], + Some(1), + None, + parse_val, + extract_children, + ) + .await?; + + assert!(mid_batch_output.truncated); + assert_eq!(mid_batch_output.values, vec!["node_D,node_M"]); + assert_eq!(mid_batch_output.missing, vec![b"node_M".to_vec()]); + + Ok(()) +} diff --git a/src/service/migrations/injectivity/repair.rs b/src/service/migrations/injectivity/repair.rs index 4d6242564..9c287edf1 100644 --- a/src/service/migrations/injectivity/repair.rs +++ b/src/service/migrations/injectivity/repair.rs @@ -218,7 +218,7 @@ async fn patch_statediffs(services: &Services, scan: &Scan) -> Result { .ready_fold(Digests::new(), |mut digests, (key, value)| { let infected = short_of(value) .filter(|state| scan.infected.contains(state)) - .and_then(|state| key.try_into().ok().map(|digest| (state, digest))); + .zip(key.try_into().ok()); if let Some((state, digest)) = infected { digests.entry(state).or_default().push(digest);