Add support to unregister from lifecycle notifications

include/envoy/server/lifecycle_notifier.h

@@ -30,6 +30,13 @@ class ServerLifecycleNotifier {
     ShutdownExit
   };
 
+  // A handle to a callback registration. Deleting this handle will unregister the callback.
+  class Handle {
+  public:
+    virtual ~Handle() = default;
+  };
+  using HandlePtr = std::unique_ptr<Handle>;
+
   /**
    * Callback invoked when the server reaches a certain lifecycle stage.
    *
@@ -49,8 +56,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,23 @@ 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) {}
+    ~LifecycleCallbackHandle() override { callbacks_.erase(it_); }
+
+  private:
+    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

@@ -182,23 +182,32 @@ TEST_P(ServerInstanceImplTest, LifecycleNotifications) {
   // Run the server in a separate thread so we can test different lifecycle stages.
   auto server_thread = Thread::threadFactoryForTest().createThread([&] {
     initialize("test/server/node_bootstrap.yaml");
-    server_->registerCallback(ServerLifecycleNotifier::Stage::Startup, [&] {
+    auto handle1 = server_->registerCallback(ServerLifecycleNotifier::Stage::Startup, [&] {
       startup = true;
       started.Notify();
     });
-    server_->registerCallback(ServerLifecycleNotifier::Stage::ShutdownExit, [&] {
+    auto handle2 = server_->registerCallback(ServerLifecycleNotifier::Stage::ShutdownExit, [&] {
       shutdown = true;
       shutdown_begin.Notify();
     });
-    server_->registerCallback(ServerLifecycleNotifier::Stage::ShutdownExit,
-                              [&](Event::PostCb completion_cb) {
-                                // Block till we're told to complete
-                                completion_block.WaitForNotification();
-                                shutdown_with_completion = true;
-                                server_->dispatcher().post(completion_cb);
-                                completion_done.Notify();
-                              });
+    auto handle3 = server_->registerCallback(ServerLifecycleNotifier::Stage::ShutdownExit,
+                                             [&](Event::PostCb completion_cb) {
+                                               // Block till we're told to complete
+                                               completion_block.WaitForNotification();
+                                               shutdown_with_completion = true;
+                                               server_->dispatcher().post(completion_cb);
+                                               completion_done.Notify();
+                                             });
+    auto handle4 =
+        server_->registerCallback(ServerLifecycleNotifier::Stage::Startup, [&] { FAIL(); });
+    handle4 = server_->registerCallback(ServerLifecycleNotifier::Stage::ShutdownExit,
+                                        [&](Event::PostCb) { FAIL(); });
+    handle4 = nullptr;
+
     server_->run();
+    handle1 = nullptr;
+    handle2 = nullptr;
+    handle3 = nullptr;
     server_ = nullptr;
     thread_local_ = nullptr;
   });