feat: impl pubsub in metasrv

src/meta-srv/Cargo.toml

@@ -58,5 +58,6 @@ uuid.workspace = true
 [dev-dependencies]
 chrono.workspace = true
 common-procedure-test = { path = "../common/procedure-test" }
+session = { path = "../session", features = ["testing"] }
 tracing = "0.1"
 tracing-subscriber = { version = "0.3", features = ["env-filter"] }

src/meta-srv/src/error.rs

@@ -21,6 +21,8 @@ use tokio::sync::mpsc::error::SendError;
 use tonic::codegen::http;
 use tonic::Code;
 
+use crate::pubsub::Message;
+
 #[derive(Debug, Snafu)]
 #[snafu(visibility(pub))]
 pub enum Error {
@@ -461,6 +463,12 @@ pub enum Error {
 
     #[snafu(display("Invalid heartbeat request: {}", err_msg))]
     InvalidHeartbeatRequest { err_msg: String, location: Location },
+
+    #[snafu(display("Failed to publish message: {:?}", source))]
+    PublishMessage {
+        source: SendError<Message>,
+        location: Location,
+    },
 }
 
 pub type Result<T> = std::result::Result<T, Error>;
@@ -506,6 +514,7 @@ impl ErrorExt for Error {
             | Error::UpdateTableMetadata { .. }
             | Error::NoEnoughAvailableDatanode { .. }
             | Error::ConvertGrpcExpr { .. }
+            | Error::PublishMessage { .. }
             | Error::Join { .. } => StatusCode::Internal,
             Error::EmptyKey { .. }
             | Error::MissingRequiredParameter { .. }

src/meta-srv/src/handler.rs

@@ -54,6 +54,7 @@ pub mod mailbox_handler;
 pub mod node_stat;
 mod on_leader_start_handler;
 mod persist_stats_handler;
+pub mod publish_heartbeat_handler;
 pub(crate) mod region_lease_handler;
 mod response_header_handler;
 

src/meta-srv/src/handler/publish_heartbeat_handler.rs

@@ -0,0 +1,50 @@
+// Copyright 2023 Greptime Team
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use api::v1::meta::{HeartbeatRequest, Role};
+use async_trait::async_trait;
+
+use crate::error::Result;
+use crate::handler::{HeartbeatAccumulator, HeartbeatHandler};
+use crate::metasrv::Context;
+use crate::pubsub::{Message, PublishRef};
+
+pub struct PublishHeartbeatHandler {
+    publish: PublishRef,
+}
+
+impl PublishHeartbeatHandler {
+    pub fn new(publish: PublishRef) -> PublishHeartbeatHandler {
+        PublishHeartbeatHandler { publish }
+    }
+}
+
+#[async_trait]
+impl HeartbeatHandler for PublishHeartbeatHandler {
+    fn is_acceptable(&self, role: Role) -> bool {
+        role == Role::Datanode
+    }
+
+    async fn handle(
+        &self,
+        req: &HeartbeatRequest,
+        _: &mut Context,
+        _: &mut HeartbeatAccumulator,
+    ) -> Result<()> {
+        let msg = Message::Heartbeat(Box::new(req.clone()));
+        self.publish.send_msg(msg).await;
+
+        Ok(())
+    }
+}

src/meta-srv/src/lib.rs

@@ -32,6 +32,7 @@ mod metrics;
 #[cfg(feature = "mock")]
 pub mod mocks;
 pub mod procedure;
+pub mod pubsub;
 pub mod selector;
 mod sequence;
 pub mod service;

src/meta-srv/src/metasrv.rs

@@ -38,6 +38,7 @@ use crate::error::{RecoverProcedureSnafu, Result};
 use crate::handler::HeartbeatHandlerGroup;
 use crate::lock::DistLockRef;
 use crate::metadata_service::MetadataServiceRef;
+use crate::pubsub::{PublishRef, SubscribeManagerRef};
 use crate::selector::{Selector, SelectorType};
 use crate::sequence::SequenceRef;
 use crate::service::mailbox::MailboxRef;
@@ -177,6 +178,7 @@ pub struct MetaSrv {
     ddl_manager: DdlManagerRef,
     table_metadata_manager: TableMetadataManagerRef,
     greptimedb_telemerty_task: Arc<GreptimeDBTelemetryTask>,
+    pubsub: Option<(PublishRef, SubscribeManagerRef)>,
 }
 
 impl MetaSrv {
@@ -196,6 +198,7 @@ impl MetaSrv {
             let procedure_manager = self.procedure_manager.clone();
             let in_memory = self.in_memory.clone();
             let leader_cached_kv_store = self.leader_cached_kv_store.clone();
+            let subscribe_manager = self.subscribe_manager().cloned();
             let mut rx = election.subscribe_leader_change();
             let task_handler = self.greptimedb_telemerty_task.clone();
             let _handle = common_runtime::spawn_bg(async move {
@@ -219,6 +222,12 @@ impl MetaSrv {
                                     });
                                 }
                                 LeaderChangeMessage::StepDown(leader) => {
+                                    if let Some(sub_manager) = subscribe_manager.clone() {
+                                        info!("Leader changed, un_subscribe all");
+                                        if let Err(e) = sub_manager.un_subscribe_all() {
+                                            error!("Failed to un_subscribe all, error: {}", e);
+                                        }
+                                    }
                                     error!("Leader :{:?} step down", leader);
                                     let _ = task_handler.stop().await.map_err(|e| {
                                         debug!(
@@ -329,6 +338,14 @@ impl MetaSrv {
         &self.table_metadata_manager
     }
 
+    pub fn publish(&self) -> Option<&PublishRef> {
+        self.pubsub.as_ref().map(|suite| &suite.0)
+    }
+
+    pub fn subscribe_manager(&self) -> Option<&SubscribeManagerRef> {
+        self.pubsub.as_ref().map(|suite| &suite.1)
+    }
+
     #[inline]
     pub fn new_ctx(&self) -> Context {
         let server_addr = self.options().server_addr.clone();
@@ -339,6 +356,7 @@ impl MetaSrv {
         let mailbox = self.mailbox.clone();
         let election = self.election.clone();
         let skip_all = Arc::new(AtomicBool::new(false));
+
         Context {
             server_addr,
             in_memory,

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

@@ -21,11 +21,13 @@ use common_grpc::channel_manager::ChannelConfig;
 use common_meta::key::{TableMetadataManager, TableMetadataManagerRef};
 use common_procedure::local::{LocalManager, ManagerConfig};
 use common_procedure::ProcedureManagerRef;
+use tokio::sync::mpsc::Sender;
 
 use crate::cluster::{MetaPeerClientBuilder, MetaPeerClientRef};
 use crate::ddl::{DdlManager, DdlManagerRef};
 use crate::error::Result;
 use crate::handler::mailbox_handler::MailboxHandler;
+use crate::handler::publish_heartbeat_handler::PublishHeartbeatHandler;
 use crate::handler::region_lease_handler::RegionLeaseHandler;
 use crate::handler::{
     CheckLeaderHandler, CollectStatsHandler, HeartbeatHandlerGroup, HeartbeatMailbox,
@@ -40,6 +42,9 @@ use crate::metasrv::{
 };
 use crate::procedure::region_failover::RegionFailoverManager;
 use crate::procedure::state_store::MetaStateStore;
+use crate::pubsub::{
+    DefaultPublish, DefaultSubscribeManager, Message, PublishRef, SubscribeManagerRef,
+};
 use crate::selector::lease_based::LeaseBasedSelector;
 use crate::sequence::Sequence;
 use crate::service::mailbox::MailboxRef;
@@ -169,6 +174,7 @@ impl MetaSrvBuilder {
             &table_metadata_manager,
         );
         let _ = ddl_manager.try_start();
+        let (publish, sub_manager) = build_publish();
 
         let handler_group = match handler_group {
             Some(handler_group) => handler_group,
@@ -216,6 +222,9 @@ impl MetaSrvBuilder {
                 group.add_handler(RegionLeaseHandler::default()).await;
                 group.add_handler(PersistStatsHandler::default()).await;
                 group
+                    .add_handler(PublishHeartbeatHandler::new(publish.clone()))
+                    .await;
+                group
             }
         };
 
@@ -237,6 +246,7 @@ impl MetaSrvBuilder {
             ddl_manager,
             table_metadata_manager,
             greptimedb_telemerty_task: get_greptimedb_telemetry_task(meta_peer_client).await,
+            pubsub: Some((publish, sub_manager)),
         })
     }
 }
@@ -315,6 +325,12 @@ fn build_ddl_manager(
     ))
 }
 
+fn build_publish() -> (PublishRef, SubscribeManagerRef) {
+    let sub_manager = Arc::new(DefaultSubscribeManager::<Sender<Message>>::default());
+    let publish = Arc::new(DefaultPublish::new(sub_manager.clone()));
+    (publish, sub_manager)
+}
+
 impl Default for MetaSrvBuilder {
     fn default() -> Self {
         Self::new()

src/meta-srv/src/pubsub.rs

@@ -0,0 +1,47 @@
+// Copyright 2023 Greptime Team
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use api::v1::meta::HeartbeatRequest;
+
+mod publish;
+mod subscribe_manager;
+mod subscriber;
+#[cfg(test)]
+mod tests;
+
+pub use publish::{DefaultPublish, Publish, PublishRef};
+pub use subscribe_manager::{
+    AddSubRequest, DefaultSubscribeManager, SubscribeManager, SubscribeManagerRef, SubscribeQuery,
+    UnSubRequest,
+};
+pub use subscriber::{Subscriber, SubscriberRef, Transport};
+
+/// Subscribed topic type, determined by the ability of meta.
+#[derive(Debug, Clone, Copy, Hash, PartialEq, Eq)]
+pub enum Topic {
+    Heartbeat,
+}
+
+#[derive(Clone, Debug)]
+pub enum Message {
+    Heartbeat(Box<HeartbeatRequest>),
+}
+
+impl Message {
+    pub fn topic(&self) -> Topic {
+        match self {
+            Message::Heartbeat(_) => Topic::Heartbeat,
+        }
+    }
+}

src/meta-srv/src/pubsub/publish.rs

@@ -0,0 +1,72 @@
+// Copyright 2023 Greptime Team
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use std::fmt::Debug;
+use std::marker::PhantomData;
+use std::sync::Arc;
+
+use common_telemetry::error;
+
+use crate::pubsub::{Message, SubscribeManager, Transport, UnSubRequest};
+
+/// This trait provides a `send_msg` method that can be used by other modules
+/// of meta to publish [Message].
+#[async_trait::async_trait]
+pub trait Publish: Send + Sync {
+    async fn send_msg(&self, message: Message);
+}
+
+pub type PublishRef = Arc<dyn Publish>;
+
+/// The default implementation of [Publish]
+pub struct DefaultPublish<M, T> {
+    subscribe_manager: Arc<M>,
+    _transport: PhantomData<T>,
+}
+
+impl<M, T> DefaultPublish<M, T> {
+    pub fn new(subscribe_manager: Arc<M>) -> Self {
+        Self {
+            subscribe_manager,
+            _transport: PhantomData,
+        }
+    }
+}
+
+#[async_trait::async_trait]
+impl<M, T> Publish for DefaultPublish<M, T>
+where
+    M: SubscribeManager<T>,
+    T: Transport + Debug,
+{
+    async fn send_msg(&self, message: Message) {
+        let sub_list = self
+            .subscribe_manager
+            .subscribers_by_topic(&message.topic());
+
+        for sub in sub_list {
+            if sub.transport_msg(message.clone()).await.is_err() {
+                // If an error occurs, we consider the subscriber offline,
+                // so un_subscribe here.
+                let req = UnSubRequest {
+                    subscriber_id: sub.id(),
+                };
+
+                if let Err(e) = self.subscribe_manager.un_subscribe(req.clone()) {
+                    error!(e; "failed to un_subscribe, req: {:?}", req);
+                }
+            }
+        }
+    }
+}