Skip to content

Commit

Permalink
Merge f6af0cd into 4fef0c2
Browse files Browse the repository at this point in the history
  • Loading branch information
snaury authored Feb 5, 2024
2 parents 4fef0c2 + f6af0cd commit 1852382
Show file tree
Hide file tree
Showing 4 changed files with 20 additions and 10 deletions.
4 changes: 2 additions & 2 deletions ydb/library/actors/util/shared_data.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand All @@ -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);
}
Expand Down
15 changes: 11 additions & 4 deletions ydb/library/actors/util/shared_data.h
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,8 @@
#include <util/system/compiler.h>
#include <util/generic/array_ref.h>

#include <atomic>

namespace NActors {

class TSharedData {
Expand All @@ -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");
Expand Down Expand Up @@ -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 {
Expand Down
5 changes: 3 additions & 2 deletions ydb/library/actors/util/shared_data_backtracing_owner.h
Original file line number Diff line number Diff line change
Expand Up @@ -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;
}
Expand Down Expand Up @@ -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 {
Expand Down
6 changes: 4 additions & 2 deletions ydb/library/actors/util/shared_data_ut.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand All @@ -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);
}
Expand Down Expand Up @@ -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);
Expand Down

0 comments on commit 1852382

Please sign in to comment.