photoncloud-monorepo/flaredb/crates/flaredb-server/src/heartbeat.rs

58 lines
2.2 KiB
Rust

use crate::store::Store;
use flaredb_proto::pdpb::pd_client::PdClient;
use flaredb_proto::pdpb::{ListRegionsRequest, RegisterStoreRequest};
use flaredb_types::RegionMeta;
use std::sync::Arc;
use tokio::time::{sleep, Duration};
/// Periodically send region/store heartbeat to PD.
pub async fn start_heartbeat(
pd_addr: String,
store: Arc<Store>,
server_addr: String,
requested_store_id: u64,
) {
tokio::spawn(async move {
let endpoint = format!("http://{}", pd_addr);
loop {
if let Ok(mut client) = PdClient::connect(endpoint.clone()).await {
if let Err(err) = client
.register_store(RegisterStoreRequest {
addr: server_addr.clone(),
store_id: requested_store_id,
})
.await
{
tracing::warn!("failed to register store with legacy PD: {}", err);
}
// list regions to keep routing fresh
if let Ok(resp) = client.list_regions(ListRegionsRequest {}).await {
let resp = resp.into_inner();
let mut metas = Vec::new();
for r in resp.regions {
let voters = if r.peers.is_empty() {
vec![store.store_id()]
} else {
r.peers.clone()
};
metas.push((
RegionMeta {
id: r.id,
start_key: r.start_key,
end_key: r.end_key,
},
voters,
));
}
if !metas.is_empty() {
let _ = store.refresh_regions(metas).await;
}
}
// send basic heartbeat info (store id + regions held)
// PD heartbeat RPC is not defined yet; placeholder refresh via list.
}
sleep(Duration::from_secs(30)).await;
}
});
}