More TChunkedMemoryPool, TChunkedMemoryAllocator, TChunkedMemoryPoolOutput to library

More TChunkedMemoryPool, TChunkedMemoryAllocator, TChunkedMemoryPoolOutput to library

wip

1 files changed, 65 insertions, 0 deletions

diff --git a/library/cpp/yt/memory/chunked_memory_pool_output.cpp b/library/cpp/yt/memory/chunked_memory_pool_output.cpp
new file mode 100644
index 0000000000..332096445a
--- /dev/null
+++ b/library/cpp/yt/memory/chunked_memory_pool_output.cpp
@@ -0,0 +1,65 @@
+#include "chunked_memory_pool_output.h"
+
+#include "chunked_memory_pool.h"
+
+#include <library/cpp/yt/memory/ref.h>
+
+namespace NYT {
+
+////////////////////////////////////////////////////////////////////////////////
+
+TChunkedMemoryPoolOutput::TChunkedMemoryPoolOutput(TChunkedMemoryPool* pool, size_t chunkSize)
+    : Pool_(pool)
+    , ChunkSize_(chunkSize)
+{ }
+
+size_t TChunkedMemoryPoolOutput::DoNext(void** ptr)
+{
+    // Check if the current chunk is exhausted.
+    if (Current_ == End_) {
+        // Emplace the (whole) last chunk, if any.
+        if (Begin_) {
+            Refs_.emplace_back(Begin_, Current_);
+        }
+        // Allocate a new chunk.
+        // Use |AllocateAligned| to get a chance to free some memory afterwards.
+        // Tune the number of bytes requested from the pool to try avoid allocations.
+        auto spareSize = Pool_->GetCurrentChunkSpareSize();
+        auto allocationSize = (spareSize == 0 ? ChunkSize_ : std::min(ChunkSize_, spareSize));
+        Begin_ = Pool_->AllocateAligned(allocationSize, /* align */ 1);
+        Current_ = Begin_;
+        End_ = Begin_ + allocationSize;
+    }
+
+    // Return the unused part of the current chunk.
+    // This could be the whole chunk allocated above.
+    *ptr = Current_;
+    auto size = End_ - Current_;
+    Current_ = End_;
+    return size;
+}
+
+void TChunkedMemoryPoolOutput::DoUndo(size_t size)
+{
+    // Just rewind the current pointer.
+    Current_ -= size;
+    YT_VERIFY(Current_ >= Begin_);
+}
+
+std::vector<TMutableRef> TChunkedMemoryPoolOutput::Finish()
+{
+    // Emplace the used part of the last chunk, if any.
+    if (Begin_) {
+        Refs_.emplace_back(Begin_, Current_);
+    }
+    // Try to free the unused part of the last chunk, if possible.
+    if (Current_ < End_) {
+        Pool_->Free(Current_, End_);
+    }
+    return std::move(Refs_);
+}
+
+////////////////////////////////////////////////////////////////////////////////
+
+} // namespace NYT
+