From ffb72136c89c3028a1e9535f1c3d5276857ca1db Mon Sep 17 00:00:00 2001
From: sakridge <sakridge@gmail.com>
Date: Tue, 4 Sep 2018 21:33:19 -0700
Subject: [PATCH] Remove account from balances table after error seen (#1120)

If balance goes to 0, then bank removes the account
from it's account table and returns no account error. Thin client
should also update the account to this state or it will
still have the cached balance from the last successful get_balance().
---
 src/thin_client.rs | 77 +++++++++++++++++++++++++++++++++++++++++-----
 1 file changed, 69 insertions(+), 8 deletions(-)

diff --git a/src/thin_client.rs b/src/thin_client.rs
index 5c77a8459d..27df1ce6c5 100644
--- a/src/thin_client.rs
+++ b/src/thin_client.rs
@@ -80,6 +80,7 @@ impl ThinClient {
             }
             Response::Account { key, account: None } => {
                 debug!("Response account {}: None ", key);
+                self.balances.remove(&key);
             }
             Response::LastId { id } => {
                 trace!("Response last_id {:?}", id);
@@ -283,22 +284,19 @@ impl ThinClient {
     ) -> io::Result<i64> {
         let now = Instant::now();
         loop {
-            let balance = match self.get_balance(&pubkey) {
-                Ok(bal) => bal,
+            match self.get_balance(&pubkey) {
+                Ok(bal) => {
+                    ThinClient::submit_poll_balance_metrics(&now.elapsed());
+                    return Ok(bal);
+                }
                 Err(e) => {
                     sleep(*polling_frequency);
                     if now.elapsed() > *timeout {
                         ThinClient::submit_poll_balance_metrics(&now.elapsed());
                         return Err(e);
                     }
-                    0
                 }
             };
-
-            if balance != 0 {
-                ThinClient::submit_poll_balance_metrics(&now.elapsed());
-                return Ok(balance);
-            }
         }
     }
 
@@ -586,4 +584,67 @@ mod tests {
         let mut client = ThinClient::new(addr, requests_socket, addr, transactions_socket);
         assert_eq!(client.transaction_count(), 0);
     }
+
+    #[test]
+    fn test_zero_balance_after_nonzero() {
+        logger::setup();
+        let leader_keypair = Keypair::new();
+        let leader = Node::new_localhost_with_pubkey(leader_keypair.pubkey());
+        let alice = Mint::new(10_000);
+        let bank = Bank::new(&alice);
+        let bob_keypair = Keypair::new();
+        let exit = Arc::new(AtomicBool::new(false));
+        let leader_data = leader.info.clone();
+        let ledger_path = tmp_ledger("zero_balance_check", &alice);
+
+        let server = Fullnode::new_with_bank(
+            leader_keypair,
+            bank,
+            0,
+            &[],
+            leader,
+            None,
+            exit.clone(),
+            Some(&ledger_path),
+            false,
+        );
+        sleep(Duration::from_millis(900));
+
+        let requests_socket = UdpSocket::bind("0.0.0.0:0").unwrap();
+        requests_socket
+            .set_read_timeout(Some(Duration::new(5, 0)))
+            .unwrap();
+        let transactions_socket = UdpSocket::bind("0.0.0.0:0").unwrap();
+        let mut client = ThinClient::new(
+            leader_data.contact_info.rpu,
+            requests_socket,
+            leader_data.contact_info.tpu,
+            transactions_socket,
+        );
+        let last_id = client.get_last_id();
+
+        // give bob 500 tokens
+        let signature = client
+            .transfer(500, &alice.keypair(), bob_keypair.pubkey(), &last_id)
+            .unwrap();
+        assert!(client.poll_for_signature(&signature).is_ok());
+
+        let balance = client.poll_get_balance(&bob_keypair.pubkey());
+        assert!(balance.is_ok());
+        assert_eq!(balance.unwrap(), 500);
+
+        // take them away
+        let signature = client
+            .transfer(500, &bob_keypair, alice.keypair().pubkey(), &last_id)
+            .unwrap();
+        assert!(client.poll_for_signature(&signature).is_ok());
+
+        // should get an error when bob's account is purged
+        let balance = client.poll_get_balance(&bob_keypair.pubkey());
+        assert!(balance.is_err());
+
+        exit.store(true, Ordering::Relaxed);
+        server.join().unwrap();
+        remove_dir_all(ledger_path).unwrap();
+    }
 }