More TChunkedMemoryPool, TChunkedMemoryAllocator, TChunkedMemoryPoolOutput to library

More TChunkedMemoryPool, TChunkedMemoryAllocator, TChunkedMemoryPoolOutput to library

wip

1 files changed, 200 insertions, 0 deletions

diff --git a/library/cpp/yt/memory/chunked_memory_pool.cpp b/library/cpp/yt/memory/chunked_memory_pool.cpp
new file mode 100644
index 00000000000..a10a6fe724f
--- /dev/null
+++ b/library/cpp/yt/memory/chunked_memory_pool.cpp
@@ -0,0 +1,200 @@
+#include "chunked_memory_pool.h"
+
+namespace NYT {
+
+////////////////////////////////////////////////////////////////////////////////
+
+TAllocationHolder::TAllocationHolder(TMutableRef ref, TRefCountedTypeCookie cookie)
+    : Ref_(ref)
+#ifdef YT_ENABLE_REF_COUNTED_TRACKING
+    , Cookie_(cookie)
+#endif
+{
+#ifdef YT_ENABLE_REF_COUNTED_TRACKING
+    if (Cookie_ != NullRefCountedTypeCookie) {
+        TRefCountedTrackerFacade::AllocateTagInstance(Cookie_);
+        TRefCountedTrackerFacade::AllocateSpace(Cookie_, Ref_.Size());
+    }
+#endif
+}
+
+TAllocationHolder::~TAllocationHolder()
+{
+#ifdef YT_ENABLE_REF_COUNTED_TRACKING
+    if (Cookie_ != NullRefCountedTypeCookie) {
+        TRefCountedTrackerFacade::FreeTagInstance(Cookie_);
+        TRefCountedTrackerFacade::FreeSpace(Cookie_, Ref_.Size());
+    }
+#endif
+}
+
+////////////////////////////////////////////////////////////////////////////////
+
+class TDefaultMemoryChunkProvider
+    : public IMemoryChunkProvider
+{
+public:
+    std::unique_ptr<TAllocationHolder> Allocate(size_t size, TRefCountedTypeCookie cookie) override
+    {
+        return std::unique_ptr<TAllocationHolder>(TAllocationHolder::Allocate<TAllocationHolder>(size, cookie));
+    }
+};
+
+const IMemoryChunkProviderPtr& GetDefaultMemoryChunkProvider()
+{
+    static const IMemoryChunkProviderPtr Result = New<TDefaultMemoryChunkProvider>();
+    return Result;
+}
+
+////////////////////////////////////////////////////////////////////////////////
+
+TChunkedMemoryPool::TChunkedMemoryPool(
+    TRefCountedTypeCookie tagCookie,
+    IMemoryChunkProviderPtr chunkProvider,
+    size_t startChunkSize)
+    : TagCookie_(tagCookie)
+    , ChunkProviderHolder_(std::move(chunkProvider))
+    , ChunkProvider_(ChunkProviderHolder_.Get())
+{
+    Initialize(startChunkSize);
+}
+
+TChunkedMemoryPool::TChunkedMemoryPool(
+    TRefCountedTypeCookie tagCookie,
+    size_t startChunkSize)
+    : TagCookie_(tagCookie)
+    , ChunkProvider_(GetDefaultMemoryChunkProvider().Get())
+{
+    Initialize(startChunkSize);
+}
+
+void TChunkedMemoryPool::Initialize(size_t startChunkSize)
+{
+    NextSmallSize_ = startChunkSize;
+    FreeZoneBegin_ = nullptr;
+    FreeZoneEnd_ = nullptr;
+}
+
+void TChunkedMemoryPool::Purge()
+{
+    Chunks_.clear();
+    OtherBlocks_.clear();
+    Size_ = 0;
+    Capacity_ = 0;
+    NextChunkIndex_ = 0;
+    FreeZoneBegin_ = nullptr;
+    FreeZoneEnd_ = nullptr;
+}
+
+char* TChunkedMemoryPool::AllocateUnalignedSlow(size_t size)
+{
+    auto* large = AllocateSlowCore(size);
+    if (large) {
+        return large;
+    }
+    return AllocateUnaligned(size);
+}
+
+char* TChunkedMemoryPool::AllocateAlignedSlow(size_t size, int align)
+{
+    // NB: Do not rely on any particular alignment of chunks.
+    auto* large = AllocateSlowCore(size + align);
+    if (large) {
+        return AlignUp(large, align);
+    }
+    return AllocateAligned(size, align);
+}
+
+char* TChunkedMemoryPool::AllocateSlowCore(size_t size)
+{
+    TMutableRef ref;
+    if (size > RegularChunkSize) {
+        auto block = ChunkProvider_->Allocate(size, TagCookie_);
+        ref = block->GetRef();
+        Size_ += size;
+        Capacity_ += ref.Size();
+        OtherBlocks_.push_back(std::move(block));
+        return ref.Begin();
+    }
+
+    YT_VERIFY(NextChunkIndex_ <= std::ssize(Chunks_));
+
+    if (NextSmallSize_ < RegularChunkSize) {
+        auto block = ChunkProvider_->Allocate(std::max(NextSmallSize_, size), TagCookie_);
+        ref = block->GetRef();
+        Capacity_ += ref.Size();
+        OtherBlocks_.push_back(std::move(block));
+        NextSmallSize_ = 2 * ref.Size();
+    } else if (NextChunkIndex_ == std::ssize(Chunks_)) {
+        auto chunk = ChunkProvider_->Allocate(RegularChunkSize, TagCookie_);
+        ref = chunk->GetRef();
+        Capacity_ += ref.Size();
+        Chunks_.push_back(std::move(chunk));
+        ++NextChunkIndex_;
+    } else {
+        ref = Chunks_[NextChunkIndex_++]->GetRef();
+    }
+
+    FreeZoneBegin_ = ref.Begin();
+    FreeZoneEnd_ = ref.End();
+
+    return nullptr;
+}
+
+void TChunkedMemoryPool::Absorb(TChunkedMemoryPool&& other)
+{
+    YT_VERIFY(ChunkProvider_ == other.ChunkProvider_);
+
+    OtherBlocks_.reserve(OtherBlocks_.size() + other.OtherBlocks_.size());
+    for (auto& block : other.OtherBlocks_) {
+        OtherBlocks_.push_back(std::move(block));
+    }
+    other.OtherBlocks_.clear();
+
+    // Suppose that
+    // - "A" is filled blocks of the current pool;
+    // - "a" is free blocks of the current pool;
+    // - "B" is filled blocks of the other pool;
+    // - "b" is free blocks of the other pool.
+    // Then, from the initial layouts "AA...Aaa...a" and "BB...Bbb...b" we obtain "BB..BAA..Aaa...abb...b".
+    Chunks_.reserve(Chunks_.size() + other.Chunks_.size());
+    size_t oldSize = Chunks_.size();
+    for (auto& chunk : other.Chunks_) {
+        Chunks_.push_back(std::move(chunk));
+    }
+    // Transform "AA...Aaa...aBB...B" => "BB...BAA...Aaa...a"
+    std::rotate(Chunks_.begin(), Chunks_.begin() + oldSize, Chunks_.begin() + oldSize + other.NextChunkIndex_);
+    if (NextChunkIndex_ == 0) {
+        FreeZoneBegin_ = other.FreeZoneBegin_;
+        FreeZoneEnd_ = other.FreeZoneEnd_;
+    }
+    NextChunkIndex_ += other.NextChunkIndex_;
+    other.Chunks_.clear();
+    other.FreeZoneBegin_ = nullptr;
+    other.FreeZoneEnd_ = nullptr;
+    other.NextChunkIndex_ = 0;
+
+    Size_ += other.Size_;
+    Capacity_ += other.Capacity_;
+    other.Size_ = 0;
+    other.Capacity_ = 0;
+}
+
+size_t TChunkedMemoryPool::GetSize() const
+{
+    return Size_;
+}
+
+size_t TChunkedMemoryPool::GetCapacity() const
+{
+    return Capacity_;
+}
+
+size_t TChunkedMemoryPool::GetCurrentChunkSpareSize() const
+{
+    return FreeZoneEnd_ - FreeZoneBegin_;
+}
+
+////////////////////////////////////////////////////////////////////////////////
+
+} // namespace NYT