Merge pull request #7896 from vespa-engine/geirst/fix-bug-in-datastore-memory-allocation

Fix bug where allocation of too large buffers in data store caused of…

4 files changed, 81 insertions, 18 deletions

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 dab853305c6..f5d76852888 100644
--- a/searchlib/src/tests/datastore/array_store/array_store_test.cpp
+++ b/searchlib/src/tests/datastore/array_store/array_store_test.cpp
@@ -36,6 +36,11 @@ struct Fixture
           refStore(),
           generation(1)
     {}
+    Fixture(const ArrayStoreConfig &storeCfg)
+        : store(storeCfg),
+          refStore(),
+          generation(1)
+    {}
     void assertAdd(const EntryVector &input) {
         EntryRef ref = add(input);
         assertGet(ref, input);
@@ -112,6 +117,9 @@ struct Fixture
 using NumberFixture = Fixture<uint32_t>;
 using StringFixture = Fixture<std::string>;
 using SmallOffsetNumberFixture = Fixture<uint32_t, EntryRefT<10>>;
+using ByteFixture = Fixture<uint8_t>;
+
+
 
 TEST("require that we test with trivial and non-trivial types")
 {
@@ -337,4 +345,16 @@ TEST_F("require that address space usage is ratio between used clusters and numb
     EXPECT_EQUAL(expLimit, f.store.addressSpaceUsage().limit());
 }
 
+TEST_F("require that offset in EntryRefT is within bounds when allocating memory buffers where wanted number of bytes is not a power of 2 and less than huge page size",
+       ByteFixture(ByteFixture::ArrayStoreType::optimizedConfigForHugePage(1023, vespalib::alloc::MemoryAllocator::HUGEPAGE_SIZE,
+                                                                           4 * 1024, 8 * 1024, ALLOC_GROW_FACTOR)))
+{
+    // The array store config used in this test is equivalent to the one multi-value attribute uses when initializing multi-value mapping.
+    // See similar test in datastore_test.cpp for more details on what happens during memory allocation.
+    for (size_t i = 0; i < 1000000; ++i) {
+        f.add({1, 2, 3});
+    }
+    f.assertStoreContent();
+}
+
 TEST_MAIN() { TEST_RUN_ALL(); }
diff --git a/searchlib/src/tests/datastore/datastore/datastore_test.cpp b/searchlib/src/tests/datastore/datastore/datastore_test.cpp
index c3de2261745..b312d2bfe55 100644
--- a/searchlib/src/tests/datastore/datastore/datastore_test.cpp
+++ b/searchlib/src/tests/datastore/datastore/datastore_test.cpp
@@ -70,19 +70,19 @@ using GrowthStats = std::vector<int>;
 constexpr float ALLOC_GROW_FACTOR = 0.4;
 constexpr size_t HUGE_PAGE_CLUSTER_SIZE = (MemoryAllocator::HUGEPAGE_SIZE / sizeof(int));
 
+template <typename DataType, typename RefType>
 class GrowStore
 {
-    using Store = DataStoreT<EntryRefT<24>>;
-    using RefType = Store::RefType;
+    using Store = DataStoreT<RefType>;
     Store _store;
-    BufferType<int> _firstType;
-    BufferType<int> _type;
+    BufferType<DataType> _firstType;
+    BufferType<DataType> _type;
     uint32_t _typeId;
 public:
-    GrowStore(size_t minClusters, size_t maxClusters, size_t numClustersForNewBuffer)
+    GrowStore(size_t clusterSize, size_t minClusters, size_t maxClusters, size_t numClustersForNewBuffer)
         : _store(),
           _firstType(1, 1, maxClusters, 0, ALLOC_GROW_FACTOR),
-          _type(1, minClusters, maxClusters, numClustersForNewBuffer, ALLOC_GROW_FACTOR),
+          _type(clusterSize, minClusters, maxClusters, numClustersForNewBuffer, ALLOC_GROW_FACTOR),
           _typeId(0)
     {
         (void) _store.addType(&_firstType);
@@ -93,20 +93,19 @@ public:
 
     GrowthStats getGrowthStats(size_t bufs) {
         GrowthStats sizes;
-        int i = 0;
-        int previ = 0;
         int prevBufferId = -1;
         while (sizes.size() < bufs) {
-            RefType iRef(_store.allocator<int>(_typeId).alloc().ref);
+            RefType iRef = (_type.getClusterSize() == 1) ?
+                           (_store.template allocator<DataType>(_typeId).alloc().ref) :
+                           (_store.template allocator<DataType>(_typeId).allocArray(_type.getClusterSize()).ref);
             int bufferId = iRef.bufferId();
             if (bufferId != prevBufferId) {
                 if (prevBufferId >= 0) {
-                    sizes.push_back(i - previ);
-                    previ = i;
+                    const auto &state = _store.getBufferState(prevBufferId);
+                    sizes.push_back(state.capacity());
                 }
                 prevBufferId = bufferId;
             }
-            ++i;
         }
         return sizes;
     }
@@ -116,7 +115,7 @@ public:
         int prevBuffer = -1;
         size_t prevAllocated = _store.getMemoryUsage().allocatedBytes();
         for (;;) {
-            RefType iRef = _store.allocator<int>(_typeId).alloc().ref;
+            RefType iRef = _store.template allocator<DataType>(_typeId).alloc().ref;
             size_t allocated = _store.getMemoryUsage().allocatedBytes();
             if (allocated != prevAllocated) {
                 sizes.push_back(i);
@@ -458,6 +457,8 @@ TEST("require that we can disable elemement hold list")
     s.trimHoldLists(101);
 }
 
+using IntGrowStore = GrowStore<int, EntryRefT<24>>;
+
 namespace {
 
 void assertGrowStats(GrowthStats expSizes,
@@ -465,9 +466,9 @@ void assertGrowStats(GrowthStats expSizes,
                      size_t expInitMemUsage,
                      size_t minClusters, size_t numClustersForNewBuffer, size_t maxClusters = 128)
 {
-    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());
+    EXPECT_EQUAL(expSizes, IntGrowStore(1, minClusters, maxClusters, numClustersForNewBuffer).getGrowthStats(expSizes.size()));
+    EXPECT_EQUAL(expFirstBufSizes, IntGrowStore(1, minClusters, maxClusters, numClustersForNewBuffer).getFirstBufGrowStats());
+    EXPECT_EQUAL(expInitMemUsage, IntGrowStore(1, minClusters, maxClusters, numClustersForNewBuffer).getMemoryUsage().allocatedBytes());
 }
 
 }
@@ -500,6 +501,43 @@ TEST("require that buffer growth works")
                             4, 0, HUGE_PAGE_CLUSTER_SIZE / 2, HUGE_PAGE_CLUSTER_SIZE * 5));
 }
 
+using RefType15 = EntryRefT<15>; // offsetSize=32768
+
+namespace {
+
+template <typename DataType>
+void assertGrowStats(GrowthStats expSizes, uint32_t clusterSize)
+{
+    uint32_t minClusters = 2048;
+    uint32_t maxClusters = RefType15::offsetSize();
+    uint32_t numClustersForNewBuffer = 2048;
+    GrowStore<DataType, RefType15> store(clusterSize, minClusters, maxClusters, numClustersForNewBuffer);
+    EXPECT_EQUAL(expSizes, store.getGrowthStats(expSizes.size()));
+}
+
+}
+
+TEST("require that offset in EntryRefT is within bounds when allocating memory buffers where wanted number of bytes is not a power of 2 and less than huge page size")
+{
+    /*
+     * When allocating new memory buffers for the data store the following happens (ref. calcAllocation() in bufferstate.cpp):
+     *   1) Calculate how many clusters to alloc.
+     *      In this case we alloc a minimum of 2048 and a maximum of 32768.
+     *   2) Calculate how many bytes to alloc: clustersToAlloc * clusterSize * elementSize.
+     *      In this case elementSize is (1 or 4) and clusterSize varies (3, 5, 7).
+     *   3) Round up bytes to alloc to match the underlying allocator (power of 2 if less than huge page size):
+     *      After this we might end up with more bytes than the offset in EntryRef can handle. In this case this is 32768.
+     *   4) Cap bytes to alloc to the max offset EntryRef can handle.
+     *      The max bytes to alloc is: maxClusters * clusterSize * elementSize.
+     */
+    assertGrowStats<uint8_t>({8192,8192,8192,16384,16384,32768,65536,65536,98304,98304,98304,98304}, 3);
+    assertGrowStats<uint8_t>({16384,16384,16384,32768,32768,65536,131072,131072,163840,163840,163840,163840}, 5);
+    assertGrowStats<uint8_t>({16384,16384,16384,32768,32768,65536,131072,131072,229376,229376,229376,229376}, 7);
+    assertGrowStats<uint32_t>({8192,8192,8192,16384,16384,32768,65536,65536,98304,98304,98304,98304}, 3);
+    assertGrowStats<uint32_t>({16384,16384,16384,32768,32768,65536,131072,131072,163840,163840,163840,163840}, 5);
+    assertGrowStats<uint32_t>({16384,16384,16384,32768,32768,65536,131072,131072,229376,229376,229376,229376}, 7);
+}
+
 }
 }
 
diff --git a/searchlib/src/vespa/searchlib/datastore/buffer_type.h b/searchlib/src/vespa/searchlib/datastore/buffer_type.h
index adeaa7f4f72..2ec66e3aae3 100644
--- a/searchlib/src/vespa/searchlib/datastore/buffer_type.h
+++ b/searchlib/src/vespa/searchlib/datastore/buffer_type.h
@@ -51,7 +51,7 @@ public:
     virtual void initializeReservedElements(void *buffer, size_t reservedElements) = 0;
     virtual size_t elementSize() const = 0;
     virtual void cleanHold(void *buffer, uint64_t offset, uint64_t len, CleanContext cleanCtx) = 0;
-    uint32_t getClusterSize() const { return _clusterSize; }
+    size_t getClusterSize() const { return _clusterSize; }
     void flushLastUsed();
     virtual void onActive(uint32_t bufferId, size_t *usedElems, size_t &deadElems, void *buffer);
     void onHold(const size_t *usedElems);
@@ -70,7 +70,7 @@ public:
     void clampMaxClusters(uint32_t maxClusters);
 
     uint32_t getActiveBuffers() const { return _activeBuffers; }
-    uint32_t getMaxClusters() const { return _maxClusters; }
+    size_t getMaxClusters() const { return _maxClusters; }
     uint32_t getNumClustersForNewBuffer() const { return _numClustersForNewBuffer; }
 };
 
diff --git a/searchlib/src/vespa/searchlib/datastore/bufferstate.cpp b/searchlib/src/vespa/searchlib/datastore/bufferstate.cpp
index 9bb6fee7b79..de4ac9f6fc2 100644
--- a/searchlib/src/vespa/searchlib/datastore/bufferstate.cpp
+++ b/searchlib/src/vespa/searchlib/datastore/bufferstate.cpp
@@ -80,6 +80,11 @@ calcAllocation(uint32_t bufferId,
     size_t allocClusters = typeHandler.calcClustersToAlloc(bufferId, elementsNeeded, resizing);
     size_t allocElements = allocClusters * typeHandler.getClusterSize();
     size_t allocBytes = roundUpToMatchAllocator(allocElements * typeHandler.elementSize());
+    size_t maxAllocBytes = typeHandler.getMaxClusters() * typeHandler.getClusterSize() * typeHandler.elementSize();
+    if (allocBytes > maxAllocBytes) {
+        // Ensure that allocated bytes does not exceed the maximum handled by this type.
+        allocBytes = maxAllocBytes;
+    }
     size_t adjustedAllocElements = (allocBytes / typeHandler.elementSize());
     return AllocResult(adjustedAllocElements, allocBytes);
 }