YT-27536: Fix (seemingly benign) race in TTwoLevelFairShareQueue

commit_hash:700b15a49621e9f79e6e6c725a1aa958f9bd132e

3 files changed, 331 insertions, 227 deletions

diff --git a/library/cpp/yt/memory/unittests/weak_ptr_ut.cpp b/library/cpp/yt/memory/unittests/weak_ptr_ut.cpp
index 98fb4d19036..bc885c267fe 100644
--- a/library/cpp/yt/memory/unittests/weak_ptr_ut.cpp
+++ b/library/cpp/yt/memory/unittests/weak_ptr_ut.cpp
@@ -308,7 +308,8 @@ TEST_F(TWeakPtrTest, IsEmpty)
 {
     TIntricateObjectWkPtr ptr;
 
-    EXPECT_TRUE(ptr == nullptr);
+    EXPECT_EQ(ptr.Get(), nullptr);
+    EXPECT_TRUE(ptr.IsExpired());
 }
 
 TEST_F(TWeakPtrTest, UpCast)
diff --git a/library/cpp/yt/memory/weak_ptr-inl.h b/library/cpp/yt/memory/weak_ptr-inl.h
new file mode 100644
index 00000000000..41090988b98
--- /dev/null
+++ b/library/cpp/yt/memory/weak_ptr-inl.h
@@ -0,0 +1,280 @@
+#ifndef WEAK_PTR_INL_H_
+#error "Direct inclusion of this file is not allowed, include weak_ptr.h"
+// For the sake of sane code completion.
+#include "weak_ptr.h"
+#endif
+
+namespace NYT {
+
+////////////////////////////////////////////////////////////////////////////////
+
+template <class T>
+TWeakPtr<T>::TWeakPtr(std::nullptr_t)
+{ }
+
+template <class T>
+TWeakPtr<T>::TWeakPtr(T* p) noexcept
+    : T_(p)
+{
+#if defined(_tsan_enabled_)
+    if (T_) {
+        RefCounter_ = GetRefCounter(T_);
+    }
+#endif
+    AcquireRef();
+}
+
+template <class T>
+TWeakPtr<T>::TWeakPtr(const TIntrusivePtr<T>& ptr) noexcept
+    : TWeakPtr(ptr.Get())
+{ }
+
+template <class T>
+template <class U, class>
+TWeakPtr<T>::TWeakPtr(const TIntrusivePtr<U>& ptr) noexcept
+    : TWeakPtr(ptr.Get())
+{
+    static_assert(
+        std::derived_from<T, TRefCountedBase>,
+        "Cast allowed only for types derived from TRefCountedBase");
+}
+
+template <class T>
+TWeakPtr<T>::TWeakPtr(const TWeakPtr& other) noexcept
+    : T_(other.T_)
+#if defined(_tsan_enabled_)
+    , RefCounter_(other.RefCounter_)
+#endif
+{
+    AcquireRef();
+}
+
+template <class T>
+template <class U, class>
+TWeakPtr<T>::TWeakPtr(const TWeakPtr<U>& other) noexcept
+    : TWeakPtr(other.Lock())
+{
+    static_assert(
+        std::derived_from<T, TRefCountedBase>,
+        "Cast allowed only for types derived from TRefCountedBase");
+}
+
+template <class T>
+TWeakPtr<T>::TWeakPtr(TWeakPtr&& other) noexcept
+{
+    other.Swap(*this);
+}
+
+template <class T>
+template <class U, class>
+TWeakPtr<T>::TWeakPtr(TWeakPtr<U>&& other) noexcept
+{
+    static_assert(
+        std::derived_from<T, TRefCountedBase>,
+        "Cast allowed only for types derived from TRefCountedBase");
+    TIntrusivePtr<U> strongOther = other.Lock();
+    if (strongOther) {
+        T_ = other.T_;
+        other.T_ = nullptr;
+
+#if defined(_tsan_enabled_)
+        RefCounter_ = other.RefCounter_;
+        other.RefCounter_ = nullptr;
+#endif
+    }
+}
+
+template <class T>
+TWeakPtr<T>::~TWeakPtr()
+{
+    ReleaseRef();
+}
+
+template <class T>
+template <class U>
+TWeakPtr<T>& TWeakPtr<T>::operator=(const TIntrusivePtr<U>& ptr) noexcept
+{
+    static_assert(
+        std::is_convertible_v<U*, T*>,
+        "U* must be convertible to T*");
+    TWeakPtr(ptr).Swap(*this);
+    return *this;
+}
+
+template <class T>
+TWeakPtr<T>& TWeakPtr<T>::operator=(const TWeakPtr& other) noexcept
+{
+    TWeakPtr(other).Swap(*this);
+    return *this;
+}
+
+template <class T>
+template <class U>
+TWeakPtr<T>& TWeakPtr<T>::operator=(const TWeakPtr<U>& other) noexcept
+{
+    static_assert(
+        std::is_convertible_v<U*, T*>,
+        "U* must be convertible to T*");
+    TWeakPtr(other).Swap(*this);
+    return *this;
+}
+
+template <class T>
+TWeakPtr<T>& TWeakPtr<T>::operator=(TWeakPtr&& other) noexcept
+{
+    other.Swap(*this);
+    return *this;
+}
+
+template <class T>
+template <class U>
+TWeakPtr<T>& TWeakPtr<T>::operator=(TWeakPtr<U>&& other) noexcept
+{
+    static_assert(
+        std::is_convertible_v<U*, T*>,
+        "U* must be convertible to T*");
+    TWeakPtr(std::move(other)).Swap(*this);
+    return *this;
+}
+
+template <class T>
+void TWeakPtr<T>::Reset() // noexcept
+{
+    TWeakPtr().Swap(*this);
+}
+
+template <class T>
+void TWeakPtr<T>::Reset(T* p) // noexcept
+{
+    TWeakPtr(p).Swap(*this);
+}
+
+template <class T>
+template <class U>
+void TWeakPtr<T>::Reset(const TIntrusivePtr<U>& ptr) // noexcept
+{
+    static_assert(
+        std::is_convertible_v<U*, T*>,
+        "U* must be convertible to T*");
+    TWeakPtr(ptr).Swap(*this);
+}
+
+template <class T>
+void TWeakPtr<T>::Swap(TWeakPtr& other) noexcept
+{
+    DoSwap(T_, other.T_);
+#if defined(_tsan_enabled_)
+    DoSwap(RefCounter_, other.RefCounter_);
+#endif
+}
+
+template <class T>
+TIntrusivePtr<T> TWeakPtr<T>::Lock() const noexcept
+{
+    return T_ && GetRefCounterImpl()->TryRef()
+        ? TIntrusivePtr<T>(T_, false)
+        : TIntrusivePtr<T>();
+}
+
+template <class T>
+bool TWeakPtr<T>::IsExpired() const noexcept
+{
+    return !T_ || (GetRefCounterImpl()->GetRefCount() == 0);
+}
+
+template <class T>
+T* TWeakPtr<T>::Get() const noexcept
+{
+    return T_;
+}
+
+template <class T>
+const TRefCounter* TWeakPtr<T>::TryGetRefCounterImpl() const
+{
+    return T_ ? GetRefCounterImpl() : nullptr;
+}
+
+template <class T>
+void TWeakPtr<T>::AcquireRef()
+{
+    if (T_) {
+        GetRefCounterImpl()->WeakRef();
+    }
+}
+
+template <class T>
+void TWeakPtr<T>::ReleaseRef()
+{
+    if (T_) {
+        // Support incomplete type.
+        if (GetRefCounterImpl()->WeakUnref()) {
+            DeallocateRefCounted(T_);
+        }
+    }
+}
+
+#if defined(_tsan_enabled_)
+template <class T>
+const TRefCounter* TWeakPtr<T>::GetRefCounterImpl() const
+{
+    return RefCounter_;
+}
+#else
+template <class T>
+const TRefCounter* TWeakPtr<T>::GetRefCounterImpl() const
+{
+    return GetRefCounter(T_);
+}
+#endif
+
+////////////////////////////////////////////////////////////////////////////////
+
+template <class T1, class T2>
+bool operator==(const TWeakPtr<T1>& lhs, const TWeakPtr<T2>& rhs)
+{
+    return lhs.Get() == rhs.Get();
+}
+
+template <class T2>
+bool operator==(std::nullptr_t, const TWeakPtr<T2>& rhs)
+{
+    return nullptr == rhs.Get();
+}
+
+template <class T1>
+bool operator==(const TWeakPtr<T1>& lhs, std::nullptr_t)
+{
+    return lhs.Get() == nullptr;
+}
+
+////////////////////////////////////////////////////////////////////////////////
+
+template <class T>
+TWeakPtr<T> MakeWeak(T* p)
+{
+    return TWeakPtr<T>(p);
+}
+
+template <class T>
+TWeakPtr<T> MakeWeak(const TIntrusivePtr<T>& p)
+{
+    return TWeakPtr<T>(p);
+}
+
+template <typename T>
+int ResetAndGetResidualRefCount(TIntrusivePtr<T>& pointer)
+{
+    auto weakPointer = MakeWeak(pointer);
+    pointer.Reset();
+    pointer = weakPointer.Lock();
+    if (pointer) {
+        // This _may_ return 0 if we are again the only holder of the pointee.
+        return pointer->GetRefCount() - 1;
+    } else {
+        return 0;
+    }
+}
+
+////////////////////////////////////////////////////////////////////////////////
+
+} // namespace NYT
diff --git a/library/cpp/yt/memory/weak_ptr.h b/library/cpp/yt/memory/weak_ptr.h
index 5298954e563..126264b0c24 100644
--- a/library/cpp/yt/memory/weak_ptr.h
+++ b/library/cpp/yt/memory/weak_ptr.h
@@ -15,210 +15,87 @@ public:
     using TUnderlying = T;
     using element_type = T;
 
-    //! Empty constructor.
+    //! Default constructor.
     TWeakPtr() = default;
 
-    TWeakPtr(std::nullptr_t)
-    { }
+    //! Null constructor.
+    TWeakPtr(std::nullptr_t);
 
     //! Constructor from an unqualified reference.
     /*!
      * Note that this constructor could be racy due to unsynchronized operations
      * on the object and on the counter.
      */
-    explicit TWeakPtr(T* p) noexcept
-        : T_(p)
-    {
-
-#if defined(_tsan_enabled_)
-        if (T_) {
-            RefCounter_ = GetRefCounter(T_);
-        }
-#endif
-        AcquireRef();
-    }
+    explicit TWeakPtr(T* p) noexcept;
 
     //! Constructor from a strong reference.
-    TWeakPtr(const TIntrusivePtr<T>& ptr) noexcept
-        : TWeakPtr(ptr.Get())
-    { }
+    TWeakPtr(const TIntrusivePtr<T>& ptr) noexcept;
 
     //! Constructor from a strong reference with an upcast.
     template <class U, class = typename std::enable_if_t<std::is_convertible_v<U*, T*>>>
-    TWeakPtr(const TIntrusivePtr<U>& ptr) noexcept
-        : TWeakPtr(ptr.Get())
-    {
-        static_assert(
-            std::derived_from<T, TRefCountedBase>,
-            "Cast allowed only for types derived from TRefCountedBase");
-    }
+    TWeakPtr(const TIntrusivePtr<U>& ptr) noexcept;
 
     //! Copy constructor.
-    TWeakPtr(const TWeakPtr& other) noexcept
-        : T_(other.T_)
-#if defined(_tsan_enabled_)
-        , RefCounter_(other.RefCounter_)
-#endif
-    {
-        AcquireRef();
-    }
+    TWeakPtr(const TWeakPtr& other) noexcept;
 
     //! Copy constructor with an upcast.
     template <class U, class = typename std::enable_if_t<std::is_convertible_v<U*, T*>>>
-    TWeakPtr(const TWeakPtr<U>& other) noexcept
-        : TWeakPtr(other.Lock())
-    {
-        static_assert(
-            std::derived_from<T, TRefCountedBase>,
-            "Cast allowed only for types derived from TRefCountedBase");
-    }
+    TWeakPtr(const TWeakPtr<U>& other) noexcept;
 
     //! Move constructor.
-    TWeakPtr(TWeakPtr&& other) noexcept
-    {
-        other.Swap(*this);
-    }
+    TWeakPtr(TWeakPtr&& other) noexcept;
 
     //! Move constructor with an upcast.
     template <class U, class = typename std::enable_if_t<std::is_convertible_v<U*, T*>>>
-    TWeakPtr(TWeakPtr<U>&& other) noexcept
-    {
-        static_assert(
-            std::derived_from<T, TRefCountedBase>,
-            "Cast allowed only for types derived from TRefCountedBase");
-        TIntrusivePtr<U> strongOther = other.Lock();
-        if (strongOther) {
-            T_ = other.T_;
-            other.T_ = nullptr;
-
-#if defined(_tsan_enabled_)
-            RefCounter_ = other.RefCounter_;
-            other.RefCounter_ = nullptr;
-#endif
-        }
-    }
+    TWeakPtr(TWeakPtr<U>&& other) noexcept;
 
     //! Destructor.
-    ~TWeakPtr()
-    {
-        ReleaseRef();
-    }
+    ~TWeakPtr();
 
     //! Assignment operator from a strong reference.
     template <class U>
-    TWeakPtr& operator=(const TIntrusivePtr<U>& ptr) noexcept
-    {
-        static_assert(
-            std::is_convertible_v<U*, T*>,
-            "U* must be convertible to T*");
-        TWeakPtr(ptr).Swap(*this);
-        return *this;
-    }
+    TWeakPtr& operator=(const TIntrusivePtr<U>& ptr) noexcept;
 
     //! Copy assignment operator.
-    TWeakPtr& operator=(const TWeakPtr& other) noexcept
-    {
-        TWeakPtr(other).Swap(*this);
-        return *this;
-    }
+    TWeakPtr& operator=(const TWeakPtr& other) noexcept;
 
     //! Copy assignment operator with an upcast.
     template <class U>
-    TWeakPtr& operator=(const TWeakPtr<U>& other) noexcept
-    {
-        static_assert(
-            std::is_convertible_v<U*, T*>,
-            "U* must be convertible to T*");
-        TWeakPtr(other).Swap(*this);
-        return *this;
-    }
+    TWeakPtr& operator=(const TWeakPtr<U>& other) noexcept;
 
     //! Move assignment operator.
-    TWeakPtr& operator=(TWeakPtr&& other) noexcept
-    {
-        other.Swap(*this);
-        return *this;
-    }
+    TWeakPtr& operator=(TWeakPtr&& other) noexcept;
 
     //! Move assignment operator with an upcast.
     template <class U>
-    TWeakPtr& operator=(TWeakPtr<U>&& other) noexcept
-    {
-        static_assert(
-            std::is_convertible_v<U*, T*>,
-            "U* must be convertible to T*");
-        TWeakPtr(std::move(other)).Swap(*this);
-        return *this;
-    }
+    TWeakPtr& operator=(TWeakPtr<U>&& other) noexcept;
 
     //! Drop the pointer.
-    void Reset() // noexcept
-    {
-        TWeakPtr().Swap(*this);
-    }
+    void Reset(); // noexcept
 
     //! Replace the pointer with a specified one.
-    void Reset(T* p) // noexcept
-    {
-        TWeakPtr(p).Swap(*this);
-    }
+    void Reset(T* p); // noexcept
 
     //! Replace the pointer with a specified one.
     template <class U>
-    void Reset(const TIntrusivePtr<U>& ptr) // noexcept
-    {
-        static_assert(
-            std::is_convertible_v<U*, T*>,
-            "U* must be convertible to T*");
-        TWeakPtr(ptr).Swap(*this);
-    }
+    void Reset(const TIntrusivePtr<U>& ptr); // noexcept
 
     //! Swap the pointer with the other one.
-    void Swap(TWeakPtr& other) noexcept
-    {
-        DoSwap(T_, other.T_);
-#if defined(_tsan_enabled_)
-        DoSwap(RefCounter_, other.RefCounter_);
-#endif
-    }
+    void Swap(TWeakPtr& other) noexcept;
 
     //! Acquire a strong reference to the pointee and return a strong pointer.
-    TIntrusivePtr<T> Lock() const noexcept
-    {
-        return T_ && RefCounter()->TryRef()
-            ? TIntrusivePtr<T>(T_, false)
-            : TIntrusivePtr<T>();
-    }
+    TIntrusivePtr<T> Lock() const noexcept;
 
-    bool IsExpired() const noexcept
-    {
-        return !T_ || (RefCounter()->GetRefCount() == 0);
-    }
+    //! Returns true if the pointee is null or already expired.
+    bool IsExpired() const noexcept;
 
-    const TRefCounter* TryGetRefCounter() const
-    {
-        return T_
-            ? RefCounter()
-            : nullptr;
-    }
+    //! Returns the underlying raw pointer. Note that it may point to an already
+    //! expired object.
+    T* Get() const noexcept;
 
 private:
-    void AcquireRef()
-    {
-        if (T_) {
-            RefCounter()->WeakRef();
-        }
-    }
-
-    void ReleaseRef()
-    {
-        if (T_) {
-            // Support incomplete type.
-            if (RefCounter()->WeakUnref()) {
-                DeallocateRefCounted(T_);
-            }
-        }
-    }
+    void AcquireRef();
+    void ReleaseRef();
 
     template <class U>
     friend class TWeakPtr;
@@ -228,36 +105,34 @@ private:
     T* T_ = nullptr;
 #if defined(_tsan_enabled_)
     const TRefCounter* RefCounter_ = nullptr;
-
-    const TRefCounter* RefCounter() const
-    {
-        return RefCounter_;
-    }
-#else
-    const TRefCounter* RefCounter() const
-    {
-        return GetRefCounter(T_);
-    }
 #endif
+
+    const TRefCounter* GetRefCounterImpl() const;
+    const TRefCounter* TryGetRefCounterImpl() const;
 };
 
 ////////////////////////////////////////////////////////////////////////////////
 
+template <class T1, class T2>
+bool operator==(const TWeakPtr<T1>& lhs, const TWeakPtr<T2>& rhs);
+
+template <class T2>
+bool operator==(std::nullptr_t, const TWeakPtr<T2>& rhs);
+
+template <class T1>
+bool operator==(const TWeakPtr<T1>& lhs, std::nullptr_t);
+
+////////////////////////////////////////////////////////////////////////////////
+
 //! Creates a weak pointer wrapper for a given raw pointer.
 //! Compared to |TWeakPtr<T>::ctor|, type inference enables omitting |T|.
 template <class T>
-TWeakPtr<T> MakeWeak(T* p)
-{
-    return TWeakPtr<T>(p);
-}
+TWeakPtr<T> MakeWeak(T* p);
 
 //! Creates a weak pointer wrapper for a given intrusive pointer.
 //! Compared to |TWeakPtr<T>::ctor|, type inference enables omitting |T|.
 template <class T>
-TWeakPtr<T> MakeWeak(const TIntrusivePtr<T>& p)
-{
-    return TWeakPtr<T>(p);
-}
+TWeakPtr<T> MakeWeak(const TIntrusivePtr<T>& p);
 
 //! A helper for acquiring weak pointer for pointee, resetting intrusive pointer and then
 //! returning the pointee reference count using the acquired weak pointer.
@@ -266,68 +141,12 @@ TWeakPtr<T> MakeWeak(const TIntrusivePtr<T>& p)
 //! NB: it is possible to rewrite this helper making it working event with intrinsic refcounted objects,
 //! but it requires much nastier integration with the intrusive pointer destruction routines.
 template <typename T>
-int ResetAndGetResidualRefCount(TIntrusivePtr<T>& pointer)
-{
-    auto weakPointer = MakeWeak(pointer);
-    pointer.Reset();
-    pointer = weakPointer.Lock();
-    if (pointer) {
-        // This _may_ return 0 if we are again the only holder of the pointee.
-        return pointer->GetRefCount() - 1;
-    } else {
-        return 0;
-    }
-}
-
-////////////////////////////////////////////////////////////////////////////////
-
-template <class T, class U>
-bool operator==(const TWeakPtr<T>& lhs, const TWeakPtr<U>& rhs)
-{
-    static_assert(
-        std::is_convertible_v<U*, T*>,
-        "U* must be convertible to T*");
-    return lhs.TryGetRefCounter() == rhs.TryGetRefCounter();
-}
-
-template <class T, class U>
-bool operator!=(const TWeakPtr<T>& lhs, const TWeakPtr<U>& rhs)
-{
-    static_assert(
-        std::is_convertible_v<U*, T*>,
-        "U* must be convertible to T*");
-    return lhs.TryGetRefCounter() != rhs.TryGetRefCounter();
-}
-
-template <class T>
-bool operator==(std::nullptr_t, const TWeakPtr<T>& rhs)
-{
-    return nullptr == rhs.TryGetRefCounter();
-}
-
-template <class T>
-bool operator!=(std::nullptr_t, const TWeakPtr<T>& rhs)
-{
-    return nullptr != rhs.TryGetRefCounter();
-}
-
-template <class T>
-bool operator==(const TWeakPtr<T>& lhs, std::nullptr_t)
-{
-    return nullptr == lhs.TryGetRefCounter();
-}
-
-template <class T>
-bool operator!=(const TWeakPtr<T>& lhs, std::nullptr_t)
-{
-    return nullptr != lhs.TryGetRefCounter();
-}
+int ResetAndGetResidualRefCount(TIntrusivePtr<T>& pointer);
 
 ////////////////////////////////////////////////////////////////////////////////
 
 } // namespace NYT
 
-
 //! A hasher for TWeakPtr.
 template <class T>
 struct THash<NYT::TWeakPtr<T>>
@@ -337,3 +156,7 @@ struct THash<NYT::TWeakPtr<T>>
         return THash<const NYT::TRefCountedBase*>()(ptr.T_);
     }
 };
+
+#define WEAK_PTR_INL_H_
+#include "weak_ptr-inl.h"
+#undef WEAK_PTR_INL_H_