From f6af0cd2935cc29c37d243da1cc9c9a742d8590c Mon Sep 17 00:00:00 2001
From: Aleksei Borzenkov <snaury@ydb.tech>
Date: Mon, 5 Feb 2024 12:07:19 +0000
Subject: [PATCH] Use std::atomic in shared data header

---
 ydb/library/actors/util/shared_data.cpp           |  4 ++--
 ydb/library/actors/util/shared_data.h             | 15 +++++++++++----
 .../actors/util/shared_data_backtracing_owner.h   |  5 +++--
 ydb/library/actors/util/shared_data_ut.cpp        |  6 ++++--
 4 files changed, 20 insertions(+), 10 deletions(-)

diff --git a/ydb/library/actors/util/shared_data.cpp b/ydb/library/actors/util/shared_data.cpp
index 51311ce7a3b3..910ae26929d8 100644
--- a/ydb/library/actors/util/shared_data.cpp
+++ b/ydb/library/actors/util/shared_data.cpp
@@ -23,8 +23,7 @@ namespace NActors {
             NActors::NMemory::TLabel<MemoryLabelSharedData>::Add(allocSize);
 
             auto* header = reinterpret_cast<THeader*>(raw + PrivateHeaderSize);
-            header->RefCount = 1;
-            header->Owner = nullptr;
+            new (header) THeader(nullptr);
 
             data = raw + OverheadSize;
             NSan::Poison(data, size);
@@ -41,6 +40,7 @@ namespace NActors {
 
             auto* header = reinterpret_cast<THeader*>(raw + PrivateHeaderSize);
             Y_DEBUG_ABORT_UNLESS(header->Owner == nullptr);
+            header->~THeader();
 
             y_deallocate(raw);
         }
diff --git a/ydb/library/actors/util/shared_data.h b/ydb/library/actors/util/shared_data.h
index bd9afb00a599..f9db1c3597f2 100644
--- a/ydb/library/actors/util/shared_data.h
+++ b/ydb/library/actors/util/shared_data.h
@@ -6,6 +6,8 @@
 #include <util/system/compiler.h>
 #include <util/generic/array_ref.h>
 
+#include <atomic>
+
 namespace NActors {
 
     class TSharedData {
@@ -25,8 +27,13 @@ namespace NActors {
         static_assert(sizeof(TPrivateHeader) == 16, "TPrivateHeader has an unexpected size");
 
         struct THeader {
-            TAtomic RefCount;
+            std::atomic<size_t> RefCount;
             IOwner* Owner;
+
+            explicit THeader(IOwner* owner)
+                : RefCount{ 1 }
+                , Owner(owner)
+            {}
         };
 
         static_assert(sizeof(THeader) == 16, "THeader has an unexpected size");
@@ -193,19 +200,19 @@ namespace NActors {
         }
 
         static bool IsPrivate(THeader* header) noexcept {
-            return 1 == AtomicGet(header->RefCount);
+            return 1 == header->RefCount.load(std::memory_order_relaxed);
         }
 
         void AddRef() noexcept {
             if (Data_) {
-                AtomicIncrement(Header()->RefCount);
+                Header()->RefCount.fetch_add(1, std::memory_order_relaxed);
             }
         }
 
         void Release() noexcept {
             if (Data_) {
                 auto* header = Header();
-                if (IsPrivate(header) || 0 == AtomicDecrement(header->RefCount)) {
+                if (1 == header->RefCount.fetch_sub(1, std::memory_order_acq_rel)) {
                     if (auto* owner = header->Owner) {
                         owner->Deallocate(Data_);
                     } else {
diff --git a/ydb/library/actors/util/shared_data_backtracing_owner.h b/ydb/library/actors/util/shared_data_backtracing_owner.h
index ea479d5fd129..21fe95ebdb64 100644
--- a/ydb/library/actors/util/shared_data_backtracing_owner.h
+++ b/ydb/library/actors/util/shared_data_backtracing_owner.h
@@ -30,8 +30,7 @@ class TBackTracingOwner : public NActors::TSharedData::IOwner {
         TSelf* btOwner = new TSelf;
         btOwner->BackTrace.Capture();
         btOwner->Info = info;
-        header->RefCount = 1;
-        header->Owner = btOwner;
+        new (header) THeader(btOwner);
         char* data = raw + OverheadSize;
         return data;
     }
@@ -59,6 +58,8 @@ class TBackTracingOwner : public NActors::TSharedData::IOwner {
 
     void Deallocate(char* data) noexcept override {
         if (!RealOwner) {
+            THeader* header = reinterpret_cast<THeader*>(data - sizeof(THeader));
+            header->~THeader();
             char* raw = data - OverheadSize;
             y_deallocate(raw);
         } else {
diff --git a/ydb/library/actors/util/shared_data_ut.cpp b/ydb/library/actors/util/shared_data_ut.cpp
index 2f7dc2ccc80e..82f0405472ef 100644
--- a/ydb/library/actors/util/shared_data_ut.cpp
+++ b/ydb/library/actors/util/shared_data_ut.cpp
@@ -80,8 +80,7 @@ namespace NActors {
             TSharedData Allocate(size_t size) {
                 char* raw = reinterpret_cast<char*>(y_allocate(sizeof(THeader) + size));
                 THeader* header = reinterpret_cast<THeader*>(raw);
-                header->RefCount = 1;
-                header->Owner = this;
+                new (header) THeader(this);
                 char* data = raw + sizeof(THeader);
                 Y_ABORT_UNLESS(Allocated_.insert(data).second);
                 return TSharedData::AttachUnsafe(data, size);
@@ -90,6 +89,8 @@ namespace NActors {
             void Deallocate(char* data) noexcept {
                 Y_ABORT_UNLESS(Allocated_.erase(data) > 0);
                 char* raw = data - sizeof(THeader);
+                THeader* header = reinterpret_cast<THeader*>(raw);
+                header->~THeader();
                 y_deallocate(raw);
                 Deallocated_.push_back(data);
             }
@@ -190,6 +191,7 @@ namespace NActors {
             // Test Detach copies correctly and doesn't affect owned data
             {
                 auto data = owner.Allocate(42);
+                ptr = data.data();
                 auto disowned = data;
                 disowned.Detach();
                 UNIT_ASSERT(owner.NextDeallocated() == nullptr);