Simplify how used elements across all active buffers are aggregated.

This also prepares for aggregating dead elements
and to take these into account when calculating how many arrays to alloc.

5 files changed, 132 insertions, 42 deletions

diff --git a/vespalib/src/tests/datastore/array_store/array_store_test.cpp b/vespalib/src/tests/datastore/array_store/array_store_test.cpp
index a1e7140fefe..417d8b80d87 100644
--- a/vespalib/src/tests/datastore/array_store/array_store_test.cpp
+++ b/vespalib/src/tests/datastore/array_store/array_store_test.cpp
@@ -150,13 +150,13 @@ TEST("require that we test with trivial and non-trivial types")
 
 TEST_F("control static sizes", NumberFixture(3)) {
 #ifdef _LIBCPP_VERSION
-    EXPECT_EQUAL(392u, sizeof(f.store));
+    EXPECT_EQUAL(400u, sizeof(f.store));
     EXPECT_EQUAL(296u, sizeof(NumberFixture::ArrayStoreType::DataStoreType));
 #else
-    EXPECT_EQUAL(424u, sizeof(f.store));
+    EXPECT_EQUAL(432u, sizeof(f.store));
     EXPECT_EQUAL(328u, sizeof(NumberFixture::ArrayStoreType::DataStoreType));
 #endif
-    EXPECT_EQUAL(64u, sizeof(NumberFixture::ArrayStoreType::SmallArrayType));
+    EXPECT_EQUAL(72u, sizeof(NumberFixture::ArrayStoreType::SmallArrayType));
     MemoryUsage usage = f.store.getMemoryUsage();
     EXPECT_EQUAL(960u, usage.allocatedBytes());
     EXPECT_EQUAL(32u, usage.usedBytes());
diff --git a/vespalib/src/tests/datastore/buffer_type/buffer_type_test.cpp b/vespalib/src/tests/datastore/buffer_type/buffer_type_test.cpp
index 96d94e9a71b..b5a0212528b 100644
--- a/vespalib/src/tests/datastore/buffer_type/buffer_type_test.cpp
+++ b/vespalib/src/tests/datastore/buffer_type/buffer_type_test.cpp
@@ -14,6 +14,7 @@ struct Setup {
     uint32_t  _minArrays;
     ElemCount _usedElems;
     ElemCount _neededElems;
+    ElemCount _deadElems;
     uint32_t  _bufferId;
     float     _allocGrowFactor;
     bool      _resizing;
@@ -21,6 +22,7 @@ struct Setup {
         : _minArrays(0),
           _usedElems(0),
           _neededElems(0),
+          _deadElems(0),
           _bufferId(1),
           _allocGrowFactor(0.5),
           _resizing(false)
@@ -33,25 +35,40 @@ struct Setup {
 };
 
 struct Fixture {
-    Setup         setup;
+    std::vector<Setup> setups;
     IntBufferType bufferType;
-    ElemCount     deadElems;
     int buffer[ARRAYS_SIZE];
     Fixture(const Setup &setup_)
-        : setup(setup_),
-          bufferType(ARRAYS_SIZE, setup._minArrays, MAX_ARRAYS, NUM_ARRAYS_FOR_NEW_BUFFER, setup._allocGrowFactor),
-          deadElems(0),
+        : setups(),
+          bufferType(ARRAYS_SIZE, setup_._minArrays, MAX_ARRAYS, NUM_ARRAYS_FOR_NEW_BUFFER, setup_._allocGrowFactor),
           buffer()
-    {}
+    {
+        setups.reserve(4);
+        setups.push_back(setup_);
+    }
     ~Fixture() {
-        bufferType.onHold(&setup._usedElems);
-        bufferType.onFree(setup._usedElems);
+        for (auto& setup : setups) {
+            bufferType.onHold(&setup._usedElems);
+            bufferType.onFree(setup._usedElems);
+        }
+    }
+    Setup& curr_setup() {
+        return setups.back();
+    }
+    void add_setup(const Setup& setup_in) {
+        // The buffer type stores pointers to ElemCount (from Setup) and we must ensure these do not move in memory.
+        assert(setups.size() < setups.capacity());
+        setups.push_back(setup_in);
     }
     void onActive() {
-        bufferType.onActive(setup._bufferId, &setup._usedElems, deadElems, &buffer[0]);
+        bufferType.onActive(curr_setup()._bufferId, &curr_setup()._usedElems, curr_setup()._deadElems, &buffer[0]);
     }
     size_t arraysToAlloc() {
-        return bufferType.calcArraysToAlloc(setup._bufferId, setup._neededElems, setup._resizing);
+        return bufferType.calcArraysToAlloc(curr_setup()._bufferId, curr_setup()._neededElems, curr_setup()._resizing);
+    }
+    void assertArraysToAlloc(size_t exp) {
+        onActive();
+        EXPECT_EQUAL(exp, arraysToAlloc());
     }
 };
 
@@ -59,8 +76,7 @@ void
 assertArraysToAlloc(size_t exp, const Setup &setup)
 {
     Fixture f(setup);
-    f.onActive();
-    EXPECT_EQUAL(exp, f.arraysToAlloc());
+    f.assertArraysToAlloc(exp);
 }
 
 TEST("require that complete arrays are allocated")
@@ -113,4 +129,14 @@ TEST("require that arrays to alloc is capped to min arrays")
     TEST_DO(assertArraysToAlloc(17, Setup().used(34 * 4).needed(4).minArrays(16)));
 }
 
+TEST("arrays to alloc considers used elements across all active buffers")
+{
+    Fixture f(Setup().used(6 * 4));
+    f.assertArraysToAlloc(6 * 0.5);
+    f.add_setup(Setup().used(8 * 4));
+    f.assertArraysToAlloc((6 + 8) * 0.5);
+    f.add_setup(Setup().used(10 * 4));
+    f.assertArraysToAlloc((6 + 8 + 10) * 0.5);
+}
+
 TEST_MAIN() { TEST_RUN_ALL(); }
diff --git a/vespalib/src/tests/datastore/datastore/datastore_test.cpp b/vespalib/src/tests/datastore/datastore/datastore_test.cpp
index f3f43fb575b..11b4f5e6631 100644
--- a/vespalib/src/tests/datastore/datastore/datastore_test.cpp
+++ b/vespalib/src/tests/datastore/datastore/datastore_test.cpp
@@ -635,7 +635,7 @@ TEST(DataStoreTest, can_set_memory_allocator)
 }
 
 TEST(DataStoreTest, control_static_sizes) {
-    EXPECT_EQ(64, sizeof(BufferTypeBase));
+    EXPECT_EQ(72, sizeof(BufferTypeBase));
     EXPECT_EQ(32, sizeof(BufferState::FreeList));
     EXPECT_EQ(1, sizeof(BufferState::State));
     EXPECT_EQ(144, sizeof(BufferState));
diff --git a/vespalib/src/vespa/vespalib/datastore/buffer_type.cpp b/vespalib/src/vespa/vespalib/datastore/buffer_type.cpp
index f25ef2c1e52..965b997a280 100644
--- a/vespalib/src/vespa/vespalib/datastore/buffer_type.cpp
+++ b/vespalib/src/vespa/vespalib/datastore/buffer_type.cpp
@@ -33,9 +33,8 @@ BufferTypeBase::BufferTypeBase(uint32_t arraySize,
       _allocGrowFactor(allocGrowFactor),
       _activeBuffers(0),
       _holdBuffers(0),
-      _activeUsedElems(0),
       _holdUsedElems(0),
-      _lastUsedElems(nullptr)
+      _aggr_counts()
 {
 }
 
@@ -50,9 +49,8 @@ BufferTypeBase::~BufferTypeBase()
 {
     assert(_activeBuffers == 0);
     assert(_holdBuffers == 0);
-    assert(_activeUsedElems == 0);
     assert(_holdUsedElems == 0);
-    assert(_lastUsedElems == nullptr);
+    assert(_aggr_counts.empty());
 }
 
 ElemCount
@@ -62,20 +60,10 @@ BufferTypeBase::getReservedElements(uint32_t bufferId) const
 }
 
 void
-BufferTypeBase::flushLastUsed()
-{
-    if (_lastUsedElems != nullptr) {
-        _activeUsedElems += *_lastUsedElems;
-        _lastUsedElems = nullptr;
-    }
-}
-
-void
 BufferTypeBase::onActive(uint32_t bufferId, ElemCount *usedElems, ElemCount &deadElems, void *buffer)
 {
-    flushLastUsed();
     ++_activeBuffers;
-    _lastUsedElems = usedElems;
+    _aggr_counts.add_buffer(usedElems);
     size_t reservedElems = getReservedElements(bufferId);
     if (reservedElems != 0u) {
         initializeReservedElements(buffer, reservedElems);
@@ -87,13 +75,9 @@ BufferTypeBase::onActive(uint32_t bufferId, ElemCount *usedElems, ElemCount &dea
 void
 BufferTypeBase::onHold(const ElemCount *usedElems)
 {
-    if (usedElems == _lastUsedElems) {
-        flushLastUsed();
-    }
     --_activeBuffers;
     ++_holdBuffers;
-    assert(_activeUsedElems >= *usedElems);
-    _activeUsedElems -= *usedElems;
+    _aggr_counts.remove_buffer(usedElems);
     _holdUsedElems += *usedElems;
 }
 
@@ -123,9 +107,11 @@ size_t
 BufferTypeBase::calcArraysToAlloc(uint32_t bufferId, ElemCount elemsNeeded, bool resizing) const
 {
     size_t reservedElems = getReservedElements(bufferId);
-    size_t usedElems = (resizing ? 0 : _activeUsedElems);
-    if (_lastUsedElems != nullptr) {
-        usedElems += *_lastUsedElems;
+    size_t usedElems = 0;
+    if (resizing) {
+        usedElems = _aggr_counts.empty() ? 0 : _aggr_counts.last_buffer().used_elems;
+    } else {
+        usedElems = _aggr_counts.all_buffers().used_elems;
     }
     assert((usedElems % _arraySize) == 0);
     size_t usedArrays = usedElems / _arraySize;
@@ -133,6 +119,7 @@ BufferTypeBase::calcArraysToAlloc(uint32_t bufferId, ElemCount elemsNeeded, bool
     size_t growArrays = (usedArrays * _allocGrowFactor);
     size_t wantedArrays = std::max((resizing ? usedArrays : 0u) + growArrays,
                                    static_cast<size_t>(_minArrays));
+
     size_t result = wantedArrays;
     if (result < neededArrays) {
         result = neededArrays;
@@ -144,6 +131,49 @@ BufferTypeBase::calcArraysToAlloc(uint32_t bufferId, ElemCount elemsNeeded, bool
     return result;
 }
 
+BufferTypeBase::AggregatedBufferCounts::AggregatedBufferCounts()
+    : _counts()
+{
+}
+
+void
+BufferTypeBase::AggregatedBufferCounts::add_buffer(const ElemCount* used_elems)
+{
+    for (const auto& elem : _counts) {
+        assert(elem.used_ptr != used_elems);
+    }
+    _counts.emplace_back(used_elems, nullptr);
+}
+
+void
+BufferTypeBase::AggregatedBufferCounts::remove_buffer(const ElemCount* used_elems)
+{
+    auto itr = std::find_if(_counts.begin(), _counts.end(),
+                            [=](const auto& elem){ return elem.used_ptr == used_elems; });
+    assert(itr != _counts.end());
+    _counts.erase(itr);
+}
+
+BufferTypeBase::BufferCounts
+BufferTypeBase::AggregatedBufferCounts::last_buffer() const
+{
+    BufferCounts result;
+    assert(!_counts.empty());
+    const auto& last = _counts.back();
+    result.used_elems += *last.used_ptr;
+    return result;
+}
+
+BufferTypeBase::BufferCounts
+BufferTypeBase::AggregatedBufferCounts::all_buffers() const
+{
+    BufferCounts result;
+    for (const auto& elem : _counts) {
+        result.used_elems += *elem.used_ptr;
+    }
+    return result;
+}
+
 template class BufferType<char>;
 template class BufferType<uint8_t>;
 template class BufferType<uint32_t>;
diff --git a/vespalib/src/vespa/vespalib/datastore/buffer_type.h b/vespalib/src/vespa/vespalib/datastore/buffer_type.h
index d52fbfec2a3..9d261ce26bc 100644
--- a/vespalib/src/vespa/vespalib/datastore/buffer_type.h
+++ b/vespalib/src/vespa/vespalib/datastore/buffer_type.h
@@ -4,6 +4,7 @@
 
 #include "atomic_entry_ref.h"
 #include <string>
+#include <vector>
 
 namespace vespalib::alloc { class MemoryAllocator; }
 
@@ -33,6 +34,7 @@ public:
         {}
         void extraBytesCleaned(size_t value);
     };
+
     BufferTypeBase(const BufferTypeBase &rhs) = delete;
     BufferTypeBase & operator=(const BufferTypeBase &rhs) = delete;
     BufferTypeBase(BufferTypeBase &&rhs) noexcept = default;
@@ -58,7 +60,6 @@ public:
     virtual size_t elementSize() const = 0;
     virtual void cleanHold(void *buffer, size_t offset, ElemCount numElems, CleanContext cleanCtx) = 0;
     size_t getArraySize() const { return _arraySize; }
-    void flushLastUsed();
     virtual void onActive(uint32_t bufferId, ElemCount *usedElems, ElemCount &deadElems, void *buffer);
     void onHold(const ElemCount *usedElems);
     virtual void onFree(ElemCount usedElems);
@@ -75,6 +76,40 @@ public:
     size_t getMaxArrays() const { return _maxArrays; }
     uint32_t getNumArraysForNewBuffer() const { return _numArraysForNewBuffer; }
 protected:
+
+    struct BufferCounts {
+        ElemCount used_elems;
+        ElemCount dead_elems;
+        BufferCounts() : used_elems(0), dead_elems(0) {}
+        BufferCounts(ElemCount used_elems_in, ElemCount dead_elems_in)
+                : used_elems(used_elems_in), dead_elems(dead_elems_in)
+        {}
+    };
+
+    /**
+     * Tracks aggregated counts for all buffers that are in state ACTIVE.
+     */
+    class AggregatedBufferCounts {
+    private:
+        struct Element {
+            const ElemCount* used_ptr;
+            const ElemCount* dead_ptr;
+            Element() noexcept : used_ptr(nullptr), dead_ptr(nullptr) {}
+            Element(const ElemCount* used_ptr_in, const ElemCount* dead_ptr_in) noexcept
+                    : used_ptr(used_ptr_in), dead_ptr(dead_ptr_in)
+            {}
+        };
+        std::vector<Element> _counts;
+
+    public:
+        AggregatedBufferCounts();
+        void add_buffer(const ElemCount* used_elems);
+        void remove_buffer(const ElemCount* used_elems);
+        BufferCounts last_buffer() const;
+        BufferCounts all_buffers() const;
+        bool empty() const { return _counts.empty(); }
+    };
+
     uint32_t _arraySize;  // Number of elements in an allocation unit
     uint32_t _minArrays;  // Minimum number of arrays to allocate in a buffer
     uint32_t _maxArrays;  // Maximum number of arrays to allocate in a buffer
@@ -83,9 +118,8 @@ protected:
     float    _allocGrowFactor;
     uint32_t _activeBuffers;
     uint32_t _holdBuffers;
-    size_t   _activeUsedElems; // Number of used elements in all but the last active buffer for this type.
     size_t   _holdUsedElems;  // Number of used elements in all held buffers for this type.
-    const ElemCount *_lastUsedElems; // Number of used elements in the last active buffer for this type.
+    AggregatedBufferCounts _aggr_counts;
 };
 
 /**