From b82002b0fabe461ff9aee40ca37a129c5e7cb263 Mon Sep 17 00:00:00 2001 From: nadya02 Date: Wed, 11 Jun 2025 17:56:39 +0300 Subject: Restore measuring the memory usage of write-table queries. commit_hash:1d28ad476551aee1d5ab64c72205347128c27c68 --- library/cpp/yt/memory/chunked_output_stream.cpp | 31 +++++++-- library/cpp/yt/memory/chunked_output_stream.h | 6 ++ library/cpp/yt/memory/memory_usage_tracker.cpp | 86 +++++++++++++++++++++++++ library/cpp/yt/memory/memory_usage_tracker.h | 33 ++++++++++ 4 files changed, 152 insertions(+), 4 deletions(-) (limited to 'library/cpp') 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 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(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(), + 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(); } +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 -- cgit v1.3