Add support to unregister from lifecycle notifications

include/envoy/server/lifecycle_notifier.h

@@ -30,6 +30,15 @@ class ServerLifecycleNotifier {
     ShutdownExit
   };
 
+  // A handle to a registration that can be used to unregister.
+  class Handle {
+  public:
+    virtual ~Handle() = default;
+
+    virtual void unregister() PURE;
+  };
+  using HandlePtr = std::unique_ptr<Handle>;
+
   /**
    * Callback invoked when the server reaches a certain lifecycle stage.
    *
@@ -49,8 +58,8 @@ class ServerLifecycleNotifier {
    * The second version which takes a completion back is currently only supported
    * for the ShutdownExit stage.
    */
-  virtual void registerCallback(Stage stage, StageCallback callback) PURE;
-  virtual void registerCallback(Stage stage, StageCallbackWithCompletion callback) PURE;
+  virtual HandlePtr registerCallback(Stage stage, StageCallback callback) PURE;
+  virtual HandlePtr registerCallback(Stage stage, StageCallbackWithCompletion callback) PURE;
 };
 
 } // namespace Server

source/server/config_validation/server.h

@@ -139,8 +139,12 @@ class ValidationInstance : Logger::Loggable<Logger::Id::main>,
   }
 
   // ServerLifecycleNotifier
-  void registerCallback(Stage, StageCallback) override {}
-  void registerCallback(Stage, StageCallbackWithCompletion) override {}
+  ServerLifecycleNotifier::HandlePtr registerCallback(Stage, StageCallback) override {
+    return nullptr;
+  }
+  ServerLifecycleNotifier::HandlePtr registerCallback(Stage, StageCallbackWithCompletion) override {
+    return nullptr;
+  }
 
 private:
   void initialize(const Options& options, Network::Address::InstanceConstSharedPtr local_address,

source/server/server.cc

@@ -547,13 +547,19 @@ void InstanceImpl::shutdownAdmin() {
   restarter_.sendParentTerminateRequest();
 }
 
-void InstanceImpl::registerCallback(Stage stage, StageCallback callback) {
-  stage_callbacks_[stage].push_back(callback);
+ServerLifecycleNotifier::HandlePtr InstanceImpl::registerCallback(Stage stage,
+                                                                  StageCallback callback) {
+  auto& callbacks = stage_callbacks_[stage];
+  return absl::make_unique<LifecycleCallbackHandle<LifecycleNotifierCallbacks>>(
+      callbacks, callbacks.insert(callbacks.end(), callback));
 }
 
-void InstanceImpl::registerCallback(Stage stage, StageCallbackWithCompletion callback) {
+ServerLifecycleNotifier::HandlePtr
+InstanceImpl::registerCallback(Stage stage, StageCallbackWithCompletion callback) {
   ASSERT(stage == Stage::ShutdownExit);
-  stage_completable_callbacks_[stage].push_back(callback);
+  auto& callbacks = stage_completable_callbacks_[stage];
+  return absl::make_unique<LifecycleCallbackHandle<LifecycleNotifierCompletionCallbacks>>(
+      callbacks, callbacks.insert(callbacks.end(), callback));
 }
 
 void InstanceImpl::notifyCallbacksForStage(Stage stage, Event::PostCb completion_cb) {

source/server/server.h

@@ -194,8 +194,9 @@ class InstanceImpl : Logger::Loggable<Logger::Id::main>,
   }
 
   // ServerLifecycleNotifier
-  void registerCallback(Stage stage, StageCallback callback) override;
-  void registerCallback(Stage stage, StageCallbackWithCompletion callback) override;
+  ServerLifecycleNotifier::HandlePtr registerCallback(Stage stage, StageCallback callback) override;
+  ServerLifecycleNotifier::HandlePtr
+  registerCallback(Stage stage, StageCallbackWithCompletion callback) override;
 
 private:
   ProtobufTypes::MessagePtr dumpBootstrapConfig();
@@ -260,8 +261,29 @@ class InstanceImpl : Logger::Loggable<Logger::Id::main>,
   Http::ContextImpl http_context_;
   std::unique_ptr<Memory::HeapShrinker> heap_shrinker_;
   const std::thread::id main_thread_id_;
-  absl::flat_hash_map<Stage, std::vector<StageCallback>> stage_callbacks_;
-  absl::flat_hash_map<Stage, std::vector<StageCallbackWithCompletion>> stage_completable_callbacks_;
+
+  using LifecycleNotifierCallbacks = std::list<StageCallback>;
+  using LifecycleNotifierCompletionCallbacks = std::list<StageCallbackWithCompletion>;
+
+  template <class T> class LifecycleCallbackHandle : public ServerLifecycleNotifier::Handle {
+  public:
+    LifecycleCallbackHandle(T& callbacks, typename T::iterator it)
+        : callbacks_(callbacks), it_(it) {}
+
+    void unregister() override {
+      ASSERT(!unregistered_);
+      unregistered_ = true;
+      callbacks_.erase(it_);
+    }
+
+  private:
+    bool unregistered_ = false;
+    T& callbacks_;
+    typename T::iterator it_;
+  };
+
+  absl::flat_hash_map<Stage, LifecycleNotifierCallbacks> stage_callbacks_;
+  absl::flat_hash_map<Stage, LifecycleNotifierCompletionCallbacks> stage_completable_callbacks_;
 };
 
 } // namespace Server

test/mocks/server/mocks.h

@@ -265,8 +265,9 @@ class MockServerLifecycleNotifier : public ServerLifecycleNotifier {
   MockServerLifecycleNotifier();
   ~MockServerLifecycleNotifier();
 
-  MOCK_METHOD2(registerCallback, void(Stage, StageCallback));
-  MOCK_METHOD2(registerCallback, void(Stage, StageCallbackWithCompletion));
+  MOCK_METHOD2(registerCallback, ServerLifecycleNotifier::HandlePtr(Stage, StageCallback));
+  MOCK_METHOD2(registerCallback,
+               ServerLifecycleNotifier::HandlePtr(Stage, StageCallbackWithCompletion));
 };
 
 class MockWorkerFactory : public WorkerFactory {

test/server/server_test.cc

@@ -198,6 +198,12 @@ TEST_P(ServerInstanceImplTest, LifecycleNotifications) {
                                 server_->dispatcher().post(completion_cb);
                                 completion_done.Notify();
                               });
+    auto handle =
+        server_->registerCallback(ServerLifecycleNotifier::Stage::Startup, [&] { FAIL(); });
+    handle->unregister();
+    handle = server_->registerCallback(ServerLifecycleNotifier::Stage::ShutdownExit,
+                                       [&](Event::PostCb) { FAIL(); });
+    handle->unregister();
     server_->run();
     server_ = nullptr;
     thread_local_ = nullptr;