Add checked_add_signed() to apply cost adjustment to cost_tracker

core/src/banking_stage.rs

@@ -1407,9 +1407,6 @@ impl BankingStage {
             ..
         } = execute_and_commit_transactions_output;
 
-        // TODO: This does not revert the cost tracker changes from all unexecuted transactions
-        // yet: For example tx that are too old will not be included in the block, but are not
-        // retryable.
         QosService::update_or_remove_transaction_costs(
             transaction_costs.iter(),
             transactions_qos_results.iter(),

runtime/src/cost_tracker.rs

@@ -9,6 +9,43 @@ use {
     std::collections::HashMap,
 };
 
+/// Copied nightly-only experimental `checked_add_signed` implementation
+/// from https://github.com/rust-lang/rust/pull/87601/files
+/// Should be removed once the feature is added to standard library.
+trait ArithmeticUtil {
+    /// Checked addition with a signed integer. Computes `self + rhs`,
+    /// returning `None` if overflow occurred.
+    fn checked_add_signed(self, rhs: i64) -> Option<Self>
+    where
+        Self: Sized;
+
+    /// Calculates `self` + `rhs` with a signed `rhs`
+    ///
+    /// Returns a tuple of the addition along with a boolean indicating
+    /// whether an arithmetic overflow would occur. If an overflow would
+    /// have occurred then the wrapped value is returned.
+    fn overflowing_add_signed(self, rhs: i64) -> (Self, bool)
+    where
+        Self: Sized;
+}
+impl ArithmeticUtil for u64 {
+    #[inline]
+    fn checked_add_signed(self, rhs: i64) -> Option<Self> {
+        let (a, b) = ArithmeticUtil::overflowing_add_signed(self, rhs);
+        if b {
+            None
+        } else {
+            Some(a)
+        }
+    }
+
+    #[inline]
+    fn overflowing_add_signed(self, rhs: i64) -> (Self, bool) {
+        let (res, overflowed) = self.overflowing_add(rhs as Self);
+        (res, overflowed ^ (rhs < 0))
+    }
+}
+
 const WRITABLE_ACCOUNTS_PER_BLOCK: usize = 512;
 
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
@@ -106,8 +143,9 @@ impl CostTracker {
                 return;
             }
 
-            let adjustment =
-                i128::from(actual_execution_units) - i128::from(estimated_execution_units);
+            let (res, _overflowed) =
+                actual_execution_units.overflowing_sub(estimated_execution_units);
+            let adjustment = res as i64;
             self.adjust_transaction_execution_cost(estimated_tx_cost, adjustment);
         }
     }
@@ -252,20 +290,29 @@ impl CostTracker {
         self.transaction_count = self.transaction_count.saturating_sub(1);
     }
 
-    fn adjust_transaction_execution_cost(&mut self, tx_cost: &TransactionCost, adjustment: i128) {
+    /// Apply execution units adjustment to cost-tracker.
+    /// If adjustment causes arithmetic overflow, it bumps costs to their limits to
+    /// prevent adding further transactions into current block.
+    fn adjust_transaction_execution_cost(&mut self, tx_cost: &TransactionCost, adjustment: i64) {
         for account_key in tx_cost.writable_accounts.iter() {
             let account_cost = self
                 .cost_by_writable_accounts
                 .entry(*account_key)
                 .or_insert(0);
-            *account_cost = u64::try_from(i128::from(*account_cost).saturating_add(adjustment))
-                .unwrap_or(*account_cost);
+            match ArithmeticUtil::checked_add_signed(*account_cost, adjustment) {
+                Some(adjusted_cost) => *account_cost = adjusted_cost,
+                None => *account_cost = self.account_cost_limit,
+            }
+        }
+        match ArithmeticUtil::checked_add_signed(self.block_cost, adjustment) {
+            Some(adjusted_cost) => self.block_cost = adjusted_cost,
+            None => self.block_cost = self.block_cost_limit,
         }
-        self.block_cost = u64::try_from(i128::from(self.block_cost).saturating_add(adjustment))
-            .unwrap_or(self.block_cost);
         if tx_cost.is_simple_vote {
-            self.vote_cost = u64::try_from(i128::from(self.vote_cost).saturating_add(adjustment))
-                .unwrap_or(self.vote_cost);
+            match ArithmeticUtil::checked_add_signed(self.vote_cost, adjustment) {
+                Some(adjusted_cost) => self.vote_cost = adjusted_cost,
+                None => self.vote_cost = self.vote_cost_limit,
+            }
         }
     }
 }
@@ -276,6 +323,7 @@ mod tests {
         super::*,
         crate::{
             bank::Bank,
+            cost_tracker::ArithmeticUtil,
             genesis_utils::{create_genesis_config, GenesisConfigInfo},
         },
         solana_sdk::{
@@ -740,7 +788,7 @@ mod tests {
 
         // adjust up
         {
-            let adjustment: i128 = 50;
+            let adjustment: i64 = 50;
             testee.adjust_transaction_execution_cost(&tx_cost, adjustment);
             expected_block_cost += 50;
             assert_eq!(expected_block_cost, testee.block_cost());
@@ -755,7 +803,7 @@ mod tests {
 
         // adjust down
         {
-            let adjustment: i128 = -50;
+            let adjustment: i64 = -50;
             testee.adjust_transaction_execution_cost(&tx_cost, adjustment);
             expected_block_cost -= 50;
             assert_eq!(expected_block_cost, testee.block_cost());
@@ -767,5 +815,56 @@ mod tests {
                     assert_eq!(expected_block_cost, *units);
                 });
         }
+
+        // adjust overflow up
+        {
+            testee.adjust_transaction_execution_cost(&tx_cost, i64::MAX);
+            testee.adjust_transaction_execution_cost(&tx_cost, i64::MAX);
+            // expect block cost set to limit
+            assert_eq!(block_max, testee.block_cost());
+            assert_eq!(expected_tx_count, testee.transaction_count());
+            testee
+                .cost_by_writable_accounts
+                .iter()
+                .for_each(|(_key, units)| {
+                    assert_eq!(account_max, *units);
+                });
+        }
+
+        // adjust overflow down
+        {
+            testee.adjust_transaction_execution_cost(&tx_cost, i64::MIN);
+            testee.adjust_transaction_execution_cost(&tx_cost, i64::MIN);
+            // expect block cost set to limit
+            assert_eq!(block_max, testee.block_cost());
+            assert_eq!(expected_tx_count, testee.transaction_count());
+            testee
+                .cost_by_writable_accounts
+                .iter()
+                .for_each(|(_key, units)| {
+                    assert_eq!(account_max, *units);
+                });
+        }
+    }
+
+    #[test]
+    fn test_checked_add_signed() {
+        assert_eq!(ArithmeticUtil::checked_add_signed(1u64, 2), Some(3));
+        assert_eq!(ArithmeticUtil::checked_add_signed(3u64, -2), Some(1));
+        assert_eq!(ArithmeticUtil::checked_add_signed(1u64, -2), None);
+        assert_eq!(ArithmeticUtil::checked_add_signed(u64::MAX - 2, 3), None);
+    }
+
+    #[test]
+    fn test_overflowing_add_signed() {
+        assert_eq!(ArithmeticUtil::overflowing_add_signed(1u64, 2), (3, false));
+        assert_eq!(
+            ArithmeticUtil::overflowing_add_signed(1u64, -2),
+            (u64::MAX, true)
+        );
+        assert_eq!(
+            ArithmeticUtil::overflowing_add_signed(u64::MAX - 2, 4),
+            (1, true)
+        );
     }
 }