YT-22072: Add try track function in memory tracker

94b357ecab85c316d9bd59bc20d7dcac779cd5db

4 files changed, 37 insertions, 0 deletions

diff --git a/yt/yt/core/misc/memory_usage_tracker.cpp b/yt/yt/core/misc/memory_usage_tracker.cpp
index cd1176ee06..0876e010d8 100644
--- a/yt/yt/core/misc/memory_usage_tracker.cpp
+++ b/yt/yt/core/misc/memory_usage_tracker.cpp
@@ -56,6 +56,13 @@ public:
     {
         return reference;
     }
+
+    TErrorOr<TSharedRef> TryTrack(
+        TSharedRef reference,
+        bool /*keepExistingTracking*/) override
+    {
+        return reference;
+    }
 };
 
 ////////////////////////////////////////////////////////////////////////////////
@@ -354,6 +361,17 @@ void TMemoryTrackedBlob::Clear()
 
 ////////////////////////////////////////////////////////////////////////////////
 
+TErrorOr<TSharedRef> TryTrackMemory(
+    const IMemoryUsageTrackerPtr& tracker,
+    TSharedRef reference,
+    bool keepExistingTracking)
+{
+    if (!tracker || !reference) {
+        return reference;
+    }
+    return tracker->TryTrack(reference, keepExistingTracking);
+}
+
 TSharedRef TrackMemory(
     const IMemoryUsageTrackerPtr& tracker,
     TSharedRef reference,
diff --git a/yt/yt/core/misc/memory_usage_tracker.h b/yt/yt/core/misc/memory_usage_tracker.h
index cacc9517c8..54b078697a 100644
--- a/yt/yt/core/misc/memory_usage_tracker.h
+++ b/yt/yt/core/misc/memory_usage_tracker.h
@@ -34,6 +34,9 @@ struct IMemoryUsageTracker
     virtual TSharedRef Track(
         TSharedRef reference,
         bool keepHolder = false) = 0;
+    virtual TErrorOr<TSharedRef> TryTrack(
+        TSharedRef reference,
+        bool keepHolder) = 0;
 };
 
 DEFINE_REFCOUNTED_TYPE(IMemoryUsageTracker)
@@ -145,10 +148,16 @@ private:
 
 ////////////////////////////////////////////////////////////////////////////////
 
+TErrorOr<TSharedRef> TryTrackMemory(
+    const IMemoryUsageTrackerPtr& tracker,
+    TSharedRef reference,
+    bool keepExistingTracking = false);
+
 TSharedRef TrackMemory(
     const IMemoryUsageTrackerPtr& tracker,
     TSharedRef reference,
     bool keepExistingTracking = false);
+
 TSharedRefArray TrackMemory(
     const IMemoryUsageTrackerPtr& tracker,
     TSharedRefArray array,
diff --git a/yt/yt/core/test_framework/test_memory_tracker.cpp b/yt/yt/core/test_framework/test_memory_tracker.cpp
index fd8b0de002..c491fd3e1c 100644
--- a/yt/yt/core/test_framework/test_memory_tracker.cpp
+++ b/yt/yt/core/test_framework/test_memory_tracker.cpp
@@ -127,6 +127,13 @@ TSharedRef TTestNodeMemoryTracker::Track(TSharedRef reference, bool keepExisting
         New<TTestTrackedReferenceHolder>(std::move(underlyingReference), std::move(guard)));
 }
 
+TErrorOr<TSharedRef> TTestNodeMemoryTracker::TryTrack(
+    TSharedRef reference,
+    bool keepHolder)
+{
+    return Track(reference, keepHolder);
+}
+
 ////////////////////////////////////////////////////////////////////////////////
 
 TTestNodeMemoryTracker::TTestTrackedReferenceHolder::TTestTrackedReferenceHolder(
diff --git a/yt/yt/core/test_framework/test_memory_tracker.h b/yt/yt/core/test_framework/test_memory_tracker.h
index 84a8d428bb..7308eb2a9d 100644
--- a/yt/yt/core/test_framework/test_memory_tracker.h
+++ b/yt/yt/core/test_framework/test_memory_tracker.h
@@ -34,6 +34,9 @@ public:
     TSharedRef Track(
         TSharedRef reference,
         bool keepHolder = false) override;
+    TErrorOr<TSharedRef> TryTrack(
+        TSharedRef reference,
+        bool keepHolder = false) override;
 private:
 
     class TTestTrackedReferenceHolder