Hazard pointers cosmetics

6 files changed, 150 insertions, 83 deletions

diff --git a/library/cpp/yt/memory/atomic_intrusive_ptr-inl.h b/library/cpp/yt/memory/atomic_intrusive_ptr-inl.h
index 1e902dbe45..3e386e4a80 100644
--- a/library/cpp/yt/memory/atomic_intrusive_ptr-inl.h
+++ b/library/cpp/yt/memory/atomic_intrusive_ptr-inl.h
@@ -50,12 +50,12 @@ TAtomicIntrusivePtr<T>& TAtomicIntrusivePtr<T>::operator=(std::nullptr_t)
 template <class T>
 TIntrusivePtr<T> TAtomicIntrusivePtr<T>::Acquire() const
 {
-    void* ptr = Ptr_.load();
+    auto ptr = Ptr_.load();
     while (true) {
-        auto [localRefs, obj] = UnpackPointer<T>(ptr);
+        auto [obj, localRefs] = TTaggedPtr<T>::Unpack(ptr);
 
         if (!obj) {
-            return TIntrusivePtr<T>();
+            return {};
         }
 
         YT_VERIFY(localRefs < ReservedRefCount);
@@ -70,20 +70,20 @@ TIntrusivePtr<T> TAtomicIntrusivePtr<T>::Acquire() const
         }
 
         // Can not Ref(obj) here because it can be destroyed.
-        if (Ptr_.compare_exchange_weak(ptr, PackPointer(obj, newLocalRefs))) {
+        if (Ptr_.compare_exchange_weak(ptr, TTaggedPtr(obj, newLocalRefs).Pack())) {
             if (Y_UNLIKELY(newLocalRefs > ReservedRefCount / 2)) {
                 Ref(obj, ReservedRefCount / 2);
 
                 // Decrease local ref count.
                 while (true) {
-                    auto [localRefs, currentObj] = UnpackPointer<T>(ptr);
+                    auto [currentObj, localRefs] = TTaggedPtr<T>::Unpack(ptr);
 
                     if (currentObj != obj || localRefs <= ReservedRefCount / 2) {
                         Unref(obj, ReservedRefCount / 2);
                         break;
                     }
 
-                    if (Ptr_.compare_exchange_weak(ptr, PackPointer(obj, localRefs - ReservedRefCount / 2))) {
+                    if (Ptr_.compare_exchange_weak(ptr, TTaggedPtr(obj, localRefs - ReservedRefCount / 2).Pack())) {
                         break;
                     }
                 }
@@ -97,7 +97,7 @@ TIntrusivePtr<T> TAtomicIntrusivePtr<T>::Acquire() const
 template <class T>
 TIntrusivePtr<T> TAtomicIntrusivePtr<T>::Exchange(TIntrusivePtr<T> other)
 {
-    auto [localRefs, obj] = UnpackPointer<T>(Ptr_.exchange(AcquireObject(other.Release(), true)));
+    auto [obj, localRefs] = TTaggedPtr<T>::Unpack(Ptr_.exchange(AcquireObject(other.Release(), true)));
     DoRelease(obj, localRefs + 1);
     return TIntrusivePtr<T>(obj, false);
 }
@@ -111,13 +111,13 @@ void TAtomicIntrusivePtr<T>::Store(TIntrusivePtr<T> other)
 template <class T>
 void TAtomicIntrusivePtr<T>::Reset()
 {
-    ReleaseObject(Ptr_.exchange(nullptr));
+    ReleaseObject(Ptr_.exchange(0));
 }
 
 template <class T>
 bool TAtomicIntrusivePtr<T>::CompareAndSwap(void*& comparePtr, T* target)
 {
-    auto targetPtr = AcquireObject(target, false);
+    auto* targetPtr = AcquireObject(target, false);
 
     auto currentPtr = Ptr_.load();
     if (UnpackPointer<T>(currentPtr).Ptr == comparePtr && Ptr_.compare_exchange_strong(currentPtr, targetPtr)) {
@@ -138,12 +138,12 @@ bool TAtomicIntrusivePtr<T>::CompareAndSwap(void*& comparePtr, TIntrusivePtr<T>
     auto targetPtr = AcquireObject(target.Release(), true);
 
     auto currentPtr = Ptr_.load();
-    if (UnpackPointer<T>(currentPtr).Ptr == comparePtr && Ptr_.compare_exchange_strong(currentPtr, targetPtr)) {
+    if (TTaggedPtr<T>::Unpack(currentPtr).Ptr == comparePtr && Ptr_.compare_exchange_strong(currentPtr, targetPtr)) {
         ReleaseObject(currentPtr);
         return true;
     }
 
-    comparePtr = UnpackPointer<T>(currentPtr).Ptr;
+    comparePtr = TTaggedPtr<T>::Unpack(currentPtr).Ptr;
 
     ReleaseObject(targetPtr);
     return false;
@@ -152,7 +152,7 @@ bool TAtomicIntrusivePtr<T>::CompareAndSwap(void*& comparePtr, TIntrusivePtr<T>
 template <class T>
 void* TAtomicIntrusivePtr<T>::Get() const
 {
-    return UnpackPointer<void>(Ptr_.load()).Ptr;
+    return TTaggedPtr<void>::Unpack(Ptr_.load()).Ptr;
 }
 
 template <class T>
@@ -162,19 +162,19 @@ TAtomicIntrusivePtr<T>::operator bool() const
 }
 
 template <class T>
-void* TAtomicIntrusivePtr<T>::AcquireObject(T* obj, bool consumeRef)
+TPackedPtr TAtomicIntrusivePtr<T>::AcquireObject(T* obj, bool consumeRef)
 {
     if (obj) {
         Ref(obj, static_cast<int>(ReservedRefCount - consumeRef));
     }
 
-    return PackPointer(obj, 0);
+    return TTaggedPtr(obj).Pack();
 }
 
 template <class T>
-void TAtomicIntrusivePtr<T>::ReleaseObject(void* packedPtr)
+void TAtomicIntrusivePtr<T>::ReleaseObject(TPackedPtr packedPtr)
 {
-    auto [localRefs, obj] = UnpackPointer<T>(packedPtr);
+    auto [obj, localRefs] = TTaggedPtr<T>::Unpack(packedPtr);
     DoRelease(obj, localRefs);
 }
 
diff --git a/library/cpp/yt/memory/atomic_intrusive_ptr.h b/library/cpp/yt/memory/atomic_intrusive_ptr.h
index f8c9a1a7fe..0137a83515 100644
--- a/library/cpp/yt/memory/atomic_intrusive_ptr.h
+++ b/library/cpp/yt/memory/atomic_intrusive_ptr.h
@@ -60,16 +60,16 @@ private:
     // LocalRefCount is incremented in Acquire method.
     // When localRefCount exceeds ReservedRefCount / 2 a new portion of refs are required globally.
     // This field is marked mutable in order to make Acquire const-qualified in accordance to its semantics.
-    mutable std::atomic<void*> Ptr_ = nullptr;
+    mutable std::atomic<TPackedPtr> Ptr_ = 0;
 
-    constexpr static int CounterBits = 64 - PtrBits;
+    constexpr static int CounterBits = PackedPtrTagBits;
     constexpr static int ReservedRefCount = (1 << CounterBits) - 1;
 
     // Consume ref if ownership is transferred.
     // AcquireObject(ptr.Release(), true)
     // AcquireObject(ptr.Get(), false)
-    static void* AcquireObject(T* obj, bool consumeRef = false);
-    static void ReleaseObject(void* packedPtr);
+    static TPackedPtr AcquireObject(T* obj, bool consumeRef = false);
+    static void ReleaseObject(TPackedPtr packedPtr);
     static void DoRelease(T* obj, int refs);
 };
 
diff --git a/library/cpp/yt/memory/ref_counted-inl.h b/library/cpp/yt/memory/ref_counted-inl.h
index 3f034ba7dc..5f1f60af7b 100644
--- a/library/cpp/yt/memory/ref_counted-inl.h
+++ b/library/cpp/yt/memory/ref_counted-inl.h
@@ -4,68 +4,78 @@
 #include "ref_counted.h"
 #endif
 
+#include "tagged_ptr.h"
+
 #include <util/system/sanitizers.h>
 
 namespace NYT {
 
 ////////////////////////////////////////////////////////////////////////////////
 
-constexpr uint16_t PtrBits = 48;
-constexpr uintptr_t PtrMask = (1ULL << PtrBits) - 1;
+// TODO(babenko): move to hazard pointers
+void RetireHazardPointer(TPackedPtr packedPtr, void (*reclaimer)(TPackedPtr));
 
-template <class T>
-Y_FORCE_INLINE void* PackPointer(T* ptr, uint16_t data)
-{
-    return reinterpret_cast<void*>((static_cast<uintptr_t>(data) << PtrBits) | reinterpret_cast<uintptr_t>(ptr));
-}
+////////////////////////////////////////////////////////////////////////////////
 
-template <class T>
-struct TPackedPointer
+namespace NDetail {
+
+////////////////////////////////////////////////////////////////////////////////
+
+template <class T, class = void>
+struct TFreeMemory
 {
-    uint16_t Data;
-    T* Ptr;
+    static void Do(void* ptr)
+    {
+#ifdef _win_
+        ::_aligned_free(ptr);
+#else
+        ::free(ptr);
+#endif
+    }
 };
 
 template <class T>
-Y_FORCE_INLINE TPackedPointer<T> UnpackPointer(void* packedPtr)
+struct TFreeMemory<T, std::void_t<typename T::TAllocator>>
 {
-    auto castedPtr = reinterpret_cast<uintptr_t>(packedPtr);
-    return {static_cast<uint16_t>(castedPtr >> PtrBits), reinterpret_cast<T*>(castedPtr & PtrMask)};
-}
+    static void Do(void* ptr)
+    {
+        using TAllocator = typename T::TAllocator;
+        TAllocator::Free(ptr);
+    }
+};
 
 ////////////////////////////////////////////////////////////////////////////////
 
 template <class T, class = void>
 struct TMemoryReleaser
 {
-    static void Do(void* ptr, uint16_t /*offset*/)
+    static void Do(void* ptr, ui16 /*offset*/)
     {
         TFreeMemory<T>::Do(ptr);
     }
 };
 
-using TDeleter = void (*)(void*);
-
-void ScheduleObjectDeletion(void* ptr, TDeleter deleter);
-
 template <class T>
 struct TMemoryReleaser<T, std::enable_if_t<T::EnableHazard>>
 {
-    static void Do(void* ptr, uint16_t offset)
+    static void Do(void* ptr, ui16 offset)
     {
         // Base pointer is used in HazardPtr as the identity of object.
-        auto* basePtr = PackPointer(static_cast<char*>(ptr) + offset, offset);
-
-        ScheduleObjectDeletion(basePtr, [] (void* ptr) {
+        auto packedPtr = TTaggedPtr<char>{static_cast<char*>(ptr) + offset, offset}.Pack();
+        RetireHazardPointer(packedPtr, [] (TPackedPtr packedPtr) {
             // Base ptr and the beginning of allocated memory region may differ.
-            auto [offset, basePtr] = UnpackPointer<char>(ptr);
-            TFreeMemory<T>::Do(basePtr - offset);
+            auto [ptr, offset] = TTaggedPtr<char>::Unpack(packedPtr);
+            TFreeMemory<T>::Do(ptr - offset);
         });
     }
 };
 
 ////////////////////////////////////////////////////////////////////////////////
 
+} // namespace NDetail
+
+////////////////////////////////////////////////////////////////////////////////
+
 Y_FORCE_INLINE int TRefCounter::GetRefCount() const noexcept
 {
     return StrongCount_.load(std::memory_order::acquire);
@@ -158,19 +168,19 @@ struct TRefCountedHelper
 
         // Fast path. Weak refs cannot appear if there are neither strong nor weak refs.
         if (refCounter->GetWeakRefCount() == 1) {
-            TMemoryReleaser<T>::Do(ptr - RefCounterOffset, RefCounterSpace);
+            NYT::NDetail::TMemoryReleaser<T>::Do(ptr - RefCounterOffset, RefCounterSpace);
             return;
         }
 
         if (refCounter->WeakUnref()) {
-            TMemoryReleaser<T>::Do(ptr - RefCounterOffset, RefCounterSpace);
+            NYT::NDetail::TMemoryReleaser<T>::Do(ptr - RefCounterOffset, RefCounterSpace);
         }
     }
 
     Y_FORCE_INLINE static void Deallocate(const T* obj)
     {
         char* ptr = reinterpret_cast<char*>(const_cast<TRefCounter*>(GetRefCounter(obj)));
-        TMemoryReleaser<T>::Do(ptr - RefCounterOffset, RefCounterSpace);
+        NYT::NDetail::TMemoryReleaser<T>::Do(ptr - RefCounterOffset, RefCounterSpace);
     }
 };
 
@@ -189,8 +199,8 @@ struct TRefCountedHelper<T, true>
 
     Y_FORCE_INLINE static void Deallocate(const TRefCountedBase* obj)
     {
-        auto* ptr = reinterpret_cast<void**>(const_cast<TRefCountedBase*>(obj));
-        auto [offset, ptrToDeleter] = UnpackPointer<void(void*, uint16_t)>(*ptr);
+        auto* ptr = reinterpret_cast<TPackedPtr*>(const_cast<TRefCountedBase*>(obj));
+        auto [ptrToDeleter, offset] = TTaggedPtr<void(void*, ui16)>::Unpack(*ptr);
 
         // The most derived type is erased here. So we cannot call TMemoryReleaser with derived type.
         ptrToDeleter(reinterpret_cast<char*>(ptr) - offset, offset);
@@ -262,17 +272,17 @@ void TRefCounted::DestroyRefCountedImpl(T* ptr)
 
     // Fast path. Weak refs cannot appear if there are neither strong nor weak refs.
     if (refCounter->GetWeakRefCount() == 1) {
-        TMemoryReleaser<T>::Do(ptr, offset);
+        NYT::NDetail::TMemoryReleaser<T>::Do(ptr, offset);
         return;
     }
 
-    YT_ASSERT(offset < std::numeric_limits<uint16_t>::max());
+    YT_ASSERT(offset < std::numeric_limits<ui16>::max());
 
-    auto* vTablePtr = reinterpret_cast<void**>(basePtr);
-    *vTablePtr = PackPointer(&TMemoryReleaser<T>::Do, offset);
+    auto* vTablePtr = reinterpret_cast<TPackedPtr*>(basePtr);
+    *vTablePtr = TTaggedPtr<void(void*, ui16)>(&NYT::NDetail::TMemoryReleaser<T>::Do, offset).Pack();
 
     if (refCounter->WeakUnref()) {
-        TMemoryReleaser<T>::Do(ptr, offset);
+        NYT::NDetail::TMemoryReleaser<T>::Do(ptr, offset);
     }
 }
 
diff --git a/library/cpp/yt/memory/ref_counted.h b/library/cpp/yt/memory/ref_counted.h
index e2ca9245be..d6f79edf8e 100644
--- a/library/cpp/yt/memory/ref_counted.h
+++ b/library/cpp/yt/memory/ref_counted.h
@@ -31,31 +31,6 @@ private:
 
 ////////////////////////////////////////////////////////////////////////////////
 
-template <class T, class = void>
-struct TFreeMemory
-{
-    static void Do(void* ptr)
-    {
-#ifdef _win_
-        ::_aligned_free(ptr);
-#else
-        ::free(ptr);
-#endif
-    }
-};
-
-template <class T>
-struct TFreeMemory<T, std::void_t<typename T::TAllocator>>
-{
-    static void Do(void* ptr)
-    {
-        using TAllocator = typename T::TAllocator;
-        TAllocator::Free(ptr);
-    }
-};
-
-////////////////////////////////////////////////////////////////////////////////
-
 class TRefCounter
 {
 public:
@@ -102,8 +77,6 @@ void DeallocateRefCounted(const T* obj);
 
 ////////////////////////////////////////////////////////////////////////////////
 
-// API
-
 template <class T>
 void Ref(T* obj, int n = 1);
 
diff --git a/library/cpp/yt/memory/tagged_ptr-inl.h b/library/cpp/yt/memory/tagged_ptr-inl.h
new file mode 100644
index 0000000000..99da61ada6
--- /dev/null
+++ b/library/cpp/yt/memory/tagged_ptr-inl.h
@@ -0,0 +1,45 @@
+#ifndef TAGGED_PTR_INL_H_
+#error "Direct inclusion of this file is not allowed, include tagged_ptr.h"
+// For the sake of sane code completion.
+#include "tagged_ptr.h"
+#endif
+
+#include <library/cpp/yt/assert/assert.h>
+
+#include <util/system/compiler.h>
+
+namespace NYT {
+
+////////////////////////////////////////////////////////////////////////////////
+
+template <class T>
+Y_FORCE_INLINE TTaggedPtr<T>::TTaggedPtr(T* ptr, ui16 tag)
+    : Ptr(ptr)
+    , Tag(tag)
+{ }
+
+template <class T>
+Y_FORCE_INLINE TTaggedPtr<T>::TTaggedPtr(T* ptr)
+    : Ptr(ptr)
+    , Tag(0)
+{ }
+
+template <class T>
+Y_FORCE_INLINE TPackedPtr TTaggedPtr<T>::Pack() const
+{
+    YT_ASSERT((reinterpret_cast<TPackedPtr>(Ptr) & PackedPtrTagMask) == 0);
+    return (static_cast<TPackedPtr>(Tag) << PackedPtrAddrsssBits) | reinterpret_cast<TPackedPtr>(Ptr);
+}
+
+template <class T>
+Y_FORCE_INLINE TTaggedPtr<T> TTaggedPtr<T>::Unpack(TPackedPtr packedPtr)
+{
+    return {
+        reinterpret_cast<T*>(packedPtr & PackedPtrAddressMask),
+        static_cast<ui16>(packedPtr >> PackedPtrAddrsssBits),
+    };
+}
+
+////////////////////////////////////////////////////////////////////////////////
+
+} // namespace NYT
diff --git a/library/cpp/yt/memory/tagged_ptr.h b/library/cpp/yt/memory/tagged_ptr.h
new file mode 100644
index 0000000000..65747fa4de
--- /dev/null
+++ b/library/cpp/yt/memory/tagged_ptr.h
@@ -0,0 +1,39 @@
+#pragma once
+
+#include <util/system/types.h>
+
+namespace NYT {
+
+////////////////////////////////////////////////////////////////////////////////
+
+using TPackedPtr = uintptr_t;
+static_assert(sizeof(TPackedPtr) == 8);
+
+constexpr size_t PackedPtrAddrsssBits = 48;
+constexpr size_t PackedPtrTagBits = 16;
+constexpr TPackedPtr PackedPtrAddressMask = (1ULL << PackedPtrAddrsssBits) - 1;
+constexpr TPackedPtr PackedPtrTagMask = ~PackedPtrAddressMask;
+
+////////////////////////////////////////////////////////////////////////////////
+
+template <class T>
+struct TTaggedPtr
+{
+    TTaggedPtr() = default;
+    TTaggedPtr(T* ptr, ui16 tag);
+    explicit TTaggedPtr(T* ptr);
+
+    T* Ptr;
+    ui16 Tag;
+
+    TPackedPtr Pack() const;
+    static TTaggedPtr<T> Unpack(TPackedPtr packedPtr);
+};
+
+////////////////////////////////////////////////////////////////////////////////
+
+} // namespace NYT
+
+#define TAGGED_PTR_INL_H_
+#include "tagged_ptr-inl.h"
+#undef TAGGED_PTR_INL_H_