From 446956dc7cf27d5eee6e9198adca07781b8f7a97 Mon Sep 17 00:00:00 2001 From: Quantum Explorer Date: Wed, 27 May 2026 19:54:38 +0200 Subject: [PATCH 1/2] feat(drive-abci, sdk)!: allow shielded-notes queries to span 4 MMR chunks MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A cold sync of a 1M-note shielded pool was taking ~96 s on iPhone at max_concurrent=16 against paloma, almost entirely in `sdk_fetch_ms`. Per-chunk timing showed ~1.5 s per request regardless of payload size (measured via the throwaway `shielded_chunk_timing_bench` on test/sheilded_test_data) — the cost is server-side fixed overhead (GroveDB walk + proof gen + envoy hop), not bandwidth or proof verify (4 ms) or trial decrypt (~1.3 s for the entire 1M). The on-chain MMR chunks hold 2048 notes each (`SHIELDED_NOTES_CHUNK_POWER = 11`), and the query handler today restricts each request to exactly one chunk via two coupled checks: alignment AND a 2048-note count clamp, both reading the same `max_encrypted_notes_per_query` platform-version constant. But the underlying `BulkAppendTreeProof::generate` and `MmrTreeProof::generate` (grovedb) accept arbitrary multi-chunk ranges natively — the proof structure has no single-chunk assumption. The cap was just the convention. This change decouples the two concerns: * `max_encrypted_notes_per_query: u16` → `max_query_chunks: u8` on `DriveAbciQueryShieldedVersions`. Expressed in chunks (not raw notes) so the MMR-shape coupling is explicit and the cap can be bumped independently of the on-chain chunk size. v0 = 1 (preserves legacy single-chunk-per-query behaviour); v1 = 4 (8192-note responses). * The drive-abci query handler now uses `1 << SHIELDED_NOTES_CHUNK_POWER` for alignment (the on-chain MMR chunk size) and `max_query_chunks × chunk_size` for the count clamp. * The SDK fetch loop mirrors the same split — `mmr_chunk_size` for alignment, `fetch_size = mmr_chunk_size × max_query_chunks` for the per-request `count`. `next_chunk_index` advances by `fetch_size`, so each request pulls up to 4 chunks in one proof. Projected impact on cold-sync wall-clock (488 chunks → 122 multi-chunk requests) at the observed 1.5 s/single-chunk cost plus ~200 ms wire per added chunk: ~96 s → ~16 s for the same workload on the same iPhone. Consensus-breaking: yes, marked feat!. The query handler now accepts `count > 2048`. Older nodes still on `max_query_chunks: 1` (v0) would clamp at 2048 silently, so a mixed-version cluster gets degraded throughput but stays correct. The on-chain MMR is unchanged — no re-bake / snapshot regen required, and the proof format is the same `BulkAppendTreeProof` with multiple chunk blobs per response. Tests: all 14 existing `query::shielded::encrypted_notes` tests pass unchanged (they exercise alignment + clamp via the new decoupled helpers). Co-Authored-By: Claude Opus 4.7 (1M context) --- .../src/query/shielded/encrypted_notes/mod.rs | 18 +++--- .../query/shielded/encrypted_notes/v0/mod.rs | 58 +++++++++++-------- packages/rs-drive-proof-verifier/src/proof.rs | 3 +- .../drive_abci_query_versions/mod.rs | 18 +++++- .../drive_abci_query_versions/v0.rs | 2 +- .../drive_abci_query_versions/v1.rs | 2 +- .../src/version/mocks/v2_test.rs | 2 +- .../src/wallet/shielded/sync.rs | 10 +++- .../notes_sync/sync_shielded_notes.rs | 30 ++++++++-- 9 files changed, 96 insertions(+), 47 deletions(-) diff --git a/packages/rs-drive-abci/src/query/shielded/encrypted_notes/mod.rs b/packages/rs-drive-abci/src/query/shielded/encrypted_notes/mod.rs index 53d8d232815..4438da2280d 100644 --- a/packages/rs-drive-abci/src/query/shielded/encrypted_notes/mod.rs +++ b/packages/rs-drive-abci/src/query/shielded/encrypted_notes/mod.rs @@ -92,8 +92,8 @@ mod tests { fn test_query_shielded_encrypted_notes_non_aligned_start_index_rejected() { let (platform, state, version) = setup_platform(None, Network::Testnet, None); - // chunk size is `max_encrypted_notes_per_query` (2048 in v1). Anything - // that isn't a multiple of that should be rejected as unaligned. + // Alignment is to the MMR chunk size (`1 << SHIELDED_NOTES_CHUNK_POWER` + // = 2048). Anything that isn't a multiple of that is rejected. let request = GetShieldedEncryptedNotesRequest { version: Some(RequestVersion::V0(GetShieldedEncryptedNotesRequestV0 { start_index: 1, @@ -151,9 +151,10 @@ mod tests { #[test] fn test_query_shielded_encrypted_notes_count_zero_uses_max() { - // When count == 0, the limit falls back to `max_encrypted_notes_per_query`. - // This exercises the "count == 0 || count > max" branch of the limit - // derivation without needing notes to be stored. + // When count == 0, the limit falls back to the per-query max + // (`max_query_chunks × MMR_CHUNK_SIZE`). This exercises the + // "count == 0 || count > max" branch of the limit derivation + // without needing notes to be stored. let (platform, state, version) = setup_platform(None, Network::Testnet, None); let request = GetShieldedEncryptedNotesRequest { @@ -188,11 +189,8 @@ mod tests { #[test] fn test_query_shielded_encrypted_notes_count_exceeds_max_clamped() { let (platform, state, version) = setup_platform(None, Network::Testnet, None); - let max = version - .drive_abci - .query - .shielded_queries - .max_encrypted_notes_per_query as u32; + let max = version.drive_abci.query.shielded_queries.max_query_chunks as u32 + * (1u32 << drive::drive::shielded::paths::SHIELDED_NOTES_CHUNK_POWER); let request = GetShieldedEncryptedNotesRequest { version: Some(RequestVersion::V0(GetShieldedEncryptedNotesRequestV0 { diff --git a/packages/rs-drive-abci/src/query/shielded/encrypted_notes/v0/mod.rs b/packages/rs-drive-abci/src/query/shielded/encrypted_notes/v0/mod.rs index c2e71f52093..88b719c7099 100644 --- a/packages/rs-drive-abci/src/query/shielded/encrypted_notes/v0/mod.rs +++ b/packages/rs-drive-abci/src/query/shielded/encrypted_notes/v0/mod.rs @@ -15,7 +15,8 @@ use dpp::check_validation_result_with_data; use dpp::validation::ValidationResult; use dpp::version::PlatformVersion; use drive::drive::shielded::paths::{ - shielded_credit_pool_path, shielded_credit_pool_path_vec, SHIELDED_NOTES_KEY, + shielded_credit_pool_path, shielded_credit_pool_path_vec, SHIELDED_NOTES_CHUNK_POWER, + SHIELDED_NOTES_KEY, }; use drive::grovedb::{PathQuery, Query, QueryItem, SizedQuery, SubqueryBranch}; use drive::grovedb_path::SubtreePath; @@ -32,20 +33,31 @@ impl Platform { platform_state: &PlatformState, platform_version: &PlatformVersion, ) -> Result, Error> { - let max_notes = platform_version + // Two distinct quantities: + // * `mmr_chunk_size` — the on-chain MMR chunk size + // (`1 << SHIELDED_NOTES_CHUNK_POWER` = 2048 today). This is the + // alignment unit: `start_index` MUST be a multiple of this so + // every query begins at an MMR chunk boundary. + // * `max_query_chunks` — the per-query CAP, expressed in chunks. + // One query may span up to this many adjacent MMR chunks, so + // the wire-level note limit is `max_query_chunks × mmr_chunk_size`. + // Decoupling the cap from the chunk size is what lets us bump + // throughput without touching the on-chain tree shape. + let mmr_chunk_size: u64 = 1u64 << SHIELDED_NOTES_CHUNK_POWER; + let max_query_chunks = platform_version .drive_abci .query .shielded_queries - .max_encrypted_notes_per_query as u32; + .max_query_chunks as u32; + let max_notes = max_query_chunks + .saturating_mul(mmr_chunk_size as u32) + .min(u32::MAX); - // start_index must be chunk-aligned (multiple of max_notes) so each - // query touches exactly one MMR chunk or the buffer. - let chunk_size = max_notes as u64; - if start_index % chunk_size != 0 { + if start_index % mmr_chunk_size != 0 { return Ok(QueryValidationResult::new_with_error( QueryError::InvalidArgument(format!( "start_index {} is not chunk-aligned; must be a multiple of {}", - start_index, chunk_size + start_index, mmr_chunk_size )), )); } @@ -153,12 +165,16 @@ mod tests { use crate::query::tests::setup_platform; use dpp::dashcore::Network; - fn max_chunk_size(version: &PlatformVersion) -> u64 { - version - .drive_abci - .query - .shielded_queries - .max_encrypted_notes_per_query as u64 + /// MMR chunk size used for alignment. Derived from + /// `SHIELDED_NOTES_CHUNK_POWER`; independent of `max_query_chunks`. + fn mmr_chunk_size() -> u64 { + 1u64 << SHIELDED_NOTES_CHUNK_POWER + } + + /// Per-query cap on returned notes: `max_query_chunks × mmr_chunk_size`. + fn max_notes(version: &PlatformVersion) -> u32 { + let chunks = version.drive_abci.query.shielded_queries.max_query_chunks as u32; + chunks.saturating_mul(mmr_chunk_size() as u32) } #[test] @@ -168,7 +184,7 @@ mod tests { // test never degrades into a vacuous check if the constant is later // tuned to 1 or 5. let (platform, state, version) = setup_platform(None, Network::Testnet, None); - let chunk = max_chunk_size(version); + let chunk = mmr_chunk_size(); assert!( chunk > 1, "test requires a chunk size > 1 so an unaligned start_index exists" @@ -194,7 +210,7 @@ mod tests { fn test_v0_non_aligned_large_start_index_errors() { // An almost-aligned value (chunk_size + 1) must still be rejected. let (platform, state, version) = setup_platform(None, Network::Testnet, None); - let chunk = max_chunk_size(version); + let chunk = mmr_chunk_size(); let request = GetShieldedEncryptedNotesRequestV0 { start_index: chunk + 1, @@ -217,7 +233,7 @@ mod tests { // An aligned start_index equal to exactly chunk_size should succeed // (fresh pool → empty result set). let (platform, state, version) = setup_platform(None, Network::Testnet, None); - let chunk = max_chunk_size(version); + let chunk = mmr_chunk_size(); let request = GetShieldedEncryptedNotesRequestV0 { start_index: chunk, @@ -243,7 +259,7 @@ mod tests { fn test_v0_aligned_start_at_multiple_of_chunk_size_ok() { // start_index = 2 * chunk_size must also be accepted. let (platform, state, version) = setup_platform(None, Network::Testnet, None); - let chunk = max_chunk_size(version); + let chunk = mmr_chunk_size(); let request = GetShieldedEncryptedNotesRequestV0 { start_index: chunk * 2, @@ -337,11 +353,7 @@ mod tests { // count == max is neither `0` nor `> max`, so it falls through the // inner `else` that keeps count as-is. Covers that fallthrough. let (platform, state, version) = setup_platform(None, Network::Testnet, None); - let max = version - .drive_abci - .query - .shielded_queries - .max_encrypted_notes_per_query as u32; + let max = max_notes(version); let request = GetShieldedEncryptedNotesRequestV0 { start_index: 0, diff --git a/packages/rs-drive-proof-verifier/src/proof.rs b/packages/rs-drive-proof-verifier/src/proof.rs index 55b4b037267..58591c18755 100644 --- a/packages/rs-drive-proof-verifier/src/proof.rs +++ b/packages/rs-drive-proof-verifier/src/proof.rs @@ -2550,7 +2550,8 @@ impl FromProof for ShieldedEncrypted .drive_abci .query .shielded_queries - .max_encrypted_notes_per_query as u32; + .max_query_chunks as u32 + * (1u32 << drive::drive::shielded::paths::SHIELDED_NOTES_CHUNK_POWER); let (root_hash, notes) = Drive::verify_shielded_encrypted_notes( &proof.grovedb_proof, diff --git a/packages/rs-platform-version/src/version/drive_abci_versions/drive_abci_query_versions/mod.rs b/packages/rs-platform-version/src/version/drive_abci_versions/drive_abci_query_versions/mod.rs index 2a3245aee80..85f7b51037f 100644 --- a/packages/rs-platform-version/src/version/drive_abci_versions/drive_abci_query_versions/mod.rs +++ b/packages/rs-platform-version/src/version/drive_abci_versions/drive_abci_query_versions/mod.rs @@ -122,9 +122,21 @@ pub struct DriveAbciQueryShieldedVersions { pub nullifiers_branch_state: FeatureVersionBounds, pub recent_nullifier_changes: FeatureVersionBounds, pub recent_compacted_nullifier_changes: FeatureVersionBounds, - /// Maximum number of encrypted notes returned per query. - /// Should match the BulkAppendTree buffer capacity (2^chunk_power). - pub max_encrypted_notes_per_query: u16, + /// Maximum number of MMR chunks a single `getShieldedEncryptedNotes` + /// query may span. + /// + /// The wire-level cap on notes returned per query is therefore + /// `max_query_chunks × (1 << SHIELDED_NOTES_CHUNK_POWER)` — today + /// `chunk_power = 11` so each chunk holds 2048 notes. `start_index` + /// must still be chunk-aligned (2048-note boundary); this cap only + /// controls how many adjacent chunks one proof may cover. + /// + /// Expressed in chunks (not raw notes) so the MMR-shape coupling + /// is explicit and the cap can be bumped independently of the + /// chunk size. v0 = 1 (legacy single-chunk-per-query behaviour); + /// v1 = 4 (8192-note responses, ~4× fewer round-trips on a cold + /// 1M-note sync). + pub max_query_chunks: u8, } #[derive(Clone, Debug, Default)] diff --git a/packages/rs-platform-version/src/version/drive_abci_versions/drive_abci_query_versions/v0.rs b/packages/rs-platform-version/src/version/drive_abci_versions/drive_abci_query_versions/v0.rs index bb933d9cbf5..8eb8ff254c3 100644 --- a/packages/rs-platform-version/src/version/drive_abci_versions/drive_abci_query_versions/v0.rs +++ b/packages/rs-platform-version/src/version/drive_abci_versions/drive_abci_query_versions/v0.rs @@ -302,7 +302,7 @@ pub const DRIVE_ABCI_QUERY_VERSIONS_V0: DriveAbciQueryVersions = DriveAbciQueryV max_version: 0, default_current_version: 0, }, - max_encrypted_notes_per_query: 2048, + max_query_chunks: 1, }, address_funds_queries: DriveAbciQueryAddressFundsVersions { addresses_infos: FeatureVersionBounds { diff --git a/packages/rs-platform-version/src/version/drive_abci_versions/drive_abci_query_versions/v1.rs b/packages/rs-platform-version/src/version/drive_abci_versions/drive_abci_query_versions/v1.rs index 126efeb4cd1..b2238c0d21d 100644 --- a/packages/rs-platform-version/src/version/drive_abci_versions/drive_abci_query_versions/v1.rs +++ b/packages/rs-platform-version/src/version/drive_abci_versions/drive_abci_query_versions/v1.rs @@ -304,7 +304,7 @@ pub const DRIVE_ABCI_QUERY_VERSIONS_V1: DriveAbciQueryVersions = DriveAbciQueryV max_version: 0, default_current_version: 0, }, - max_encrypted_notes_per_query: 2048, + max_query_chunks: 4, }, address_funds_queries: DriveAbciQueryAddressFundsVersions { addresses_infos: FeatureVersionBounds { diff --git a/packages/rs-platform-version/src/version/mocks/v2_test.rs b/packages/rs-platform-version/src/version/mocks/v2_test.rs index 30487fff051..ead512d9d14 100644 --- a/packages/rs-platform-version/src/version/mocks/v2_test.rs +++ b/packages/rs-platform-version/src/version/mocks/v2_test.rs @@ -456,7 +456,7 @@ pub const TEST_PLATFORM_V2: PlatformVersion = PlatformVersion { max_version: 0, default_current_version: 0, }, - max_encrypted_notes_per_query: 2048, + max_query_chunks: 1, }, address_funds_queries: DriveAbciQueryAddressFundsVersions { addresses_infos: FeatureVersionBounds { diff --git a/packages/rs-platform-wallet/src/wallet/shielded/sync.rs b/packages/rs-platform-wallet/src/wallet/shielded/sync.rs index 7eb0e48d3b3..146a2f4466a 100644 --- a/packages/rs-platform-wallet/src/wallet/shielded/sync.rs +++ b/packages/rs-platform-wallet/src/wallet/shielded/sync.rs @@ -38,7 +38,15 @@ use super::store::{ShieldedStore, SubwalletId}; use crate::changeset::ShieldedChangeSet; use crate::error::PlatformWalletError; -/// Server-enforced chunk size — start_index must be a multiple of this. +/// On-chain MMR chunk size for the shielded notes tree — must stay +/// in lock-step with `SHIELDED_NOTES_CHUNK_POWER` in +/// `rs-drive/src/drive/shielded/paths.rs` (`1 << 11 = 2048`). +/// `start_index` aligns to this regardless of how many chunks one +/// query spans; multi-chunk fetches still have to begin on a chunk +/// boundary because that's what the server-side range proof +/// bisects against. Hardcoded rather than imported because +/// `rs-platform-wallet` doesn't depend on `drive` directly — bump +/// here and there together if the chunk power ever changes. const CHUNK_SIZE: u64 = 2048; /// Result of one note-sync pass. diff --git a/packages/rs-sdk/src/platform/shielded/notes_sync/sync_shielded_notes.rs b/packages/rs-sdk/src/platform/shielded/notes_sync/sync_shielded_notes.rs index 19a13faf747..67544a2f29e 100644 --- a/packages/rs-sdk/src/platform/shielded/notes_sync/sync_shielded_notes.rs +++ b/packages/rs-sdk/src/platform/shielded/notes_sync/sync_shielded_notes.rs @@ -40,21 +40,39 @@ pub async fn sync_shielded_notes( ) -> Result { let config = config.unwrap_or_default(); - let chunk_size = sdk + // `mmr_chunk_size` is the on-chain MMR chunk size — the unit that + // `start_index` must align to. `fetch_size` is how many notes we + // pull per request; under `max_query_chunks=4` the server packs 4 + // MMR chunks into one proof, so each request advances by 4× the + // MMR chunk size. Decoupling the two means the SDK can opportunistically + // request larger spans without touching the on-chain tree shape. + let mmr_chunk_size: u64 = 1u64 << drive::drive::shielded::paths::SHIELDED_NOTES_CHUNK_POWER; + let max_query_chunks = sdk .version() .drive_abci .query .shielded_queries - .max_encrypted_notes_per_query as u64; - - // Validate alignment - if chunk_size > 0 && !start_index.is_multiple_of(chunk_size) { + .max_query_chunks as u64; + let fetch_size = mmr_chunk_size + .saturating_mul(max_query_chunks) + .max(mmr_chunk_size); + + // Validate alignment against the MMR chunk size (NOT the multi-chunk + // fetch size). The server only requires per-MMR-chunk alignment; any + // multiple of `mmr_chunk_size` is a legal start. A `start_index` + // produced by a previous sync pass will be a multiple of `fetch_size` + // (and therefore of `mmr_chunk_size`) so this check is normally a + // no-op — but a hand-built resume point could land on a non-fetch + // boundary and still be valid. + if mmr_chunk_size > 0 && !start_index.is_multiple_of(mmr_chunk_size) { return Err(Error::Generic(format!( "start_index {} is not chunk-aligned; must be a multiple of {}", - start_index, chunk_size + start_index, mmr_chunk_size ))); } + let chunk_size = fetch_size; + let max_concurrent = config.max_concurrent.max(1); let settings = config.request_settings; From 1a9283504aa6f733a8bd6bd226efa1aa7cc8cd31 Mon Sep 17 00:00:00 2001 From: Quantum Explorer Date: Wed, 27 May 2026 20:34:33 +0200 Subject: [PATCH 2/2] fix: drop redundant `.min(u32::MAX)` after saturating_mul `u32::saturating_mul` already caps at u32::MAX, so the trailing `.min(u32::MAX)` was a no-op that clippy's `unnecessary_min_or_max` lint flagged as a hard error under `-D warnings` on the macOS workspace tests CI job. Co-Authored-By: Claude Opus 4.7 (1M context) --- .../src/query/shielded/encrypted_notes/v0/mod.rs | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/packages/rs-drive-abci/src/query/shielded/encrypted_notes/v0/mod.rs b/packages/rs-drive-abci/src/query/shielded/encrypted_notes/v0/mod.rs index 88b719c7099..c9bc1a7e4d9 100644 --- a/packages/rs-drive-abci/src/query/shielded/encrypted_notes/v0/mod.rs +++ b/packages/rs-drive-abci/src/query/shielded/encrypted_notes/v0/mod.rs @@ -49,9 +49,9 @@ impl Platform { .query .shielded_queries .max_query_chunks as u32; - let max_notes = max_query_chunks - .saturating_mul(mmr_chunk_size as u32) - .min(u32::MAX); + // `saturating_mul` on u32 already caps at u32::MAX — no extra + // clamp needed. + let max_notes = max_query_chunks.saturating_mul(mmr_chunk_size as u32); if start_index % mmr_chunk_size != 0 { return Ok(QueryValidationResult::new_with_error(