diff --git a/src/meta-srv/src/handler.rs b/src/meta-srv/src/handler.rs index 082f902ee1..50b53f4970 100644 --- a/src/meta-srv/src/handler.rs +++ b/src/meta-srv/src/handler.rs @@ -13,7 +13,7 @@ // limitations under the License. use std::cmp::Ordering; -use std::collections::{BTreeMap, HashSet}; +use std::collections::{BTreeMap, HashMap, HashSet}; use std::fmt::{Debug, Display}; use std::ops::Bound; use std::sync::{Arc, Mutex}; @@ -119,6 +119,8 @@ pub struct HeartbeatAccumulator { pub stat: Option, pub inactive_region_ids: HashSet, pub region_lease: Option, + /// Generic heartbeat response extensions accumulated by handlers. + pub extensions: HashMap>, } impl HeartbeatAccumulator { @@ -405,7 +407,7 @@ impl HeartbeatHandlerGroup { region_lease: acc.region_lease, mailbox_message, heartbeat_config, - extensions: Default::default(), + extensions: std::mem::take(&mut acc.extensions), }; Ok(res) } @@ -888,16 +890,21 @@ mod tests { use std::sync::Arc; use std::time::Duration; - use api::v1::meta::{MailboxMessage, Role}; + use api::v1::meta::{HeartbeatRequest, MailboxMessage, RequestHeader, Role}; use common_meta::kv_backend::memory::MemoryKvBackend; use common_meta::sequence::SequenceBuilder; use tokio::sync::mpsc; - use super::{HeartbeatHandlerGroupBuilder, PusherId, Pushers}; + use super::{ + HandleControl, HeartbeatAccumulator, HeartbeatHandler, HeartbeatHandlerGroupBuilder, + PusherId, Pushers, + }; use crate::error; use crate::handler::collect_stats_handler::CollectStatsHandler; use crate::handler::response_header_handler::ResponseHeaderHandler; + use crate::handler::test_utils::TestEnv; use crate::handler::{HeartbeatHandlerGroup, HeartbeatMailbox, Pusher}; + use crate::metasrv::Context; use crate::service::mailbox::{Channel, MailboxReceiver, MailboxRef}; #[tokio::test] @@ -1316,4 +1323,44 @@ mod tests { drop(pusher); deregister_signal_tx.changed().await.unwrap(); } + + #[tokio::test] + async fn test_custom_handler_extension_reaches_response() { + struct CustomExtensionHandler; + + #[async_trait::async_trait] + impl HeartbeatHandler for CustomExtensionHandler { + fn is_acceptable(&self, role: Role) -> bool { + role == Role::Frontend + } + + async fn handle( + &self, + _req: &HeartbeatRequest, + _ctx: &mut Context, + acc: &mut HeartbeatAccumulator, + ) -> crate::error::Result { + acc.extensions + .insert("custom.key".to_string(), b"custom-value".to_vec()); + Ok(HandleControl::Continue) + } + } + + let mut builder = + HeartbeatHandlerGroupBuilder::new(Pushers::default()).add_default_handlers(); + builder.add_handler_last(CustomExtensionHandler); + + let group = builder.build().unwrap(); + + let req = HeartbeatRequest { + header: Some(RequestHeader::new(1, Role::Frontend, Default::default())), + ..Default::default() + }; + let ctx = TestEnv::new().ctx(); + + let res = group.handle(req, ctx).await.unwrap(); + + let extensions = res.extensions; + assert_eq!(extensions.get("custom.key").unwrap(), b"custom-value"); + } }