Make TAtomicIntrusivePtr lsan-friendly

commit_hash:c4a4db14dd6f9b82fb65377014112bf0a131d3e0

3 files changed, 134 insertions, 26 deletions

diff --git a/library/cpp/yt/memory/atomic_intrusive_ptr-inl.h b/library/cpp/yt/memory/atomic_intrusive_ptr-inl.h
index 14f2e10365..c605916e2d 100644
--- a/library/cpp/yt/memory/atomic_intrusive_ptr-inl.h
+++ b/library/cpp/yt/memory/atomic_intrusive_ptr-inl.h
@@ -16,35 +16,114 @@ TAtomicIntrusivePtr<T>::TAtomicIntrusivePtr(std::nullptr_t)
 { }
 
 template <class T>
+TAtomicIntrusivePtr<T>& TAtomicIntrusivePtr<T>::operator=(TIntrusivePtr<T> other)
+{
+    Store(std::move(other));
+    return *this;
+}
+
+template <class T>
+TAtomicIntrusivePtr<T>& TAtomicIntrusivePtr<T>::operator=(std::nullptr_t)
+{
+    Reset();
+    return *this;
+}
+
+template <class T>
+TAtomicIntrusivePtr<T>::operator bool() const
+{
+    return Get();
+}
+
+#ifdef _lsan_enabled_
+
+template <class T>
 TAtomicIntrusivePtr<T>::TAtomicIntrusivePtr(TIntrusivePtr<T> other)
-    : Ptr_(AcquireObject(other.Release(), /*consumeRef*/ true))
+    : Ptr_(std::move(other))
 { }
 
 template <class T>
 TAtomicIntrusivePtr<T>::TAtomicIntrusivePtr(TAtomicIntrusivePtr&& other)
-    : Ptr_(other.Ptr_.load(std::memory_order::relaxed))
+    : Ptr_(std::move(other))
+{ }
+
+template <class T>
+TAtomicIntrusivePtr<T>::~TAtomicIntrusivePtr() = default;
+
+template <class T>
+TIntrusivePtr<T> TAtomicIntrusivePtr<T>::Acquire() const
 {
-    other.Ptr_.store(nullptr, std::memory_order::relaxed);
+    auto guard = Guard(Lock_);
+    return Ptr_;
 }
 
 template <class T>
-TAtomicIntrusivePtr<T>::~TAtomicIntrusivePtr()
+TIntrusivePtr<T> TAtomicIntrusivePtr<T>::Exchange(TIntrusivePtr<T> other)
 {
-    ReleaseObject(Ptr_.load());
+    auto guard = Guard(Lock_);
+    Ptr_.Swap(other);
+    return other;
 }
 
 template <class T>
-TAtomicIntrusivePtr<T>& TAtomicIntrusivePtr<T>::operator=(TIntrusivePtr<T> other)
+void TAtomicIntrusivePtr<T>::Store(TIntrusivePtr<T> other)
 {
-    Store(std::move(other));
-    return *this;
+    Exchange(std::move(other));
 }
 
 template <class T>
-TAtomicIntrusivePtr<T>& TAtomicIntrusivePtr<T>::operator=(std::nullptr_t)
+void TAtomicIntrusivePtr<T>::Reset()
 {
-    Reset();
-    return *this;
+    Exchange(nullptr);
+    ReleaseObject(Ptr_.exchange(0));
+}
+
+template <class T>
+bool TAtomicIntrusivePtr<T>::CompareAndSwap(TRawPtr& comparePtr, T* target)
+{
+    auto guard = Guard(Lock_);
+    if (Ptr_.Get() != comparePtr) {
+        comparePtr = Ptr_.Get();
+        return false;
+    }
+    auto targetPtr = TIntrusivePtr<T>(target, /*addReference*/ false);
+    Ptr_.Swap(targetPtr);
+    guard.Release();
+    // targetPtr will die here.
+    return true;
+}
+
+template <class T>
+bool TAtomicIntrusivePtr<T>::CompareAndSwap(TRawPtr& comparePtr, TIntrusivePtr<T> target)
+{
+    return CompareAndSwap(comparePtr, target.Release());
+}
+
+template <class T>
+typename TAtomicIntrusivePtr<T>::TRawPtr TAtomicIntrusivePtr<T>::Get() const
+{
+    auto guard = Guard(Lock_);
+    return Ptr_.Get();
+}
+
+#else
+
+template <class T>
+TAtomicIntrusivePtr<T>::TAtomicIntrusivePtr(TIntrusivePtr<T> other)
+    : Ptr_(AcquireObject(other.Release(), /*consumeRef*/ true))
+{ }
+
+template <class T>
+TAtomicIntrusivePtr<T>::TAtomicIntrusivePtr(TAtomicIntrusivePtr&& other)
+    : Ptr_(other.Ptr_.load(std::memory_order::relaxed))
+{
+    other.Ptr_.store(nullptr, std::memory_order::relaxed);
+}
+
+template <class T>
+TAtomicIntrusivePtr<T>::~TAtomicIntrusivePtr()
+{
+    ReleaseObject(Ptr_.load());
 }
 
 template <class T>
@@ -159,12 +238,6 @@ typename TAtomicIntrusivePtr<T>::TRawPtr TAtomicIntrusivePtr<T>::Get() const
 }
 
 template <class T>
-TAtomicIntrusivePtr<T>::operator bool() const
-{
-    return Get();
-}
-
-template <class T>
 TPackedPtr TAtomicIntrusivePtr<T>::AcquireObject(T* obj, bool consumeRef)
 {
     if (obj) {
@@ -189,6 +262,8 @@ void TAtomicIntrusivePtr<T>::DoRelease(T* obj, int refs)
     }
 }
 
+#endif
+
 ////////////////////////////////////////////////////////////////////////////////
 
 template <class T>
diff --git a/library/cpp/yt/memory/atomic_intrusive_ptr.h b/library/cpp/yt/memory/atomic_intrusive_ptr.h
index 72405053c0..91d220f7a9 100644
--- a/library/cpp/yt/memory/atomic_intrusive_ptr.h
+++ b/library/cpp/yt/memory/atomic_intrusive_ptr.h
@@ -2,6 +2,12 @@
 
 #include "intrusive_ptr.h"
 
+#include <util/system/compiler.h>
+
+#ifdef _lsan_enabled_
+#include <util/system/spinlock.h>
+#endif
+
 namespace NYT {
 
 ////////////////////////////////////////////////////////////////////////////////
@@ -54,6 +60,10 @@ private:
     template <class U>
     friend bool operator!=(const TIntrusivePtr<U>& lhs, const TAtomicIntrusivePtr<U>& rhs);
 
+#ifdef _lsan_enabled_
+    ::TSpinLock Lock_;
+    TIntrusivePtr<T> Ptr_;
+#else
     // Keeps packed pointer (localRefCount, objectPtr).
     // Atomic ptr holds N references, where N = ReservedRefCount - localRefCount.
     // LocalRefCount is incremented in Acquire method.
@@ -65,11 +75,10 @@ private:
     constexpr static int ReservedRefCount = (1 << CounterBits) - 1;
 
     // Consume ref if ownership is transferred.
-    // AcquireObject(ptr.Release(), true)
-    // AcquireObject(ptr.Get(), false)
     static TPackedPtr AcquireObject(T* obj, bool consumeRef = false);
     static void ReleaseObject(TPackedPtr packedPtr);
     static void DoRelease(T* obj, int refs);
+#endif
 };
 
 ////////////////////////////////////////////////////////////////////////////////
diff --git a/library/cpp/yt/memory/unittests/atomic_intrusive_ptr_ut.cpp b/library/cpp/yt/memory/unittests/atomic_intrusive_ptr_ut.cpp
index 68489fcdf7..e43dca33c6 100644
--- a/library/cpp/yt/memory/unittests/atomic_intrusive_ptr_ut.cpp
+++ b/library/cpp/yt/memory/unittests/atomic_intrusive_ptr_ut.cpp
@@ -3,6 +3,7 @@
 #include <library/cpp/yt/memory/new.h>
 #include <library/cpp/yt/memory/ref_counted.h>
 #include <library/cpp/yt/memory/atomic_intrusive_ptr.h>
+#include <library/cpp/yt/memory/leaky_singleton.h>
 
 namespace NYT {
 namespace {
@@ -168,7 +169,7 @@ private:
 
 ////////////////////////////////////////////////////////////////////////////////
 
-TEST(TAtomicPtrTest, Empty)
+TEST(TIntrusiveAtomicPtrTest, Empty)
 {
     TIntricateObjectPtr emptyPointer;
     EXPECT_EQ(nullptr, emptyPointer.Get());
@@ -177,7 +178,7 @@ TEST(TAtomicPtrTest, Empty)
 // Reserved ref count.
 constexpr int RRC = 65535;
 
-TEST(TAtomicPtrTest, Basic)
+TEST(TIntrusiveAtomicPtrTest, Basic)
 {
     TIntricateObject object;
 
@@ -216,7 +217,7 @@ TEST(TAtomicPtrTest, Basic)
     EXPECT_THAT(object, HasRefCounts(2 + RRC, 2 + RRC, 2));
 }
 
-TEST(TAtomicPtrTest, BasicConst)
+TEST(TIntrusiveAtomicPtrTest, BasicConst)
 {
     const TIntricateObject object;
 
@@ -255,7 +256,7 @@ TEST(TAtomicPtrTest, BasicConst)
     EXPECT_THAT(object, HasRefCounts(2 + RRC, 2 + RRC, 2));
 }
 
-TEST(TAtomicPtrTest, Acquire)
+TEST(TIntrusiveAtomicPtrTest, Acquire)
 {
     TIntricateObject object;
     {
@@ -281,7 +282,7 @@ TEST(TAtomicPtrTest, Acquire)
     EXPECT_THAT(object, HasRefCounts(RRC + RRC / 2, RRC + RRC / 2, 1));
 }
 
-TEST(TAtomicPtrTest, AcquireConst)
+TEST(TIntrusiveAtomicPtrTest, AcquireConst)
 {
     const TIntricateObject object;
     {
@@ -307,7 +308,7 @@ TEST(TAtomicPtrTest, AcquireConst)
     EXPECT_THAT(object, HasRefCounts(RRC + RRC / 2, RRC + RRC / 2, 1));
 }
 
-TEST(TAtomicPtrTest, CAS)
+TEST(TIntrusiveAtomicPtrTest, CAS)
 {
     TIntricateObject o1;
     TIntricateObject o2;
@@ -337,7 +338,7 @@ TEST(TAtomicPtrTest, CAS)
     EXPECT_THAT(o2, HasRefCounts(RRC, RRC, 1));
 }
 
-TEST(TAtomicPtrTest, CASConst)
+TEST(TIntrusiveAtomicPtrTest, CASConst)
 {
     const TIntricateObject o1;
     const TIntricateObject o2;
@@ -367,6 +368,29 @@ TEST(TAtomicPtrTest, CASConst)
     EXPECT_THAT(o2, HasRefCounts(RRC, RRC, 1));
 }
 
+TEST(TIntrusiveAtomicPtrTest, LSan)
+{
+    struct S final
+    { };
+
+    struct TSingleton
+    {
+        TSingleton()
+        {
+            for (auto& ptr : Ptrs) {
+                ptr.Store(New<S>());
+                // Clobber pointer bits to prevent LSan from tracing the pointer.
+                ptr.Acquire();
+            }
+        }
+
+        // LSan has some issues detecting leaks when just one allocation is made.
+        std::array<TAtomicIntrusivePtr<S>, 100> Ptrs;
+    };
+
+    LeakySingleton<TSingleton>();
+}
+
 ////////////////////////////////////////////////////////////////////////////////
 
 } // namespace