tina-core 0.0.2

#![cfg(feature = "client-tonic")]
//! tonic gRPC client

use std::{
    fmt::Debug,
    str::FromStr,
    sync::{Arc, RwLock},
};

use dashmap::{mapref::one::RefMut, DashMap};
use prost::Message;
use tonic::{client::Grpc, metadata::MetadataValue, GrpcMethod, Status};
use tracing_subscriber::{
    registry::{LookupSpan, SpanData},
    Registry,
};

use crate::{
    app_error_from, app_system_error,
    tina::{
        constant::Constants,
        data::{app_error::AppError, AppResult},
        grpc::{FromGrpcResponse, IntoGrpcRequest},
        log::AsyncLoggerFields,
    },
};

use super::GrpcClientProps;

/// gRPC 请求
pub type GrpcRequest<M> = tonic::Request<M>;
/// gRPC 响应
pub type GrpcResponse<M> = tonic::Response<M>;
/// gRPC channel
pub type GrpcChannel = tonic::transport::Channel;

#[derive(Default)]
/// 基于tonic的gRPC client
pub struct GrpcClient {
    pub(crate) channels: DashMap<GrpcClientProps, Grpc<GrpcChannel>>,
}

impl Debug for GrpcClient {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("GrpcClient").finish()
    }
}

#[allow(dead_code)]
impl GrpcClient {
    /// 获取服务Client实例
    pub async fn get_channel(&self, props: &GrpcClientProps) -> AppResult<RefMut<'_, GrpcClientProps, Grpc<GrpcChannel>>> {
        match self.channels.get_mut(props) {
            Some(v) => Ok(v),
            None => self.new_channel(props).await,
        }
    }
    /// 新建channel
    pub(crate) async fn new_channel(&self, props: &GrpcClientProps) -> AppResult<RefMut<'_, GrpcClientProps, Grpc<GrpcChannel>>> {
        let address = props.address.trim();
        let address = match address.starts_with(Constants::HTTP) {
            true => address.to_string(),
            false => match address.starts_with(Constants::HTTPS) {
                true => address.to_string(),
                false => format!("{}{}", Constants::HTTP, address),
            },
        };
        let endpoint = GrpcChannel::builder(address.parse().map_err(app_error_from!())?);
        let channel = endpoint.connect().await.map_err(app_error_from!())?;
        self.channels.insert(props.clone(), Grpc::new(channel));
        self.channels.get_mut(props).ok_or_else(|| app_system_error!("No channel cache found"))
    }
    /// 移除channel
    pub(crate) fn remove_channel(&self, props: &GrpcClientProps) {
        self.channels.remove(props);
    }
    /// 请求
    #[instrument(name = "send grpc", level = "debug", fields(service_name, method_name))]
    pub async fn send<Req, ReqM, Res, ResM>(
        &self,
        props: &GrpcClientProps,
        data: Req,
        service_name: &'static str,
        method_name: &'static str,
    ) -> AppResult<Res>
    where
        Req: IntoGrpcRequest<Request = tonic::Request<ReqM>> + Debug + Send + Sync + 'static,
        ReqM: Message + Debug + Send + Sync + 'static,
        Res: FromGrpcResponse<Response = tonic::Response<ResM>, Rejection = AppError> + Debug + Send + Sync + 'static,
        ResM: Message + Default + Debug + Send + Sync + 'static,
    {
        let (request_id, span_id): (Option<String>, Option<u64>) = tracing::dispatcher::get_default(|dispatcher| {
            let span = dispatcher.current_span();
            if let Some(id) = span.id() {
                let span_id = id.into_u64();
                if let Some(v) = dispatcher.downcast_ref::<Registry>() {
                    if let Some(v) = v.span_data(id) {
                        let v = v.extensions();
                        if let Some(v) = v.get::<Arc<RwLock<AsyncLoggerFields>>>() {
                            if let Ok(v) = v.read() {
                                let values = v.lookup_field_values("request_id");
                                return match values.is_empty() {
                                    true => (None, Some(span_id)),
                                    false => (Some(values.join(",")), Some(span_id)),
                                };
                            }
                        }
                        return (None, Some(span_id));
                    }
                }
            }

            (None, None)
        });
        let mut grpc = self.get_channel(props).await?;
        {
            grpc.ready().await.map_err(app_error_from!())?;
        }
        let codec = tonic::codec::ProstCodec::default();
        let service_path = format!("/{service_name}/{method_name}");
        let path = http::uri::PathAndQuery::from_str(service_path.as_str()).map_err(app_error_from!())?;
        tracing::debug!("grpc send call: path = {service_path}, req = {data:?}");
        let mut req = data.into_grpc_request().await;

        if let Some(request_id) = request_id {
            let request_id_value = match MetadataValue::try_from(request_id.as_str()) {
                Ok(v) => Some(v),
                Err(err) => {
                    tracing::error!("parse request id header value failed, err: {err:?}, value: {request_id}");
                    None
                }
            };
            if let Some(v) = request_id_value {
                req.metadata_mut().insert(Constants::REQUEST_ID_HEADER_VALUE, v);
            }
        }

        if let Some(span_id) = span_id {
            req.metadata_mut().insert(Constants::SPAN_ID_HEADER, MetadataValue::from(span_id));
        }

        req.extensions_mut().insert(GrpcMethod::new(service_name, method_name));
        let r: Result<tonic::Response<ResM>, Status> = grpc.unary(req, path, codec).await;
        tracing::trace!("grpc receive call: path = {service_path}, res = {r:?}");
        match r {
            Ok(res) => match Res::from_grpc_response(res).await {
                Ok(v) => {
                    tracing::debug!("grpc receive ok call: path = {service_path}, res = {v:?}");
                    Ok(v)
                }
                Err(err) => {
                    tracing::error!("grpc receive err call failed: path = {service_path}, err = {err:?}");
                    Err(err)
                }
            },
            Err(err) => {
                tracing::error!("grpc receive status error call: path = {service_path}, err = {err:?}");
                let err = AppError::from(err);
                Err(app_error_from!(err))
            }
        }
    }
}