From d00453f7472c5c8f534a94cd2af25488010c15fd Mon Sep 17 00:00:00 2001
From: Jack May <jack@solana.com>
Date: Fri, 25 Sep 2020 11:08:10 -0700
Subject: [PATCH] Drain the entire compute budget (#12478)

---
 programs/bpf/benches/bpf_loader.rs  | 3 ++-
 programs/bpf_loader/src/lib.rs      | 3 ++-
 programs/bpf_loader/src/syscalls.rs | 2 +-
 runtime/src/message_processor.rs    | 3 ++-
 4 files changed, 7 insertions(+), 4 deletions(-)

diff --git a/programs/bpf/benches/bpf_loader.rs b/programs/bpf/benches/bpf_loader.rs
index cc08db8f21..7c8885ed17 100644
--- a/programs/bpf/benches/bpf_loader.rs
+++ b/programs/bpf/benches/bpf_loader.rs
@@ -286,8 +286,9 @@ pub struct MockComputeMeter {
 }
 impl ComputeMeter for MockComputeMeter {
     fn consume(&mut self, amount: u64) -> Result<(), InstructionError> {
+        let exceeded = self.remaining < amount;
         self.remaining = self.remaining.saturating_sub(amount);
-        if self.remaining == 0 {
+        if exceeded {
             return Err(InstructionError::ComputationalBudgetExceeded);
         }
         Ok(())
diff --git a/programs/bpf_loader/src/lib.rs b/programs/bpf_loader/src/lib.rs
index 5e0319444d..f0e1a548c9 100644
--- a/programs/bpf_loader/src/lib.rs
+++ b/programs/bpf_loader/src/lib.rs
@@ -287,8 +287,9 @@ mod tests {
     }
     impl ComputeMeter for MockComputeMeter {
         fn consume(&mut self, amount: u64) -> Result<(), InstructionError> {
+            let exceeded = self.remaining < amount;
             self.remaining = self.remaining.saturating_sub(amount);
-            if self.remaining == 0 {
+            if exceeded {
                 return Err(InstructionError::ComputationalBudgetExceeded);
             }
             Ok(())
diff --git a/programs/bpf_loader/src/syscalls.rs b/programs/bpf_loader/src/syscalls.rs
index 53733211c9..e9f85ba161 100644
--- a/programs/bpf_loader/src/syscalls.rs
+++ b/programs/bpf_loader/src/syscalls.rs
@@ -1269,7 +1269,7 @@ mod tests {
         let addr = string.as_ptr() as *const _ as u64;
 
         let compute_meter: Rc<RefCell<dyn ComputeMeter>> =
-            Rc::new(RefCell::new(MockComputeMeter { remaining: 4 }));
+            Rc::new(RefCell::new(MockComputeMeter { remaining: 3 }));
         let log = Rc::new(RefCell::new(vec![]));
         let logger: Rc<RefCell<dyn Logger>> =
             Rc::new(RefCell::new(MockLogger { log: log.clone() }));
diff --git a/runtime/src/message_processor.rs b/runtime/src/message_processor.rs
index e4fd8ac29d..8dd67e1e3a 100644
--- a/runtime/src/message_processor.rs
+++ b/runtime/src/message_processor.rs
@@ -187,8 +187,9 @@ pub struct ThisComputeMeter {
 }
 impl ComputeMeter for ThisComputeMeter {
     fn consume(&mut self, amount: u64) -> Result<(), InstructionError> {
+        let exceeded = self.remaining < amount;
         self.remaining = self.remaining.saturating_sub(amount);
-        if self.remaining == 0 {
+        if exceeded {
             return Err(InstructionError::ComputationalBudgetExceeded);
         }
         Ok(())