From d0580a51332db9e2f8e2d24799ef22942dc16abb Mon Sep 17 00:00:00 2001
From: Tor Egge <Tor.Egge@broadpark.no>
Date: Fri, 12 Feb 2021 23:04:51 +0100
Subject: Use mmap file allocator in dense tensor attribute if huge flag is set
 in config.

---
 .../tensorattribute/tensorattribute_test.cpp       | 23 ++++++++++++++++++++++
 .../dense_tensor_store/dense_tensor_store_test.cpp |  3 ++-
 2 files changed, 25 insertions(+), 1 deletion(-)

(limited to 'searchlib/src/tests')

diff --git a/searchlib/src/tests/attribute/tensorattribute/tensorattribute_test.cpp b/searchlib/src/tests/attribute/tensorattribute/tensorattribute_test.cpp
index f1d910f2635..8efccf9b51b 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.setHuge(true);
+        }
         if (_traits.use_mock_index) {
             _index_factory = std::make_unique<MockNearestNeighborIndexFactory>();
         } else {
@@ -998,4 +1008,17 @@ TEST_F("NN blueprint handles strong filter triggering brute force search", Neare
     EXPECT_FALSE(bp->may_approximate());
 }
 
+TEST("Dense tensor attribute with huge 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 b5d4fb3130f..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
@@ -2,6 +2,7 @@
 #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>
@@ -27,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);
-- 
cgit v1.2.3