4 files changed, 152 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..bf62412e00d 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,
+    ISimpleMemoryUsageTrackerPtr memoryUsageTracker,
     size_t initialReserveSize,
     size_t maxReserveSize)
-    : MaxReserveSize_(RoundUpToPage(maxReserveSize))
+    : MemoryUsageTracker_(std::move(memoryUsageTracker))
+    , CurrentChunkMemoryUsageGuard_(TSimpleMemoryUsageTrackerGuard::Build(MemoryUsageTracker_))
+    , MaxReserveSize_(RoundUpToPage(maxReserveSize))
     , CurrentReserveSize_(RoundUpToPage(initialReserveSize))
     , CurrentChunk_(tagCookie, /*size*/ 0)
 {
@@ -23,7 +26,11 @@ 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 +56,13 @@ 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 +79,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 +96,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 +104,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 +112,15 @@ 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 +128,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());
 }
 
 ////////////////////////////////////////////////////////////////////////////////
diff --git a/library/cpp/yt/memory/chunked_output_stream.h b/library/cpp/yt/memory/chunked_output_stream.h
index ddb8235bd84..e6043976f75 100644
--- a/library/cpp/yt/memory/chunked_output_stream.h
+++ b/library/cpp/yt/memory/chunked_output_stream.h
@@ -22,6 +22,7 @@ class TChunkedOutputStream
 public:
     explicit TChunkedOutputStream(
         TRefCountedTypeCookie tagCookie = GetRefCountedTypeCookie<TDefaultChunkedOutputStreamTag>(),
+        ISimpleMemoryUsageTrackerPtr memoryUsageTracker = GetNullSimpleMemoryUsageTracker(),
         size_t initialReserveSize = 4_KB,
         size_t maxReserveSize = 64_KB);
 
@@ -47,6 +48,9 @@ public:
     void Advance(size_t size);
 
 private:
+    const ISimpleMemoryUsageTrackerPtr MemoryUsageTracker_;
+
+    TSimpleMemoryUsageTrackerGuard CurrentChunkMemoryUsageGuard_;
     size_t MaxReserveSize_;
     size_t CurrentReserveSize_;
 
@@ -57,6 +61,8 @@ private:
 
     void ReserveNewChunk(size_t spaceRequired);
 
+    void UpdateCurrentChunkMemoryUsage();
+
     void DoWrite(const void* buf, size_t len) override;
     size_t DoNext(void** ptr) override;
     void DoUndo(size_t len) override;
diff --git a/library/cpp/yt/memory/memory_usage_tracker.cpp b/library/cpp/yt/memory/memory_usage_tracker.cpp
index 6f0c90a9d1a..4055bf62c5e 100644
--- a/library/cpp/yt/memory/memory_usage_tracker.cpp
+++ b/library/cpp/yt/memory/memory_usage_tracker.cpp
@@ -17,15 +17,101 @@ public:
 
     void Release(i64 /*size*/) override
     { }
+
+    TSharedRef Track(TSharedRef reference, bool /*keepExistingTracking*/) override
+    {
+        return reference;
+    }
 };
 
 ////////////////////////////////////////////////////////////////////////////////
 
+TSimpleMemoryUsageTrackerGuard::TSimpleMemoryUsageTrackerGuard(TSimpleMemoryUsageTrackerGuard&& other)
+{
+    MoveFrom(std::move(other));
+}
+
+TSimpleMemoryUsageTrackerGuard::~TSimpleMemoryUsageTrackerGuard()
+{
+    Release();
+}
+
+TSimpleMemoryUsageTrackerGuard& TSimpleMemoryUsageTrackerGuard::operator=(TSimpleMemoryUsageTrackerGuard&& other)
+{
+    if (this != &other) {
+        Release();
+        MoveFrom(std::move(other));
+    }
+    return *this;
+}
+
+void TSimpleMemoryUsageTrackerGuard::MoveFrom(TSimpleMemoryUsageTrackerGuard&& other)
+{
+    Tracker_ = other.Tracker_;
+    AcquiredSize_ = other.AcquiredSize_;
+
+    other.Tracker_ = nullptr;
+    other.AcquiredSize_ = 0;
+}
+
+TSimpleMemoryUsageTrackerGuard TSimpleMemoryUsageTrackerGuard::Build(ISimpleMemoryUsageTrackerPtr tracker)
+{
+    if (!tracker) {
+        return {};
+    }
+
+    TSimpleMemoryUsageTrackerGuard guard;
+    guard.Tracker_ = tracker;
+    return guard;
+}
+
+void TSimpleMemoryUsageTrackerGuard::Release()
+{
+    if (Tracker_) {
+        if (AcquiredSize_) {
+            Tracker_->Release(AcquiredSize_);
+        }
+
+        Tracker_.Reset();
+        AcquiredSize_ = 0;
+    }
+}
+
+void TSimpleMemoryUsageTrackerGuard::SetSize(i64 size)
+{
+    if (!Tracker_) {
+        return;
+    }
+
+    YT_VERIFY(size >= 0);
+
+    if (size > AcquiredSize_) {
+        Tracker_->Acquire(size - AcquiredSize_);
+    } else if (size < AcquiredSize_) {
+        Tracker_->Release(AcquiredSize_ - size);
+    }
+
+    AcquiredSize_ = size;
+}
+
+////////////////////////////////////////////////////////////////////////////////
+
 ISimpleMemoryUsageTrackerPtr GetNullSimpleMemoryUsageTracker()
 {
     return LeakyRefCountedSingleton<TNulSimplelMemoryUsageTracker>();
 }
 
+TSharedRef TrackMemory(
+    const ISimpleMemoryUsageTrackerPtr& tracker,
+    TSharedRef reference,
+    bool keepExistingTracking)
+{
+    if (!tracker || !reference) {
+        return reference;
+    }
+    return tracker->Track(std::move(reference), keepExistingTracking);
+}
+
 ////////////////////////////////////////////////////////////////////////////////
 
 } // namespace NYT
diff --git a/library/cpp/yt/memory/memory_usage_tracker.h b/library/cpp/yt/memory/memory_usage_tracker.h
index 58b4f78ee6b..92849a06ef1 100644
--- a/library/cpp/yt/memory/memory_usage_tracker.h
+++ b/library/cpp/yt/memory/memory_usage_tracker.h
@@ -11,14 +11,47 @@ struct ISimpleMemoryUsageTracker
 {
     virtual bool Acquire(i64 size) = 0;
     virtual void Release(i64 size) = 0;
+    virtual TSharedRef Track(TSharedRef reference, bool keepExistingTracking) = 0;
 };
 
 DEFINE_REFCOUNTED_TYPE(ISimpleMemoryUsageTracker)
 
 ////////////////////////////////////////////////////////////////////////////////
 
+class TSimpleMemoryUsageTrackerGuard
+    : private TNonCopyable
+{
+public:
+    TSimpleMemoryUsageTrackerGuard() = default;
+    TSimpleMemoryUsageTrackerGuard(const TSimpleMemoryUsageTrackerGuard& other) = delete;
+    TSimpleMemoryUsageTrackerGuard(TSimpleMemoryUsageTrackerGuard&& other);
+    ~TSimpleMemoryUsageTrackerGuard();
+
+    TSimpleMemoryUsageTrackerGuard& operator=(const TSimpleMemoryUsageTrackerGuard& other) = delete;
+    TSimpleMemoryUsageTrackerGuard& operator=(TSimpleMemoryUsageTrackerGuard&& other);
+
+    static TSimpleMemoryUsageTrackerGuard Build(ISimpleMemoryUsageTrackerPtr tracker);
+
+    void SetSize(i64 size);
+
+    void Release();
+
+private:
+    ISimpleMemoryUsageTrackerPtr Tracker_;
+    i64 AcquiredSize_ = 0;
+    void MoveFrom(TSimpleMemoryUsageTrackerGuard&& other);
+
+};
+
+////////////////////////////////////////////////////////////////////////////////
+
 ISimpleMemoryUsageTrackerPtr GetNullSimpleMemoryUsageTracker();
 
+TSharedRef TrackMemory(
+    const ISimpleMemoryUsageTrackerPtr& tracker,
+    TSharedRef reference,
+    bool keepExistingTracking = false);
+
 ////////////////////////////////////////////////////////////////////////////////
 
 } // namespace NYT