Intermediate changes

3 files changed, 119 insertions, 11 deletions

diff --git a/library/cpp/yt/memory/atomic_intrusive_ptr-inl.h b/library/cpp/yt/memory/atomic_intrusive_ptr-inl.h
index 3e386e4a80..ae8704ae19 100644
--- a/library/cpp/yt/memory/atomic_intrusive_ptr-inl.h
+++ b/library/cpp/yt/memory/atomic_intrusive_ptr-inl.h
@@ -115,7 +115,7 @@ void TAtomicIntrusivePtr<T>::Reset()
 }
 
 template <class T>
-bool TAtomicIntrusivePtr<T>::CompareAndSwap(void*& comparePtr, T* target)
+bool TAtomicIntrusivePtr<T>::CompareAndSwap(TRawPtr& comparePtr, T* target)
 {
     auto* targetPtr = AcquireObject(target, false);
 
@@ -132,7 +132,7 @@ bool TAtomicIntrusivePtr<T>::CompareAndSwap(void*& comparePtr, T* target)
 }
 
 template <class T>
-bool TAtomicIntrusivePtr<T>::CompareAndSwap(void*& comparePtr, TIntrusivePtr<T> target)
+bool TAtomicIntrusivePtr<T>::CompareAndSwap(TRawPtr& comparePtr, TIntrusivePtr<T> target)
 {
     // TODO(lukyan): Make helper for packed owning ptr?
     auto targetPtr = AcquireObject(target.Release(), true);
@@ -150,7 +150,7 @@ bool TAtomicIntrusivePtr<T>::CompareAndSwap(void*& comparePtr, TIntrusivePtr<T>
 }
 
 template <class T>
-void* TAtomicIntrusivePtr<T>::Get() const
+typename TAtomicIntrusivePtr<T>::TRawPtr TAtomicIntrusivePtr<T>::Get() const
 {
     return TTaggedPtr<void>::Unpack(Ptr_.load()).Ptr;
 }
diff --git a/library/cpp/yt/memory/atomic_intrusive_ptr.h b/library/cpp/yt/memory/atomic_intrusive_ptr.h
index 0137a83515..ae0a307182 100644
--- a/library/cpp/yt/memory/atomic_intrusive_ptr.h
+++ b/library/cpp/yt/memory/atomic_intrusive_ptr.h
@@ -34,11 +34,13 @@ public:
     void Store(TIntrusivePtr<T> other);
 
     void Reset();
-    bool CompareAndSwap(void*& comparePtr, T* target);
-    bool CompareAndSwap(void*& comparePtr, TIntrusivePtr<T> target);
+
+    using TRawPtr = std::conditional_t<std::is_const_v<T>, const void*, void*>;
+    bool CompareAndSwap(TRawPtr& comparePtr, T* target);
+    bool CompareAndSwap(TRawPtr& comparePtr, TIntrusivePtr<T> target);
 
     // Result is suitable only for comparison. Not dereference.
-    void* Get() const;
+    TRawPtr Get() const;
 
     explicit operator bool() const;
 
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 27c87d379b..f6a26e56b1 100644
--- a/library/cpp/yt/memory/unittests/atomic_intrusive_ptr_ut.cpp
+++ b/library/cpp/yt/memory/unittests/atomic_intrusive_ptr_ut.cpp
@@ -24,16 +24,16 @@ using ::testing::StrictMock;
 struct TIntricateObject
     : private TNonCopyable
 {
-    int Increments = 0;
-    int Decrements = 0;
-    int Zeros = 0;
+    mutable int Increments = 0;
+    mutable int Decrements = 0;
+    mutable int Zeros = 0;
 
-    void Ref(int n)
+    void Ref(int n) const
     {
         Increments += n;
     }
 
-    void Unref(int n)
+    void Unref(int n) const
     {
         Decrements += n;
         if (Increments == Decrements) {
@@ -43,6 +43,7 @@ struct TIntricateObject
 };
 
 typedef TIntrusivePtr<TIntricateObject> TIntricateObjectPtr;
+typedef TIntrusivePtr<const TIntricateObject> TConstIntricateObjectPtr;
 
 void Ref(TIntricateObject* obj, int n = 1)
 {
@@ -54,6 +55,16 @@ void Unref(TIntricateObject* obj, int n = 1)
     obj->Unref(n);
 }
 
+void Ref(const TIntricateObject* obj, int n = 1)
+{
+    obj->Ref(n);
+}
+
+void Unref(const TIntricateObject* obj, int n = 1)
+{
+    obj->Unref(n);
+}
+
 MATCHER_P3(HasRefCounts, increments, decrements, zeros,
     "Reference counter " \
     "was incremented " + ::testing::PrintToString(increments) + " times, " +
@@ -205,6 +216,45 @@ TEST(TAtomicPtrTest, Basic)
     EXPECT_THAT(object, HasRefCounts(2 + RRC, 2 + RRC, 2));
 }
 
+TEST(TAtomicPtrTest, BasicConst)
+{
+    const TIntricateObject object;
+
+    EXPECT_THAT(object, HasRefCounts(0, 0, 0));
+
+    {
+        TConstIntricateObjectPtr owningPointer(&object);
+        EXPECT_THAT(object, HasRefCounts(1, 0, 0));
+        EXPECT_EQ(&object, owningPointer.Get());
+    }
+
+    EXPECT_THAT(object, HasRefCounts(1, 1, 1));
+
+
+    {
+        TConstIntricateObjectPtr owningPointer(&object);
+        TAtomicIntrusivePtr<const TIntricateObject> atomicPointer(owningPointer);
+
+        EXPECT_THAT(object, HasRefCounts(2 + RRC, 1, 1));
+        EXPECT_EQ(&object, owningPointer.Get());
+
+
+        auto p1 = atomicPointer.Acquire();
+
+        EXPECT_THAT(object, HasRefCounts(2 + RRC, 1, 1));
+
+        p1.Reset();
+
+        EXPECT_THAT(object, HasRefCounts(2 + RRC, 2, 1));
+
+        owningPointer.Reset();
+
+        EXPECT_THAT(object, HasRefCounts(2 + RRC, 3, 1));
+    }
+
+    EXPECT_THAT(object, HasRefCounts(2 + RRC, 2 + RRC, 2));
+}
+
 TEST(TAtomicPtrTest, Acquire)
 {
     TIntricateObject object;
@@ -231,6 +281,32 @@ TEST(TAtomicPtrTest, Acquire)
     EXPECT_THAT(object, HasRefCounts(RRC + RRC / 2, RRC + RRC / 2, 1));
 }
 
+TEST(TAtomicPtrTest, AcquireConst)
+{
+    const TIntricateObject object;
+    {
+        TAtomicIntrusivePtr<const TIntricateObject> atomicPtr{TConstIntricateObjectPtr(&object)};
+        EXPECT_THAT(object, HasRefCounts(RRC, 0, 0));
+
+        for (int i = 0; i < RRC / 2; ++i) {
+            {
+                auto tmp = atomicPtr.Acquire();
+                EXPECT_THAT(object, HasRefCounts(RRC, i, 0));
+            }
+            EXPECT_THAT(object, HasRefCounts(RRC, i + 1, 0));
+        }
+
+        {
+            auto tmp = atomicPtr.Acquire();
+            EXPECT_THAT(object, HasRefCounts( RRC + RRC / 2, RRC - 1, 0));
+        }
+
+        EXPECT_THAT(object, HasRefCounts(RRC + RRC / 2, RRC, 0));
+    }
+
+    EXPECT_THAT(object, HasRefCounts(RRC + RRC / 2, RRC + RRC / 2, 1));
+}
+
 TEST(TAtomicPtrTest, CAS)
 {
     TIntricateObject o1;
@@ -261,6 +337,36 @@ TEST(TAtomicPtrTest, CAS)
     EXPECT_THAT(o2, HasRefCounts(RRC, RRC, 1));
 }
 
+TEST(TAtomicPtrTest, CASConst)
+{
+    const TIntricateObject o1;
+    const TIntricateObject o2;
+    {
+
+        TAtomicIntrusivePtr<const TIntricateObject> atomicPtr{TConstIntricateObjectPtr(&o1)};
+        EXPECT_THAT(o1, HasRefCounts(RRC, 0, 0));
+
+        TConstIntricateObjectPtr p2(&o2);
+        EXPECT_THAT(o2, HasRefCounts(1, 0, 0));
+
+        const void* rawPtr = &o1;
+        EXPECT_TRUE(atomicPtr.CompareAndSwap(rawPtr, std::move(p2)));
+        EXPECT_EQ(rawPtr, &o1);
+
+        EXPECT_THAT(o1, HasRefCounts(RRC, RRC, 1));
+        EXPECT_THAT(o2, HasRefCounts(RRC, 0, 0));
+
+        rawPtr = nullptr;
+        EXPECT_FALSE(atomicPtr.CompareAndSwap(rawPtr, TConstIntricateObjectPtr(&o1)));
+        EXPECT_EQ(rawPtr, &o2);
+
+        EXPECT_THAT(o1, HasRefCounts(2 * RRC, 2 * RRC, 2));
+        EXPECT_THAT(o2, HasRefCounts(RRC, 0, 0));
+    }
+
+    EXPECT_THAT(o2, HasRefCounts(RRC, RRC, 1));
+}
+
 ////////////////////////////////////////////////////////////////////////////////
 
 } // namespace