Keep using same memory allocator when resizing rcu vector.

3 files changed, 76 insertions, 3 deletions

diff --git a/vespalib/src/tests/util/rcuvector/rcuvector_test.cpp b/vespalib/src/tests/util/rcuvector/rcuvector_test.cpp
index cf84ab03a25..14802a60ff1 100644
--- a/vespalib/src/tests/util/rcuvector/rcuvector_test.cpp
+++ b/vespalib/src/tests/util/rcuvector/rcuvector_test.cpp
@@ -1,11 +1,17 @@
 // Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 
 #include <vespa/vespalib/testkit/testapp.h>
+#include <vespa/vespalib/test/memory_allocator_observer.h>
 #include <vespa/vespalib/util/rcuvector.h>
 #include <vespa/vespalib/util/size_literals.h>
 
 using namespace vespalib;
 
+using vespalib::alloc::Alloc;
+using vespalib::alloc::MemoryAllocator;
+using MyMemoryAllocator = vespalib::alloc::test::MemoryAllocatorObserver;
+using AllocStats = MyMemoryAllocator::Stats;
+
 bool
 assertUsage(const MemoryUsage & exp, const MemoryUsage & act)
 {
@@ -288,4 +294,70 @@ TEST("test small expand")
     g.trimHoldLists(2);
 }
 
+struct Fixture {
+    using generation_t = GenerationHandler::generation_t;
+
+    AllocStats stats;
+    std::unique_ptr<MemoryAllocator> allocator;
+    Alloc initial_alloc;
+    GenerationHolder g;
+    RcuVectorBase<int> arr;
+
+    Fixture();
+    ~Fixture();
+    void transfer_and_trim(generation_t transfer_gen, generation_t trim_gen)
+    {
+        g.transferHoldLists(transfer_gen);
+        g.trimHoldLists(trim_gen);
+    }
+};
+
+Fixture::Fixture()
+    : stats(),
+      allocator(std::make_unique<MyMemoryAllocator>(stats)),
+      initial_alloc(Alloc::alloc_with_allocator(allocator.get())),
+      g(),
+      arr(g, initial_alloc)
+{
+    arr.reserve(100);
+}
+
+Fixture::~Fixture() = default;
+
+TEST_F("require that memory allocator can be set", Fixture)
+{
+    EXPECT_EQUAL(AllocStats(2, 0), f.stats);
+    f.transfer_and_trim(1, 2);
+    EXPECT_EQUAL(AllocStats(2, 1), f.stats);
+}
+
+TEST_F("require that memory allocator is preserved across reset", Fixture)
+{
+    f.arr.reset();
+    f.arr.reserve(100);
+    EXPECT_EQUAL(AllocStats(4, 1), f.stats);
+    f.transfer_and_trim(1, 2);
+    EXPECT_EQUAL(AllocStats(4, 3), f.stats);
+}
+
+TEST_F("require that created replacement vector uses same memory allocator", Fixture)
+{
+    auto arr2 = f.arr.create_replacement_vector();
+    EXPECT_EQUAL(AllocStats(2, 0), f.stats);
+    arr2.reserve(100);
+    EXPECT_EQUAL(AllocStats(3, 0), f.stats);
+    f.transfer_and_trim(1, 2);
+    EXPECT_EQUAL(AllocStats(3, 1), f.stats);
+}
+
+TEST_F("require that ensure_size and shrink use same memory allocator", Fixture)
+{
+    f.arr.ensure_size(2000);
+    EXPECT_EQUAL(AllocStats(3, 0), f.stats);
+    f.arr.shrink(1000);
+    EXPECT_EQUAL(AllocStats(4, 0), f.stats);
+    f.transfer_and_trim(1, 2);
+    EXPECT_EQUAL(AllocStats(4, 3), f.stats);
+}
+
 TEST_MAIN() { TEST_RUN_ALL(); }
diff --git a/vespalib/src/vespa/vespalib/util/rcuvector.h b/vespalib/src/vespa/vespalib/util/rcuvector.h
index dd4fa660279..00d050fa8d1 100644
--- a/vespalib/src/vespa/vespalib/util/rcuvector.h
+++ b/vespalib/src/vespa/vespalib/util/rcuvector.h
@@ -122,6 +122,7 @@ public:
     void reset();
     void shrink(size_t newSize) __attribute__((noinline));
     void replaceVector(ArrayType replacement);
+    ArrayType create_replacement_vector() const { return _data.create(); }
 };
 
 template <typename T>
diff --git a/vespalib/src/vespa/vespalib/util/rcuvector.hpp b/vespalib/src/vespa/vespalib/util/rcuvector.hpp
index 3c455149dfd..7c70539da0e 100644
--- a/vespalib/src/vespa/vespalib/util/rcuvector.hpp
+++ b/vespalib/src/vespa/vespalib/util/rcuvector.hpp
@@ -42,7 +42,7 @@ template <typename T>
 void
 RcuVectorBase<T>::reset() {
     // Assumes no readers at this moment
-    ArrayType().swap(_data);
+    _data.reset();
     _data.reserve(16);
 }
 
@@ -52,7 +52,7 @@ RcuVectorBase<T>::~RcuVectorBase() = default;
 template <typename T>
 void
 RcuVectorBase<T>::expand(size_t newCapacity) {
-    ArrayType tmpData;
+    auto tmpData = create_replacement_vector();
     tmpData.reserve(newCapacity);
     for (const T & v : _data) {
         tmpData.push_back_fast(v);
@@ -91,7 +91,7 @@ RcuVectorBase<T>::shrink(size_t newSize)
         return;
     }
     if (!_data.try_unreserve(wantedCapacity)) {
-        ArrayType tmpData;
+        auto tmpData = create_replacement_vector();
         tmpData.reserve(wantedCapacity);
         tmpData.resize(newSize);
         for (uint32_t i = 0; i < newSize; ++i) {