Merge pull request #21159 from vespa-engine/toregge/move-ownership-of-memory-allocator-to-attribute-vector

Move ownership of memory allocator to AttributeVector.

5 files changed, 40 insertions, 18 deletions

diff --git a/searchlib/src/vespa/searchlib/attribute/attributevector.cpp b/searchlib/src/vespa/searchlib/attribute/attributevector.cpp
index a2ac482ebf3..08721b7302e 100644
--- a/searchlib/src/vespa/searchlib/attribute/attributevector.cpp
+++ b/searchlib/src/vespa/searchlib/attribute/attributevector.cpp
@@ -23,6 +23,7 @@
 #include <vespa/searchlib/util/file_settings.h>
 #include <vespa/searchlib/util/logutil.h>
 #include <vespa/vespalib/util/exceptions.h>
+#include <vespa/vespalib/util/mmap_file_allocator_factory.h>
 #include <vespa/vespalib/util/size_literals.h>
 #include <thread>
 
@@ -107,6 +108,34 @@ AttributeVector::ValueModifier::~ValueModifier() {
     }
 }
 
+namespace {
+
+bool
+allow_paged(const search::attribute::Config& config)
+{
+    if (!config.paged()) {
+        return false;
+    }
+    using Type = search::attribute::BasicType::Type;
+    if (config.basicType() == Type::REFERENCE || config.basicType() == Type::PREDICATE) {
+        return false;
+    }
+    if (config.basicType() == Type::TENSOR) {
+        return (!config.tensorType().is_error() && config.tensorType().is_dense());
+    }
+    return true;
+}
+
+std::unique_ptr<vespalib::alloc::MemoryAllocator>
+make_memory_allocator(const vespalib::string& name, const search::attribute::Config& config)
+{
+    if (allow_paged(config)) {
+        return vespalib::alloc::MmapFileAllocatorFactory::instance().make_memory_allocator(name);
+    }
+    return {};
+}
+
+}
 
 AttributeVector::AttributeVector(vespalib::stringref baseFileName, const Config &c)
     : _baseFileName(baseFileName),
@@ -124,7 +153,9 @@ AttributeVector::AttributeVector(vespalib::stringref baseFileName, const Config
       _compactLidSpaceGeneration(0u),
       _hasEnum(false),
       _loaded(false),
-      _isUpdateableInMemoryOnly(attribute::isUpdateableInMemoryOnly(getName(), getConfig()))
+      _isUpdateableInMemoryOnly(attribute::isUpdateableInMemoryOnly(getName(), getConfig())),
+      _nextStatUpdateTime(),
+      _memory_allocator(make_memory_allocator(_baseFileName.getAttributeName(), c))
 {
 }
 
diff --git a/searchlib/src/vespa/searchlib/attribute/attributevector.h b/searchlib/src/vespa/searchlib/attribute/attributevector.h
index acd00413568..2bf2f3a6ed6 100644
--- a/searchlib/src/vespa/searchlib/attribute/attributevector.h
+++ b/searchlib/src/vespa/searchlib/attribute/attributevector.h
@@ -377,6 +377,7 @@ protected:
     virtual vespalib::MemoryUsage getEnumStoreValuesMemoryUsage() const;
     virtual void populate_address_space_usage(AddressSpaceUsage& usage) const;
 
+    const std::shared_ptr<vespalib::alloc::MemoryAllocator>& get_memory_allocator() const noexcept { return _memory_allocator; }
 public:
     DECLARE_IDENTIFIABLE_ABSTRACT(AttributeVector);
     bool isLoaded() const { return _loaded; }
@@ -584,6 +585,7 @@ private:
     bool                                  _loaded;
     bool                                  _isUpdateableInMemoryOnly;
     vespalib::steady_time                 _nextStatUpdateTime;
+    std::shared_ptr<vespalib::alloc::MemoryAllocator> _memory_allocator;
 
 ////// Locking strategy interface. only available from the Guards.
     /**
diff --git a/searchlib/src/vespa/searchlib/tensor/dense_tensor_attribute.cpp b/searchlib/src/vespa/searchlib/tensor/dense_tensor_attribute.cpp
index 6dd630a6426..d376fb020be 100644
--- a/searchlib/src/vespa/searchlib/tensor/dense_tensor_attribute.cpp
+++ b/searchlib/src/vespa/searchlib/tensor/dense_tensor_attribute.cpp
@@ -99,17 +99,6 @@ BlobSequenceReader::is_present() {
     return true;
 }
 
-
-
-std::unique_ptr<vespalib::alloc::MemoryAllocator>
-make_memory_allocator(const vespalib::string& name, bool swappable)
-{
-    if (swappable) {
-        return vespalib::alloc::MmapFileAllocatorFactory::instance().make_memory_allocator(name);
-    }
-    return {};
-}
-
 }
 
 void
@@ -161,7 +150,7 @@ DenseTensorAttribute::populate_address_space_usage(AddressSpaceUsage& usage) con
 DenseTensorAttribute::DenseTensorAttribute(vespalib::stringref baseFileName, const Config& cfg,
                                            const NearestNeighborIndexFactory& index_factory)
     : TensorAttribute(baseFileName, cfg, _denseTensorStore),
-      _denseTensorStore(cfg.tensorType(), make_memory_allocator(getName(), cfg.paged())),
+      _denseTensorStore(cfg.tensorType(), get_memory_allocator()),
       _index()
 {
     if (cfg.hnsw_index_params().has_value()) {
diff --git a/searchlib/src/vespa/searchlib/tensor/dense_tensor_store.cpp b/searchlib/src/vespa/searchlib/tensor/dense_tensor_store.cpp
index ed3fb737b7d..6435ba6f27c 100644
--- a/searchlib/src/vespa/searchlib/tensor/dense_tensor_store.cpp
+++ b/searchlib/src/vespa/searchlib/tensor/dense_tensor_store.cpp
@@ -45,7 +45,7 @@ DenseTensorStore::TensorSizeCalc::TensorSizeCalc(const ValueType &type)
     _aligned_size = my_align(buf_size, alignment);
 }
 
-DenseTensorStore::BufferType::BufferType(const TensorSizeCalc &tensorSizeCalc, std::unique_ptr<vespalib::alloc::MemoryAllocator> allocator)
+DenseTensorStore::BufferType::BufferType(const TensorSizeCalc &tensorSizeCalc, std::shared_ptr<vespalib::alloc::MemoryAllocator> allocator)
     : vespalib::datastore::BufferType<char>(tensorSizeCalc.alignedSize(), MIN_BUFFER_ARRAYS, RefType::offsetSize()),
       _allocator(std::move(allocator))
 {}
@@ -65,7 +65,7 @@ DenseTensorStore::BufferType::get_memory_allocator() const
     return _allocator.get();
 }
 
-DenseTensorStore::DenseTensorStore(const ValueType &type, std::unique_ptr<vespalib::alloc::MemoryAllocator> allocator)
+DenseTensorStore::DenseTensorStore(const ValueType &type, std::shared_ptr<vespalib::alloc::MemoryAllocator> allocator)
     : TensorStore(_concreteStore),
       _concreteStore(),
       _tensorSizeCalc(type),
diff --git a/searchlib/src/vespa/searchlib/tensor/dense_tensor_store.h b/searchlib/src/vespa/searchlib/tensor/dense_tensor_store.h
index 47932fbff7e..7176edbcf08 100644
--- a/searchlib/src/vespa/searchlib/tensor/dense_tensor_store.h
+++ b/searchlib/src/vespa/searchlib/tensor/dense_tensor_store.h
@@ -37,9 +37,9 @@ public:
     class BufferType : public vespalib::datastore::BufferType<char>
     {
         using CleanContext = vespalib::datastore::BufferType<char>::CleanContext;
-        std::unique_ptr<vespalib::alloc::MemoryAllocator> _allocator;
+        std::shared_ptr<vespalib::alloc::MemoryAllocator> _allocator;
     public:
-        BufferType(const TensorSizeCalc &tensorSizeCalc, std::unique_ptr<vespalib::alloc::MemoryAllocator> allocator);
+        BufferType(const TensorSizeCalc &tensorSizeCalc, std::shared_ptr<vespalib::alloc::MemoryAllocator> allocator);
         ~BufferType() override;
         void cleanHold(void *buffer, size_t offset, ElemCount numElems, CleanContext cleanCtx) override;
         const vespalib::alloc::MemoryAllocator* get_memory_allocator() const override;
@@ -55,7 +55,7 @@ private:
     TensorStore::EntryRef
     setDenseTensor(const TensorType &tensor);
 public:
-    DenseTensorStore(const ValueType &type, std::unique_ptr<vespalib::alloc::MemoryAllocator> allocator);
+    DenseTensorStore(const ValueType &type, std::shared_ptr<vespalib::alloc::MemoryAllocator> allocator);
     ~DenseTensorStore() override;
 
     const ValueType &type() const { return _type; }