Restore measuring the memory usage of write-table queries.

commit_hash:1d28ad476551aee1d5ab64c72205347128c27c68

1 files changed, 86 insertions, 0 deletions

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