YT-19007: Implement NewWithOffloadedDtor

4 files changed, 120 insertions, 112 deletions

diff --git a/library/cpp/yt/memory/new-inl.h b/library/cpp/yt/memory/new-inl.h
index 282a2a2515..7d80b81d65 100644
--- a/library/cpp/yt/memory/new-inl.h
+++ b/library/cpp/yt/memory/new-inl.h
@@ -4,6 +4,8 @@
 #include "new.h"
 #endif
 
+#include "ref_tracked.h"
+
 #include <library/cpp/yt/malloc//malloc.h>
 
 namespace NYT {
@@ -56,9 +58,9 @@ class TRefCountedWrapperWithDeleter final
 {
 public:
     template <class... TArgs>
-    explicit TRefCountedWrapperWithDeleter(const TDeleter& deleter, TArgs&&... args)
+    explicit TRefCountedWrapperWithDeleter(TDeleter deleter, TArgs&&... args)
         : T(std::forward<TArgs>(args)...)
-        , Deleter_(deleter)
+        , Deleter_(std::move(deleter))
     { }
 
     ~TRefCountedWrapperWithDeleter() = default;
@@ -69,7 +71,7 @@ public:
     }
 
 private:
-    TDeleter Deleter_;
+    const TDeleter Deleter_;
 };
 
 template <class T>
@@ -164,7 +166,7 @@ Y_FORCE_INLINE TIntrusivePtr<T> SafeConstruct(void* ptr, As&&... args)
 {
     try {
         auto* instance = TConstructHelper<T>::Construct(ptr, std::forward<As>(args)...);
-        return TIntrusivePtr<T>(instance, false);
+        return TIntrusivePtr<T>(instance, /*addReference*/ false);
     } catch (const std::exception& ex) {
         // Do not forget to free the memory.
         TFreeMemory<T>::Do(ptr);
@@ -252,19 +254,18 @@ Y_FORCE_INLINE TIntrusivePtr<T> NewWithExtraSpace(
 
 ////////////////////////////////////////////////////////////////////////////////
 
-// Support for polymorphic only
 template <class T, class TDeleter, class... As>
-Y_FORCE_INLINE TIntrusivePtr<T> NewWithDelete(const TDeleter& deleter, As&&... args)
+Y_FORCE_INLINE TIntrusivePtr<T> NewWithDeleter(TDeleter deleter, As&&... args)
 {
     using TWrapper = TRefCountedWrapperWithDeleter<T, TDeleter>;
     void* ptr = NYT::NDetail::AllocateConstSizeAligned<sizeof(TWrapper), alignof(TWrapper)>();
 
     auto* instance = NYT::NDetail::NewEpilogue<TWrapper>(
         ptr,
-        deleter,
+        std::move(deleter),
         std::forward<As>(args)...);
 
-    return TIntrusivePtr<T>(instance, false);
+    return TIntrusivePtr<T>(instance, /*addReference*/ false);
 }
 
 ////////////////////////////////////////////////////////////////////////////////
@@ -285,7 +286,7 @@ Y_FORCE_INLINE TIntrusivePtr<T> NewWithLocation(
     Y_UNUSED(location);
 #endif
 
-    return TIntrusivePtr<T>(instance, false);
+    return TIntrusivePtr<T>(instance, /*addReference*/ false);
 }
 
 ////////////////////////////////////////////////////////////////////////////////
diff --git a/library/cpp/yt/memory/new.h b/library/cpp/yt/memory/new.h
index ea96fd60e0..7108b0e7a0 100644
--- a/library/cpp/yt/memory/new.h
+++ b/library/cpp/yt/memory/new.h
@@ -1,7 +1,6 @@
 #pragma once
 
 #include "intrusive_ptr.h"
-#include "ref_tracked.h"
 
 #include <library/cpp/yt/misc/source_location.h>
 
@@ -93,9 +92,9 @@ TIntrusivePtr<T> NewWithExtraSpace(size_t extraSpaceSize, As&&... args);
 template <class T, class... As, class = typename THasAllocator<T>::TTrue>
 TIntrusivePtr<T> NewWithExtraSpace(typename T::TAllocator* allocator, size_t extraSpaceSize, As&&... args);
 
-//! Allocates a new instance of |T| with user deleter.
+//! Allocates a new instance of |T| with a custom #deleter.
 template <class T, class TDeleter, class... As>
-TIntrusivePtr<T> NewWithDelete(const TDeleter& deleter, As&&... args);
+TIntrusivePtr<T> NewWithDeleter(TDeleter deleter, As&&... args);
 
 //! Allocates a new instance of |T|.
 //! The allocation is additionally marked with #location.
diff --git a/library/cpp/yt/memory/ref_counted-inl.h b/library/cpp/yt/memory/ref_counted-inl.h
index 5f1f60af7b..fcd64abb17 100644
--- a/library/cpp/yt/memory/ref_counted-inl.h
+++ b/library/cpp/yt/memory/ref_counted-inl.h
@@ -72,78 +72,39 @@ struct TMemoryReleaser<T, std::enable_if_t<T::EnableHazard>>
 
 ////////////////////////////////////////////////////////////////////////////////
 
-} // namespace NDetail
-
-////////////////////////////////////////////////////////////////////////////////
-
-Y_FORCE_INLINE int TRefCounter::GetRefCount() const noexcept
-{
-    return StrongCount_.load(std::memory_order::acquire);
-}
-
-Y_FORCE_INLINE void TRefCounter::Ref(int n) const noexcept
-{
-    // It is safe to use relaxed here, since new reference is always created from another live reference.
-    StrongCount_.fetch_add(n, std::memory_order::relaxed);
-
-    YT_ASSERT(WeakCount_.load(std::memory_order::relaxed) > 0);
-}
-
-Y_FORCE_INLINE bool TRefCounter::TryRef() const noexcept
+template <class T>
+Y_FORCE_INLINE void DestroyRefCountedImpl(T* obj)
 {
-    auto value = StrongCount_.load(std::memory_order::relaxed);
-    YT_ASSERT(WeakCount_.load(std::memory_order::relaxed) > 0);
+    // No standard way to statically calculate the base offset even if T is final.
+    // static_cast<TFinalDerived*>(virtualBasePtr) does not work.
+    auto* basePtr = static_cast<TRefCountedBase*>(obj);
+    auto offset = reinterpret_cast<uintptr_t>(basePtr) - reinterpret_cast<uintptr_t>(obj);
+    auto* refCounter = GetRefCounter(obj);
 
-    while (value != 0 && !StrongCount_.compare_exchange_weak(value, value + 1));
-    return value != 0;
-}
+    // No virtual call when T is final.
+    obj->~T();
 
-Y_FORCE_INLINE bool TRefCounter::Unref(int n) const
-{
-    // We must properly synchronize last access to object with it destruction.
-    // Otherwise compiler might reorder access to object past this decrement.
-    //
-    // See http://www.boost.org/doc/libs/1_55_0/doc/html/atomic/usage_examples.html#boost_atomic.usage_examples.example_reference_counters
-    //
-    auto oldStrongCount = StrongCount_.fetch_sub(n, std::memory_order::release);
-    YT_ASSERT(oldStrongCount >= n);
-    if (oldStrongCount == n) {
-        std::atomic_thread_fence(std::memory_order::acquire);
-        NSan::Acquire(&StrongCount_);
-        return true;
-    } else {
-        return false;
+    // Fast path. Weak refs cannot appear if there are neither strong nor weak refs.
+    if (refCounter->GetWeakRefCount() == 1) {
+        NYT::NDetail::TMemoryReleaser<T>::Do(obj, offset);
+        return;
     }
-}
 
-Y_FORCE_INLINE int TRefCounter::GetWeakRefCount() const noexcept
-{
-    return WeakCount_.load(std::memory_order::acquire);
-}
+    YT_ASSERT(offset < std::numeric_limits<ui16>::max());
 
-Y_FORCE_INLINE void TRefCounter::WeakRef() const noexcept
-{
-    auto oldWeakCount = WeakCount_.fetch_add(1, std::memory_order::relaxed);
-    YT_ASSERT(oldWeakCount > 0);
-}
+    auto* vTablePtr = reinterpret_cast<TPackedPtr*>(basePtr);
+    *vTablePtr = TTaggedPtr<void(void*, ui16)>(&NYT::NDetail::TMemoryReleaser<T>::Do, offset).Pack();
 
-Y_FORCE_INLINE bool TRefCounter::WeakUnref() const
-{
-    auto oldWeakCount = WeakCount_.fetch_sub(1, std::memory_order::release);
-    YT_ASSERT(oldWeakCount > 0);
-    if (oldWeakCount == 1) {
-        std::atomic_thread_fence(std::memory_order::acquire);
-        NSan::Acquire(&WeakCount_);
-        return true;
-    } else {
-        return false;
+    if (refCounter->WeakUnref()) {
+        NYT::NDetail::TMemoryReleaser<T>::Do(obj, offset);
     }
 }
 
 ////////////////////////////////////////////////////////////////////////////////
 
+// Specialization for final classes.
 template <class T, bool = std::is_base_of_v<TRefCountedBase, T>>
-struct TRefCountedHelper
+struct TRefCountedTraits
 {
     static_assert(
         std::is_final_v<T>,
@@ -184,8 +145,9 @@ struct TRefCountedHelper
     }
 };
 
+// Specialization for classes derived from TRefCountedBase.
 template <class T>
-struct TRefCountedHelper<T, true>
+struct TRefCountedTraits<T, true>
 {
     Y_FORCE_INLINE static const TRefCounter* GetRefCounter(const T* obj)
     {
@@ -209,22 +171,92 @@ struct TRefCountedHelper<T, true>
 
 ////////////////////////////////////////////////////////////////////////////////
 
+} // namespace NDetail
+
+////////////////////////////////////////////////////////////////////////////////
+
+Y_FORCE_INLINE int TRefCounter::GetRefCount() const noexcept
+{
+    return StrongCount_.load(std::memory_order::acquire);
+}
+
+Y_FORCE_INLINE void TRefCounter::Ref(int n) const noexcept
+{
+    // It is safe to use relaxed here, since new reference is always created from another live reference.
+    StrongCount_.fetch_add(n, std::memory_order::relaxed);
+
+    YT_ASSERT(WeakCount_.load(std::memory_order::relaxed) > 0);
+}
+
+Y_FORCE_INLINE bool TRefCounter::TryRef() const noexcept
+{
+    auto value = StrongCount_.load(std::memory_order::relaxed);
+    YT_ASSERT(WeakCount_.load(std::memory_order::relaxed) > 0);
+
+    while (value != 0 && !StrongCount_.compare_exchange_weak(value, value + 1));
+    return value != 0;
+}
+
+Y_FORCE_INLINE bool TRefCounter::Unref(int n) const
+{
+    // We must properly synchronize last access to object with it destruction.
+    // Otherwise compiler might reorder access to object past this decrement.
+    //
+    // See http://www.boost.org/doc/libs/1_55_0/doc/html/atomic/usage_examples.html#boost_atomic.usage_examples.example_reference_counters
+    //
+    auto oldStrongCount = StrongCount_.fetch_sub(n, std::memory_order::release);
+    YT_ASSERT(oldStrongCount >= n);
+    if (oldStrongCount == n) {
+        std::atomic_thread_fence(std::memory_order::acquire);
+        NSan::Acquire(&StrongCount_);
+        return true;
+    } else {
+        return false;
+    }
+}
+
+Y_FORCE_INLINE int TRefCounter::GetWeakRefCount() const noexcept
+{
+    return WeakCount_.load(std::memory_order::acquire);
+}
+
+Y_FORCE_INLINE void TRefCounter::WeakRef() const noexcept
+{
+    auto oldWeakCount = WeakCount_.fetch_add(1, std::memory_order::relaxed);
+    YT_ASSERT(oldWeakCount > 0);
+}
+
+Y_FORCE_INLINE bool TRefCounter::WeakUnref() const
+{
+    auto oldWeakCount = WeakCount_.fetch_sub(1, std::memory_order::release);
+    YT_ASSERT(oldWeakCount > 0);
+    if (oldWeakCount == 1) {
+        std::atomic_thread_fence(std::memory_order::acquire);
+        NSan::Acquire(&WeakCount_);
+        return true;
+    } else {
+        return false;
+    }
+}
+
+////////////////////////////////////////////////////////////////////////////////
+
 template <class T>
 Y_FORCE_INLINE const TRefCounter* GetRefCounter(const T* obj)
 {
-    return TRefCountedHelper<T>::GetRefCounter(obj);
+    return NYT::NDetail::TRefCountedTraits<T>::GetRefCounter(obj);
 }
 
 template <class T>
 Y_FORCE_INLINE void DestroyRefCounted(const T* obj)
 {
-    TRefCountedHelper<T>::Destroy(obj);
+    NYT::NDetail::TRefCountedTraits<T>::Destroy(obj);
 }
 
 template <class T>
 Y_FORCE_INLINE void DeallocateRefCounted(const T* obj)
 {
-    TRefCountedHelper<T>::Deallocate(obj);
+    NYT::NDetail::TRefCountedTraits<T>::Deallocate(obj);
 }
 
 ////////////////////////////////////////////////////////////////////////////////
@@ -258,32 +290,9 @@ Y_FORCE_INLINE void TRefCounted::WeakUnref() const
 }
 
 template <class T>
-void TRefCounted::DestroyRefCountedImpl(T* ptr)
+void TRefCounted::DestroyRefCountedImpl(T* obj)
 {
-    // No standard way to statically calculate the base offset even if T is final.
-    // static_cast<TFinalDerived*>(virtualBasePtr) does not work.
-
-    auto* basePtr = static_cast<TRefCountedBase*>(ptr);
-    auto offset = reinterpret_cast<uintptr_t>(basePtr) - reinterpret_cast<uintptr_t>(ptr);
-    auto* refCounter = GetRefCounter(ptr);
-
-    // No virtual call when T is final.
-    ptr->~T();
-
-    // Fast path. Weak refs cannot appear if there are neither strong nor weak refs.
-    if (refCounter->GetWeakRefCount() == 1) {
-        NYT::NDetail::TMemoryReleaser<T>::Do(ptr, offset);
-        return;
-    }
-
-    YT_ASSERT(offset < std::numeric_limits<ui16>::max());
-
-    auto* vTablePtr = reinterpret_cast<TPackedPtr*>(basePtr);
-    *vTablePtr = TTaggedPtr<void(void*, ui16)>(&NYT::NDetail::TMemoryReleaser<T>::Do, offset).Pack();
-
-    if (refCounter->WeakUnref()) {
-        NYT::NDetail::TMemoryReleaser<T>::Do(ptr, offset);
-    }
+    NYT::NDetail::DestroyRefCountedImpl<T>(obj);
 }
 
 ////////////////////////////////////////////////////////////////////////////////
diff --git a/library/cpp/yt/memory/ref_counted.h b/library/cpp/yt/memory/ref_counted.h
index 6327141cd8..6ef546de0d 100644
--- a/library/cpp/yt/memory/ref_counted.h
+++ b/library/cpp/yt/memory/ref_counted.h
@@ -16,17 +16,15 @@ class TRefCountedBase
 public:
     TRefCountedBase() = default;
 
-    // Make destructor protected
-    virtual ~TRefCountedBase() noexcept = default;
-
-    virtual void DestroyRefCounted() = 0;
-
-private:
     TRefCountedBase(const TRefCountedBase&) = delete;
     TRefCountedBase(TRefCountedBase&&) = delete;
 
     TRefCountedBase& operator=(const TRefCountedBase&) = delete;
     TRefCountedBase& operator=(TRefCountedBase&&) = delete;
+
+    virtual ~TRefCountedBase() noexcept = default;
+
+    virtual void DestroyRefCounted() = 0;
 };
 
 ////////////////////////////////////////////////////////////////////////////////
@@ -85,16 +83,17 @@ void Unref(T* obj, int n = 1);
 
 ////////////////////////////////////////////////////////////////////////////////
 
-struct TRefCounted
+class TRefCounted
     : public TRefCountedBase
     , public TRefCounter
 {
+public:
     void Unref() const;
-
     void WeakUnref() const;
 
+protected:
     template <class T>
-    static void DestroyRefCountedImpl(T* ptr);
+    static void DestroyRefCountedImpl(T* obj);
 };
 
 ////////////////////////////////////////////////////////////////////////////////
@@ -145,15 +144,15 @@ using TRefCountedPtr = TIntrusivePtr<TRefCounted>;
 #define DEFINE_REFCOUNTED_TYPE(type) \
     [[maybe_unused]] YT_ATTRIBUTE_USED Y_FORCE_INLINE const ::NYT::TRefCounter* GetRefCounter(const type* obj) \
     { \
-        return ::NYT::TRefCountedHelper<type>::GetRefCounter(obj); \
+        return ::NYT::NDetail::TRefCountedTraits<type>::GetRefCounter(obj); \
     } \
     [[maybe_unused]] YT_ATTRIBUTE_USED Y_FORCE_INLINE void DestroyRefCounted(const type* obj) \
     { \
-        ::NYT::TRefCountedHelper<type>::Destroy(obj); \
+        ::NYT::NDetail::TRefCountedTraits<type>::Destroy(obj); \
     } \
     [[maybe_unused]] YT_ATTRIBUTE_USED Y_FORCE_INLINE void DeallocateRefCounted(const type* obj) \
     { \
-        ::NYT::TRefCountedHelper<type>::Deallocate(obj); \
+        ::NYT::NDetail::TRefCountedTraits<type>::Deallocate(obj); \
     }
 
 //! Provides weak implementations for Ref/Unref overloads.