From ea192b3c83745aad924dcdc4eab2beefb8c25df5 Mon Sep 17 00:00:00 2001
From: "mergify[bot]" <37929162+mergify[bot]@users.noreply.github.com>
Date: Mon, 19 Jul 2021 20:15:40 +0000
Subject: [PATCH] prevent excess allocation with AccountsIndexIterator (#18605)
 (#18642)

(cherry picked from commit 0bd8710d3454416a3d13f8d1c90342a61bf1a5b8)

Co-authored-by: Jeff Washington (jwash) <75863576+jeffwashington@users.noreply.github.com>
---
 runtime/src/accounts_index.rs | 25 ++++++++++++-------------
 1 file changed, 12 insertions(+), 13 deletions(-)

diff --git a/runtime/src/accounts_index.rs b/runtime/src/accounts_index.rs
index 82f7073af2..b7d08cb130 100644
--- a/runtime/src/accounts_index.rs
+++ b/runtime/src/accounts_index.rs
@@ -566,19 +566,18 @@ impl<'a, T: 'static + Clone> Iterator for AccountsIndexIterator<'a, T> {
             return None;
         }
 
-        let chunk: Vec<(Pubkey, AccountMapEntry<T>)> = self
-            .account_maps
-            .iter()
-            .map(|i| {
-                i.read()
-                    .unwrap()
-                    .range((self.start_bound, self.end_bound))
-                    .map(|(pubkey, account_map_entry)| (*pubkey, account_map_entry.clone()))
-                    .collect::<Vec<_>>()
-            })
-            .flatten()
-            .take(ITER_BATCH_SIZE)
-            .collect();
+        let mut chunk: Vec<(Pubkey, AccountMapEntry<T>)> = Vec::with_capacity(ITER_BATCH_SIZE);
+        'outer: for i in self.account_maps.iter() {
+            for (pubkey, account_map_entry) in
+                i.read().unwrap().range((self.start_bound, self.end_bound))
+            {
+                if chunk.len() >= ITER_BATCH_SIZE {
+                    break 'outer;
+                }
+                let item = (*pubkey, account_map_entry.clone());
+                chunk.push(item);
+            }
+        }
 
         if chunk.is_empty() {
             self.is_finished = true;