Merge pull request #21203 from vespa-engine/toregge/keep-memory-allocator-in-growable-bitvector

Keep memory allocator in growable bitvector.

7 files changed, 38 insertions, 25 deletions

diff --git a/searchlib/src/tests/common/bitvector/bitvector_test.cpp b/searchlib/src/tests/common/bitvector/bitvector_test.cpp
index faeebf9984a..a9850756ac6 100644
--- a/searchlib/src/tests/common/bitvector/bitvector_test.cpp
+++ b/searchlib/src/tests/common/bitvector/bitvector_test.cpp
@@ -8,11 +8,16 @@
 #include <vespa/searchlib/common/bitvectoriterator.h>
 #include <vespa/searchlib/fef/termfieldmatchdata.h>
 #include <vespa/searchlib/fef/termfieldmatchdataarray.h>
+#include <vespa/vespalib/test/memory_allocator_observer.h>
 #include <vespa/vespalib/util/rand48.h>
 #include <algorithm>
 
 using namespace search;
 
+using vespalib::alloc::Alloc;
+using vespalib::alloc::test::MemoryAllocatorObserver;
+using AllocStats = MemoryAllocatorObserver::Stats;
+
 namespace {
 
 std::string
@@ -653,5 +658,26 @@ TEST("requireThatGrowWorks")
     g.trimHoldLists(2);
 }
 
+TEST("require that growable bit vectors keeps memory allocator")
+{
+    AllocStats stats;
+    auto memory_allocator = std::make_unique<MemoryAllocatorObserver>(stats);
+    Alloc init_alloc = Alloc::alloc_with_allocator(memory_allocator.get());
+    vespalib::GenerationHolder g;
+    GrowableBitVector v(200, 200, g, &init_alloc);
+    EXPECT_EQUAL(AllocStats(1, 0), stats);
+    v.resize(1);
+    EXPECT_EQUAL(AllocStats(2, 1), stats);
+    v.reserve(2000);
+    EXPECT_EQUAL(AllocStats(3, 1), stats);
+    v.extend(4000);
+    EXPECT_EQUAL(AllocStats(4, 1), stats);
+    v.shrink(200);
+    EXPECT_EQUAL(AllocStats(4, 1), stats);
+    v.resize(1);
+    EXPECT_EQUAL(AllocStats(5, 2), stats);
+    g.transferHoldLists(1);
+    g.trimHoldLists(2);
+}
 
 TEST_MAIN() { TEST_RUN_ALL(); }
diff --git a/searchlib/src/vespa/searchlib/common/allocatedbitvector.cpp b/searchlib/src/vespa/searchlib/common/allocatedbitvector.cpp
index d428ce52953..ab56317db6f 100644
--- a/searchlib/src/vespa/searchlib/common/allocatedbitvector.cpp
+++ b/searchlib/src/vespa/searchlib/common/allocatedbitvector.cpp
@@ -59,10 +59,10 @@ AllocatedBitVector::AllocatedBitVector(Index numberOfElements, Alloc buffer, siz
 {
 }
 
-AllocatedBitVector::AllocatedBitVector(Index numberOfElements, Index capacityBits, const void * rhsBuf, size_t rhsSize) :
+AllocatedBitVector::AllocatedBitVector(Index numberOfElements, Index capacityBits, const void * rhsBuf, size_t rhsSize, const Alloc* init_alloc) :
     BitVector(),
     _capacityBits(capacityBits),
-    _alloc(allocatePaddedAndAligned(0, numberOfElements, capacityBits))
+    _alloc(allocatePaddedAndAligned(0, numberOfElements, capacityBits, init_alloc))
 {
     _capacityBits = computeCapacity(_capacityBits, _alloc.size());
     init(_alloc.get(), 0, numberOfElements);
@@ -104,17 +104,9 @@ AllocatedBitVector::AllocatedBitVector(const BitVector & rhs, std::pair<Index, I
 AllocatedBitVector::~AllocatedBitVector() = default;
 
 void
-AllocatedBitVector::cleanup()
-{
-    init(nullptr, 0, 0);
-    Alloc().swap(_alloc);
-    _capacityBits = 0;
-}
-
-void
 AllocatedBitVector::resize(Index newLength)
 {
-    _alloc = allocatePaddedAndAligned(newLength);
+    _alloc = allocatePaddedAndAligned(0, newLength, newLength, &_alloc);
     _capacityBits = computeCapacity(newLength, _alloc.size());
     init(_alloc.get(), 0, newLength);
     clear();
@@ -145,7 +137,7 @@ AllocatedBitVector::grow(Index newSize, Index newCapacity)
     assert(newCapacity >= newSize);
     GenerationHeldBase::UP ret;
     if (newCapacity != capacity()) {
-        AllocatedBitVector tbv(newSize, newCapacity, _alloc.get(), size());
+        AllocatedBitVector tbv(newSize, newCapacity, _alloc.get(), size(), &_alloc);
         if (newSize > size()) {
             tbv.clearBitAndMaintainCount(size());  // Clear old guard bit.
         }
diff --git a/searchlib/src/vespa/searchlib/common/allocatedbitvector.h b/searchlib/src/vespa/searchlib/common/allocatedbitvector.h
index c57ca93ac02..2223c7b4701 100644
--- a/searchlib/src/vespa/searchlib/common/allocatedbitvector.h
+++ b/searchlib/src/vespa/searchlib/common/allocatedbitvector.h
@@ -35,7 +35,7 @@ public:
      * Creates a new bitvector with size of numberOfElements bits and at least a capacity of capacity.
      * Copies what it can from the original vector. This is used for extending vector.
      */
-    AllocatedBitVector(Index numberOfElements, Index capacity, const void * rhsBuf, size_t rhsSize);
+    AllocatedBitVector(Index numberOfElements, Index capacity, const void * rhsBuf, size_t rhsSize, const Alloc* init_alloc);
 
     AllocatedBitVector(const BitVector &other);
     AllocatedBitVector(const AllocatedBitVector &other);
@@ -74,12 +74,6 @@ private:
     }
 
     AllocatedBitVector(const BitVector &other, std::pair<Index, Index> size_capacity);
-
-    /**
-     * Prepare for potential reuse where new value might be filled in by
-     * Read method.
-     */
-    void cleanup();
 };
 
 } // namespace search
diff --git a/searchlib/src/vespa/searchlib/common/bitvector.cpp b/searchlib/src/vespa/searchlib/common/bitvector.cpp
index a3a1b9bed66..d2e6d9027ea 100644
--- a/searchlib/src/vespa/searchlib/common/bitvector.cpp
+++ b/searchlib/src/vespa/searchlib/common/bitvector.cpp
@@ -47,13 +47,13 @@ using vespalib::GenerationHeldAlloc;
 using vespalib::GenerationHolder;
 
 Alloc
-BitVector::allocatePaddedAndAligned(Index start, Index end, Index capacity)
+BitVector::allocatePaddedAndAligned(Index start, Index end, Index capacity, const Alloc* init_alloc)
 {
     assert(capacity >= end);
     uint32_t words = numActiveWords(start, capacity);
     words += (-words & 15); // Pad to 64 byte alignment
     const size_t sz(words * sizeof(Word));
-    Alloc alloc = Alloc::alloc(sz, MMAP_LIMIT);
+    Alloc alloc = (init_alloc != nullptr) ? init_alloc->create(sz) : Alloc::alloc(sz, MMAP_LIMIT);
     assert(alloc.size()/sizeof(Word) >= words);
     // Clear padding
     size_t usedBytes = numBytes(end - start);
diff --git a/searchlib/src/vespa/searchlib/common/bitvector.h b/searchlib/src/vespa/searchlib/common/bitvector.h
index 184669829cb..c1d447047d1 100644
--- a/searchlib/src/vespa/searchlib/common/bitvector.h
+++ b/searchlib/src/vespa/searchlib/common/bitvector.h
@@ -281,7 +281,7 @@ protected:
     static Alloc allocatePaddedAndAligned(Index start, Index end) {
         return allocatePaddedAndAligned(start, end, end);
     }
-    static Alloc allocatePaddedAndAligned(Index start, Index end, Index capacity);
+    static Alloc allocatePaddedAndAligned(Index start, Index end, Index capacity, const Alloc* init_alloc = nullptr);
 
 private:
     friend PartialBitVector;
diff --git a/searchlib/src/vespa/searchlib/common/growablebitvector.cpp b/searchlib/src/vespa/searchlib/common/growablebitvector.cpp
index 177266c5618..7c29945292e 100644
--- a/searchlib/src/vespa/searchlib/common/growablebitvector.cpp
+++ b/searchlib/src/vespa/searchlib/common/growablebitvector.cpp
@@ -10,8 +10,9 @@ using vespalib::GenerationHeldBase;
 using vespalib::GenerationHolder;
 
 GrowableBitVector::GrowableBitVector(Index newSize, Index newCapacity,
-                                     GenerationHolder &generationHolder)
-    : AllocatedBitVector(newSize, newCapacity, nullptr, 0),
+                                     GenerationHolder &generationHolder,
+                                     const Alloc* init_alloc)
+    : AllocatedBitVector(newSize, newCapacity, nullptr, 0, init_alloc),
       _generationHolder(generationHolder)
 {
     assert(newSize <= newCapacity);
diff --git a/searchlib/src/vespa/searchlib/common/growablebitvector.h b/searchlib/src/vespa/searchlib/common/growablebitvector.h
index 08c9fede8f2..71261e130b6 100644
--- a/searchlib/src/vespa/searchlib/common/growablebitvector.h
+++ b/searchlib/src/vespa/searchlib/common/growablebitvector.h
@@ -9,7 +9,7 @@ namespace search {
 class GrowableBitVector : public AllocatedBitVector
 {
 public:
-    GrowableBitVector(Index newSize, Index newCapacity, GenerationHolder &generationHolder);
+    GrowableBitVector(Index newSize, Index newCapacity, GenerationHolder &generationHolder, const Alloc* init_alloc = nullptr);
 
     /** Will return true if a a buffer is held */
     bool reserve(Index newCapacity);