From 897843b06aa286329b9f0222b37dd6490ee1c624 Mon Sep 17 00:00:00 2001
From: Abraham Egnor <abraham.egnor@mongodb.com>
Date: Thu, 26 Oct 2023 10:15:58 -0400
Subject: [PATCH] RUST-1779 Fix a memory leak in cleanup tracking (#979)

---
 src/client.rs |  9 +++------
 src/id_set.rs | 54 +++++++++++++++++++++++++++++++++++----------------
 2 files changed, 40 insertions(+), 23 deletions(-)

diff --git a/src/client.rs b/src/client.rs
index 1ee9a134e..f558d624a 100644
--- a/src/client.rs
+++ b/src/client.rs
@@ -503,12 +503,9 @@ impl Client {
             // await points in between.
             let id = id_rx.await.unwrap();
             // If the cleanup channel is closed, that task was dropped.
-            let cleanup = if let Ok(f) = cleanup_rx.await {
-                f
-            } else {
-                return;
-            };
-            cleanup.await;
+            if let Ok(cleanup) = cleanup_rx.await {
+                cleanup.await;
+            }
             if let Some(client) = weak.upgrade() {
                 client
                     .inner
diff --git a/src/id_set.rs b/src/id_set.rs
index d7723799b..14d4cc90c 100644
--- a/src/id_set.rs
+++ b/src/id_set.rs
@@ -1,42 +1,62 @@
-use std::collections::HashMap;
-
-/// A set that provides removal tokens when an item is added.
+/// A container that provides removal tokens when an item is added.
 #[derive(Debug, Clone)]
 pub(crate) struct IdSet<T> {
-    values: HashMap<u32, T>,
-    // Incrementing a counter is not the best source of tokens - it can
-    // cause poor hash behavior - but efficiency is not an immediate concern.
-    next_id: u32,
+    values: Vec<Entry<T>>,
+    free: Vec<usize>,
+}
+
+#[derive(Debug, Clone)]
+struct Entry<T> {
+    generation: u32,
+    value: Option<T>,
 }
 
 #[derive(Debug, Clone, PartialEq, Eq)]
-pub(crate) struct Id(u32);
+pub(crate) struct Id {
+    index: usize,
+    generation: u32,
+}
 
 impl<T> IdSet<T> {
     pub(crate) fn new() -> Self {
         Self {
-            values: HashMap::new(),
-            next_id: 0,
+            values: vec![],
+            free: vec![],
         }
     }
 
     pub(crate) fn insert(&mut self, value: T) -> Id {
-        let id = self.next_id;
-        self.next_id += 1;
-        self.values.insert(id, value);
-        Id(id)
+        let value = Some(value);
+        if let Some(index) = self.free.pop() {
+            let generation = self.values[index].generation + 1;
+            self.values[index] = Entry { generation, value };
+            Id { index, generation }
+        } else {
+            let generation = 0;
+            self.values.push(Entry { generation, value });
+            Id {
+                index: self.values.len() - 1,
+                generation,
+            }
+        }
     }
 
     pub(crate) fn remove(&mut self, id: &Id) {
-        self.values.remove(&id.0);
+        if let Some(entry) = self.values.get_mut(id.index) {
+            if entry.generation == id.generation {
+                entry.value = None;
+                self.free.push(id.index);
+            }
+        }
     }
 
     #[cfg(all(test, mongodb_internal_tracking_arc))]
     pub(crate) fn values(&self) -> impl Iterator<Item = &T> {
-        self.values.values()
+        self.values.iter().filter_map(|e| e.value.as_ref())
     }
 
     pub(crate) fn extract(&mut self) -> Vec<T> {
-        self.values.drain().map(|(_, v)| v).collect()
+        self.free.clear();
+        self.values.drain(..).filter_map(|e| e.value).collect()
     }
 }