rust/greptimedb
1
0
Fork 0
mirror of https://github.com/GreptimeTeam/greptimedb.git synced 2026-01-17 10:42:55 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: use the write runtime to handle the heartbeats (#4177)

Browse Source
This commit is contained in:
Jeremyhi
2024-06-20 16:45:07 +08:00
committed by GitHub
parent 4306cba866
commit 5bcd7a14bb
1 changed files with 26 additions and 18 deletions
Download Patch File Download Diff File Expand all files Collapse all files

44
src/meta-srv/src/service/heartbeat.rs
View File

@@ -24,12 +24,13 @@ use common_telemetry::{debug, error, info, warn};
use futures::StreamExt;
use once_cell::sync::OnceCell;
use tokio::sync::mpsc;
use tokio::sync::mpsc::Sender;
use tokio_stream::wrappers::ReceiverStream;
use tonic::{Request, Response, Streaming};
use crate::error;
use crate::error::Result;
use crate::handler::Pusher;
use crate::handler::{HeartbeatHandlerGroup, Pusher};
use crate::metasrv::{Context, Metasrv};
use crate::service::{GrpcResult, GrpcStream};
@@ -45,31 +46,24 @@ impl heartbeat_server::Heartbeat for Metasrv {
let (tx, rx) = mpsc::channel(128);
let handler_group = self.handler_group().clone();
let ctx = self.new_ctx();
let _handle = common_runtime::spawn_bg(async move {
let _handle = common_runtime::spawn_write(async move {
let mut pusher_key = None;
while let Some(msg) = in_stream.next().await {
let mut is_not_leader = false;
match msg {
Ok(req) => {
let header = match req.header.as_ref() {
Some(header) => header,
None => {
let err = error::MissingRequestHeaderSnafu {}.build();
error!("Exit on malformed request: MissingRequestHeader");
let _ = tx.send(Err(err.into())).await;
break;
}
};
debug!("Receiving heartbeat request: {:?}", req);
let Some(header) = req.header.as_ref() else {
error!("Exit on malformed request: MissingRequestHeader");
let _ = tx
.send(Err(error::MissingRequestHeaderSnafu {}.build().into()))
.await;
break;
};
if pusher_key.is_none() {
let node_id = get_node_id(header);
let role = header.role() as i32;
let key = format!("{}-{}", role, node_id);
let pusher = Pusher::new(tx.clone(), header);
handler_group.register(&key, pusher).await;
pusher_key = Some(key);
pusher_key = register_pusher(&handler_group, header, tx.clone()).await;
}
let res = handler_group
@@ -80,6 +74,7 @@ impl heartbeat_server::Heartbeat for Metasrv {
is_not_leader = res.as_ref().map_or(false, |r| r.is_not_leader());
debug!("Sending heartbeat response: {:?}", res);
if tx.send(res).await.is_err() {
info!("ReceiverStream was dropped; shutting down");
break;
@@ -168,6 +163,19 @@ fn get_node_id(header: &RequestHeader) -> u64 {
}
}
async fn register_pusher(
handler_group: &HeartbeatHandlerGroup,
header: &RequestHeader,
sender: Sender<std::result::Result<HeartbeatResponse, tonic::Status>>,
) -> Option<String> {
let role = header.role() as i32;
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;
Some(key)
}
#[cfg(test)]
mod tests {
use std::sync::Arc;