#pragma once #include "aligned_page_pool.h" #include "mkql_mem_info.h" #include #include #include #include #include #include #include #include #include #include #include #include namespace NKikimr { namespace NMiniKQL { const ui64 MKQL_ALIGNMENT = 16; struct TAllocPageHeader { ui64 Capacity; ui64 Offset; ui64 UseCount; ui64 Deallocated; TAlignedPagePool* MyAlloc; TAllocPageHeader* Link; }; using TMemorySubPoolIdx = ui32; enum class EMemorySubPool: TMemorySubPoolIdx { Default = 0, Temporary = 1, Count }; constexpr ui32 MaxPageUserData = TAlignedPagePool::POOL_PAGE_SIZE - sizeof(TAllocPageHeader); static_assert(sizeof(TAllocPageHeader) % MKQL_ALIGNMENT == 0, "Incorrect size of header"); struct TAllocState : public TAlignedPagePool { struct TListEntry { TListEntry *Left = nullptr; TListEntry *Right = nullptr; void Link(TListEntry* root) noexcept; void Unlink() noexcept; void InitLinks() noexcept { Left = Right = this; } void Clear() noexcept { Left = Right = nullptr; } bool IsUnlinked() const noexcept { return !Left && !Right; } }; #ifndef NDEBUG std::unordered_map> ActiveMemInfo; #endif bool SupportsSizedAllocators = false; void* LargeAlloc(size_t size) { return Alloc(size); } void LargeFree(void* ptr, size_t size) noexcept { Free(ptr, size); } using TCurrentPages = std::array; static TAllocPageHeader EmptyPageHeader; static TCurrentPages EmptyCurrentPages; std::array CurrentPages = EmptyCurrentPages; TListEntry OffloadedBlocksRoot; TListEntry GlobalPAllocList; TListEntry* CurrentPAllocList; TListEntry ArrowBlocksRoot; std::unordered_set ArrowBuffers; bool EnableArrowTracking = true; void* MainContext = nullptr; void* CurrentContext = nullptr; struct TLockInfo { i32 OriginalRefs; i32 Locks; }; bool UseRefLocking = false; std::unordered_map LockedObjectsRefs; ::NKikimr::NUdf::TBoxedValueLink Root; NKikimr::NUdf::TBoxedValueLink* GetRoot() noexcept { return &Root; } explicit TAllocState(const TSourceLocation& location, const TAlignedPagePoolCounters& counters, bool supportsSizedAllocators); void KillAllBoxed(); void InvalidateMemInfo(); size_t GetDeallocatedInPages() const; static void CleanupPAllocList(TListEntry* root); static void CleanupArrowList(TListEntry* root); void LockObject(::NKikimr::NUdf::TUnboxedValuePod value); void UnlockObject(::NKikimr::NUdf::TUnboxedValuePod value); }; extern Y_POD_THREAD(TAllocState*) TlsAllocState; class TPAllocScope { public: TPAllocScope() { PAllocList.InitLinks(); Attach(); } ~TPAllocScope() { Cleanup(); Detach(); } void Attach() { Y_ABORT_UNLESS(!Prev); Prev = TlsAllocState->CurrentPAllocList; Y_ABORT_UNLESS(Prev); TlsAllocState->CurrentPAllocList = &PAllocList; } void Detach() { if (Prev) { Y_ABORT_UNLESS(TlsAllocState->CurrentPAllocList == &PAllocList); TlsAllocState->CurrentPAllocList = Prev; Prev = nullptr; } } void Cleanup() { TAllocState::CleanupPAllocList(&PAllocList); } private: TAllocState::TListEntry PAllocList; 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; ui64 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"); constexpr size_t ArrowAlignment = 64; struct TMkqlArrowHeader { TAllocState::TListEntry Entry; ui64 Size; char Padding[ArrowAlignment - sizeof(TAllocState::TListEntry) - sizeof(ui64)]; }; static_assert(sizeof(TMkqlArrowHeader) == ArrowAlignment); class TScopedAlloc { public: explicit TScopedAlloc(const TSourceLocation& location, const TAlignedPagePoolCounters& counters = TAlignedPagePoolCounters(), bool supportsSizedAllocators = false, bool initiallyAcquired = true) : InitiallyAcquired_(initiallyAcquired) , MyState_(location, counters, supportsSizedAllocators) { MyState_.MainContext = PgInitializeMainContext(); if (InitiallyAcquired_) { Acquire(); } } ~TScopedAlloc() { if (!InitiallyAcquired_) { Acquire(); } MyState_.KillAllBoxed(); Release(); PgDestroyMainContext(MyState_.MainContext); } TAllocState& Ref() { return MyState_; } void Acquire(); void Release(); size_t GetUsed() const { return MyState_.GetUsed(); } size_t GetPeakUsed() const { return MyState_.GetPeakUsed(); } size_t GetAllocated() const { return MyState_.GetAllocated(); } size_t GetPeakAllocated() const { return MyState_.GetPeakAllocated(); } size_t GetLimit() const { return MyState_.GetLimit(); } void SetLimit(size_t limit) { MyState_.SetLimit(limit); } void DisableStrictAllocationCheck() { MyState_.DisableStrictAllocationCheck(); } void ReleaseFreePages() { MyState_.ReleaseFreePages(); } void InvalidateMemInfo() { MyState_.InvalidateMemInfo(); } bool IsAttached() const { return AttachedCount_ > 0; } void SetGUCSettings(const TGUCSettings::TPtr& GUCSettings) { Acquire(); PgSetGUCSettings(MyState_.MainContext, GUCSettings); Release(); } void SetMaximumLimitValueReached(bool IsReached) { MyState_.SetMaximumLimitValueReached(IsReached); } private: const bool InitiallyAcquired_; TAllocState MyState_; size_t AttachedCount_ = 0; TAllocState* PrevState_ = nullptr; }; class TPagedArena { public: TPagedArena(TAlignedPagePool* pagePool) noexcept : PagePool_(pagePool) , CurrentPages_(TAllocState::EmptyCurrentPages) {} TPagedArena(const TPagedArena&) = delete; TPagedArena(TPagedArena&& other) noexcept : PagePool_(other.PagePool_) , CurrentPages_(other.CurrentPages_) { other.CurrentPages_ = TAllocState::EmptyCurrentPages; } void operator=(const TPagedArena&) = delete; void operator=(TPagedArena&& other) noexcept { Clear(); PagePool_ = other.PagePool_; CurrentPages_ = other.CurrentPages_; other.CurrentPages_ = TAllocState::EmptyCurrentPages; } ~TPagedArena() noexcept { Clear(); } void* Alloc(size_t sz, const EMemorySubPool pagePool = EMemorySubPool::Default) { auto& currentPage = CurrentPages_[(TMemorySubPoolIdx)pagePool]; if (Y_LIKELY(currentPage->Offset + sz <= currentPage->Capacity)) { void* ret = (char*)currentPage + currentPage->Offset; currentPage->Offset = AlignUp(currentPage->Offset + sz, MKQL_ALIGNMENT); return ret; } return AllocSlow(sz, pagePool); } void Clear() noexcept; private: void* AllocSlow(const size_t sz, const EMemorySubPool pagePool); private: TAlignedPagePool* PagePool_; TAllocState::TCurrentPages CurrentPages_ = TAllocState::EmptyCurrentPages; }; void* MKQLAllocSlow(size_t sz, TAllocState* state, const EMemorySubPool mPool); inline void* MKQLAllocFastDeprecated(size_t sz, TAllocState* state, const EMemorySubPool mPool) { Y_DEBUG_ABORT_UNLESS(state); #ifdef PROFILE_MEMORY_ALLOCATIONS auto ret = (TAllocState::TListEntry*)malloc(sizeof(TAllocState::TListEntry) + sz); if (!ret) { throw TMemoryLimitExceededException(); } ret->Link(&state->OffloadedBlocksRoot); return ret + 1; #endif auto currPage = state->CurrentPages[(TMemorySubPoolIdx)mPool]; if (Y_LIKELY(currPage->Offset + sz <= currPage->Capacity)) { void* ret = (char*)currPage + currPage->Offset; currPage->Offset = AlignUp(currPage->Offset + sz, MKQL_ALIGNMENT); ++currPage->UseCount; return ret; } return MKQLAllocSlow(sz, state, mPool); } inline void* MKQLAllocFastWithSize(size_t sz, TAllocState* state, const EMemorySubPool mPool) { Y_DEBUG_ABORT_UNLESS(state); bool useMemalloc = state->SupportsSizedAllocators && sz > MaxPageUserData; #ifdef PROFILE_MEMORY_ALLOCATIONS useMemalloc = true; #endif if (useMemalloc) { state->OffloadAlloc(sizeof(TAllocState::TListEntry) + sz); auto ret = (TAllocState::TListEntry*)malloc(sizeof(TAllocState::TListEntry) + sz); if (!ret) { throw TMemoryLimitExceededException(); } ret->Link(&state->OffloadedBlocksRoot); return ret + 1; } auto currPage = state->CurrentPages[(TMemorySubPoolIdx)mPool]; if (Y_LIKELY(currPage->Offset + sz <= currPage->Capacity)) { void* ret = (char*)currPage + currPage->Offset; currPage->Offset = AlignUp(currPage->Offset + sz, MKQL_ALIGNMENT); ++currPage->UseCount; return ret; } return MKQLAllocSlow(sz, state, mPool); } void MKQLFreeSlow(TAllocPageHeader* header, TAllocState *state, const EMemorySubPool mPool) noexcept; inline void MKQLFreeDeprecated(const void* mem, const EMemorySubPool mPool) noexcept { if (!mem) { return; } #ifdef PROFILE_MEMORY_ALLOCATIONS TAllocState *state = TlsAllocState; Y_DEBUG_ABORT_UNLESS(state); auto entry = (TAllocState::TListEntry*)(mem) - 1; entry->Unlink(); free(entry); return; #endif TAllocPageHeader* header = (TAllocPageHeader*)TAllocState::GetPageStart(mem); Y_DEBUG_ABORT_UNLESS(header->MyAlloc == TlsAllocState, "%s", (TStringBuilder() << "wrong allocator was used; " "allocated with: " << header->MyAlloc->GetDebugInfo() << " freed with: " << TlsAllocState->GetDebugInfo()).data()); if (Y_LIKELY(--header->UseCount != 0)) { return; } MKQLFreeSlow(header, TlsAllocState, mPool); } inline void MKQLFreeFastWithSize(const void* mem, size_t sz, TAllocState* state, const EMemorySubPool mPool) noexcept { if (!mem) { return; } Y_DEBUG_ABORT_UNLESS(state); bool useFree = state->SupportsSizedAllocators && sz > MaxPageUserData; #ifdef PROFILE_MEMORY_ALLOCATIONS useFree = true; #endif if (useFree) { auto entry = (TAllocState::TListEntry*)(mem) - 1; entry->Unlink(); free(entry); state->OffloadFree(sizeof(TAllocState::TListEntry) + sz); return; } TAllocPageHeader* header = (TAllocPageHeader*)TAllocState::GetPageStart(mem); Y_DEBUG_ABORT_UNLESS(header->MyAlloc == state, "%s", (TStringBuilder() << "wrong allocator was used; " "allocated with: " << header->MyAlloc->GetDebugInfo() << " freed with: " << TlsAllocState->GetDebugInfo()).data()); if (Y_LIKELY(--header->UseCount != 0)) { header->Deallocated += sz; return; } MKQLFreeSlow(header, state, mPool); } inline void* MKQLAllocDeprecated(size_t sz, const EMemorySubPool mPool) { return MKQLAllocFastDeprecated(sz, TlsAllocState, mPool); } inline void* MKQLAllocWithSize(size_t sz, const EMemorySubPool mPool) { return MKQLAllocFastWithSize(sz, TlsAllocState, mPool); } inline void MKQLFreeWithSize(const void* mem, size_t sz, const EMemorySubPool mPool) noexcept { MKQLFreeFastWithSize(mem, sz, TlsAllocState, mPool); } inline void MKQLRegisterObject(NUdf::TBoxedValue* value) noexcept { value->Link(TlsAllocState->GetRoot()); } inline void MKQLUnregisterObject(NUdf::TBoxedValue* value) noexcept { value->Unlink(); } void* MKQLArrowAllocate(ui64 size); void* MKQLArrowReallocate(const void* mem, ui64 prevSize, ui64 size); void MKQLArrowFree(const void* mem, ui64 size); void MKQLArrowUntrack(const void* mem); template struct TWithMiniKQLAlloc { static constexpr EMemorySubPool MemoryPool = MemoryPoolExt; static void FreeWithSize(const void* mem, const size_t sz) { NMiniKQL::MKQLFreeWithSize(mem, sz, MemoryPool); } static void* AllocWithSize(const size_t sz) { return NMiniKQL::MKQLAllocWithSize(sz, MemoryPool); } void* operator new(size_t sz) { return NMiniKQL::MKQLAllocWithSize(sz, MemoryPool); } void* operator new[](size_t sz) { return NMiniKQL::MKQLAllocWithSize(sz, MemoryPool); } void operator delete(void *mem, std::size_t sz) noexcept { NMiniKQL::MKQLFreeWithSize(mem, sz, MemoryPool); } void operator delete[](void *mem, std::size_t sz) noexcept { NMiniKQL::MKQLFreeWithSize(mem, sz, MemoryPool); } }; template T* AllocateOn(TAllocState* state, Args&&... args) { void* addr = MKQLAllocFastWithSize(sizeof(T), state, T::MemoryPool); return ::new(addr) T(std::forward(args)...); static_assert(std::is_base_of, T>::value, "Class must inherit TWithMiniKQLAlloc."); } template struct TMKQLAllocator { typedef Type value_type; typedef Type* pointer; typedef const Type* const_pointer; typedef Type& reference; typedef const Type& const_reference; typedef size_t size_type; typedef ptrdiff_t difference_type; TMKQLAllocator() noexcept = default; ~TMKQLAllocator() noexcept = default; template TMKQLAllocator(const TMKQLAllocator&) noexcept {} template struct rebind { typedef TMKQLAllocator other; }; template bool operator==(const TMKQLAllocator&) const { return true; } template bool operator!=(const TMKQLAllocator&) const { return false; } static pointer allocate(size_type n, const void* = nullptr) { return static_cast(MKQLAllocWithSize(n * sizeof(value_type), MemoryPool)); } static void deallocate(const_pointer p, size_type n) noexcept { MKQLFreeWithSize(p, n * sizeof(value_type), MemoryPool); } }; using TWithDefaultMiniKQLAlloc = TWithMiniKQLAlloc; using TWithTemporaryMiniKQLAlloc = TWithMiniKQLAlloc; template struct TMKQLHugeAllocator { typedef Type value_type; typedef Type* pointer; typedef const Type* const_pointer; typedef Type& reference; typedef const Type& const_reference; typedef size_t size_type; typedef ptrdiff_t difference_type; TMKQLHugeAllocator() noexcept = default; ~TMKQLHugeAllocator() noexcept = default; template TMKQLHugeAllocator(const TMKQLHugeAllocator&) noexcept {} template struct rebind { typedef TMKQLHugeAllocator other; }; template bool operator==(const TMKQLHugeAllocator&) const { return true; } template bool operator!=(const TMKQLHugeAllocator&) const { return false; } static pointer allocate(size_type n, const void* = nullptr) { size_t size = Max(n * sizeof(value_type), TAllocState::POOL_PAGE_SIZE); return static_cast(TlsAllocState->GetBlock(size)); } static void deallocate(const_pointer p, size_type n) noexcept { size_t size = Max(n * sizeof(value_type), TAllocState::POOL_PAGE_SIZE); TlsAllocState->ReturnBlock(const_cast(p), size); } }; template class TPagedList { public: static_assert(sizeof(T) <= TAlignedPagePool::POOL_PAGE_SIZE, "Too big object"); static constexpr size_t OBJECTS_PER_PAGE = TAlignedPagePool::POOL_PAGE_SIZE / sizeof(T); class TIterator; class TConstIterator; TPagedList(TAlignedPagePool& pool) : Pool(pool) , IndexInLastPage(OBJECTS_PER_PAGE) {} TPagedList(const TPagedList&) = delete; TPagedList(TPagedList&&) = delete; ~TPagedList() { Clear(); } void Add(T&& value) { if (IndexInLastPage < OBJECTS_PER_PAGE) { auto ptr = ObjectAt(Pages.back(), IndexInLastPage); new(ptr) T(std::move(value)); ++IndexInLastPage; return; } auto ptr = Pool.GetPage(); IndexInLastPage = 1; Pages.push_back(ptr); new(ptr) T(std::move(value)); } void Clear() { for (ui32 i = 0; i + 1 < Pages.size(); ++i) { for (ui32 objIndex = 0; objIndex < OBJECTS_PER_PAGE; ++objIndex) { ObjectAt(Pages[i], objIndex)->~T(); } Pool.ReturnPage(Pages[i]); } if (!Pages.empty()) { for (ui32 objIndex = 0; objIndex < IndexInLastPage; ++objIndex) { ObjectAt(Pages.back(), objIndex)->~T(); } Pool.ReturnPage(Pages.back()); } TPages().swap(Pages); IndexInLastPage = OBJECTS_PER_PAGE; } const T& operator[](size_t i) const { const auto table = i / OBJECTS_PER_PAGE; const auto index = i % OBJECTS_PER_PAGE; return *ObjectAt(Pages[table], index); } size_t Size() const { return Pages.empty() ? 0 : ((Pages.size() - 1) * OBJECTS_PER_PAGE + IndexInLastPage); } TConstIterator Begin() const { return TConstIterator(this, 0, 0); } TConstIterator begin() const { return Begin(); } TConstIterator End() const { if (IndexInLastPage == OBJECTS_PER_PAGE) { return TConstIterator(this, Pages.size(), 0); } return TConstIterator(this, Pages.size() - 1, IndexInLastPage); } TConstIterator end() const { return End(); } TIterator Begin() { return TIterator(this, 0, 0); } TIterator begin() { return Begin(); } TIterator End() { if (IndexInLastPage == OBJECTS_PER_PAGE) { return TIterator(this, Pages.size(), 0); } return TIterator(this, Pages.size() - 1, IndexInLastPage); } TIterator end() { return End(); } class TIterator { public: using TOwner = TPagedList; TIterator() : Owner(nullptr) , PageNo(0) , PageIndex(0) {} TIterator(const TIterator&) = default; TIterator& operator=(const TIterator&) = default; TIterator(TOwner* owner, size_t pageNo, size_t pageIndex) : Owner(owner) , PageNo(pageNo) , PageIndex(pageIndex) {} T& operator*() { Y_DEBUG_ABORT_UNLESS(PageIndex < OBJECTS_PER_PAGE); Y_DEBUG_ABORT_UNLESS(PageNo < Owner->Pages.size()); Y_DEBUG_ABORT_UNLESS(PageNo + 1 < Owner->Pages.size() || PageIndex < Owner->IndexInLastPage); return *Owner->ObjectAt(Owner->Pages[PageNo], PageIndex); } TIterator& operator++() { if (++PageIndex == OBJECTS_PER_PAGE) { ++PageNo; PageIndex = 0; } return *this; } bool operator==(const TIterator& other) const { return PageNo == other.PageNo && PageIndex == other.PageIndex; } bool operator!=(const TIterator& other) const { return !operator==(other); } private: TOwner* Owner; size_t PageNo; size_t PageIndex; }; class TConstIterator { public: using TOwner = TPagedList; TConstIterator() : Owner(nullptr) , PageNo(0) , PageIndex(0) {} TConstIterator(const TConstIterator&) = default; TConstIterator& operator=(const TConstIterator&) = default; TConstIterator(const TOwner* owner, size_t pageNo, size_t pageIndex) : Owner(owner) , PageNo(pageNo) , PageIndex(pageIndex) {} const T& operator*() { Y_DEBUG_ABORT_UNLESS(PageIndex < OBJECTS_PER_PAGE); Y_DEBUG_ABORT_UNLESS(PageNo < Owner->Pages.size()); Y_DEBUG_ABORT_UNLESS(PageNo + 1 < Owner->Pages.size() || PageIndex < Owner->IndexInLastPage); return *Owner->ObjectAt(Owner->Pages[PageNo], PageIndex); } TConstIterator& operator++() { if (++PageIndex == OBJECTS_PER_PAGE) { ++PageNo; PageIndex = 0; } return *this; } bool operator==(const TConstIterator& other) const { return PageNo == other.PageNo && PageIndex == other.PageIndex; } bool operator!=(const TConstIterator& other) const { return !operator==(other); } private: const TOwner* Owner; size_t PageNo; size_t PageIndex; }; private: static const T* ObjectAt(const void* page, size_t objectIndex) { return reinterpret_cast(static_cast(page) + objectIndex * sizeof(T)); } static T* ObjectAt(void* page, size_t objectIndex) { return reinterpret_cast(static_cast(page) + objectIndex * sizeof(T)); } TAlignedPagePool& Pool; using TPages = std::vector>; TPages Pages; size_t IndexInLastPage; }; inline void TBoxedValueWithFree::operator delete(void *mem) noexcept { auto size = ((TMkqlPAllocHeader*)mem)->Size + sizeof(TMkqlPAllocHeader); MKQLFreeWithSize(mem, size, EMemorySubPool::Default); } } // NMiniKQL } // NKikimr