mirror of
https://github.com/rustfs/rustfs.git
synced 2026-08-11 23:56:53 +00:00
houseme
87 lines
3.0 KiB
Rust
87 lines
3.0 KiB
Rust
// Copyright 2024 RustFS 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 crate::rpc::{TONIC_RPC_PREFIX, gen_signature_headers};
|
|
use http::Method;
|
|
use rustfs_common::GLOBAL_CONN_MAP;
|
|
use rustfs_protos::{create_new_channel, proto_gen::node_service::node_service_client::NodeServiceClient};
|
|
use std::error::Error;
|
|
use tonic::{service::interceptor::InterceptedService, transport::Channel};
|
|
use tracing::debug;
|
|
|
|
/// 3. Subsequent calls will attempt fresh connections
|
|
/// 4. If node is still down, connection will fail fast (3s timeout)
|
|
pub async fn node_service_time_out_client(
|
|
addr: &String,
|
|
interceptor: TonicInterceptor,
|
|
) -> Result<NodeServiceClient<InterceptedService<Channel, TonicInterceptor>>, Box<dyn Error>> {
|
|
// Try to get cached channel
|
|
let cached_channel = { GLOBAL_CONN_MAP.read().await.get(addr).cloned() };
|
|
|
|
let channel = match cached_channel {
|
|
Some(channel) => {
|
|
debug!("Using cached gRPC channel for: {}", addr);
|
|
channel
|
|
}
|
|
None => {
|
|
// No cached connection, create new one
|
|
create_new_channel(addr).await?
|
|
}
|
|
};
|
|
|
|
Ok(NodeServiceClient::with_interceptor(channel, interceptor))
|
|
}
|
|
|
|
pub async fn node_service_time_out_client_no_auth(
|
|
addr: &String,
|
|
) -> Result<NodeServiceClient<InterceptedService<Channel, TonicInterceptor>>, Box<dyn Error>> {
|
|
node_service_time_out_client(addr, TonicInterceptor::NoOp(NoOpInterceptor)).await
|
|
}
|
|
|
|
pub struct TonicSignatureInterceptor;
|
|
|
|
impl tonic::service::Interceptor for TonicSignatureInterceptor {
|
|
fn call(&mut self, mut req: tonic::Request<()>) -> Result<tonic::Request<()>, tonic::Status> {
|
|
let headers = gen_signature_headers(TONIC_RPC_PREFIX, &Method::GET);
|
|
req.metadata_mut().as_mut().extend(headers);
|
|
Ok(req)
|
|
}
|
|
}
|
|
|
|
pub fn gen_tonic_signature_interceptor() -> TonicSignatureInterceptor {
|
|
TonicSignatureInterceptor
|
|
}
|
|
|
|
pub struct NoOpInterceptor;
|
|
|
|
impl tonic::service::Interceptor for NoOpInterceptor {
|
|
fn call(&mut self, req: tonic::Request<()>) -> Result<tonic::Request<()>, tonic::Status> {
|
|
Ok(req)
|
|
}
|
|
}
|
|
|
|
pub enum TonicInterceptor {
|
|
Signature(TonicSignatureInterceptor),
|
|
NoOp(NoOpInterceptor),
|
|
}
|
|
|
|
impl tonic::service::Interceptor for TonicInterceptor {
|
|
fn call(&mut self, req: tonic::Request<()>) -> Result<tonic::Request<()>, tonic::Status> {
|
|
match self {
|
|
TonicInterceptor::Signature(interceptor) => interceptor.call(req),
|
|
TonicInterceptor::NoOp(interceptor) => interceptor.call(req),
|
|
}
|
|
}
|
|
}
|