Add archival version option to Bank RPC, README and test cases (#1284)

* simple archival rpc * archival rpc test * try to fix readme with using dict params
Sovereign-Labs · Jan 9, 2024 · 5884991 · 5884991
1 parent 7d97267
commit 5884991
Show file tree

Hide file tree

Showing 13 changed files with 213 additions and 33 deletions.
diff --git a/examples/demo-rollup/README.md b/examples/demo-rollup/README.md
@@ -253,7 +253,7 @@ $ cargo run --bin sov-cli rpc submit-batch by-address sov1l6n2cku82yfqld30lanm2n
 
 #### Verify the Token Supply
 ```bash,test-ci,bashtestmd:compare-output
-$ curl -X POST -H "Content-Type: application/json" -d '{"jsonrpc":"2.0","method":"bank_supplyOf","params":["sov1zdwj8thgev2u3yyrrlekmvtsz4av4tp3m7dm5mx5peejnesga27svq9m72"],"id":1}' http://127.0.0.1:12345
+$ curl -X POST -H "Content-Type: application/json" -d '{"jsonrpc":"2.0","method":"bank_supplyOf","params":{"token_address":"sov1zdwj8thgev2u3yyrrlekmvtsz4av4tp3m7dm5mx5peejnesga27svq9m72"},"id":1}' http://127.0.0.1:12345
 {"jsonrpc":"2.0","result":{"amount":1000},"id":1}
 ```
 

diff --git a/examples/demo-rollup/README_CELESTIA.md b/examples/demo-rollup/README_CELESTIA.md
@@ -299,7 +299,7 @@ This command will use your default private key.
 #### 4. Verify the Token Supply
 
 ```bash,test-ci,bashtestmd:compare-output
-$ curl -X POST -H "Content-Type: application/json" -d '{"jsonrpc":"2.0","method":"bank_supplyOf","params":["sov1zdwj8thgev2u3yyrrlekmvtsz4av4tp3m7dm5mx5peejnesga27svq9m72"],"id":1}' http://127.0.0.1:12345
+$ curl -X POST -H "Content-Type: application/json" -d '{"jsonrpc":"2.0","method":"bank_supplyOf","params":{"token_address":"sov1zdwj8thgev2u3yyrrlekmvtsz4av4tp3m7dm5mx5peejnesga27svq9m72"},"id":1}' http://127.0.0.1:12345
 {"jsonrpc":"2.0","result":{"amount":1000},"id":1}
 ```
 

diff --git a/examples/demo-rollup/stf/src/tests/stf_tests.rs b/examples/demo-rollup/stf/src/tests/stf_tests.rs
@@ -81,7 +81,7 @@ fn test_demo_values_in_db() {
         let mut working_set = WorkingSet::new(storage);
         let resp = runtime
             .bank
-            .supply_of(get_default_token_address(), &mut working_set)
+            .supply_of(None, get_default_token_address(), &mut working_set)
             .unwrap();
         assert_eq!(resp, sov_bank::TotalSupplyResponse { amount: Some(1000) });
 
@@ -143,7 +143,7 @@ fn test_demo_values_in_cache() {
 
     let resp = runtime
         .bank
-        .supply_of(get_default_token_address(), &mut working_set)
+        .supply_of(None, get_default_token_address(), &mut working_set)
         .unwrap();
     assert_eq!(resp, sov_bank::TotalSupplyResponse { amount: Some(1000) });
 
@@ -210,7 +210,7 @@ fn test_demo_values_not_in_db() {
 
         let resp = runtime
             .bank
-            .supply_of(get_default_token_address(), &mut working_set)
+            .supply_of(None, get_default_token_address(), &mut working_set)
             .unwrap();
         assert_eq!(resp, sov_bank::TotalSupplyResponse { amount: Some(1000) });
 

diff --git a/examples/demo-rollup/stf/src/tests/tx_revert_tests.rs b/examples/demo-rollup/stf/src/tests/tx_revert_tests.rs
@@ -91,6 +91,7 @@ fn test_tx_revert() {
         let resp = runtime
             .bank
             .balance_of(
+                None,
                 get_default_private_key().default_address(),
                 get_default_token_address(),
                 &mut working_set,

diff --git a/examples/demo-rollup/tests/bank/mod.rs b/examples/demo-rollup/tests/bank/mod.rs
@@ -5,11 +5,12 @@ use demo_stf::genesis_config::GenesisPaths;
 use demo_stf::runtime::RuntimeCall;
 use jsonrpsee::core::client::{Subscription, SubscriptionClientT};
 use jsonrpsee::rpc_params;
+use sov_bank::Coins;
 use sov_mock_da::MockDaSpec;
 use sov_modules_api::default_context::DefaultContext;
 use sov_modules_api::default_signature::private_key::DefaultPrivateKey;
 use sov_modules_api::transaction::Transaction;
-use sov_modules_api::{PrivateKey, Spec};
+use sov_modules_api::{Address, PrivateKey, Spec};
 use sov_modules_stf_blueprint::kernels::basic::BasicKernelGenesisPaths;
 use sov_sequencer::utils::SimpleClient;
 use sov_stf_runner::RollupProverConfig;
@@ -40,21 +41,13 @@ async fn bank_tx_tests() -> Result<(), anyhow::Error> {
     // If the rollup throws an error, return it and stop trying to send the transaction
     tokio::select! {
         err = rollup_task => err?,
-        res = send_test_create_token_tx(port) => res?,
+        res = send_test_bank_txs(port) => res?,
     };
     Ok(())
 }
 
-async fn send_test_create_token_tx(rpc_address: SocketAddr) -> Result<(), anyhow::Error> {
-    let key = DefaultPrivateKey::generate();
+async fn build_create_token_tx(key: &DefaultPrivateKey, nonce: u64) -> Transaction<DefaultContext> {
     let user_address: <DefaultContext as Spec>::Address = key.to_address();
-
-    let token_address = sov_bank::get_token_address::<DefaultContext>(
-        TOKEN_NAME,
-        user_address.as_ref(),
-        TOKEN_SALT,
-    );
-
     let msg = RuntimeCall::<DefaultContext, MockDaSpec>::bank(sov_bank::CallMessage::<
         DefaultContext,
     >::CreateToken {
@@ -67,16 +60,57 @@ async fn send_test_create_token_tx(rpc_address: SocketAddr) -> Result<(), anyhow
     let chain_id = 0;
     let gas_tip = 0;
     let gas_limit = 0;
-    let nonce = 0;
-    let tx = Transaction::<DefaultContext>::new_signed_tx(
+    Transaction::<DefaultContext>::new_signed_tx(
         &key,
         msg.try_to_vec().unwrap(),
         chain_id,
         gas_tip,
         gas_limit,
         nonce,
+    )
+}
+
+async fn build_transfer_token_tx(
+    key: &DefaultPrivateKey,
+    token_address: Address,
+    recipient: <DefaultContext as Spec>::Address,
+    amount: u64,
+    nonce: u64,
+) -> Transaction<DefaultContext> {
+    let msg = RuntimeCall::<DefaultContext, MockDaSpec>::bank(sov_bank::CallMessage::<
+        DefaultContext,
+    >::Transfer {
+        to: recipient,
+        coins: Coins {
+            amount,
+            token_address,
+        },
+    });
+    let chain_id = 0;
+    let gas_tip = 0;
+    let gas_limit = 0;
+    Transaction::<DefaultContext>::new_signed_tx(
+        &key,
+        msg.try_to_vec().unwrap(),
+        chain_id,
+        gas_tip,
+        gas_limit,
+        nonce,
+    )
+}
+
+async fn send_test_bank_txs(rpc_address: SocketAddr) -> Result<(), anyhow::Error> {
+    let key = DefaultPrivateKey::generate();
+    let user_address: <DefaultContext as Spec>::Address = key.to_address();
+
+    let token_address = sov_bank::get_token_address::<DefaultContext>(
+        TOKEN_NAME,
+        user_address.as_ref(),
+        TOKEN_SALT,
     );
 
+    let tx = build_create_token_tx(&key, 0).await;
+
     let port = rpc_address.port();
     let client = SimpleClient::new("localhost", port).await?;
 
@@ -96,10 +130,106 @@ async fn send_test_create_token_tx(rpc_address: SocketAddr) -> Result<(), anyhow
 
     let balance_response = sov_bank::BankRpcClient::<DefaultContext>::balance_of(
         client.http(),
+        None,
+        user_address,
+        token_address,
+    )
+    .await?;
+    assert_eq!(balance_response.amount.unwrap_or_default(), 1000);
+
+    let recipient_key = DefaultPrivateKey::generate();
+    let recipient_address: <DefaultContext as Spec>::Address = recipient_key.to_address();
+
+    let tx = build_transfer_token_tx(
+        &key,
+        token_address.clone(),
+        recipient_address.clone(),
+        100,
+        1,
+    )
+    .await;
+
+    let mut slot_processed_subscription: Subscription<u64> = client
+        .ws()
+        .subscribe(
+            "ledger_subscribeSlots",
+            rpc_params![],
+            "ledger_unsubscribeSlots",
+        )
+        .await?;
+
+    client.send_transaction(tx).await?;
+
+    // Wait until the rollup has processed the next slot
+    let _ = slot_processed_subscription.next().await;
+
+    let balance_response = sov_bank::BankRpcClient::<DefaultContext>::balance_of(
+        client.http(),
+        None,
+        user_address,
+        token_address,
+    )
+    .await?;
+    assert_eq!(balance_response.amount.unwrap_or_default(), 900);
+
+    let tx = build_transfer_token_tx(
+        &key,
+        token_address.clone(),
+        recipient_address.clone(),
+        200,
+        2,
+    )
+    .await;
+
+    let mut slot_processed_subscription: Subscription<u64> = client
+        .ws()
+        .subscribe(
+            "ledger_subscribeSlots",
+            rpc_params![],
+            "ledger_unsubscribeSlots",
+        )
+        .await?;
+
+    client.send_transaction(tx).await?;
+
+    // Wait until the rollup has processed the next slot
+    let _ = slot_processed_subscription.next().await;
+
+    let balance_response = sov_bank::BankRpcClient::<DefaultContext>::balance_of(
+        client.http(),
+        None,
+        user_address,
+        token_address,
+    )
+    .await?;
+    assert_eq!(balance_response.amount.unwrap_or_default(), 700);
+
+    let balance_response = sov_bank::BankRpcClient::<DefaultContext>::balance_of(
+        client.http(),
+        Some(3),
+        user_address,
+        token_address,
+    )
+    .await?;
+    assert_eq!(balance_response.amount.unwrap_or_default(), 900);
+
+    let balance_response = sov_bank::BankRpcClient::<DefaultContext>::balance_of(
+        client.http(),
+        Some(4),
+        user_address,
+        token_address,
+    )
+    .await?;
+    assert_eq!(balance_response.amount.unwrap_or_default(), 700);
+
+    let balance_response = sov_bank::BankRpcClient::<DefaultContext>::balance_of(
+        client.http(),
+        Some(2),
         user_address,
         token_address,
     )
     .await?;
     assert_eq!(balance_response.amount.unwrap_or_default(), 1000);
+
     Ok(())
 }
diff --git a/module-system/RPC_WALKTHROUGH.md b/module-system/RPC_WALKTHROUGH.md
@@ -126,3 +126,38 @@ fn main() {
 
 }
 ```
+
+## Enabling Archival queries for RPC
+* We use `working_set: &mut WorkingSet<C>` in order to query state. `WorkingSet` has a function `working_set.set_archival_version(v)` where v is of type `u64` and represents the block height. 
+* Once the `set_archival_version` is called, the working_set is configured to query against the state at height `v`.
+* To modify an RPC query of the form
+```rust
+pub fn balance_of(
+         &self,
+         user_address: C::Address,
+         token_address: C::Address,
+         working_set: &mut WorkingSet<C>,
+     ) -> RpcResult<BalanceResponse> {
+    Ok(BalanceResponse {
+        amount: self.get_balance_of(user_address, token_address, working_set),
+    })
+}
+```
+We need to make the following changes
+```rust
+pub fn balance_of(
+         &self,
+         version: Option<u64>,
+         user_address: C::Address,
+         token_address: C::Address,
+         working_set: &mut WorkingSet<C>,
+     ) -> RpcResult<BalanceResponse> {
+    if let Some(v) = version {
+        working_set.set_archival_version(v)
+    }
+    Ok(BalanceResponse {
+        amount: self.get_balance_of(user_address, token_address, working_set),
+    })
+}
+```
+* NOTE: `set_archival_version` handles configuring `WorkingSet` for both JMT state as well as accessory state
diff --git a/module-system/module-implementations/sov-bank/src/query.rs b/module-system/module-implementations/sov-bank/src/query.rs
@@ -26,10 +26,14 @@ impl<C: sov_modules_api::Context> Bank<C> {
     /// stored at the address `token_address`.
     pub fn balance_of(
         &self,
+        version: Option<u64>,
         user_address: C::Address,
         token_address: C::Address,
         working_set: &mut WorkingSet<C>,
     ) -> RpcResult<BalanceResponse> {
+        if let Some(v) = version {
+            working_set.set_archival_version(v)
+        }
         Ok(BalanceResponse {
             amount: self.get_balance_of(user_address, token_address, working_set),
         })
@@ -39,9 +43,13 @@ impl<C: sov_modules_api::Context> Bank<C> {
     /// Rpc method that returns the supply of a token stored at the address `token_address`.
     pub fn supply_of(
         &self,
+        version: Option<u64>,
         token_address: C::Address,
         working_set: &mut WorkingSet<C>,
     ) -> RpcResult<TotalSupplyResponse> {
+        if let Some(v) = version {
+            working_set.set_archival_version(v)
+        }
         Ok(TotalSupplyResponse {
             amount: self.get_total_supply_of(&token_address, working_set),
         })

diff --git a/module-system/module-implementations/sov-bank/tests/burn_test.rs b/module-system/module-implementations/sov-bank/tests/burn_test.rs
@@ -48,7 +48,8 @@ fn burn_deployed_tokens() {
     assert!(working_set.events().is_empty());
 
     let query_total_supply = |working_set: &mut WorkingSet<DefaultContext>| -> Option<u64> {
-        let total_supply: TotalSupplyResponse = bank.supply_of(token_address, working_set).unwrap();
+        let total_supply: TotalSupplyResponse =
+            bank.supply_of(None, token_address, working_set).unwrap();
         total_supply.amount
     };
 

diff --git a/module-system/module-implementations/sov-bank/tests/freeze_test.rs b/module-system/module-implementations/sov-bank/tests/freeze_test.rs
@@ -132,12 +132,12 @@ fn freeze_token() {
         minter_address: new_holder,
     };
 
-    let query_total_supply = |token_address: Address,
-                              working_set: &mut WorkingSet<DefaultContext>|
-     -> Option<u64> {
-        let total_supply: TotalSupplyResponse = bank.supply_of(token_address, working_set).unwrap();
-        total_supply.amount
-    };
+    let query_total_supply =
+        |token_address: Address, working_set: &mut WorkingSet<DefaultContext>| -> Option<u64> {
+            let total_supply: TotalSupplyResponse =
+                bank.supply_of(None, token_address, working_set).unwrap();
+            total_supply.amount
+        };
 
     let minted = bank.call(mint_message, &minter_context, &mut working_set);
     assert!(minted.is_err());

diff --git a/module-system/module-implementations/sov-bank/tests/mint_test.rs b/module-system/module-implementations/sov-bank/tests/mint_test.rs
@@ -42,12 +42,12 @@ fn mint_token() {
     // No events at the moment. If there are, needs to be checked
     assert!(working_set.events().is_empty());
 
-    let query_total_supply = |token_address: Address,
-                              working_set: &mut WorkingSet<DefaultContext>|
-     -> Option<u64> {
-        let total_supply: TotalSupplyResponse = bank.supply_of(token_address, working_set).unwrap();
-        total_supply.amount
-    };
+    let query_total_supply =
+        |token_address: Address, working_set: &mut WorkingSet<DefaultContext>| -> Option<u64> {
+            let total_supply: TotalSupplyResponse =
+                bank.supply_of(None, token_address, working_set).unwrap();
+            total_supply.amount
+        };
 
     let query_user_balance =
         |user_address: Address, working_set: &mut WorkingSet<DefaultContext>| -> Option<u64> {

diff --git a/module-system/module-implementations/sov-bank/tests/transfer_test.rs b/module-system/module-implementations/sov-bank/tests/transfer_test.rs
@@ -40,7 +40,8 @@ fn transfer_initial_token() {
         };
 
     let query_total_supply = |working_set: &mut WorkingSet<DefaultContext>| -> Option<u64> {
-        let total_supply: TotalSupplyResponse = bank.supply_of(token_address, working_set).unwrap();
+        let total_supply: TotalSupplyResponse =
+            bank.supply_of(None, token_address, working_set).unwrap();
         total_supply.amount
     };
 
@@ -282,7 +283,8 @@ fn transfer_deployed_token() {
         };
 
     let query_total_supply = |working_set: &mut WorkingSet<DefaultContext>| -> Option<u64> {
-        let total_supply: TotalSupplyResponse = bank.supply_of(token_address, working_set).unwrap();
+        let total_supply: TotalSupplyResponse =
+            bank.supply_of(None, token_address, working_set).unwrap();
         total_supply.amount
     };