YT-23989: Track memory in TChunkedOutputStream

commit_hash:1fcce66b4e0a0a5e5fb55aba38889e9bf5b42a85

1 files changed, 19 insertions, 4 deletions

diff --git a/library/cpp/yt/memory/chunked_output_stream.cpp b/library/cpp/yt/memory/chunked_output_stream.cpp
index 90ec36fe23f..d2654209477 100644
--- a/library/cpp/yt/memory/chunked_output_stream.cpp
+++ b/library/cpp/yt/memory/chunked_output_stream.cpp
@@ -8,9 +8,12 @@ namespace NYT {
 
 TChunkedOutputStream::TChunkedOutputStream(
     TRefCountedTypeCookie tagCookie,
+    IMemoryUsageTrackerPtr memoryUsageTracker,
     size_t initialReserveSize,
     size_t maxReserveSize)
-    : MaxReserveSize_(RoundUpToPage(maxReserveSize))
+    : MemoryUsageTracker_(std::move(memoryUsageTracker))
+    , CurrentChunkMemoryUsageGuard_(TMemoryUsageTrackerGuard::Build(MemoryUsageTracker_))
+    , MaxReserveSize_(RoundUpToPage(maxReserveSize))
     , CurrentReserveSize_(RoundUpToPage(initialReserveSize))
     , CurrentChunk_(tagCookie, /*size*/ 0)
 {
@@ -23,7 +26,8 @@ TChunkedOutputStream::TChunkedOutputStream(
 
 std::vector<TSharedRef> TChunkedOutputStream::Finish()
 {
-    FinishedChunks_.push_back(TSharedRef::FromBlob(std::move(CurrentChunk_)));
+    FinishedChunks_.push_back(TrackMemory(MemoryUsageTracker_, TSharedRef::FromBlob(std::move(CurrentChunk_))));
+    CurrentChunkMemoryUsageGuard_.Release();
 
     YT_ASSERT(CurrentChunk_.IsEmpty());
     FinishedSize_ = 0;
@@ -49,9 +53,10 @@ void TChunkedOutputStream::ReserveNewChunk(size_t spaceRequired)
 {
     YT_ASSERT(CurrentChunk_.Size() == CurrentChunk_.Capacity());
     FinishedSize_ += CurrentChunk_.Size();
-    FinishedChunks_.push_back(TSharedRef::FromBlob(std::move(CurrentChunk_)));
+    FinishedChunks_.push_back(TrackMemory(MemoryUsageTracker_, TSharedRef::FromBlob(std::move(CurrentChunk_))));
     CurrentReserveSize_ = std::min(2 * CurrentReserveSize_, MaxReserveSize_);
     CurrentChunk_.Reserve(std::max(RoundUpToPage(spaceRequired), CurrentReserveSize_));
+    UpdateCurrentChunkMemoryUsage();
 }
 
 void TChunkedOutputStream::DoWrite(const void* buffer, size_t length)
@@ -68,6 +73,7 @@ void TChunkedOutputStream::DoWrite(const void* buffer, size_t length)
         ReserveNewChunk(spaceRequired);
         CurrentChunk_.Append(static_cast<const char*>(buffer) + spaceAvailable, spaceRequired);
     }
+    UpdateCurrentChunkMemoryUsage();
 }
 
 size_t TChunkedOutputStream::DoNext(void** ptr)
@@ -84,6 +90,7 @@ size_t TChunkedOutputStream::DoNext(void** ptr)
     YT_ASSERT(spaceAvailable > 0);
     *ptr = CurrentChunk_.End();
     CurrentChunk_.Resize(CurrentChunk_.Capacity(), /*initializeStorage*/ false);
+    UpdateCurrentChunkMemoryUsage();
     return spaceAvailable;
 }
 
@@ -91,6 +98,7 @@ void TChunkedOutputStream::DoUndo(size_t len)
 {
     YT_VERIFY(CurrentChunk_.Size() >= len);
     CurrentChunk_.Resize(CurrentChunk_.Size() - len);
+    UpdateCurrentChunkMemoryUsage();
 }
 
 char* TChunkedOutputStream::Preallocate(size_t size)
@@ -98,12 +106,13 @@ char* TChunkedOutputStream::Preallocate(size_t size)
     size_t available = CurrentChunk_.Capacity() - CurrentChunk_.Size();
     if (available < size) {
         FinishedSize_ += CurrentChunk_.Size();
-        FinishedChunks_.push_back(TSharedRef::FromBlob(std::move(CurrentChunk_)));
+        FinishedChunks_.push_back(TrackMemory(MemoryUsageTracker_, TSharedRef::FromBlob(std::move(CurrentChunk_))));
 
         CurrentReserveSize_ = std::min(2 * CurrentReserveSize_, MaxReserveSize_);
 
         CurrentChunk_.Reserve(std::max(RoundUpToPage(size), CurrentReserveSize_));
     }
+    UpdateCurrentChunkMemoryUsage();
     return CurrentChunk_.End();
 }
 
@@ -111,6 +120,12 @@ void TChunkedOutputStream::Advance(size_t size)
 {
     YT_ASSERT(CurrentChunk_.Size() + size <= CurrentChunk_.Capacity());
     CurrentChunk_.Resize(CurrentChunk_.Size() + size, false);
+    UpdateCurrentChunkMemoryUsage();
+}
+
+void TChunkedOutputStream::UpdateCurrentChunkMemoryUsage()
+{
+    CurrentChunkMemoryUsageGuard_.SetSize(CurrentChunk_.Capacity());
 }
 
 ////////////////////////////////////////////////////////////////////////////////