YT-23859: Various memory tracking improvements

commit_hash:c87fb9e7c5cb8a896dd4b758ad73c5c540339df7

6 files changed, 29 insertions, 6 deletions

diff --git a/library/cpp/yt/memory/chunked_memory_allocator.h b/library/cpp/yt/memory/chunked_memory_allocator.h
index 4804a969cf7..d5e56c9f706 100644
--- a/library/cpp/yt/memory/chunked_memory_allocator.h
+++ b/library/cpp/yt/memory/chunked_memory_allocator.h
@@ -31,7 +31,9 @@ public:
             chunkSize,
             maxSmallBlockSizeRatio,
             GetRefCountedTypeCookie<TTag>())
-    { }
+    {
+        static_assert(sizeof(TTag) <= 1);
+    }
 
     //! Allocates #sizes bytes without any alignment.
     TSharedMutableRef AllocateUnaligned(i64 size);
diff --git a/library/cpp/yt/memory/chunked_memory_pool-inl.h b/library/cpp/yt/memory/chunked_memory_pool-inl.h
index 00733235f49..0faad070e9c 100644
--- a/library/cpp/yt/memory/chunked_memory_pool-inl.h
+++ b/library/cpp/yt/memory/chunked_memory_pool-inl.h
@@ -71,7 +71,9 @@ inline TChunkedMemoryPool::TChunkedMemoryPool(
     : TChunkedMemoryPool(
         GetRefCountedTypeCookie<TTag>(),
         startChunkSize)
-{ }
+{
+    static_assert(sizeof(TTag) <= 1);
+}
 
 inline char* TChunkedMemoryPool::AllocateUnaligned(size_t size)
 {
diff --git a/library/cpp/yt/memory/ref-inl.h b/library/cpp/yt/memory/ref-inl.h
index a2074f97311..13541a285ab 100644
--- a/library/cpp/yt/memory/ref-inl.h
+++ b/library/cpp/yt/memory/ref-inl.h
@@ -134,6 +134,7 @@ Y_FORCE_INLINE TSharedRef::operator TRef() const
 template <class TTag>
 Y_FORCE_INLINE TSharedRef TSharedRef::FromString(TString str)
 {
+    static_assert(sizeof(TTag) <= 1);
     return FromString(std::move(str), GetRefCountedTypeCookie<TTag>());
 }
 
@@ -145,6 +146,7 @@ Y_FORCE_INLINE TSharedRef TSharedRef::FromString(TString str)
 template <class TTag>
 Y_FORCE_INLINE TSharedRef TSharedRef::FromString(std::string str)
 {
+    static_assert(sizeof(TTag) <= 1);
     return FromString(std::move(str), GetRefCountedTypeCookie<TTag>());
 }
 
@@ -161,6 +163,7 @@ Y_FORCE_INLINE TStringBuf TSharedRef::ToStringBuf() const
 template <class TTag>
 Y_FORCE_INLINE TSharedRef TSharedRef::MakeCopy(TRef ref)
 {
+    static_assert(sizeof(TTag) <= 1);
     return MakeCopy(ref, GetRefCountedTypeCookie<TTag>());
 }
 
@@ -224,6 +227,7 @@ Y_FORCE_INLINE TSharedMutableRef TSharedMutableRef::AllocatePageAligned(size_t s
 template <class TTag>
 Y_FORCE_INLINE TSharedMutableRef TSharedMutableRef::MakeCopy(TRef ref)
 {
+    static_assert(sizeof(TTag) <= 1);
     return MakeCopy(ref, GetRefCountedTypeCookie<TTag>());
 }
 
@@ -243,12 +247,14 @@ Y_FORCE_INLINE TSharedMutableRef TSharedMutableRef::Slice(void* begin, void* end
 template <class TTag>
 Y_FORCE_INLINE TSharedMutableRef TSharedMutableRef::Allocate(size_t size, TSharedMutableRefAllocateOptions options)
 {
+    static_assert(sizeof(TTag) <= 1);
     return Allocate(size, options, GetRefCountedTypeCookie<TTag>());
 }
 
 template <class TTag>
 Y_FORCE_INLINE TSharedMutableRef TSharedMutableRef::AllocatePageAligned(size_t size, TSharedMutableRefAllocateOptions options)
 {
+    static_assert(sizeof(TTag) <= 1);
     return AllocatePageAligned(size, options, GetRefCountedTypeCookie<TTag>());
 }
 
diff --git a/yt/yt/client/table_client/row_buffer.h b/yt/yt/client/table_client/row_buffer.h
index 05ab9635c4f..cd8b5205af9 100644
--- a/yt/yt/client/table_client/row_buffer.h
+++ b/yt/yt/client/table_client/row_buffer.h
@@ -38,7 +38,9 @@ public:
         , Pool_(
             TTag(),
             startChunkSize)
-    { }
+    {
+        static_assert(sizeof(TTag) <= 1);
+    }
 
     template <class TTag>
     TRowBuffer(
@@ -49,7 +51,9 @@ public:
         , Pool_(
             GetRefCountedTypeCookie<TTag>(),
             std::move(chunkProvider))
-    { }
+    {
+        static_assert(sizeof(TTag) <= 1);
+    }
 
     TChunkedMemoryPool* GetPool();
 
diff --git a/yt/yt/client/table_client/wire_protocol.cpp b/yt/yt/client/table_client/wire_protocol.cpp
index bc7aceee619..b80de65602c 100644
--- a/yt/yt/client/table_client/wire_protocol.cpp
+++ b/yt/yt/client/table_client/wire_protocol.cpp
@@ -1071,12 +1071,14 @@ public:
         NCompression::ECodec codecId,
         TTableSchemaPtr schema,
         bool schemaful,
+        IMemoryChunkProviderPtr memoryChunkProvider,
         const NLogging::TLogger& logger,
         TWireProtocolOptions options)
         : CompressedBlocks_(compressedBlocks)
         , Codec_(NCompression::GetCodec(codecId))
         , Schema_(std::move(schema))
         , Schemaful_(schemaful)
+        , MemoryChunkProvider_(std::move(memoryChunkProvider))
         , Logger(logger.WithTag("ReaderId: %v", TGuid::Create()))
         , Options_(std::move(options))
     {
@@ -1106,7 +1108,10 @@ public:
             BlockIndex_,
             uncompressedBlock.Size());
 
-        auto rowBuffer = New<TRowBuffer>(TWireProtocolReaderTag(), ReaderBufferChunkSize);
+        auto rowBuffer = New<TRowBuffer>(
+            GetRefCountedTypeCookie<TWireProtocolReaderTag>(),
+            MemoryChunkProvider_,
+            ReaderBufferChunkSize);
         WireReader_ = CreateWireProtocolReader(uncompressedBlock, std::move(rowBuffer), Options_);
 
         if (!SchemaChecked_) {
@@ -1167,7 +1172,8 @@ private:
     const std::vector<TSharedRef> CompressedBlocks_;
     NCompression::ICodec* const Codec_;
     const TTableSchemaPtr Schema_;
-    bool Schemaful_;
+    const bool Schemaful_;
+    const IMemoryChunkProviderPtr MemoryChunkProvider_;
     const NLogging::TLogger Logger;
     const TWireProtocolOptions Options_;
 
@@ -1182,6 +1188,7 @@ IWireProtocolRowsetReaderPtr CreateWireProtocolRowsetReader(
     NCompression::ECodec codecId,
     TTableSchemaPtr schema,
     bool schemaful,
+    IMemoryChunkProviderPtr memoryChunkProvider,
     const NLogging::TLogger& logger,
     TWireProtocolOptions options)
 {
@@ -1190,6 +1197,7 @@ IWireProtocolRowsetReaderPtr CreateWireProtocolRowsetReader(
         codecId,
         std::move(schema),
         schemaful,
+        std::move(memoryChunkProvider),
         logger,
         std::move(options));
 }
diff --git a/yt/yt/client/table_client/wire_protocol.h b/yt/yt/client/table_client/wire_protocol.h
index 493d03b494a..eec6bb7ed82 100644
--- a/yt/yt/client/table_client/wire_protocol.h
+++ b/yt/yt/client/table_client/wire_protocol.h
@@ -313,6 +313,7 @@ IWireProtocolRowsetReaderPtr CreateWireProtocolRowsetReader(
     NCompression::ECodec codecId,
     NTableClient::TTableSchemaPtr schema,
     bool schemaful,
+    IMemoryChunkProviderPtr memoryChunkProvider,
     const NLogging::TLogger& logger,
     TWireProtocolOptions options = {});