2308dd0
use std::time::Duration;
use tonic::metadata::{Ascii, MetadataValue};
use tonic::service::Interceptor;
use tonic::service::interceptor::InterceptedService;
use tonic::transport::{Channel, Endpoint};
use tonic::{Request, Status};
use uuid::Uuid;
use crate::error::{KenShellError, KenShellResult};
use crate::proto;
/// 64 MiB：KenShell 的 blob/archive 流与 pack 分片都按这个量级返回。
const DEFAULT_MESSAGE_SIZE: usize = 64 * 1024 * 1024;
#[derive(Clone, Default)]
pub struct RpcInterceptor {
key: Option<MetadataValue<Ascii>>,
}
impl RpcInterceptor {
pub fn new(rpc_key: &str) -> KenShellResult<Self> {
let key =
if rpc_key.is_empty() {
None
} else {
Some(MetadataValue::try_from(rpc_key).map_err(|_| {
KenShellError::Config("rpc_key is not valid header value".into())
})?)
};
Ok(Self { key })
}
}
impl Interceptor for RpcInterceptor {
fn call(&mut self, mut request: Request<()>) -> Result<Request<()>, Status> {
if let Some(key) = &self.key {
request
.metadata_mut()
.insert(crate::RPC_KEY_HEADER, key.clone());
}
if let Ok(value) = MetadataValue::try_from(Uuid::now_v7().to_string()) {
request
.metadata_mut()
.insert(crate::REQUEST_ID_HEADER, value);
}
Ok(request)
}
}
pub type Intercepted = InterceptedService<Channel, RpcInterceptor>;
#[derive(Clone)]
pub struct KenShellClient {
channel: Channel,
interceptor: RpcInterceptor,
request_timeout: Duration,
connect_timeout: Duration,
}
macro_rules! accessor {
($name:ident, $client:ty) => {
pub fn $name(&self) -> $client {
<$client>::new(self.intercepted())
.max_decoding_message_size(DEFAULT_MESSAGE_SIZE)
.max_encoding_message_size(DEFAULT_MESSAGE_SIZE)
}
};
}
impl KenShellClient {
pub fn connect(endpoint: &str, request_timeout: Duration) -> KenShellResult<Self> {
Self::connect_with_key(endpoint, "", request_timeout)
}
pub fn connect_with_key(
endpoint: &str,
rpc_key: &str,
request_timeout: Duration,
) -> KenShellResult<Self> {
if endpoint.is_empty() {
return Err(KenShellError::Config("kenshell endpoint is empty".into()));
}
let connect_timeout = Duration::from_secs(3);
let channel = Endpoint::from_shared(endpoint.to_string())
.map_err(|error| KenShellError::Config(error.to_string()))?
.connect_timeout(connect_timeout)
.timeout(request_timeout)
.connect_lazy();
Ok(Self {
channel,
interceptor: RpcInterceptor::new(rpc_key)?,
request_timeout,
connect_timeout,
})
}
pub fn channel(&self) -> Channel {
self.channel.clone()
}
pub fn connect_timeout(&self) -> Duration {
self.connect_timeout
}
pub fn request_timeout(&self) -> Duration {
self.request_timeout
}
pub fn set_timeout<T>(&self, request: &mut Request<T>) {
request.set_timeout(self.request_timeout);
}
fn intercepted(&self) -> Intercepted {
InterceptedService::new(self.channel.clone(), self.interceptor.clone())
}
accessor!(
repository,
proto::v1::repository_service_client::RepositoryServiceClient<Intercepted>
);
accessor!(
blob,
proto::v1::blob_service_client::BlobServiceClient<Intercepted>
);
accessor!(
tree,
proto::v1::tree_service_client::TreeServiceClient<Intercepted>
);
accessor!(
commit,
proto::v1::commit_service_client::CommitServiceClient<Intercepted>
);
accessor!(
diff,
proto::v1::diff_service_client::DiffServiceClient<Intercepted>
);
accessor!(
refs,
proto::v1::ref_service_client::RefServiceClient<Intercepted>
);
accessor!(
command,
proto::v1::command_service_client::CommandServiceClient<Intercepted>
);
accessor!(
remote,
proto::v1::remote_service_client::RemoteServiceClient<Intercepted>
);
accessor!(
job,
proto::v1::job_service_client::JobServiceClient<Intercepted>
);
accessor!(
lfs,
proto::v1::lfs_service_client::LfsServiceClient<Intercepted>
);
accessor!(
archive,
proto::v1::archive_service_client::ArchiveServiceClient<Intercepted>
);
accessor!(
analysis,
proto::v1::analysis_service_client::AnalysisServiceClient<Intercepted>
);
accessor!(
search,
proto::v1::search_service_client::SearchServiceClient<Intercepted>
);
accessor!(
conflict,
proto::v1::conflict_service_client::ConflictServiceClient<Intercepted>
);
accessor!(
comment,
proto::v1::code_comment_service_client::CodeCommentServiceClient<Intercepted>
);
accessor!(
repository_admin,
proto::internal::v1::repository_admin_service_client::RepositoryAdminServiceClient<
Intercepted,
>
);
accessor!(
node_admin,
proto::internal::v1::node_admin_service_client::NodeAdminServiceClient<Intercepted>
);
accessor!(
volume_admin,
proto::internal::v1::volume_admin_service_client::VolumeAdminServiceClient<Intercepted>
);
accessor!(
lfs_admin,
proto::internal::v1::lfs_admin_service_client::LfsAdminServiceClient<Intercepted>
);
accessor!(
push_lock_admin,
proto::internal::v1::push_lock_admin_service_client::PushLockAdminServiceClient<
Intercepted,
>
);
accessor!(
maintenance_admin,
proto::internal::v1::maintenance_admin_service_client::MaintenanceAdminServiceClient<
Intercepted,
>
);
}