Add a trait for BankClient add-on methods

Also add Context parameters to the methods can freely move between
the BanksClient and BanksClientExt trait without requiring code
changes.

runtime/src/banks_service.rs

@@ -194,10 +194,7 @@ mod tests {
     use super::*;
     use crate::genesis_utils::create_genesis_config;
     use solana_sdk::{
-        banks_client::{get_balance, get_recent_blockhash, process_transaction},
-        message::Message,
-        pubkey::Pubkey,
-        signature::Signer,
+        banks_client::BanksClientExt, message::Message, pubkey::Pubkey, signature::Signer,
         system_instruction,
     };
     use tarpc::context;
@@ -219,12 +216,20 @@ mod tests {
 
         Runtime::new()?.block_on(async {
             let mut banks_client = start_local_service(&bank_forks).await?;
-            let recent_blockhash = get_recent_blockhash(&mut banks_client).await?;
+            let recent_blockhash = banks_client
+                .get_recent_blockhash(context::current())
+                .await?;
             let transaction = Transaction::new(&[&genesis.mint_keypair], message, recent_blockhash);
-            process_transaction(&mut banks_client, transaction)
+            banks_client
+                .process_transaction(context::current(), transaction)
                 .await
                 .unwrap();
-            assert_eq!(get_balance(&mut banks_client, bob_pubkey).await?, 1);
+            assert_eq!(
+                banks_client
+                    .get_balance(context::current(), bob_pubkey)
+                    .await?,
+                1
+            );
             Ok(())
         })
     }
@@ -269,7 +274,12 @@ mod tests {
                     .await?;
             }
             assert_eq!(status, Some(Ok(())));
-            assert_eq!(get_balance(&mut banks_client, bob_pubkey).await?, 1);
+            assert_eq!(
+                banks_client
+                    .get_balance(context::current(), bob_pubkey)
+                    .await?,
+                1
+            );
             Ok(())
         })
     }

sdk/Cargo.toml

@@ -31,6 +31,7 @@ default = [
 
 [dependencies]
 assert_matches = { version = "1.3.0", optional = true }
+async-trait = "0.1.36"
 bincode = "1.2.1"
 bs58 = "0.3.1"
 bv = { version = "0.11.1", features = ["serde"] }

sdk/src/banks_client.rs

@@ -8,8 +8,9 @@ use crate::{
     transaction::{self, Transaction},
     transport,
 };
+use async_trait::async_trait;
 use std::io::{self, Error, ErrorKind};
-use tarpc::{client, context, serde_transport::tcp};
+use tarpc::{client, context::Context, serde_transport::tcp};
 use tokio::net::ToSocketAddrs;
 use tokio_serde::formats::Bincode;
 
@@ -28,31 +29,46 @@ pub trait Banks {
     async fn get_account(pubkey: Pubkey) -> Option<Account>;
 }
 
-pub async fn start_tcp_client<T: ToSocketAddrs>(addr: T) -> io::Result<BanksClient> {
-    let transport = tcp::connect(addr, Bincode::default()).await?;
-    BanksClient::new(client::Config::default(), transport).spawn()
+#[async_trait]
+pub trait BanksClientExt {
+    async fn get_recent_blockhash(&mut self, _: Context) -> io::Result<Hash>;
+    async fn process_transaction(
+        &mut self,
+        _: Context,
+        transaction: Transaction,
+    ) -> transport::Result<()>;
+    async fn get_balance(&mut self, _: Context, pubkey: Pubkey) -> io::Result<u64>;
 }
 
-pub async fn get_recent_blockhash(banks_client: &mut BanksClient) -> io::Result<Hash> {
-    Ok(banks_client.get_fees(context::current()).await?.1)
-}
+#[async_trait]
+impl BanksClientExt for BanksClient {
+    async fn get_recent_blockhash(&mut self, context: Context) -> io::Result<Hash> {
+        Ok(self.get_fees(context).await?.1)
+    }
 
-pub async fn process_transaction(
-    banks_client: &mut BanksClient,
-    transaction: Transaction,
-) -> transport::Result<()> {
-    let result = banks_client
-        .send_and_confirm_transaction(context::current(), transaction)
-        .await?;
-    match result {
-        None => Err(Error::new(ErrorKind::TimedOut, "invalid blockhash or fee-payer").into()),
-        Some(transaction_result) => Ok(transaction_result?),
+    async fn process_transaction(
+        &mut self,
+        context: Context,
+        transaction: Transaction,
+    ) -> transport::Result<()> {
+        let result = self
+            .send_and_confirm_transaction(context, transaction)
+            .await?;
+        match result {
+            None => Err(Error::new(ErrorKind::TimedOut, "invalid blockhash or fee-payer").into()),
+            Some(transaction_result) => Ok(transaction_result?),
+        }
+    }
+
+    async fn get_balance(&mut self, context: Context, pubkey: Pubkey) -> io::Result<u64> {
+        let account = self.get_account(context, pubkey).await?;
+        Ok(account.map(|x| x.lamports).unwrap_or(0))
     }
 }
 
-pub async fn get_balance(banks_client: &mut BanksClient, pubkey: Pubkey) -> io::Result<u64> {
-    let account = banks_client.get_account(context::current(), pubkey).await?;
-    Ok(account.map(|x| x.lamports).unwrap_or(0))
+pub async fn start_tcp_client<T: ToSocketAddrs>(addr: T) -> io::Result<BanksClient> {
+    let transport = tcp::connect(addr, Bincode::default()).await?;
+    BanksClient::new(client::Config::default(), transport).spawn()
 }
 
 #[cfg(test)]

tokens/src/thin_client.rs

@@ -1,6 +1,6 @@
 use solana_sdk::{
     account::Account,
-    banks_client::{get_balance, BanksClient},
+    banks_client::{BanksClient, BanksClientExt},
     clock::Slot,
     fee_calculator::FeeCalculator,
     hash::Hash,
@@ -105,7 +105,7 @@ impl ThinClient {
     }
 
     pub async fn get_balance(&mut self, pubkey: &Pubkey) -> io::Result<u64> {
-        get_balance(&mut self.client, *pubkey).await
+        self.client.get_balance(context::current(), *pubkey).await
     }
 
     pub async fn get_account(&mut self, pubkey: &Pubkey) -> io::Result<Option<Account>> {