dynamic_modules: add cluster lifecycle events to bootstrap (#43950)

## Description

This PR added server lifecycle callbacks including
`on_cluster_server_initialized`, `on_cluster_drain_started`, and
`on_cluster_shutdown` to the dynamic module cluster extension, enabling
modules to respond to server initialization, drain, and shutdown events.

---

**Commit Message:** dynamic_modules: added server lifecycle callbacks
for dynamic module clusters
**Additional Description:** Added server lifecycle callbacks to the
dynamic module cluster extension.
**Risk Level:** Low
**Testing:** Added Tests
**Docs Changes:** Added
**Release Notes:** N/A

Signed-off-by: Rohit Agrawal <rohit.agrawal@databricks.com>

Author: agrawroh

source/extensions/bootstrap/dynamic_modules/abi_impl.cc

@@ -575,4 +575,12 @@ bool envoy_dynamic_module_callback_bootstrap_extension_remove_admin_handler(
   return admin->removeHandler(prefix_str);
 }
 
+// -------------------- Cluster Lifecycle Callbacks --------------------
+
+bool envoy_dynamic_module_callback_bootstrap_extension_enable_cluster_lifecycle(
+    envoy_dynamic_module_type_bootstrap_extension_config_envoy_ptr extension_config_envoy_ptr) {
+  auto* config = static_cast<DynamicModuleBootstrapExtensionConfig*>(extension_config_envoy_ptr);
+  return config->enableClusterLifecycle();
+}
+
 } // extern "C"

source/extensions/bootstrap/dynamic_modules/extension_config.cc

@@ -48,6 +48,36 @@ void DynamicModuleBootstrapExtensionConfig::signalInitComplete() {
                    "traffic");
 }
 
+bool DynamicModuleBootstrapExtensionConfig::enableClusterLifecycle() {
+  if (cluster_lifecycle_enabled_) {
+    return false;
+  }
+  cluster_lifecycle_enabled_ = true;
+  cluster_update_callbacks_handle_ =
+      context_.clusterManager().addThreadLocalClusterUpdateCallbacks(*this);
+  // Register a shutdown callback to release the handle before the underlying TLS data is
+  // destroyed. The TLS shutdown happens in terminate() after ShutdownExit callbacks fire.
+  cluster_lifecycle_shutdown_handle_ = context_.lifecycleNotifier().registerCallback(
+      Server::ServerLifecycleNotifier::Stage::ShutdownExit,
+      [this]() { cluster_update_callbacks_handle_.reset(); });
+  return true;
+}
+
+void DynamicModuleBootstrapExtensionConfig::onClusterAddOrUpdate(
+    absl::string_view cluster_name, Upstream::ThreadLocalClusterCommand&) {
+  if (in_module_config_ != nullptr && on_bootstrap_extension_cluster_add_or_update_ != nullptr) {
+    on_bootstrap_extension_cluster_add_or_update_(thisAsVoidPtr(), in_module_config_,
+                                                  {cluster_name.data(), cluster_name.size()});
+  }
+}
+
+void DynamicModuleBootstrapExtensionConfig::onClusterRemoval(const std::string& cluster_name) {
+  if (in_module_config_ != nullptr && on_bootstrap_extension_cluster_removal_ != nullptr) {
+    on_bootstrap_extension_cluster_removal_(thisAsVoidPtr(), in_module_config_,
+                                            {cluster_name.data(), cluster_name.size()});
+  }
+}
+
 void DynamicModuleBootstrapExtensionConfig::onScheduled(uint64_t event_id) {
   if (in_module_config_ != nullptr && on_bootstrap_extension_config_scheduled_ != nullptr) {
     on_bootstrap_extension_config_scheduled_(thisAsVoidPtr(), in_module_config_, event_id);
@@ -247,6 +277,20 @@ newDynamicModuleBootstrapExtensionConfig(
     return on_admin_request.status();
   }
 
+  auto on_cluster_add_or_update =
+      dynamic_module->getFunctionPointer<OnBootstrapExtensionClusterAddOrUpdateType>(
+          "envoy_dynamic_module_on_bootstrap_extension_cluster_add_or_update");
+  if (!on_cluster_add_or_update.ok()) {
+    return on_cluster_add_or_update.status();
+  }
+
+  auto on_cluster_removal =
+      dynamic_module->getFunctionPointer<OnBootstrapExtensionClusterRemovalType>(
+          "envoy_dynamic_module_on_bootstrap_extension_cluster_removal");
+  if (!on_cluster_removal.ok()) {
+    return on_cluster_removal.status();
+  }
+
   auto config = std::make_shared<DynamicModuleBootstrapExtensionConfig>(
       extension_name, extension_config, metrics_namespace, std::move(dynamic_module),
       main_thread_dispatcher, context, stats_store);
@@ -276,6 +320,8 @@ newDynamicModuleBootstrapExtensionConfig(
   config->on_bootstrap_extension_http_callout_done_ = on_http_callout_done.value();
   config->on_bootstrap_extension_timer_fired_ = on_timer_fired.value();
   config->on_bootstrap_extension_admin_request_ = on_admin_request.value();
+  config->on_bootstrap_extension_cluster_add_or_update_ = on_cluster_add_or_update.value();
+  config->on_bootstrap_extension_cluster_removal_ = on_cluster_removal.value();
 
   return config;
 }

source/extensions/bootstrap/dynamic_modules/extension_config.h

@@ -9,6 +9,7 @@
 #include "envoy/stats/scope.h"
 #include "envoy/stats/stats.h"
 #include "envoy/stats/store.h"
+#include "envoy/upstream/cluster_manager.h"
 
 #include "source/common/common/logger.h"
 #include "source/common/http/message_impl.h"
@@ -49,6 +50,10 @@ using OnBootstrapExtensionTimerFiredType =
     decltype(&envoy_dynamic_module_on_bootstrap_extension_timer_fired);
 using OnBootstrapExtensionAdminRequestType =
     decltype(&envoy_dynamic_module_on_bootstrap_extension_admin_request);
+using OnBootstrapExtensionClusterAddOrUpdateType =
+    decltype(&envoy_dynamic_module_on_bootstrap_extension_cluster_add_or_update);
+using OnBootstrapExtensionClusterRemovalType =
+    decltype(&envoy_dynamic_module_on_bootstrap_extension_cluster_removal);
 
 class DynamicModuleBootstrapExtension;
 
@@ -58,6 +63,7 @@ class DynamicModuleBootstrapExtension;
  */
 class DynamicModuleBootstrapExtensionConfig
     : public std::enable_shared_from_this<DynamicModuleBootstrapExtensionConfig>,
+      public Upstream::ClusterUpdateCallbacks,
       public Logger::Loggable<Logger::Id::dynamic_modules> {
 public:
   /**
@@ -112,6 +118,23 @@ class DynamicModuleBootstrapExtensionConfig
    */
   void signalInitComplete();
 
+  /**
+   * Enables cluster lifecycle event notifications. When enabled, the module will receive
+   * on_bootstrap_extension_cluster_add_or_update and on_bootstrap_extension_cluster_removal
+   * callbacks when clusters are added, updated, or removed.
+   *
+   * This must be called on the main thread after the server is initialized, since the
+   * ClusterManager is not available during bootstrap extension creation.
+   *
+   * @return true if the callbacks were successfully registered, false if already registered.
+   */
+  bool enableClusterLifecycle();
+
+  // Upstream::ClusterUpdateCallbacks
+  void onClusterAddOrUpdate(absl::string_view cluster_name,
+                            Upstream::ThreadLocalClusterCommand& get_cluster) override;
+  void onClusterRemoval(const std::string& cluster_name) override;
+
   // The corresponding in-module configuration.
   envoy_dynamic_module_type_bootstrap_extension_config_module_ptr in_module_config_ = nullptr;
 
@@ -130,6 +153,9 @@ class DynamicModuleBootstrapExtensionConfig
   OnBootstrapExtensionHttpCalloutDoneType on_bootstrap_extension_http_callout_done_ = nullptr;
   OnBootstrapExtensionTimerFiredType on_bootstrap_extension_timer_fired_ = nullptr;
   OnBootstrapExtensionAdminRequestType on_bootstrap_extension_admin_request_ = nullptr;
+  OnBootstrapExtensionClusterAddOrUpdateType on_bootstrap_extension_cluster_add_or_update_ =
+      nullptr;
+  OnBootstrapExtensionClusterRemovalType on_bootstrap_extension_cluster_removal_ = nullptr;
 
   // The dynamic module.
   Extensions::DynamicModules::DynamicModulePtr dynamic_module_;
@@ -374,6 +400,14 @@ class DynamicModuleBootstrapExtensionConfig
   std::vector<ModuleGaugeVecHandle> gauge_vecs_;
   std::vector<ModuleHistogramHandle> histograms_;
   std::vector<ModuleHistogramVecHandle> histogram_vecs_;
+
+  // Cluster lifecycle callback handle. Set when the module enables cluster lifecycle events
+  // via enableClusterLifecycle(). Reset during shutdown to avoid use-after-free since the
+  // underlying TLS data is destroyed before the config.
+  Upstream::ClusterUpdateCallbacksHandlePtr cluster_update_callbacks_handle_;
+  // Handle for the shutdown lifecycle callback that cleans up cluster_update_callbacks_handle_.
+  Server::ServerLifecycleNotifier::HandlePtr cluster_lifecycle_shutdown_handle_;
+  bool cluster_lifecycle_enabled_ = false;
 };
 
 using DynamicModuleBootstrapExtensionConfigSharedPtr =

source/extensions/dynamic_modules/abi/abi.h

@@ -7326,6 +7326,44 @@ void envoy_dynamic_module_on_bootstrap_extension_timer_fired(
     envoy_dynamic_module_type_bootstrap_extension_config_module_ptr extension_config_module_ptr,
     envoy_dynamic_module_type_bootstrap_extension_timer_module_ptr timer_ptr);
 
+/**
+ * envoy_dynamic_module_on_bootstrap_extension_cluster_add_or_update is called when a cluster is
+ * added to or updated in the ClusterManager.
+ *
+ * This is only called if the module has opted in to receiving cluster lifecycle events via
+ * envoy_dynamic_module_callback_bootstrap_extension_enable_cluster_lifecycle. The callback is
+ * registered on the main thread and invoked on the main thread.
+ *
+ * @param extension_config_envoy_ptr is the pointer to the DynamicModuleBootstrapExtensionConfig
+ * object.
+ * @param extension_config_module_ptr is the pointer to the in-module bootstrap extension
+ * configuration created by envoy_dynamic_module_on_bootstrap_extension_config_new.
+ * @param cluster_name is the name of the cluster that was added or updated.
+ */
+void envoy_dynamic_module_on_bootstrap_extension_cluster_add_or_update(
+    envoy_dynamic_module_type_bootstrap_extension_config_envoy_ptr extension_config_envoy_ptr,
+    envoy_dynamic_module_type_bootstrap_extension_config_module_ptr extension_config_module_ptr,
+    envoy_dynamic_module_type_envoy_buffer cluster_name);
+
+/**
+ * envoy_dynamic_module_on_bootstrap_extension_cluster_removal is called when a cluster is
+ * removed from the ClusterManager.
+ *
+ * This is only called if the module has opted in to receiving cluster lifecycle events via
+ * envoy_dynamic_module_callback_bootstrap_extension_enable_cluster_lifecycle. The callback is
+ * registered on the main thread and invoked on the main thread.
+ *
+ * @param extension_config_envoy_ptr is the pointer to the DynamicModuleBootstrapExtensionConfig
+ * object.
+ * @param extension_config_module_ptr is the pointer to the in-module bootstrap extension
+ * configuration created by envoy_dynamic_module_on_bootstrap_extension_config_new.
+ * @param cluster_name is the name of the cluster that was removed.
+ */
+void envoy_dynamic_module_on_bootstrap_extension_cluster_removal(
+    envoy_dynamic_module_type_bootstrap_extension_config_envoy_ptr extension_config_envoy_ptr,
+    envoy_dynamic_module_type_bootstrap_extension_config_module_ptr extension_config_module_ptr,
+    envoy_dynamic_module_type_envoy_buffer cluster_name);
+
 // =============================================================================
 // Bootstrap Extension Callbacks
 // =============================================================================
@@ -7893,6 +7931,27 @@ bool envoy_dynamic_module_callback_bootstrap_extension_remove_admin_handler(
     envoy_dynamic_module_type_bootstrap_extension_config_envoy_ptr extension_config_envoy_ptr,
     envoy_dynamic_module_type_module_buffer path_prefix);
 
+// -------------------- Bootstrap Extension Callbacks - Cluster Lifecycle --------------------
+
+/**
+ * envoy_dynamic_module_callback_bootstrap_extension_enable_cluster_lifecycle is called by the
+ * module to opt in to receiving cluster lifecycle events
+ * (envoy_dynamic_module_on_bootstrap_extension_cluster_add_or_update and
+ * envoy_dynamic_module_on_bootstrap_extension_cluster_removal).
+ *
+ * This must be called on the main thread, typically during or after
+ * envoy_dynamic_module_on_bootstrap_extension_server_initialized, since the ClusterManager is
+ * not available until that point.
+ *
+ * This should be called at most once. Subsequent calls are no-ops and return false.
+ *
+ * @param extension_config_envoy_ptr is the pointer to the DynamicModuleBootstrapExtensionConfig
+ * object.
+ * @return true if the cluster lifecycle callbacks were successfully registered, false otherwise.
+ */
+bool envoy_dynamic_module_callback_bootstrap_extension_enable_cluster_lifecycle(
+    envoy_dynamic_module_type_bootstrap_extension_config_envoy_ptr extension_config_envoy_ptr);
+
 // =============================================================================
 // Common Host Types (shared by cluster and standalone load balancer extensions)
 // =============================================================================

source/extensions/dynamic_modules/abi_impl.cc

@@ -355,6 +355,14 @@ __attribute__((weak)) void envoy_dynamic_module_callback_bootstrap_extension_tim
                "not implemented in this context");
 }
 
+__attribute__((weak)) bool
+envoy_dynamic_module_callback_bootstrap_extension_enable_cluster_lifecycle(
+    envoy_dynamic_module_type_bootstrap_extension_config_envoy_ptr) {
+  IS_ENVOY_BUG("envoy_dynamic_module_callback_bootstrap_extension_enable_cluster_lifecycle: "
+               "not implemented in this context");
+  return false;
+}
+
 // ---------------------- Cluster extension callbacks ------------------------
 // These are weak symbols that provide default stub implementations. The actual implementations
 // are provided in the cluster extension abi_impl.cc when the cluster extension is used.

source/extensions/dynamic_modules/sdk/rust/src/bootstrap.rs

@@ -219,6 +219,17 @@ pub trait EnvoyBootstrapExtensionConfig {
   ///
   /// This must be called on the main thread.
   fn remove_admin_handler(&self, path_prefix: &str) -> bool;
+
+  /// Enable cluster lifecycle event notifications. When enabled, the module will receive
+  /// [`BootstrapExtensionConfig::on_cluster_add_or_update`] and
+  /// [`BootstrapExtensionConfig::on_cluster_removal`] callbacks when clusters are added, updated,
+  /// or removed from the ClusterManager.
+  ///
+  /// This must be called on the main thread, typically during or after `on_server_initialized`,
+  /// since the ClusterManager is not available until that point.
+  ///
+  /// This should be called at most once. Subsequent calls are no-ops and return `false`.
+  fn enable_cluster_lifecycle(&self) -> bool;
 }
 
 /// EnvoyBootstrapExtension is the Envoy-side bootstrap extension.
@@ -332,6 +343,36 @@ pub trait BootstrapExtensionConfig: Send + Sync {
   ) -> (u32, String) {
     (404, String::new())
   }
+
+  /// This is called when a cluster is added to or updated in the ClusterManager.
+  ///
+  /// This is only called if the module has opted in via
+  /// [`EnvoyBootstrapExtensionConfig::enable_cluster_lifecycle`]. The callback is invoked on the
+  /// main thread.
+  ///
+  /// * `envoy_extension_config` can be used to interact with the underlying Envoy config object.
+  /// * `cluster_name` is the name of the cluster that was added or updated.
+  fn on_cluster_add_or_update(
+    &self,
+    _envoy_extension_config: &mut dyn EnvoyBootstrapExtensionConfig,
+    _cluster_name: &str,
+  ) {
+  }
+
+  /// This is called when a cluster is removed from the ClusterManager.
+  ///
+  /// This is only called if the module has opted in via
+  /// [`EnvoyBootstrapExtensionConfig::enable_cluster_lifecycle`]. The callback is invoked on the
+  /// main thread.
+  ///
+  /// * `envoy_extension_config` can be used to interact with the underlying Envoy config object.
+  /// * `cluster_name` is the name of the cluster that was removed.
+  fn on_cluster_removal(
+    &self,
+    _envoy_extension_config: &mut dyn EnvoyBootstrapExtensionConfig,
+    _cluster_name: &str,
+  ) {
+  }
 }
 
 /// A completion callback that must be invoked exactly once to signal that an asynchronous
@@ -1026,6 +1067,12 @@ impl EnvoyBootstrapExtensionConfig for EnvoyBootstrapExtensionConfigImpl {
       )
     }
   }
+
+  fn enable_cluster_lifecycle(&self) -> bool {
+    unsafe {
+      abi::envoy_dynamic_module_callback_bootstrap_extension_enable_cluster_lifecycle(self.raw)
+    }
+  }
 }
 
 // Implementation of EnvoyBootstrapExtension
@@ -1433,3 +1480,59 @@ pub unsafe extern "C" fn envoy_dynamic_module_on_bootstrap_extension_admin_reque
 
   status_code
 }
+
+/// Event hook called by Envoy when a cluster is added to or updated in the ClusterManager.
+///
+/// # Safety
+///
+/// This is an FFI function called by Envoy. All pointer arguments must be valid as guaranteed
+/// by the Envoy dynamic module ABI.
+#[no_mangle]
+pub unsafe extern "C" fn envoy_dynamic_module_on_bootstrap_extension_cluster_add_or_update(
+  envoy_ptr: abi::envoy_dynamic_module_type_bootstrap_extension_config_envoy_ptr,
+  extension_config_ptr: abi::envoy_dynamic_module_type_bootstrap_extension_config_module_ptr,
+  cluster_name: abi::envoy_dynamic_module_type_envoy_buffer,
+) {
+  let extension_config = extension_config_ptr as *const *const dyn BootstrapExtensionConfig;
+  let extension_config = unsafe { &**extension_config };
+
+  let cluster_name_str = unsafe {
+    std::str::from_utf8_unchecked(std::slice::from_raw_parts(
+      cluster_name.ptr as *const u8,
+      cluster_name.length,
+    ))
+  };
+
+  extension_config.on_cluster_add_or_update(
+    &mut EnvoyBootstrapExtensionConfigImpl::new(envoy_ptr),
+    cluster_name_str,
+  );
+}
+
+/// Event hook called by Envoy when a cluster is removed from the ClusterManager.
+///
+/// # Safety
+///
+/// This is an FFI function called by Envoy. All pointer arguments must be valid as guaranteed
+/// by the Envoy dynamic module ABI.
+#[no_mangle]
+pub unsafe extern "C" fn envoy_dynamic_module_on_bootstrap_extension_cluster_removal(
+  envoy_ptr: abi::envoy_dynamic_module_type_bootstrap_extension_config_envoy_ptr,
+  extension_config_ptr: abi::envoy_dynamic_module_type_bootstrap_extension_config_module_ptr,
+  cluster_name: abi::envoy_dynamic_module_type_envoy_buffer,
+) {
+  let extension_config = extension_config_ptr as *const *const dyn BootstrapExtensionConfig;
+  let extension_config = unsafe { &**extension_config };
+
+  let cluster_name_str = unsafe {
+    std::str::from_utf8_unchecked(std::slice::from_raw_parts(
+      cluster_name.ptr as *const u8,
+      cluster_name.length,
+    ))
+  };
+
+  extension_config.on_cluster_removal(
+    &mut EnvoyBootstrapExtensionConfigImpl::new(envoy_ptr),
+    cluster_name_str,
+  );
+}