Improve buffer allocation strategy in data store by matching underlying allocators.

This should reduce the amount of memory wasted in allocations.
  1) heap allocation: buffer size is power of 2 to match vespamalloc.
  2) mmap allocation: buffer size is multiple of huge page size (2MB) to match mmap allocator.

7 files changed, 95 insertions, 42 deletions

diff --git a/searchlib/src/tests/attribute/enumstore/enumstore_test.cpp b/searchlib/src/tests/attribute/enumstore/enumstore_test.cpp
index fa59baa2bc5..daff432d68d 100644
--- a/searchlib/src/tests/attribute/enumstore/enumstore_test.cpp
+++ b/searchlib/src/tests/attribute/enumstore/enumstore_test.cpp
@@ -371,8 +371,10 @@ EnumStoreTest::testCompaction(bool hasPostings, bool disableReEnumerate)
 {
     // entrySize = 15 before alignment
     uint32_t entrySize = EnumStoreType::alignEntrySize(15);
-    uint32_t bufferSize = entrySize * 5;
-    EnumStoreType ses(bufferSize, hasPostings);
+    uint32_t initBufferSize = entrySize * 5;
+    EnumStoreType ses(initBufferSize, hasPostings);
+    // Note: Sizes of underlying data store buffers are power of 2.
+    uint32_t adjustedBufferSize = vespalib::roundUp2inN(initBufferSize) - RESERVED_BYTES;
     EnumIndex idx;
     std::vector<EnumIndex> indices;
     typename EnumStoreType::Type t = "foo";
@@ -385,18 +387,19 @@ EnumStoreTest::testCompaction(bool hasPostings, bool disableReEnumerate)
 
     // fill with unique values
     for (uint32_t i = 0; i < 5; ++i) {
-        EXPECT_TRUE(ses.getRemaining() == bufferSize - i * entrySize);
+        size_t expRemaining = adjustedBufferSize - i * entrySize;
+        EXPECT_EQUAL(expRemaining, ses.getRemaining());
         ses.addEnum(uniques[i].c_str(), idx);
         ses.incRefCount(idx);
         EXPECT_TRUE(ses.getRefCount(idx));
         indices.push_back(idx);
     }
-    EXPECT_EQUAL(0u, ses.getRemaining());
-    EXPECT_EQUAL(0u, ses.getBuffer(0).remaining());
+    EXPECT_EQUAL(32u, ses.getRemaining());
+    EXPECT_EQUAL(32u, ses.getBuffer(0).remaining());
     EXPECT_EQUAL(entrySize * 5 + RESERVED_BYTES, ses.getBuffer(0).size());
     EXPECT_EQUAL(RESERVED_BYTES, ses.getBuffer(0).getDeadElems());
     uint32_t failEntrySize = ses.getEntrySize("enum05");
-    EXPECT_TRUE(failEntrySize > ses.getRemaining());
+    EXPECT_EQUAL(16u, failEntrySize);
 
     // change from enum00 -> enum01
     ses.decRefCount(indices[0]);
@@ -525,7 +528,9 @@ EnumStoreTest::testReset(bool hasPostings)
     }
 
     ses.reset(builder);
-    EXPECT_EQUAL(RESERVED_BYTES, ses.getRemaining());
+    // Note: Sizes of underlying data store buffers are power of 2.
+    EXPECT_EQUAL(524288u, ses.getCapacity());
+    EXPECT_EQUAL(204272u, ses.getRemaining());
 
     // check for old unique strings
     for (StringVector::iterator iter = uniques.begin(); iter != uniques.end(); ++iter) {
@@ -597,7 +602,8 @@ EnumStoreTest::testHoldListAndGeneration()
         }
     }
 
-    EXPECT_EQUAL(0u, ses.getRemaining());
+    // Note: Sizes of underlying data store buffers are power of 2.
+    EXPECT_EQUAL(432u, ses.getRemaining());
     EXPECT_EQUAL(RESERVED_BYTES, ses.getBuffer(0).getDeadElems());
 
     // remove all uniques
@@ -657,7 +663,8 @@ EnumStoreTest::testMemoryUsage()
     // usage before inserting enums
     MemoryUsage usage = ses.getMemoryUsage();
     EXPECT_EQUAL(ses.getNumUniques(), uint32_t(0));
-    EXPECT_EQUAL(enumStoreAlign(200u) + RESERVED_BYTES, usage.allocatedBytes());
+    // Note: Sizes of underlying data store buffers are power of 2.
+    EXPECT_EQUAL(vespalib::roundUp2inN(enumStoreAlign(200u) + RESERVED_BYTES), usage.allocatedBytes());
     EXPECT_EQUAL(RESERVED_BYTES, usage.usedBytes());
     EXPECT_EQUAL(RESERVED_BYTES, usage.deadBytes());
     EXPECT_EQUAL(0u, usage.allocatedBytesOnHold());
@@ -672,7 +679,8 @@ EnumStoreTest::testMemoryUsage()
     // usage after inserting enums
     usage = ses.getMemoryUsage();
     EXPECT_EQUAL(ses.getNumUniques(), num);
-    EXPECT_EQUAL(enumStoreAlign(200u) + RESERVED_BYTES, usage.allocatedBytes());
+    // Note: Sizes of underlying data store buffers are power of 2.
+    EXPECT_EQUAL(vespalib::roundUp2inN(enumStoreAlign(200u) + RESERVED_BYTES), usage.allocatedBytes());
     EXPECT_EQUAL(num * entrySize + RESERVED_BYTES, usage.usedBytes());
     EXPECT_EQUAL(RESERVED_BYTES, usage.deadBytes());
     EXPECT_EQUAL(0u, usage.allocatedBytesOnHold());
@@ -689,7 +697,8 @@ EnumStoreTest::testMemoryUsage()
     // usage after removing enums
     usage = ses.getMemoryUsage();
     EXPECT_EQUAL(ses.getNumUniques(), num / 2);
-    EXPECT_EQUAL(enumStoreAlign(200u) + RESERVED_BYTES, usage.allocatedBytes());
+    // Note: Sizes of underlying data store buffers are power of 2.
+    EXPECT_EQUAL(vespalib::roundUp2inN(enumStoreAlign(200u) + RESERVED_BYTES), usage.allocatedBytes());
     EXPECT_EQUAL(num * entrySize + RESERVED_BYTES, usage.usedBytes());
     EXPECT_EQUAL((num / 2) * entrySize + RESERVED_BYTES, usage.deadBytes());
     EXPECT_EQUAL(0u, usage.allocatedBytesOnHold());
diff --git a/searchlib/src/tests/btree/btree_test.cpp b/searchlib/src/tests/btree/btree_test.cpp
index 5795385250f..1f39c7315e8 100644
--- a/searchlib/src/tests/btree/btree_test.cpp
+++ b/searchlib/src/tests/btree/btree_test.cpp
@@ -1022,6 +1022,15 @@ Test::requireThatTreeIteratorAssignWorks()
     }
 }
 
+size_t
+adjustAllocatedBytes(size_t nodeCount, size_t nodeSize)
+{
+    // Note: Sizes of underlying data store buffers are power of 2.
+    size_t allocatedBytes = vespalib::roundUp2inN(nodeCount * nodeSize);
+    size_t adjustedNodeCount = allocatedBytes / nodeSize;
+    return adjustedNodeCount * nodeSize;
+}
+
 void
 Test::requireThatMemoryUsageIsCalculated()
 {
@@ -1041,8 +1050,8 @@ Test::requireThatMemoryUsageIsCalculated()
     MemoryUsage mu;
     const uint32_t initialInternalNodes = 128u;
     const uint32_t initialLeafNodes = 128u;
-    mu.incAllocatedBytes(sizeof(INode) * initialInternalNodes);
-    mu.incAllocatedBytes(sizeof(LNode) * initialLeafNodes);
+    mu.incAllocatedBytes(adjustAllocatedBytes(initialInternalNodes, sizeof(INode)));
+    mu.incAllocatedBytes(adjustAllocatedBytes(initialLeafNodes, sizeof(LNode)));
     mu.incUsedBytes(sizeof(INode));
     mu.incDeadBytes(sizeof(INode));
     EXPECT_TRUE(assertMemoryUsage(mu, tm.getMemoryUsage()));
@@ -1071,8 +1080,8 @@ Test::requireThatMemoryUsageIsCalculated()
     gh.incGeneration();
     tm.trimHoldLists(gh.getFirstUsedGeneration());
     mu = MemoryUsage();
-    mu.incAllocatedBytes(sizeof(INode) * initialInternalNodes);
-    mu.incAllocatedBytes(sizeof(LNode) * initialLeafNodes);
+    mu.incAllocatedBytes(adjustAllocatedBytes(initialInternalNodes, sizeof(INode)));
+    mu.incAllocatedBytes(adjustAllocatedBytes(initialLeafNodes, sizeof(LNode)));
     mu.incUsedBytes(sizeof(INode) * 2);
     mu.incDeadBytes(sizeof(INode) * 2);
     mu.incUsedBytes(sizeof(LNode));
diff --git a/searchlib/src/tests/datastore/array_store/array_store_test.cpp b/searchlib/src/tests/datastore/array_store/array_store_test.cpp
index fff4445890b..dab853305c6 100644
--- a/searchlib/src/tests/datastore/array_store/array_store_test.cpp
+++ b/searchlib/src/tests/datastore/array_store/array_store_test.cpp
@@ -316,7 +316,7 @@ TEST_F("require that used, onHold and dead memory usage is tracked for large arr
 TEST_F("require that address space usage is ratio between used clusters and number of possible clusters", NumberFixture(3))
 {
     f.add({2,2});
-    f.add({4,4,4});
+    f.add({3,3,3});
     // 1 cluster is reserved (buffer 0, offset 0).
     EXPECT_EQUAL(3u, f.store.addressSpaceUsage().used());
     EXPECT_EQUAL(1u, f.store.addressSpaceUsage().dead());
@@ -324,11 +324,15 @@ TEST_F("require that address space usage is ratio between used clusters and numb
     /*
      * Expected limit is sum of allocated clusters for active buffers and
      * potentially allocated clusters for free buffers. If all buffers were
-     * free then the limit would be 4 Gi.  Then we subtract clusters for 4
-     * buffers that are not free, and add their actual number of allocated
-     * clusters (16 clusters per buffer).
+     * free then the limit would be 4 Gi.
+     * Then we subtract clusters for 4 buffers that are not free (arraySize=1,2,3 + largeArray),
+     * and add their actual number of allocated clusters (16 clusters per buffer).
+     * Note: arraySize=3 has 21 clusters as allocated buffer is rounded up to power of 2:
+     *   16 * 3 * sizeof(int) = 192 -> 256.
+     *   allocated elements = 256 / sizeof(int) = 64.
+     *   limit = 64 / 3 = 21.
      */
-    size_t expLimit = fourgig - 4 * F1::EntryRefType::offsetSize() + 4 * 16;
+    size_t expLimit = fourgig - 4 * F1::EntryRefType::offsetSize() + 3 * 16 + 21;
     EXPECT_EQUAL(static_cast<double>(2)/ expLimit, f.store.addressSpaceUsage().usage());
     EXPECT_EQUAL(expLimit, f.store.addressSpaceUsage().limit());
 }
diff --git a/searchlib/src/tests/datastore/datastore/datastore_test.cpp b/searchlib/src/tests/datastore/datastore/datastore_test.cpp
index 2463439c47c..c3de2261745 100644
--- a/searchlib/src/tests/datastore/datastore/datastore_test.cpp
+++ b/searchlib/src/tests/datastore/datastore/datastore_test.cpp
@@ -11,6 +11,8 @@ LOG_SETUP("datastore_test");
 namespace search {
 namespace datastore {
 
+using vespalib::alloc::MemoryAllocator;
+
 struct IntReclaimer
 {
     static void reclaim(int *) {}
@@ -65,21 +67,22 @@ public:
 
 using GrowthStats = std::vector<int>;
 
-constexpr float ALLOC_GROW_FACTOR = 0.5;
+constexpr float ALLOC_GROW_FACTOR = 0.4;
+constexpr size_t HUGE_PAGE_CLUSTER_SIZE = (MemoryAllocator::HUGEPAGE_SIZE / sizeof(int));
 
 class GrowStore
 {
-    using Store = DataStoreT<EntryRefT<22>>;
+    using Store = DataStoreT<EntryRefT<24>>;
     using RefType = Store::RefType;
     Store _store;
     BufferType<int> _firstType;
     BufferType<int> _type;
     uint32_t _typeId;
 public:
-    GrowStore(size_t minSize, size_t minSwitch)
+    GrowStore(size_t minClusters, size_t maxClusters, size_t numClustersForNewBuffer)
         : _store(),
-          _firstType(1, 1, 64, 0, ALLOC_GROW_FACTOR),
-          _type(1, minSize, 64, minSwitch, ALLOC_GROW_FACTOR),
+          _firstType(1, 1, maxClusters, 0, ALLOC_GROW_FACTOR),
+          _type(1, minClusters, maxClusters, numClustersForNewBuffer, ALLOC_GROW_FACTOR),
           _typeId(0)
     {
         (void) _store.addType(&_firstType);
@@ -460,11 +463,11 @@ namespace {
 void assertGrowStats(GrowthStats expSizes,
                      GrowthStats expFirstBufSizes,
                      size_t expInitMemUsage,
-                     size_t minSize, size_t minSwitch)
+                     size_t minClusters, size_t numClustersForNewBuffer, size_t maxClusters = 128)
 {
-    EXPECT_EQUAL(expSizes, GrowStore(minSize, minSwitch).getGrowthStats(expSizes.size()));
-    EXPECT_EQUAL(expFirstBufSizes, GrowStore(minSize, minSwitch).getFirstBufGrowStats());
-    EXPECT_EQUAL(expInitMemUsage, GrowStore(minSize, minSwitch).getMemoryUsage().allocatedBytes());
+    EXPECT_EQUAL(expSizes, GrowStore(minClusters, maxClusters, numClustersForNewBuffer).getGrowthStats(expSizes.size()));
+    EXPECT_EQUAL(expFirstBufSizes, GrowStore(minClusters, maxClusters, numClustersForNewBuffer).getFirstBufGrowStats());
+    EXPECT_EQUAL(expInitMemUsage, GrowStore(minClusters, maxClusters, numClustersForNewBuffer).getMemoryUsage().allocatedBytes());
 }
 
 }
@@ -472,23 +475,29 @@ void assertGrowStats(GrowthStats expSizes,
 TEST("require that buffer growth works")
 {
     // Always switch to new buffer, min size 4
-    TEST_DO(assertGrowStats({ 4, 4, 4, 6, 9, 13, 20, 30, 45, 64 },
+    TEST_DO(assertGrowStats({ 4, 4, 4, 4, 8, 16, 16, 32, 64, 64 },
                             { 4 }, 20, 4, 0));
     // Resize if buffer size is less than 4, min size 0
-    TEST_DO(assertGrowStats({ 3, 3, 3, 4, 6, 9, 14, 21, 31, 47 },
-                            { 0, 1, 2, 3 }, 4, 0, 4));
+    TEST_DO(assertGrowStats({ 4, 4, 4, 4, 8, 16, 16, 32, 64, 64 },
+                            { 0, 1, 2, 4 }, 4, 0, 4));
     // Always switch to new buffer, min size 16
-    TEST_DO(assertGrowStats({ 16, 16, 16, 24, 36, 54, 64, 64, 64 },
+    TEST_DO(assertGrowStats({ 16, 16, 16, 32, 32, 64, 128, 128, 128 },
                             { 16 }, 68, 16, 0));
     // Resize if buffer size is less than 16, min size 0
-    TEST_DO(assertGrowStats({ 19, 19, 19, 28, 42, 63, 64, 64, 64 },
-                            { 0, 1, 2, 3, 4, 6, 9, 13, 19 }, 4, 0, 16));
+    TEST_DO(assertGrowStats({ 16, 16, 16, 32, 32, 64, 128, 128, 128 },
+                            { 0, 1, 2, 4, 8, 16 }, 4, 0, 16));
     // Resize if buffer size is less than 16, min size 4
-    TEST_DO(assertGrowStats({ 19, 19, 19, 28, 42, 63, 64, 64, 64 },
-                            { 4, 6, 9, 13, 19 }, 20, 4, 16));
+    TEST_DO(assertGrowStats({ 16, 16, 16, 32, 32, 64, 128, 128, 128 },
+                            { 4, 8, 16 }, 20, 4, 16));
     // Always switch to new buffer, min size 0
-    TEST_DO(assertGrowStats({ 1, 1, 1, 1, 2, 3, 4, 6, 9 },
+    TEST_DO(assertGrowStats({ 1, 1, 1, 1, 1, 2, 2, 4, 8, 8, 16, 32 },
                             { 0, 1 }, 4, 0, 0));
+
+    // Buffers with sizes larger than the huge page size of the mmap allocator.
+    ASSERT_EQUAL(524288u, HUGE_PAGE_CLUSTER_SIZE);
+    TEST_DO(assertGrowStats({ 262144, 262144, 262144, 524288, 524288, 524288 * 2, 524288 * 3, 524288 * 4, 524288 * 5, 524288 * 5 },
+                            { 0, 1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096, 8192, 16384, 32768, 65536, 131072, 262144 },
+                            4, 0, HUGE_PAGE_CLUSTER_SIZE / 2, HUGE_PAGE_CLUSTER_SIZE * 5));
 }
 
 }
diff --git a/searchlib/src/tests/memoryindex/memoryindex/memoryindex_test.cpp b/searchlib/src/tests/memoryindex/memoryindex/memoryindex_test.cpp
index 77a687796b3..9de6ac9f310 100644
--- a/searchlib/src/tests/memoryindex/memoryindex/memoryindex_test.cpp
+++ b/searchlib/src/tests/memoryindex/memoryindex/memoryindex_test.cpp
@@ -381,7 +381,7 @@ TEST("requireThatNumDocsAndDocIdLimitIsReturned")
 
 TEST("requireThatWeUnderstandTheMemoryFootprint")
 {
-    constexpr size_t BASE_SIZE = 118860u;
+    constexpr size_t BASE_SIZE = 188172u;
     {
         Setup setup;
         Index index(setup);
diff --git a/searchlib/src/vespa/searchlib/attribute/enumstorebase.h b/searchlib/src/vespa/searchlib/attribute/enumstorebase.h
index 8cb5fe596b7..f74345a8806 100644
--- a/searchlib/src/vespa/searchlib/attribute/enumstorebase.h
+++ b/searchlib/src/vespa/searchlib/attribute/enumstorebase.h
@@ -288,6 +288,9 @@ public:
     uint32_t getRemaining() const {
         return _store.getBufferState(_store.getActiveBufferId(TYPE_ID)).remaining();
     }
+    uint32_t getCapacity() const {
+        return _store.getBufferState(_store.getActiveBufferId(TYPE_ID)).capacity();
+    }
     MemoryUsage getMemoryUsage() const;
     MemoryUsage getTreeMemoryUsage() const { return _enumDict->getTreeMemoryUsage(); }
 
diff --git a/searchlib/src/vespa/searchlib/datastore/bufferstate.cpp b/searchlib/src/vespa/searchlib/datastore/bufferstate.cpp
index c4ef2124189..9bb6fee7b79 100644
--- a/searchlib/src/vespa/searchlib/datastore/bufferstate.cpp
+++ b/searchlib/src/vespa/searchlib/datastore/bufferstate.cpp
@@ -4,6 +4,7 @@
 #include <limits>
 
 using vespalib::alloc::Alloc;
+using vespalib::alloc::MemoryAllocator;
 
 namespace search::datastore {
 
@@ -30,7 +31,7 @@ BufferState::BufferState()
       _typeId(0),
       _clusterSize(0),
       _compacting(false),
-      _buffer(Alloc::alloc())
+      _buffer(Alloc::alloc(0, MemoryAllocator::HUGEPAGE_SIZE))
 {
 }
 
@@ -53,6 +54,23 @@ struct AllocResult {
     AllocResult(size_t elements_, size_t bytes_) : elements(elements_), bytes(bytes_) {}
 };
 
+size_t
+roundUpToMatchAllocator(size_t sz)
+{
+    if (sz == 0) {
+        return 0;
+    }
+    // We round up the wanted number of bytes to allocate to match
+    // the underlying allocator to ensure little to no waste of allocated memory.
+    if (sz < MemoryAllocator::HUGEPAGE_SIZE) {
+        // Match heap allocator in vespamalloc.
+        return vespalib::roundUp2inN(sz);
+    } else {
+        // Match mmap allocator.
+        return MemoryAllocator::roundUpToHugePages(sz);
+    }
+}
+
 AllocResult
 calcAllocation(uint32_t bufferId,
                BufferTypeBase &typeHandler,
@@ -61,8 +79,9 @@ calcAllocation(uint32_t bufferId,
 {
     size_t allocClusters = typeHandler.calcClustersToAlloc(bufferId, elementsNeeded, resizing);
     size_t allocElements = allocClusters * typeHandler.getClusterSize();
-    size_t allocBytes = allocElements * typeHandler.elementSize();
-    return AllocResult(allocElements, allocBytes);
+    size_t allocBytes = roundUpToMatchAllocator(allocElements * typeHandler.elementSize());
+    size_t adjustedAllocElements = (allocBytes / typeHandler.elementSize());
+    return AllocResult(adjustedAllocElements, allocBytes);
 }
 
 }