impl: recursive rocksdb multi get

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.

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<V, K> {
+	/// Values successfully fetched and parsed during traversal.
+	pub values: Vec<V>,
+
+	/// Keys requested during traversal that were missing from the database.
+	pub missing: Vec<K>,
+
+	/// 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<K, V, P, F, I>(
+	self: &Arc<Self>,
+	roots: I,
+	max_nodes: Option<usize>,
+	max_depth: Option<usize>,
+	parse_value: P,
+	extract_children: F,
+) -> Result<RecursiveGetOutput<V, K>>
+where
+	K: AsRef<[u8]> + Ord + Hash + Clone + Send + Sync + 'static,
+	V: Send + 'static,
+	P: Fn(&[u8]) -> Result<V> + Send + Sync + 'static,
+	F: Fn(&V, &mut Vec<K>) + Send + Sync + 'static,
+	I: IntoIterator<Item = K> + 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")
+		}
+	})?
+}

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,
 };

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> {
+		String::from_utf8(slice.to_vec()).map_err(|e| std::io::Error::other(e).into())
+	};
+
+	let extract_children = |val: &String, sink: &mut Vec<Vec<u8>>| {
+		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(())
+}

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);