79 lines
2.4 KiB
Rust
79 lines
2.4 KiB
Rust
use crate::cluster::Cluster;
|
|
use crate::tso::TsoOracle;
|
|
use flaredb_proto::pdpb::pd_server::Pd;
|
|
use flaredb_proto::pdpb::{
|
|
GetRegionRequest, GetRegionResponse, ListRegionsRequest, ListRegionsResponse,
|
|
MoveRegionRequest, MoveRegionResponse, RegisterStoreRequest, RegisterStoreResponse,
|
|
};
|
|
use std::sync::Arc;
|
|
use tonic::{Request, Response, Status};
|
|
|
|
pub struct PdServiceImpl {
|
|
_tso_oracle: Arc<TsoOracle>,
|
|
cluster: Cluster,
|
|
}
|
|
|
|
impl PdServiceImpl {
|
|
pub fn new(tso_oracle: Arc<TsoOracle>) -> Self {
|
|
Self {
|
|
_tso_oracle: tso_oracle,
|
|
cluster: Cluster::new(),
|
|
}
|
|
}
|
|
|
|
#[allow(dead_code)]
|
|
async fn list_regions(
|
|
&self,
|
|
_request: Request<ListRegionsRequest>,
|
|
) -> Result<Response<ListRegionsResponse>, Status> {
|
|
let (regions, stores) = self.cluster.list_regions();
|
|
Ok(Response::new(ListRegionsResponse { regions, stores }))
|
|
}
|
|
|
|
#[allow(dead_code)]
|
|
async fn move_region(
|
|
&self,
|
|
_request: Request<MoveRegionRequest>,
|
|
) -> Result<Response<MoveRegionResponse>, Status> {
|
|
Ok(Response::new(MoveRegionResponse { accepted: true }))
|
|
}
|
|
}
|
|
|
|
#[tonic::async_trait]
|
|
impl Pd for PdServiceImpl {
|
|
async fn register_store(
|
|
&self,
|
|
request: Request<RegisterStoreRequest>,
|
|
) -> Result<Response<RegisterStoreResponse>, Status> {
|
|
let req = request.into_inner();
|
|
let requested_store_id = (req.store_id != 0).then_some(req.store_id);
|
|
let store_id = self.cluster.register_store(req.addr, requested_store_id);
|
|
Ok(Response::new(RegisterStoreResponse {
|
|
store_id,
|
|
cluster_id: 1, // fixed for now
|
|
}))
|
|
}
|
|
|
|
async fn get_region(
|
|
&self,
|
|
request: Request<GetRegionRequest>,
|
|
) -> Result<Response<GetRegionResponse>, Status> {
|
|
let req = request.into_inner();
|
|
if let Some((region, leader)) = self.cluster.get_region_for_key(&req.key) {
|
|
Ok(Response::new(GetRegionResponse {
|
|
region: Some(region),
|
|
leader: Some(leader),
|
|
}))
|
|
} else {
|
|
Err(Status::not_found("Region not initialized yet"))
|
|
}
|
|
}
|
|
|
|
async fn list_regions(
|
|
&self,
|
|
_request: Request<ListRegionsRequest>,
|
|
) -> Result<Response<ListRegionsResponse>, Status> {
|
|
let (regions, stores) = self.cluster.list_regions();
|
|
Ok(Response::new(ListRegionsResponse { regions, stores }))
|
|
}
|
|
}
|