YQL-19767: Track allocations.

commit_hash:67081edb88c3baf060c3179c3f21775fad73b530

5 files changed, 90 insertions, 13 deletions

diff --git a/yql/essentials/minikql/comp_nodes/mkql_unwrap.cpp b/yql/essentials/minikql/comp_nodes/mkql_unwrap.cpp
index 307609aa082..75d0f4ba581 100644
--- a/yql/essentials/minikql/comp_nodes/mkql_unwrap.cpp
+++ b/yql/essentials/minikql/comp_nodes/mkql_unwrap.cpp
@@ -20,6 +20,10 @@ public:
     }
 
     NUdf::TUnboxedValuePod DoCalculate(TComputationContext& compCtx) const {
+        return DoCalculateImpl(compCtx).Release();
+    }
+
+    NUdf::TUnboxedValue DoCalculateImpl(TComputationContext& compCtx) const {
         auto value = Optional()->GetValue(compCtx);
         if (value) {
             return value.GetOptionalValue();
diff --git a/yql/essentials/minikql/comp_nodes/mkql_weakmember.cpp b/yql/essentials/minikql/comp_nodes/mkql_weakmember.cpp
index 2d535459104..1c058b21a36 100644
--- a/yql/essentials/minikql/comp_nodes/mkql_weakmember.cpp
+++ b/yql/essentials/minikql/comp_nodes/mkql_weakmember.cpp
@@ -26,6 +26,11 @@ public:
     }
 
     NUdf::TUnboxedValuePod DoCalculate(TComputationContext& ctx) const {
+        auto result = DoCalculateImpl(ctx);
+        return result.Release();
+    }
+
+    NUdf::TUnboxedValue DoCalculateImpl(TComputationContext& ctx) const {
         if (const auto& restDict = RestDict->GetValue(ctx)) {
             if (const auto& tryMember = restDict.Lookup(MemberName)) {
                 return SimpleValueFromYson(SchemeType, tryMember.AsStringRef());
@@ -46,7 +51,7 @@ public:
                         stringStream.DoWrite(ref.Data(), size);
                         return stringStream.Value();
                     } else if (SchemeType == NUdf::EDataSlot::String) {
-                        return tryMember.Release();
+                        return tryMember;
                     } else {
                         return {};
                     }
diff --git a/yql/essentials/minikql/mkql_alloc.cpp b/yql/essentials/minikql/mkql_alloc.cpp
index 299e691f168..53e73d696ab 100644
--- a/yql/essentials/minikql/mkql_alloc.cpp
+++ b/yql/essentials/minikql/mkql_alloc.cpp
@@ -27,9 +27,15 @@ void TAllocState::TListEntry::Unlink() noexcept {
 
 TAllocState::TAllocState(const TSourceLocation& location, const NKikimr::TAlignedPagePoolCounters &counters, bool supportsSizedAllocators)
     : TAlignedPagePool(location, counters)
+#ifndef NDEBUG
+    , DefaultMemInfo(MakeIntrusive<TMemoryUsageInfo>("default"))
+#endif
     , SupportsSizedAllocators(supportsSizedAllocators)
     , CurrentPAllocList(&GlobalPAllocList)
 {
+#ifndef NDEBUG
+    ActiveMemInfo.emplace(DefaultMemInfo.Get(), DefaultMemInfo);
+#endif
     GetRoot()->InitLinks();
     OffloadedBlocksRoot.InitLinks();
     GlobalPAllocList.InitLinks();
diff --git a/yql/essentials/minikql/mkql_alloc.h b/yql/essentials/minikql/mkql_alloc.h
index 24bbbb8e9ed..323a8c1c20a 100644
--- a/yql/essentials/minikql/mkql_alloc.h
+++ b/yql/essentials/minikql/mkql_alloc.h
@@ -13,6 +13,7 @@
 #include <unordered_map>
 #include <atomic>
 #include <memory>
+#include <source_location>
 
 namespace NKikimr {
 
@@ -43,6 +44,25 @@ static_assert(sizeof(TAllocPageHeader) % MKQL_ALIGNMENT == 0, "Incorrect size of
 
 struct TMkqlArrowHeader;
 
+#ifndef NDEBUG
+using TAllocLocation = std::source_location;
+#else
+struct TAllocLocation
+{
+    const char* file_name() const {
+        return "";
+    }
+
+    std::uint_least32_t line() const {
+        return 0;
+    }
+
+    static TAllocLocation current() {
+        return {};
+    }
+};
+#endif
+
 struct TAllocState : public TAlignedPagePool
 {
     struct TListEntry {
@@ -57,6 +77,7 @@ struct TAllocState : public TAlignedPagePool
     };
 
 #ifndef NDEBUG
+    TIntrusivePtr<TMemoryUsageInfo> DefaultMemInfo;
     std::unordered_map<TMemoryUsageInfo*, TIntrusivePtr<TMemoryUsageInfo>> ActiveMemInfo;
 #endif
     bool SupportsSizedAllocators = false;
@@ -307,7 +328,10 @@ private:
 
 void* MKQLAllocSlow(size_t sz, TAllocState* state, const EMemorySubPool mPool);
 
-inline void* MKQLAllocFastDeprecated(size_t sz, TAllocState* state, const EMemorySubPool mPool) {
+inline void* MKQLAllocFastDeprecated(size_t sz, TAllocState* state, const EMemorySubPool mPool, const TAllocLocation& location = TAllocLocation::current()) {
+#ifdef NDEBUG
+    Y_UNUSED(location);
+#endif
     Y_DEBUG_ABORT_UNLESS(state);
 
 #if defined(ALLOW_DEFAULT_ALLOCATOR)
@@ -318,6 +342,9 @@ inline void* MKQLAllocFastDeprecated(size_t sz, TAllocState* state, const EMemor
         }
 
         ret->Link(&state->OffloadedBlocksRoot);
+#ifndef NDEBUG
+        state->DefaultMemInfo->Take(ret + 1, sz, { location.file_name(), (int)location.line() });
+#endif
         return ret + 1;
     }
 #endif
@@ -327,13 +354,23 @@ inline void* MKQLAllocFastDeprecated(size_t sz, TAllocState* state, const EMemor
         void* ret = (char*)currPage + currPage->Offset;
         currPage->Offset = AlignUp(currPage->Offset + sz, MKQL_ALIGNMENT);
         ++currPage->UseCount;
+#ifndef NDEBUG
+        state->DefaultMemInfo->Take(ret, sz, { location.file_name(), (int)location.line() });
+#endif
         return ret;
     }
 
-    return MKQLAllocSlow(sz, state, mPool);
+    auto ret = MKQLAllocSlow(sz, state, mPool);
+#ifndef NDEBUG
+    state->DefaultMemInfo->Take(ret, sz, { location.file_name(), (int)location.line() });
+#endif
+    return ret;
 }
 
-inline void* MKQLAllocFastWithSize(size_t sz, TAllocState* state, const EMemorySubPool mPool) {
+inline void* MKQLAllocFastWithSize(size_t sz, TAllocState* state, const EMemorySubPool mPool, const TAllocLocation& location = TAllocLocation::current()) {
+#ifdef NDEBUG
+    Y_UNUSED(location);
+#endif
     Y_DEBUG_ABORT_UNLESS(state);
 
     bool useMalloc = state->SupportsSizedAllocators && sz > MaxPageUserData;
@@ -349,6 +386,9 @@ inline void* MKQLAllocFastWithSize(size_t sz, TAllocState* state, const EMemoryS
         }
 
         ret->Link(&state->OffloadedBlocksRoot);
+#ifndef NDEBUG
+        state->DefaultMemInfo->Take(ret + 1, sz, { location.file_name(), (int)location.line() });
+#endif
         return ret + 1;
     }
 
@@ -357,10 +397,17 @@ inline void* MKQLAllocFastWithSize(size_t sz, TAllocState* state, const EMemoryS
         void* ret = (char*)currPage + currPage->Offset;
         currPage->Offset = AlignUp(currPage->Offset + sz, MKQL_ALIGNMENT);
         ++currPage->UseCount;
+#ifndef NDEBUG
+        state->DefaultMemInfo->Take(ret, sz, { location.file_name(), (int)location.line() });
+#endif
         return ret;
     }
 
-    return MKQLAllocSlow(sz, state, mPool);
+    auto ret = MKQLAllocSlow(sz, state, mPool);
+#ifndef NDEBUG
+    state->DefaultMemInfo->Take(ret, sz, { location.file_name(), (int)location.line() });
+#endif
+    return ret;
 }
 
 void MKQLFreeSlow(TAllocPageHeader* header, TAllocState *state, const EMemorySubPool mPool) noexcept;
@@ -370,6 +417,10 @@ inline void MKQLFreeDeprecated(const void* mem, const EMemorySubPool mPool) noex
         return;
     }
 
+#ifndef NDEBUG
+    TlsAllocState->DefaultMemInfo->Return(mem);
+#endif
+
 #if defined(ALLOW_DEFAULT_ALLOCATOR)
     if (Y_UNLIKELY(TAllocState::IsDefaultAllocatorUsed())) {
         TAllocState *state = TlsAllocState;
@@ -398,6 +449,9 @@ inline void MKQLFreeFastWithSize(const void* mem, size_t sz, TAllocState* state,
     }
 
     Y_DEBUG_ABORT_UNLESS(state);
+#ifndef NDEBUG
+    state->DefaultMemInfo->Return(mem, sz);
+#endif
 
     bool useFree = state->SupportsSizedAllocators && sz > MaxPageUserData;
 #if defined(ALLOW_DEFAULT_ALLOCATOR)
@@ -423,12 +477,12 @@ inline void MKQLFreeFastWithSize(const void* mem, size_t sz, TAllocState* state,
     MKQLFreeSlow(header, state, mPool);
 }
 
-inline void* MKQLAllocDeprecated(size_t sz, const EMemorySubPool mPool) {
-    return MKQLAllocFastDeprecated(sz, TlsAllocState, mPool);
+inline void* MKQLAllocDeprecated(size_t sz, const EMemorySubPool mPool, const TAllocLocation& location = TAllocLocation::current()) {
+    return MKQLAllocFastDeprecated(sz, TlsAllocState, mPool, location);
 }
 
-inline void* MKQLAllocWithSize(size_t sz, const EMemorySubPool mPool) {
-    return MKQLAllocFastWithSize(sz, TlsAllocState, mPool);
+inline void* MKQLAllocWithSize(size_t sz, const EMemorySubPool mPool, const TAllocLocation& location = TAllocLocation::current()) {
+    return MKQLAllocFastWithSize(sz, TlsAllocState, mPool, location);
 }
 
 inline void MKQLFreeWithSize(const void* mem, size_t sz, const EMemorySubPool mPool) noexcept {
@@ -478,6 +532,14 @@ struct TWithMiniKQLAlloc {
 };
 
 template <typename T, typename... Args>
+T* AllocateOn(const TAllocLocation& location, TAllocState* state, Args&&... args)
+{
+    void* addr = MKQLAllocFastWithSize(sizeof(T), state, T::MemoryPool, location);
+    return ::new(addr) T(std::forward<Args>(args)...);
+    static_assert(std::is_base_of<TWithMiniKQLAlloc<T::MemoryPool>, T>::value, "Class must inherit TWithMiniKQLAlloc.");
+}
+
+template <typename T, typename... Args>
 T* AllocateOn(TAllocState* state, Args&&... args)
 {
     void* addr = MKQLAllocFastWithSize(sizeof(T), state, T::MemoryPool);
diff --git a/yql/essentials/minikql/mkql_string_util_ut.cpp b/yql/essentials/minikql/mkql_string_util_ut.cpp
index 9826ee0ee10..f0d5545ab73 100644
--- a/yql/essentials/minikql/mkql_string_util_ut.cpp
+++ b/yql/essentials/minikql/mkql_string_util_ut.cpp
@@ -9,10 +9,10 @@ using namespace NKikimr::NMiniKQL;
 Y_UNIT_TEST_SUITE(TMiniKQLStringUtils) {
     Y_UNIT_TEST(SubstringWithLargeOffset) {
         TScopedAlloc alloc(__LOCATION__);
-        const auto big = MakeStringNotFilled(NUdf::TUnboxedValuePod::OffsetLimit << 1U);
-        const auto sub0 = SubString(big, 1U, 42U);
-        const auto sub1 = SubString(big, NUdf::TUnboxedValuePod::OffsetLimit - 1U, 42U);
-        const auto sub2 = SubString(big, NUdf::TUnboxedValuePod::OffsetLimit, 42U);
+        const auto big = MakeStringNotFilled(/*size=*/NUdf::TUnboxedValuePod::OffsetLimit << 1U);
+        const auto sub0 = NUdf::TUnboxedValue(SubString(big, 1U, 42U));
+        const auto sub1 = NUdf::TUnboxedValue(SubString(big, NUdf::TUnboxedValuePod::OffsetLimit - 1U, 42U));
+        const auto sub2 = NUdf::TUnboxedValue(SubString(big, NUdf::TUnboxedValuePod::OffsetLimit, 42U));
 
         UNIT_ASSERT(sub0.AsStringValue().Data() == sub1.AsStringValue().Data());
         UNIT_ASSERT(sub1.AsStringValue().Data() != sub2.AsStringValue().Data());