max num batches for quota (aptos-labs#7176)

config/src/config/quorum_store_config.rs

@@ -57,6 +57,7 @@ pub struct QuorumStoreConfig {
     pub batch_expiry_grace_rounds: Round,
     pub memory_quota: usize,
     pub db_quota: usize,
+    pub batch_quota: usize,
     pub mempool_txn_pull_max_bytes: u64,
     pub back_pressure: QuorumStoreBackPressureConfig,
     pub num_workers_for_remote_fragments: usize,
@@ -81,6 +82,7 @@ impl Default for QuorumStoreConfig {
             batch_expiry_grace_rounds: 5,
             memory_quota: 120_000_000,
             db_quota: 300_000_000,
+            batch_quota: 300_000,
             mempool_txn_pull_max_bytes: 4 * 1024 * 1024,
             back_pressure: QuorumStoreBackPressureConfig::default(),
             // number of batch coordinators to handle QS Fragment messages, should be >= 1

consensus/src/quorum_store/batch_store.rs

@@ -64,28 +64,39 @@ enum StorageMode {
 struct QuotaManager {
     memory_balance: usize,
     db_balance: usize,
+    batch_balance: usize,
     memory_quota: usize,
     db_quota: usize,
+    batch_quota: usize,
 }
 
 impl QuotaManager {
-    fn new(max_db_balance: usize, max_memory_balance: usize) -> Self {
-        assert!(max_db_balance >= max_memory_balance);
+    fn new(db_quota: usize, memory_quota: usize, batch_quota: usize) -> Self {
+        assert!(db_quota >= memory_quota);
         Self {
             memory_balance: 0,
             db_balance: 0,
-            memory_quota: max_memory_balance,
-            db_quota: max_db_balance,
+            batch_balance: 0,
+            memory_quota,
+            db_quota,
+            batch_quota,
         }
     }
 
     pub(crate) fn update_quota(&mut self, num_bytes: usize) -> anyhow::Result<StorageMode> {
+        if self.batch_balance + 1 > self.batch_quota {
+            counters::EXCEEDED_BATCH_QUOTA_COUNT.inc();
+            bail!("Batch quota exceeded ");
+        }
+
         if self.memory_balance + num_bytes <= self.memory_quota {
             self.memory_balance += num_bytes;
             self.db_balance += num_bytes;
+            self.batch_balance += 1;
             Ok(StorageMode::MemoryAndPersisted)
         } else if self.db_balance + num_bytes <= self.db_quota {
             self.db_balance += num_bytes;
+            self.batch_balance += 1;
             Ok(StorageMode::PersistedOnly)
         } else {
             counters::EXCEEDED_STORAGE_QUOTA_COUNT.inc();
@@ -94,13 +105,12 @@ impl QuotaManager {
     }
 
     pub(crate) fn free_quota(&mut self, num_bytes: usize, storage_mode: StorageMode) {
+        self.db_balance -= num_bytes;
+        self.batch_balance -= 1;
         match storage_mode {
-            StorageMode::PersistedOnly => {
-                self.db_balance -= num_bytes;
-            },
+            StorageMode::PersistedOnly => {},
             StorageMode::MemoryAndPersisted => {
                 self.memory_balance -= num_bytes;
-                self.db_balance -= num_bytes;
             },
         }
     }
@@ -128,6 +138,7 @@ pub struct BatchStore<T> {
     expiry_grace_rounds: Round,
     memory_quota: usize,
     db_quota: usize,
+    batch_quota: usize,
     batch_requester: BatchRequester<T>,
     validator_signer: ValidatorSigner,
     validator_verifier: ValidatorVerifier,
@@ -144,6 +155,7 @@ impl<T: QuorumStoreSender + Clone + Send + Sync + 'static> BatchStore<T> {
         expiry_grace_rounds: Round,
         memory_quota: usize,
         db_quota: usize,
+        batch_quota: usize,
         batch_requester: BatchRequester<T>,
         validator_signer: ValidatorSigner,
         validator_verifier: ValidatorVerifier,
@@ -162,6 +174,7 @@ impl<T: QuorumStoreSender + Clone + Send + Sync + 'static> BatchStore<T> {
             expiry_grace_rounds,
             memory_quota,
             db_quota,
+            batch_quota,
             batch_requester,
             validator_signer,
             validator_verifier,
@@ -253,7 +266,11 @@ impl<T: QuorumStoreSender + Clone + Send + Sync + 'static> BatchStore<T> {
             if self
                 .peer_quota
                 .entry(author)
-                .or_insert(QuotaManager::new(self.db_quota, self.memory_quota))
+                .or_insert(QuotaManager::new(
+                    self.db_quota,
+                    self.memory_quota,
+                    self.batch_quota,
+                ))
                 .update_quota(value.num_bytes)?
                 == StorageMode::PersistedOnly
             {

consensus/src/quorum_store/counters.rs

@@ -356,6 +356,15 @@ pub static EXCEEDED_STORAGE_QUOTA_COUNT: Lazy<IntCounter> = Lazy::new(|| {
     .unwrap()
 });
 
+/// Count of the exceeded batch quota.
+pub static EXCEEDED_BATCH_QUOTA_COUNT: Lazy<IntCounter> = Lazy::new(|| {
+    register_int_counter!(
+        "quorum_store_exceeded_batch_quota_count",
+        "Count of the exceeded batch quota."
+    )
+    .unwrap()
+});
+
 /// Count of the number of batch request sent to other nodes.
 pub static GET_BATCH_FROM_DB_COUNT: Lazy<IntCounter> = Lazy::new(|| {
     register_int_counter!(

consensus/src/quorum_store/quorum_store_builder.rs

@@ -257,6 +257,7 @@ impl InnerBuilder {
             self.config.batch_expiry_grace_rounds,
             self.config.memory_quota,
             self.config.db_quota,
+            self.config.batch_quota,
             batch_requester,
             signer,
             self.verifier.clone(),

consensus/src/quorum_store/tests/batch_store_test.rs

@@ -43,6 +43,7 @@ fn batch_store_for_test_no_db(memory_quota: usize) -> Arc<BatchStore<MockQuorumS
         0,            // grace period rounds
         memory_quota, // memory_quota
         1000,         // db quota
+        1000,         //batch quota
         requester,
         signers[0].clone(),
         validator_verifier,