ri/grpc/
client.rs

1//! Copyright 2025-2026 Wenze Wei. All Rights Reserved.
2//!
3//! This file is part of Ri.
4//! The Ri project belongs to the Dunimd Team.
5//!
6//! Licensed under the Apache License, Version 2.0 (the "License");
7//! You may not use this file except in compliance with the License.
8//! You may obtain a copy of the License at
9//!
10//!     http://www.apache.org/licenses/LICENSE-2.0
11//!
12//! Unless required by applicable law or agreed to in writing, software
13//! distributed under the License is distributed on an "AS IS" BASIS,
14//! WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15//! See the License for the specific language governing permissions and
16//! limitations under the License.
17
18//! # gRPC Client Implementation
19
20use super::*;
21use std::time::Duration;
22use std::sync::atomic::{AtomicU64, Ordering};
23
24pub struct RiGrpcClient {
25    channel: Option<tonic::transport::Channel>,
26    endpoint: String,
27    timeout: Duration,
28    stats: Arc<RwLock<RiGrpcStats>>,
29    request_id: Arc<AtomicU64>,
30    connected: Arc<RwLock<bool>>,
31    retry_count: u32,
32    retry_delay: Duration,
33}
34
35#[cfg(feature = "pyo3")]
36#[pyclass]
37pub struct RiGrpcClientPy {
38    inner: RiGrpcClient,
39}
40
41#[cfg(feature = "pyo3")]
42#[pymethods]
43impl RiGrpcClientPy {
44    #[new]
45    fn new(endpoint: String) -> Self {
46        Self {
47            inner: RiGrpcClient::new(endpoint),
48        }
49    }
50
51    #[pyo3(signature = (timeout_secs=30))]
52    fn with_timeout(&mut self, timeout_secs: u64) {
53        self.inner.timeout = Duration::from_secs(timeout_secs);
54    }
55
56    #[pyo3(signature = (count=3, delay_ms=100))]
57    fn with_retry(&mut self, count: u32, delay_ms: u64) {
58        self.inner.retry_count = count;
59        self.inner.retry_delay = Duration::from_millis(delay_ms);
60    }
61
62    fn get_stats(&self) -> RiGrpcStats {
63        self.inner.get_stats()
64    }
65
66    fn is_connected(&self) -> bool {
67        self.inner.channel.is_some()
68    }
69
70    fn get_endpoint(&self) -> String {
71        self.inner.endpoint.clone()
72    }
73
74    fn connect(&mut self) -> PyResult<()> {
75        let rt = crate::py_runtime::py_runtime();
76        
77        rt.block_on(async {
78            self.inner.connect().await
79        }).map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string()))
80    }
81
82    fn disconnect(&mut self) {
83        let rt = crate::py_runtime::py_runtime();
84        rt.block_on(async {
85            self.inner.disconnect().await
86        });
87    }
88
89    #[pyo3(signature = (service_name, method, data))]
90    fn call(&mut self, service_name: String, method: String, data: Vec<u8>) -> PyResult<Vec<u8>> {
91        let rt = crate::py_runtime::py_runtime();
92        
93        rt.block_on(async {
94            self.inner.call(&service_name, &method, &data).await
95        }).map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string()))
96    }
97}
98
99impl RiGrpcClient {
100    pub fn new(endpoint: String) -> Self {
101        Self {
102            channel: None,
103            endpoint,
104            timeout: Duration::from_secs(30),
105            stats: Arc::new(RwLock::new(RiGrpcStats::new())),
106            request_id: Arc::new(AtomicU64::new(0)),
107            connected: Arc::new(RwLock::new(false)),
108            retry_count: 3,
109            retry_delay: Duration::from_millis(100),
110        }
111    }
112
113    pub fn with_timeout(mut self, timeout: Duration) -> Self {
114        self.timeout = timeout;
115        self
116    }
117
118    pub fn with_retry(mut self, count: u32, delay: Duration) -> Self {
119        self.retry_count = count;
120        self.retry_delay = delay;
121        self
122    }
123
124    pub fn get_stats(&self) -> RiGrpcStats {
125        self.stats.try_read()
126            .map(|guard| guard.clone())
127            .unwrap_or_else(|_| RiGrpcStats::new())
128    }
129
130    pub async fn connect(&mut self) -> RiResult<()> {
131        let endpoint = tonic::transport::Endpoint::from_shared(self.endpoint.clone())
132            .map_err(|e| GrpcError::ConnectionFailed {
133                message: format!("Invalid endpoint: {}", e)
134            })?
135            .connect_timeout(self.timeout)
136            .timeout(self.timeout);
137
138        let channel = endpoint.connect()
139            .await
140            .map_err(|e| GrpcError::ConnectionFailed {
141                message: format!("Connection failed: {}", e)
142            })?;
143
144        self.channel = Some(channel);
145        *self.connected.write().await = true;
146
147        tracing::info!("gRPC client connected to {}", self.endpoint);
148        Ok(())
149    }
150
151    pub async fn is_connected(&self) -> bool {
152        *self.connected.read().await && self.channel.is_some()
153    }
154
155    fn generate_request_id(&self) -> u64 {
156        self.request_id.fetch_add(1, Ordering::SeqCst)
157    }
158
159    pub async fn call(&mut self, service_name: &str, method: &str, data: &[u8]) -> RiResult<Vec<u8>> {
160        let channel = match &self.channel {
161            Some(ch) => ch.clone(),
162            None => {
163                return Err(GrpcError::Client {
164                    message: "Not connected to gRPC server".to_string()
165                }.into());
166            }
167        };
168
169        if !*self.connected.read().await {
170            return Err(GrpcError::Client {
171                message: "gRPC client not connected".to_string()
172            }.into());
173        }
174
175        let request_id = self.generate_request_id();
176        let path = format!("/{}/{}", service_name, method);
177
178        tracing::debug!("gRPC call: {} (request_id={})", path, request_id);
179
180        let mut last_error: Option<RiError> = None;
181        for attempt in 0..=self.retry_count {
182            if attempt > 0 {
183                tokio::time::sleep(self.retry_delay).await;
184                tracing::warn!("Retrying gRPC call (attempt {}/{})", attempt, self.retry_count);
185            }
186
187            match Self::execute_unary_call(channel.clone(), &path, data).await {
188                Ok(response) => {
189                    let mut stats = self.stats.write().await;
190                    stats.record_request(data.len());
191                    stats.record_response(response.len());
192                    return Ok(response);
193                }
194                Err(e) => {
195                    last_error = Some(e.clone());
196                    let mut stats = self.stats.write().await;
197                    stats.record_error();
198                    
199                    if !Self::is_retryable_error(&e) {
200                        return Err(e);
201                    }
202                }
203            }
204        }
205
206        Err(last_error.unwrap_or_else(|| GrpcError::Client {
207            message: "Unknown error after retries".to_string()
208        }.into()))
209    }
210
211    async fn execute_unary_call(
212        channel: tonic::transport::Channel,
213        path: &str,
214        data: &[u8],
215    ) -> RiResult<Vec<u8>> {
216        use tonic::client::Grpc;
217        use tonic::codec::ProstCodec;
218        
219        let codec = ProstCodec::<Vec<u8>, Vec<u8>>::new();
220        let mut client = Grpc::new(channel);
221        
222        let request = tonic::Request::new(data.to_vec());
223        let path_and_query: http::uri::PathAndQuery = path.parse()
224            .map_err(|e| GrpcError::Client {
225                message: format!("Invalid path: {}", e)
226            })?;
227        
228        let response = client.unary(request, path_and_query, codec)
229            .await
230            .map_err(|e| GrpcError::Client {
231                message: format!("RPC call failed: {}", e)
232            })?;
233
234        Ok(response.into_inner())
235    }
236
237    fn is_retryable_error(error: &RiError) -> bool {
238        let error_str = error.to_string();
239        error_str.contains("UNAVAILABLE") ||
240        error_str.contains("DEADLINE_EXCEEDED") ||
241        error_str.contains("RESOURCE_EXHAUSTED")
242    }
243
244    pub async fn disconnect(&mut self) {
245        self.channel.take();
246        *self.connected.write().await = false;
247        tracing::info!("gRPC client disconnected from {}", self.endpoint);
248    }
249}
250
251impl Drop for RiGrpcClient {
252    fn drop(&mut self) {
253        if self.channel.is_some() {
254            let rt = crate::py_runtime::py_runtime();
255            rt.block_on(async {
256                self.disconnect().await;
257            });
258        }
259    }
260}
261
262impl Default for RiGrpcClient {
263    fn default() -> Self {
264        Self::new("http://127.0.0.1:50051".to_string())
265    }
266}
267
268impl Clone for RiGrpcClient {
269    fn clone(&self) -> Self {
270        Self {
271            channel: self.channel.clone(),
272            endpoint: self.endpoint.clone(),
273            timeout: self.timeout,
274            stats: self.stats.clone(),
275            request_id: self.request_id.clone(),
276            connected: self.connected.clone(),
277            retry_count: self.retry_count,
278            retry_delay: self.retry_delay,
279        }
280    }
281}
ri/grpc/client.rs

ri/grpc/
client.rs
