Added poll_balance_with_timeout method (#1062)

* Added poll_balance_with_timeout method

- updated bench-tps, fullnode and wallet to use this method instead
  of repeatedly calling poll_get_balance()

* Address review comments

- Revert some changes to use wrapper poll_get_balance()

* Reverting bench-tps to use poll_get_balance

- The original code is checking if the balance has been updated,
  instead of just retrieving the balance. The logic is different
  than poll_balance_with_timeout()

* Reverting wallet to use poll_get_balance

- The break condition in the loop is different than poll_balance_with_timeout().
  It's checking if the balance has been updated.

src/bin/bench-tps.rs

@@ -143,7 +143,9 @@ fn send_barrier_transaction(barrier_client: &mut ThinClient, last_id: &mut Hash,
             );
 
             // Sanity check that the client balance is still 1
-            let balance = barrier_client.poll_get_balance(&id.pubkey()).unwrap_or(-1);
+            let balance = barrier_client
+                .poll_get_balance(&id.pubkey())
+                .expect("Failed to get balance");
             if balance != 1 {
                 panic!("Expected an account balance of 1 (balance: {}", balance);
             }
@@ -285,15 +287,16 @@ fn airdrop_tokens(client: &mut ThinClient, leader: &NodeInfo, id: &Keypair, tx_c
             airdrop_amount, drone_addr
         );
 
-        let previous_balance = starting_balance;
         request_airdrop(&drone_addr, &id.pubkey(), airdrop_amount as u64).unwrap();
 
         // TODO: return airdrop Result from Drone instead of polling the
         //       network
-        let mut current_balance = previous_balance;
+        let mut current_balance = starting_balance;
         for _ in 0..20 {
             sleep(Duration::from_millis(500));
-            current_balance = client.poll_get_balance(&id.pubkey()).unwrap();
+            current_balance = client
+                .poll_get_balance(&id.pubkey())
+                .unwrap_or(starting_balance);
             if starting_balance != current_balance {
                 break;
             }

src/bin/fullnode.rs

@@ -18,6 +18,7 @@ use solana::wallet::request_airdrop;
 use std::fs::File;
 use std::net::{IpAddr, Ipv4Addr, SocketAddr};
 use std::process::exit;
+use std::time::Duration;
 
 fn main() -> () {
     logger::setup();
@@ -114,14 +115,13 @@ fn main() -> () {
             )
         });
 
-        // Try multiple times to confirm a non-zero balance.  |poll_get_balance| currently times
-        // out after 1 second, and sometimes this is not enough time while the network is
-        // booting
-        let balance_ok = (0..30).any(|i| {
-            let balance = client.poll_get_balance(&leader_pubkey).unwrap_or(0);
-            eprintln!("new balance is {} (attempt #{})", balance, i);
-            balance > 0
-        });
+        let balance_ok = client
+            .poll_balance_with_timeout(
+                &leader_pubkey,
+                &Duration::from_millis(100),
+                &Duration::from_secs(30),
+            )
+            .unwrap() > 0;
         assert!(balance_ok, "0 balance, airdrop failed?");
     }
 

src/thin_client.rs

@@ -242,37 +242,49 @@ impl ThinClient {
         self.last_id.expect("some last_id")
     }
 
-    pub fn poll_get_balance(&mut self, pubkey: &Pubkey) -> io::Result<i64> {
-        let mut balance_result;
-        let mut balance_value = -1;
-        let now = Instant::now();
-        loop {
-            balance_result = self.get_balance(pubkey);
-            if balance_result.is_ok() {
-                balance_value = *balance_result.as_ref().unwrap();
-            }
-            if balance_value > 0 || now.elapsed().as_secs() > 1 {
-                break;
-            }
-            sleep(Duration::from_millis(100));
-        }
+    pub fn submit_poll_balance_metrics(elapsed: &Duration) {
         metrics::submit(
             influxdb::Point::new("thinclient")
                 .add_tag("op", influxdb::Value::String("get_balance".to_string()))
                 .add_field(
                     "duration_ms",
-                    influxdb::Value::Integer(timing::duration_as_ms(&now.elapsed()) as i64),
+                    influxdb::Value::Integer(timing::duration_as_ms(elapsed) as i64),
                 )
                 .to_owned(),
         );
-        if balance_value >= 0 {
-            Ok(balance_value)
-        } else {
-            assert!(balance_result.is_err());
-            balance_result
+    }
+
+    pub fn poll_balance_with_timeout(
+        &mut self,
+        pubkey: &Pubkey,
+        polling_frequency: &Duration,
+        timeout: &Duration,
+    ) -> io::Result<i64> {
+        let now = Instant::now();
+        loop {
+            let balance = match self.get_balance(&pubkey) {
+                Ok(bal) => bal,
+                Err(e) => {
+                    sleep(*polling_frequency);
+                    if now.elapsed() > *timeout {
+                        ThinClient::submit_poll_balance_metrics(&now.elapsed());
+                        return Err(e);
+                    }
+                    -1
+                }
+            };
+
+            if balance >= 0 {
+                ThinClient::submit_poll_balance_metrics(&now.elapsed());
+                return Ok(balance);
+            }
         }
     }
 
+    pub fn poll_get_balance(&mut self, pubkey: &Pubkey) -> io::Result<i64> {
+        self.poll_balance_with_timeout(pubkey, &Duration::from_millis(100), &Duration::from_secs(1))
+    }
+
     /// Poll the server to confirm a transaction.
     pub fn poll_for_signature(&mut self, signature: &Signature) -> io::Result<()> {
         let now = Instant::now();