[yt/cpp/mapreduce] YT-19268: Lock memory for parallel writer

2 files changed, 27 insertions, 0 deletions

diff --git a/yt/cpp/mapreduce/interface/io-inl.h b/yt/cpp/mapreduce/interface/io-inl.h
index c35ebb74811..542749c679b 100644
--- a/yt/cpp/mapreduce/interface/io-inl.h
+++ b/yt/cpp/mapreduce/interface/io-inl.h
@@ -740,6 +740,11 @@ struct IWriterImplBase
         }
     }
 
+    virtual size_t GetBufferMemoryUsage() const
+    {
+        return 0;
+    }
+
     virtual size_t GetTableCount() const = 0;
     virtual void FinishTable(size_t tableIndex) = 0;
     virtual void Abort()
@@ -816,6 +821,11 @@ public:
         }
     }
 
+    size_t GetBufferMemoryUsage() const
+    {
+        return DoGetBufferMemoryUsage();
+    }
+
 protected:
     template <class U>
     void DoAddRow(const U& row, size_t tableIndex = 0, size_t rowWeight = 0)
@@ -869,6 +879,11 @@ protected:
         Writer_->AddRowBatch(std::move(rowBatch), tableIndex, rowBatchWeight);
     }
 
+    size_t DoGetBufferMemoryUsage() const
+    {
+        return Writer_->GetBufferMemoryUsage();
+    }
+
     ::TIntrusivePtr<IWriterImpl> GetWriterImpl()
     {
         return Writer_;
diff --git a/yt/cpp/mapreduce/interface/io.h b/yt/cpp/mapreduce/interface/io.h
index e2b20a18029..dbd3a2cb707 100644
--- a/yt/cpp/mapreduce/interface/io.h
+++ b/yt/cpp/mapreduce/interface/io.h
@@ -168,6 +168,11 @@ public:
     /// By default it does nothing, but implementations are welcome to override this method.
     virtual void Abort()
     { }
+
+    virtual size_t GetBufferMemoryUsage() const
+    {
+        return 0;
+    }
 };
 
 /// @brief Interface to deal with multiple raw output streams.
@@ -191,6 +196,11 @@ public:
     /// By default it does nothing, but implementations are welcome to override this method.
     virtual void Abort()
     { }
+
+    virtual size_t GetBufferMemoryUsage() const
+    {
+        return 0;
+    }
 };
 
 ////////////////////////////////////////////////////////////////////////////////
@@ -378,6 +388,8 @@ public:
 
     /// Stop writing data as soon as possible (without flushing data, e.g. before aborting parent transaction).
     void Finish();
+
+    size_t GetBufferMemoryUsage() const;
 };
 
 ////////////////////////////////////////////////////////////////////////////////