ri/grpc/
mod.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 Support
19
20use crate::core::RiResult;
21use async_trait::async_trait;
22use std::sync::Arc;
23use tokio::sync::RwLock;
24#[cfg(feature = "grpc")]
25use std::collections::HashMap as FxHashMap;
26
27#[cfg(feature = "pyo3")]
28use pyo3::prelude::*;
29
30#[cfg(feature = "grpc")]
31mod server;
32#[cfg(feature = "grpc")]
33mod client;
34
35#[cfg(feature = "grpc")]
36pub use server::RiGrpcServer;
37#[cfg(feature = "grpc")]
38pub use client::RiGrpcClient;
39
40#[cfg(all(feature = "grpc", feature = "pyo3"))]
41pub use server::RiGrpcServerPy;
42#[cfg(all(feature = "grpc", feature = "pyo3"))]
43pub use client::RiGrpcClientPy;
44
45#[derive(Debug, Clone, PartialEq, Eq)]
46#[cfg_attr(feature = "pyo3", pyclass)]
47pub struct RiGrpcConfig {
48    pub addr: String,
49    pub port: u16,
50    pub max_concurrent_requests: u32,
51    pub enable_tls: bool,
52    pub cert_path: Option<String>,
53    pub key_path: Option<String>,
54    pub enable_auth: bool,
55    pub auth_token: Option<String>,
56}
57
58#[cfg(feature = "pyo3")]
59#[pymethods]
60impl RiGrpcConfig {
61    #[new]
62    fn new() -> Self {
63        Self::default()
64    }
65    
66    #[getter]
67    fn get_addr(&self) -> String {
68        self.addr.clone()
69    }
70    
71    #[setter]
72    fn set_addr(&mut self, addr: String) {
73        self.addr = addr;
74    }
75    
76    #[getter]
77    fn get_port(&self) -> u16 {
78        self.port
79    }
80    
81    #[setter]
82    fn set_port(&mut self, port: u16) {
83        self.port = port;
84    }
85    
86    #[getter]
87    fn get_max_concurrent_requests(&self) -> u32 {
88        self.max_concurrent_requests
89    }
90    
91    #[setter]
92    fn set_max_concurrent_requests(&mut self, max_concurrent_requests: u32) {
93        self.max_concurrent_requests = max_concurrent_requests;
94    }
95    
96    #[getter]
97    fn get_enable_tls(&self) -> bool {
98        self.enable_tls
99    }
100    
101    #[setter]
102    fn set_enable_tls(&mut self, enable_tls: bool) {
103        self.enable_tls = enable_tls;
104    }
105    
106    #[getter]
107    fn get_cert_path(&self) -> Option<String> {
108        self.cert_path.clone()
109    }
110    
111    #[setter]
112    fn set_cert_path(&mut self, cert_path: Option<String>) {
113        self.cert_path = cert_path;
114    }
115    
116    #[getter]
117    fn get_key_path(&self) -> Option<String> {
118        self.key_path.clone()
119    }
120    
121    #[setter]
122    fn set_key_path(&mut self, key_path: Option<String>) {
123        self.key_path = key_path;
124    }
125}
126
127impl Default for RiGrpcConfig {
128    fn default() -> Self {
129        Self {
130            addr: "127.0.0.1".to_string(),
131            port: 50051,
132            max_concurrent_requests: 100,
133            enable_tls: false,
134            cert_path: None,
135            key_path: None,
136            enable_auth: false,
137            auth_token: None,
138        }
139    }
140}
141
142#[cfg(feature = "grpc")]
143#[async_trait]
144pub trait RiGrpcService: Send + Sync {
145    async fn handle_request(&self, method: &str, data: &[u8]) -> RiResult<Vec<u8>>;
146    fn service_name(&self) -> &'static str;
147}
148
149#[cfg(feature = "grpc")]
150#[derive(Clone)]
151pub struct RiGrpcServiceRegistry {
152    pub services: Arc<RwLock<FxHashMap<String, Arc<dyn RiGrpcService>>>>,
153}
154
155#[cfg(feature = "grpc")]
156impl RiGrpcServiceRegistry {
157    pub fn new() -> Self {
158        Self {
159            services: Arc::new(RwLock::new(FxHashMap::default())),
160        }
161    }
162
163    pub fn register_service(&mut self, service: Arc<dyn RiGrpcService>) {
164        let name = service.service_name().to_string();
165        let mut services = self.services.blocking_write();
166        services.insert(name, service);
167    }
168
169    pub fn list_services(&self) -> Vec<String> {
170        let services = self.services.blocking_read();
171        services.keys().cloned().collect()
172    }
173}
174
175#[cfg(feature = "grpc")]
176impl Default for RiGrpcServiceRegistry {
177    fn default() -> Self {
178        Self::new()
179    }
180}
181
182#[cfg(all(feature = "grpc", feature = "pyo3"))]
183#[pyclass]
184pub struct RiGrpcServiceRegistryPy {
185    registry: RiGrpcServiceRegistry,
186}
187
188#[cfg(all(feature = "grpc", feature = "pyo3"))]
189#[pymethods]
190impl RiGrpcServiceRegistryPy {
191    #[new]
192    fn new() -> Self {
193        Self {
194            registry: RiGrpcServiceRegistry::new(),
195        }
196    }
197    
198    fn register(&mut self, service_name: &str, handler: Py<PyAny>) {
199        let service = RiGrpcPythonService::new(service_name, handler);
200        self.registry.register_service(Arc::new(service));
201    }
202    
203    fn list_services(&self) -> Vec<String> {
204        self.registry.list_services()
205    }
206}
207
208#[cfg(all(feature = "grpc", feature = "pyo3"))]
209impl Default for RiGrpcServiceRegistryPy {
210    fn default() -> Self {
211        Self::new()
212    }
213}
214
215#[derive(Debug, Clone)]
216#[cfg_attr(feature = "pyo3", pyclass)]
217pub struct RiGrpcStats {
218    pub requests_received: u64,
219    pub requests_completed: u64,
220    pub requests_failed: u64,
221    pub bytes_received: u64,
222    pub bytes_sent: u64,
223    pub active_connections: u64,
224}
225
226#[cfg(feature = "pyo3")]
227#[pymethods]
228impl RiGrpcStats {
229    #[getter]
230    fn get_requests_received(&self) -> u64 {
231        self.requests_received
232    }
233
234    #[getter]
235    fn get_requests_completed(&self) -> u64 {
236        self.requests_completed
237    }
238
239    #[getter]
240    fn get_requests_failed(&self) -> u64 {
241        self.requests_failed
242    }
243
244    #[getter]
245    fn get_bytes_received(&self) -> u64 {
246        self.bytes_received
247    }
248
249    #[getter]
250    fn get_bytes_sent(&self) -> u64 {
251        self.bytes_sent
252    }
253
254    #[getter]
255    fn get_active_connections(&self) -> u64 {
256        self.active_connections
257    }
258}
259
260impl RiGrpcStats {
261    pub fn new() -> Self {
262        Self {
263            requests_received: 0,
264            requests_completed: 0,
265            requests_failed: 0,
266            bytes_received: 0,
267            bytes_sent: 0,
268            active_connections: 0,
269        }
270    }
271
272    pub fn record_request(&mut self, size: usize) {
273        self.requests_received += 1;
274        self.bytes_received += size as u64;
275        self.active_connections += 1;
276    }
277
278    pub fn record_response(&mut self, size: usize) {
279        self.requests_completed += 1;
280        self.bytes_sent += size as u64;
281        if self.active_connections > 0 {
282            self.active_connections -= 1;
283        }
284    }
285
286    pub fn record_error(&mut self) {
287        self.requests_failed += 1;
288        if self.active_connections > 0 {
289            self.active_connections -= 1;
290        }
291    }
292}
293
294impl Default for RiGrpcStats {
295    fn default() -> Self {
296        Self::new()
297    }
298}
299
300#[cfg(all(feature = "grpc", feature = "pyo3"))]
301#[pyclass]
302pub struct RiGrpcPythonService {
303    service_name: String,
304    handler: Py<PyAny>,
305}
306
307#[cfg(all(feature = "grpc", feature = "pyo3"))]
308impl RiGrpcPythonService {
309    pub fn new(service_name: &str, handler: Py<PyAny>) -> Self {
310        Self {
311            service_name: service_name.to_string(),
312            handler,
313        }
314    }
315}
316
317#[cfg(all(feature = "grpc", feature = "pyo3"))]
318#[async_trait]
319impl RiGrpcService for RiGrpcPythonService {
320    async fn handle_request(&self, method: &str, data: &[u8]) -> RiResult<Vec<u8>> {
321        let method_str = method.to_string();
322        let data_vec = data.to_vec();
323        
324        let result = pyo3::Python::attach(|py| {
325            self.handler.call1(py, (method_str, data_vec))
326        });
327        
328        match result {
329            Ok(obj) => {
330                let result_vec = pyo3::Python::attach(|py| {
331                    obj.extract::<Vec<u8>>(py)
332                });
333                match result_vec {
334                    Ok(bytes) => Ok(bytes),
335                    Err(e) => Err(RiError::Other(format!("Failed to extract response bytes: {:?}", e))),
336                }
337            }
338            Err(e) => Err(RiError::Other(format!("Python handler error: {:?}", e))),
339        }
340    }
341    
342    fn service_name(&self) -> &'static str {
343        Box::leak(self.service_name.clone().into_boxed_str())
344    }
345}
346
347#[derive(Debug, thiserror::Error)]
348pub enum GrpcError {
349    #[error("Server error: {message}")]
350    Server { message: String },
351    #[error("Client error: {message}")]
352    Client { message: String },
353    #[error("Service not found: {service_name}")]
354    ServiceNotFound { service_name: String },
355    #[error("Connection failed: {message}")]
356    ConnectionFailed { message: String },
357    #[error("Request timeout")]
358    Timeout,
359}
360
361impl From<GrpcError> for RiError {
362    fn from(error: GrpcError) -> Self {
363        RiError::Other(format!("gRPC error: {}", error))
364    }
365}
366
367use crate::core::RiError;
ri/grpc/mod.rs

ri/grpc/
mod.rs
