diff --git a/bobbin/crates/xrpc/src/enrich.rs b/bobbin/crates/xrpc/src/enrich.rs index f7512ac8..b77bad0a 100644 --- a/bobbin/crates/xrpc/src/enrich.rs +++ b/bobbin/crates/xrpc/src/enrich.rs @@ -1,4 +1,4 @@ -use std::collections::HashSet; +use std::collections::{HashMap, HashSet}; use axum::{ Json, @@ -29,6 +29,7 @@ pub const TYPE_VIEWER: &str = "sh.tangled.query.enrichResponse#viewer"; pub const TYPE_MINIDOC: &str = "com.bad-example.identity.miniDoc"; const KNOWN_TYPES: [&str; 4] = [TYPE_COUNT, TYPE_DISTINCT_AUTHORS, TYPE_VIEWER, TYPE_MINIDOC]; +const MINIDOC_CONCURRENCY: usize = 32; /// a payload type nsid, with an optional #fragment for lexicon defs. the raw /// string is kept because it echoes into the data map as the payload key @@ -180,6 +181,8 @@ pub async fn enrich( State(state): State, Json(input): Json, ) -> Result, XrpcError> { + let mut seen_descriptors = HashSet::new(); + let mut descriptors = Vec::new(); for descriptor in &input.enrich { if !KNOWN_TYPES.contains(&descriptor.ty.as_str()) { return Err(descriptor_error(&descriptor.source, "unknown enrich type")); @@ -191,18 +194,38 @@ pub async fn enrich( "viewer payloads require a viewer param", )); } + if seen_descriptors.insert((&descriptor.source, &descriptor.ty)) { + descriptors.push(descriptor); + } } + let sources = input + .sources + .as_ref() + .map(|sources| { + sources + .iter() + .enumerate() + .map(|(i, path)| { + RecordPath::parse(path).map_err(|e| { + XrpcError::InvalidParams(format!("sources[{i}] {path:?}: {e}")) + }) + }) + .collect::, _>>() + }) + .transpose()?; + let inner = run_inner(&state, &input.xrpc, input.params.unwrap_or_default()).await?; + if descriptors.is_empty() { + return Ok(Json(json!({ "output": inner, "data": {} }))); + } let mut refs: Vec = Vec::new(); let mut seen: HashSet = HashSet::new(); - match &input.sources { + match &sources { Some(sources) => { - for (i, path) in sources.iter().enumerate() { - let path = RecordPath::parse(path) - .map_err(|e| XrpcError::InvalidParams(format!("sources[{i}] {path:?}: {e}")))?; - for node in walk_path(&path, [&inner]) { + for path in sources { + for node in walk_path(path, [&inner]) { collect_ref(node, &mut refs, &mut seen); } } @@ -222,8 +245,8 @@ pub async fn enrich( } let mut data = Map::new(); - let mut minidoc_targets: Vec<(Did, LinkSource)> = Vec::new(); - for descriptor in &input.enrich { + let mut minidoc_targets: HashMap, Vec> = HashMap::new(); + for descriptor in descriptors { for reference in &refs { let Some(subject) = applicable_subject(descriptor, reference)? else { continue; @@ -256,7 +279,10 @@ pub async fn enrich( } TYPE_MINIDOC => { if let Some(target) = repo_did(&subject) { - minidoc_targets.push((target, descriptor.source.clone())); + let sources = minidoc_targets.entry(target).or_default(); + if !sources.contains(&descriptor.source) { + sources.push(descriptor.source.clone()); + } } } _ => unreachable!("validated up front"), @@ -265,14 +291,16 @@ pub async fn enrich( } let docs = resolve_minidocs(&state, minidoc_targets).await; - for (target, source, doc) in docs { - put( - &mut data, - target.as_str(), - source.as_str(), - TYPE_MINIDOC, - doc, - ); + for (target, sources, doc) in docs { + for source in sources { + put( + &mut data, + target.as_str(), + source.as_str(), + TYPE_MINIDOC, + doc.clone(), + ); + } } Ok(Json(json!({ "output": inner, "data": data }))) @@ -296,20 +324,20 @@ fn put(data: &mut Map, reference: &str, source: &str, ty: &str, p /// we drop failures, the client falls back to resolveMiniDoc for misses async fn resolve_minidocs( state: &AppState, - targets: Vec<(Did, LinkSource)>, -) -> Vec<(Did, LinkSource, Value)> { + targets: HashMap, Vec>, +) -> Vec<(Did, Vec, Value)> { futures::stream::iter(targets) - .map(|(did, source)| async move { + .map(|(did, sources)| async move { let doc = state .slingshot .resolve_mini_doc(&AtIdentifier::Did(did.clone())) .await .ok() .and_then(|bytes| serde_json::from_slice::(&bytes).ok()); - (did, source, doc) + (did, sources, doc) }) - .buffered(crate::FETCH_CONCURRENCY) - .filter_map(|(did, source, doc)| async move { doc.map(|doc| (did, source, doc)) }) + .buffer_unordered(MINIDOC_CONCURRENCY) + .filter_map(|(did, sources, doc)| async move { doc.map(|doc| (did, sources, doc)) }) .collect() .await } diff --git a/bobbin/crates/xrpc/tests/enrich.rs b/bobbin/crates/xrpc/tests/enrich.rs index b5b6985d..2c8301c0 100644 --- a/bobbin/crates/xrpc/tests/enrich.rs +++ b/bobbin/crates/xrpc/tests/enrich.rs @@ -539,12 +539,14 @@ async fn minidoc_payloads_resolve_record_authors() { "handle": "a.example.com", "pds": "https://pds.example.com" }))) + .expect(1) .mount(&h.server) .await; Mock::given(method("GET")) .and(path("/xrpc/com.bad-example.identity.resolveMiniDoc")) .and(query_param("identifier", "did:plc:b")) .respond_with(ResponseTemplate::new(404)) + .expect(1) .mount(&h.server) .await; @@ -553,7 +555,11 @@ async fn minidoc_payloads_resolve_record_authors() { app.oneshot(enrich_request(json!({ "xrpc": "sh.tangled.graph.listFollows", "params": { "subject": owner.as_ref() }, - "enrich": [{ "source": "sh.tangled.graph.follow:.repo", "type": MINIDOC }] + "enrich": [ + { "source": "sh.tangled.graph.follow:.repo", "type": MINIDOC }, + { "source": "sh.tangled.graph.follow:.repo", "type": MINIDOC }, + { "source": "sh.tangled.feed.star:.repo", "type": MINIDOC } + ] }))) .await .unwrap(), @@ -565,6 +571,10 @@ async fn minidoc_payloads_resolve_record_authors() { body["data"]["did:plc:a"]["sh.tangled.graph.follow:.repo"][MINIDOC]["handle"], json!("a.example.com") ); + assert_eq!( + body["data"]["did:plc:a"]["sh.tangled.feed.star:.repo"][MINIDOC]["handle"], + json!("a.example.com") + ); // resolution failures are dropped, the client falls back for misses assert!(body["data"]["did:plc:b"].is_null(), "{body}"); // the profile owner authored nothing here, so it earns no minidoc