Merge pull request #28146 from vespa-engine/toregge/use-premmapped-areas-in-mmap-file-allocator

Use premmapped areas for smaller allocations than _small_limit.

3 files changed, 177 insertions, 21 deletions

diff --git a/vespalib/src/tests/util/mmap_file_allocator/mmap_file_allocator_test.cpp b/vespalib/src/tests/util/mmap_file_allocator/mmap_file_allocator_test.cpp
index ef16998902e..61faf4e5168 100644
--- a/vespalib/src/tests/util/mmap_file_allocator/mmap_file_allocator_test.cpp
+++ b/vespalib/src/tests/util/mmap_file_allocator/mmap_file_allocator_test.cpp
@@ -12,6 +12,7 @@ namespace {
 
 vespalib::string basedir("mmap-file-allocator-dir");
 vespalib::string hello("hello");
+vespalib::string world("world");
 
 struct MyAlloc
 {
@@ -36,7 +37,24 @@ struct MyAlloc
 
 }
 
-class MmapFileAllocatorTest : public ::testing::Test
+struct AllocatorSetup {
+    uint32_t small_limit;
+    uint32_t premmap_size;
+
+    AllocatorSetup(uint32_t small_limit_in, uint32_t premmap_size_in)
+        : small_limit(small_limit_in),
+          premmap_size(premmap_size_in)
+    {
+    }
+};
+
+std::ostream& operator<<(std::ostream& os, const AllocatorSetup setup)
+{
+    os << "small" << setup.small_limit << "premm" << setup.premmap_size;
+    return os;
+}
+
+class MmapFileAllocatorTest : public ::testing::TestWithParam<AllocatorSetup>
 {
 protected:
     MmapFileAllocator _allocator;
@@ -47,49 +65,69 @@ public:
 };
 
 MmapFileAllocatorTest::MmapFileAllocatorTest()
-    : _allocator(basedir)
+    : _allocator(basedir, GetParam().small_limit, GetParam().premmap_size)
 {
 }
 
 MmapFileAllocatorTest::~MmapFileAllocatorTest() = default;
 
-TEST_F(MmapFileAllocatorTest, zero_sized_allocation_is_handled)
+INSTANTIATE_TEST_SUITE_P(MmapFileAllocatorMultiTest,
+                         MmapFileAllocatorTest,
+                         testing::Values(AllocatorSetup(0, 1_Mi), AllocatorSetup(512, 1_Mi), AllocatorSetup(128_Ki, 1_Mi)), testing::PrintToStringParamName());
+
+
+
+TEST_P(MmapFileAllocatorTest, zero_sized_allocation_is_handled)
 {
     MyAlloc buf(_allocator, _allocator.alloc(0));
     EXPECT_EQ(nullptr, buf.data);
     EXPECT_EQ(0u, buf.size);
 }
 
-TEST_F(MmapFileAllocatorTest, mmap_file_allocator_works)
+TEST_P(MmapFileAllocatorTest, mmap_file_allocator_works)
 {
-    MyAlloc buf(_allocator, _allocator.alloc(4));
-    EXPECT_LE(4u, buf.size);
+    MyAlloc buf(_allocator, _allocator.alloc(400));
+    EXPECT_LE(400u, buf.size);
     EXPECT_TRUE(buf.data != nullptr);
     memcpy(buf.data, "1st", 4);
-    MyAlloc buf2(_allocator, _allocator.alloc(5));
+    MyAlloc buf2(_allocator, _allocator.alloc(600));
     EXPECT_LE(5u, buf2.size);
     EXPECT_TRUE(buf2.data != nullptr);
     EXPECT_TRUE(buf.data != buf2.data);
     memcpy(buf2.data, "fine", 5);
-    EXPECT_EQ(0u, _allocator.resize_inplace(buf.asPair(), 5));
-    EXPECT_EQ(0u, _allocator.resize_inplace(buf.asPair(), 3));
+    EXPECT_EQ(0u, _allocator.resize_inplace(buf.asPair(), 500));
+    EXPECT_EQ(0u, _allocator.resize_inplace(buf.asPair(), 300));
     EXPECT_NE(0u, _allocator.get_end_offset());
-    int result = msync(buf.data, buf.size, MS_SYNC);
-    EXPECT_EQ(0, result);
-    result = msync(buf2.data, buf2.size, MS_SYNC);
-    EXPECT_EQ(0, result);
+    if (GetParam().small_limit == 0) {
+        int result = msync(buf.data, buf.size, MS_SYNC);
+        EXPECT_EQ(0, result);
+        result = msync(buf2.data, buf2.size, MS_SYNC);
+        EXPECT_EQ(0, result);
+    }
 }
 
-TEST_F(MmapFileAllocatorTest, reuse_file_offset_works)
+TEST_P(MmapFileAllocatorTest, reuse_file_offset_works)
 {
+    constexpr size_t size_400 = 400;
+    constexpr size_t size_600 = 600;
+    assert(hello.size() + 1 <= size_400);
+    assert(world.size() + 1 <= size_600);
     {
-        MyAlloc buf(_allocator, _allocator.alloc(hello.size() + 1));
+        MyAlloc buf(_allocator, _allocator.alloc(size_400));
         memcpy(buf.data, hello.c_str(), hello.size() + 1);
     }
     {
-        MyAlloc buf(_allocator, _allocator.alloc(hello.size() + 1));
+        MyAlloc buf(_allocator, _allocator.alloc(size_400));
         EXPECT_EQ(0, memcmp(buf.data, hello.c_str(), hello.size() + 1));
     }
+    {
+        MyAlloc buf(_allocator, _allocator.alloc(size_600));
+        memcpy(buf.data, world.c_str(), world.size() + 1);
+    }
+    {
+        MyAlloc buf(_allocator, _allocator.alloc(size_600));
+        EXPECT_EQ(0, memcmp(buf.data, world.c_str(), world.size() + 1));
+    }
 }
 
 GTEST_MAIN_RUN_ALL_TESTS()
diff --git a/vespalib/src/vespa/vespalib/util/mmap_file_allocator.cpp b/vespalib/src/vespa/vespalib/util/mmap_file_allocator.cpp
index 3d0b9debeda..f1ea0034474 100644
--- a/vespalib/src/vespa/vespalib/util/mmap_file_allocator.cpp
+++ b/vespalib/src/vespa/vespalib/util/mmap_file_allocator.cpp
@@ -16,11 +16,21 @@ namespace fs = std::filesystem;
 namespace vespalib::alloc {
 
 MmapFileAllocator::MmapFileAllocator(const vespalib::string& dir_name)
+    : MmapFileAllocator(dir_name, default_small_limit, default_premmap_size)
+{
+}
+
+MmapFileAllocator::MmapFileAllocator(const vespalib::string& dir_name, uint32_t small_limit, uint32_t premmap_size)
     : _dir_name(dir_name),
+      _small_limit(small_limit),
+      _premmap_size(premmap_size),
       _file(_dir_name + "/swapfile"),
       _end_offset(0),
       _allocations(),
-      _freelist()
+      _freelist(),
+      _small_allocations(),
+      _small_freelist(),
+      _premmapped_areas()
 {
     fs::create_directories(fs::path(_dir_name));
     _file.open(O_RDWR | O_CREAT | O_TRUNC, false);
@@ -28,6 +38,20 @@ MmapFileAllocator::MmapFileAllocator(const vespalib::string& dir_name)
 
 MmapFileAllocator::~MmapFileAllocator()
 {
+    assert(_small_allocations.empty());
+    assert(_allocations.size() == _premmapped_areas.size());
+    for (auto& area : _premmapped_areas) {
+        auto offset = area.first;
+        auto ptr = area.second;
+        auto itr = _allocations.find(ptr);
+        assert(itr != _allocations.end());
+        assert(itr->first == ptr);
+        assert(itr->second.offset == offset);
+        auto size = itr->second.size;
+        _small_freelist.remove_premmapped_area(offset, size);
+        free_large({ptr, size});
+    }
+    _premmapped_areas.clear();
     assert(_allocations.empty());
     _file.close();
     _file.unlink();
@@ -53,6 +77,18 @@ MmapFileAllocator::alloc(size_t sz) const
     if (sz == 0) {
         return PtrAndSize(); // empty allocation
     }
+    static constexpr size_t alignment = 64;
+    sz = (sz + alignment - 1) & -alignment; // round sz to a multiple of alignment
+    if (sz >= _small_limit) {
+        return alloc_large(sz);
+    } else {
+        return alloc_small(sz);
+    }
+}
+
+PtrAndSize
+MmapFileAllocator::alloc_large(size_t sz) const
+{
     sz = round_up_to_page_size(sz);
     uint64_t offset = alloc_area(sz);
     void *buf = mmap(nullptr, sz, PROT_READ | PROT_WRITE, MAP_SHARED, _file.getFileDescriptor(), offset);
@@ -74,6 +110,45 @@ MmapFileAllocator::alloc(size_t sz) const
     return PtrAndSize(buf, sz);
 }
 
+void*
+MmapFileAllocator::map_premapped_offset_to_ptr(uint64_t offset, size_t size) const
+{
+    auto itr = _premmapped_areas.lower_bound(offset);
+    if (itr == _premmapped_areas.end() || itr->first > offset) {
+        assert(itr != _premmapped_areas.begin());
+        --itr;
+    }
+    auto aitr = _allocations.find(itr->second);
+    assert(aitr != _allocations.end());
+    assert(aitr->first == itr->second);
+    assert(offset >= aitr->second.offset);
+    assert(offset + size <= aitr->second.offset + aitr->second.size);
+    return static_cast<char*>(itr->second) + (offset - aitr->second.offset);
+}
+
+PtrAndSize
+MmapFileAllocator::alloc_small(size_t sz) const
+{
+    uint64_t offset = _small_freelist.alloc(sz);
+    if (offset == FileAreaFreeList::bad_offset) {
+        auto new_premmap = alloc_large(_premmap_size);
+        assert(new_premmap.size() >= _premmap_size);
+        auto itr = _allocations.find(new_premmap.get());
+        assert(itr != _allocations.end());
+        assert(itr->first == new_premmap.get());
+        _small_freelist.add_premmapped_area(itr->second.offset, itr->second.size);
+        auto ins_res = _premmapped_areas.emplace(itr->second.offset, new_premmap.get());
+        assert(ins_res.second);
+        offset = _small_freelist.alloc(sz);
+        assert(offset != FileAreaFreeList::bad_offset);
+    }
+    auto ptr = map_premapped_offset_to_ptr(offset, sz);
+    // Register allocation
+    auto ins_res = _small_allocations.insert(std::make_pair(ptr, SizeAndOffset(sz, offset)));
+    assert(ins_res.second);
+    return {ptr, sz};
+}
+
 void
 MmapFileAllocator::free(PtrAndSize alloc) const noexcept
 {
@@ -82,13 +157,30 @@ MmapFileAllocator::free(PtrAndSize alloc) const noexcept
         return; // empty allocation
     }
     assert(alloc.get() != nullptr);
+    if (alloc.size() >= _small_limit) {
+        free_large(alloc);
+    } else {
+        free_small(alloc);
+    }
+}
+
+uint64_t
+MmapFileAllocator::remove_allocation(PtrAndSize alloc, Allocations& allocations) const noexcept
+{
     // Check that matching allocation is registered
-    auto itr = _allocations.find(alloc.get());
-    assert(itr != _allocations.end());
+    auto itr = allocations.find(alloc.get());
+    assert(itr != allocations.end());
     assert(itr->first == alloc.get());
     assert(itr->second.size == alloc.size());
     auto offset = itr->second.offset;
-    _allocations.erase(itr);
+    allocations.erase(itr);
+    return offset;
+}
+
+void
+MmapFileAllocator::free_large(PtrAndSize alloc) const noexcept
+{
+    auto offset = remove_allocation(alloc, _allocations);
     int retval = madvise(alloc.get(), alloc.size(), MADV_DONTNEED);
     assert(retval == 0);
     retval = munmap(alloc.get(), alloc.size());
@@ -96,6 +188,13 @@ MmapFileAllocator::free(PtrAndSize alloc) const noexcept
     _freelist.free(offset, alloc.size());
 }
 
+void
+MmapFileAllocator::free_small(PtrAndSize alloc) const noexcept
+{
+    auto offset = remove_allocation(alloc, _small_allocations);
+    _small_freelist.free(offset, alloc.size());
+}
+
 size_t
 MmapFileAllocator::resize_inplace(PtrAndSize, size_t) const
 {
diff --git a/vespalib/src/vespa/vespalib/util/mmap_file_allocator.h b/vespalib/src/vespa/vespalib/util/mmap_file_allocator.h
index 883c7e49848..c79dc8682ba 100644
--- a/vespalib/src/vespa/vespalib/util/mmap_file_allocator.h
+++ b/vespalib/src/vespa/vespalib/util/mmap_file_allocator.h
@@ -7,6 +7,7 @@
 #include <vespa/vespalib/io/fileutil.h>
 #include <vespa/vespalib/stllike/hash_map.h>
 #include <vespa/vespalib/stllike/string.h>
+#include <vespa/vespalib/util/size_literals.h>
 #include <map>
 
 namespace vespalib::alloc {
@@ -15,6 +16,9 @@ namespace vespalib::alloc {
  * Class handling memory allocations backed by one or more files.
  * Not reentrant or thread safe. Should not be destructed before all allocations
  * have been freed.
+ *
+ * Memory allocations smaller than _small_limit use portions of
+ * premmapped areas to reduce the total number of memory mappings.
  */
 class MmapFileAllocator : public MemoryAllocator {
     struct SizeAndOffset {
@@ -30,14 +34,29 @@ class MmapFileAllocator : public MemoryAllocator {
         {
         }
     };
+    using Allocations = hash_map<void *, SizeAndOffset>;
     const vespalib::string _dir_name;
+    const uint32_t   _small_limit;
+    const uint32_t   _premmap_size;
     mutable File     _file;
     mutable uint64_t _end_offset;
-    mutable hash_map<void *, SizeAndOffset> _allocations;
+    mutable Allocations _allocations;
     mutable FileAreaFreeList _freelist;
+    mutable Allocations _small_allocations;
+    mutable FileAreaFreeList _small_freelist;
+    mutable std::map<uint64_t, void*> _premmapped_areas;
     uint64_t alloc_area(size_t sz) const;
+    PtrAndSize alloc_large(size_t size) const;
+    PtrAndSize alloc_small(size_t size) const;
+    void free_large(PtrAndSize alloc) const noexcept;
+    void free_small(PtrAndSize alloc) const noexcept;
+    void* map_premapped_offset_to_ptr(uint64_t offset, size_t size) const;
+    uint64_t remove_allocation(PtrAndSize alloc, Allocations& allocations) const noexcept;
 public:
+    static constexpr uint32_t default_small_limit =  0;
+    static constexpr uint32_t default_premmap_size = 1_Mi;
     MmapFileAllocator(const vespalib::string& dir_name);
+    MmapFileAllocator(const vespalib::string& dir_name, uint32_t small_limit, uint32_t premmap_size);
     ~MmapFileAllocator();
     PtrAndSize alloc(size_t sz) const override;
     void free(PtrAndSize alloc) const noexcept override;