From 2dfc46c71e9efd5aa76f7c56a876ca14e45444b7 Mon Sep 17 00:00:00 2001
From: Andrew Fitzgerald <apfitzge@gmail.com>
Date: Tue, 28 Mar 2023 12:58:39 -0700
Subject: [PATCH] Forwarder: separate update_data_budget (#30920)

---
 core/src/banking_stage/forwarder.rs | 27 ++++++++++++++++-----------
 1 file changed, 16 insertions(+), 11 deletions(-)

diff --git a/core/src/banking_stage/forwarder.rs b/core/src/banking_stage/forwarder.rs
index c41aad0974..4699ee001e 100644
--- a/core/src/banking_stage/forwarder.rs
+++ b/core/src/banking_stage/forwarder.rs
@@ -162,17 +162,7 @@ impl Forwarder {
             None => return (Ok(()), 0, None),
         };
 
-        const INTERVAL_MS: u64 = 100;
-        // 12 MB outbound limit per second
-        const MAX_BYTES_PER_SECOND: usize = 12_000_000;
-        const MAX_BYTES_PER_INTERVAL: usize = MAX_BYTES_PER_SECOND * INTERVAL_MS as usize / 1000;
-        const MAX_BYTES_BUDGET: usize = MAX_BYTES_PER_INTERVAL * 5;
-        self.data_budget.update(INTERVAL_MS, |bytes| {
-            std::cmp::min(
-                bytes.saturating_add(MAX_BYTES_PER_INTERVAL),
-                MAX_BYTES_BUDGET,
-            )
-        });
+        self.update_data_budget();
 
         let packet_vec: Vec<_> = forwardable_packets
             .filter_map(|p| {
@@ -227,6 +217,21 @@ impl Forwarder {
 
         (Ok(()), packet_vec_len, Some(leader_pubkey))
     }
+
+    /// Re-fill the data budget if enough time has passed
+    fn update_data_budget(&self) {
+        const INTERVAL_MS: u64 = 100;
+        // 12 MB outbound limit per second
+        const MAX_BYTES_PER_SECOND: usize = 12_000_000;
+        const MAX_BYTES_PER_INTERVAL: usize = MAX_BYTES_PER_SECOND * INTERVAL_MS as usize / 1000;
+        const MAX_BYTES_BUDGET: usize = MAX_BYTES_PER_INTERVAL * 5;
+        self.data_budget.update(INTERVAL_MS, |bytes| {
+            std::cmp::min(
+                bytes.saturating_add(MAX_BYTES_PER_INTERVAL),
+                MAX_BYTES_BUDGET,
+            )
+        });
+    }
 }
 
 #[cfg(test)]