Refactor and improve TSharedRangeHolder

7 files changed, 128 insertions, 28 deletions

diff --git a/library/cpp/yt/memory/public.h b/library/cpp/yt/memory/public.h
index 86965457f8..fb1546dd59 100644
--- a/library/cpp/yt/memory/public.h
+++ b/library/cpp/yt/memory/public.h
@@ -9,6 +9,7 @@ namespace NYT {
 class TChunkedMemoryPool;
 
 DECLARE_REFCOUNTED_STRUCT(IMemoryChunkProvider)
+DECLARE_REFCOUNTED_STRUCT(TSharedRangeHolder)
 
 ////////////////////////////////////////////////////////////////////////////////
 
diff --git a/library/cpp/yt/memory/ref-inl.h b/library/cpp/yt/memory/ref-inl.h
index b3738ab79b..a57bb83a8d 100644
--- a/library/cpp/yt/memory/ref-inl.h
+++ b/library/cpp/yt/memory/ref-inl.h
@@ -371,6 +371,25 @@ public:
         return Begin() + Size_;
     }
 
+
+    // TSharedRangeHolder overrides.
+    std::optional<size_t> GetTotalByteSize() const override
+    {
+        size_t result = 0;
+        for (size_t index = 0; index < Size(); ++index) {
+            const auto& part = (*this)[index];
+            if (!part) {
+                continue;
+            }
+            auto partSize = part.GetHolder()->GetTotalByteSize();
+            if (!partSize) {
+                return std::nullopt;
+            }
+            result += *partSize;
+        }
+        return result;
+    }
+
 private:
     friend class TSharedRefArrayBuilder;
 
diff --git a/library/cpp/yt/memory/ref.cpp b/library/cpp/yt/memory/ref.cpp
index d167247c20..06ec7f9ffe 100644
--- a/library/cpp/yt/memory/ref.cpp
+++ b/library/cpp/yt/memory/ref.cpp
@@ -29,6 +29,8 @@ public:
         : Blob_(std::move(blob))
     { }
 
+    // NB: TBlob's capacity is dynamic; cannot provide GetTotalByteSize implementation.
+
 private:
     const TBlob Blob_;
 };
@@ -63,6 +65,12 @@ public:
         return String_;
     }
 
+    // TSharedRangeHolder overrides.
+    std::optional<size_t> GetTotalByteSize() const override
+    {
+        return String_.capacity();
+    }
+
 private:
     const TString String_;
 #ifdef YT_ENABLE_REF_COUNTED_TRACKING
@@ -96,7 +104,10 @@ protected:
     TRefCountedTypeCookie Cookie_;
 #endif
 
-    void Initialize(size_t size, TSharedMutableRefAllocateOptions options, TRefCountedTypeCookie cookie)
+    void Initialize(
+        size_t size,
+        TSharedMutableRefAllocateOptions options,
+        TRefCountedTypeCookie cookie)
     {
         Size_ = size;
         Cookie_ = cookie;
@@ -117,7 +128,10 @@ class TDefaultAllocationHolder
     , public TWithExtraSpace<TDefaultAllocationHolder>
 {
 public:
-    TDefaultAllocationHolder(size_t size, TSharedMutableRefAllocateOptions options, TRefCountedTypeCookie cookie)
+    TDefaultAllocationHolder(
+        size_t size,
+        TSharedMutableRefAllocateOptions options,
+        TRefCountedTypeCookie cookie)
     {
         if (options.ExtendToUsableSize) {
             if (auto usableSize = GetUsableSpaceSize(); usableSize != 0) {
@@ -131,6 +145,12 @@ public:
     {
         return static_cast<char*>(GetExtraSpacePtr());
     }
+
+    // TSharedRangeHolder overrides.
+    std::optional<size_t> GetTotalByteSize() const override
+    {
+        return Size_;
+    }
 };
 
 ////////////////////////////////////////////////////////////////////////////////
@@ -139,7 +159,10 @@ class TPageAlignedAllocationHolder
     : public TAllocationHolderBase<TPageAlignedAllocationHolder>
 {
 public:
-    TPageAlignedAllocationHolder(size_t size, TSharedMutableRefAllocateOptions options, TRefCountedTypeCookie cookie)
+    TPageAlignedAllocationHolder(
+        size_t size,
+        TSharedMutableRefAllocateOptions options,
+        TRefCountedTypeCookie cookie)
         : Begin_(static_cast<char*>(::aligned_malloc(size, GetPageSize())))
     {
         Initialize(size, options, cookie);
@@ -155,6 +178,12 @@ public:
         return Begin_;
     }
 
+    // TSharedRangeHolder overrides.
+    std::optional<size_t> GetTotalByteSize() const override
+    {
+        return AlignUp(Size_, GetPageSize());
+    }
+
 private:
     char* const Begin_;
 };
diff --git a/library/cpp/yt/memory/shared_range.cpp b/library/cpp/yt/memory/shared_range.cpp
index c147ba5041..4d4bd40d18 100644
--- a/library/cpp/yt/memory/shared_range.cpp
+++ b/library/cpp/yt/memory/shared_range.cpp
@@ -10,6 +10,11 @@ TSharedRangeHolderPtr TSharedRangeHolder::Clone(const TSharedRangeHolderCloneOpt
     return this;
 }
 
+std::optional<size_t> TSharedRangeHolder::GetTotalByteSize() const
+{
+    return std::nullopt;
+}
+
 ////////////////////////////////////////////////////////////////////////////////
 
 TSharedRangeHolderPtr MakeCompositeSharedRangeHolder(std::vector<TSharedRangeHolderPtr> holders)
@@ -17,27 +22,42 @@ TSharedRangeHolderPtr MakeCompositeSharedRangeHolder(std::vector<TSharedRangeHol
     struct THolder
         : public TSharedRangeHolder
     {
-        std::vector<TSharedRangeHolderPtr> Holders;
+        std::vector<TSharedRangeHolderPtr> Subholders;
 
         TSharedRangeHolderPtr Clone(const TSharedRangeHolderCloneOptions& options) override
         {
             auto newHolder = New<THolder>();
-            newHolder->Holders.reserve(Holders.size());
-            for (const auto& holder : Holders) {
-                if (!holder) {
+            newHolder->Subholders.reserve(Subholders.size());
+            for (const auto& subholder : Subholders) {
+                if (!subholder) {
                     continue;
                 }
-                if (auto cloned = holder->Clone(options)) {
-                    newHolder->Holders.push_back(cloned);
+                if (auto clonedSubholder = subholder->Clone(options)) {
+                    newHolder->Subholders.push_back(clonedSubholder);
                 }
             }
-
             return newHolder;
         }
+
+        std::optional<size_t> GetTotalByteSize() const override
+        {
+            size_t result = 0;
+            for (const auto& subholder : Subholders) {
+                if (!subholder) {
+                    continue;
+                }
+                auto subsize = subholder->GetTotalByteSize();
+                if (!subsize) {
+                    return std::nullopt;
+                }
+                result += *subsize;
+            }
+            return result;
+        }
     };
 
     auto holder = New<THolder>();
-    holder->Holders = std::move(holders);
+    holder->Subholders = std::move(holders);
     return holder;
 }
 
diff --git a/library/cpp/yt/memory/shared_range.h b/library/cpp/yt/memory/shared_range.h
index c67e40fa05..780a44ae8b 100644
--- a/library/cpp/yt/memory/shared_range.h
+++ b/library/cpp/yt/memory/shared_range.h
@@ -1,11 +1,14 @@
 #pragma once
 
+#include "public.h"
 #include "intrusive_ptr.h"
 #include "range.h"
 #include "ref_counted.h"
 
 #include <library/cpp/yt/assert/assert.h>
 
+#include <optional>
+
 namespace NYT {
 
 ////////////////////////////////////////////////////////////////////////////////
@@ -20,12 +23,21 @@ struct TSharedRangeHolderCloneOptions
     bool KeepMemoryReferenceTracking = true;
 };
 
-DECLARE_REFCOUNTED_STRUCT(TSharedRangeHolder)
-
 struct TSharedRangeHolder
     : public TRefCounted
 {
-    virtual TSharedRangeHolderPtr Clone(const TSharedRangeHolderCloneOptions& /*options*/);
+    //! Clones the holder possibly adjusting its flavor based on #options.
+    /*!
+     *  The default implementation just returns this.
+     */
+    virtual TSharedRangeHolderPtr Clone(const TSharedRangeHolderCloneOptions& options);
+
+    //! Returns the (estimated) total number of bytes being held or |null| if unable to estimate.
+    /*!
+     *  The returned value is static and never changes.
+     *  The default implementation returns |null|.
+     */
+    virtual std::optional<size_t> GetTotalByteSize() const;
 };
 
 DEFINE_REFCOUNTED_TYPE(TSharedRangeHolder)
diff --git a/library/cpp/yt/yson_string/string-inl.h b/library/cpp/yt/yson_string/string-inl.h
index 5c41629cc0..09db98881f 100644
--- a/library/cpp/yt/yson_string/string-inl.h
+++ b/library/cpp/yt/yson_string/string-inl.h
@@ -26,6 +26,35 @@ bool Equals(const TLeft& lhs, const TRight& rhs)
         lhs.GetType() == rhs.GetType();
 }
 
+class TYsonStringHolder
+    : public TSharedRangeHolder
+    , public TWithExtraSpace<TYsonStringHolder>
+{
+public:
+    explicit TYsonStringHolder(size_t size)
+        : Size_(size)
+    { }
+
+    char* GetData()
+    {
+        return static_cast<char*>(GetExtraSpacePtr());
+    }
+
+    // TSharedRangeHolder overrides.
+    std::optional<size_t> GetTotalByteSize() const override
+    {
+        return Size_ + sizeof(TYsonStringHolder);
+    }
+
+    static TIntrusivePtr<TYsonStringHolder> Allocate(size_t size)
+    {
+        return NewWithExtraSpace<TYsonStringHolder>(size, size);
+    }
+
+private:
+    const size_t Size_;
+};
+
 } // namespace NDetail
 
 inline bool operator == (const TYsonString& lhs, const TYsonString& rhs)
diff --git a/library/cpp/yt/yson_string/string.cpp b/library/cpp/yt/yson_string/string.cpp
index c92d4d6e42..dec428c31e 100644
--- a/library/cpp/yt/yson_string/string.cpp
+++ b/library/cpp/yt/yson_string/string.cpp
@@ -71,23 +71,13 @@ TYsonString::TYsonString()
 TYsonString::TYsonString(const TYsonStringBuf& ysonStringBuf)
 {
     if (ysonStringBuf) {
-        struct TCapturedYsonStringPayload
-            : public TSharedRangeHolder
-            , public TWithExtraSpace<TCapturedYsonStringPayload>
-        {
-            char* GetData()
-            {
-                return static_cast<char*>(GetExtraSpacePtr());
-            }
-        };
-
         auto data = ysonStringBuf.AsStringBuf();
-        auto payload = NewWithExtraSpace<TCapturedYsonStringPayload>(data.length());
-        ::memcpy(payload->GetData(), data.data(), data.length());
-        Payload_ = payload;
-        Begin_ = payload->GetData();
+        auto holder = NDetail::TYsonStringHolder::Allocate(data.length());
+        ::memcpy(holder->GetData(), data.data(), data.length());
+        Begin_ = holder->GetData();
         Size_ = data.Size();
         Type_ = ysonStringBuf.GetType();
+        Payload_ = std::move(holder);
     } else {
         Begin_ = nullptr;
         Size_ = 0;