YQL-16867 keep size of allocated block

6 files changed, 63 insertions, 46 deletions

diff --git a/ydb/library/yql/minikql/mkql_alloc.cpp b/ydb/library/yql/minikql/mkql_alloc.cpp
index 57dd6a1af0..11952678cc 100644
--- a/ydb/library/yql/minikql/mkql_alloc.cpp
+++ b/ydb/library/yql/minikql/mkql_alloc.cpp
@@ -36,7 +36,8 @@ TAllocState::TAllocState(const TSourceLocation& location, const NKikimr::TAligne
 void TAllocState::CleanupPAllocList(TListEntry* root) {
     for (auto curr = root->Right; curr != root; ) {
         auto next = curr->Right;
-        MKQLFreeDeprecated(curr); // may free items from OffloadedBlocksRoot
+        auto size = ((TMkqlPAllocHeader*)curr)->Size;
+        MKQLFreeWithSize(curr, size, EMemorySubPool::Default); // may free items from OffloadedBlocksRoot
         curr = next;
     }
 
diff --git a/ydb/library/yql/minikql/mkql_alloc.h b/ydb/library/yql/minikql/mkql_alloc.h
index f6aeec9120..96e4abae62 100644
--- a/ydb/library/yql/minikql/mkql_alloc.h
+++ b/ydb/library/yql/minikql/mkql_alloc.h
@@ -139,6 +139,29 @@ private:
     TAllocState::TListEntry* Prev = nullptr;
 };
 
+// TListEntry and IBoxedValue use the same place
+static_assert(sizeof(NUdf::IBoxedValue) == sizeof(TAllocState::TListEntry));
+
+class TBoxedValueWithFree : public NUdf::TBoxedValueBase {
+public:
+    void operator delete(void *mem) noexcept;
+};
+
+struct TMkqlPAllocHeader {
+    union {
+        TAllocState::TListEntry Entry;
+        TBoxedValueWithFree Boxed;
+    } U;
+
+    size_t Size;
+    void* Self; // should be placed right before pointer to allocated area, see GetMemoryChunkContext
+};
+
+static_assert(sizeof(TMkqlPAllocHeader) == 
+    sizeof(size_t) +
+    sizeof(TAllocState::TListEntry) +
+    sizeof(void*), "Padding is not allowed");
+
 class TScopedAlloc {
 public:
     explicit TScopedAlloc(const TSourceLocation& location,
@@ -289,7 +312,7 @@ inline void* MKQLAllocFastWithSize(size_t sz, TAllocState* state, const EMemoryS
 
 void MKQLFreeSlow(TAllocPageHeader* header, TAllocState *state, const EMemorySubPool mPool) noexcept;
 
-inline void MKQLFreeDeprecated(const void* mem, const EMemorySubPool mPool = EMemorySubPool::Default) noexcept {
+inline void MKQLFreeDeprecated(const void* mem, const EMemorySubPool mPool) noexcept {
     if (!mem) {
         return;
     }
@@ -346,7 +369,7 @@ inline void MKQLFreeFastWithSize(const void* mem, size_t sz, TAllocState* state,
     MKQLFreeSlow(header, state, mPool);
 }
 
-inline void* MKQLAllocDeprecated(size_t sz, const EMemorySubPool mPool = EMemorySubPool::Default) {
+inline void* MKQLAllocDeprecated(size_t sz, const EMemorySubPool mPool) {
     return MKQLAllocFastDeprecated(sz, TlsAllocState, mPool);
 }
 
@@ -658,6 +681,10 @@ private:
     size_t IndexInLastPage;
 };
 
+inline void TBoxedValueWithFree::operator delete(void *mem) noexcept {
+    auto size = ((TMkqlPAllocHeader*)mem)->Size;
+    return MKQLFreeWithSize(mem, size, EMemorySubPool::Default);
+}
 
 } // NMiniKQL
 
diff --git a/ydb/library/yql/parser/pg_wrapper/comp_factory.cpp b/ydb/library/yql/parser/pg_wrapper/comp_factory.cpp
index ee9ebd3834..c8dbff384c 100644
--- a/ydb/library/yql/parser/pg_wrapper/comp_factory.cpp
+++ b/ydb/library/yql/parser/pg_wrapper/comp_factory.cpp
@@ -104,20 +104,20 @@ NUdf::TUnboxedValue CreatePgString(i32 typeLen, ui32 targetTypeId, TStringBuf da
 }
 
 void *MkqlAllocSetAlloc(MemoryContext context, Size size) {
-    auto fullSize = size + PallocHdrSize;
-    auto ptr = (char *)MKQLAllocDeprecated(fullSize);
-    auto ret = (void*)(ptr + PallocHdrSize);
-    *(MemoryContext *)(((char *)ret) - sizeof(void *)) = context;
-    ((TAllocState::TListEntry*)ptr)->Link(TlsAllocState->CurrentPAllocList);
-    return ret;
+    auto fullSize = size + sizeof(TMkqlPAllocHeader);
+    auto header = (TMkqlPAllocHeader*)MKQLAllocWithSize(fullSize, EMemorySubPool::Default);
+    header->Size = size;
+    header->U.Entry.Link(TlsAllocState->CurrentPAllocList);
+    header->Self = context;
+    return header + 1;
 }
 
 void MkqlAllocSetFree(MemoryContext context, void* pointer) {
     if (pointer) {
-        auto original = (void*)((char*)pointer - PallocHdrSize);
+        auto header = ((TMkqlPAllocHeader*)pointer) - 1;
         // remove this block from list
-        ((TAllocState::TListEntry*)original)->Unlink();
-        MKQLFreeDeprecated(original);
+        header->U.Entry.Unlink();
+        MKQLFreeWithSize(header, header->Size, EMemorySubPool::Default);
     }
 }
 
@@ -127,12 +127,13 @@ void* MkqlAllocSetRealloc(MemoryContext context, void* pointer, Size size) {
         return nullptr;
     }
 
-    void* ret = MkqlAllocSetAlloc(context, size);
+    auto ret = MkqlAllocSetAlloc(context, size);
     if (pointer) {
-        memmove(ret, pointer, size);
+        auto header = ((TMkqlPAllocHeader*)pointer) - 1;
+        memmove(ret, pointer, header->Size);
+        MkqlAllocSetFree(context, pointer);
     }
 
-    MkqlAllocSetFree(context, pointer);
     return ret;
 }
 
@@ -1107,10 +1108,10 @@ NUdf::TUnboxedValuePod ConvertFromPgValue(NUdf::TUnboxedValuePod value, TMaybe<N
     case NUdf::EDataSlot::String:
     case NUdf::EDataSlot::Utf8:
         if (IsCString) {
-            auto x = (const char*)value.AsBoxed().Get() + PallocHdrSize;
+            auto x = (const char*)PointerDatumFromPod(value);
             return MakeString(TStringBuf(x));
         } else {
-            auto x = (const text*)((const char*)value.AsBoxed().Get() + PallocHdrSize);
+            auto x = (const text*)PointerDatumFromPod(value);
             return MakeString(GetVarBuf(x));
         }
     default:
@@ -3533,12 +3534,12 @@ public:
     }
 
     NUdf::TStringRef AsCStringBuffer(const NUdf::TUnboxedValue& value) const override {
-        auto x = (const char*)value.AsBoxed().Get() + PallocHdrSize;
+        auto x = (const char*)PointerDatumFromPod(value);
         return { x, strlen(x) + 1};
     }
 
     NUdf::TStringRef AsTextBuffer(const NUdf::TUnboxedValue& value) const override {
-        auto x = (const text*)((const char*)value.AsBoxed().Get() + PallocHdrSize);
+        auto x = (const text*)PointerDatumFromPod(value);
         return { (const char*)x, GetFullVarSize(x) };
     }
 
@@ -3557,7 +3558,7 @@ public:
     }
 
     NUdf::TStringRef AsFixedStringBuffer(const NUdf::TUnboxedValue& value, ui32 length) const override {
-        auto x = (const char*)value.AsBoxed().Get() + PallocHdrSize;
+        auto x = (const char*)PointerDatumFromPod(value);
         return { x, length };
     }
 };
diff --git a/ydb/library/yql/parser/pg_wrapper/utils.h b/ydb/library/yql/parser/pg_wrapper/utils.h
index 6965edde13..344d89d270 100644
--- a/ydb/library/yql/parser/pg_wrapper/utils.h
+++ b/ydb/library/yql/parser/pg_wrapper/utils.h
@@ -5,11 +5,6 @@
 
 namespace NYql {
 
-// allow to construct TListEntry in the space for IBoxedValue
-static_assert(sizeof(NKikimr::NUdf::IBoxedValue) >= sizeof(NKikimr::NMiniKQL::TAllocState::TListEntry));
-
-constexpr size_t PallocHdrSize = sizeof(void*) + sizeof(NKikimr::NUdf::IBoxedValue);
-
 inline NKikimr::NUdf::TUnboxedValuePod ScalarDatumToPod(Datum datum) {
     return NKikimr::NUdf::TUnboxedValuePod((ui64)datum);
 }
@@ -22,27 +17,19 @@ inline Datum ScalarDatumFromItem(const NKikimr::NUdf::TBlockItem& value) {
     return (Datum)value.As<ui64>();
 }
 
-class TBoxedValueWithFree : public NKikimr::NUdf::TBoxedValueBase {
-public:
-    void operator delete(void *mem) noexcept {
-        return NKikimr::NMiniKQL::MKQLFreeDeprecated(mem);
-    }
-};
-
 inline NKikimr::NUdf::TUnboxedValuePod PointerDatumToPod(Datum datum) {
-    auto original = (char*)datum - PallocHdrSize;
+    auto header = ((NKikimr::NMiniKQL::TMkqlPAllocHeader*)datum) - 1;
     // remove this block from list
-    ((NKikimr::NMiniKQL::TAllocState::TListEntry*)original)->Unlink();
-
-    auto raw = (NKikimr::NUdf::IBoxedValue*)original;
-    new(raw) TBoxedValueWithFree();
+    header->U.Entry.Unlink();
+    NKikimr::NUdf::IBoxedValue* raw = &header->U.Boxed;
+    new(raw) NKikimr::NMiniKQL::TBoxedValueWithFree();
     NKikimr::NUdf::IBoxedValuePtr ref(raw);
     return NKikimr::NUdf::TUnboxedValuePod(std::move(ref));
 }
 
 inline NKikimr::NUdf::TUnboxedValuePod OwnedPointerDatumToPod(Datum datum) {
-    auto original = (char*)datum - PallocHdrSize;
-    auto raw = (NKikimr::NUdf::IBoxedValue*)original;
+    auto header = ((NKikimr::NMiniKQL::TMkqlPAllocHeader*)datum) - 1;
+    NKikimr::NUdf::IBoxedValue* raw = &header->U.Boxed;
     NKikimr::NUdf::IBoxedValuePtr ref(raw);
     return NKikimr::NUdf::TUnboxedValuePod(std::move(ref));
 }
@@ -50,11 +37,12 @@ inline NKikimr::NUdf::TUnboxedValuePod OwnedPointerDatumToPod(Datum datum) {
 class TVPtrHolder {
 public:
     TVPtrHolder() {
-        new(Dummy) TBoxedValueWithFree();
+        new(Dummy) NKikimr::NMiniKQL::TBoxedValueWithFree();
     }
 
     static bool IsBoxedVPtr(Datum ptr) {
-        return *(const uintptr_t*)((char*)ptr - PallocHdrSize) == *(const uintptr_t*)Instance.Dummy;
+        auto header = ((NKikimr::NMiniKQL::TMkqlPAllocHeader*)ptr) - 1;
+        return *(const uintptr_t*)&header->U.Boxed == *(const uintptr_t*)Instance.Dummy;
     }
 
 private:
@@ -77,7 +65,7 @@ inline NKikimr::NUdf::TUnboxedValuePod AnyDatumToPod(Datum datum, bool passByVal
 }
 
 inline Datum PointerDatumFromPod(const NKikimr::NUdf::TUnboxedValuePod& value) {
-    return (Datum)(((const char*)value.AsBoxed().Get()) + PallocHdrSize);
+    return (Datum)(((const NKikimr::NMiniKQL::TMkqlPAllocHeader*)value.AsBoxed().Get()) + 1);
 }
 
 inline Datum PointerDatumFromItem(const NKikimr::NUdf::TBlockItem& value) {
diff --git a/ydb/library/yql/public/udf/service/exception_policy/udf_service.cpp b/ydb/library/yql/public/udf/service/exception_policy/udf_service.cpp
index f4c5930b58..bc5cafa331 100644
--- a/ydb/library/yql/public/udf/service/exception_policy/udf_service.cpp
+++ b/ydb/library/yql/public/udf/service/exception_policy/udf_service.cpp
@@ -3,11 +3,11 @@
 #include <ydb/library/yql/minikql/mkql_terminator.h>
 
 extern "C" void* UdfAllocate(ui64 size) {
-    return ::NKikimr::NMiniKQL::MKQLAllocDeprecated(size);
+    return ::NKikimr::NMiniKQL::MKQLAllocDeprecated(size, ::NKikimr::NMiniKQL::EMemorySubPool::Default);
 }
 
 extern "C" void UdfFree(const void* mem) {
-    return ::NKikimr::NMiniKQL::MKQLFreeDeprecated(mem);
+    return ::NKikimr::NMiniKQL::MKQLFreeDeprecated(mem, ::NKikimr::NMiniKQL::EMemorySubPool::Default);
 }
 
 extern "C" [[noreturn]] void UdfTerminate(const char* message) {
diff --git a/ydb/library/yql/public/udf/service/terminate_policy/udf_service.cpp b/ydb/library/yql/public/udf/service/terminate_policy/udf_service.cpp
index e795a8f04a..f5b228cb09 100644
--- a/ydb/library/yql/public/udf/service/terminate_policy/udf_service.cpp
+++ b/ydb/library/yql/public/udf/service/terminate_policy/udf_service.cpp
@@ -3,11 +3,11 @@
 #include <ydb/library/yql/minikql/mkql_terminator.h>
 
 extern "C" void* UdfAllocate(ui64 size) {
-    return ::NKikimr::NMiniKQL::MKQLAllocDeprecated(size);
+    return ::NKikimr::NMiniKQL::MKQLAllocDeprecated(size, ::NKikimr::NMiniKQL::EMemorySubPool::Default);
 }
 
 extern "C" void UdfFree(const void* mem) {
-    return ::NKikimr::NMiniKQL::MKQLFreeDeprecated(mem);
+    return ::NKikimr::NMiniKQL::MKQLFreeDeprecated(mem, ::NKikimr::NMiniKQL::EMemorySubPool::Default);
 }
 
 extern "C" [[noreturn]] void UdfTerminate(const char* message) {