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 1510524309..d2c383d345 100644 --- a/programs/bpf_loader/src/syscalls.rs +++ b/programs/bpf_loader/src/syscalls.rs @@ -1260,7 +1260,7 @@ mod tests { let addr = string.as_ptr() as *const _ as u64; let compute_meter: Rc> = - Rc::new(RefCell::new(MockComputeMeter { remaining: 3 })); + Rc::new(RefCell::new(MockComputeMeter { remaining: 2 })); let log = Rc::new(RefCell::new(vec![])); let logger: Rc> = Rc::new(RefCell::new(MockLogger { log: log.clone() })); diff --git a/runtime/src/message_processor.rs b/runtime/src/message_processor.rs index 5ab1c86f2e..7c8c3db351 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(())