Get rid of NYTAlloc::GetAllocationSize usages

5 files changed, 29 insertions, 19 deletions

diff --git a/library/cpp/yt/malloc/malloc.cpp b/library/cpp/yt/malloc/malloc.cpp
index 808afacdfb..4a35758cab 100644
--- a/library/cpp/yt/malloc/malloc.cpp
+++ b/library/cpp/yt/malloc/malloc.cpp
@@ -4,16 +4,14 @@
 
 ////////////////////////////////////////////////////////////////////////////////
 
-Y_WEAK extern "C" size_t nallocx(size_t size, int /* flags */) noexcept
+Y_WEAK extern "C" size_t nallocx(size_t size, int /*flags*/) noexcept
 {
     return size;
 }
 
-#ifndef _win_
-Y_WEAK extern "C" size_t malloc_usable_size(void* /* ptr */) noexcept
+Y_WEAK extern "C" size_t malloc_usable_size(void* /*ptr*/) noexcept
 {
     return 0;
 }
-#endif
 
 ////////////////////////////////////////////////////////////////////////////////
diff --git a/library/cpp/yt/memory/new-inl.h b/library/cpp/yt/memory/new-inl.h
index 0a84818516..5d5fcc45c4 100644
--- a/library/cpp/yt/memory/new-inl.h
+++ b/library/cpp/yt/memory/new-inl.h
@@ -6,6 +6,8 @@
 
 #include <library/cpp/ytalloc/api/ytalloc.h>
 
+#include <library/cpp/yt/malloc//malloc.h>
+
 namespace NYT {
 
 ////////////////////////////////////////////////////////////////////////////////
@@ -305,6 +307,16 @@ void* TWithExtraSpace<T>::GetExtraSpacePtr()
     return static_cast<T*>(this) + 1;
 }
 
+template <class T>
+size_t TWithExtraSpace<T>::GetUsableSpaceSize() const
+{
+#ifdef _win_
+    return 0;
+#else
+    return malloc_usable_size(const_cast<T*>(static_cast<const T*>(this))) - sizeof(T);
+#endif
+}
+
 ////////////////////////////////////////////////////////////////////////////////
 
 } // namespace NYT
diff --git a/library/cpp/yt/memory/new.h b/library/cpp/yt/memory/new.h
index 2db45e0465..ea96fd60e0 100644
--- a/library/cpp/yt/memory/new.h
+++ b/library/cpp/yt/memory/new.h
@@ -116,6 +116,7 @@ class TWithExtraSpace
 protected:
     const void* GetExtraSpacePtr() const;
     void* GetExtraSpacePtr();
+    size_t GetUsableSpaceSize() const;
 };
 
 ////////////////////////////////////////////////////////////////////////////////
diff --git a/library/cpp/yt/memory/ref.cpp b/library/cpp/yt/memory/ref.cpp
index 1a6b1830f7..257a8c9050 100644
--- a/library/cpp/yt/memory/ref.cpp
+++ b/library/cpp/yt/memory/ref.cpp
@@ -77,13 +77,6 @@ class TAllocationHolderBase
     : public TRefCounted
 {
 public:
-    TAllocationHolderBase(size_t size, TRefCountedTypeCookie cookie)
-        : Size_(size)
-#ifdef YT_ENABLE_REF_COUNTED_TRACKING
-        , Cookie_(cookie)
-#endif
-    { }
-
     ~TAllocationHolderBase()
     {
 #ifdef YT_ENABLE_REF_COUNTED_TRACKING
@@ -98,13 +91,15 @@ public:
     }
 
 protected:
-    const size_t Size_;
+    size_t Size_;
 #ifdef YT_ENABLE_REF_COUNTED_TRACKING
-    const TRefCountedTypeCookie Cookie_;
+    TRefCountedTypeCookie Cookie_;
 #endif
 
-    void Initialize(TSharedMutableRefAllocateOptions options)
+    void Initialize(size_t size, TSharedMutableRefAllocateOptions options, TRefCountedTypeCookie cookie)
     {
+        Size_ = size;
+        Cookie_ = cookie;
         if (options.InitializeStorage) {
             ::memset(static_cast<TDerived*>(this)->GetBegin(), 0, Size_);
         }
@@ -123,9 +118,13 @@ class TDefaultAllocationHolder
 {
 public:
     TDefaultAllocationHolder(size_t size, TSharedMutableRefAllocateOptions options, TRefCountedTypeCookie cookie)
-        : TAllocationHolderBase(size, cookie)
     {
-        Initialize(options);
+        if (options.ExtendToUsableSize) {
+            if (auto usableSize = GetUsableSpaceSize(); usableSize != 0) {
+                size = usableSize;
+            }
+        }
+        Initialize(size, options, cookie);
     }
 
     char* GetBegin()
@@ -141,10 +140,9 @@ class TPageAlignedAllocationHolder
 {
 public:
     TPageAlignedAllocationHolder(size_t size, TSharedMutableRefAllocateOptions options, TRefCountedTypeCookie cookie)
-        : TAllocationHolderBase(size, cookie)
-        , Begin_(static_cast<char*>(NYTAlloc::AllocatePageAligned(size)))
+        : Begin_(static_cast<char*>(NYTAlloc::AllocatePageAligned(size)))
     {
-        Initialize(options);
+        Initialize(size, options, cookie);
     }
 
     ~TPageAlignedAllocationHolder()
diff --git a/library/cpp/yt/memory/ref.h b/library/cpp/yt/memory/ref.h
index 2003ba0c0d..cc7942176c 100644
--- a/library/cpp/yt/memory/ref.h
+++ b/library/cpp/yt/memory/ref.h
@@ -170,6 +170,7 @@ private:
 struct TSharedMutableRefAllocateOptions
 {
     bool InitializeStorage = true;
+    bool ExtendToUsableSize = false;
 };
 
 //! A reference to a mutable range of memory with shared ownership.