126 atomic balances

src/accountant.rs

@@ -14,6 +14,7 @@ use rayon::prelude::*;
 use signature::{KeyPair, PublicKey, Signature};
 use std::collections::hash_map::Entry::Occupied;
 use std::collections::{HashMap, HashSet, VecDeque};
+use std::sync::atomic::{AtomicIsize, Ordering};
 use std::result;
 use std::sync::RwLock;
 use transaction::Transaction;
@@ -30,18 +31,18 @@ pub enum AccountingError {
 pub type Result<T> = result::Result<T, AccountingError>;
 
 /// Commit funds to the 'to' party.
-fn apply_payment(balances: &RwLock<HashMap<PublicKey, RwLock<i64>>>, payment: &Payment) {
+fn apply_payment(balances: &RwLock<HashMap<PublicKey, AtomicIsize>>, payment: &Payment) {
     if balances.read().unwrap().contains_key(&payment.to) {
         let bals = balances.read().unwrap();
-        *bals[&payment.to].write().unwrap() += payment.tokens;
+        bals[&payment.to].fetch_add(payment.tokens as isize, Ordering::Relaxed);
     } else {
         let mut bals = balances.write().unwrap();
-        bals.insert(payment.to, RwLock::new(payment.tokens));
+        bals.insert(payment.to, AtomicIsize::new(payment.tokens as isize));
     }
 }
 
 pub struct Accountant {
-    balances: RwLock<HashMap<PublicKey, RwLock<i64>>>,
+    balances: RwLock<HashMap<PublicKey, AtomicIsize>>,
     pending: RwLock<HashMap<Signature, Plan>>,
     last_ids: RwLock<VecDeque<(Hash, RwLock<HashSet<Signature>>)>>,
     time_sources: RwLock<HashSet<PublicKey>>,
@@ -127,27 +128,37 @@ impl Accountant {
     /// funds and isn't a duplicate.
     pub fn process_verified_transaction_debits(&self, tr: &Transaction) -> Result<()> {
         let bals = self.balances.read().unwrap();
-
-        // Hold a write lock before the condition check, so that a debit can't occur
-        // between checking the balance and the withdraw.
         let option = bals.get(&tr.from);
+
         if option.is_none() {
             return Err(AccountingError::AccountNotFound);
         }
-        let mut bal = option.unwrap().write().unwrap();
 
         if !self.reserve_signature_with_last_id(&tr.sig, &tr.data.last_id) {
             return Err(AccountingError::InvalidTransferSignature);
         }
 
-        if *bal < tr.data.tokens {
-            self.forget_signature_with_last_id(&tr.sig, &tr.data.last_id);
-            return Err(AccountingError::InsufficientFunds);
-        }
+        loop {
+            let bal = option.unwrap();
+            let current = bal.load(Ordering::Relaxed) as i64;
+
+            if current < tr.data.tokens {
+                self.forget_signature_with_last_id(&tr.sig, &tr.data.last_id);
+                return Err(AccountingError::InsufficientFunds);
+            }
 
-        *bal -= tr.data.tokens;
+            let result = bal.compare_exchange(
+                current as isize,
+                (current - tr.data.tokens) as isize,
+                Ordering::Relaxed,
+                Ordering::Relaxed,
+            );
 
-        Ok(())
+            match result {
+                Ok(_) => return Ok(()),
+                Err(_) => continue,
+            };
+        }
     }
 
     pub fn process_verified_transaction_credits(&self, tr: &Transaction) {
@@ -300,7 +311,7 @@ impl Accountant {
 
     pub fn get_balance(&self, pubkey: &PublicKey) -> Option<i64> {
         let bals = self.balances.read().unwrap();
-        bals.get(pubkey).map(|x| *x.read().unwrap())
+        bals.get(pubkey).map(|x| x.load(Ordering::Relaxed) as i64)
     }
 }
 

src/accountant_skel.rs

@@ -774,11 +774,11 @@ mod bench {
                 // Seed the 'from' account.
                 let rando0 = KeyPair::new();
                 let tr = Transaction::new(&mint.keypair(), rando0.pubkey(), 1_000, last_id);
-                acc.process_verified_transaction(&tr).unwrap();
+                let _ = acc.process_verified_transaction(&tr);
 
                 let rando1 = KeyPair::new();
                 let tr = Transaction::new(&rando0, rando1.pubkey(), 2, last_id);
-                acc.process_verified_transaction(&tr).unwrap();
+                let _ = acc.process_verified_transaction(&tr);
 
                 // Finally, return a transaction that's unique
                 Transaction::new(&rando0, rando1.pubkey(), 1, last_id)
@@ -803,7 +803,10 @@ mod bench {
         drop(skel.historian.sender);
         let entries: Vec<Entry> = skel.historian.receiver.iter().collect();
         assert_eq!(entries.len(), 1);
-        assert_eq!(entries[0].events.len(), txs as usize);
+        assert_eq!(
+            entries[0].events.len(),
+            txs as usize
+        );
 
         println!("{} tps", tps);
     }