From 994e31032febea660d9338ade66d24169b5af823 Mon Sep 17 00:00:00 2001
From: Andrew Noyes <andrew@weaselab.dev>
Date: Thu, 11 Sep 2025 11:32:59 -0400
Subject: [PATCH] Fix data race in freeing control block

---
 src/reference.hpp        | 27 ++++++++++---------
 tests/test_reference.cpp | 58 ++++++++++++++++++++++++++++++++++++++++
 2 files changed, 73 insertions(+), 12 deletions(-)

diff --git a/src/reference.hpp b/src/reference.hpp
index e90a21a..7a191e8 100644
--- a/src/reference.hpp
+++ b/src/reference.hpp
@@ -172,19 +172,24 @@ private:
    */
   void release() noexcept {
     if (control_block) {
+      // Prevent weak count from going to 0 concurrently
+      control_block->increment_weak();
+
       uint32_t prev_strong = control_block->decrement_strong();
 
       // If this was the last strong reference, destroy the object
       if (prev_strong == 1) {
         T *obj = get();
         obj->~T();
+      }
 
-        // Check if there are any weak references
-        uint32_t current_weak =
-            control_block->weak_count.load(std::memory_order_acquire);
-        if (current_weak == 0) {
-          std::free(control_block);
-        }
+      // Now release our temporary weak reference
+      uint32_t prev_weak = control_block->decrement_weak();
+
+      // If this was the last weak reference and we destroyed the object,
+      // we can free the control block
+      if (prev_weak == 1 && prev_strong == 1) {
+        std::free(control_block);
       }
     }
   }
@@ -314,16 +319,14 @@ private:
    */
   void release() noexcept {
     if (control_block) {
+      // Prevent strong count from going to 0 concurrently
+      auto strong = lock();
       uint32_t prev_weak = control_block->decrement_weak();
 
       // If this was the last weak reference, check if we need to free control
       // block
-      if (prev_weak == 1) {
-        uint32_t current_strong =
-            control_block->strong_count.load(std::memory_order_acquire);
-        if (current_strong == 0) {
-          std::free(control_block);
-        }
+      if (prev_weak == 1 && !strong) {
+        std::free(control_block);
       }
     }
   }
diff --git a/tests/test_reference.cpp b/tests/test_reference.cpp
index 1236772..7cd7d84 100644
--- a/tests/test_reference.cpp
+++ b/tests/test_reference.cpp
@@ -1,3 +1,4 @@
+#include <barrier>
 #include <doctest/doctest.h>
 #include <latch>
 #include <thread>
@@ -156,6 +157,63 @@ TEST_CASE("Ref thread safety") {
   }
 }
 
+TEST_CASE("Control block cleanup race condition test") {
+  // This test specifically targets the race condition where both
+  // the last strong reference and last weak reference are destroyed
+  // simultaneously, potentially causing double-free of control block
+
+  const int test_iterations = 10000;
+
+  // Shared state for passing references between threads
+  Ref<TestObject> ptr1;
+  WeakRef<TestObject> ptr2;
+  auto setup = [&]() {
+    ptr1 = make_ref<TestObject>(0);
+    ptr2 = ptr1;
+  };
+
+  // Barrier for synchronization - 2 participants (main thread + worker thread)
+  std::barrier sync_barrier{2};
+
+  std::thread worker_thread([&]() {
+    for (int iter = 0; iter < test_iterations; ++iter) {
+      // Wait for main thread to create the references
+      sync_barrier.arrive_and_wait();
+
+      // Worker thread destroys the weak reference simultaneously with main
+      // thread
+      ptr2.reset();
+
+      // Wait for next iteration
+      sync_barrier.arrive_and_wait();
+    }
+  });
+
+  for (int iter = 0; iter < test_iterations; ++iter) {
+    // Create references
+    setup();
+
+    // Both threads are ready - synchronize for simultaneous destruction
+    sync_barrier.arrive_and_wait();
+
+    // Main thread destroys the strong reference at the same time
+    // as worker thread destroys the weak reference
+    ptr1.reset();
+
+    // Wait for both destructions to complete
+    sync_barrier.arrive_and_wait();
+
+    // Clean up for next iteration
+    ptr1.reset();
+    ptr2.reset();
+  }
+
+  worker_thread.join();
+
+  // If we reach here without segfault/double-free, the test passes
+  // The bug would manifest as a crash or memory corruption
+}
+
 TEST_CASE("WeakRef prevents circular references") {
   SUBCASE("simple weak reference lifecycle") {
     WeakRef<TestObject> weak_ref;