Merge pull request #13190 from vespa-engine/toregge/add-get-transient-memory-usage-to-attribute-initializer

Add get_transient_memory_usage() method to AttibuteInitializer.

3 files changed, 99 insertions, 6 deletions

diff --git a/searchcore/src/tests/proton/attribute/attribute_initializer/attribute_initializer_test.cpp b/searchcore/src/tests/proton/attribute/attribute_initializer/attribute_initializer_test.cpp
index b28e04f659f..264d6d397e1 100644
--- a/searchcore/src/tests/proton/attribute/attribute_initializer/attribute_initializer_test.cpp
+++ b/searchcore/src/tests/proton/attribute/attribute_initializer/attribute_initializer_test.cpp
@@ -27,6 +27,7 @@ namespace {
 const Config int32_sv(BasicType::Type::INT32);
 const Config int16_sv(BasicType::Type::INT16);
 const Config int32_array(BasicType::Type::INT32, CollectionType::Type::ARRAY);
+const Config int32_wset(BasicType::Type::INT32, CollectionType::Type::WSET);
 const Config string_wset(BasicType::Type::STRING, CollectionType::Type::WSET);
 const Config predicate(BasicType::Type::PREDICATE);
 const CollectionType wset2(CollectionType::Type::WSET, false, true);
@@ -48,6 +49,13 @@ Config getTensor(const vespalib::string &spec)
     return ret;
 }
 
+Config get_int32_wset_fs()
+{
+    Config ret(int32_wset);
+    ret.setFastSearch(true);
+    return ret;
+}
+
 void
 saveAttr(const vespalib::string &name, const Config &cfg, SerialNum serialNum, SerialNum createSerialNum)
 {
@@ -64,6 +72,12 @@ saveAttr(const vespalib::string &name, const Config &cfg, SerialNum serialNum, S
     av->addDoc(docId);
     assert(docId == 1u);
     av->clearDoc(docId);
+    if (cfg.basicType().type() == BasicType::Type::INT32 &&
+        cfg.collectionType().type() == CollectionType::Type::WSET) {
+        auto &iav = dynamic_cast<search::IntegerAttribute &>(*av);
+        iav.append(docId, 10, 1);
+        iav.append(docId, 11, 1);
+    }
     av->save();
     writer->markValidSnapshot(serialNum);
 }
@@ -198,6 +212,30 @@ TEST("require that too old attribute is not loaded")
     EXPECT_EQUAL(1u, av->getNumDocs());
 }
 
+TEST("require that transient memory usage is reported for first time posting list attribute load after enabling posting lists")
+{
+    saveAttr("a", int32_wset, 10, 2);
+    Fixture f;
+    auto avi = f.createInitializer({"a", get_int32_wset_fs()}, 5);
+    EXPECT_EQUAL(40u, avi->get_transient_memory_usage());
+}
+
+TEST("require that transient memory usage is reported for normal posting list attribute load")
+{
+    saveAttr("a", get_int32_wset_fs(), 10, 2);
+    Fixture f;
+    auto avi = f.createInitializer({"a", get_int32_wset_fs()}, 5);
+    EXPECT_EQUAL(24u, avi->get_transient_memory_usage());
+}
+
+TEST("require that transient memory usage is reported for attribute load without posting list")
+{
+    saveAttr("a", int32_wset, 10, 2);
+    Fixture f;
+    auto avi = f.createInitializer({"a", int32_wset}, 5);
+    EXPECT_EQUAL(0u, avi->get_transient_memory_usage());
+}
+
 }
 
 TEST_MAIN()
diff --git a/searchcore/src/vespa/searchcore/proton/attribute/attribute_initializer.cpp b/searchcore/src/vespa/searchcore/proton/attribute/attribute_initializer.cpp
index 0ccebd34a39..88e4d7a2191 100644
--- a/searchcore/src/vespa/searchcore/proton/attribute/attribute_initializer.cpp
+++ b/searchcore/src/vespa/searchcore/proton/attribute/attribute_initializer.cpp
@@ -8,9 +8,12 @@
 #include <vespa/vespalib/data/fileheader.h>
 #include <vespa/vespalib/stllike/asciistream.h>
 #include <vespa/vespalib/io/fileutil.h>
+#include <vespa/searchcommon/attribute/persistent_predicate_params.h>
 #include <vespa/searchlib/util/fileutil.h>
 #include <vespa/searchlib/attribute/attribute_header.h>
 #include <vespa/searchlib/attribute/attributevector.h>
+#include <vespa/searchlib/attribute/loadedenumvalue.h>
+#include <vespa/searchlib/attribute/loadedvalue.h>
 #include <vespa/fastos/file.h>
 
 #include <vespa/log/log.h>
@@ -149,16 +152,30 @@ logAttributeWrongType(const AttributeVector::SP &attr, const AttributeHeader &he
 
 }
 
+void
+AttributeInitializer::readHeader()
+{
+    if (!_attrDir->empty()) {
+        search::SerialNum serialNum = _attrDir->getFlushedSerialNum();
+        vespalib::string attrFileName = _attrDir->getAttributeFileName(serialNum);
+        if (serialNum != 0) {
+            _header = std::make_unique<const AttributeHeader>(extractHeader(attrFileName));
+            if (_header->getCreateSerialNum() <= _currentSerialNum && headerTypeOK(*_header, _spec.getConfig()) && (serialNum >= _currentSerialNum)) {
+                _header_ok = true;
+            }
+        }
+    }
+}
+
 AttributeVector::SP
 AttributeInitializer::tryLoadAttribute() const
 {
     search::SerialNum serialNum = _attrDir->getFlushedSerialNum();
     vespalib::string attrFileName = _attrDir->getAttributeFileName(serialNum);
     AttributeVector::SP attr = _factory.create(attrFileName, _spec.getConfig());
-    if (serialNum != 0) {
-        AttributeHeader header = extractHeader(attrFileName);
-        if (header.getCreateSerialNum() > _currentSerialNum || !headerTypeOK(header, attr->getConfig()) || (serialNum < _currentSerialNum)) {
-            setupEmptyAttribute(attr, serialNum, header);
+    if (serialNum != 0 && _header) {
+        if (!_header_ok) {
+            setupEmptyAttribute(attr, serialNum, *_header);
             return attr;
         }
         if (!loadAttribute(attr, serialNum)) {
@@ -223,8 +240,11 @@ AttributeInitializer::AttributeInitializer(const std::shared_ptr<AttributeDirect
       _documentSubDbName(documentSubDbName),
       _spec(spec),
       _currentSerialNum(currentSerialNum),
-      _factory(factory)
+      _factory(factory),
+      _header(),
+      _header_ok(false)
 {
+    readHeader();
 }
 
 AttributeInitializer::~AttributeInitializer() = default;
@@ -239,4 +259,35 @@ AttributeInitializer::init() const
     }
 }
 
+size_t
+AttributeInitializer::get_transient_memory_usage() const
+{
+    if (_header_ok) {
+        auto &header = *_header;
+        if (_spec.getConfig().fastSearch()) {
+            if (header.getEnumerated()) {
+                return sizeof(search::attribute::LoadedEnumAttribute) * header.get_total_value_count();
+            } else {
+                switch (_spec.getConfig().basicType().type()) {
+                case BasicType::Type::INT8:
+                    return sizeof(search::attribute::LoadedValue<int8_t>) * header.get_total_value_count();
+                case BasicType::Type::INT16:
+                    return sizeof(search::attribute::LoadedValue<int16_t>) * header.get_total_value_count();
+                case BasicType::Type::INT32:
+                    return sizeof(search::attribute::LoadedValue<int32_t>) * header.get_total_value_count();
+                case BasicType::Type::INT64:
+                    return sizeof(search::attribute::LoadedValue<int64_t>) * header.get_total_value_count();
+                case BasicType::Type::FLOAT:
+                    return sizeof(search::attribute::LoadedValue<float>) * header.get_total_value_count();
+                case BasicType::Type::DOUBLE:
+                    return sizeof(search::attribute::LoadedValue<double>) * header.get_total_value_count();
+                default:
+                    return 0u;
+                }
+            }
+        }
+    }
+    return 0u;
+}
+
 } // namespace proton
diff --git a/searchcore/src/vespa/searchcore/proton/attribute/attribute_initializer.h b/searchcore/src/vespa/searchcore/proton/attribute/attribute_initializer.h
index 988f4d06d16..a2277fd8134 100644
--- a/searchcore/src/vespa/searchcore/proton/attribute/attribute_initializer.h
+++ b/searchcore/src/vespa/searchcore/proton/attribute/attribute_initializer.h
@@ -6,7 +6,6 @@
 #include "attribute_initializer_result.h"
 #include <vespa/vespalib/stllike/string.h>
 #include <vespa/searchlib/common/serialnum.h>
-#include <vespa/searchcommon/attribute/persistent_predicate_params.h>
 
 namespace search::attribute { class AttributeHeader; }
 
@@ -30,6 +29,10 @@ private:
     const AttributeSpec             _spec;
     const uint64_t                  _currentSerialNum;
     const IAttributeFactory        &_factory;
+    std::unique_ptr<const search::attribute::AttributeHeader> _header;
+    bool                            _header_ok;
+
+    void readHeader();
 
     AttributeVectorSP tryLoadAttribute() const;
 
@@ -47,6 +50,7 @@ public:
 
     AttributeInitializerResult init() const;
     uint64_t getCurrentSerialNum() const { return _currentSerialNum; }
+    size_t get_transient_memory_usage() const;
 };
 
 } // namespace proton