Merge pull request #16159 from vespa-engine/toregge/propagate-attribute-usage-to-resource-usage-tracker

Track attribute resource usage.

6 files changed, 325 insertions, 9 deletions

diff --git a/searchcore/src/tests/proton/persistenceengine/CMakeLists.txt b/searchcore/src/tests/proton/persistenceengine/CMakeLists.txt
index d9f82549d40..3480038144f 100644
--- a/searchcore/src/tests/proton/persistenceengine/CMakeLists.txt
+++ b/searchcore/src/tests/proton/persistenceengine/CMakeLists.txt
@@ -4,6 +4,7 @@ vespa_add_executable(searchcore_persistenceengine_test_app TEST
     persistenceengine_test.cpp
     DEPENDS
     searchcore_persistenceengine
+    searchcore_attribute
     searchcore_pcommon
     searchcore_proton_metrics
 )
diff --git a/searchcore/src/tests/proton/persistenceengine/resource_usage_tracker/CMakeLists.txt b/searchcore/src/tests/proton/persistenceengine/resource_usage_tracker/CMakeLists.txt
index d1a70c5be14..171724f2077 100644
--- a/searchcore/src/tests/proton/persistenceengine/resource_usage_tracker/CMakeLists.txt
+++ b/searchcore/src/tests/proton/persistenceengine/resource_usage_tracker/CMakeLists.txt
@@ -4,6 +4,7 @@ vespa_add_executable(searchcore_resource_usage_tracker_test_app TEST
     resource_usage_tracker_test.cpp
     DEPENDS
     searchcore_persistenceengine
+    searchcore_attribute
     GTest::GTest
 )
 vespa_add_test(NAME searchcore_resource_usage_tracker_test_app COMMAND searchcore_resource_usage_tracker_test_app)
diff --git a/searchcore/src/tests/proton/persistenceengine/resource_usage_tracker/resource_usage_tracker_test.cpp b/searchcore/src/tests/proton/persistenceengine/resource_usage_tracker/resource_usage_tracker_test.cpp
index 7391f9fb3a2..4144b87c61d 100644
--- a/searchcore/src/tests/proton/persistenceengine/resource_usage_tracker/resource_usage_tracker_test.cpp
+++ b/searchcore/src/tests/proton/persistenceengine/resource_usage_tracker/resource_usage_tracker_test.cpp
@@ -2,13 +2,18 @@
 
 #include <vespa/persistence/spi/resource_usage_listener.h>
 #include <vespa/persistence/spi/resource_usage.h>
+#include <vespa/searchcore/proton/attribute/attribute_usage_stats.h>
+#include <vespa/searchcore/proton/attribute/i_attribute_usage_listener.h>
 #include <vespa/searchcore/proton/persistenceengine/resource_usage_tracker.h>
 #include <vespa/searchcore/proton/test/disk_mem_usage_notifier.h>
 #include <vespa/vespalib/gtest/gtest.h>
 #include <vespa/vespalib/util/idestructorcallback.h>
+#include <atomic>
 
+using storage::spi::AttributeResourceUsage;
 using storage::spi::ResourceUsage;
 using proton::test::DiskMemUsageNotifier;
+using proton::AttributeUsageStats;
 using proton::DiskMemUsageState;
 using proton::ResourceUsageTracker;
 
@@ -16,9 +21,19 @@ namespace {
 
 struct MyResourceUsageListener : public storage::spi::ResourceUsageListener
 {
-    using storage::spi::ResourceUsageListener::ResourceUsageListener;
+    std::atomic<size_t> _update_count;
 
-    std::vector<double> get_usage_vector() const { return { get_usage().get_disk_usage(), get_usage().get_memory_usage() }; }
+    MyResourceUsageListener()
+        : storage::spi::ResourceUsageListener(),
+          _update_count(0u)
+    {
+    }
+
+    void update_resource_usage(const ResourceUsage& resource_usage) override {
+        storage::spi::ResourceUsageListener::update_resource_usage(resource_usage);
+        ++_update_count;
+    }
+    size_t get_update_count() const { return _update_count; }
 };
 
 }
@@ -46,17 +61,19 @@ public:
         _notifier.notify(DiskMemUsageState({ 0.8, disk_usage }, { 0.8, memory_usage }));
     }
 
+    ResourceUsage get_usage() { return _listener->get_usage(); }
+    size_t get_update_count() const { return _listener->get_update_count(); }
 };
 
 ResourceUsageTrackerTest::~ResourceUsageTrackerTest() = default;
 
 TEST_F(ResourceUsageTrackerTest, resource_usage_is_forwarded_to_listener)
 {
-    EXPECT_EQ((std::vector<double>{ 0.0, 0.0 }), _listener->get_usage_vector());
+    EXPECT_EQ(ResourceUsage(0.0, 0.0), get_usage());
     auto register_guard = _tracker->set_listener(*_listener);
-    EXPECT_EQ((std::vector<double>{ 0.5, 0.4 }), _listener->get_usage_vector());
+    EXPECT_EQ(ResourceUsage(0.5, 0.4), get_usage());
     notify(0.75, 0.25);
-    EXPECT_EQ((std::vector<double>{ 0.75, 0.25 }), _listener->get_usage_vector());
+    EXPECT_EQ(ResourceUsage(0.75, 0.25), get_usage());
 }
 
 TEST_F(ResourceUsageTrackerTest, forwarding_depends_on_register_guard)
@@ -64,14 +81,14 @@ TEST_F(ResourceUsageTrackerTest, forwarding_depends_on_register_guard)
     auto register_guard = _tracker->set_listener(*_listener);
     register_guard.reset();
     notify(0.75, 0.25);
-    EXPECT_EQ((std::vector<double>{ 0.5, 0.4 }), _listener->get_usage_vector());
+    EXPECT_EQ(ResourceUsage(0.5, 0.4), get_usage());
 }
 
 TEST_F(ResourceUsageTrackerTest, no_forwarding_to_deleted_listener)
 {
     _listener->set_register_guard(_tracker->set_listener(*_listener));
     notify(0.75, 0.25);
-    EXPECT_EQ((std::vector<double>{ 0.75, 0.25 }), _listener->get_usage_vector());
+    EXPECT_EQ(ResourceUsage(0.75, 0.25), get_usage());
     _listener.reset();
     notify(0.2, 0.1);
 }
@@ -82,4 +99,119 @@ TEST_F(ResourceUsageTrackerTest, register_guard_handles_deleted_tracker)
     _tracker.reset();
 }
 
+namespace {
+
+struct NamedAttribute
+{
+    vespalib::string subdb;
+    vespalib::string attribute;
+
+    NamedAttribute(const vespalib::string& subdb_in, const vespalib::string& attribute_in)
+        : subdb(subdb_in),
+          attribute(attribute_in)
+    {
+    }
+};
+
+NamedAttribute ready_a1("0.ready", "a1");
+NamedAttribute notready_a1("2.notready", "a1");
+NamedAttribute ready_a2("0.ready", "a2");
+
+constexpr size_t usage_limit = 1024;
+
+struct AttributeUsageStatsBuilder
+{
+    AttributeUsageStats stats;
+
+    AttributeUsageStatsBuilder()
+        : stats()
+    {
+    }
+
+    ~AttributeUsageStatsBuilder();
+
+    AttributeUsageStatsBuilder& reset() { stats = AttributeUsageStats(); return *this; }
+    AttributeUsageStatsBuilder& merge(const NamedAttribute& named_attribute, size_t used_enum_store, size_t used_multivalue);
+
+    AttributeUsageStats build() { return stats; }
+
+};
+
+AttributeUsageStatsBuilder::~AttributeUsageStatsBuilder() = default;
+
+AttributeUsageStatsBuilder&
+AttributeUsageStatsBuilder::merge(const NamedAttribute& named_attribute, size_t used_enum_store, size_t used_multivalue)
+{
+    vespalib::AddressSpace enum_store_usage(used_enum_store, 0, usage_limit);
+    vespalib::AddressSpace multivalue_usage(used_multivalue, 0, usage_limit);
+    search::AddressSpaceUsage as_usage(enum_store_usage, multivalue_usage);
+    stats.merge(as_usage, named_attribute.attribute, named_attribute.subdb);
+    return *this;
+}
+
+double rel_usage(size_t usage) noexcept {
+    return (double) usage / (double) usage_limit;
+}
+
+ResourceUsage make_resource_usage(const vespalib::string& enum_store_name, size_t used_enum_store, const vespalib::string &multivalue_name, size_t used_multivalue)
+{
+    AttributeResourceUsage enum_store_usage(rel_usage(used_enum_store), enum_store_name);
+    AttributeResourceUsage multivalue_usage(rel_usage(used_multivalue), multivalue_name);
+    return ResourceUsage(0.0, 0.0, enum_store_usage, multivalue_usage);
+}
+
+}
+
+TEST_F(ResourceUsageTrackerTest, aggregates_attribute_usage)
+{
+    notify(0.0, 0.0);
+    auto register_guard = _tracker->set_listener(*_listener);
+    auto aul1 = _tracker->make_attribute_usage_listener("doctype1");
+    auto aul2 = _tracker->make_attribute_usage_listener("doctype2");
+    AttributeUsageStatsBuilder b1;
+    AttributeUsageStatsBuilder b2;
+    b1.merge(ready_a1, 10, 20).merge(ready_a2, 5, 30);
+    b2.merge(ready_a1, 15, 15);
+    aul1->notify_attribute_usage(b1.build());
+    aul2->notify_attribute_usage(b2.build());
+    EXPECT_EQ(make_resource_usage("doctype2.0.ready.a1", 15, "doctype1.0.ready.a2", 30), get_usage());
+    b1.merge(notready_a1, 5, 31);
+    aul1->notify_attribute_usage(b1.build());
+    EXPECT_EQ(make_resource_usage("doctype2.0.ready.a1", 15, "doctype1.2.notready.a1", 31), get_usage());
+    b1.reset().merge(ready_a1, 10, 20).merge(ready_a2, 5, 30);
+    aul1->notify_attribute_usage(b1.build());
+    EXPECT_EQ(make_resource_usage("doctype2.0.ready.a1", 15, "doctype1.0.ready.a2", 30), get_usage());
+    aul2.reset();
+    EXPECT_EQ(make_resource_usage("doctype1.0.ready.a1", 10, "doctype1.0.ready.a2", 30), get_usage());
+    aul1.reset();
+    EXPECT_EQ(make_resource_usage("", 0, "", 0), get_usage());
+    aul2 = _tracker->make_attribute_usage_listener("doctype2");
+    aul2->notify_attribute_usage(b2.build());
+    EXPECT_EQ(make_resource_usage("doctype2.0.ready.a1", 15, "doctype2.0.ready.a1", 15), get_usage());
+}
+
+TEST_F(ResourceUsageTrackerTest, can_skip_scan_when_aggregating_attributes)
+{
+    notify(0.0, 0.0);
+    auto register_guard = _tracker->set_listener(*_listener);
+    auto aul1 = _tracker->make_attribute_usage_listener("doctype1");
+    auto aul2 = _tracker->make_attribute_usage_listener("doctype2");
+    AttributeUsageStatsBuilder b1;
+    AttributeUsageStatsBuilder b2;
+    b1.merge(ready_a1, 20, 20).merge(ready_a2, 5, 30);
+    b2.merge(ready_a1, 15, 15);
+    aul1->notify_attribute_usage(b1.build());
+    EXPECT_EQ(make_resource_usage("doctype1.0.ready.a1", 20, "doctype1.0.ready.a2", 30), get_usage());
+    EXPECT_EQ(2u, get_update_count());
+    aul1->notify_attribute_usage(b1.build());
+    EXPECT_EQ(make_resource_usage("doctype1.0.ready.a1", 20, "doctype1.0.ready.a2", 30), get_usage());
+    EXPECT_EQ(2u, get_update_count()); // usage for doctype1 has not changed
+    aul2->notify_attribute_usage(b2.build());
+    EXPECT_EQ(make_resource_usage("doctype1.0.ready.a1", 20, "doctype1.0.ready.a2", 30), get_usage());
+    EXPECT_EQ(2u, get_update_count()); // usage for doctype2 is less than usage for doctype1
+    aul2.reset();
+    aul1.reset();
+    EXPECT_EQ(4u, get_update_count()); // never skip scan when removing document type
+}
+
 GTEST_MAIN_RUN_ALL_TESTS()
diff --git a/searchcore/src/vespa/searchcore/proton/attribute/i_attribute_usage_listener.h b/searchcore/src/vespa/searchcore/proton/attribute/i_attribute_usage_listener.h
new file mode 100644
index 00000000000..6c7ae1248b8
--- /dev/null
+++ b/searchcore/src/vespa/searchcore/proton/attribute/i_attribute_usage_listener.h
@@ -0,0 +1,17 @@
+// Copyright Verizon Media. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+
+#pragma once
+
+namespace proton {
+
+/*
+ * Interface class for listening to attribute usage changes.
+ */
+class IAttributeUsageListener
+{
+public:
+    virtual ~IAttributeUsageListener() = default;
+    virtual void notify_attribute_usage(const AttributeUsageStats &attribute_usage) = 0;
+};
+
+}
diff --git a/searchcore/src/vespa/searchcore/proton/persistenceengine/resource_usage_tracker.cpp b/searchcore/src/vespa/searchcore/proton/persistenceengine/resource_usage_tracker.cpp
index 5636666cfbf..9791b55f7b7 100644
--- a/searchcore/src/vespa/searchcore/proton/persistenceengine/resource_usage_tracker.cpp
+++ b/searchcore/src/vespa/searchcore/proton/persistenceengine/resource_usage_tracker.cpp
@@ -1,12 +1,15 @@
 // Copyright Verizon Media. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 
 #include "resource_usage_tracker.h"
+#include <vespa/searchcore/proton/attribute/i_attribute_usage_listener.h>
 #include <vespa/searchcore/proton/server/disk_mem_usage_state.h>
 #include <vespa/searchcore/proton/server/i_disk_mem_usage_notifier.h>
 #include <vespa/persistence/spi/i_resource_usage_listener.h>
 #include <vespa/vespalib/util/idestructorcallback.h>
+#include <vespa/vespalib/stllike/hash_map.hpp>
 #include <cassert>
 
+using storage::spi::AttributeResourceUsage;
 using storage::spi::ResourceUsage;
 
 namespace proton {
@@ -32,13 +35,46 @@ ResourceUsageTracker::ListenerGuard::~ListenerGuard()
     }
 }
 
+class ResourceUsageTracker::AttributeUsageListener : public IAttributeUsageListener
+{
+    std::shared_ptr<ResourceUsageTracker> _tracker;
+    vespalib::string                      _document_type;
+
+public:
+    AttributeUsageListener(std::shared_ptr<ResourceUsageTracker> tracker, const vespalib::string& document_type);
+
+    ~AttributeUsageListener() override;
+    void notify_attribute_usage(const AttributeUsageStats &attribute_usage) override;
+};
+
+ResourceUsageTracker::AttributeUsageListener::AttributeUsageListener(std::shared_ptr<ResourceUsageTracker> tracker, const vespalib::string &document_type)
+    : IAttributeUsageListener(),
+      _tracker(std::move(tracker)),
+      _document_type(document_type)
+{
+}
+
+ResourceUsageTracker::AttributeUsageListener::~AttributeUsageListener()
+{
+    _tracker->remove_document_type(_document_type);
+}
+
+void
+ResourceUsageTracker::AttributeUsageListener::notify_attribute_usage(const AttributeUsageStats &attribute_usage)
+{
+    _tracker->notify_attribute_usage(_document_type, attribute_usage);
+}
+
 ResourceUsageTracker::ResourceUsageTracker(IDiskMemUsageNotifier& disk_mem_usage_notifier)
     : std::enable_shared_from_this<ResourceUsageTracker>(),
       IDiskMemUsageListener(),
       _lock(),
       _resource_usage(),
       _listener(nullptr),
-      _disk_mem_usage_notifier(disk_mem_usage_notifier)
+      _disk_mem_usage_notifier(disk_mem_usage_notifier),
+      _attribute_usage(),
+      _attribute_enum_store_max_document_type(),
+      _attribute_multivalue_max_document_type()
 {
     _disk_mem_usage_notifier.addDiskMemUsageListener(this);
 }
@@ -52,7 +88,7 @@ void
 ResourceUsageTracker::notifyDiskMemUsage(DiskMemUsageState state)
 {
     std::lock_guard guard(_lock);
-    _resource_usage = ResourceUsage(state.diskState().usage(), state.memoryState().usage());
+    _resource_usage = ResourceUsage(state.diskState().usage(), state.memoryState().usage(), _resource_usage.get_attribute_enum_store_usage(), _resource_usage.get_attribute_multivalue_usage());
     if (_listener != nullptr) {
         _listener->update_resource_usage(_resource_usage);
     }
@@ -75,5 +111,123 @@ ResourceUsageTracker::remove_listener()
     _listener = nullptr;
 }
 
+void
+ResourceUsageTracker::remove_document_type(const vespalib::string &document_type)
+{
+    std::lock_guard guard(_lock);
+    _attribute_usage.erase(document_type);
+    if (scan_attribute_usage(true, guard) && _listener != nullptr) {
+        _listener->update_resource_usage(_resource_usage);
+    }
+}
+
+namespace {
+
+bool same_usage(const AddressSpaceUsageStats &lhs, const AddressSpaceUsageStats &rhs) {
+    return ((lhs.getUsage().usage() == rhs.getUsage().usage()) &&
+            (lhs.getAttributeName() == rhs.getAttributeName()) &&
+            (lhs.getSubDbName() == rhs.getSubDbName()));
+}
+
+bool can_skip_scan(double max, double old_max, bool same_document_type) noexcept {
+    return (!same_document_type && (max <= old_max));
+}
+
+}
+
+void
+ResourceUsageTracker::notify_attribute_usage(const vespalib::string &document_type, const AttributeUsageStats &attribute_usage)
+{
+    std::lock_guard guard(_lock);
+    auto& old_usage = _attribute_usage[document_type];
+    if (same_usage(old_usage.enumStoreUsage(), attribute_usage.enumStoreUsage()) && 
+        same_usage(old_usage.multiValueUsage(), attribute_usage.multiValueUsage())) {
+        return; // usage for document type has not changed
+    }
+    old_usage = attribute_usage;
+    double enum_store_max = attribute_usage.enumStoreUsage().getUsage().usage();
+    double multivalue_max = attribute_usage.multiValueUsage().getUsage().usage();
+    double old_enum_store_max = _resource_usage.get_attribute_enum_store_usage().get_usage();
+    double old_multivalue_max = _resource_usage.get_attribute_multivalue_usage().get_usage();
+
+    if (can_skip_scan(enum_store_max, old_enum_store_max, document_type == _attribute_enum_store_max_document_type) &&
+        can_skip_scan(multivalue_max, old_multivalue_max, document_type == _attribute_multivalue_max_document_type)) {
+        return; // usage for document type is less than or equal to usage for other document types
+    }
+    if (scan_attribute_usage(false, guard) && _listener != nullptr) {
+        _listener->update_resource_usage(_resource_usage);
+    }
+}
+
+namespace {
+
+class MaxAttributeUsage
+{
+    const AddressSpaceUsageStats* _max;
+    const vespalib::string*       _document_type;
+    double                        _max_usage;
+
+    vespalib::string get_name() const {
+        return *_document_type + "." + _max->getSubDbName() + "." + _max->getAttributeName();
+    }
+
+public:
+    MaxAttributeUsage()
+        : _max(nullptr),
+          _document_type(nullptr),
+          _max_usage(0.0)
+    {
+    }
+
+    void sample(const vespalib::string& document_type, const AddressSpaceUsageStats& usage) {
+        if (_max == nullptr || usage.getUsage().usage() > _max_usage) {
+            _max = &usage;
+            _document_type = &document_type;
+            _max_usage = usage.getUsage().usage();
+        }
+    }
+    
+    AttributeResourceUsage get_max_resource_usage() {
+        if (_max != nullptr) {
+            return AttributeResourceUsage(_max_usage, get_name());
+        } else {
+            return AttributeResourceUsage();
+        }
+    }
+
+    const vespalib::string get_document_type() const { return _document_type != nullptr ? *_document_type : ""; }
 };
 
+}
+
+bool
+ResourceUsageTracker::scan_attribute_usage(bool force_changed, std::lock_guard<std::mutex>&)
+{
+    MaxAttributeUsage enum_store_max;
+    MaxAttributeUsage multivalue_max;
+    for (const auto& kv : _attribute_usage) {
+        enum_store_max.sample(kv.first, kv.second.enumStoreUsage());
+        multivalue_max.sample(kv.first, kv.second.multiValueUsage());
+    }
+    ResourceUsage new_resource_usage(_resource_usage.get_disk_usage(),
+                                     _resource_usage.get_memory_usage(),
+                                     enum_store_max.get_max_resource_usage(),
+                                     multivalue_max.get_max_resource_usage());
+
+    bool changed = (new_resource_usage != _resource_usage) ||
+                   force_changed;
+    if (changed) {
+        _resource_usage = std::move(new_resource_usage);
+        _attribute_enum_store_max_document_type = enum_store_max.get_document_type();
+        _attribute_multivalue_max_document_type = multivalue_max.get_document_type();
+    }
+    return changed;
+}
+
+std::unique_ptr<IAttributeUsageListener>
+ResourceUsageTracker::make_attribute_usage_listener(const vespalib::string &document_type)
+{
+    return std::make_unique<AttributeUsageListener>(shared_from_this(), document_type);
+}
+
+};
diff --git a/searchcore/src/vespa/searchcore/proton/persistenceengine/resource_usage_tracker.h b/searchcore/src/vespa/searchcore/proton/persistenceengine/resource_usage_tracker.h
index e41435a1a83..879e60aeeac 100644
--- a/searchcore/src/vespa/searchcore/proton/persistenceengine/resource_usage_tracker.h
+++ b/searchcore/src/vespa/searchcore/proton/persistenceengine/resource_usage_tracker.h
@@ -2,8 +2,10 @@
 
 #pragma once
 
+#include <vespa/searchcore/proton/attribute/attribute_usage_stats.h>
 #include <vespa/searchcore/proton/server/i_disk_mem_usage_listener.h>
 #include <vespa/persistence/spi/resource_usage.h>
+#include <vespa/vespalib/stllike/hash_map.h>
 #include <mutex>
 #include <memory>
 #include <vector>
@@ -14,6 +16,7 @@ namespace vespalib { class IDestructorCallback; }
 namespace proton {
 
 class DiskMemUsageState;
+class IAttributeUsageListener;
 class IDiskMemUsageNotifier;
 
 /*
@@ -22,16 +25,24 @@ class IDiskMemUsageNotifier;
 class ResourceUsageTracker : public std::enable_shared_from_this<ResourceUsageTracker>, public IDiskMemUsageListener
 {
     class ListenerGuard;
+    class AttributeUsageListener;
     std::mutex                  _lock;
     storage::spi::ResourceUsage _resource_usage;
     storage::spi::IResourceUsageListener* _listener;
     IDiskMemUsageNotifier&      _disk_mem_usage_notifier;
+    vespalib::hash_map<vespalib::string, AttributeUsageStats> _attribute_usage;
+    vespalib::string            _attribute_enum_store_max_document_type;
+    vespalib::string            _attribute_multivalue_max_document_type;
     void remove_listener();
+    void remove_document_type(const vespalib::string &document_type);
+    void notify_attribute_usage(const vespalib::string &document_type, const AttributeUsageStats &attribute_usage);
+    bool scan_attribute_usage(bool force_changed, std::lock_guard<std::mutex>&);
 public:
     ResourceUsageTracker(IDiskMemUsageNotifier& notifier);
     ~ResourceUsageTracker() override;
     void notifyDiskMemUsage(DiskMemUsageState state) override;
     std::unique_ptr<vespalib::IDestructorCallback> set_listener(storage::spi::IResourceUsageListener& listener);
+    std::unique_ptr<IAttributeUsageListener> make_attribute_usage_listener(const vespalib::string &document_type);
 };
 
 }