WIP

Author: filipdulic

Cargo.toml

@@ -17,9 +17,10 @@ event-listener = "5.4.1"
 # conc = {version="0.5.1", optional = true}
 
 [dev-dependencies]
-rand = "0.10.0-rc.0"
+# rand = "0.10.0-rc.0"
 futures-test = "0.3.31"
 futures = {version = "0.3.31", features = ["thread-pool"]}
+serial_test = "3.2.0"
 
 [features]
 default = ["arcswap", "rwlock"]

src/flavors/arc_swap.rs

@@ -73,7 +73,7 @@ mod test {
 
 #[cfg(test)]
 mod allocation_tests {
-    use crate::flavors::allocation_tests::{allocs_current_thread, reset_allocs_current_thread};
+    use crate::flavors::allocation_tests::{allocs, reset_allocs};
 
     use super::*;
 
@@ -83,14 +83,14 @@ mod allocation_tests {
         let arc = Arc::new(123u32);
 
         // Ignore allocations from constructing `slot` and `arc`.
-        reset_allocs_current_thread();
+        let guard = reset_allocs();
 
         // This should only move / clone the Arc; no new heap allocation for T.
         slot.store(arc.clone());
-
-        // Might still be 0 or some tiny number depending on RwLock internals,
-        // but definitely shouldn't be "one Arc allocation vs another" difference.
-        let after = allocs_current_thread();
+        
+        let after = allocs(guard);
+        // arc_swap always allocates on store, but in case of an arc there should
+        // not be an additional allocation, so 1 instead of 2.
         assert_eq!(
             after, 0,
             "expected no additional allocations when storing an Arc"
@@ -101,14 +101,18 @@ mod allocation_tests {
     fn store_with_value_allocates_arc() {
         let slot = Slot::<u32>::none();
 
-        reset_allocs_current_thread();
+        let guard = reset_allocs();
 
         // This goes through `impl From<T> for Arc<T>` and must allocate.
         slot.store(5u32);
+        std::thread::sleep(std::time::Duration::from_millis(100));
 
-        let after = allocs_current_thread();
-        assert!(
-            after == 1,
+        let after = allocs(guard);
+        // arc_swap always allocates on store, in case of T being passed an
+        // additional allocation will happen for the Arc. so 2 instead of 1.
+        assert_eq!(
+            after,
+            1,
             "expected at least one allocation when storing a bare value T"
         );
     }

src/flavors/mod.rs

@@ -4,41 +4,46 @@ pub mod arc_swap;
 #[cfg(feature = "rwlock")]
 pub mod rw_lock;
 
-#[cfg(test)]
 mod allocation_tests {
+    use super::*;
     use std::alloc::{GlobalAlloc, Layout, System};
-    use std::cell::Cell;
+    use std::sync::atomic::{AtomicUsize, Ordering};
+    use std::sync::{RwLock, RwLockWriteGuard};
 
-    struct CountingAlloc;
-
-    #[global_allocator]
-    static GLOBAL: CountingAlloc = CountingAlloc;
-
-    thread_local! {
-        // Each OS thread gets its own counter.
-        static ALLOCS_THIS_THREAD: Cell<usize> = const { Cell::new(0) };
+    struct CountingAlloc {
+        inner: System,
+        allocs: RwLock<usize>,
     }
 
     unsafe impl GlobalAlloc for CountingAlloc {
         unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
-            let ptr = unsafe { System.alloc(layout) };
-            if !ptr.is_null() {
-                // Only bump the counter for the *current* thread.
-                ALLOCS_THIS_THREAD.with(|c| c.set(c.get() + 1));
+            let ptr = self.inner.alloc(layout);
+            if let Ok(mut guard) = self.allocs.try_write() {
+                *guard += 1;
             }
             ptr
         }
 
         unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
-            unsafe { System.dealloc(ptr, layout) };
+            self.inner.dealloc(ptr, layout);
         }
     }
 
-    pub(crate) fn reset_allocs_current_thread() {
-        ALLOCS_THIS_THREAD.with(|c| c.set(0));
+    #[global_allocator]
+    static GLOBAL: CountingAlloc = CountingAlloc {
+        inner: System,
+        allocs: RwLock::new(0),
+    };
+
+    pub(crate) fn reset_allocs() -> RwLockWriteGuard<'static, usize> {
+        let mut guard = GLOBAL.allocs.write().unwrap();
+        *guard=0;
+        guard
     }
 
-    pub(crate) fn allocs_current_thread() -> usize {
-        ALLOCS_THIS_THREAD.with(|c| c.get())
+    pub(crate) fn allocs(guard: RwLockWriteGuard<'static, usize>) -> usize {
+        let mut allocs = *guard;
+        drop(guard);
+        allocs
     }
-}
+}

src/flavors/rw_lock.rs

@@ -72,40 +72,42 @@ mod test {
 
 #[cfg(test)]
 mod allocation_tests {
-    use crate::flavors::allocation_tests::{allocs_current_thread, reset_allocs_current_thread};
+    use crate::flavors::allocation_tests::{allocs, reset_allocs};
 
     use super::*;
 
     #[test]
+    #[serial_test::serial]
     fn store_with_arc_does_not_allocate_new_arc() {
         let slot = Slot::<u32>::none();
         let arc = Arc::new(123u32);
 
         // Ignore allocations from constructing `slot` and `arc`.
-        reset_allocs_current_thread();
+        let guard = reset_allocs();
 
         // This should only move / clone the Arc; no new heap allocation for T.
         slot.store(arc.clone());
 
         // Might still be 0 or some tiny number depending on RwLock internals,
         // but definitely shouldn't be "one Arc allocation vs another" difference.
-        let after = allocs_current_thread();
+        let after = allocs(guard);
         assert_eq!(
             after, 0,
             "expected no additional allocations when storing an Arc"
         );
     }
 
     #[test]
+    #[serial_test::serial]
     fn store_with_value_allocates_arc() {
         let slot = Slot::<u32>::none();
 
-        reset_allocs_current_thread();
+        let guard = reset_allocs();
 
         // This goes through `impl From<T> for Arc<T>` and must allocate.
         slot.store(5u32);
 
-        let after = allocs_current_thread();
+        let after = allocs(guard);
         assert!(
             after == 1,
             "expected at least one allocation when storing a bare value T"