diff --git a/crates/node/src/database/mod.rs b/crates/node/src/database/mod.rs index 3047a96..1443589 100644 --- a/crates/node/src/database/mod.rs +++ b/crates/node/src/database/mod.rs @@ -27,6 +27,7 @@ pub trait DatabaseBackend: Send + Sync { &self, tag: NoteTag, cursor: u64, + limit: Option, ) -> Result, DatabaseError>; /// Fetch notes matching ANY of a set of tags, in a single DB snapshot. @@ -34,10 +35,14 @@ pub trait DatabaseBackend: Send + Sync { /// This is the preferred multi-tag query — running per-tag queries back /// to back reopens a race where a concurrent INSERT can land between two /// per-tag queries and get leapfrogged by the cursor advance. + /// + /// `limit` caps the number of returned notes. When `None`, the server's + /// default batch size applies. Values above the server cap are clamped. async fn fetch_notes_by_tags( &self, tags: &[NoteTag], cursor: u64, + limit: Option, ) -> Result, DatabaseError>; /// Get statistics about the database @@ -94,8 +99,9 @@ impl Database { &self, tag: NoteTag, cursor: u64, + limit: Option, ) -> Result, DatabaseError> { - self.backend.fetch_notes(tag, cursor).await + self.backend.fetch_notes(tag, cursor, limit).await } /// Fetch notes matching ANY of a set of tags, in a single DB snapshot. @@ -103,8 +109,9 @@ impl Database { &self, tags: &[NoteTag], cursor: u64, + limit: Option, ) -> Result, DatabaseError> { - self.backend.fetch_notes_by_tags(tags, cursor).await + self.backend.fetch_notes_by_tags(tags, cursor, limit).await } /// Get statistics about the database @@ -147,7 +154,7 @@ mod tests { db.store_note(¬e).await.unwrap(); // Cursor is now seq-based; 0 fetches everything. - let fetched_notes = db.fetch_notes(TAG_LOCAL_ANY.into(), 0).await.unwrap(); + let fetched_notes = db.fetch_notes(TAG_LOCAL_ANY.into(), 0, None).await.unwrap(); assert_eq!(fetched_notes.len(), 1); assert_eq!(fetched_notes[0].header.id(), note.header.id()); assert!(fetched_notes[0].seq > 0); @@ -184,7 +191,7 @@ mod tests { db.store_note(&second).await.unwrap(); // Fetch everything and assert INSERT order = read order = seq ascending. - let fetched = db.fetch_notes(TAG_LOCAL_ANY.into(), 0).await.unwrap(); + let fetched = db.fetch_notes(TAG_LOCAL_ANY.into(), 0, None).await.unwrap(); assert_eq!(fetched.len(), 2); assert!( fetched[0].seq < fetched[1].seq, @@ -197,7 +204,7 @@ mod tests { // Cursor between the two seqs returns only the second. let mid_cursor = u64::try_from(fetched[0].seq).expect("seq is non-negative"); - let after_first = db.fetch_notes(TAG_LOCAL_ANY.into(), mid_cursor).await.unwrap(); + let after_first = db.fetch_notes(TAG_LOCAL_ANY.into(), mid_cursor, None).await.unwrap(); assert_eq!(after_first.len(), 1); assert_eq!(after_first[0].details, vec![2]); } @@ -244,8 +251,8 @@ mod tests { } while writers.join_next().await.is_some() {} - let fetched_a = db.fetch_notes(TAG_A.into(), 0).await.unwrap(); - let fetched_b = db.fetch_notes(TAG_B.into(), 0).await.unwrap(); + let fetched_a = db.fetch_notes(TAG_A.into(), 0, None).await.unwrap(); + let fetched_b = db.fetch_notes(TAG_B.into(), 0, None).await.unwrap(); assert_eq!( fetched_a.len() + fetched_b.len(), 40, @@ -272,14 +279,14 @@ mod tests { db.store_note(¬e).await.unwrap(); // cursor=0 is strictly before any assigned seq → should return the note - let fetched = db.fetch_notes(TAG_LOCAL_ANY.into(), 0).await.unwrap(); + let fetched = db.fetch_notes(TAG_LOCAL_ANY.into(), 0, None).await.unwrap(); assert_eq!(fetched.len(), 1); let stored_seq = fetched[0].seq; assert!(stored_seq > 0, "expected seq > 0, got {stored_seq}"); // cursor = the note's own seq → strictly-greater filter excludes it let cursor = u64::try_from(stored_seq).expect("seq is non-negative"); - let after = db.fetch_notes(TAG_LOCAL_ANY.into(), cursor).await.unwrap(); + let after = db.fetch_notes(TAG_LOCAL_ANY.into(), cursor, None).await.unwrap(); assert_eq!(after.len(), 0); } @@ -307,7 +314,7 @@ mod tests { db.store_note(¬e1).await.unwrap(); // Client's first fetch sees note1, advances cursor using the returned seq. - let first = db.fetch_notes(TAG_LOCAL_ANY.into(), 0).await.unwrap(); + let first = db.fetch_notes(TAG_LOCAL_ANY.into(), 0, None).await.unwrap(); assert_eq!(first.len(), 1); let cursor = u64::try_from(first[0].seq).expect("seq is non-negative"); @@ -322,7 +329,7 @@ mod tests { // With seq cursor: note2 has seq > cursor → returned. // With old timestamp cursor: note2.ts == cursor → filtered by strict `>` → LOST. - let second = db.fetch_notes(TAG_LOCAL_ANY.into(), cursor).await.unwrap(); + let second = db.fetch_notes(TAG_LOCAL_ANY.into(), cursor, None).await.unwrap(); assert_eq!( second.len(), 1, @@ -366,7 +373,7 @@ mod tests { // === Simulate the old per-tag loop === // Step 1: fetch tag A. - let a_result = db.fetch_notes(TAG_A.into(), 0).await.unwrap(); + let a_result = db.fetch_notes(TAG_A.into(), 0, None).await.unwrap(); assert_eq!(a_result.len(), 1); // Between the per-tag queries, two concurrent writes commit in order: @@ -391,7 +398,7 @@ mod tests { .unwrap(); // Step 2: fetch tag B. - let b_result = db.fetch_notes(TAG_B.into(), 0).await.unwrap(); + let b_result = db.fetch_notes(TAG_B.into(), 0, None).await.unwrap(); assert_eq!(b_result.len(), 1); // Client advances cursor to max(seq) across both results, mirroring the @@ -404,8 +411,8 @@ mod tests { .unwrap_or(0); // Client's next per-tag fetches with the advanced cursor. - let retry_a = db.fetch_notes(TAG_A.into(), rcursor).await.unwrap(); - let retry_b = db.fetch_notes(TAG_B.into(), rcursor).await.unwrap(); + let retry_a = db.fetch_notes(TAG_A.into(), rcursor, None).await.unwrap(); + let retry_b = db.fetch_notes(TAG_B.into(), rcursor, None).await.unwrap(); // Per-tag loop sees only 2 of the 3 notes — the interleaved tag A // insert with the details=[2] payload is missing. @@ -420,7 +427,8 @@ mod tests { ); // === The fix: single-snapshot multi-tag query === - let snapshot = db.fetch_notes_by_tags(&[TAG_A.into(), TAG_B.into()], 0).await.unwrap(); + let snapshot = + db.fetch_notes_by_tags(&[TAG_A.into(), TAG_B.into()], 0, None).await.unwrap(); assert_eq!( snapshot.len(), 3, @@ -454,7 +462,7 @@ mod tests { // A realistic "legacy" cursor — microseconds since the epoch, currently // ~1.76×10^15. Well above the 10^12 threshold. let legacy_cursor: u64 = 1_760_000_000_000_000; - let fetched = db.fetch_notes(TAG_LOCAL_ANY.into(), legacy_cursor).await.unwrap(); + let fetched = db.fetch_notes(TAG_LOCAL_ANY.into(), legacy_cursor, None).await.unwrap(); assert_eq!( fetched.len(), 1, @@ -463,7 +471,7 @@ mod tests { // Sanity check: a non-legacy cursor above the note's seq should NOT trigger the reset. let normal_cursor: u64 = 1_000; - let empty = db.fetch_notes(TAG_LOCAL_ANY.into(), normal_cursor).await.unwrap(); + let empty = db.fetch_notes(TAG_LOCAL_ANY.into(), normal_cursor, None).await.unwrap(); assert_eq!(empty.len(), 0, "normal cursor > seq should filter correctly"); } @@ -493,7 +501,7 @@ mod tests { } // First fetch from cursor=0 returns exactly BATCH_SIZE rows. - let first = db.fetch_notes(TAG_LOCAL_ANY.into(), 0).await.unwrap(); + let first = db.fetch_notes(TAG_LOCAL_ANY.into(), 0, None).await.unwrap(); assert_eq!( i64::try_from(first.len()).unwrap(), FETCH_NOTES_BATCH_SIZE, @@ -502,13 +510,13 @@ mod tests { // Advance cursor to max(seq) and refetch — remaining rows come back. let advanced: u64 = first.iter().map(|n| u64::try_from(n.seq).unwrap()).max().unwrap(); - let second = db.fetch_notes(TAG_LOCAL_ANY.into(), advanced).await.unwrap(); + let second = db.fetch_notes(TAG_LOCAL_ANY.into(), advanced, None).await.unwrap(); assert_eq!(second.len(), extra, "second batch must contain the remaining {extra} rows"); // Third fetch drains nothing (nothing left). let third_cursor: u64 = second.iter().map(|n| u64::try_from(n.seq).unwrap()).max().unwrap_or(advanced); - let third = db.fetch_notes(TAG_LOCAL_ANY.into(), third_cursor).await.unwrap(); + let third = db.fetch_notes(TAG_LOCAL_ANY.into(), third_cursor, None).await.unwrap(); assert_eq!(third.len(), 0, "drained"); // Stats reflect every row written. diff --git a/crates/node/src/database/sqlite/mod.rs b/crates/node/src/database/sqlite/mod.rs index d67f131..d385486 100644 --- a/crates/node/src/database/sqlite/mod.rs +++ b/crates/node/src/database/sqlite/mod.rs @@ -130,20 +130,26 @@ impl DatabaseBackend for SqliteDatabase { Ok(()) } - #[tracing::instrument(skip(self), fields(operation = "db.fetch_notes"))] async fn fetch_notes( &self, tag: NoteTag, cursor: u64, + limit: Option, ) -> Result, DatabaseError> { - self.fetch_notes_by_tags(&[tag], cursor).await + self.fetch_notes_by_tags(&[tag], cursor, limit).await } - #[tracing::instrument(skip(self, tags), fields(operation = "db.fetch_notes_by_tags"))] + #[tracing::instrument(skip(self, tags), fields( + operation = "db.fetch_notes_by_tags", + tag_count = tags.len(), + cursor = cursor, + notes_returned = tracing::field::Empty, + ))] async fn fetch_notes_by_tags( &self, tags: &[NoteTag], cursor: u64, + limit: Option, ) -> Result, DatabaseError> { let timer = self.metrics.db_fetch_notes(); @@ -169,6 +175,11 @@ impl DatabaseBackend for SqliteDatabase { let tag_values: Vec = tags.iter().map(|t| i64::from(t.as_u32())).collect(); + // Client-requested limit, clamped to the server cap. + let effective_limit = limit + .filter(|&l| l > 0) + .map_or(FETCH_NOTES_BATCH_SIZE, |l| i64::from(l).min(FETCH_NOTES_BATCH_SIZE)); + // Single query for all tags runs in ONE DB snapshot, so a concurrent // INSERT can't land between per-tag queries and get leapfrogged by the // cursor advance. This closes the second half of the pagination race @@ -183,7 +194,7 @@ impl DatabaseBackend for SqliteDatabase { .filter(tag.eq_any(&tag_values)) .filter(seq.gt(cursor_i64)) .order(seq.asc()) - .limit(FETCH_NOTES_BATCH_SIZE) + .limit(effective_limit) .load::(conn)?; Ok(fetched_notes) }) @@ -197,6 +208,7 @@ impl DatabaseBackend for SqliteDatabase { stored_notes.push(stored_note); } + tracing::Span::current().record("notes_returned", stored_notes.len()); timer.finish("ok"); Ok(stored_notes) diff --git a/crates/node/src/node/grpc/mod.rs b/crates/node/src/node/grpc/mod.rs index f48aead..5134b66 100644 --- a/crates/node/src/node/grpc/mod.rs +++ b/crates/node/src/node/grpc/mod.rs @@ -33,9 +33,9 @@ use crate::metrics::MetricsGrpc; /// `fetch_notes` request. Guards against two concerns: /// - Server CPU: deduplicating `request_data.tags` via `BTreeSet` is `O(n log n)`; a client /// sending millions of tags can burn a worker. -/// - SQLite `IN (...)`: the underlying driver caps bound variables at +/// - `SQLite` `IN (...)`: the underlying driver caps bound variables at /// `SQLITE_MAX_VARIABLE_NUMBER` (32766 on recent builds, lower on older); blow that and the -/// query errors. Well below the SQLite cap so we have headroom for future query-plan changes. +/// query errors. Well below the `SQLite` cap so we have headroom for future query-plan changes. /// /// A realistic wallet tracks O(10) to O(100) tags; 128 is generous without /// being an attack surface. @@ -206,7 +206,7 @@ impl miden_note_transport_proto::miden_note_transport::miden_note_transport_serv // matching rows in one consistent snapshot. let stored_notes = self .database - .fetch_notes_by_tags(&tags, cursor) + .fetch_notes_by_tags(&tags, cursor, request_data.limit) .await .map_err(|e| tonic::Status::internal(format!("Failed to fetch notes: {e:?}")))?; @@ -307,7 +307,7 @@ mod tests { let server = test_server().await; let tags = vec![0u32; MAX_TAGS_PER_FETCH_REQUEST + 1]; - let request = tonic::Request::new(FetchNotesRequest { tags, cursor: 0 }); + let request = tonic::Request::new(FetchNotesRequest { tags, cursor: 0, limit: None }); let result = server.fetch_notes(request).await; let status = result.expect_err("expected InvalidArgument"); @@ -327,7 +327,7 @@ mod tests { let server = test_server().await; let tags = vec![0u32; MAX_TAGS_PER_FETCH_REQUEST]; - let request = tonic::Request::new(FetchNotesRequest { tags, cursor: 0 }); + let request = tonic::Request::new(FetchNotesRequest { tags, cursor: 0, limit: None }); let result = server.fetch_notes(request).await; let response = result.expect("request at the cap must succeed").into_inner(); diff --git a/crates/node/src/node/grpc/streaming.rs b/crates/node/src/node/grpc/streaming.rs index 2962bc5..210956b 100644 --- a/crates/node/src/node/grpc/streaming.rs +++ b/crates/node/src/node/grpc/streaming.rs @@ -82,7 +82,7 @@ impl NoteStreamerManager { let mut updates = vec![]; for (tag, tag_data) in &self.tags { - let snotes = self.database.fetch_notes(*tag, tag_data.cursor).await?; + let snotes = self.database.fetch_notes(*tag, tag_data.cursor, None).await?; let mut cursor = tag_data.cursor; for snote in &snotes { // Advance cursor using the DB-assigned monotonic `seq` diff --git a/crates/proto/src/generated/miden_note_transport.rs b/crates/proto/src/generated/miden_note_transport.rs index f9da3ea..329abbb 100644 --- a/crates/proto/src/generated/miden_note_transport.rs +++ b/crates/proto/src/generated/miden_note_transport.rs @@ -26,6 +26,10 @@ pub struct FetchNotesRequest { pub tags: ::prost::alloc::vec::Vec, #[prost(fixed64, tag = "2")] pub cursor: u64, + /// Maximum number of notes to return. When unset the server applies its + /// own batch cap (currently 500). Values above the server cap are clamped. + #[prost(uint32, optional, tag = "3")] + pub limit: ::core::option::Option, } /// API response for fetching notes #[derive(Clone, PartialEq, ::prost::Message)] diff --git a/proto/proto/miden_note_transport.proto b/proto/proto/miden_note_transport.proto index fbfc801..656d96f 100644 --- a/proto/proto/miden_note_transport.proto +++ b/proto/proto/miden_note_transport.proto @@ -26,6 +26,9 @@ message SendNoteResponse { } message FetchNotesRequest { repeated fixed32 tags = 1; fixed64 cursor = 2; + // Maximum number of notes to return. When unset the server applies its + // own batch cap (currently 500). Values above the server cap are clamped. + optional uint32 limit = 3; } // API response for fetching notes