flare_rpc_core/client/
client.rs

1use std::marker::PhantomData;
2use tonic::transport::Channel;
3use crate::discover::{RpcDiscovery, ServiceError};
4
5pub trait GrpcClient {
6    fn new(channel: Channel) -> Self;
7}
8
9/// RPC 客户端工厂
10/// T: tonic 生成的客户端类型 (例如: echo_client::EchoClient<Channel>)
11/// D: 服务发现实现
12pub struct RpcClient<T, D>
13where
14    T: GrpcClient,
15    D: RpcDiscovery,
16{
17    service_name: String,
18    discovery: D,
19    protocol: String,
20    _marker: PhantomData<T>,
21}
22
23impl<T, D> RpcClient<T, D>
24where
25    T: GrpcClient,
26    D: RpcDiscovery,
27{
28    pub fn new(service_name: impl Into<String>, discovery: D) -> Self {
29        Self {
30            service_name: service_name.into(),
31            discovery,
32            protocol: "http".to_string(),
33            _marker: PhantomData,
34        }
35    }
36
37    pub fn with_protocol(mut self, protocol: impl Into<String>) -> Self {
38        self.protocol = protocol.into();
39        self
40    }
41
42    /// 获取一个可用的客户端
43    pub async fn client(&self) -> Result<T, ServiceError> {
44        let endpoint = self.discovery.discover(&self.service_name).await?;
45        
46        let url = format!("{}://{}:{}", self.protocol, endpoint.address, endpoint.port);
47        let channel = Channel::from_shared(url)
48            .map_err(|e| ServiceError::ConnectionError(e.to_string()))?
49            .connect()
50            .await
51            .map_err(|e| ServiceError::ConnectionError(e.to_string()))?;
52
53        Ok(T::new(channel))
54    }
55
56    /// 获取服务发现实例的引用
57    pub fn discovery(&self) -> &D {
58        &self.discovery
59    }
60}