Retain the original address vector of a UtxosChanged mutation along t…

…he full chain of notifiers up to the IndexProcessor
tiram88 · Feb 5, 2024 · 1390eb7 · 1390eb7
1 parent e898e62
commit 1390eb7
Showing 1 changed file with 26 additions and 22 deletions.
diff --git a/notify/src/address/tracker.rs b/notify/src/address/tracker.rs
@@ -539,52 +539,56 @@ impl Tracker {
 
     pub fn register<T: Indexer>(&self, indexes: &T, mut addresses: Vec<Address>) -> Result<Vec<Address>> {
         let mut rollback: bool = false;
-        let mut added = Vec::with_capacity(addresses.len());
-        'outer: for chunk in &addresses.drain(..).chunks(Self::ADDRESS_CHUNK_SIZE) {
+        {
             let mut inner = self.inner.write();
-            for address in chunk {
-                let spk = pay_to_address_script(&address);
+            addresses.retain(|address| {
+                let spk = pay_to_address_script(address);
                 match inner.get_or_insert(spk) {
                     Ok(index) => {
                         if indexes.insert(index) {
-                            added.push(address);
                             inner.inc_count(index);
+                            true
+                        } else {
+                            false
                         }
                     }
                     Err(Error::MaxCapacityReached) => {
                         // Rollback registration
                         rollback = true;
-                        break 'outer;
+                        false
                     }
                 }
-            }
+            });
         }
         match rollback {
-            false => Ok(added),
+            false => Ok(addresses),
             true => {
-                let _ = self.unregister(indexes, added);
+                let _ = self.unregister(indexes, addresses);
                 Err(Error::MaxCapacityReached)
             }
         }
     }
 
     pub fn unregister<T: Indexer>(&self, indexes: &T, mut addresses: Vec<Address>) -> Vec<Address> {
-        let mut removed = Vec::with_capacity(addresses.len());
-        if !indexes.is_empty() {
-            for chunk in &addresses.drain(..).chunks(Self::ADDRESS_CHUNK_SIZE) {
-                let mut inner = self.inner.write();
-                for address in chunk {
-                    let spk = pay_to_address_script(&address);
-                    if let Some((index, _)) = inner.get(&spk) {
-                        if indexes.remove(index) {
-                            removed.push(address);
-                            inner.dec_count(index);
-                        }
+        if indexes.is_empty() {
+            vec![]
+        } else {
+            let mut inner = self.inner.write();
+            addresses.retain(|address| {
+                let spk = pay_to_address_script(address);
+                if let Some((index, _)) = inner.get(&spk) {
+                    if indexes.remove(index) {
+                        inner.dec_count(index);
+                        true
+                    } else {
+                        false
                     }
+                } else {
+                    false
                 }
-            }
+            });
+            addresses
         }
-        removed
     }
 
     pub fn unregister_indexes(&self, indexes: &Indexes) {