From 2cdd103874a649282612bd05f94023fc0556f6b2 Mon Sep 17 00:00:00 2001
From: Weny Xu <wenymedia@gmail.com>
Date: Tue, 8 Oct 2024 17:02:06 +0800
Subject: [PATCH] feat: introduce `HeartbeatHandlerGroupBuilderCustomizer`
 (#4803)

* feat: introduce `HeartbeatHandlerGroupBuilderFinalizer`

* chore: rename to `HeartbeatHandlerGroupBuilderCustomizer`
---
 src/meta-srv/src/handler.rs         | 42 +++++++++++++++++++++--------
 src/meta-srv/src/metasrv/builder.rs |  2 +-
 2 files changed, 32 insertions(+), 12 deletions(-)

diff --git a/src/meta-srv/src/handler.rs b/src/meta-srv/src/handler.rs
index 9d82920c08dd..b45a1ac7e3ad 100644
--- a/src/meta-srv/src/handler.rs
+++ b/src/meta-srv/src/handler.rs
@@ -527,11 +527,22 @@ impl HeartbeatHandlerGroupBuilder {
     }
 
     /// Builds the group of heartbeat handlers.
-    pub fn build(self) -> HeartbeatHandlerGroup {
-        HeartbeatHandlerGroup {
+    ///
+    /// Applies the customizer if it exists.
+    pub fn build(mut self) -> Result<HeartbeatHandlerGroup> {
+        if let Some(customizer) = self
+            .plugins
+            .as_ref()
+            .and_then(|plugins| plugins.get::<HeartbeatHandlerGroupBuilderCustomizerRef>())
+        {
+            debug!("Customizing the heartbeat handler group builder");
+            customizer.customize(&mut self)?;
+        }
+
+        Ok(HeartbeatHandlerGroup {
             handlers: self.handlers.into_iter().collect(),
             pushers: self.pushers,
-        }
+        })
     }
 
     /// Adds the handler after the specified handler.
@@ -582,6 +593,14 @@ impl HeartbeatHandlerGroupBuilder {
     }
 }
 
+pub type HeartbeatHandlerGroupBuilderCustomizerRef =
+    Arc<dyn HeartbeatHandlerGroupBuilderCustomizer>;
+
+/// The customizer of the [`HeartbeatHandlerGroupBuilder`].
+pub trait HeartbeatHandlerGroupBuilderCustomizer: Send + Sync {
+    fn customize(&self, builder: &mut HeartbeatHandlerGroupBuilder) -> Result<()>;
+}
+
 #[cfg(test)]
 mod tests {
 
@@ -670,7 +689,8 @@ mod tests {
     fn test_handler_group_builder() {
         let group = HeartbeatHandlerGroupBuilder::new(Pushers::default())
             .add_default_handlers()
-            .build();
+            .build()
+            .unwrap();
 
         let handlers = group.handlers;
         assert_eq!(12, handlers.len());
@@ -706,7 +726,7 @@ mod tests {
             )
             .unwrap();
 
-        let group = builder.build();
+        let group = builder.build().unwrap();
         let handlers = group.handlers;
         assert_eq!(13, handlers.len());
 
@@ -739,7 +759,7 @@ mod tests {
             .add_handler_before("ResponseHeaderHandler", CollectStatsHandler::default())
             .unwrap();
 
-        let group = builder.build();
+        let group = builder.build().unwrap();
         let handlers = group.handlers;
         assert_eq!(13, handlers.len());
 
@@ -772,7 +792,7 @@ mod tests {
             .add_handler_after("MailboxHandler", CollectStatsHandler::default())
             .unwrap();
 
-        let group = builder.build();
+        let group = builder.build().unwrap();
         let handlers = group.handlers;
         assert_eq!(13, handlers.len());
 
@@ -805,7 +825,7 @@ mod tests {
             .add_handler_after("CollectStatsHandler", ResponseHeaderHandler)
             .unwrap();
 
-        let group = builder.build();
+        let group = builder.build().unwrap();
         let handlers = group.handlers;
         assert_eq!(13, handlers.len());
 
@@ -838,7 +858,7 @@ mod tests {
             .replace_handler("MailboxHandler", CollectStatsHandler::default())
             .unwrap();
 
-        let group = builder.build();
+        let group = builder.build().unwrap();
         let handlers = group.handlers;
         assert_eq!(12, handlers.len());
 
@@ -870,7 +890,7 @@ mod tests {
             .replace_handler("CollectStatsHandler", ResponseHeaderHandler)
             .unwrap();
 
-        let group = builder.build();
+        let group = builder.build().unwrap();
         let handlers = group.handlers;
         assert_eq!(12, handlers.len());
 
@@ -902,7 +922,7 @@ mod tests {
             .replace_handler("ResponseHeaderHandler", CollectStatsHandler::default())
             .unwrap();
 
-        let group = builder.build();
+        let group = builder.build().unwrap();
         let handlers = group.handlers;
         assert_eq!(12, handlers.len());
 
diff --git a/src/meta-srv/src/metasrv/builder.rs b/src/meta-srv/src/metasrv/builder.rs
index 88d7e5bd6d05..050f82517457 100644
--- a/src/meta-srv/src/metasrv/builder.rs
+++ b/src/meta-srv/src/metasrv/builder.rs
@@ -363,7 +363,7 @@ impl MetasrvBuilder {
                     .with_region_failure_handler(region_failover_handler)
                     .with_region_lease_handler(Some(region_lease_handler))
                     .add_default_handlers()
-                    .build()
+                    .build()?
             }
         };