fix: correct reference counting for Slice (#3998)

4paradigm · Dec 6, 2024 · cf44b96 · cf44b96
1 parent e6796f0
commit cf44b96
Show file tree

Hide file tree

Showing 2 changed files with 32 additions and 26 deletions.
diff --git a/hybridse/include/base/fe_slice.h b/hybridse/include/base/fe_slice.h
@@ -23,6 +23,7 @@
 #include <string.h>
 
 #include <string>
+#include <utility>
 
 #include "base/raw_buffer.h"
 
@@ -100,6 +101,12 @@ class Slice {
         return ((size_ >= x.size_) && (memcmp(data_, x.data_, x.size_) == 0));
     }
 
+ protected:
+    void _swap(Slice &slice) {
+        std::swap(data_, slice.data_);
+        std::swap(size_, slice.size_);
+    }
+
  private:
     uint32_t size_;
     const char *data_;
@@ -154,16 +161,17 @@ class RefCountedSlice : public Slice {
  private:
     RefCountedSlice(int8_t *data, size_t size, bool managed)
         : Slice(reinterpret_cast<const char *>(data), size),
-          ref_cnt_(managed ? new int(1) : nullptr) {}
+          ref_cnt_(managed ? new std::atomic<int32_t>(1) : nullptr) {}
 
     RefCountedSlice(const char *data, size_t size, bool managed)
-        : Slice(data, size), ref_cnt_(managed ? new int(1) : nullptr) {}
+        : Slice(data, size), ref_cnt_(managed ? new std::atomic<int32_t>(1) : nullptr) {}
 
-    void Release();
+    void _release();
 
-    void Update(const RefCountedSlice &slice);
+    void _copy(const RefCountedSlice &slice);
+    void _swap(RefCountedSlice &slice);
 
-    int32_t *ref_cnt_;
+    std::atomic<int32_t> *ref_cnt_;
 };
 
 }  // namespace base

diff --git a/hybridse/src/base/fe_slice.cc b/hybridse/src/base/fe_slice.cc
@@ -15,56 +15,54 @@
  */
 
 #include "base/fe_slice.h"
+#include <atomic>
 
 namespace hybridse {
 namespace base {
 
-RefCountedSlice::~RefCountedSlice() { Release(); }
+RefCountedSlice::~RefCountedSlice() { _release(); }
 
-void RefCountedSlice::Release() {
+void RefCountedSlice::_release() {
     if (this->ref_cnt_ != nullptr) {
-        auto& cnt = *this->ref_cnt_;
-        cnt -= 1;
-        if (cnt == 0 && buf() != nullptr) {
-            // memset in case the buf is still used after free
-            memset(buf(), 0, size());
+        if (std::atomic_fetch_add_explicit(ref_cnt_, -1, std::memory_order_release) == 1) {
+            std::atomic_thread_fence(std::memory_order_acquire);
+            assert(buf());
             free(buf());
             delete this->ref_cnt_;
         }
     }
 }
 
-void RefCountedSlice::Update(const RefCountedSlice& slice) {
+void RefCountedSlice::_copy(const RefCountedSlice& slice) {
     reset(slice.data(), slice.size());
-    this->ref_cnt_ = slice.ref_cnt_;
+    ref_cnt_ = slice.ref_cnt_;
     if (this->ref_cnt_ != nullptr) {
-        (*this->ref_cnt_) += 1;
+        std::atomic_fetch_add_explicit(ref_cnt_, 1, std::memory_order_relaxed);
     }
 }
 
-RefCountedSlice::RefCountedSlice(const RefCountedSlice& slice) {
-    this->Update(slice);
+void RefCountedSlice::_swap(RefCountedSlice& slice) {
+    std::swap(ref_cnt_, slice.ref_cnt_);
+    Slice::_swap(slice);
 }
 
-RefCountedSlice::RefCountedSlice(RefCountedSlice&& slice) {
-    this->Update(slice);
+RefCountedSlice::RefCountedSlice(const RefCountedSlice& slice) {
+    _copy(slice);
 }
 
+RefCountedSlice::RefCountedSlice(RefCountedSlice&& slice) { _swap(slice); }
+
 RefCountedSlice& RefCountedSlice::operator=(const RefCountedSlice& slice) {
     if (&slice == this) {
         return *this;
     }
-    this->Release();
-    this->Update(slice);
+    _release();
+    _copy(slice);
     return *this;
 }
 
 RefCountedSlice& RefCountedSlice::operator=(RefCountedSlice&& slice) {
-    if (&slice == this) {
-        return *this;
-    }
-    this->Release();
-    this->Update(slice);
+    _swap(slice);
     return *this;
 }