From 7ac8031fc8b4def8712516e52025030390281f12 Mon Sep 17 00:00:00 2001
From: Jack Thomson <jackabt@amazon.com>
Date: Fri, 17 Jan 2025 17:11:07 +0000
Subject: [PATCH] Use fast-counter for size tracking

Update to track the size using the fast-counter crate which scales
better to a higher core count.
---
 Cargo.lock |  9 ++++++++-
 Cargo.toml |  1 +
 src/map.rs | 34 ++++++++++++++++++----------------
 3 files changed, 27 insertions(+), 17 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 29adda1..fbece89 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1,6 +1,6 @@
 # This file is automatically @generated by Cargo.
 # It is not intended for manual editing.
-version = 3
+version = 4
 
 [[package]]
 name = "ahash"
@@ -211,12 +211,19 @@ version = "1.10.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "11157ac094ffbdde99aa67b23417ebdd801842852b500e395a45a9c0aac03e4a"
 
+[[package]]
+name = "fast-counter"
+version = "1.0.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "2bf000f5d36bbfe94dd055f7fc5ec15e67cc674fd867a109b1c56577c37d8234"
+
 [[package]]
 name = "flurry"
 version = "0.5.2"
 dependencies = [
  "ahash",
  "criterion",
+ "fast-counter",
  "num_cpus",
  "parking_lot",
  "rand",
diff --git a/Cargo.toml b/Cargo.toml
index 373f224..58ca377 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -25,6 +25,7 @@ num_cpus = "1.12.0"
 rayon = {version = "1.3", optional = true}
 serde = {version = "1.0.105", optional = true}
 seize = "0.3.3"
+fast-counter = "1.0.0"
 
 [dependencies.ahash]
 version = "0.8.4"
diff --git a/src/map.rs b/src/map.rs
index 557b22b..89ace7d 100644
--- a/src/map.rs
+++ b/src/map.rs
@@ -1,3 +1,4 @@
+use fast_counter::ConcurrentCounter;
 use seize::Linked;
 
 use crate::iter::*;
@@ -92,7 +93,7 @@ pub struct HashMap<K, V, S = crate::DefaultHashBuilder> {
     /// The next table index (plus one) to split while resizing.
     transfer_index: AtomicIsize,
 
-    count: AtomicIsize,
+    counter: ConcurrentCounter,
 
     /// Table initialization and resizing control.  When negative, the
     /// table is being initialized or resized: -1 for initialization,
@@ -279,7 +280,7 @@ impl<K, V, S> HashMap<K, V, S> {
             table: Atomic::null(),
             next_table: Atomic::null(),
             transfer_index: AtomicIsize::new(0),
-            count: AtomicIsize::new(0),
+            counter: ConcurrentCounter::new(num_cpus()),
             size_ctl: AtomicIsize::new(0),
             build_hasher: hash_builder,
             collector: Collector::new(),
@@ -348,6 +349,12 @@ impl<K, V, S> HashMap<K, V, S> {
 
     /// Returns the number of entries in the map.
     ///
+    /// Note that the returned value is _NOT_ an
+    /// atomic snapshot; invocation in the absence of concurrent
+    /// updates returns an accurate result, but concurrent updates that
+    /// occur while the sum is being calculated might not be
+    /// incorporated.
+    ///
     /// # Examples
     ///
     /// ```
@@ -360,7 +367,7 @@ impl<K, V, S> HashMap<K, V, S> {
     /// assert!(map.pin().len() == 2);
     /// ```
     pub fn len(&self) -> usize {
-        let n = self.count.load(Ordering::Relaxed);
+        let n = self.counter.sum();
         if n < 0 {
             0
         } else {
@@ -1132,23 +1139,18 @@ where
     }
 
     fn add_count(&self, n: isize, resize_hint: Option<usize>, guard: &Guard<'_>) {
-        // TODO: implement the Java CounterCell business here
-
-        use std::cmp;
-        let mut count = match n.cmp(&0) {
-            cmp::Ordering::Greater => self.count.fetch_add(n, Ordering::SeqCst) + n,
-            cmp::Ordering::Less => self.count.fetch_sub(n.abs(), Ordering::SeqCst) - n,
-            cmp::Ordering::Equal => self.count.load(Ordering::SeqCst),
-        };
+        if n != 0 {
+            self.counter.add(n)
+        }
 
         // if resize_hint is None, it means the caller does not want us to consider a resize.
         // if it is Some(n), the caller saw n entries in a bin
-        if resize_hint.is_none() {
+        // TODO: use the resize hint
+        let Some(_saw_bin_length) = resize_hint else {
             return;
-        }
+        };
 
-        // TODO: use the resize hint
-        let _saw_bin_length = resize_hint.unwrap();
+        let mut count = self.counter.sum();
 
         loop {
             let sc = self.size_ctl.load(Ordering::SeqCst);
@@ -1208,7 +1210,7 @@ where
             }
 
             // another resize may be needed!
-            count = self.count.load(Ordering::SeqCst);
+            count = self.counter.sum();
         }
     }