Intermediate changes

4 files changed, 295 insertions, 0 deletions

diff --git a/library/cpp/yt/memory/chunked_memory_pool_allocator-inl.h b/library/cpp/yt/memory/chunked_memory_pool_allocator-inl.h
new file mode 100644
index 00000000000..8c221c82da0
--- /dev/null
+++ b/library/cpp/yt/memory/chunked_memory_pool_allocator-inl.h
@@ -0,0 +1,61 @@
+#ifndef CHUNKED_MEMORY_POOL_ALLOCATOR_INL_H_
+#error "Direct inclusion of this file is not allowed, include chunked_memory_pool_allocator.h"
+// For the sake of sane code completion.
+#include "chunked_memory_pool_allocator.h"
+#endif
+
+#include <util/system/align.h>
+
+namespace NYT {
+
+////////////////////////////////////////////////////////////////////////////////
+
+template <class T>
+TChunkedMemoryPoolAllocator<T> TChunkedMemoryPoolAllocator<T>::select_on_container_copy_construction(
+    const TChunkedMemoryPoolAllocator& allocator) noexcept
+{
+    return allocator;
+}
+
+template <class T>
+TChunkedMemoryPoolAllocator<T>::TChunkedMemoryPoolAllocator(
+    TChunkedMemoryPool* pool) noexcept
+    : Pool_(pool)
+{ }
+
+template <class T>
+template <class U>
+TChunkedMemoryPoolAllocator<T>::TChunkedMemoryPoolAllocator(
+    const TChunkedMemoryPoolAllocator<U>& other) noexcept
+    : Pool_(other.Pool_)
+{ }
+
+template <class T>
+T* TChunkedMemoryPoolAllocator<T>::allocate(std::size_t count)
+{
+    return reinterpret_cast<T*>(Pool_->AllocateAligned(count * sizeof(T), alignof(T)));
+}
+
+template <class T>
+void TChunkedMemoryPoolAllocator<T>::deallocate(
+    T* /*pointer*/,
+    std::size_t /*count*/) const noexcept
+{ }
+
+template <class T>
+bool TChunkedMemoryPoolAllocator<T>::operator==(
+    const TChunkedMemoryPoolAllocator& other) const noexcept
+{
+    return Pool_ == other.Pool_;
+}
+
+template <class T>
+bool TChunkedMemoryPoolAllocator<T>::operator!=(
+    const TChunkedMemoryPoolAllocator& other) const noexcept
+{
+    return !(*this == other);
+}
+
+////////////////////////////////////////////////////////////////////////////////
+
+} // namespace NYT
diff --git a/library/cpp/yt/memory/chunked_memory_pool_allocator.h b/library/cpp/yt/memory/chunked_memory_pool_allocator.h
new file mode 100644
index 00000000000..937bfeef6d6
--- /dev/null
+++ b/library/cpp/yt/memory/chunked_memory_pool_allocator.h
@@ -0,0 +1,61 @@
+#pragma once
+
+#include "chunked_memory_pool.h"
+
+#include <cstddef>
+#include <type_traits>
+
+namespace NYT {
+
+////////////////////////////////////////////////////////////////////////////////
+
+// Made to use TChunkedMemoryPool with std containers.
+// Does not own pool so lifetime should be no greater than lifetime of the pool.
+// Redirects allocations to pool.
+// Deallocations are noop.
+template <class T>
+class TChunkedMemoryPoolAllocator
+{
+public:
+    using value_type = T;
+
+    // All defines for AllocatorAwareContainer are deliberately set to the same as default.
+    using is_always_equal = std::false_type;
+
+    using propagate_on_container_copy_assignment = std::false_type;
+    using propagate_on_container_move_assignment = std::false_type;
+    using propagate_on_container_swap = std::false_type;
+
+    static TChunkedMemoryPoolAllocator select_on_container_copy_construction(
+        const TChunkedMemoryPoolAllocator& allocator) noexcept;
+
+    explicit TChunkedMemoryPoolAllocator(TChunkedMemoryPool* pool) noexcept;
+
+    template <class U>
+    TChunkedMemoryPoolAllocator(const TChunkedMemoryPoolAllocator<U>& other) noexcept;
+
+    [[nodiscard]] T* allocate(std::size_t count);
+    void deallocate(T* pointer, std::size_t count) const noexcept;
+
+    // https://en.cppreference.com/w/cpp/named_req/Allocator
+    // a1 == a2
+    // * true only if the storage allocated by the allocator a1 can be deallocated through a2.
+    // * Establishes reflexive, symmetric, and transitive relationship.
+    // * Does not throw exceptions.
+    bool operator==(const TChunkedMemoryPoolAllocator& other) const noexcept;
+    bool operator!=(const TChunkedMemoryPoolAllocator& other) const noexcept;
+
+private:
+    template <class U>
+    friend class TChunkedMemoryPoolAllocator;
+
+    TChunkedMemoryPool* const Pool_;
+};
+
+////////////////////////////////////////////////////////////////////////////////
+
+} // namespace NYT
+
+#define CHUNKED_MEMORY_POOL_ALLOCATOR_INL_H_
+#include "chunked_memory_pool_allocator-inl.h"
+#undef CHUNKED_MEMORY_POOL_ALLOCATOR_INL_H_
diff --git a/library/cpp/yt/memory/unittests/chunked_memory_pool_allocator_ut.cpp b/library/cpp/yt/memory/unittests/chunked_memory_pool_allocator_ut.cpp
new file mode 100644
index 00000000000..214ee80643f
--- /dev/null
+++ b/library/cpp/yt/memory/unittests/chunked_memory_pool_allocator_ut.cpp
@@ -0,0 +1,172 @@
+#include <library/cpp/testing/gtest/gtest.h>
+
+#include <library/cpp/yt/memory/chunked_memory_pool.h>
+#include <library/cpp/yt/memory/chunked_memory_pool_allocator.h>
+
+#include <random>
+#include <set>
+#include <vector>
+
+namespace NYT {
+namespace {
+
+////////////////////////////////////////////////////////////////////////////////
+
+TEST(TChunkedMemoryPoolAllocatorTest, SimpleSet)
+{
+    TChunkedMemoryPool pool;
+    TChunkedMemoryPoolAllocator<int> allocator(&pool);
+
+    using TSet = std::set<int, std::less<int>, TChunkedMemoryPoolAllocator<int>>;
+    TSet set(allocator);
+
+    constexpr int seed = 0;
+    std::mt19937 generator(seed);
+    auto valueDistribution = std::uniform_int_distribution<int>(
+        -10000,
+        10000);
+
+    constexpr int iterationCount = 10000;
+    constexpr int nodeSize = sizeof(TSet::node_type);
+
+    ASSERT_EQ(pool.GetSize(), 0ull);
+    size_t previousPoolSize = 0;
+    for (int iteration = 0; iteration < iterationCount; ++iteration) {
+        bool inserted = set.insert(valueDistribution(generator)).second;
+
+        ASSERT_GE(pool.GetSize(), set.size() * nodeSize);
+        if (inserted) {
+            ASSERT_GT(pool.GetSize(), previousPoolSize);
+            previousPoolSize = pool.GetSize();
+        } else {
+            ASSERT_EQ(pool.GetSize(), previousPoolSize);
+        }
+    }
+}
+
+TEST(TChunkedMemoryPoolAllocatorTest, SimpleVector)
+{
+    TChunkedMemoryPool pool;
+    TChunkedMemoryPoolAllocator<int> allocator(&pool);
+
+    using TVector = std::vector<int, TChunkedMemoryPoolAllocator<int>>;
+    TVector vector(allocator);
+
+    constexpr int seed = 0;
+    std::mt19937 generator(seed);
+    auto valueDistribution = std::uniform_int_distribution<int>(
+        -10000,
+        10000);
+
+    constexpr int iterationCount = 10000;
+
+    ASSERT_EQ(pool.GetSize(), 0ull);
+    for (int iteration = 0; iteration < iterationCount; ++iteration) {
+        vector.push_back(valueDistribution(generator));
+
+        ASSERT_GE(pool.GetSize(), vector.size());
+    }
+}
+
+TEST(TChunkedMemoryPoolAllocatorTest, DestructorCalled)
+{
+    TChunkedMemoryPool pool;
+    TChunkedMemoryPoolAllocator<int> allocator(&pool);
+
+    class TSetInDestructor
+    {
+    public:
+        explicit TSetInDestructor(bool* flag)
+            : Flag_(flag)
+        { }
+
+        ~TSetInDestructor()
+        {
+            *Flag_ = true;
+        }
+
+        bool operator<(const TSetInDestructor& other) const
+        {
+            return Flag_ < other.Flag_;
+        }
+
+    private:
+        bool* const Flag_;
+    };
+
+    using TSet = std::set<TSetInDestructor, std::less<TSetInDestructor>, TChunkedMemoryPoolAllocator<TSetInDestructor>>;
+
+    bool destructorCalled = false;
+    auto setInDestructor = TSetInDestructor{&destructorCalled};
+    {
+        auto set = std::make_unique<TSet>(allocator);
+        set->insert(setInDestructor);
+    }
+
+    ASSERT_TRUE(destructorCalled);
+}
+
+TEST(TChunkedMemoryPoolAllocatorTest, Copy)
+{
+    TChunkedMemoryPool pool;
+    TChunkedMemoryPoolAllocator<int> allocator(&pool);
+
+    using TSet = std::set<int, std::less<int>, TChunkedMemoryPoolAllocator<int>>;
+    TSet set(allocator);
+
+    constexpr int iterationCount = 1000;
+    for (int iteration = 0; iteration < iterationCount; ++iteration) {
+        set.insert(iteration);
+    }
+
+    auto poolSize = pool.GetSize();
+
+    auto copy = set;
+
+    ASSERT_EQ(std::ssize(set), iterationCount);
+    ASSERT_EQ(std::ssize(copy), iterationCount);
+    ASSERT_EQ(pool.GetSize(), poolSize * 2);
+}
+
+TEST(TChunkedMemoryPoolAllocatorTest, Move)
+{
+    TChunkedMemoryPool pool;
+    TChunkedMemoryPoolAllocator<void> allocator(&pool);
+
+    using TSet = std::set<int, std::less<int>, TChunkedMemoryPoolAllocator<int>>;
+    TSet set(allocator);
+
+    constexpr int iterationCount = 1000;
+    for (int iteration = 0; iteration < iterationCount; ++iteration) {
+        set.insert(iteration);
+    }
+
+    auto size = pool.GetSize();
+
+    auto movedTo = std::move(set);
+
+    ASSERT_EQ(std::ssize(movedTo), iterationCount);
+    ASSERT_EQ(pool.GetSize(), size);
+}
+
+TEST(TChunkedMemoryPoolAllocatorTest, NoLeak)
+{
+    TChunkedMemoryPool pool;
+    TChunkedMemoryPoolAllocator<void> allocator(&pool);
+
+    using TSet = std::set<int, std::less<int>, TChunkedMemoryPoolAllocator<int>>;
+    auto* set = pool.AllocateUninitialized<TSet>(1);
+    new(set) TSet(allocator);
+
+    constexpr int iterationCount = 1000;
+    for (int iteration = 0; iteration < iterationCount; ++iteration) {
+        set->insert(iteration);
+    }
+
+    ASSERT_GE(pool.GetSize(), sizeof(int) * iterationCount);
+}
+
+////////////////////////////////////////////////////////////////////////////////
+
+} // namespace
+} // namespace NYT
diff --git a/library/cpp/yt/memory/unittests/ya.make b/library/cpp/yt/memory/unittests/ya.make
index 54f8f4f1a5e..0f475fcf38c 100644
--- a/library/cpp/yt/memory/unittests/ya.make
+++ b/library/cpp/yt/memory/unittests/ya.make
@@ -9,6 +9,7 @@ ENDIF()
 SRCS(
     atomic_intrusive_ptr_ut.cpp
     chunked_memory_pool_ut.cpp
+    chunked_memory_pool_allocator_ut.cpp
     chunked_memory_pool_output_ut.cpp
     free_list_ut.cpp
     intrusive_ptr_ut.cpp