Merge pull request #19026 from vespa-engine/arnej/new-removelocal-method

add extra method for use by unregister RPC api

3 files changed, 92 insertions, 4 deletions

diff --git a/slobrok/src/tests/local_rpc_monitor_map/local_rpc_monitor_map_test.cpp b/slobrok/src/tests/local_rpc_monitor_map/local_rpc_monitor_map_test.cpp
index 40f8baa32a9..a69a77d3640 100644
--- a/slobrok/src/tests/local_rpc_monitor_map/local_rpc_monitor_map_test.cpp
+++ b/slobrok/src/tests/local_rpc_monitor_map/local_rpc_monitor_map_test.cpp
@@ -240,16 +240,33 @@ TEST_F(LocalRpcMonitorMapTest, local_add_already_up) {
     ASSERT_TRUE(handler_deleted);
 }
 
-TEST_F(LocalRpcMonitorMapTest, local_add_already_down) {
+TEST_F(LocalRpcMonitorMapTest, local_add_unknown_comes_up) {
     std::unique_ptr<OkState> state;
     bool handler_deleted;
     add_mapping(mapping, false);
     map.addLocal(mapping, std::make_unique<MyAddLocalHandler>(state, handler_deleted));
-    monitor_log.expect({});
+    monitor_log.expect({MonitorCall::stop(mapping), MonitorCall::start(mapping, true)});
     map_log.expect({});
+    ASSERT_FALSE(state);
+    map.up(mapping);
     ASSERT_TRUE(state);
     EXPECT_TRUE(state->ok());
     ASSERT_TRUE(handler_deleted);
+    map_log.expect({MapCall::add(mapping)});
+}
+
+TEST_F(LocalRpcMonitorMapTest, local_add_unknown_goes_down) {
+    std::unique_ptr<OkState> state;
+    bool handler_deleted;
+    add_mapping(mapping, false);
+    map.addLocal(mapping, std::make_unique<MyAddLocalHandler>(state, handler_deleted));
+    monitor_log.expect({MonitorCall::stop(mapping), MonitorCall::start(mapping, true)});
+    map_log.expect({});
+    ASSERT_FALSE(state);
+    map.down(mapping);
+    ASSERT_TRUE(state);
+    EXPECT_FALSE(state->ok());
+    ASSERT_TRUE(handler_deleted);
 }
 
 TEST_F(LocalRpcMonitorMapTest, local_add_conflict) {
diff --git a/slobrok/src/vespa/slobrok/server/local_rpc_monitor_map.cpp b/slobrok/src/vespa/slobrok/server/local_rpc_monitor_map.cpp
index 73d355b7645..55d0e744743 100644
--- a/slobrok/src/vespa/slobrok/server/local_rpc_monitor_map.cpp
+++ b/slobrok/src/vespa/slobrok/server/local_rpc_monitor_map.cpp
@@ -10,6 +10,26 @@ namespace slobrok {
 
 #pragma GCC diagnostic ignored "-Winline"
 
+namespace {
+
+struct ChainedAddLocalCompletionHandler : LocalRpcMonitorMap::AddLocalCompletionHandler {
+    std::unique_ptr<AddLocalCompletionHandler> first;
+    std::unique_ptr<AddLocalCompletionHandler> second;
+
+    ChainedAddLocalCompletionHandler(std::unique_ptr<AddLocalCompletionHandler> f,
+                                     std::unique_ptr<AddLocalCompletionHandler> s)
+        : first(std::move(f)), second(std::move(s))
+    {}
+
+    void doneHandler(OkState result) override {
+        first->doneHandler(result);
+        second->doneHandler(result);
+    }
+    ~ChainedAddLocalCompletionHandler() override {}
+};
+
+}
+
 void LocalRpcMonitorMap::DelayedTasks::PerformTask() {
     std::vector<Event> todo;
     std::swap(todo, _queue);
@@ -89,11 +109,22 @@ void LocalRpcMonitorMap::addLocal(const ServiceMapping &mapping,
         mapping.name.c_str(), mapping.spec.c_str());
     auto old = _map.find(mapping.name);
     if (old != _map.end()) {
-        const PerService & exists = old->second;
+        PerService & exists = old->second;
         if (exists.spec == mapping.spec) {
             LOG(debug, "added mapping %s->%s was already present",
                 mapping.name.c_str(), mapping.spec.c_str());
-            inflight->doneHandler(OkState(0, "already registered"));
+            if (exists.up) {
+                inflight->doneHandler(OkState(0, "already registered"));
+            } else if (exists.inflight) {
+                auto newInflight = std::make_unique<ChainedAddLocalCompletionHandler>(
+                    std::move(exists.inflight),
+                    std::move(inflight));
+                exists.inflight = std::move(newInflight);
+            } else {
+                _mappingMonitor->stop(mapping);
+                exists.inflight = std::move(inflight);
+                _mappingMonitor->start(mapping, true);
+            }
             return;
         }
         LOG(warning, "tried addLocal for mapping %s->%s, but already had conflicting mapping %s->%s",
@@ -105,6 +136,43 @@ void LocalRpcMonitorMap::addLocal(const ServiceMapping &mapping,
     addToMap(mapping, localService(mapping, std::move(inflight)), true);
 }
 
+void LocalRpcMonitorMap::removeLocal(const ServiceMapping &mapping) {
+    LOG(debug, "try local remove: mapping %s->%s",
+        mapping.name.c_str(), mapping.spec.c_str());
+    auto old = _map.find(mapping.name);
+    if (old == _map.end()) {
+        return; // already removed, OK
+    }
+    PerService & exists = old->second;
+    if (exists.spec != mapping.spec) {
+        LOG(warning, "tried removeLocal for mapping %s->%s, but already had conflicting mapping %s->%s",
+            mapping.name.c_str(), mapping.spec.c_str(),
+            mapping.name.c_str(), exists.spec.c_str());
+        return; // unregister for old, conflicting mapping
+    }
+    if (exists.localOnly) {
+        // we can just remove it
+        auto removed = removeFromMap(old);
+        if (removed.inflight) {
+            auto target = std::move(removed.inflight);
+            target->doneHandler(OkState(13, "removed during initialization"));
+        }
+        if (removed.up) {
+            _dispatcher.remove(removed.mapping);            
+        }
+        return;
+    }
+    // also exists in consensus map, so we can't just remove it
+    // instead, pretend it's down and delay next ping
+    _mappingMonitor->stop(mapping);
+    if (exists.up) {
+        exists.up = false;
+        _dispatcher.remove(mapping);            
+    }
+    _mappingMonitor->start(mapping, false);
+    return;
+}
+
 void LocalRpcMonitorMap::add(const ServiceMapping &mapping) {
     _delayedTasks.handleLater(Event::add(mapping));
 }
diff --git a/slobrok/src/vespa/slobrok/server/local_rpc_monitor_map.h b/slobrok/src/vespa/slobrok/server/local_rpc_monitor_map.h
index 33a1109d915..920d54a405f 100644
--- a/slobrok/src/vespa/slobrok/server/local_rpc_monitor_map.h
+++ b/slobrok/src/vespa/slobrok/server/local_rpc_monitor_map.h
@@ -133,6 +133,9 @@ public:
     void addLocal(const ServiceMapping &mapping,
                   std::unique_ptr<AddLocalCompletionHandler> inflight);
 
+    /** for use by unregister API */
+    void removeLocal(const ServiceMapping &mapping);
+
     void add(const ServiceMapping &mapping) override;
     void remove(const ServiceMapping &mapping) override;