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 +++++++++ yt/yt/core/misc/memory_usage_tracker.cpp | 11 --- yt/yt/core/misc/memory_usage_tracker.h | 5 -- .../column_converters/string_column_converter.cpp | 2 + yt/yt/library/formats/arrow_parser.cpp | 1 + 8 files changed, 155 insertions(+), 20 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 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 diff --git a/yt/yt/core/misc/memory_usage_tracker.cpp b/yt/yt/core/misc/memory_usage_tracker.cpp index 58b18307b90..b8808209393 100644 --- a/yt/yt/core/misc/memory_usage_tracker.cpp +++ b/yt/yt/core/misc/memory_usage_tracker.cpp @@ -373,17 +373,6 @@ TErrorOr TryTrackMemory( return tracker->TryTrack(std::move(reference), keepExistingTracking); } -TSharedRef TrackMemory( - const IMemoryUsageTrackerPtr& tracker, - TSharedRef reference, - bool keepExistingTracking) -{ - if (!tracker || !reference) { - return reference; - } - return tracker->Track(std::move(reference), keepExistingTracking); -} - TSharedRefArray TrackMemory( const IMemoryUsageTrackerPtr& tracker, TSharedRefArray array, diff --git a/yt/yt/core/misc/memory_usage_tracker.h b/yt/yt/core/misc/memory_usage_tracker.h index ece48dfbb54..0409d5db4b3 100644 --- a/yt/yt/core/misc/memory_usage_tracker.h +++ b/yt/yt/core/misc/memory_usage_tracker.h @@ -166,11 +166,6 @@ TErrorOr TryTrackMemory( 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/library/column_converters/string_column_converter.cpp b/yt/yt/library/column_converters/string_column_converter.cpp index 04b3854c15e..a76d9c9fd22 100644 --- a/yt/yt/library/column_converters/string_column_converter.cpp +++ b/yt/yt/library/column_converters/string_column_converter.cpp @@ -6,6 +6,7 @@ #include #include +#include #include @@ -91,6 +92,7 @@ private: DictionaryByteSize_ = 0; DirectBuffer_ = std::make_unique( GetRefCountedTypeCookie(), + GetNullMemoryUsageTracker(), // TODO(nadya02): YT-25327. 256_KB, 1_MB); diff --git a/yt/yt/library/formats/arrow_parser.cpp b/yt/yt/library/formats/arrow_parser.cpp index eb0bfb10914..5b554ef9f5c 100644 --- a/yt/yt/library/formats/arrow_parser.cpp +++ b/yt/yt/library/formats/arrow_parser.cpp @@ -1546,6 +1546,7 @@ public: auto bufferForStringLikeValues = std::make_shared( GetRefCountedTypeCookie(), + GetNullMemoryUsageTracker(), // TODO(nadya02): YT-25327. 256_KB, 1_MB); -- cgit v1.3