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 f5f280739..1f6022851 100644 --- a/crates/discovery-core/Cargo.toml +++ b/crates/discovery-core/Cargo.toml @@ -14,6 +14,7 @@ starknet-providers = "0.16" tracing = "0.1" serde = { version = "1.0", features = ["derive"] } url = "2" +futures = "0.3" zeroize = "1" [dev-dependencies] diff --git a/crates/discovery-core/src/discovery/cursor.rs b/crates/discovery-core/src/discovery/cursor.rs new file mode 100644 index 000000000..7ed58e9f1 --- /dev/null +++ b/crates/discovery-core/src/discovery/cursor.rs @@ -0,0 +1,116 @@ +//! Cursor types for paginated discovery. +//! +//! These cursors track progress across paginated discovery calls, allowing +//! callers to resume discovery from where they left off. + +use std::collections::HashMap; + +use serde::{Deserialize, Serialize}; +use starknet_types_core::felt::Felt; + +/// Top-level cursor for channel discovery (shared by incoming and outgoing). +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +pub struct DiscoveryCursor { + /// All channels have been enumerated. Set by the discovery service once + /// the sentinel channel is reached. When `true`, no further channel + /// discovery is attempted — only channels already in the cursor are + /// processed. + #[serde(default)] + pub channel_discovery_complete: bool, + + /// Total number of channels (cached from `get_num_of_channels` for incoming). + /// Used as optimization to avoid redundant RPC calls. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub total_n_channels: Option, + + /// Last fully processed channel index. `None` = start from index 0. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub last_channel_index: Option, + + /// Channels with pending subchannel/note discovery. + /// - Incoming: keyed by sender address. + /// - Outgoing: keyed by recipient address. + #[serde(default, skip_serializing_if = "HashMap::is_empty")] + pub channels: HashMap, +} + +impl DiscoveryCursor { + /// Returns `true` when all discovery levels are complete: channels, + /// subchannels within each channel, and notes within each subchannel. + pub fn is_complete(&self) -> bool { + self.channel_discovery_complete && self.all_channels_processed() + } + + /// Returns `true` when every channel currently in the cursor has + /// completed subchannel and note discovery. Also returns `true` when + /// the cursor has no channels (vacuously). + /// + /// Used by sync orchestrators to decide whether to discover new + /// channels vs. process pending subchannel/note work. + pub fn all_channels_processed(&self) -> bool { + self.channels.values().all(ChannelCursor::is_complete) + } +} + +/// Cursor state for a single channel (shared by incoming and outgoing). +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct ChannelCursor { + // TODO: Consider encrypting/masking channel_key in the serialized cursor + // to avoid exposing it in plaintext (sensitive value). + /// The channel key for this channel. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub channel_key: Option, + + /// All subchannels have been enumerated. Set by the discovery service + /// once the sentinel subchannel is reached. When `true`, no further + /// subchannel discovery is attempted — only subchannels already in the + /// cursor are processed. + #[serde(default)] + pub subchannel_discovery_complete: bool, + + /// Last fully processed subchannel index. `None` = start from index 0. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub last_subchannel_index: Option, + + /// Subchannels with pending note discovery, keyed by token address. + #[serde(default, skip_serializing_if = "HashMap::is_empty")] + pub subchannels: HashMap, +} + +impl ChannelCursor { + /// Returns `true` when subchannel discovery is complete and all + /// subchannels have finished note discovery. + pub fn is_complete(&self) -> bool { + self.subchannel_discovery_complete + && self + .subchannels + .values() + .all(|sc| sc.note_discovery_complete) + } +} + +/// Cursor state for a single subchannel (shared by incoming and outgoing). +/// +/// For incoming (linear scan): only `last_note_index` is used. +/// For outgoing (exponential search): `last_note_index` = lo, `max_note_index` = hi. +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +pub struct SubchannelCursor { + /// All notes in this subchannel have been discovered. Set when the + /// note discovery scan completes without budget exhaustion. + #[serde(default)] + pub note_discovery_complete: bool, + + /// Last note index where a note exists. + /// - Incoming: last scanned index. + /// - Outgoing: lower bound (lo) for exponential search. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub last_note_index: Option, + + /// - 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/mod.rs b/crates/discovery-core/src/discovery/mod.rs index b1155481a..c352343da 100644 --- a/crates/discovery-core/src/discovery/mod.rs +++ b/crates/discovery-core/src/discovery/mod.rs @@ -5,21 +5,32 @@ use thiserror::Error; use crate::privacy_pool::decryption::DecryptionError; use crate::storage_backend::StorageError; +pub mod cursor; +pub use cursor::{ChannelCursor, DiscoveryCursor, SubchannelCursor}; pub mod incoming_channels; pub mod notes; pub mod subchannels; /// Cost for `get_num_of_channels` (1 storage slot read). -const COST_NUM_CHANNELS: usize = 1; +pub const COST_NUM_CHANNELS: usize = 1; /// Cost for `get_channel_info` (3 storage slot reads). -const COST_CHANNEL_INFO: usize = 3; +pub const COST_CHANNEL_INFO: usize = 3; /// Cost for `get_subchannel_info` (2 storage slot reads). -const COST_SUBCHANNEL_INFO: usize = 2; +pub const COST_SUBCHANNEL_INFO: usize = 2; -/// Cost for `get_note` (1 storage slot read). -const COST_NOTE: usize = 1; +/// Cost for `get_note` + `nullifier_exists` (2 storage slot reads). +pub const COST_NOTE: usize = 2; + +/// Cost for outgoing channel info (3 storage slot reads: salt + enc_recipient_addr + public_key). +pub const COST_OUTGOING_CHANNEL_INFO: usize = 3; + +/// Cost for a single note existence probe (1 `get_note` read, no nullifier check). +pub const COST_NOTE_PROBING: usize = 1; + +/// Cost for a single `get_public_key` (1 storage slot read). +pub const COST_PUBLIC_KEY: usize = 1; /// Errors that can occur during channel discovery. #[derive(Debug, Error)] @@ -36,4 +47,10 @@ pub enum DiscoveryError { #[source] source: DecryptionError, }, + /// A spawned task panicked. + #[error("spawned task panicked: {0}")] + TaskPanicked(String), + /// Invalid cursor data provided by client. + #[error("invalid cursor: {0}")] + InvalidCursor(String), } diff --git a/crates/discovery-core/src/discovery/notes.rs b/crates/discovery-core/src/discovery/notes.rs index 8b9d01bbc..89fdd9a93 100644 --- a/crates/discovery-core/src/discovery/notes.rs +++ b/crates/discovery-core/src/discovery/notes.rs @@ -267,7 +267,7 @@ mod tests { .await .expect("Alice's channel should have at least one subchannel"); - // Budget exhausted before starting (COST_NOTE = 1) + // Budget exhausted before starting (COST_NOTE = 2) let budget = IoBudget::new(0); let result = discover_notes(&backend, channel_key, token, 0, &budget) .await diff --git a/crates/discovery-core/src/io_budget.rs b/crates/discovery-core/src/io_budget.rs index 143dbc665..1ed74971c 100644 --- a/crates/discovery-core/src/io_budget.rs +++ b/crates/discovery-core/src/io_budget.rs @@ -43,6 +43,41 @@ impl IoBudget { }) .is_ok() } + + /// Atomically consumes as many whole items as the budget allows. + /// + /// Returns `(num_consumed_items, budget_exhausted)`: + /// - `num_consumed_items`: number of items consumed (0..=`max_items`). + /// - `budget_exhausted`: `true` when the budget limited how many items could + /// be consumed (i.e. `consumed < max_items`). + /// Callers use this to derive `has_more` for pagination cursors. + /// + /// When `max_items == 0`, returns `(0, false)` — the request is trivially + /// satisfied, not budget-limited. + /// + /// # Panics + /// + /// Panics if `cost_per_item == 0`. All callers use compile-time cost constants; + /// a zero cost is always a programmer bug. + pub fn consume_up_to(&self, max_items: usize, cost_per_item: usize) -> (usize, bool) { + assert!( + cost_per_item > 0, + "cost_per_item must be positive; zero-cost items are not supported" + ); + if max_items == 0 { + return (0, false); + } + let num_consumed_items = self + .remaining + .fetch_update(Ordering::SeqCst, Ordering::SeqCst, |current| { + let num_items = (current / cost_per_item).min(max_items); + current.checked_sub(num_items * cost_per_item) + }) + .map(|old| (old / cost_per_item).min(max_items)) + .unwrap_or(0); + let budget_exhausted = num_consumed_items < max_items; + (num_consumed_items, budget_exhausted) + } } #[cfg(test)] @@ -90,6 +125,69 @@ 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 → not budget-limited (got what we asked for) + assert_eq!(budget.consume_up_to(2, 3), (2, false)); + assert_eq!(budget.remaining(), 3); + // 3 / 3 = 1 item, asked for 5 → budget-limited + assert_eq!(budget.consume_up_to(5, 3), (1, true)); + assert_eq!(budget.remaining(), 0); + } + + #[test] + fn test_consume_up_to_partial_budget() { + let budget = IoBudget::new(7); + // 7 / 3 = 2 items, asked for 10 → budget-limited + assert_eq!(budget.consume_up_to(10, 3), (2, true)); + assert_eq!(budget.remaining(), 1); + // 1 / 3 = 0 items → budget-limited + assert_eq!(budget.consume_up_to(10, 3), (0, true)); + assert_eq!(budget.remaining(), 1); // unchanged + } + + #[test] + fn test_consume_up_to_zero_budget() { + let budget = IoBudget::new(0); + // Asked for 5 but budget is 0 → budget-limited + assert_eq!(budget.consume_up_to(5, 3), (0, true)); + assert_eq!(budget.remaining(), 0); + } + + #[test] + #[should_panic(expected = "cost_per_item must be positive")] + fn test_consume_up_to_zero_cost_panics() { + let budget = IoBudget::new(10); + budget.consume_up_to(5, 0); + } + + #[test] + fn test_consume_up_to_zero_max_items() { + let budget = IoBudget::new(10); + // Asked for 0 → trivially satisfied, not budget-limited + assert_eq!(budget.consume_up_to(0, 3), (0, false)); + 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).0)); + } + + 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-service/src/rpc_backend.rs b/crates/discovery-service/src/rpc_backend.rs index 9c9588a64..69013ff16 100644 --- a/crates/discovery-service/src/rpc_backend.rs +++ b/crates/discovery-service/src/rpc_backend.rs @@ -39,6 +39,9 @@ impl From for StorageError { } } +/// Default maximum number of storage slots per JSON-RPC batch request. +const DEFAULT_MAX_BATCH_SIZE: usize = 50; + /// Configuration for the connection pool. #[derive(Debug, Clone)] pub struct PoolConfig { @@ -50,6 +53,9 @@ pub struct PoolConfig { pub request_timeout_secs: u64, /// Maximum idle connections per host. pub pool_max_idle_per_host: usize, + /// Maximum number of storage slots per JSON-RPC batch request. + /// Larger `read_slots` calls are automatically chunked. + pub max_batch_size: usize, } impl Default for PoolConfig { @@ -59,6 +65,7 @@ impl Default for PoolConfig { connect_timeout_secs: 30, request_timeout_secs: 60, pool_max_idle_per_host: 10, + max_batch_size: DEFAULT_MAX_BATCH_SIZE, } } } @@ -102,6 +109,7 @@ struct RpcBackendInner { provider: JsonRpcClient, contract_address: Felt, head: RwLock>, + max_batch_size: usize, } /// RPC-based storage backend that reads from a StarkNet node via JSON-RPC. @@ -135,6 +143,7 @@ impl RpcBackend { provider, contract_address: config.contract_address, head: RwLock::new(None), + max_batch_size: config.pool_config.max_batch_size, }), }) } @@ -159,28 +168,11 @@ pub struct RpcSnapshot { block_id: BlockId, } -#[async_trait] -impl RawStorageAccess for RpcSnapshot { - async fn read_slot(&self, slot: Felt) -> Result { - self.backend - .inner - .provider - .get_storage_at(self.backend.inner.contract_address, slot, self.block_id) - .await - .map_err(|e| RpcBackendError::Request(e.to_string()).into()) - } - - async fn read_slots(&self, slots: Vec) -> Result, StorageError> { - if slots.is_empty() { - return Ok(vec![]); - } - if slots.len() == 1 { - return Ok(vec![self.read_slot(slots[0]).await?]); - } - +impl RpcSnapshot { + /// Executes a single JSON-RPC batch request for the given slots. + async fn batch_read(&self, slots: &[Felt]) -> Result, StorageError> { let contract_address = self.backend.inner.contract_address; - // Build batch request let requests: Vec = slots .iter() .map(|&slot| { @@ -192,7 +184,6 @@ impl RawStorageAccess for RpcSnapshot { }) .collect(); - // Execute batch let responses = self .backend .inner @@ -201,7 +192,6 @@ impl RawStorageAccess for RpcSnapshot { .await .map_err(|e| RpcBackendError::Request(e.to_string()))?; - // Extract results responses .into_iter() .map(|resp| match resp { @@ -212,6 +202,34 @@ impl RawStorageAccess for RpcSnapshot { } } +#[async_trait] +impl RawStorageAccess for RpcSnapshot { + async fn read_slot(&self, slot: Felt) -> Result { + self.backend + .inner + .provider + .get_storage_at(self.backend.inner.contract_address, slot, self.block_id) + .await + .map_err(|e| RpcBackendError::Request(e.to_string()).into()) + } + + async fn read_slots(&self, slots: Vec) -> Result, StorageError> { + if slots.is_empty() { + return Ok(vec![]); + } + if slots.len() == 1 { + return Ok(vec![self.read_slot(slots[0]).await?]); + } + + let mut results = Vec::with_capacity(slots.len()); + for chunk in slots.chunks(self.backend.inner.max_batch_size) { + let chunk_results = self.batch_read(chunk).await?; + results.extend(chunk_results); + } + Ok(results) + } +} + #[async_trait] impl ChainState for RpcBackend { async fn get_head(&self) -> Option {