From 2557ee589b3dcbe5e61595335ae5dbe5b1e8c266 Mon Sep 17 00:00:00 2001
From: Andrew Fitzgerald <apfitzge@gmail.com>
Date: Tue, 27 Aug 2024 08:10:51 -0500
Subject: [PATCH] use u64 to track write-version, manual iteration

---
 accounts-db/src/accounts_db.rs                    | 15 ++++++---------
 .../src/accounts_db/geyser_plugin_utils.rs        | 10 ++++------
 2 files changed, 10 insertions(+), 15 deletions(-)

diff --git a/accounts-db/src/accounts_db.rs b/accounts-db/src/accounts_db.rs
index 6bffa2c65a4f44..a5fd0a704be76f 100644
--- a/accounts-db/src/accounts_db.rs
+++ b/accounts-db/src/accounts_db.rs
@@ -90,6 +90,7 @@ use {
         hash::Hash,
         pubkey::Pubkey,
         rent_collector::RentCollector,
+        saturating_add_assign,
         timing::AtomicInterval,
         transaction::SanitizedTransaction,
     },
@@ -6811,14 +6812,9 @@ impl AccountsDb {
         accounts_and_meta_to_store: &impl StorableAccounts<'b>,
         mut txn_iter: impl Iterator<Item = &'a SanitizedTransaction>,
     ) -> Vec<AccountInfo> {
-        let mut write_version_producer = if self.accounts_update_notifier.is_some() {
-            let current_version = self
-                .write_version
-                .fetch_add(accounts_and_meta_to_store.len() as u64, Ordering::AcqRel);
-            current_version..current_version + accounts_and_meta_to_store.len() as u64
-        } else {
-            0..0 // dummy empty iterator
-        };
+        let mut current_write_version = self
+            .write_version
+            .fetch_add(accounts_and_meta_to_store.len() as u64, Ordering::AcqRel);
 
         let (account_infos, cached_accounts) = (0..accounts_and_meta_to_store.len())
             .map(|index| {
@@ -6834,8 +6830,9 @@ impl AccountsDb {
                         &account_shared_data,
                         txn,
                         pubkey,
-                        &mut write_version_producer,
+                        current_write_version,
                     );
+                    saturating_add_assign!(current_write_version, 1);
 
                     let cached_account =
                         self.accounts_cache.store(slot, pubkey, account_shared_data);
diff --git a/accounts-db/src/accounts_db/geyser_plugin_utils.rs b/accounts-db/src/accounts_db/geyser_plugin_utils.rs
index 6f94f4300de675..cedebafa4f08e7 100644
--- a/accounts-db/src/accounts_db/geyser_plugin_utils.rs
+++ b/accounts-db/src/accounts_db/geyser_plugin_utils.rs
@@ -57,23 +57,21 @@ impl AccountsDb {
         notify_stats.report();
     }
 
-    pub fn notify_account_at_accounts_update<P>(
+    pub fn notify_account_at_accounts_update(
         &self,
         slot: Slot,
         account: &AccountSharedData,
         txn: &Option<&SanitizedTransaction>,
         pubkey: &Pubkey,
-        write_version_producer: &mut P,
-    ) where
-        P: Iterator<Item = u64>,
-    {
+        write_version: u64,
+    ) {
         if let Some(accounts_update_notifier) = &self.accounts_update_notifier {
             accounts_update_notifier.notify_account_update(
                 slot,
                 account,
                 txn,
                 pubkey,
-                write_version_producer.next().unwrap(),
+                write_version,
             );
         }
     }