Merge pull request #18692 from vespa-engine/arnej/add-union-service-map-use-proxy

Arnej/add union service map use proxy

7 files changed, 330 insertions, 0 deletions

diff --git a/slobrok/CMakeLists.txt b/slobrok/CMakeLists.txt
index 00ddf7296ca..153a029e74d 100644
--- a/slobrok/CMakeLists.txt
+++ b/slobrok/CMakeLists.txt
@@ -25,4 +25,5 @@ vespa_define_module(
     src/tests/standalone
     src/tests/startsome
     src/tests/startup
+    src/tests/union_service_map
 )
diff --git a/slobrok/src/tests/union_service_map/CMakeLists.txt b/slobrok/src/tests/union_service_map/CMakeLists.txt
new file mode 100644
index 00000000000..523294742f2
--- /dev/null
+++ b/slobrok/src/tests/union_service_map/CMakeLists.txt
@@ -0,0 +1,9 @@
+# Copyright Verizon Media. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+vespa_add_executable(slobrok_union_service_map_test_app TEST
+    SOURCES
+    union_service_map_test.cpp
+    DEPENDS
+    slobrok_slobrokserver
+    GTest::GTest
+)
+vespa_add_test(NAME slobrok_union_service_map_test_app COMMAND slobrok_union_service_map_test_app)
diff --git a/slobrok/src/tests/union_service_map/union_service_map_test.cpp b/slobrok/src/tests/union_service_map/union_service_map_test.cpp
new file mode 100644
index 00000000000..5f1f70fb9fb
--- /dev/null
+++ b/slobrok/src/tests/union_service_map/union_service_map_test.cpp
@@ -0,0 +1,180 @@
+// Copyright Verizon Media. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+
+#include <vespa/slobrok/server/union_service_map.h>
+#include <vespa/vespalib/gtest/gtest.h>
+#include <vespa/vespalib/util/stringfmt.h>
+
+using namespace vespalib;
+using namespace slobrok;
+using vespalib::make_string_short::fmt;
+
+enum class Event { NONE, ADD, REMOVE, UPDATE };
+
+struct MapObserver : public MapListener {
+    MapObserver();
+    virtual ~MapObserver();
+    void add(const ServiceMapping &mapping) override;
+    void remove(const ServiceMapping &mapping) override;
+    void update(const ServiceMapping &old_mapping,
+                const ServiceMapping &new_mapping) override;
+
+    Event last_event = Event::NONE;
+    ServiceMapping last_add = {{}, {}};
+    ServiceMapping last_remove = {{}, {}};
+
+    void clear() { last_event = Event::NONE; }
+};
+
+MapObserver::MapObserver() = default;
+MapObserver::~MapObserver() = default;
+
+void MapObserver::add(const ServiceMapping &mapping) {
+    last_event = Event::ADD;
+    last_add = mapping;
+}
+
+void MapObserver::remove(const ServiceMapping &mapping) {
+    last_event = Event::REMOVE;
+    last_remove = mapping;
+}
+
+void MapObserver::update(const ServiceMapping &old_mapping,
+                         const ServiceMapping &new_mapping)
+{
+    last_event = Event::UPDATE;
+    last_remove = old_mapping;
+    last_add = new_mapping;
+}
+
+TEST(UnionServiceMapTest, forwards_simple_requests) {
+    ProxyMapSource source;
+    UnionServiceMap unionizer;
+    MapObserver observer;
+    unionizer.registerListener(observer);
+    source.registerListener(unionizer);
+
+    EXPECT_EQ(observer.last_event, Event::NONE);
+
+    ServiceMapping one{"foo/1", "bar/1"};
+    source.add(one);
+    EXPECT_EQ(observer.last_event, Event::ADD);
+    EXPECT_EQ(observer.last_add, one);
+    ServiceMapping two{"foo/2", "bar/2"};
+    source.add(two);
+    EXPECT_EQ(observer.last_event, Event::ADD);
+    EXPECT_EQ(observer.last_add, two);
+
+    source.remove(one);
+    EXPECT_EQ(observer.last_event, Event::REMOVE);
+    EXPECT_EQ(observer.last_remove, one);
+
+    ServiceMapping two_q{"foo/2", "qux/2"};
+    source.update(two, two_q);
+    // update implemented ass remove+add:
+    EXPECT_EQ(observer.last_event, Event::ADD);
+    EXPECT_EQ(observer.last_remove, two);
+    EXPECT_EQ(observer.last_add, two_q);
+}
+
+TEST(UnionServiceMapTest, handles_refcount) {
+    ProxyMapSource source1;
+    ProxyMapSource source2;
+    ProxyMapSource source3;
+    UnionServiceMap unionizer;
+    MapObserver observer;
+    unionizer.registerListener(observer);
+    source1.registerListener(unionizer);
+    source2.registerListener(unionizer);
+    source3.registerListener(unionizer);
+
+    EXPECT_EQ(observer.last_event, Event::NONE);
+    ServiceMapping one{"foo/1", "bar/1"};
+    source1.add(one);
+    EXPECT_EQ(observer.last_event, Event::ADD);
+    EXPECT_EQ(observer.last_add, one);
+    observer.clear();
+    EXPECT_EQ(observer.last_event, Event::NONE);
+    source2.add(one);
+    EXPECT_EQ(observer.last_event, Event::NONE);
+    source3.add(one);
+    EXPECT_EQ(observer.last_event, Event::NONE);
+    ServiceMapping two{"foo/2", "bar/2"};
+    source1.add(two);
+    EXPECT_EQ(observer.last_event, Event::ADD);
+    EXPECT_EQ(observer.last_add, two);
+    observer.clear();
+    EXPECT_EQ(observer.last_event, Event::NONE);
+    source2.add(two);
+    EXPECT_EQ(observer.last_event, Event::NONE);
+
+    source1.remove(one);
+    EXPECT_EQ(observer.last_event, Event::NONE);
+    source2.remove(one);
+    EXPECT_EQ(observer.last_event, Event::NONE);
+
+    source1.remove(two);
+    EXPECT_EQ(observer.last_event, Event::NONE);
+    source2.remove(two);
+    EXPECT_EQ(observer.last_event, Event::REMOVE);
+    EXPECT_EQ(observer.last_remove, two);
+
+    observer.clear();
+    EXPECT_EQ(observer.last_event, Event::NONE);
+    source3.remove(one);
+    EXPECT_EQ(observer.last_event, Event::REMOVE);
+    EXPECT_EQ(observer.last_remove, one);
+}
+
+TEST(UnionServiceMapTest, handles_conflicts) {
+    ProxyMapSource source1;
+    ProxyMapSource source2;
+    ProxyMapSource source3;
+    UnionServiceMap unionizer;
+    MapObserver observer;
+    unionizer.registerListener(observer);
+    source1.registerListener(unionizer);
+    source2.registerListener(unionizer);
+    source3.registerListener(unionizer);
+
+    EXPECT_EQ(observer.last_event, Event::NONE);
+    ServiceMapping one{"foo/1", "bar/1"};
+    source1.add(one);
+    EXPECT_EQ(observer.last_event, Event::ADD);
+    EXPECT_EQ(observer.last_add, one);
+    observer.clear();
+    source2.add(one);
+    EXPECT_EQ(observer.last_event, Event::NONE);
+
+    ServiceMapping two{"foo/2", "bar/2"};
+    source1.add(two);
+    EXPECT_EQ(observer.last_event, Event::ADD);
+    EXPECT_EQ(observer.last_add, two);
+    observer.clear();
+    source2.add(two);
+    EXPECT_EQ(observer.last_event, Event::NONE);
+
+    ServiceMapping one_q{"foo/1", "qux/1"};
+    source3.add(one_q);
+    EXPECT_EQ(observer.last_event, Event::REMOVE);
+    EXPECT_EQ(observer.last_remove, one);
+
+    ServiceMapping two_q{"foo/2", "qux/2"};
+    source3.add(two_q);
+    EXPECT_EQ(observer.last_event, Event::REMOVE);
+    EXPECT_EQ(observer.last_remove, two);
+
+    source3.remove(one_q);
+    EXPECT_EQ(observer.last_event, Event::ADD);
+    EXPECT_EQ(observer.last_add, one);
+
+    observer.clear();
+    source1.remove(two);
+    EXPECT_EQ(observer.last_event, Event::NONE);
+    source2.remove(two);
+    EXPECT_EQ(observer.last_event, Event::ADD);
+    EXPECT_EQ(observer.last_add, two_q);
+}
+
+
+GTEST_MAIN_RUN_ALL_TESTS()
+
diff --git a/slobrok/src/vespa/slobrok/server/CMakeLists.txt b/slobrok/src/vespa/slobrok/server/CMakeLists.txt
index 5b9059af932..153ce681f86 100644
--- a/slobrok/src/vespa/slobrok/server/CMakeLists.txt
+++ b/slobrok/src/vespa/slobrok/server/CMakeLists.txt
@@ -26,6 +26,7 @@ vespa_add_library(slobrok_slobrokserver
     service_map_history.cpp
     service_mapping.cpp
     slobrokserver.cpp
+    union_service_map.cpp
     INSTALL lib64
     DEPENDS
     slobrok
diff --git a/slobrok/src/vespa/slobrok/server/service_mapping.h b/slobrok/src/vespa/slobrok/server/service_mapping.h
index 6561c120284..6c540b8e9b5 100644
--- a/slobrok/src/vespa/slobrok/server/service_mapping.h
+++ b/slobrok/src/vespa/slobrok/server/service_mapping.h
@@ -12,6 +12,10 @@ struct ServiceMapping {
     vespalib::string spec;
     ServiceMapping(const vespalib::string & name_, const vespalib::string & spec_) noexcept : name(name_), spec(spec_) { }
     ~ServiceMapping();
+
+    bool operator== (const ServiceMapping &other) const {
+        return name == other.name && spec == other.spec;
+    }
 };
 
 typedef std::vector<ServiceMapping> ServiceMappingList;
diff --git a/slobrok/src/vespa/slobrok/server/union_service_map.cpp b/slobrok/src/vespa/slobrok/server/union_service_map.cpp
new file mode 100644
index 00000000000..f3d63f2a087
--- /dev/null
+++ b/slobrok/src/vespa/slobrok/server/union_service_map.cpp
@@ -0,0 +1,93 @@
+// Copyright Verizon Media. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+
+#include "union_service_map.h"
+#include <vespa/log/log.h>
+
+LOG_SETUP(".slobrok.server.union_service_map");
+
+namespace slobrok {
+
+UnionServiceMap::UnionServiceMap() = default;
+UnionServiceMap::~UnionServiceMap() = default;
+
+void UnionServiceMap::add(const ServiceMapping &mapping)
+{
+    const vespalib::string &key = mapping.name;
+    auto iter = _mappings.find(key);
+    if (iter == _mappings.end()) {
+        _mappings[key].emplace_back(mapping.spec, 1u);
+        _proxy.add(mapping);
+    } else {
+        Mappings &values = iter->second;
+        for (CountedSpec &old : values) {
+            if (old.spec == mapping.spec) {
+                ++old.count;
+                return;
+            }
+        }
+        if (values.size() == 1u) {
+            LOG(warning, "Multiple specs seen for name '%s', un-publishing", key.c_str());
+            _proxy.remove(ServiceMapping{key, values[0].spec});
+        }
+        values.emplace_back(mapping.spec, 1u);
+    }
+}
+
+void UnionServiceMap::remove(const ServiceMapping &mapping)
+{
+    const vespalib::string &key = mapping.name;
+    auto iter = _mappings.find(key);
+    if (iter == _mappings.end()) {
+        LOG(error, "Broken invariant: did not find %s in mappings", key.c_str());
+        return;
+    }
+    Mappings &values = iter->second;
+    bool found = false;
+    for (CountedSpec &old : values) {
+        if (old.spec == mapping.spec) {
+            if (--old.count > 0u) return;
+            found = true;
+        }
+    }
+    if (! found) {
+        LOG(error, "Broken invariant: did not find %s->%s in mappings",
+            key.c_str(), mapping.spec.c_str());
+        return;
+    }
+    size_t old_size = values.size();
+    std::erase_if(values, [] (const CountedSpec &v) { return v.count == 0; });
+    if (values.size() == 1u) {
+        LOG_ASSERT(old_size == 2u);
+        LOG(info, "Had multiple mappings for %s, but now only %s remains",
+            key.c_str(), values[0].spec.c_str());
+        _proxy.add(ServiceMapping{key, values[0].spec});
+    }
+    if (values.size() == 0u) {
+        LOG_ASSERT(old_size == 1u);
+        LOG(debug, "Last reference for %s -> %s removed",
+            key.c_str(), mapping.spec.c_str());
+        _proxy.remove(mapping);
+        _mappings.erase(iter);
+    }
+}
+
+void UnionServiceMap::update(const ServiceMapping &old_mapping,
+                             const ServiceMapping &new_mapping)
+{
+    LOG_ASSERT(old_mapping.name == new_mapping.name);
+    remove(old_mapping);
+    add(new_mapping);
+}
+
+void UnionServiceMap::registerListener(MapListener &listener)
+{
+    _proxy.registerListener(listener);
+}
+
+void UnionServiceMap::unregisterListener(MapListener &listener)
+{
+    _proxy.unregisterListener(listener);
+}
+
+} // namespace slobrok
+
diff --git a/slobrok/src/vespa/slobrok/server/union_service_map.h b/slobrok/src/vespa/slobrok/server/union_service_map.h
new file mode 100644
index 00000000000..b96167da26c
--- /dev/null
+++ b/slobrok/src/vespa/slobrok/server/union_service_map.h
@@ -0,0 +1,42 @@
+// Copyright Verizon Media. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+
+#pragma once
+
+#include "map_listener.h"
+#include "map_source.h"
+#include "proxy_map_source.h"
+
+#include <map>
+#include <vector>
+
+namespace slobrok {
+
+/**
+ * Listens to events from multiple maps and publishes the union of them.
+ **/
+class UnionServiceMap
+    : public MapSource, public MapListener
+{
+private:
+    struct CountedSpec {
+        vespalib::string spec;
+        size_t count;
+    };
+    using Mappings = std::vector<CountedSpec>;
+    std::map<vespalib::string, Mappings> _mappings;
+    ProxyMapSource _proxy;
+public:
+    UnionServiceMap();
+    virtual ~UnionServiceMap();
+
+    void add(const ServiceMapping &mapping) override;
+    void remove(const ServiceMapping &mapping) override;
+    void update(const ServiceMapping &old_mapping,
+                const ServiceMapping &new_mapping) override;
+
+    void registerListener(MapListener &listener) override;
+    void unregisterListener(MapListener &listener) override;
+};
+
+} // namespace slobrok
+