A swappable attribute will use a file backed memory allocator.

6 files changed, 55 insertions, 12 deletions

diff --git a/searchlib/src/tests/attribute/tensorattribute/tensorattribute_test.cpp b/searchlib/src/tests/attribute/tensorattribute/tensorattribute_test.cpp
index b27c26d1139..56233722f69 100644
--- a/searchlib/src/tests/attribute/tensorattribute/tensorattribute_test.cpp
+++ b/searchlib/src/tests/attribute/tensorattribute/tensorattribute_test.cpp
@@ -26,6 +26,7 @@
 #include <vespa/vespalib/io/fileutil.h>
 #include <vespa/vespalib/test/insertion_operators.h>
 #include <vespa/vespalib/testkit/test_kit.h>
+#include <vespa/vespalib/util/mmap_file_allocator_factory.h>
 #include <vespa/searchlib/util/bufferwriter.h>
 
 #include <vespa/log/log.h>
@@ -256,6 +257,7 @@ struct FixtureTraits {
     bool use_direct_tensor_attribute = false;
     bool enable_hnsw_index = false;
     bool use_mock_index = false;
+    bool use_mmap_file_allocator = false;
 
     FixtureTraits dense() && {
         use_dense_tensor_attribute = true;
@@ -263,6 +265,11 @@ struct FixtureTraits {
         return *this;
     }
 
+    FixtureTraits mmap_file_allocator() && {
+        use_mmap_file_allocator = true;
+        return *this;
+    }
+
     FixtureTraits hnsw() && {
         use_dense_tensor_attribute = true;
         enable_hnsw_index = true;
@@ -327,6 +334,9 @@ struct Fixture {
         if (_cfg.tensorType().is_dense()) {
             _denseTensors = true;
         }
+        if (_traits.use_mmap_file_allocator) {
+            _cfg.setSwappable(true);
+        }
         if (_traits.use_mock_index) {
             _index_factory = std::make_unique<MockNearestNeighborIndexFactory>();
         } else {
@@ -1000,4 +1010,17 @@ TEST_F("NN blueprint handles strong filter triggering brute force search", Neare
     EXPECT_FALSE(bp->may_approximate());
 }
 
+TEST("Dense tensor attribute with swappable flag uses mmap file allocator")
+{
+    vespalib::string basedir("mmap-file-allocator-factory-dir");
+    vespalib::alloc::MmapFileAllocatorFactory::instance().setup(basedir);
+    {
+        Fixture f(vec_2d_spec, FixtureTraits().dense().mmap_file_allocator());
+        vespalib::string allocator_dir(basedir + "/0.my_attr");
+        EXPECT_TRUE(vespalib::isDirectory(allocator_dir));
+    }
+    vespalib::alloc::MmapFileAllocatorFactory::instance().setup("");
+    vespalib::rmdir(basedir, true);
+}
+
 TEST_MAIN() { TEST_RUN_ALL(); }
diff --git a/searchlib/src/tests/tensor/dense_tensor_store/dense_tensor_store_test.cpp b/searchlib/src/tests/tensor/dense_tensor_store/dense_tensor_store_test.cpp
index eb8d89edc07..75e231a815c 100644
--- a/searchlib/src/tests/tensor/dense_tensor_store/dense_tensor_store_test.cpp
+++ b/searchlib/src/tests/tensor/dense_tensor_store/dense_tensor_store_test.cpp
@@ -1,5 +1,8 @@
 // Copyright 2017 Yahoo Holdings. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+#include <vespa/log/log.h>
+LOG_SETUP("dense_tensor_store_test");
 #include <vespa/vespalib/testkit/test_kit.h>
+#include <vespa/vespalib/util/memory_allocator.h>
 #include <vespa/searchlib/tensor/dense_tensor_store.h>
 #include <vespa/eval/eval/simple_value.h>
 #include <vespa/eval/eval/tensor_spec.h>
@@ -7,9 +10,6 @@
 #include <vespa/eval/eval/value_type.h>
 #include <vespa/eval/eval/test/value_compare.h>
 
-#include <vespa/log/log.h>
-LOG_SETUP("dense_tensor_store_test");
-
 using search::tensor::DenseTensorStore;
 using vespalib::eval::SimpleValue;
 using vespalib::eval::TensorSpec;
@@ -28,7 +28,7 @@ struct Fixture
 {
     DenseTensorStore store;
     Fixture(const vespalib::string &tensorType)
-        : store(ValueType::from_spec(tensorType))
+        : store(ValueType::from_spec(tensorType), {})
     {}
     void assertSetAndGetTensor(const TensorSpec &tensorSpec) {
         Value::UP expTensor = makeTensor(tensorSpec);
diff --git a/searchlib/src/vespa/searchlib/tensor/dense_tensor_attribute.cpp b/searchlib/src/vespa/searchlib/tensor/dense_tensor_attribute.cpp
index be15967ab10..8cc4acefd0f 100644
--- a/searchlib/src/vespa/searchlib/tensor/dense_tensor_attribute.cpp
+++ b/searchlib/src/vespa/searchlib/tensor/dense_tensor_attribute.cpp
@@ -10,6 +10,8 @@
 #include <vespa/searchlib/attribute/load_utils.h>
 #include <vespa/searchlib/attribute/readerbase.h>
 #include <vespa/vespalib/data/slime/inserter.h>
+#include <vespa/vespalib/util/memory_allocator.h>
+#include <vespa/vespalib/util/mmap_file_allocator_factory.h>
 
 #include <vespa/log/log.h>
 LOG_SETUP(".searchlib.tensor.dense_tensor_attribute");
@@ -72,6 +74,15 @@ can_use_index_save_file(const search::attribute::Config &config, const search::a
     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
@@ -114,7 +125,7 @@ DenseTensorAttribute::memory_usage() const
 DenseTensorAttribute::DenseTensorAttribute(vespalib::stringref baseFileName, const Config& cfg,
                                            const NearestNeighborIndexFactory& index_factory)
     : TensorAttribute(baseFileName, cfg, _denseTensorStore),
-      _denseTensorStore(cfg.tensorType()),
+      _denseTensorStore(cfg.tensorType(), make_memory_allocator(getName(), cfg.swappable())),
       _index()
 {
     if (cfg.hnsw_index_params().has_value()) {
diff --git a/searchlib/src/vespa/searchlib/tensor/dense_tensor_attribute.h b/searchlib/src/vespa/searchlib/tensor/dense_tensor_attribute.h
index 0884b776ca5..7eb8f2510f7 100644
--- a/searchlib/src/vespa/searchlib/tensor/dense_tensor_attribute.h
+++ b/searchlib/src/vespa/searchlib/tensor/dense_tensor_attribute.h
@@ -29,7 +29,7 @@ private:
 public:
     DenseTensorAttribute(vespalib::stringref baseFileName, const Config& cfg,
                          const NearestNeighborIndexFactory& index_factory = DefaultNearestNeighborIndexFactory());
-    virtual ~DenseTensorAttribute();
+    ~DenseTensorAttribute() override;
     // Implements AttributeVector and ITensorAttribute
     uint32_t clearDoc(DocId docId) override;
     void setTensor(DocId docId, const vespalib::eval::Value &tensor) override;
diff --git a/searchlib/src/vespa/searchlib/tensor/dense_tensor_store.cpp b/searchlib/src/vespa/searchlib/tensor/dense_tensor_store.cpp
index e6da839da2e..13796d35dec 100644
--- a/searchlib/src/vespa/searchlib/tensor/dense_tensor_store.cpp
+++ b/searchlib/src/vespa/searchlib/tensor/dense_tensor_store.cpp
@@ -40,8 +40,9 @@ DenseTensorStore::TensorSizeCalc::alignedSize() const
     return my_align(bufSize(), DENSE_TENSOR_ALIGNMENT);
 }
 
-DenseTensorStore::BufferType::BufferType(const TensorSizeCalc &tensorSizeCalc)
-    : vespalib::datastore::BufferType<char>(tensorSizeCalc.alignedSize(), MIN_BUFFER_ARRAYS, RefType::offsetSize())
+DenseTensorStore::BufferType::BufferType(const TensorSizeCalc &tensorSizeCalc, std::unique_ptr<vespalib::alloc::MemoryAllocator> allocator)
+    : vespalib::datastore::BufferType<char>(tensorSizeCalc.alignedSize(), MIN_BUFFER_ARRAYS, RefType::offsetSize()),
+      _allocator(std::move(allocator))
 {}
 
 DenseTensorStore::BufferType::~BufferType() = default;
@@ -53,11 +54,17 @@ DenseTensorStore::BufferType::cleanHold(void *buffer, size_t offset,
     memset(static_cast<char *>(buffer) + offset, 0, numElems);
 }
 
-DenseTensorStore::DenseTensorStore(const ValueType &type)
+const vespalib::alloc::MemoryAllocator*
+DenseTensorStore::BufferType::get_memory_allocator() const
+{
+    return _allocator.get();
+}
+
+DenseTensorStore::DenseTensorStore(const ValueType &type, std::unique_ptr<vespalib::alloc::MemoryAllocator> allocator)
     : TensorStore(_concreteStore),
       _concreteStore(),
       _tensorSizeCalc(type),
-      _bufferType(_tensorSizeCalc),
+      _bufferType(_tensorSizeCalc, std::move(allocator)),
       _type(type),
       _emptySpace()
 {
diff --git a/searchlib/src/vespa/searchlib/tensor/dense_tensor_store.h b/searchlib/src/vespa/searchlib/tensor/dense_tensor_store.h
index 638f602d613..dad28642e67 100644
--- a/searchlib/src/vespa/searchlib/tensor/dense_tensor_store.h
+++ b/searchlib/src/vespa/searchlib/tensor/dense_tensor_store.h
@@ -36,10 +36,12 @@ public:
     class BufferType : public vespalib::datastore::BufferType<char>
     {
         using CleanContext = vespalib::datastore::BufferType<char>::CleanContext;
+        std::unique_ptr<vespalib::alloc::MemoryAllocator> _allocator;
     public:
-        BufferType(const TensorSizeCalc &tensorSizeCalc);
+        BufferType(const TensorSizeCalc &tensorSizeCalc, std::unique_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;
     };
 private:
     DataStoreType _concreteStore;
@@ -55,7 +57,7 @@ private:
     setDenseTensor(const TensorType &tensor);
 
 public:
-    DenseTensorStore(const ValueType &type);
+    DenseTensorStore(const ValueType &type, std::unique_ptr<vespalib::alloc::MemoryAllocator> allocator);
     ~DenseTensorStore() override;
 
     const ValueType &type() const { return _type; }