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 = tokio::runtime::Runtime::new()
76            .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string()))?;
77        
78        rt.block_on(async {
79            self.inner.connect().await
80        }).map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string()))
81    }
82
83    fn disconnect(&mut self) {
84        if let Ok(rt) = tokio::runtime::Runtime::new() {
85            rt.block_on(async {
86                self.inner.disconnect().await
87            });
88        }
89    }
90
91    #[pyo3(signature = (service_name, method, data))]
92    fn call(&mut self, service_name: String, method: String, data: Vec<u8>) -> PyResult<Vec<u8>> {
93        let rt = tokio::runtime::Runtime::new()
94            .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string()))?;
95        
96        rt.block_on(async {
97            self.inner.call(&service_name, &method, &data).await
98        }).map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string()))
99    }
100}
101
102impl RiGrpcClient {
103    pub fn new(endpoint: String) -> Self {
104        Self {
105            channel: None,
106            endpoint,
107            timeout: Duration::from_secs(30),
108            stats: Arc::new(RwLock::new(RiGrpcStats::new())),
109            request_id: Arc::new(AtomicU64::new(0)),
110            connected: Arc::new(RwLock::new(false)),
111            retry_count: 3,
112            retry_delay: Duration::from_millis(100),
113        }
114    }
115
116    pub fn with_timeout(mut self, timeout: Duration) -> Self {
117        self.timeout = timeout;
118        self
119    }
120
121    pub fn with_retry(mut self, count: u32, delay: Duration) -> Self {
122        self.retry_count = count;
123        self.retry_delay = delay;
124        self
125    }
126
127    pub fn get_stats(&self) -> RiGrpcStats {
128        self.stats.try_read()
129            .map(|guard| guard.clone())
130            .unwrap_or_else(|_| RiGrpcStats::new())
131    }
132
133    pub async fn connect(&mut self) -> RiResult<()> {
134        let endpoint = tonic::transport::Endpoint::from_shared(self.endpoint.clone())
135            .map_err(|e| GrpcError::ConnectionFailed {
136                message: format!("Invalid endpoint: {}", e)
137            })?
138            .connect_timeout(self.timeout)
139            .timeout(self.timeout);
140
141        let channel = endpoint.connect()
142            .await
143            .map_err(|e| GrpcError::ConnectionFailed {
144                message: format!("Connection failed: {}", e)
145            })?;
146
147        self.channel = Some(channel);
148        *self.connected.write().await = true;
149
150        tracing::info!("gRPC client connected to {}", self.endpoint);
151        Ok(())
152    }
153
154    pub async fn is_connected(&self) -> bool {
155        *self.connected.read().await && self.channel.is_some()
156    }
157
158    fn generate_request_id(&self) -> u64 {
159        self.request_id.fetch_add(1, Ordering::SeqCst)
160    }
161
162    pub async fn call(&mut self, service_name: &str, method: &str, data: &[u8]) -> RiResult<Vec<u8>> {
163        let channel = match &self.channel {
164            Some(ch) => ch.clone(),
165            None => {
166                return Err(GrpcError::Client {
167                    message: "Not connected to gRPC server".to_string()
168                }.into());
169            }
170        };
171
172        if !*self.connected.read().await {
173            return Err(GrpcError::Client {
174                message: "gRPC client not connected".to_string()
175            }.into());
176        }
177
178        let request_id = self.generate_request_id();
179        let path = format!("/{}/{}", service_name, method);
180
181        tracing::debug!("gRPC call: {} (request_id={})", path, request_id);
182
183        let mut last_error: Option<RiError> = None;
184        for attempt in 0..=self.retry_count {
185            if attempt > 0 {
186                tokio::time::sleep(self.retry_delay).await;
187                tracing::warn!("Retrying gRPC call (attempt {}/{})", attempt, self.retry_count);
188            }
189
190            match Self::execute_unary_call(channel.clone(), &path, data).await {
191                Ok(response) => {
192                    let mut stats = self.stats.write().await;
193                    stats.record_request(data.len());
194                    stats.record_response(response.len());
195                    return Ok(response);
196                }
197                Err(e) => {
198                    last_error = Some(e.clone());
199                    let mut stats = self.stats.write().await;
200                    stats.record_error();
201                    
202                    if !Self::is_retryable_error(&e) {
203                        return Err(e);
204                    }
205                }
206            }
207        }
208
209        Err(last_error.unwrap_or_else(|| GrpcError::Client {
210            message: "Unknown error after retries".to_string()
211        }.into()))
212    }
213
214    async fn execute_unary_call(
215        channel: tonic::transport::Channel,
216        path: &str,
217        data: &[u8],
218    ) -> RiResult<Vec<u8>> {
219        use tonic::client::Grpc;
220        use tonic::codec::ProstCodec;
221        
222        let codec = ProstCodec::<Vec<u8>, Vec<u8>>::new();
223        let mut client = Grpc::new(channel);
224        
225        let request = tonic::Request::new(data.to_vec());
226        let path_and_query: http::uri::PathAndQuery = path.parse()
227            .map_err(|e| GrpcError::Client {
228                message: format!("Invalid path: {}", e)
229            })?;
230        
231        let response = client.unary(request, path_and_query, codec)
232            .await
233            .map_err(|e| GrpcError::Client {
234                message: format!("RPC call failed: {}", e)
235            })?;
236
237        Ok(response.into_inner())
238    }
239
240    fn is_retryable_error(error: &RiError) -> bool {
241        let error_str = error.to_string();
242        error_str.contains("UNAVAILABLE") ||
243        error_str.contains("DEADLINE_EXCEEDED") ||
244        error_str.contains("RESOURCE_EXHAUSTED")
245    }
246
247    pub async fn disconnect(&mut self) {
248        self.channel.take();
249        *self.connected.write().await = false;
250        tracing::info!("gRPC client disconnected from {}", self.endpoint);
251    }
252}
253
254impl Drop for RiGrpcClient {
255    fn drop(&mut self) {
256        if self.channel.is_some() {
257            if let Ok(rt) = tokio::runtime::Runtime::new() {
258                rt.block_on(async {
259                    self.disconnect().await;
260                });
261            }
262        }
263    }
264}
265
266impl Default for RiGrpcClient {
267    fn default() -> Self {
268        Self::new("http://127.0.0.1:50051".to_string())
269    }
270}
271
272impl Clone for RiGrpcClient {
273    fn clone(&self) -> Self {
274        Self {
275            channel: self.channel.clone(),
276            endpoint: self.endpoint.clone(),
277            timeout: self.timeout,
278            stats: self.stats.clone(),
279            request_id: self.request_id.clone(),
280            connected: self.connected.clone(),
281            retry_count: self.retry_count,
282            retry_delay: self.retry_delay,
283        }
284    }
285}
ri/grpc/client.rs

ri/grpc/
client.rs
