refactor: introduce HeartbeatHandlerGroupBuilder (#4785)

src/meta-srv/src/handler.rs

@@ -22,12 +22,28 @@ use api::v1::meta::{
  HeartbeatRequest, HeartbeatResponse, MailboxMessage, RegionLease, RequestHeader,
  ResponseHeader, Role, PROTOCOL_VERSION,
 };
+use check_leader_handler::CheckLeaderHandler;
+use collect_cluster_info_handler::{
+ CollectDatanodeClusterInfoHandler, CollectFlownodeClusterInfoHandler,
+ CollectFrontendClusterInfoHandler,
+};
+use collect_stats_handler::CollectStatsHandler;
+use common_base::Plugins;
 use common_meta::datanode::Stat;
 use common_meta::instruction::{Instruction, InstructionReply};
 use common_meta::sequence::Sequence;
 use common_telemetry::{debug, info, warn};
 use dashmap::DashMap;
+use extract_stat_handler::ExtractStatHandler;
+use failure_handler::RegionFailureHandler;
+use filter_inactive_region_stats::FilterInactiveRegionStatsHandler;
 use futures::future::join_all;
+use keep_lease_handler::{DatanodeKeepLeaseHandler, FlownodeKeepLeaseHandler};
+use mailbox_handler::MailboxHandler;
+use on_leader_start_handler::OnLeaderStartHandler;
+use publish_heartbeat_handler::PublishHeartbeatHandler;
+use region_lease_handler::RegionLeaseHandler;
+use response_header_handler::ResponseHeaderHandler;
 use snafu::{OptionExt, ResultExt};
 use store_api::storage::RegionId;
 use tokio::sync::mpsc::Sender;
@@ -36,6 +52,7 @@ use tokio::sync::{oneshot, Notify, RwLock};
 use crate::error::{self, DeserializeFromJsonSnafu, Result, UnexpectedInstructionReplySnafu};
 use crate::metasrv::Context;
 use crate::metrics::{METRIC_META_HANDLER_EXECUTE, METRIC_META_HEARTBEAT_CONNECTION_NUM};
+use crate::pubsub::PublisherRef;
 use crate::service::mailbox::{
  BroadcastChannel, Channel, Mailbox, MailboxReceiver, MailboxRef, MessageId,
 };
@@ -96,6 +113,7 @@ impl HeartbeatAccumulator {
  }
 }
 
+/// The pusher of the heartbeat response.
 pub struct Pusher {
  sender: Sender<std::result::Result<HeartbeatResponse, tonic::Status>>,
  res_header: ResponseHeader,
@@ -131,6 +149,7 @@ impl Pusher {
  }
 }
 
+/// The group of heartbeat pushers.
 #[derive(Clone, Default)]
 pub struct Pushers(Arc<RwLock<BTreeMap<String, Pusher>>>);
 
@@ -203,50 +222,57 @@ impl NameCachedHandler {
  }
 }
 
-#[derive(Clone, Default)]
+pub type HeartbeatHandlerGroupRef = Arc<HeartbeatHandlerGroup>;
+
+/// The group of heartbeat handlers.
+#[derive(Default)]
 pub struct HeartbeatHandlerGroup {
- handlers: Arc<RwLock<Vec<NameCachedHandler>>>,
+ handlers: Vec<NameCachedHandler>,
  pushers: Pushers,
 }
 
 impl HeartbeatHandlerGroup {
  pub(crate) fn new(pushers: Pushers) -> Self {
  Self {
- handlers: Arc::new(RwLock::new(vec![])),
+ handlers: vec![],
  pushers,
  }
  }
 
- pub async fn add_handler(&self, handler: impl HeartbeatHandler + 'static) {
- let mut handlers = self.handlers.write().await;
- handlers.push(NameCachedHandler::new(handler));
+ fn add_handler(&mut self, handler: impl HeartbeatHandler + 'static) {
+ self.handlers.push(NameCachedHandler::new(handler));
  }
 
- pub async fn register(&self, key: impl AsRef<str>, pusher: Pusher) {
+ /// Registers the heartbeat response [`Pusher`] with the given key to the group.
+ pub async fn register_pusher(&self, key: impl AsRef<str>, pusher: Pusher) {
  let key = key.as_ref();
  METRIC_META_HEARTBEAT_CONNECTION_NUM.inc();
  info!("Pusher register: {}", key);
  let _ = self.pushers.insert(key.to_string(), pusher).await;
  }
 
- pub async fn deregister(&self, key: impl AsRef<str>) -> Option<Pusher> {
+ /// Deregisters the heartbeat response [`Pusher`] with the given key from the group.
+ ///
+ /// Returns the [`Pusher`] if it exists.
+ pub async fn deregister_push(&self, key: impl AsRef<str>) -> Option<Pusher> {
  let key = key.as_ref();
  METRIC_META_HEARTBEAT_CONNECTION_NUM.dec();
  info!("Pusher unregister: {}", key);
  self.pushers.remove(key).await
  }
 
+ /// Returns the [`Pushers`] of the group.
  pub fn pushers(&self) -> Pushers {
  self.pushers.clone()
  }
 
+ /// Handles the heartbeat request.
  pub async fn handle(
  &self,
  req: HeartbeatRequest,
  mut ctx: Context,
  ) -> Result<HeartbeatResponse> {
  let mut acc = HeartbeatAccumulator::default();
- let handlers = self.handlers.read().await;
  let role = req
  .header
  .as_ref()
@@ -255,7 +281,7 @@ impl HeartbeatHandlerGroup {
  err_msg: format!("invalid role: {:?}", req.header),
  })?;
 
- for NameCachedHandler { name, handler } in handlers.iter() {
+ for NameCachedHandler { name, handler } in self.handlers.iter() {
  if !handler.is_acceptable(role) {
  continue;
  }
@@ -426,6 +452,84 @@ impl Mailbox for HeartbeatMailbox {
  }
 }
 
+/// The builder to build the group of heartbeat handlers.
+pub struct HeartbeatHandlerGroupBuilder {
+ /// The handler to handle region failure.
+ region_failure_handler: Option<RegionFailureHandler>,
+
+ /// The handler to handle region lease.
+ region_lease_handler: RegionLeaseHandler,
+
+ /// The plugins.
+ plugins: Option<Plugins>,
+
+ /// The heartbeat response pushers.
+ pushers: Pushers,
+}
+
+impl HeartbeatHandlerGroupBuilder {
+ pub fn new(pushers: Pushers, region_lease_handler: RegionLeaseHandler) -> Self {
+ Self {
+ region_failure_handler: None,
+ region_lease_handler,
+ plugins: None,
+ pushers,
+ }
+ }
+
+ /// Sets the [`RegionFailureHandler`].
+ pub fn with_region_failure_handler(mut self, handler: Option<RegionFailureHandler>) -> Self {
+ self.region_failure_handler = handler;
+ self
+ }
+
+ /// Sets the [`Plugins`].
+ pub fn with_plugins(mut self, plugins: Option<Plugins>) -> Self {
+ self.plugins = plugins;
+ self
+ }
+
+ /// Builds the group of heartbeat handlers.
+ pub fn build(self) -> HeartbeatHandlerGroup {
+ // Extract the `PublishHeartbeatHandler` from the plugins.
+ let publish_heartbeat_handler = if let Some(plugins) = self.plugins {
+ plugins
+ .get::<PublisherRef>()
+ .map(|publish| PublishHeartbeatHandler::new(publish.clone()))
+ } else {
+ None
+ };
+
+ // TODO(weny): Considers classifying handlers
+ // to make it easier for upper layers to customize handler groups.
+ let mut group = HeartbeatHandlerGroup::new(self.pushers);
+ group.add_handler(ResponseHeaderHandler);
+ // `KeepLeaseHandler` should preferably be in front of `CheckLeaderHandler`,
+ // because even if the current meta-server node is no longer the leader it can
+ // still help the datanode to keep lease.
+ group.add_handler(DatanodeKeepLeaseHandler);
+ group.add_handler(FlownodeKeepLeaseHandler);
+ group.add_handler(CheckLeaderHandler);
+ group.add_handler(OnLeaderStartHandler);
+ group.add_handler(ExtractStatHandler);
+ group.add_handler(CollectDatanodeClusterInfoHandler);
+ group.add_handler(CollectFrontendClusterInfoHandler);
+ group.add_handler(CollectFlownodeClusterInfoHandler);
+ group.add_handler(MailboxHandler);
+ group.add_handler(self.region_lease_handler);
+ group.add_handler(FilterInactiveRegionStatsHandler);
+ if let Some(region_failure_handler) = self.region_failure_handler {
+ group.add_handler(region_failure_handler);
+ }
+ if let Some(publish_heartbeat_handler) = publish_heartbeat_handler {
+ group.add_handler(publish_heartbeat_handler);
+ }
+ group.add_handler(CollectStatsHandler::default());
+
+ group
+ }
+}
+
 #[cfg(test)]
 mod tests {
 
@@ -489,7 +593,7 @@ mod tests {
  let pusher: Pusher = Pusher::new(pusher_tx, &res_header);
  let handler_group = HeartbeatHandlerGroup::default();
  handler_group
- .register(format!("{}-{}", Role::Datanode as i32, datanode_id), pusher)
+ .register_pusher(format!("{}-{}", Role::Datanode as i32, datanode_id), pusher)
  .await;
 
  let kv_backend = Arc::new(MemoryKvBackend::new());
@@ -519,21 +623,21 @@ mod tests {
 
  #[tokio::test]
  async fn test_handler_name() {
- let group = HeartbeatHandlerGroup::default();
- group.add_handler(ResponseHeaderHandler).await;
- group.add_handler(DatanodeKeepLeaseHandler).await;
- group.add_handler(FlownodeKeepLeaseHandler).await;
- group.add_handler(CheckLeaderHandler).await;
- group.add_handler(OnLeaderStartHandler).await;
- group.add_handler(ExtractStatHandler).await;
- group.add_handler(CollectDatanodeClusterInfoHandler).await;
- group.add_handler(CollectFrontendClusterInfoHandler).await;
- group.add_handler(CollectFlownodeClusterInfoHandler).await;
- group.add_handler(MailboxHandler).await;
- group.add_handler(FilterInactiveRegionStatsHandler).await;
- group.add_handler(CollectStatsHandler::default()).await;
-
- let handlers = group.handlers.read().await;
+ let mut group = HeartbeatHandlerGroup::default();
+ group.add_handler(ResponseHeaderHandler);
+ group.add_handler(DatanodeKeepLeaseHandler);
+ group.add_handler(FlownodeKeepLeaseHandler);
+ group.add_handler(CheckLeaderHandler);
+ group.add_handler(OnLeaderStartHandler);
+ group.add_handler(ExtractStatHandler);
+ group.add_handler(CollectDatanodeClusterInfoHandler);
+ group.add_handler(CollectFrontendClusterInfoHandler);
+ group.add_handler(CollectFlownodeClusterInfoHandler);
+ group.add_handler(MailboxHandler);
+ group.add_handler(FilterInactiveRegionStatsHandler);
+ group.add_handler(CollectStatsHandler::default());
+
+ let handlers = group.handlers;
 
  assert_eq!(12, handlers.len());
 

src/meta-srv/src/metasrv.rs

@@ -52,7 +52,7 @@ use crate::error::{
  StopProcedureManagerSnafu,
 };
 use crate::failure_detector::PhiAccrualFailureDetectorOptions;
-use crate::handler::HeartbeatHandlerGroup;
+use crate::handler::HeartbeatHandlerGroupRef;
 use crate::lease::lookup_datanode_peer;
 use crate::lock::DistLockRef;
 use crate::procedure::region_migration::manager::RegionMigrationManagerRef;
@@ -366,7 +366,7 @@ pub struct Metasrv {
  selector: SelectorRef,
  // The flow selector is used to select a target flownode.
  flow_selector: SelectorRef,
- handler_group: HeartbeatHandlerGroup,
+ handler_group: HeartbeatHandlerGroupRef,
  election: Option<ElectionRef>,
  lock: DistLockRef,
  procedure_manager: ProcedureManagerRef,
@@ -562,7 +562,7 @@ impl Metasrv {
  &self.flow_selector
  }
 
- pub fn handler_group(&self) -> &HeartbeatHandlerGroup {
+ pub fn handler_group(&self) -> &HeartbeatHandlerGroupRef {
  &self.handler_group
  }
 

src/meta-srv/src/metasrv/builder.rs

@@ -46,22 +46,11 @@ use crate::cluster::{MetaPeerClientBuilder, MetaPeerClientRef};
 use crate::error::{self, Result};
 use crate::flow_meta_alloc::FlowPeerAllocator;
 use crate::greptimedb_telemetry::get_greptimedb_telemetry_task;
-use crate::handler::check_leader_handler::CheckLeaderHandler;
-use crate::handler::collect_cluster_info_handler::{
- CollectDatanodeClusterInfoHandler, CollectFlownodeClusterInfoHandler,
- CollectFrontendClusterInfoHandler,
-};
-use crate::handler::collect_stats_handler::CollectStatsHandler;
-use crate::handler::extract_stat_handler::ExtractStatHandler;
 use crate::handler::failure_handler::RegionFailureHandler;
-use crate::handler::filter_inactive_region_stats::FilterInactiveRegionStatsHandler;
-use crate::handler::keep_lease_handler::{DatanodeKeepLeaseHandler, FlownodeKeepLeaseHandler};
-use crate::handler::mailbox_handler::MailboxHandler;
-use crate::handler::on_leader_start_handler::OnLeaderStartHandler;
-use crate::handler::publish_heartbeat_handler::PublishHeartbeatHandler;
 use crate::handler::region_lease_handler::RegionLeaseHandler;
-use crate::handler::response_header_handler::ResponseHeaderHandler;
-use crate::handler::{HeartbeatHandlerGroup, HeartbeatMailbox, Pushers};
+use crate::handler::{
+ HeartbeatHandlerGroup, HeartbeatHandlerGroupBuilder, HeartbeatMailbox, Pushers,
+};
 use crate::lease::MetaPeerLookupService;
 use crate::lock::memory::MemLock;
 use crate::lock::DistLockRef;
@@ -70,7 +59,6 @@ use crate::metasrv::{
 };
 use crate::procedure::region_migration::manager::RegionMigrationManager;
 use crate::procedure::region_migration::DefaultContextFactory;
-use crate::pubsub::PublisherRef;
 use crate::region::supervisor::{
  HeartbeatAcceptor, RegionFailureDetectorControl, RegionSupervisor, RegionSupervisorTicker,
  DEFAULT_TICK_INTERVAL,
@@ -364,41 +352,16 @@ impl MetasrvBuilder {
  let handler_group = match handler_group {
  Some(handler_group) => handler_group,
  None => {
- let publish_heartbeat_handler = plugins
- .clone()
- .and_then(|plugins| plugins.get::<PublisherRef>())
- .map(|publish| PublishHeartbeatHandler::new(publish.clone()));
-
  let region_lease_handler = RegionLeaseHandler::new(
  distributed_time_constants::REGION_LEASE_SECS,
  table_metadata_manager.clone(),
  memory_region_keeper.clone(),
  );
 
- let group = HeartbeatHandlerGroup::new(pushers);
- group.add_handler(ResponseHeaderHandler).await;
- // `KeepLeaseHandler` should preferably be in front of `CheckLeaderHandler`,
- // because even if the current meta-server node is no longer the leader it can
- // still help the datanode to keep lease.
- group.add_handler(DatanodeKeepLeaseHandler).await;
- group.add_handler(FlownodeKeepLeaseHandler).await;
- group.add_handler(CheckLeaderHandler).await;
- group.add_handler(OnLeaderStartHandler).await;
- group.add_handler(ExtractStatHandler).await;
- group.add_handler(CollectDatanodeClusterInfoHandler).await;
- group.add_handler(CollectFrontendClusterInfoHandler).await;
- group.add_handler(CollectFlownodeClusterInfoHandler).await;
- group.add_handler(MailboxHandler).await;
- group.add_handler(region_lease_handler).await;
- group.add_handler(FilterInactiveRegionStatsHandler).await;
- if let Some(region_failover_handler) = region_failover_handler {
- group.add_handler(region_failover_handler).await;
- }
- if let Some(publish_heartbeat_handler) = publish_heartbeat_handler {
- group.add_handler(publish_heartbeat_handler).await;
- }
- group.add_handler(CollectStatsHandler::default()).await;
- group
+ HeartbeatHandlerGroupBuilder::new(pushers, region_lease_handler)
+ .with_plugins(plugins.clone())
+ .with_region_failure_handler(region_failover_handler)
+ .build()
  }
  };
 
@@ -417,7 +380,7 @@ impl MetasrvBuilder {
  selector,
  // TODO(jeremy): We do not allow configuring the flow selector.
  flow_selector: Arc::new(RoundRobinSelector::new(SelectTarget::Flownode)),
- handler_group,
+ handler_group: Arc::new(handler_group),
  election,
  lock,
  procedure_manager,

src/meta-srv/src/service/heartbeat.rs

@@ -113,7 +113,7 @@ impl heartbeat_server::Heartbeat for Metasrv {
  );
 
  if let Some(key) = pusher_key {
- let _ = handler_group.deregister(&key).await;
+ let _ = handler_group.deregister_push(&key).await;
  }
  });
 
@@ -177,7 +177,7 @@ async fn register_pusher(
  let node_id = get_node_id(header);
  let key = format!("{}-{}", role, node_id);
  let pusher = Pusher::new(sender, header);
- handler_group.register(&key, pusher).await;
+ handler_group.register_pusher(&key, pusher).await;
  Some(key)
 }