feat: introduce HeartbeatHandlerGroupBuilderCustomizer (#4803)

* feat: introduce `HeartbeatHandlerGroupBuilderFinalizer` * chore: rename to `HeartbeatHandlerGroupBuilderCustomizer`
GreptimeTeam · Oct 8, 2024 · 2cdd103 · 2cdd103
1 parent 4dea4ca
commit 2cdd103
Show file tree

Hide file tree

Showing 2 changed files with 32 additions and 12 deletions.
diff --git 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
@@ -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()?
             }
         };