diff --git a/Cargo.lock b/Cargo.lock index 6f3a0adaf..619c73c20 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -528,6 +528,7 @@ name = "discovery-core" version = "0.1.0" dependencies = [ "async-trait", + "futures", "num-traits", "serde", "serde_json", @@ -764,6 +765,21 @@ version = "2.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e6d5a32815ae3f33302d95fdcb2ce17862f8c65363dcfd29360480ba1001fc9c" +[[package]] +name = "futures" +version = "0.3.31" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "65bc07b1a8bc7c85c5f2e110c476c7389b4554ba72af57d8445ea63a576b0876" +dependencies = [ + "futures-channel", + "futures-core", + "futures-executor", + "futures-io", + "futures-sink", + "futures-task", + "futures-util", +] + [[package]] name = "futures-channel" version = "0.3.31" @@ -771,6 +787,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2dff15bf788c671c1934e366d07e30c1814a8ef514e1af724a602e8a2fbe1b10" dependencies = [ "futures-core", + "futures-sink", ] [[package]] @@ -779,6 +796,23 @@ version = "0.3.31" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "05f29059c0c2090612e8d742178b0580d2dc940c837851ad723096f87af6663e" +[[package]] +name = "futures-executor" +version = "0.3.31" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1e28d1d997f585e54aebc3f97d39e72338912123a67330d723fdbb564d646c9f" +dependencies = [ + "futures-core", + "futures-task", + "futures-util", +] + +[[package]] +name = "futures-io" +version = "0.3.31" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9e5c1b78ca4aae1ac06c48a526a655760685149f0d465d21f37abfe57ce075c6" + [[package]] name = "futures-macro" version = "0.3.31" @@ -808,10 +842,13 @@ version = "0.3.31" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9fa08315bb612088cc391249efdc3bc77536f16c91f6cf495e6fbe85b20a4a81" dependencies = [ + "futures-channel", "futures-core", + "futures-io", "futures-macro", "futures-sink", "futures-task", + "memchr", "pin-project-lite", "pin-utils", "slab", diff --git a/crates/discovery-core/Cargo.toml b/crates/discovery-core/Cargo.toml index 565588116..ec9bebe46 100644 --- a/crates/discovery-core/Cargo.toml +++ b/crates/discovery-core/Cargo.toml @@ -11,7 +11,7 @@ starknet-core = "0.16" starknet-crypto = "0.8" starknet-types-core = { version = "0.2", features = ["curve", "serde"] } starknet-providers = "0.16" -tokio = { version = "1", features = ["rt"] } +futures = "0.3" tracing = "0.1" serde = { version = "1.0", features = ["derive"] } url = "2" diff --git a/crates/discovery-core/src/discovery/cursor.rs b/crates/discovery-core/src/discovery/cursor.rs index fbc998172..63339b833 100644 --- a/crates/discovery-core/src/discovery/cursor.rs +++ b/crates/discovery-core/src/discovery/cursor.rs @@ -65,9 +65,11 @@ pub struct SubchannelCursor { #[serde(default, skip_serializing_if = "Option::is_none")] pub last_note_index: Option, - /// First index where no note exists (sentinel). - /// - Incoming: not used. - /// - Outgoing: upper bound (hi) for exponential search; `Some` = bisection phase. + /// - Incoming: last index confirmed to exist by exponential probe. Linear + /// scan reads notes up to this index. Kept after scan — used to bound the + /// next exponential probe range (`max_note_index * 2`). Re-probe triggers + /// when `last_note_index == max_note_index`. + /// - Outgoing: first index confirmed empty (hi); `Some` = bisection phase. #[serde(default, skip_serializing_if = "Option::is_none")] pub max_note_index: Option, } diff --git a/crates/discovery-core/src/discovery/incoming_channels.rs b/crates/discovery-core/src/discovery/incoming_channels.rs index 317a1ad15..0ef65d428 100644 --- a/crates/discovery-core/src/discovery/incoming_channels.rs +++ b/crates/discovery-core/src/discovery/incoming_channels.rs @@ -104,44 +104,40 @@ pub async fn discover_incoming_channels( }); } - // Discover and decrypt each channel. - // Cap pre-allocation: total_n_channels may come from an untrusted cursor, - // so a malicious value must not cause OOM via Vec::with_capacity. - const MAX_CAPACITY: usize = 1024; - let capacity = usize::try_from(total_n_channels.saturating_sub(start_index)) - .unwrap_or(0) - .min(MAX_CAPACITY); - let mut channels = Vec::with_capacity(capacity); - let mut index = start_index; - let mut out_of_budget = false; - - loop { - // Check if we've processed all channels - if index >= total_n_channels { - break; - } + // Batch-read as many channels as budget allows in a single RPC call. + let remaining = usize::try_from(total_n_channels.saturating_sub(start_index)).unwrap_or(0); + let batch_size = budget.consume_up_to(remaining, COST_CHANNEL_INFO); + if batch_size == 0 { + return Ok(ChannelDiscoveryResult { + channels: vec![], + last_index: None, + has_more: true, + }); + } - // Consume budget for get_channel_info - if !budget.consume(COST_CHANNEL_INFO) { - out_of_budget = true; - break; - } + let encrypted_batch = privacy_pool + .get_channel_info_batch(recipient_addr, start_index, batch_size) + .await?; - let encrypted = privacy_pool.get_channel_info(recipient_addr, index).await?; + let mut channels = Vec::with_capacity(batch_size); + for (i, encrypted) in encrypted_batch.into_iter().enumerate() { + let index = start_index + + u64::try_from(i) + .map_err(|_| DiscoveryError::InvalidCursor("channel index overflow".into()))?; let info = decrypt_channel_info(&encrypted, private_key) .map_err(|source| DiscoveryError::Decryption { index, source })?; channels.push(IncomingChannel { index, info }); - index += 1; } let last_index = channels.last().map(|c| c.index); + let has_more = batch_size < remaining; Ok(ChannelDiscoveryResult { channels, last_index, - has_more: out_of_budget, + has_more, }) } diff --git a/crates/discovery-core/src/discovery/last_note_index.rs b/crates/discovery-core/src/discovery/last_note_index.rs index 93e20032b..67829bbfc 100644 --- a/crates/discovery-core/src/discovery/last_note_index.rs +++ b/crates/discovery-core/src/discovery/last_note_index.rs @@ -3,8 +3,8 @@ //! Finds the last note index in a subchannel via exponential search + bisection. //! No decryption, no nullifier checks — just probes note existence. //! -//! We assume an exponential distribution of note counts (most users have very -//! few notes), so exponential search finds the boundary quickly in the common case. +//! [`exponential_ascend`] is shared between notes discovery (finds `max_note_index` +//! for a linear scan) and [`find_last_note_index_paginated`] (outgoing channels). use std::future::Future; @@ -17,11 +17,23 @@ use crate::io_budget::IoBudget; use crate::privacy_pool::hashes::compute_note_id; use crate::privacy_pool::views::IViews; +/// Result of a batched exponential probe. +pub struct ExponentialProbeResult { + /// Last note found to exist: `(index, packed_amount)`. + /// `None` if the first probe missed (empty subchannel) or budget exhausted. + pub last_found_note: Option<(u64, Felt)>, + /// First index where no note exists (`None` if all probes hit = need more). + /// Used by [`find_last_note_index_paginated`] as bisection upper bound. + pub first_empty_index: Option, + /// Whether the probe ran out of budget before finding a boundary. + pub budget_exhausted: bool, +} + /// Finds the last note index via [`exponential_ascend`] + [`bisect_boundary`]. /// /// Runs in two phases, resumable via cursor: -/// 1. **Ascending**: exponential probing to find bounds (`lo`, `hi`). -/// 2. **Bisection**: binary search for exact boundary. +/// 1. **Ascending**: batched exponential probing to find bounds (`lo`, `hi`). +/// 2. **Bisection**: sequential binary search for exact boundary. /// /// Returns `(last_index, has_more)`. When `has_more` is `false`, the search is complete. pub async fn find_last_note_index_paginated( @@ -31,7 +43,31 @@ pub async fn find_last_note_index_paginated( cursor: &mut SubchannelCursor, budget: &IoBudget, ) -> Result<(Option, bool), DiscoveryError> { - // Returns (Some(idx), _) if note exists, (None, false) if empty, (None, true) if out of budget. + // Phase 1: Ascending (find bounds) + if cursor.max_note_index.is_none() { + let start = cursor.last_note_index.map_or(0, |lo| lo + 1); + let result = exponential_ascend(pool, channel_key, token, start, u64::MAX, budget).await?; + + if let Some((index, _)) = result.last_found_note { + cursor.last_note_index = Some(index); + } + if let Some(index) = result.first_empty_index { + cursor.max_note_index = Some(index); + } + + // Budget exhausted or all probes hit — need more probing. + if cursor.max_note_index.is_none() { + let has_more = result.budget_exhausted || result.last_found_note.is_some(); + return Ok((cursor.last_note_index, has_more)); + } + + // Empty subchannel: offset-0 probe found sentinel. + if cursor.last_note_index.is_none() { + return Ok((None, false)); + } + } + + // Phase 2: Bisection (narrow down to exact boundary, sequential) let probe = |idx: u64| { let budget = budget.clone(); async move { @@ -46,61 +82,78 @@ pub async fn find_last_note_index_paginated( } } }; - - // Phase 1: Ascending (find bounds) - if cursor.max_note_index.is_none() { - let found_sentinel = exponential_ascend(&probe, cursor).await?; - if !found_sentinel { - return Ok((cursor.last_note_index, true)); - } - // Empty subchannel: first probe found sentinel - if cursor.last_note_index.is_none() { - return Ok((None, false)); - } - } - - // Phase 2: Bisection (narrow down to exact boundary) let complete = bisect_boundary(&probe, cursor).await?; Ok((cursor.last_note_index, !complete)) } -/// Exponential ascending: probes at `lo+step`, `lo+2*step`, `lo+4*step`... -/// until finding an empty slot or exhausting budget. +/// Probes note existence at exponentially increasing indices in a single batch. /// -/// Updates cursor in place. Returns `true` if sentinel found, `false` if budget exhausted. -// TODO: Issue probes in batches to amortise RPC round-trip latency. -async fn exponential_ascend( - probe: F, - cursor: &mut SubchannelCursor, -) -> Result -where - F: Fn(u64) -> Fut, - Fut: Future, bool), DiscoveryError>>, -{ - // Compute step from lo: 2^k where k = trailing_zeros(lo + 1). - // Reconstructs the exponential sequence: 0, 1, 3, 7, 15, ... - let mut step = match cursor.last_note_index { - None | Some(0) => 1, - Some(n) => 1u64 << (n + 1).trailing_zeros(), - }; - let mut probe_at = cursor - .last_note_index - .map_or(0, |lo| lo.saturating_add(step)); - - loop { - match probe(probe_at).await? { - (Some(_), _) => { - cursor.last_note_index = Some(probe_at); - probe_at = probe_at.saturating_add(step); - step = step.saturating_mul(2); - } - (None, false) => { - cursor.max_note_index = Some(probe_at); - return Ok(true); - } - (None, true) => return Ok(false), +/// From `start`, probes at offsets `0, 2^0, 2^1, ..., 2^k` where +/// `start + 2^k <= upper_limit`. Offset 0 checks `start` itself — needed +/// when `start` is the only note. +/// +/// `upper_limit` caps the search range: +/// - First discovery (no prior knowledge): `u64::MAX` +/// - Re-probe after scan: `max_note_index * 2` (bounded growth) +/// +/// Max probes = 1 + floor(log2(upper_limit - start)) + 1. +/// +/// Issues exactly one batch of probes via `pool.get_notes_batch()`. Callers +/// handle iteration via cursor-based pagination. +pub async fn exponential_ascend( + pool: &S, + channel_key: Felt, + token: Felt, + start_index: u64, + upper_index_bound: u64, + budget: &IoBudget, +) -> Result { + let range = upper_index_bound.saturating_sub(start_index); + + // Offsets: [0, 1, 2, 4, 8, ..., 2^k] where 2^k <= range. + let offsets: Vec = std::iter::once(0) + .chain( + (0..64) + .map(|exp| 1u64 << exp) + .take_while(|&off| off <= range), + ) + .collect(); + + // Consume as many probes as budget allows (capped by batch_budget internally). + let batch_size = budget.consume_up_to(offsets.len(), COST_NOTE_PROBING); + if batch_size == 0 { + return Ok(ExponentialProbeResult { + last_found_note: None, + first_empty_index: None, + budget_exhausted: true, + }); + } + + let note_ids: Vec<_> = offsets[..batch_size] + .iter() + .map(|&off| compute_note_id(channel_key, token, start_index + off)) + .collect(); + let results = pool.get_notes_batch(¬e_ids).await?; + + let mut last_found_note: Option<(u64, Felt)> = None; + let mut first_empty_index: Option = None; + for (i, &packed) in results.iter().enumerate() { + let idx = start_index + offsets[i]; + if packed != Felt::ZERO { + last_found_note = Some((idx, packed)); + } else { + first_empty_index = Some(idx); + break; } } + + let budget_exhausted = first_empty_index.is_none() && batch_size < offsets.len(); + + Ok(ExponentialProbeResult { + last_found_note, + first_empty_index, + budget_exhausted, + }) } /// Binary search between `lo` (exists) and `hi` (absent) to find @@ -141,61 +194,98 @@ where #[cfg(test)] mod tests { use super::*; + use crate::privacy_pool::hashes::compute_note_id; + use crate::privacy_pool::storage_slots; use crate::storage_backend::MockBackend; use crate::test_fixtures::{get_channel_key, get_subchannel_token, load_devnet_fixture}; - /// Helper: create a probe function that returns exists for indices < last_index. - fn mock_probe( - last_index: Option, - ) -> impl Fn(u64) -> std::future::Ready, bool), DiscoveryError>> { - move |idx| { - let exists = last_index.is_some_and(|last| idx <= last); - std::future::ready(Ok(if exists { - (Some(idx), false) - } else { - (None, false) - })) + const CK: Felt = Felt::from_hex_unchecked("0x12345"); + const TK: Felt = Felt::from_hex_unchecked("0x67890"); + + /// Creates a mock backend with notes at indices 0..=last_index. + fn mock_with_notes(last_index: Option) -> MockBackend { + let mut backend = MockBackend::empty(); + if let Some(last) = last_index { + for i in 0..=last { + let note_id = compute_note_id(CK, TK, i); + let slot = storage_slots::notes(note_id); + backend.insert(slot, Felt::ONE); // non-zero = exists + } } + backend } #[tokio::test] async fn test_exponential_ascend_empty() { - let mut cursor = SubchannelCursor::default(); - let found = exponential_ascend(mock_probe(None), &mut cursor) + let backend = mock_with_notes(None); + let budget = IoBudget::new(100); + let result = exponential_ascend(&backend, CK, TK, 0, u64::MAX, &budget) .await .unwrap(); - assert!(found, "should find sentinel immediately"); - assert_eq!(cursor.last_note_index, None); - assert_eq!(cursor.max_note_index, Some(0)); + + assert_eq!(result.last_found_note, None); + assert_eq!(result.first_empty_index, Some(0)); + assert!(!result.budget_exhausted); } #[tokio::test] async fn test_exponential_ascend_one_element() { - let mut cursor = SubchannelCursor::default(); - let found = exponential_ascend(mock_probe(Some(0)), &mut cursor) + let backend = mock_with_notes(Some(0)); + let budget = IoBudget::new(100); + let result = exponential_ascend(&backend, CK, TK, 0, u64::MAX, &budget) .await .unwrap(); - assert!(found); - assert_eq!(cursor.last_note_index, Some(0)); - assert_eq!(cursor.max_note_index, Some(1)); + + assert_eq!(result.last_found_note, Some((0, Felt::ONE))); + assert_eq!(result.first_empty_index, Some(1)); + assert!(!result.budget_exhausted); } #[tokio::test] async fn test_exponential_ascend_multiple_elements() { - // Elements at 0, 1, 2, 3, 4 (last_index = 4) - let mut cursor = SubchannelCursor::default(); - let found = exponential_ascend(mock_probe(Some(4)), &mut cursor) + // Elements at 0..=4. Probes: offset 0→0 (hit), 1→1 (hit), 2→2 (hit), + // 4→4 (hit), 8→8 (miss) + let backend = mock_with_notes(Some(4)); + let budget = IoBudget::new(100); + let result = exponential_ascend(&backend, CK, TK, 0, u64::MAX, &budget) .await .unwrap(); - assert!(found); - // Probes: 0 (exists), 1 (exists), 3 (exists), 7 (empty) - assert_eq!(cursor.last_note_index, Some(3)); - assert_eq!(cursor.max_note_index, Some(7)); + + assert_eq!(result.last_found_note, Some((4, Felt::ONE))); + assert_eq!(result.first_empty_index, Some(8)); + assert!(!result.budget_exhausted); + } + + #[tokio::test] + async fn test_exponential_ascend_budget_exhausted() { + let backend = mock_with_notes(Some(100)); + // Budget for only 2 probes: offsets 0, 1 + let budget = IoBudget::new(2 * COST_NOTE_PROBING); + let result = exponential_ascend(&backend, CK, TK, 0, u64::MAX, &budget) + .await + .unwrap(); + + assert_eq!(result.last_found_note, Some((1, Felt::ONE))); + assert_eq!(result.first_empty_index, None); + assert!(result.budget_exhausted); + } + + /// Helper: mock bisection probe. + fn mock_probe( + last_index: Option, + ) -> impl Fn(u64) -> std::future::Ready, bool), DiscoveryError>> { + move |idx| { + let exists = last_index.is_some_and(|last| idx <= last); + std::future::ready(Ok(if exists { + (Some(idx), false) + } else { + (None, false) + })) + } } #[tokio::test] async fn test_bisect_boundary_adjacent() { - // lo=0, hi=1 → no bisection needed let mut cursor = SubchannelCursor { last_note_index: Some(0), max_note_index: Some(1), @@ -209,7 +299,6 @@ mod tests { #[tokio::test] async fn test_bisect_boundary_gap() { - // lo=3, hi=7, actual last is 4 let mut cursor = SubchannelCursor { last_note_index: Some(3), max_note_index: Some(7), @@ -289,7 +378,8 @@ mod tests { let mut cursor = SubchannelCursor::default(); - // Budget for 1 probe: finds note at 0, exhausted before probing 1 + // Budget for 1 probe: batch gets offset 0 only (note at 0 exists), + // but no first_empty found — budget_exhausted. let budget = IoBudget::new(COST_NOTE_PROBING); let (last_index, has_more) = find_last_note_index_paginated(&backend, channel_key, token, &mut cursor, &budget) diff --git a/crates/discovery-core/src/discovery/notes.rs b/crates/discovery-core/src/discovery/notes.rs index d3f50cab8..f5e09b24a 100644 --- a/crates/discovery-core/src/discovery/notes.rs +++ b/crates/discovery-core/src/discovery/notes.rs @@ -1,16 +1,13 @@ //! Notes discovery for a subchannel (channel_key + token pair). //! -//! This module provides functionality to discover and decrypt notes -//! within a specific subchannel. +//! Two-phase algorithm: +//! 1. **Exponential probe**: batched probes at exponentially increasing indices +//! to find `max_note_index` (the last index confirmed to exist). +//! 2. **Linear scan**: batch-reads amounts + nullifiers for all notes in +//! `start..=max_note_index`. No sentinel checking — the contiguous invariant +//! guarantees all notes in range exist. //! -//! # Sequential vs Parallel Scanning -//! -//! Notes within a subchannel are scanned sequentially because: -//! - Request budget is limited and pagination is cursor-based -//! - Parallel requests may not all complete within budget -//! - Sequential scanning avoids gaps in discovered indices -//! -//! However, parallelization is possible at higher levels: +//! Parallelization is possible at higher levels: //! - Multiple subchannel scans (for notes) can run in parallel //! - Multiple channel scans (for subchannels) can run in parallel @@ -18,8 +15,9 @@ use serde::{Deserialize, Serialize}; use starknet_types_core::felt::Felt; use super::cursor::SubchannelCursor; +use super::last_note_index::exponential_ascend; use super::DiscoveryError; -use super::COST_NOTE; +use super::{COST_NOTE, COST_NOTE_PROBING}; use crate::io_budget::IoBudget; use crate::privacy_pool::decryption::{decrypt_note_amount, unpack_note_amount}; use crate::privacy_pool::hashes::{compute_note_id, compute_nullifier}; @@ -49,123 +47,204 @@ pub struct NotesDiscoveryResult { /// `cursor.last_note_index = result.last_index`. pub last_index: Option, /// Whether there may be more notes to discover. - /// `true` if stopped due to budget exhaustion, `false` if sentinel was found. + /// `true` if stopped due to budget exhaustion, `false` if all notes scanned. pub has_more: bool, } -/// Discovers and decrypts notes for a given channel key and token. -/// -/// # Algorithm -/// -/// For each note index starting from `start_index`: -/// 1. Compute `note_id = hash(NOTE_ID_TAG, channel_key, token, index, 0)` -/// 2. Fetch `packed_amount` from storage -/// 3. If `packed_amount == 0`, stop (sentinel - no more notes) -/// 4. Decrypt to get `(amount, salt)` -/// 5. Compute nullifier and check existence — skip spent notes +/// Discovers notes with cursor-based pagination using batched RPC calls. /// -/// # Arguments +/// Two-phase algorithm: +/// 1. **Exponential probe** (when `max_note_index` is unknown or scan caught up): +/// batched probes to find the last existing note index. +/// 2. **Linear scan** (when `last_note_index < max_note_index`): batch-reads +/// amounts + nullifiers for notes in range. /// -/// * `privacy_pool` - Storage backend implementing the IViews trait. -/// * `channel_key` - The channel key. -/// * `token` - The token address for this subchannel. -/// * `start_index` - Starting index (inclusive). For incremental discovery, pass -/// `last_index + 1` from previous result. -/// * `decryption_key` - The owner's private (viewing) key for nullifier computation. -/// * `budget` - I/O budget to limit storage operations. -/// -/// # Returns +/// `max_note_index` is kept after scan — used to bound the next exponential +/// probe range. Re-probe triggers when `last_note_index == max_note_index`. +pub async fn discover_notes_paginated( + pool: &S, + channel_key: Felt, + token: Felt, + cursor: &mut SubchannelCursor, + decryption_key: &SecretFelt, + budget: &IoBudget, +) -> Result<(Vec, bool), DiscoveryError> { + let start_index = cursor.last_note_index.map_or(0, |i| i + 1); + + let need_probe = + cursor.max_note_index.is_none() || cursor.max_note_index.is_some_and(|m| start_index >= m); + + // Phase 1: Exponential probe + if need_probe { + let upper_index_bound = match cursor.max_note_index { + None => u64::MAX, + Some(m) => m.saturating_mul(2).max(start_index), + }; + let result = exponential_ascend( + pool, + channel_key, + token, + start_index, + upper_index_bound, + budget, + ) + .await?; + + let Some((last_found_index, last_found_packed_amount)) = result.last_found_note else { + // No notes found at start_index or beyond or out of budget. + return Ok((Vec::new(), result.budget_exhausted)); + }; + + cursor.max_note_index = Some(last_found_index); + + // Single-note case: probe confirmed only start_index exists. + // Packed amount is already available — only a nullifier check is needed. + if last_found_index == start_index { + if !budget.consume(COST_NOTE_PROBING) { + // Budget exhausted before nullifier check. + return Ok((Vec::new(), true)); + } + let note = scan_single_note( + pool, + channel_key, + token, + start_index, + last_found_packed_amount, + decryption_key, + ) + .await?; + cursor.last_note_index = Some(start_index); + // If the note is unspent, return it. + return Ok((note.into_iter().collect(), false)); + } + } + + let Some(max_note_index) = cursor.max_note_index else { + // No max_note_index — nothing to scan. + return Ok((Vec::new(), false)); + }; + + // Phase 2: Linear batch scan + if start_index > max_note_index { + // Scan already past max — nothing to do. + return Ok((Vec::new(), false)); + } + + let result = discover_notes( + pool, + channel_key, + token, + start_index, + max_note_index, + decryption_key, + budget, + ) + .await?; + + if let Some(last_index) = result.last_index { + cursor.last_note_index = Some(last_index); + } + Ok((result.notes, result.has_more)) +} + +/// Linear scan of notes in `start_index..=end_index`. /// -/// A `NotesDiscoveryResult` containing all discovered notes and metadata -/// for incremental discovery. -// TODO: Iterate in batches doing multiple RPC requests under the hood -// (get_note + nullifier_exists per note), until reaching a non-existent note, -// then bisect to find the exact boundary. -pub async fn discover_notes( - privacy_pool: &PrivacyPool, +/// Reads amounts + nullifiers in batches. No sentinel checking — the contiguous +/// invariant guarantees all notes in range exist. +// TODO: Consider fetching nullifiers first, then only fetch unspent note amounts +async fn discover_notes( + pool: &S, channel_key: Felt, token: Felt, start_index: u64, + end_index: u64, decryption_key: &SecretFelt, budget: &IoBudget, ) -> Result { - let mut notes = Vec::new(); let mut index = start_index; - let mut out_of_budget = false; + let mut notes = Vec::new(); let mut last_scanned_index: Option = None; loop { - // Consume budget for get_note + nullifier_exists - if !budget.consume(COST_NOTE) { - out_of_budget = true; - break; + if index > end_index { + return Ok(NotesDiscoveryResult { + notes, + last_index: last_scanned_index, + has_more: false, + }); } - let note_id = compute_note_id(channel_key, token, index); - let packed_amount = privacy_pool.get_note(note_id).await?; - - // Sentinel: contract stores zero for non-existent notes - if packed_amount == Felt::ZERO { - break; + let remaining = usize::try_from(end_index - index + 1).unwrap_or(usize::MAX); + let batch_size = budget.consume_up_to(remaining, COST_NOTE); + if batch_size == 0 { + return Ok(NotesDiscoveryResult { + notes, + last_index: last_scanned_index, + has_more: true, + }); } - let (salt, enc_amount) = unpack_note_amount(packed_amount); - - // TODO: Open notes (salt == 1) store the amount in plaintext, - // so enc_amount is already the actual amount - no decryption needed. - let amount = decrypt_note_amount(enc_amount, salt, channel_key, token, index); - - let nullifier = compute_nullifier(channel_key, token, index, decryption_key); - let is_spent = privacy_pool.nullifier_exists(nullifier).await?; - - last_scanned_index = Some(index); - - if !is_spent { - notes.push(DecryptedNote { - index, - note_id, - amount, - salt, - }); + let batch_end = index + + u64::try_from(batch_size) + .map_err(|_| DiscoveryError::InvalidCursor("batch size overflow".into()))?; + + let note_ids: Vec<_> = (index..batch_end) + .map(|i| compute_note_id(channel_key, token, i)) + .collect(); + let batch_nullifiers: Vec<_> = (index..batch_end) + .map(|i| compute_nullifier(channel_key, token, i, decryption_key)) + .collect(); + + let (packed_amounts, nullifier_exists) = pool + .get_note_and_nullifier_batch(¬e_ids, &batch_nullifiers) + .await?; + + for (j, idx) in (index..batch_end).enumerate() { + last_scanned_index = Some(idx); + + if !nullifier_exists[j] { + let (salt, enc_amount) = unpack_note_amount(packed_amounts[j]); + // TODO: Open notes (salt == 1) store the amount in plaintext, + // so enc_amount is already the actual amount - no decryption needed. + let amount = decrypt_note_amount(enc_amount, salt, channel_key, token, idx); + notes.push(DecryptedNote { + index: idx, + note_id: note_ids[j], + amount, + salt, + }); + } } - index += 1; - } - Ok(NotesDiscoveryResult { - notes, - last_index: last_scanned_index, - has_more: out_of_budget, - }) + index = batch_end; + } } -/// Discovers notes with cursor-based pagination. +/// Processes a single note whose packed amount is already known from a probe. /// -/// Delegates to [`discover_notes`] and updates the cursor. -/// -/// Returns discovered notes and `has_more`. `has_more = false` means -/// the sentinel was found (no more notes in this subchannel). -pub async fn discover_notes_paginated( +/// Only needs a nullifier check (cost 1 instead of `COST_NOTE` = 2). +/// Returns `Some(note)` if unspent, `None` if spent. +async fn scan_single_note( pool: &S, channel_key: Felt, token: Felt, - cursor: &mut SubchannelCursor, + note_index: u64, + packed_amount: Felt, decryption_key: &SecretFelt, - budget: &IoBudget, -) -> Result<(Vec, bool), DiscoveryError> { - let start_index = cursor.last_note_index.map_or(0, |i| i + 1); - let result = discover_notes( - pool, - channel_key, - token, - start_index, - decryption_key, - budget, - ) - .await?; - - cursor.last_note_index = result.last_index.or(cursor.last_note_index); - - Ok((result.notes, result.has_more)) +) -> Result, DiscoveryError> { + let nullifier = compute_nullifier(channel_key, token, note_index, decryption_key); + if pool.nullifier_exists(nullifier).await? { + return Ok(None); + } + let (salt, enc_amount) = unpack_note_amount(packed_amount); + let amount = decrypt_note_amount(enc_amount, salt, channel_key, token, note_index); + let note_id = compute_note_id(channel_key, token, note_index); + Ok(Some(DecryptedNote { + index: note_index, + note_id, + amount, + salt, + })) } #[cfg(test)] @@ -174,6 +253,29 @@ mod tests { use crate::storage_backend::MockBackend; use crate::test_fixtures::{get_channel_key, get_subchannel_token, load_devnet_fixture}; + /// Helper: runs discover_notes_paginated with a fresh default cursor and + /// returns the result along with the cursor for inspection. + async fn discover_with_fresh_cursor( + backend: &MockBackend, + channel_key: Felt, + token: Felt, + decryption_key: &SecretFelt, + budget: &IoBudget, + ) -> (Vec, bool, SubchannelCursor) { + let mut cursor = SubchannelCursor::default(); + let (notes, has_more) = discover_notes_paginated( + backend, + channel_key, + token, + &mut cursor, + decryption_key, + budget, + ) + .await + .unwrap(); + (notes, has_more, cursor) + } + #[tokio::test] async fn test_discover_no_notes() { let backend = MockBackend::empty(); @@ -182,13 +284,11 @@ mod tests { let budget = IoBudget::new(100); let zero_key = SecretFelt::new(Felt::ZERO); - let result = discover_notes(&backend, channel_key, token, 0, &zero_key, &budget) - .await - .unwrap(); + let (notes, has_more, _cursor) = + discover_with_fresh_cursor(&backend, channel_key, token, &zero_key, &budget).await; - assert_eq!(result.notes.len(), 0); - assert_eq!(result.last_index, None); - assert!(!result.has_more); + assert_eq!(notes.len(), 0); + assert!(!has_more); } #[tokio::test] @@ -196,7 +296,6 @@ mod tests { let fixture = load_devnet_fixture(); let backend = MockBackend::new(fixture.slots); - // Discover channel -> subchannel -> notes let channel_key = get_channel_key( &backend, fixture.constants.alice_address, @@ -211,18 +310,17 @@ mod tests { let budget = IoBudget::new(100); let key = SecretFelt::new(fixture.constants.alice_viewing_key); - let result = discover_notes(&backend, channel_key, token, 0, &key, &budget) - .await - .unwrap(); + let (notes, has_more, cursor) = + discover_with_fresh_cursor(&backend, channel_key, token, &key, &budget).await; // Alice deposited 100 STRK, transferred 50 to Bob. // The transfer consumed the deposit and wrote a change note (50 STRK) // at index 0. This note is unspent. - assert_eq!(result.notes.len(), 1, "1 unspent change note"); - assert_eq!(result.notes[0].index, 0); - assert!(result.notes[0].amount > 0, "Note amount should be positive"); - assert_eq!(result.last_index, Some(0)); - assert!(!result.has_more); + assert_eq!(notes.len(), 1, "1 unspent change note"); + assert_eq!(notes[0].index, 0); + assert!(notes[0].amount > 0, "Note amount should be positive"); + assert_eq!(cursor.last_note_index, Some(0)); + assert!(!has_more); } #[tokio::test] @@ -230,7 +328,6 @@ mod tests { let fixture = load_devnet_fixture(); let backend = MockBackend::new(fixture.slots); - // Discover Bob's incoming channel let channel_key = get_channel_key( &backend, fixture.constants.bob_address, @@ -245,14 +342,13 @@ mod tests { let budget = IoBudget::new(100); let key = SecretFelt::new(fixture.constants.bob_viewing_key); - let result = discover_notes(&backend, channel_key, token, 0, &key, &budget) - .await - .unwrap(); + let (notes, has_more, cursor) = + discover_with_fresh_cursor(&backend, channel_key, token, &key, &budget).await; // Bob withdrew his 50 STRK note → nullifier exists → filtered out - assert_eq!(result.notes.len(), 0, "Bob's note is spent"); - assert_eq!(result.last_index, Some(0), "note 0 was scanned"); - assert!(!result.has_more); + assert_eq!(notes.len(), 0, "Bob's note is spent"); + assert_eq!(cursor.last_note_index, Some(0), "note 0 was scanned"); + assert!(!has_more); } #[tokio::test] @@ -260,7 +356,6 @@ mod tests { let fixture = load_devnet_fixture(); let backend = MockBackend::new(fixture.slots); - // Discover Alice's channel and subchannel let channel_key = get_channel_key( &backend, fixture.constants.alice_address, @@ -276,20 +371,21 @@ mod tests { // First discovery — Alice has 1 unspent change note (50 STRK at index 0) let budget = IoBudget::new(100); let key = SecretFelt::new(fixture.constants.alice_viewing_key); - let result1 = discover_notes(&backend, channel_key, token, 0, &key, &budget) - .await - .unwrap(); - assert_eq!(result1.notes.len(), 1, "1 unspent change note"); - assert!(!result1.has_more); - let last_index = result1.last_index.unwrap(); + let mut cursor = SubchannelCursor::default(); + let (notes, has_more) = + discover_notes_paginated(&backend, channel_key, token, &mut cursor, &key, &budget) + .await + .unwrap(); + assert_eq!(notes.len(), 1, "1 unspent change note"); + assert!(!has_more); - // Incremental discovery starting from last_index + 1 - should find 0 new notes - let result2 = discover_notes(&backend, channel_key, token, last_index + 1, &key, &budget) - .await - .unwrap(); - assert_eq!(result2.notes.len(), 0); - assert_eq!(result2.last_index, None); - assert!(!result2.has_more); + // Incremental discovery — should find 0 new notes (sentinel at index 1) + let (notes2, has_more2) = + discover_notes_paginated(&backend, channel_key, token, &mut cursor, &key, &budget) + .await + .unwrap(); + assert_eq!(notes2.len(), 0); + assert!(!has_more2); } #[tokio::test] @@ -297,7 +393,6 @@ mod tests { let fixture = load_devnet_fixture(); let backend = MockBackend::new(fixture.slots); - // Discover Alice's channel and subchannel let channel_key = get_channel_key( &backend, fixture.constants.alice_address, @@ -310,16 +405,33 @@ mod tests { .await .expect("Alice's channel should have at least one subchannel"); - // Budget exhausted before starting (COST_NOTE = 2) + // Budget exhausted before starting let budget = IoBudget::new(0); let key = SecretFelt::new(fixture.constants.alice_viewing_key); - let result = discover_notes(&backend, channel_key, token, 0, &key, &budget) - .await - .unwrap(); + let (notes, has_more, _cursor) = + discover_with_fresh_cursor(&backend, channel_key, token, &key, &budget).await; + + assert_eq!(notes.len(), 0); + assert!(has_more); + } + + #[tokio::test] + async fn test_exponential_probe_empty_subchannel() { + // Empty subchannel: exponential probe finds sentinel at offset 0. + // With batch_budget=16, the probe batch includes 1 probe (offset 0 = empty). + // Cost = 1 (single probe). + let backend = MockBackend::empty(); + let channel_key = Felt::from_hex_unchecked("0x12345"); + let token = Felt::from_hex_unchecked("0x67890"); + let budget = IoBudget::new(100).with_batch_budget(1); - assert_eq!(result.notes.len(), 0); - assert_eq!(result.last_index, None); - assert!(result.has_more); + let zero_key = SecretFelt::new(Felt::ZERO); + let (notes, has_more, _cursor) = + discover_with_fresh_cursor(&backend, channel_key, token, &zero_key, &budget).await; + + assert_eq!(notes.len(), 0); + assert!(!has_more, "sentinel found, not budget exhaustion"); + assert_eq!(budget.remaining(), 99); } #[tokio::test] @@ -349,6 +461,44 @@ mod tests { assert!(!has_more, "sentinel should be found"); } + #[tokio::test] + async fn test_budget_exhaustion_then_resume_skips_initial_probe() { + let fixture = load_devnet_fixture(); + let backend = MockBackend::new(fixture.slots); + + let channel_key = get_channel_key( + &backend, + fixture.constants.alice_address, + &fixture.constants.alice_viewing_key, + ) + .await + .unwrap(); + let token = get_subchannel_token(&backend, channel_key).await.unwrap(); + + let key = SecretFelt::new(fixture.constants.alice_viewing_key); + let mut cursor = SubchannelCursor::default(); + + // First call with enough budget to discover notes. + let budget = IoBudget::new(100); + let (notes, has_more) = + discover_notes_paginated(&backend, channel_key, token, &mut cursor, &key, &budget) + .await + .unwrap(); + assert_eq!(notes.len(), 1); + assert!(!has_more); + + // Now resume — cursor.last_note_index is Some(0), so start_index = 1. + // The initial probe should not be skipped (no cached probe at index 1). + // It should find sentinel at index 1 and return immediately. + let budget2 = IoBudget::new(100); + let (notes2, has_more2) = + discover_notes_paginated(&backend, channel_key, token, &mut cursor, &key, &budget2) + .await + .unwrap(); + assert_eq!(notes2.len(), 0); + assert!(!has_more2); + } + #[tokio::test] async fn test_paginated_budget_limited() { let fixture = load_devnet_fixture(); @@ -365,16 +515,20 @@ mod tests { let mut cursor = SubchannelCursor::default(); let key = SecretFelt::new(fixture.constants.bob_viewing_key); - // Budget for exactly 1 note (COST_NOTE=2: get_note + nullifier_exists) - let budget = IoBudget::new(COST_NOTE); + // Bob has 1 note at index 0. With batch_budget=2: + // Exponential probe: 2 probes (offsets 0, 1) → hit at 0, miss at 1. + // Cost = 2 * COST_NOTE_PROBING = 2. + // Single-note optimization: 1 nullifier check = COST_NOTE_PROBING = 1. + // Total = 3. + let budget = IoBudget::new(3).with_batch_budget(2); let (notes, has_more) = discover_notes_paginated(&backend, channel_key, token, &mut cursor, &key, &budget) .await .unwrap(); - // Bob's note 0 is spent → filtered out, but still scanned + // Bob's note 0 is spent → filtered out assert_eq!(notes.len(), 0, "Bob's note is spent"); - assert!(has_more, "sentinel not reached"); + assert!(!has_more, "single-note optimization handled it"); assert_eq!(cursor.last_note_index, Some(0)); } } diff --git a/crates/discovery-core/src/io_budget.rs b/crates/discovery-core/src/io_budget.rs index 143dbc665..ef944fe66 100644 --- a/crates/discovery-core/src/io_budget.rs +++ b/crates/discovery-core/src/io_budget.rs @@ -7,6 +7,9 @@ use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::Arc; +/// Default maximum budget units per batch in discovery operations. +const DEFAULT_BATCH_BUDGET: usize = 16; + /// Thread-safe I/O budget counter. /// /// Used to limit the number of storage operations during discovery. @@ -17,16 +20,30 @@ use std::sync::Arc; #[derive(Debug, Clone)] pub struct IoBudget { remaining: Arc, + /// Maximum budget units per batch in discovery operations. + batch_budget: usize, } impl IoBudget { - /// Creates a new budget with the given limit. + /// Creates a new budget with the given limit and default batch budget (16). pub fn new(limit: usize) -> Self { Self { remaining: Arc::new(AtomicUsize::new(limit)), + batch_budget: DEFAULT_BATCH_BUDGET, } } + /// Sets the batch budget (max budget units per batch). Returns `self` for chaining. + pub fn with_batch_budget(mut self, batch_budget: usize) -> Self { + self.batch_budget = batch_budget; + self + } + + /// Returns the batch budget (max budget units per batch). + pub fn batch_budget(&self) -> usize { + self.batch_budget + } + /// Returns the current remaining budget. pub fn remaining(&self) -> usize { self.remaining.load(Ordering::Relaxed) @@ -43,6 +60,33 @@ impl IoBudget { }) .is_ok() } + + /// Atomically consumes as many whole items as the budget allows. + /// + /// Returns the number of items consumed (0..=`max_items`). + /// Each item costs `cost_per_item` units. Capped by `batch_budget / cost_per_item`. + /// Returns 0 if `cost_per_item == 0`, `max_items == 0`, or the budget is + /// insufficient for even one item. + pub fn consume_up_to(&self, max_items: usize, cost_per_item: usize) -> usize { + if cost_per_item == 0 || max_items == 0 { + return 0; + } + let cap = max_items.min(self.batch_budget / cost_per_item); + if cap == 0 { + return 0; + } + self.remaining + .fetch_update(Ordering::SeqCst, Ordering::SeqCst, |current| { + let n = (current / cost_per_item).min(cap); + if n == 0 { + None + } else { + Some(current - n * cost_per_item) + } + }) + .map(|old| (old / cost_per_item).min(cap)) + .unwrap_or(0) + } } #[cfg(test)] @@ -90,6 +134,67 @@ mod tests { assert_eq!(budget.remaining(), 0); } + #[test] + fn test_consume_up_to_exact_budget() { + let budget = IoBudget::new(9); + // 9 / 3 = 3 items, but max is 2 + assert_eq!(budget.consume_up_to(2, 3), 2); + assert_eq!(budget.remaining(), 3); + // 3 / 3 = 1 item + assert_eq!(budget.consume_up_to(5, 3), 1); + assert_eq!(budget.remaining(), 0); + } + + #[test] + fn test_consume_up_to_partial_budget() { + let budget = IoBudget::new(7); + // 7 / 3 = 2 items (with 1 leftover) + assert_eq!(budget.consume_up_to(10, 3), 2); + assert_eq!(budget.remaining(), 1); + // 1 / 3 = 0 items + assert_eq!(budget.consume_up_to(10, 3), 0); + assert_eq!(budget.remaining(), 1); // unchanged + } + + #[test] + fn test_consume_up_to_zero_budget() { + let budget = IoBudget::new(0); + assert_eq!(budget.consume_up_to(5, 3), 0); + assert_eq!(budget.remaining(), 0); + } + + #[test] + fn test_consume_up_to_zero_cost() { + let budget = IoBudget::new(10); + assert_eq!(budget.consume_up_to(5, 0), 0); + assert_eq!(budget.remaining(), 10); // unchanged + } + + #[test] + fn test_consume_up_to_zero_max_items() { + let budget = IoBudget::new(10); + assert_eq!(budget.consume_up_to(0, 3), 0); + assert_eq!(budget.remaining(), 10); // unchanged + } + + #[test] + fn test_consume_up_to_concurrent() { + let budget = IoBudget::new(100); + let mut handles = vec![]; + + // 10 threads each trying to consume up to 5 items at cost 2 = 10 per thread + for _ in 0..10 { + let budget = budget.clone(); + handles.push(thread::spawn(move || budget.consume_up_to(5, 2))); + } + + let total: usize = handles.into_iter().map(|h| h.join().unwrap()).sum(); + + // 100 / 2 = 50 items total possible + assert_eq!(total, 50); + assert_eq!(budget.remaining(), 0); + } + #[test] fn test_concurrent_consume() { let budget = IoBudget::new(1000); diff --git a/crates/discovery-core/src/privacy_pool/views.rs b/crates/discovery-core/src/privacy_pool/views.rs index 8eff11a40..5d09267a9 100644 --- a/crates/discovery-core/src/privacy_pool/views.rs +++ b/crates/discovery-core/src/privacy_pool/views.rs @@ -53,6 +53,31 @@ pub trait IViews: Send + Sync { /// Returns the compliance public key. async fn get_compliance_public_key(&self) -> Result; + + /// Batch-reads channel info for `count` consecutive channels starting at `start_index`. + /// + /// Returns a `Vec` of length `count`, fetched in a single `read_slots` call. + async fn get_channel_info_batch( + &self, + recipient_addr: Felt, + start_index: u64, + count: usize, + ) -> Result, StorageError>; + + /// Batch-reads packed note values for the given note IDs. + /// + /// Returns a `Vec` matching the input length. Zero = note doesn't exist. + async fn get_notes_batch(&self, note_ids: &[Felt]) -> Result, StorageError>; + + /// Batch-reads packed note amounts and nullifier existence. + /// + /// Returns `(packed_amounts, nullifier_exists)`. + /// Both vectors match the lengths of their respective inputs. + async fn get_note_and_nullifier_batch( + &self, + note_ids: &[Felt], + nullifiers: &[Felt], + ) -> Result<(Vec, Vec), StorageError>; } /// Checks that `values` has exactly `expected` elements, returning @@ -182,4 +207,85 @@ impl IViews for T { let slot = storage_slots::compliance_public_key(); self.read_slot(slot).await } + + #[tracing::instrument( + name = "get_channel_info_batch", + level = "debug", + skip(self), + fields(count) + )] + async fn get_channel_info_batch( + &self, + recipient_addr: Felt, + start_index: u64, + count: usize, + ) -> Result, StorageError> { + let mut slots = Vec::with_capacity(count * 3); + for i in 0..count { + let idx = start_index + i as u64; + let s = storage_slots::recipient_channels_element(recipient_addr, idx); + slots.push(s.ephemeral_pubkey); + slots.push(s.enc_channel_key); + slots.push(s.enc_sender_addr); + } + let values = self.read_slots(slots).await?; + check_slots_len(&values, count * 3)?; + let mut result = Vec::with_capacity(count); + for chunk in values.chunks_exact(3) { + result.push(EncChannelInfo { + ephemeral_pubkey: chunk[0], + enc_channel_key: chunk[1], + enc_sender_addr: chunk[2], + }); + } + Ok(result) + } + + #[tracing::instrument( + name = "get_notes_batch", + level = "debug", + skip(self, note_ids), + fields(count = note_ids.len()) + )] + async fn get_notes_batch(&self, note_ids: &[Felt]) -> Result, StorageError> { + let slots: Vec<_> = note_ids + .iter() + .map(|&nid| storage_slots::notes(nid)) + .collect(); + let values = self.read_slots(slots).await?; + check_slots_len(&values, note_ids.len())?; + Ok(values) + } + + #[tracing::instrument( + name = "get_note_and_nullifier_batch", + level = "debug", + skip(self, note_ids, nullifiers) + )] + async fn get_note_and_nullifier_batch( + &self, + note_ids: &[Felt], + nullifiers: &[Felt], + ) -> Result<(Vec, Vec), StorageError> { + let n_notes = note_ids.len(); + let n_nullifiers = nullifiers.len(); + let total = n_notes + n_nullifiers; + + let mut slots = Vec::with_capacity(total); + for &nid in note_ids { + slots.push(storage_slots::notes(nid)); + } + for &nul in nullifiers { + slots.push(storage_slots::nullifiers(nul)); + } + + let values = self.read_slots(slots).await?; + check_slots_len(&values, total)?; + + let packed_amounts = values[..n_notes].to_vec(); + let nullifier_exists: Vec = + values[n_notes..].iter().map(|v| *v != Felt::ZERO).collect(); + + Ok((packed_amounts, nullifier_exists)) + } } diff --git a/crates/discovery-core/src/sync/incoming_state.rs b/crates/discovery-core/src/sync/incoming_state.rs index 2982b99e8..d048efda5 100644 --- a/crates/discovery-core/src/sync/incoming_state.rs +++ b/crates/discovery-core/src/sync/incoming_state.rs @@ -3,13 +3,13 @@ //! Composes paginated channel, subchannel, and note discovery into a //! single [`sync_incoming_state`] call that advances a [`DiscoveryCursor`]. //! -//! Channel and subchannel processing is parallelised via [`tokio::task::JoinSet`]. +//! Channel and subchannel processing is concurrent via [`FuturesUnordered`]. use std::collections::HashMap; +use futures::stream::{FuturesUnordered, StreamExt}; use serde::{Deserialize, Serialize}; use starknet_types_core::felt::Felt; -use tokio::task::JoinSet; use crate::discovery::cursor::{ChannelCursor, DiscoveryCursor, SubchannelCursor}; use crate::discovery::incoming_channels::discover_incoming_channels_paginated; @@ -62,20 +62,17 @@ struct SubchannelResult { /// Each level uses cursor-based pagination. Fully-discovered subchannels /// and channels are pruned from the cursor so subsequent calls skip them. /// -/// Channel and subchannel processing is parallelised via [`JoinSet`]. +/// Channel and subchannel processing is concurrent via [`FuturesUnordered`]. // TODO: Handle open notes — notes with salt==OPEN_NOTE_SALT(1) have plaintext // amounts, non-zero token and depositor fields (see Cairo objects.cairo // Note struct) -pub async fn sync_incoming_state( +pub async fn sync_incoming_state( pool: &S, recipient: Felt, decryption_key: &SecretFelt, mut cursor: DiscoveryCursor, budget: &IoBudget, -) -> Result -where - S: IViews + Clone + Send + Sync + 'static, -{ +) -> Result { discover_incoming_channels_paginated(pool, recipient, decryption_key, &mut cursor, budget) .await?; @@ -86,25 +83,21 @@ where }); } - // TODO(security): Cap cursor.channels size before spawning tasks — the - // HashMap is deserialized from an untrusted request with no size limit. - // An attacker can send 50K+ entries within the 2MB body limit, each - // spawning a tokio task → OOM / scheduler exhaustion. + // TODO(security): Cap cursor.channels size — the HashMap is deserialized + // from an untrusted request with no size limit. An attacker can send 50K+ + // entries within the 2MB body limit → OOM. // Fix: reject or truncate cursor.channels to a max size (e.g., 256). let channels = std::mem::take(&mut cursor.channels); - let mut join_set = JoinSet::new(); - for (sender_addr, ch_cursor) in channels { - let pool = pool.clone(); - let budget = budget.clone(); - let key = decryption_key.clone(); - join_set.spawn(async move { - process_channel(&pool, sender_addr, ch_cursor, &key, &budget).await - }); - } + let mut futs: FuturesUnordered<_> = channels + .into_iter() + .map(|(sender_addr, ch_cursor)| { + process_channel(pool, sender_addr, ch_cursor, decryption_key, budget) + }) + .collect(); let mut channels_output: HashMap = HashMap::new(); - while let Some(join_result) = join_set.join_next().await { - let result = join_result.map_err(|e| DiscoveryError::TaskPanicked(e.to_string()))??; + while let Some(result) = futs.next().await { + let result = result?; channels_output.insert(result.sender_addr, result.output); if let Some(ch_cursor) = result.cursor { cursor.channels.insert(result.sender_addr, ch_cursor); @@ -117,44 +110,34 @@ where }) } -/// Processes a single channel: discovers subchannels, then spawns note -/// discovery for each subchannel concurrently via [`JoinSet`]. -async fn process_channel( +/// Processes a single channel: discovers subchannels, then runs note +/// discovery for each subchannel concurrently via [`FuturesUnordered`]. +async fn process_channel( pool: &S, sender_addr: Felt, mut cursor: ChannelCursor, decryption_key: &SecretFelt, budget: &IoBudget, -) -> Result -where - S: IViews + Clone + Send + Sync + 'static, -{ +) -> Result { let channel_key = cursor.channel_key.ok_or_else(|| { DiscoveryError::InvalidCursor("channel_key is required for incoming channel".into()) })?; discover_subchannels_paginated(pool, channel_key, &mut cursor, budget).await?; - let subchannels: Vec<(Felt, SubchannelCursor)> = std::mem::take(&mut cursor.subchannels) + // TODO(security): Cap cursor.subchannels size — same unbounded-HashMap + // attack vector as cursor.channels, multiplied per channel. + let subchannels = std::mem::take(&mut cursor.subchannels); + let mut futs: FuturesUnordered<_> = subchannels .into_iter() + .map(|(token, sc_cursor)| { + process_subchannel(pool, channel_key, token, sc_cursor, decryption_key, budget) + }) .collect(); - // TODO(security): Cap cursor.subchannels size before spawning tasks — - // same unbounded-HashMap attack vector as cursor.channels above, - // multiplied per channel (N channels × M subchannels = N×M tasks). - let mut join_set = JoinSet::new(); - for (token, sc_cursor) in subchannels { - let pool = pool.clone(); - let budget = budget.clone(); - let key = decryption_key.clone(); - join_set.spawn(async move { - process_subchannel(&pool, channel_key, token, sc_cursor, &key, &budget).await - }); - } - let mut subchannel_notes: HashMap> = HashMap::new(); - while let Some(join_result) = join_set.join_next().await { - let result = join_result.map_err(|e| DiscoveryError::TaskPanicked(e.to_string()))??; + while let Some(result) = futs.next().await { + let result = result?; subchannel_notes.insert(result.token, result.notes); if let Some(sc_cursor) = result.cursor { cursor.subchannels.insert(result.token, sc_cursor); @@ -202,7 +185,9 @@ async fn process_subchannel( #[cfg(test)] mod tests { use super::*; - use crate::discovery::{COST_CHANNEL_INFO, COST_NOTE, COST_NUM_CHANNELS, COST_SUBCHANNEL_INFO}; + use crate::discovery::{ + COST_CHANNEL_INFO, COST_NOTE_PROBING, COST_NUM_CHANNELS, COST_SUBCHANNEL_INFO, + }; use crate::storage_backend::MockBackend; use crate::test_fixtures::load_devnet_fixture; @@ -220,8 +205,7 @@ mod tests { /// | 2 | 3 | Channel 0 discovered | /// | 3 | 2 | Subchannel 0 (STRK) discovered | /// | 4 | 2 | Subchannel sentinel → total cached | - /// | 5 | 1 | Subchannels skipped (cached) + note 0 discovered | - /// | 6 | 1 | Subchannels skipped + note sentinel → all done | + /// | 5 | 4 | Subchannels skipped + batched note scan → all done | #[tokio::test] async fn test_sync_incoming_state_step_by_step_pagination() { let f = load_devnet_fixture(); @@ -295,7 +279,7 @@ mod tests { // Step 4: budget = COST_SUBCHANNEL_INFO (2) // Reads subchannel index 1 → sentinel (salt=0). - // Budget=0, note discovery fails (needs COST_NOTE=1). + // Budget=0, note discovery fails (needs COST_NOTE_PROBING=1 for initial probe). let budget = IoBudget::new(COST_SUBCHANNEL_INFO); let out = sync_incoming_state(&backend, recipient, &decryption_key, out.cursor, &budget) .await @@ -312,18 +296,20 @@ mod tests { "step 4: subchannel still in cursor (notes not started)" ); - // Step 5: budget = COST_NOTE (2) - // Subchannels skipped (total cached). Discovers note 0 (1 get_note + - // 1 nullifier_exists). Note is spent → filtered out. Budget=0, can't - // check sentinel. - let budget = IoBudget::new(COST_NOTE); + // Step 5: budget = 2 * COST_NOTE_PROBING + COST_NOTE_PROBING (3) + // Subchannels skipped (total cached). Notes discovery: + // Exponential probe: 2 probes (offsets 0, 1) → hit at 0, miss at 1. Cost = 2. + // Single-note optimization: 1 nullifier check. Cost = 1. + // Bob's note is spent → filtered. Total = 3. + // batch_budget=2 limits the probe batch to 2 probes, leaving 1 for nullifier. + let budget = IoBudget::new(2 * COST_NOTE_PROBING + COST_NOTE_PROBING).with_batch_budget(2); let out = sync_incoming_state(&backend, recipient, &decryption_key, out.cursor, &budget) .await .unwrap(); assert!( - out.cursor.channels.contains_key(&sender_addr), - "step 5: channel still in cursor (note sentinel not checked)" + out.cursor.channels.is_empty(), + "step 5: cursor empty (all discovery complete)" ); let notes = &out.channels[&sender_addr].subchannels[&subchannel_token]; assert_eq!( @@ -331,24 +317,5 @@ mod tests { 0, "step 5: note discovered but spent (filtered out)" ); - - // Step 6: budget = COST_NOTE (2) - // Subchannels skipped (total cached). Reads note index 1 → sentinel - // (packed_amount=0 → break, no nullifier check). 1 budget unit unused. - // All done: subchannel + channel pruned from cursor. - let budget = IoBudget::new(COST_NOTE); - let out = sync_incoming_state(&backend, recipient, &decryption_key, out.cursor, &budget) - .await - .unwrap(); - - assert!( - out.cursor.channels.is_empty(), - "step 6: cursor empty (all discovery complete)" - ); - assert_eq!( - out.channels[&sender_addr].subchannels[&subchannel_token].len(), - 0, - "step 6: no new notes discovered" - ); } } diff --git a/crates/discovery-core/src/sync/outgoing_state.rs b/crates/discovery-core/src/sync/outgoing_state.rs index 20962c79b..4f399ecc1 100644 --- a/crates/discovery-core/src/sync/outgoing_state.rs +++ b/crates/discovery-core/src/sync/outgoing_state.rs @@ -3,13 +3,13 @@ //! Composes paginated outgoing channel, subchannel, and note-index discovery //! into a single [`sync_outgoing_state`] call that advances a [`DiscoveryCursor`]. //! -//! Channel and subchannel processing is parallelised via [`tokio::task::JoinSet`]. +//! Channel and subchannel processing is concurrent via [`FuturesUnordered`]. use std::collections::HashMap; +use futures::stream::{FuturesUnordered, StreamExt}; use serde::{Deserialize, Serialize}; use starknet_types_core::felt::Felt; -use tokio::task::JoinSet; use crate::discovery::cursor::{ChannelCursor, DiscoveryCursor, SubchannelCursor}; use crate::discovery::last_note_index::find_last_note_index_paginated; @@ -57,17 +57,14 @@ struct OutgoingSubchannelResult { /// Each level uses cursor-based pagination. Fully-discovered subchannels /// and channels are pruned from the cursor so subsequent calls skip them. /// -/// Channel and subchannel processing is parallelised via [`JoinSet`]. -pub async fn sync_outgoing_state( +/// Channel and subchannel processing is concurrent via [`FuturesUnordered`]. +pub async fn sync_outgoing_state( pool: &S, sender_addr: Felt, viewing_key: &SecretFelt, mut cursor: DiscoveryCursor, budget: &IoBudget, -) -> Result -where - S: IViews + Clone + Send + Sync + 'static, -{ +) -> Result { discover_outgoing_channels_paginated(pool, sender_addr, viewing_key, &mut cursor, budget) .await?; @@ -79,18 +76,16 @@ where } let channels = std::mem::take(&mut cursor.channels); - let mut join_set = JoinSet::new(); - for (recipient_addr, ch_cursor) in channels { - let pool = pool.clone(); - let budget = budget.clone(); - join_set.spawn(async move { - process_outgoing_channel(&pool, recipient_addr, ch_cursor, &budget).await - }); - } + let mut futs: FuturesUnordered<_> = channels + .into_iter() + .map(|(recipient_addr, ch_cursor)| { + process_outgoing_channel(pool, recipient_addr, ch_cursor, budget) + }) + .collect(); let mut channels_output: HashMap = HashMap::new(); - while let Some(join_result) = join_set.join_next().await { - let result = join_result.map_err(|e| DiscoveryError::TaskPanicked(e.to_string()))??; + while let Some(result) = futs.next().await { + let result = result?; channels_output.insert(result.recipient_addr, result.output); if let Some(ch_cursor) = result.cursor { cursor.channels.insert(result.recipient_addr, ch_cursor); @@ -103,39 +98,31 @@ where }) } -/// Processes a single outgoing channel: discovers subchannels, then spawns -/// note-index probing for each subchannel concurrently via [`JoinSet`]. -async fn process_outgoing_channel( +/// Processes a single outgoing channel: discovers subchannels, then runs +/// note-index probing for each subchannel concurrently via [`FuturesUnordered`]. +async fn process_outgoing_channel( pool: &S, recipient_addr: Felt, mut cursor: ChannelCursor, budget: &IoBudget, -) -> Result -where - S: IViews + Clone + Send + Sync + 'static, -{ +) -> Result { let channel_key = cursor.channel_key.ok_or_else(|| { DiscoveryError::InvalidCursor("channel_key is required for outgoing channel".into()) })?; discover_subchannels_paginated(pool, channel_key, &mut cursor, budget).await?; - let subchannels: Vec<(Felt, SubchannelCursor)> = std::mem::take(&mut cursor.subchannels) + let subchannels = std::mem::take(&mut cursor.subchannels); + let mut futs: FuturesUnordered<_> = subchannels .into_iter() + .map(|(token, sc_cursor)| { + process_outgoing_subchannel(pool, channel_key, token, sc_cursor, budget) + }) .collect(); - let mut join_set = JoinSet::new(); - for (token, sc_cursor) in subchannels { - let pool = pool.clone(); - let budget = budget.clone(); - join_set.spawn(async move { - process_outgoing_subchannel(&pool, channel_key, token, sc_cursor, &budget).await - }); - } - let mut subchannel_results: HashMap> = HashMap::new(); - while let Some(join_result) = join_set.join_next().await { - let result = join_result.map_err(|e| DiscoveryError::TaskPanicked(e.to_string()))??; + while let Some(result) = futs.next().await { + let result = result?; subchannel_results.insert(result.token, result.last_note_index); if let Some(sc_cursor) = result.cursor { cursor.subchannels.insert(result.token, sc_cursor);