feat(rpc): add streaming SyncAccountVaultV2 endpoint

crates/rpc/src/server/api.rs

@@ -90,6 +90,7 @@ mod submit_proven_tx_batch;
 mod subscription;
 mod sync_account_storage_maps;
 mod sync_account_vault;
+mod sync_account_vault_v2;
 mod sync_chain_mmr;
 mod sync_notes;
 mod sync_nullifiers;

crates/rpc/src/server/api/sync_account_vault_v2.rs

@@ -0,0 +1,174 @@
+use std::num::NonZeroUsize;
+use std::ops::RangeInclusive;
+use std::time::Duration;
+
+use miden_node_proto::decode::{read_account_id, read_block_range};
+use miden_node_proto::generated as proto;
+use miden_node_store::{AccountVaultValue, AccountVaultValuesPage, StateView};
+use miden_node_utils::tracing::{miden_instrument, miden_span_record};
+use miden_protocol::Word;
+use miden_protocol::account::AccountId;
+use miden_protocol::block::BlockNumber;
+use tokio::sync::mpsc;
+use tokio::sync::mpsc::error::SendTimeoutError;
+use tokio_stream::wrappers::ReceiverStream;
+use tonic::Status;
+use tracing::Instrument;
+
+use super::{
+    RpcInvalidBlockRange,
+    RpcService,
+    database_error_to_status,
+    invalid_block_range_to_status,
+};
+use crate::{COMPONENT, LOG_TARGET};
+
+/// Database rows fetched per page. This bounds internal work and memory, not encoded response size.
+const DB_PAGE_SIZE: NonZeroUsize = NonZeroUsize::new(256).unwrap();
+/// Stream items buffered before backpressure pauses the database producer.
+const STREAM_BUFFER_SIZE: usize = 32;
+/// Maximum time a stream producer waits for a stalled client to accept one update.
+const SEND_TIMEOUT: Duration = Duration::from_secs(10);
+
+type Input = (AccountId, RangeInclusive<BlockNumber>);
+
+#[tonic::async_trait]
+impl proto::server::rpc_api::SyncAccountVaultV2 for RpcService {
+    type Input = Input;
+    type Item = AccountVaultValue;
+    type ItemStream = ReceiverStream<tonic::Result<Self::Item>>;
+
+    fn decode(request: proto::rpc::SyncAccountVaultV2Request) -> tonic::Result<Self::Input> {
+        let account_id =
+            read_account_id::<proto::rpc::SyncAccountVaultV2Request, Status>(request.account_id)?;
+        let range = read_block_range::<Status>(request.block_range, "SyncAccountVaultV2Request")?;
+        let block_range = range
+            .into_inclusive_range::<RpcInvalidBlockRange>()
+            .map_err(invalid_block_range_to_status)?;
+
+        Ok((account_id, block_range))
+    }
+
+    fn encode(item: Self::Item) -> tonic::Result<proto::rpc::AccountVaultUpdate> {
+        let vault_key: Word = item.vault_key.into();
+        Ok(proto::rpc::AccountVaultUpdate {
+            vault_key: Some(vault_key.into()),
+            asset: item.asset.map(Into::into),
+            block_num: item.block_num.as_u32(),
+        })
+    }
+
+    #[miden_instrument(
+        target = COMPONENT,
+        name = "sync_account_vault_v2",
+        err,
+    )]
+    async fn handle(
+        &self,
+        (account_id, block_range): Self::Input,
+        _metadata: &tonic::metadata::MetadataMap,
+        _extensions: &tonic::codegen::http::Extensions,
+    ) -> tonic::Result<Self::ItemStream> {
+        miden_span_record!(
+            account.id = %account_id,
+            block_range.from = %block_range.start(),
+            block_range.to = %block_range.end(),
+        );
+
+        tracing::debug!(target: LOG_TARGET, "Streaming account vault updates");
+
+        if !account_id.is_public() {
+            return Err(Status::invalid_argument(format!("account {account_id} is not public")));
+        }
+
+        // Keep this view for the finite stream's lifetime. Besides fixing the chain-tip view used
+        // for validation, this pins the history generation so pruning cannot remove rows between
+        // internal database pages. Cancellation and the bounded send timeout release the view if
+        // the client stops consuming the stream.
+        let view = self.state.view();
+        let first_page = view
+            .sync_account_vault_v2_page(account_id, block_range.clone(), None, DB_PAGE_SIZE)
+            .await
+            .map_err(|err| database_error_to_status(&err))?;
+
+        // Reserve a slot for a terminal error so a full data buffer cannot turn a timeout or
+        // database failure into an apparently successful end-of-stream.
+        let (tx, rx) = mpsc::channel(STREAM_BUFFER_SIZE + 1);
+        let terminal_permit = tx
+            .clone()
+            .try_reserve_owned()
+            .expect("a newly created vault sync channel must have capacity");
+        VaultSyncProducer {
+            view,
+            account_id,
+            block_range,
+            page: first_page,
+            tx,
+            terminal_permit: Some(terminal_permit),
+        }
+        .spawn();
+
+        Ok(ReceiverStream::new(rx))
+    }
+}
+
+struct VaultSyncProducer {
+    view: StateView,
+    account_id: AccountId,
+    block_range: RangeInclusive<BlockNumber>,
+    page: AccountVaultValuesPage,
+    tx: mpsc::Sender<tonic::Result<AccountVaultValue>>,
+    terminal_permit: Option<mpsc::OwnedPermit<tonic::Result<AccountVaultValue>>>,
+}
+
+impl VaultSyncProducer {
+    fn spawn(self) {
+        tokio::spawn(self.run().instrument(tracing::Span::current()));
+    }
+
+    async fn run(mut self) {
+        loop {
+            let next_cursor = self.page.next_cursor.take();
+            for value in std::mem::take(&mut self.page.values) {
+                match self.tx.send_timeout(Ok(value), SEND_TIMEOUT).await {
+                    Ok(()) => {},
+                    Err(SendTimeoutError::Closed(_)) => return,
+                    Err(SendTimeoutError::Timeout(_)) => {
+                        self.send_terminal_error(Status::deadline_exceeded(
+                            "account vault sync client stopped consuming updates",
+                        ));
+                        return;
+                    },
+                }
+            }
+
+            let Some(cursor) = next_cursor else {
+                return;
+            };
+
+            self.page = match self
+                .view
+                .sync_account_vault_v2_page(
+                    self.account_id,
+                    self.block_range.clone(),
+                    Some(cursor),
+                    DB_PAGE_SIZE,
+                )
+                .await
+            {
+                Ok(page) => page,
+                Err(err) => {
+                    self.send_terminal_error(database_error_to_status(&err));
+                    return;
+                },
+            };
+        }
+    }
+
+    fn send_terminal_error(&mut self, status: Status) {
+        self.terminal_permit
+            .take()
+            .expect("terminal permit is consumed at most once")
+            .send(Err(status));
+    }
+}

crates/rpc/src/tests.rs

@@ -43,6 +43,12 @@ use miden_protocol::account::{
     AccountUpdateDetails,
     AssetCallbackFlag,
 };
+use miden_protocol::asset::{Asset, FungibleAsset};
+use miden_protocol::block::BlockNumber;
+use miden_protocol::testing::account_id::{
+    ACCOUNT_ID_PUBLIC_FUNGIBLE_FAUCET,
+    ACCOUNT_ID_PUBLIC_FUNGIBLE_FAUCET_1,
+};
 use miden_protocol::testing::noop_auth_component::NoopAuthComponent;
 use miden_protocol::transaction::{ProvenTransaction, TxAccountUpdate};
 use miden_protocol::utils::serde::Serializable;
@@ -1190,12 +1196,16 @@ async fn get_limits_endpoint() {
         QueryParamNoteTagLimit::LIMIT
     );
 
-    // SyncAccountVault and SyncAccountStorageMaps accept a singular account_id, not a repeated
+    // The account vault and storage-map endpoints accept a singular account_id, not a repeated
     // list, so they do not have list parameter limits.
     assert!(
         !limits.endpoints.contains_key("SyncAccountVault"),
         "SyncAccountVault should not have list parameter limits"
     );
+    assert!(
+        !limits.endpoints.contains_key("SyncAccountVaultV2"),
+        "SyncAccountVaultV2 should not have list parameter limits"
+    );
     assert!(
         !limits.endpoints.contains_key("SyncAccountStorageMaps"),
         "SyncAccountStorageMaps should not have list parameter limits"
@@ -1341,6 +1351,15 @@ async fn sync_endpoints_reject_block_to_beyond_chain_tip() {
         .expect_err("sync_account_vault should reject block_to beyond chain tip");
     assert_beyond_tip(&status, "sync_account_vault");
 
+    let status = rpc_client
+        .sync_account_vault_v2(proto::rpc::SyncAccountVaultV2Request {
+            block_range: block_range(),
+            account_id: account_id(),
+        })
+        .await
+        .expect_err("sync_account_vault_v2 should reject block_to beyond chain tip");
+    assert_beyond_tip(&status, "sync_account_vault_v2");
+
     let status = rpc_client
         .sync_transactions(proto::rpc::SyncTransactionsRequest {
             block_range: block_range(),
@@ -1350,3 +1369,83 @@ async fn sync_endpoints_reject_block_to_beyond_chain_tip() {
         .expect_err("sync_transactions should reject block_to beyond chain tip");
     assert_beyond_tip(&status, "sync_transactions");
 }
+
+#[tokio::test]
+async fn sync_account_vault_v2_validates_requests_and_completes_empty_stream() {
+    let (mut rpc_client, _rpc_addr, _store) = start_rpc().await;
+    let public_account = AccountId::dummy(
+        [0; 15],
+        AccountIdVersion::Version1,
+        AccountType::Public,
+        AssetCallbackFlag::Disabled,
+    );
+
+    let status = rpc_client
+        .sync_account_vault_v2(proto::rpc::SyncAccountVaultV2Request {
+            block_range: None,
+            account_id: Some(public_account.into()),
+        })
+        .await
+        .expect_err("sync_account_vault_v2 should require a block range");
+    assert_eq!(status.code(), tonic::Code::InvalidArgument);
+
+    let private_account = AccountId::dummy(
+        [1; 15],
+        AccountIdVersion::Version1,
+        AccountType::Private,
+        AssetCallbackFlag::Disabled,
+    );
+    let status = rpc_client
+        .sync_account_vault_v2(proto::rpc::SyncAccountVaultV2Request {
+            block_range: Some(proto::rpc::BlockRange { block_from: 0, block_to: 0 }),
+            account_id: Some(private_account.into()),
+        })
+        .await
+        .expect_err("sync_account_vault_v2 should reject private accounts");
+    assert_eq!(status.code(), tonic::Code::InvalidArgument);
+
+    let mut stream = rpc_client
+        .sync_account_vault_v2(proto::rpc::SyncAccountVaultV2Request {
+            block_range: Some(proto::rpc::BlockRange { block_from: 0, block_to: 0 }),
+            account_id: Some(public_account.into()),
+        })
+        .await
+        .expect("sync_account_vault_v2 should accept a public account at the chain tip")
+        .into_inner();
+    assert_eq!(stream.message().await.expect("stream should complete successfully"), None);
+}
+
+#[tokio::test]
+async fn sync_account_vault_v2_streams_squashed_updates() {
+    let (mut rpc_client, _rpc_addr, store) = start_rpc().await;
+    let account_id = AccountId::try_from(ACCOUNT_ID_PUBLIC_FUNGIBLE_FAUCET).unwrap();
+    let other_faucet = AccountId::try_from(ACCOUNT_ID_PUBLIC_FUNGIBLE_FAUCET_1).unwrap();
+    let asset_a = Asset::Fungible(FungibleAsset::new(account_id, 100).unwrap());
+    let asset_b = Asset::Fungible(FungibleAsset::new(other_faucet, 200).unwrap());
+    miden_node_store::test_support::seed_account_vault(
+        &store.data_directory_path().join("miden-store.sqlite3"),
+        account_id,
+        BlockNumber::GENESIS,
+        &[(asset_a.id(), Some(asset_a)), (asset_b.id(), Some(asset_b))],
+    );
+
+    let mut stream = rpc_client
+        .sync_account_vault_v2(proto::rpc::SyncAccountVaultV2Request {
+            block_range: Some(proto::rpc::BlockRange { block_from: 0, block_to: 0 }),
+            account_id: Some(account_id.into()),
+        })
+        .await
+        .expect("sync_account_vault_v2 should return a stream")
+        .into_inner();
+
+    let mut assets = Vec::new();
+    while let Some(update) = stream.message().await.expect("stream should complete successfully") {
+        assert_eq!(update.block_num, 0);
+        assets.push(Asset::try_from(update.asset.expect("seeded values are additions")).unwrap());
+    }
+    assets.sort_by_key(Asset::id);
+
+    let mut expected = vec![asset_a, asset_b];
+    expected.sort_by_key(Asset::id);
+    assert_eq!(assets, expected);
+}

crates/store/src/db/mod.rs

@@ -120,14 +120,28 @@ impl DerefMut for Db {
 /// Describes the value of an asset for an account ID at `block_num` specifically.
 ///
 /// If `asset` is `None`, the asset was removed.
-#[derive(Debug, Clone)]
+#[derive(Debug, Clone, PartialEq, Eq)]
 pub struct AccountVaultValue {
     pub block_num: BlockNumber,
     pub vault_key: AssetId,
     /// None if the asset was removed
     pub asset: Option<Asset>,
 }
 
+/// Stable cursor used to read a squashed account-vault delta in bounded database pages.
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct AccountVaultCursor {
+    pub(crate) block_num: BlockNumber,
+    pub(crate) vault_key: AssetId,
+}
+
+/// A bounded page of squashed account-vault updates.
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub struct AccountVaultValuesPage {
+    pub values: Vec<AccountVaultValue>,
+    pub next_cursor: Option<AccountVaultCursor>,
+}
+
 impl AccountVaultValue {
     pub fn from_raw_row(row: (i64, Vec<u8>, Option<Vec<u8>>)) -> Result<Self, DatabaseError> {
         let (block_num, vault_key, asset) = row;
@@ -784,6 +798,27 @@ impl Db {
         .await
     }
 
+    /// Selects one final update per vault key changed in `block_range`.
+    pub async fn select_account_vault_updates_v2(
+        &self,
+        account_id: AccountId,
+        block_range: ScopedBlockRange,
+        cursor: Option<AccountVaultCursor>,
+        page_size: NonZeroUsize,
+    ) -> Result<AccountVaultValuesPage> {
+        let block_range = block_range.into_inner();
+        self.transact("account vault sync v2", move |conn| {
+            queries::select_account_vault_updates_v2(
+                conn,
+                account_id,
+                block_range,
+                cursor,
+                page_size,
+            )
+        })
+        .await
+    }
+
     /// Returns the script for a note by its root.
     pub async fn select_note_script_by_root(&self, root: Word) -> Result<Option<NoteScript>> {
         self.transact("note script by root", move |conn| {