More assertions in refcounters

commit_hash:0c6147f498fdf817921889ac84a321f84f2d1059

3 files changed, 32 insertions, 8 deletions

diff --git a/library/cpp/yt/memory/atomic_intrusive_ptr-inl.h b/library/cpp/yt/memory/atomic_intrusive_ptr-inl.h
index 21cd8b9412..14f2e10365 100644
--- a/library/cpp/yt/memory/atomic_intrusive_ptr-inl.h
+++ b/library/cpp/yt/memory/atomic_intrusive_ptr-inl.h
@@ -17,7 +17,7 @@ TAtomicIntrusivePtr<T>::TAtomicIntrusivePtr(std::nullptr_t)
 
 template <class T>
 TAtomicIntrusivePtr<T>::TAtomicIntrusivePtr(TIntrusivePtr<T> other)
-    : Ptr_(AcquireObject(other.Release(), true))
+    : Ptr_(AcquireObject(other.Release(), /*consumeRef*/ true))
 { }
 
 template <class T>
@@ -100,7 +100,7 @@ TIntrusivePtr<T> TAtomicIntrusivePtr<T>::Acquire() const
 template <class T>
 TIntrusivePtr<T> TAtomicIntrusivePtr<T>::Exchange(TIntrusivePtr<T> other)
 {
-    auto [obj, localRefs] = TTaggedPtr<T>::Unpack(Ptr_.exchange(AcquireObject(other.Release(), true)));
+    auto [obj, localRefs] = TTaggedPtr<T>::Unpack(Ptr_.exchange(AcquireObject(other.Release(), /*consumeRef*/ true)));
     DoRelease(obj, localRefs + 1);
     return TIntrusivePtr<T>(obj, false);
 }
@@ -108,7 +108,7 @@ TIntrusivePtr<T> TAtomicIntrusivePtr<T>::Exchange(TIntrusivePtr<T> other)
 template <class T>
 void TAtomicIntrusivePtr<T>::Store(TIntrusivePtr<T> other)
 {
-    ReleaseObject(Ptr_.exchange(AcquireObject(other.Release(), true)));
+    ReleaseObject(Ptr_.exchange(AcquireObject(other.Release(), /*consumeRef*/ true)));
 }
 
 template <class T>
@@ -120,7 +120,7 @@ void TAtomicIntrusivePtr<T>::Reset()
 template <class T>
 bool TAtomicIntrusivePtr<T>::CompareAndSwap(TRawPtr& comparePtr, T* target)
 {
-    auto* targetPtr = AcquireObject(target, false);
+    auto* targetPtr = AcquireObject(target, /*consumeRef*/ false);
 
     auto currentPtr = Ptr_.load();
     if (UnpackPointer<T>(currentPtr).Ptr == comparePtr && Ptr_.compare_exchange_strong(currentPtr, targetPtr)) {
@@ -138,7 +138,7 @@ template <class T>
 bool TAtomicIntrusivePtr<T>::CompareAndSwap(TRawPtr& comparePtr, TIntrusivePtr<T> target)
 {
     // TODO(lukyan): Make helper for packed owning ptr?
-    auto targetPtr = AcquireObject(target.Release(), true);
+    auto targetPtr = AcquireObject(target.Release(), /*consumeRef*/ true);
 
     auto currentPtr = Ptr_.load();
     if (TTaggedPtr<T>::Unpack(currentPtr).Ptr == comparePtr && Ptr_.compare_exchange_strong(currentPtr, targetPtr)) {
@@ -168,7 +168,7 @@ template <class T>
 TPackedPtr TAtomicIntrusivePtr<T>::AcquireObject(T* obj, bool consumeRef)
 {
     if (obj) {
-        Ref(obj, static_cast<int>(ReservedRefCount - consumeRef));
+        Ref(obj, ReservedRefCount - static_cast<int>(consumeRef));
     }
 
     return TTaggedPtr(obj).Pack();
@@ -185,7 +185,7 @@ template <class T>
 void TAtomicIntrusivePtr<T>::DoRelease(T* obj, int refs)
 {
     if (obj) {
-        Unref(obj, static_cast<int>(ReservedRefCount - refs));
+        Unref(obj, ReservedRefCount - refs);
     }
 }
 
diff --git a/library/cpp/yt/memory/ref_counted-inl.h b/library/cpp/yt/memory/ref_counted-inl.h
index f86f634ffd..9ffdf328eb 100644
--- a/library/cpp/yt/memory/ref_counted-inl.h
+++ b/library/cpp/yt/memory/ref_counted-inl.h
@@ -182,8 +182,24 @@ Y_FORCE_INLINE int TRefCounter::GetRefCount() const noexcept
 
 Y_FORCE_INLINE void TRefCounter::Ref(int n) const noexcept
 {
+    YT_ASSERT(n >= 0);
+
     // 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);
+    auto value = StrongCount_.fetch_add(n, std::memory_order::relaxed);
+    YT_ASSERT(value > 0);
+    YT_ASSERT(value <= std::numeric_limits<int>::max() - n);
+
+    YT_ASSERT(WeakCount_.load(std::memory_order::relaxed) > 0);
+}
+
+Y_FORCE_INLINE void TRefCounter::DangerousRef(int n) const noexcept
+{
+    YT_ASSERT(n >= 0);
+
+    // Relaxed is fine as per lukyan@, the caller guarantees object liveness.
+    auto value = StrongCount_.fetch_add(n, std::memory_order::relaxed);
+    YT_ASSERT(value >= 0);
+    YT_ASSERT(value <= std::numeric_limits<int>::max() - n);
 
     YT_ASSERT(WeakCount_.load(std::memory_order::relaxed) > 0);
 }
@@ -191,6 +207,7 @@ Y_FORCE_INLINE void TRefCounter::Ref(int n) const noexcept
 Y_FORCE_INLINE bool TRefCounter::TryRef() const noexcept
 {
     auto value = StrongCount_.load(std::memory_order::relaxed);
+    YT_ASSERT(value >= 0 && value < std::numeric_limits<int>::max());
     YT_ASSERT(WeakCount_.load(std::memory_order::relaxed) > 0);
 
     while (value != 0 && !StrongCount_.compare_exchange_weak(value, value + 1));
@@ -199,6 +216,8 @@ Y_FORCE_INLINE bool TRefCounter::TryRef() const noexcept
 
 Y_FORCE_INLINE bool TRefCounter::Unref(int n) const
 {
+    YT_ASSERT(n >= 0);
+
     // We must properly synchronize last access to object with it destruction.
     // Otherwise compiler might reorder access to object past this decrement.
     //
diff --git a/library/cpp/yt/memory/ref_counted.h b/library/cpp/yt/memory/ref_counted.h
index 6ef546de0d..ed55aee6ee 100644
--- a/library/cpp/yt/memory/ref_counted.h
+++ b/library/cpp/yt/memory/ref_counted.h
@@ -40,8 +40,13 @@ public:
     int GetRefCount() const noexcept;
 
     //! Increments the strong reference counter.
+    //! The current strong RC must be positive.
     void Ref(int n = 1) const noexcept;
 
+    //! Increments the strong reference counter.
+    //! The current strong RC may be zero.
+    void DangerousRef(int n = 1) const noexcept;
+
     //! Increments the strong reference counter if it is not null.
     bool TryRef() const noexcept;