feat: gas hook on ExecutionFinalResult

examples/src/ref_finance.rs

@@ -203,7 +203,6 @@ async fn create_custom_ft(
         .transact()
         .await?
         .into_result()?;
-    ();
 
     Ok(ft)
 }

workspaces/src/error/impls.rs

@@ -1,5 +1,5 @@
-use std::borrow::Cow;
 use std::fmt;
+use std::{borrow::Cow, sync::PoisonError};
 
 use crate::result::ExecutionFailure;
 
@@ -135,6 +135,12 @@ impl std::error::Error for Error {
     }
 }
 
+impl<T> From<PoisonError<T>> for Error {
+    fn from(value: PoisonError<T>) -> Self {
+        Error::custom(ErrorKind::Other, value.to_string())
+    }
+}
+
 impl SandboxErrorCode {
     pub(crate) fn message<T>(self, msg: T) -> Error
     where

workspaces/src/network/variants.rs

@@ -42,9 +42,16 @@ where
     T: DevNetwork + TopLevelAccountCreator + 'static,
 {
     pub async fn create_tla(&self, id: AccountId, sk: SecretKey) -> Result<Execution<Account>> {
-        self.workspace
+        let res = self
+            .workspace
             .create_tla(self.clone().coerce(), id, sk)
-            .await
+            .await?;
+
+        self.tx_callbacks.iter().for_each(|meter| {
+            meter(res.details.total_gas_burnt).unwrap();
+        });
+
+        Ok(res)
     }
 
     pub async fn create_tla_and_deploy(
@@ -53,9 +60,16 @@ where
         sk: SecretKey,
         wasm: &[u8],
     ) -> Result<Execution<Contract>> {
-        self.workspace
+        let res = self
+            .workspace
             .create_tla_and_deploy(self.clone().coerce(), id, sk, wasm)
-            .await
+            .await?;
+
+        self.tx_callbacks.iter().for_each(|meter| {
+            meter(res.details.total_gas_burnt).unwrap();
+        });
+
+        Ok(res)
     }
 
     pub async fn dev_generate(&self) -> (AccountId, SecretKey) {

workspaces/src/operations.rs

@@ -236,13 +236,28 @@ impl Transaction {
     }
 
     async fn transact_raw(self) -> Result<FinalExecutionOutcomeView> {
-        send_batch_tx_and_retry(
+        let view = send_batch_tx_and_retry(
             self.worker.client(),
             &self.signer,
             &self.receiver_id,
             self.actions?,
         )
-        .await
+        .await?;
+
+        if !self.worker.tx_callbacks.is_empty() {
+            let total_gas_burnt = view.transaction_outcome.outcome.gas_burnt
+                + view
+                    .receipts_outcome
+                    .iter()
+                    .map(|t| t.outcome.gas_burnt)
+                    .sum::<u64>();
+
+            self.worker.tx_callbacks.iter().for_each(|meter| {
+                meter(Gas::from_gas(total_gas_burnt)).unwrap();
+            });
+        }
+
+        Ok(view)
     }
 
     /// Process the transaction, and return the result of the execution.
@@ -335,7 +350,8 @@ impl CallTransaction {
     /// object and return us the execution details, along with any errors if the transaction
     /// failed in any process along the way.
     pub async fn transact(self) -> Result<ExecutionFinalResult> {
-        self.worker
+        let txn = self
+            .worker
             .client()
             .call(
                 &self.signer,
@@ -347,7 +363,13 @@ impl CallTransaction {
             )
             .await
             .map(ExecutionFinalResult::from_view)
-            .map_err(crate::error::Error::from)
+            .map_err(crate::error::Error::from)?;
+
+        self.worker.tx_callbacks.iter().for_each(|meter| {
+            meter(txn.total_gas_burnt).unwrap();
+        });
+
+        Ok(txn)
     }
 
     /// Send the transaction to the network to be processed. This will be done asynchronously
@@ -446,10 +468,15 @@ impl<'a, 'b> CreateAccountTransaction<'a, 'b> {
 
         let signer = InMemorySigner::from_secret_key(id, sk);
         let account = Account::new(signer, self.worker.clone());
+        let details = ExecutionFinalResult::from_view(outcome);
+
+        self.worker.tx_callbacks.iter().for_each(|meter| {
+            meter(details.total_gas_burnt).unwrap();
+        });
 
         Ok(Execution {
             result: account,
-            details: ExecutionFinalResult::from_view(outcome),
+            details,
         })
     }
 }

workspaces/src/types/gas_meter.rs

@@ -0,0 +1,69 @@
+use std::sync::{Arc, Mutex};
+
+use super::Gas;
+use crate::result::Result;
+use crate::Worker;
+
+/// A hook that is called on every transaction that is sent to the network.
+/// This is useful for debugging purposes, or for tracking the amount of gas
+/// that is being used.
+pub type GasHook = Arc<dyn Fn(Gas) -> Result<()> + Send + Sync>;
+
+/// Allows you to meter the amount of gas consumed by transaction(s).
+/// Note: This only works with transactions that resolve to [`crate::result::ExecutionFinalResult`]
+/// Example
+/// ```rust, ignore, no_run
+/// let mut worker = near_workspaces::sandbox().await?;
+/// let meter = GasMeter::now(&mut worker);
+///
+/// let wasm = std::fs::read(STATUS_MSG_WASM_FILEPATH)?;
+/// let contract = worker.dev_deploy(&wasm).await?;
+///
+/// contract
+///    .call("set_status")
+///    .args_json(json!({
+///        "message": "hello_world",
+///    }))
+///    .transact()
+///    .await?;
+///
+/// println!("Total Gas consumed: {}", meter.elapsed()?);
+/// ```
+pub struct GasMeter {
+    gas: Arc<Mutex<Gas>>,
+}
+
+impl GasMeter {
+    /// Create a new gas meter with 0 gas consumed.
+    pub fn now<T: ?Sized>(worker: &mut Worker<T>) -> Self {
+        let meter = Self {
+            gas: Arc::new(Mutex::new(Gas::from_gas(0))),
+        };
+
+        let gas_consumed = Arc::downgrade(&Arc::clone(&meter.gas));
+        worker.tx_callbacks.push(Arc::new(move |gas: Gas| {
+            // upgrades if meter is still alive, else noop.
+            _ = gas_consumed.upgrade().map(|consumed| {
+                let mut consumed = consumed.lock().expect("GasMeter lock poisoned");
+                *consumed = Gas::from_gas(consumed.as_gas() + gas.as_gas());
+            });
+
+            Ok(())
+        }));
+
+        meter
+    }
+
+    /// Get the total amount of gas consumed.
+    pub fn elapsed(&self) -> Result<Gas> {
+        let meter = self.gas.lock()?;
+        Ok(*meter)
+    }
+
+    /// Reset the gas consumed to 0.
+    pub fn reset(&self) -> Result<()> {
+        let mut meter = self.gas.lock()?;
+        *meter = Gas::from_gas(0);
+        Ok(())
+    }
+}

workspaces/src/types/mod.rs

@@ -5,6 +5,7 @@
 pub(crate) mod account;
 pub(crate) mod block;
 pub(crate) mod chunk;
+pub(crate) mod gas_meter;
 
 #[cfg(feature = "interop_sdk")]
 mod sdk;
@@ -27,6 +28,8 @@ use crate::result::Result;
 pub use self::account::{AccountDetails, AccountDetailsPatch};
 pub use self::chunk::{Chunk, ChunkHeader};
 
+pub use self::gas_meter::{GasHook, GasMeter};
+
 /// Nonce is a unit used to determine the order of transactions in the pool.
 pub type Nonce = u64;
 

workspaces/src/worker/impls.rs

@@ -33,6 +33,7 @@ impl<T: ?Sized> Clone for Worker<T> {
     fn clone(&self) -> Self {
         Self {
             workspace: self.workspace.clone(),
+            tx_callbacks: self.tx_callbacks.clone(),
         }
     }
 }

workspaces/src/worker/mod.rs

@@ -5,6 +5,7 @@ use std::sync::Arc;
 
 use crate::network::builder::NetworkBuilder;
 use crate::network::{Betanet, Custom, Mainnet, Sandbox, Testnet};
+use crate::types::GasHook;
 use crate::{Network, Result};
 
 /// The `Worker` type allows us to interact with any NEAR related networks, such
@@ -13,6 +14,7 @@ use crate::{Network, Result};
 /// deploying a contract, or interacting with transactions.
 pub struct Worker<T: ?Sized> {
     pub(crate) workspace: Arc<T>,
+    pub(crate) tx_callbacks: Vec<GasHook>,
 }
 
 impl<T> Worker<T>
@@ -22,6 +24,7 @@ where
     pub(crate) fn new(network: T) -> Self {
         Self {
             workspace: Arc::new(network),
+            tx_callbacks: vec![],
         }
     }
 }
@@ -30,6 +33,7 @@ impl<T: Network + 'static> Worker<T> {
     pub(crate) fn coerce(self) -> Worker<dyn Network> {
         Worker {
             workspace: self.workspace,
+            tx_callbacks: self.tx_callbacks,
         }
     }
 }