diff --git a/Cargo.lock b/Cargo.lock index 3b35ddb..98e02c3 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1058,6 +1058,7 @@ dependencies = [ "axum", "axum-extra", "axum-metrics", + "base64 0.22.1", "bincode 1.3.3", "clap", "ctrlc", diff --git a/constellation/Cargo.toml b/constellation/Cargo.toml index 18cd525..4c26982 100644 --- a/constellation/Cargo.toml +++ b/constellation/Cargo.toml @@ -10,6 +10,7 @@ askama = { version = "0.12.1", features = ["serde-json"] } axum = "0.8.1" axum-extra = { version = "0.10.0", features = ["query", "typed-header"] } axum-metrics = "0.2" +base64 = "0.22.1" bincode = "1.3.3" clap = { workspace = true } ctrlc = "3.4.5" diff --git a/constellation/src/lib.rs b/constellation/src/lib.rs index 3452ed5..a2b4e98 100644 --- a/constellation/src/lib.rs +++ b/constellation/src/lib.rs @@ -22,7 +22,7 @@ pub enum ActionableEvent { DeleteAccount(Did), } -#[derive(Debug, PartialEq, Eq, Hash, Clone, Serialize, Deserialize)] +#[derive(Debug, Hash, Clone, Serialize, Deserialize, PartialEq, Eq, PartialOrd, Ord)] pub struct Did(pub String); impl> From for Did { @@ -31,7 +31,7 @@ impl> From for Did { } } -#[derive(Debug, PartialEq, Serialize, Deserialize, Clone)] +#[derive(Debug, Serialize, Deserialize, Clone, PartialEq, Eq, PartialOrd, Ord)] pub struct RecordId { pub did: Did, pub collection: String, diff --git a/constellation/src/storage/mem_store.rs b/constellation/src/storage/mem_store.rs index 54965ed..52658ad 100644 --- a/constellation/src/storage/mem_store.rs +++ b/constellation/src/storage/mem_store.rs @@ -1,6 +1,7 @@ use super::{ LinkReader, LinkStorage, Order, PagedAppendingCollection, PagedOrderedCollection, StorageStats, }; +use crate::storage::{decode_m2m_cursor, encode_m2m_cursor}; use crate::{ActionableEvent, CountsByCount, Did, RecordId}; use anyhow::Result; use links::CollectedLink; @@ -250,11 +251,6 @@ impl LinkReader for MemStorage { next: None, }); - // struct MemStorageData { - // dids: HashMap, - // targets: HashMap>, - // links: HashMap>>, - // } let data = self.0.lock().unwrap(); let Some(sources) = data.targets.get(&Target::new(target)) else { @@ -315,16 +311,48 @@ impl LinkReader for MemStorage { }) .collect::>(); - items.sort_by(|a: &(RecordId, String), b| a.1.cmp(&b.1)); + // first try to sort by subject, then by did, collection and finally rkey + items.sort_by(|a, b| { + if a.1 == b.1 { + a.0.cmp(&b.0) + } else { + a.1.cmp(&b.1) + } + }); + // Parse cursor if provided (malformed cursor silently ignored) + let after_cursor = after.and_then(|a| decode_m2m_cursor(&a).ok()); + + // Apply cursor: skip everything up to and including the cursor position items = items .into_iter() - .skip_while(|item| after.as_ref().map(|a| &item.1 <= a).unwrap_or(false)) - .take(limit as usize) + .skip_while(|item| { + let Some((after_did, after_rkey, after_subject)) = &after_cursor else { + return false; + }; + + if &item.1 == after_subject { + // Same subject — compare by RecordId to find our position + let cursor_id = RecordId { + did: Did(after_did.clone()), + collection: collection.to_string(), + rkey: after_rkey.clone(), + }; + item.0.cmp(&cursor_id).is_le() + } else { + // Different subject — compare subjects directly + item.1.cmp(after_subject).is_le() + } + }) + .take(limit as usize + 1) .collect(); - let next = if items.len() as u64 >= limit { - items.last().map(|item| item.1.clone()) + // Build the new cursor from last item, if needed + let next = if items.len() as u64 > limit { + items.truncate(limit as usize); + items + .last() + .map(|item| encode_m2m_cursor(&item.0.did.0, &item.0.rkey, &item.1)) } else { None }; diff --git a/constellation/src/storage/mod.rs b/constellation/src/storage/mod.rs index ed2014c..c72a5b0 100644 --- a/constellation/src/storage/mod.rs +++ b/constellation/src/storage/mod.rs @@ -6,6 +6,11 @@ use std::collections::{HashMap, HashSet}; pub mod mem_store; pub use mem_store::MemStorage; +use anyhow::anyhow; + +use base64::engine::general_purpose as b64; +use base64::Engine as _; + #[cfg(feature = "rocks")] pub mod rocks_store; #[cfg(feature = "rocks")] @@ -156,6 +161,35 @@ pub trait LinkReader: Clone + Send + Sync + 'static { fn get_stats(&self) -> Result; } +// Shared helpers + +/// Decode a base64 cursor into its component parts (did, rkey, subject). +/// The subject is placed last because it may contain '|' characters. +pub(crate) fn decode_m2m_cursor(cursor: &str) -> Result<(String, String, String)> { + let decoded = String::from_utf8(b64::URL_SAFE.decode(cursor)?)?; + let mut parts = decoded.splitn(3, '|').map(String::from); + + // Using .next() to pull each part out of the iterator in order. + // This avoids collecting into a Vec just to index and clone back out. + let did = parts + .next() + .ok_or_else(|| anyhow!("missing did in cursor"))?; + let rkey = parts + .next() + .ok_or_else(|| anyhow!("missing rkey in cursor"))?; + let subject = parts + .next() + .ok_or_else(|| anyhow!("missing subject in cursor"))?; + + Ok((did, rkey, subject)) +} + +/// Encode cursor components into a base64 string. +pub(crate) fn encode_m2m_cursor(did: &str, rkey: &str, subject: &str) -> String { + let raw = format!("{did}|{rkey}|{subject}"); + b64::URL_SAFE.encode(&raw) +} + #[cfg(test)] mod tests { use super::*; @@ -1904,5 +1938,86 @@ mod tests { assert_eq!(result.items.len(), 2); assert!(result.items.iter().all(|(_, subject)| subject == "b.com")); assert!(result.items.iter().all(|(r, _)| r.did.0 == "did:plc:asdf")); + + // Pagination edge cases: we have 4 flat items + + // Case 1: limit > items (limit=10, items=4) -> next should be None + let result = storage.get_many_to_many( + "a.com", + "app.t.c", + ".abc.uri", + ".def.uri", + 10, + None, + &HashSet::new(), + &HashSet::new(), + )?; + assert_eq!(result.items.len(), 4); + assert_eq!(result.next, None, "next should be None when items < limit"); + + // Case 2: limit == items (limit=4, items=4) -> next should be None + let result = storage.get_many_to_many( + "a.com", + "app.t.c", + ".abc.uri", + ".def.uri", + 4, + None, + &HashSet::new(), + &HashSet::new(), + )?; + assert_eq!(result.items.len(), 4); + assert_eq!( + result.next, None, + "next should be None when items == limit (no more pages)" + ); + + // Case 3: limit < items (limit=3, items=4) -> next should be Some + let result = storage.get_many_to_many( + "a.com", + "app.t.c", + ".abc.uri", + ".def.uri", + 3, + None, + &HashSet::new(), + &HashSet::new(), + )?; + assert_eq!(result.items.len(), 3); + assert!( + result.next.is_some(), + "next should be Some when items > limit" + ); + + // Verify second page returns remaining item with no cursor. + // This now works correctly because we use a composite cursor that includes + // (target, did, rkey), allowing pagination even when multiple records share + // the same target string. + let result2 = storage.get_many_to_many( + "a.com", + "app.t.c", + ".abc.uri", + ".def.uri", + 3, + result.next, + &HashSet::new(), + &HashSet::new(), + )?; + assert_eq!( + result2.items.len(), + 1, + "second page should have 1 remaining item" + ); + assert_eq!(result2.next, None, "next should be None on final page"); + + // Verify we got all 4 unique items across both pages (no duplicates, no gaps) + let mut all_rkeys: Vec<_> = result.items.iter().map(|(r, _)| r.rkey.clone()).collect(); + all_rkeys.extend(result2.items.iter().map(|(r, _)| r.rkey.clone())); + all_rkeys.sort(); + assert_eq!( + all_rkeys, + vec!["asdf", "asdf2", "fdsa", "fdsa2"], + "should have all 4 records across both pages" + ); }); } diff --git a/constellation/src/storage/rocks_store.rs b/constellation/src/storage/rocks_store.rs index cf8c878..3e0ab44 100644 --- a/constellation/src/storage/rocks_store.rs +++ b/constellation/src/storage/rocks_store.rs @@ -2,6 +2,7 @@ use super::{ ActionableEvent, LinkReader, LinkStorage, Order, PagedAppendingCollection, PagedOrderedCollection, StorageStats, }; +use crate::storage::{decode_m2m_cursor, encode_m2m_cursor}; use crate::{CountsByCount, Did, RecordId}; use anyhow::{bail, Result}; use bincode::Options as BincodeOptions; @@ -1138,7 +1139,8 @@ impl LinkReader for RocksStorage { let target_key = TargetKey(Target(target.to_string()), collection.clone(), path); - let after = after.map(|s| s.parse::().map(TargetId)).transpose()?; + // Parse cursor if provided (malformed cursor silently ignored) + let after_cursor = after.and_then(|a| decode_m2m_cursor(&a).ok()); let Some(target_id) = self.target_id_table.get_id_val(&self.db, &target_key)? else { eprintln!("Target not found for {target_key:?}"); @@ -1212,10 +1214,6 @@ impl LinkReader for RocksStorage { continue; }; - // pagination logic mirrors what is currently done in get_many_to_many_counts - if after.as_ref().map(|a| fwd_target <= *a).unwrap_or(false) { - continue; - } let page_is_full = grouped_links.len() as u64 >= limit; if page_is_full { let current_max = grouped_links.keys().next_back().unwrap(); @@ -1224,17 +1222,53 @@ impl LinkReader for RocksStorage { } } + // link to be added + let record_id = RecordId { + did, + collection: collection.0.clone(), + rkey: rkey.0, + }; + + // pagination: + if after_cursor.is_some() { + // extract composite-cursor parts + let Some((after_did, after_rkey, after_subject)) = &after_cursor else { + continue; + }; + + let Some(fwd_target_key) = self + .target_id_table + .get_val_from_id(&self.db, fwd_target.0)? + else { + eprintln!("failed to look up target from target_id {fwd_target:?}"); + continue; + }; + + // first try and compare by subject only + if &fwd_target_key.0 .0 != after_subject + && fwd_target_key.0 .0.cmp(after_subject).is_le() + { + continue; + } + + // then, if needed, we compare by record id + let cursor_id = RecordId { + did: Did(after_did.clone()), + collection: collection.0.clone(), + rkey: after_rkey.clone(), + }; + if record_id.cmp(&cursor_id).is_le() { + continue; + } + } + // pagination, continued let mut should_evict = false; let entry = grouped_links.entry(fwd_target.clone()).or_insert_with(|| { should_evict = page_is_full; Vec::default() }); - entry.push(RecordId { - did, - collection: collection.0.clone(), - rkey: rkey.0, - }); + entry.push(record_id); if should_evict { grouped_links.pop_last(); @@ -1258,8 +1292,12 @@ impl LinkReader for RocksStorage { .for_each(|r| items.push((r.clone(), target_string.clone()))); } - let next = if grouped_links.len() as u64 >= limit { - grouped_links.keys().next_back().map(|k| format!("{}", k.0)) + // Build new cursor from last the item, if needed + let next = if items.len() as u64 > limit { + items.truncate(limit as usize); + items + .last() + .map(|item| encode_m2m_cursor(&item.0.did.0, &item.0.rkey, &item.1)) } else { None };