1#[cfg(not(reifydb_single_threaded))]
5use std::{collections::HashMap, sync::mpsc};
6
7#[cfg(not(reifydb_single_threaded))]
8use reifydb_client::{GrpcClient, WireFormat};
9#[cfg(not(reifydb_single_threaded))]
10use reifydb_runtime::sync::mutex::Mutex;
11#[cfg(not(reifydb_single_threaded))]
12use reifydb_value::error::Diagnostic;
13use reifydb_value::error::Error;
14#[cfg(not(reifydb_single_threaded))]
15use reifydb_value::{params::Params, value::frame::frame::Frame};
16#[cfg(not(reifydb_single_threaded))]
17use tokio::runtime::Handle;
18
19#[cfg(not(reifydb_single_threaded))]
20type CacheKey = (String, Option<String>);
21
22#[cfg(not(reifydb_single_threaded))]
23pub struct RemoteRegistry {
24 handle: Mutex<Option<Handle>>,
25 clients: Mutex<HashMap<CacheKey, GrpcClient>>,
26}
27
28#[cfg(not(reifydb_single_threaded))]
29impl RemoteRegistry {
30 pub fn new(handle: Handle) -> Self {
31 Self {
32 handle: Mutex::new(Some(handle)),
33 clients: Mutex::new(HashMap::new()),
34 }
35 }
36
37 pub fn shutdown(&self) {
38 self.clients.lock().clear();
39 *self.handle.lock() = None;
40 }
41
42 fn handle(&self) -> Result<Handle, Error> {
43 self.handle.lock().clone().ok_or_else(|| {
44 Error(Box::new(Diagnostic {
45 code: "REMOTE_002".to_string(),
46 message: "remote runtime has been shut down".to_string(),
47 ..Default::default()
48 }))
49 })
50 }
51
52 pub fn forward_query(
53 &self,
54 address: &str,
55 rql: &str,
56 params: Params,
57 token: Option<&str>,
58 ) -> Result<Vec<Frame>, Error> {
59 let params_opt = match ¶ms {
60 Params::None => None,
61 _ => Some(params),
62 };
63
64 let client = self.get_or_connect(address, token)?;
65 match self.run_query(&client, rql, params_opt.clone()) {
66 Ok(frames) => Ok(frames),
67 Err(e) if is_transport_error(&e) => {
68 self.evict(address, token);
69 let client = self.get_or_connect(address, token)?;
70 self.run_query(&client, rql, params_opt)
71 }
72 Err(e) => Err(e),
73 }
74 }
75
76 fn run_query(&self, client: &GrpcClient, rql: &str, params: Option<Params>) -> Result<Vec<Frame>, Error> {
77 let client = client.clone();
78 let rql = rql.to_string();
79 let (tx, rx) = mpsc::sync_channel(1);
80
81 self.handle()?.spawn(async move {
82 let result = client.query(&rql, params).await;
83 let _ = tx.send(result);
84 });
85
86 rx.recv().map_err(|_| {
87 Error(Box::new(Diagnostic {
88 code: "REMOTE_002".to_string(),
89 message: "remote query channel closed".to_string(),
90 ..Default::default()
91 }))
92 })?
93 }
94
95 fn get_or_connect(&self, address: &str, token: Option<&str>) -> Result<GrpcClient, Error> {
96 let key = cache_key(address, token);
97 if let Some(c) = self.clients.lock().get(&key) {
98 return Ok(c.clone());
99 }
100 let client = self.connect(address, token)?;
101 self.clients.lock().entry(key).or_insert_with(|| client.clone());
102 Ok(client)
103 }
104
105 fn evict(&self, address: &str, token: Option<&str>) {
106 self.clients.lock().remove(&cache_key(address, token));
107 }
108
109 #[cfg(test)]
110 fn cache_len(&self) -> usize {
111 self.clients.lock().len()
112 }
113
114 fn connect(&self, address: &str, token: Option<&str>) -> Result<GrpcClient, Error> {
115 let address_owned = address.to_string();
116 let (tx, rx) = mpsc::sync_channel(1);
117
118 self.handle()?.spawn(async move {
119 let result = GrpcClient::connect(&address_owned, WireFormat::Proto).await;
120 let _ = tx.send(result);
121 });
122
123 let mut client = rx.recv().map_err(|_| {
124 Error(Box::new(Diagnostic {
125 code: "REMOTE_002".to_string(),
126 message: "remote connect channel closed".to_string(),
127 ..Default::default()
128 }))
129 })??;
130 if let Some(token) = token {
131 client.authenticate(token);
132 }
133 Ok(client)
134 }
135}
136
137#[cfg(not(reifydb_single_threaded))]
138fn cache_key(address: &str, token: Option<&str>) -> CacheKey {
139 (address.to_string(), token.map(str::to_string))
140}
141
142#[cfg(not(reifydb_single_threaded))]
143fn is_transport_error(err: &Error) -> bool {
144 err.0.code.starts_with("GRPC_")
145}
146
147pub fn is_remote_query(err: &Error) -> bool {
148 err.0.code == "REMOTE_001"
149}
150
151pub fn extract_remote_address(err: &Error) -> Option<String> {
152 err.0.notes.iter().find_map(|n| n.strip_prefix("Remote gRPC address: ")).map(|s| s.to_string())
153}
154
155pub fn extract_remote_token(err: &Error) -> Option<String> {
156 err.0.notes.iter().find_map(|n| n.strip_prefix("Remote token: ")).map(|s| s.to_string())
157}
158
159#[cfg(test)]
160mod tests {
161 use reifydb_runtime::{Runtime, RuntimeConfig, pool::PoolConfig};
162 use reifydb_value::{error::Diagnostic, fragment::Fragment};
163
164 use super::*;
165
166 fn make_remote_error(address: &str) -> Error {
167 Error(Box::new(Diagnostic {
168 code: "REMOTE_001".to_string(),
169 message: format!(
170 "Remote namespace 'remote_ns': source 'users' is on remote instance at {}",
171 address
172 ),
173 notes: vec![
174 "Namespace 'remote_ns' is configured as a remote namespace".to_string(),
175 format!("Remote gRPC address: {}", address),
176 ],
177 fragment: Fragment::None,
178 ..Default::default()
179 }))
180 }
181
182 #[test]
183 fn test_is_remote_query_true() {
184 let err = make_remote_error("http://localhost:50051");
185 assert!(is_remote_query(&err));
186 }
187
188 #[test]
189 fn test_is_remote_query_false() {
190 let err = Error(Box::new(Diagnostic {
191 code: "CATALOG_001".to_string(),
192 message: "Table not found".to_string(),
193 fragment: Fragment::None,
194 ..Default::default()
195 }));
196 assert!(!is_remote_query(&err));
197 }
198
199 #[test]
200 fn test_extract_remote_address() {
201 let err = make_remote_error("http://localhost:50051");
202 assert_eq!(extract_remote_address(&err), Some("http://localhost:50051".to_string()));
203 }
204
205 #[test]
206 fn test_extract_remote_address_missing() {
207 let err = Error(Box::new(Diagnostic {
208 code: "REMOTE_001".to_string(),
209 message: "Some error".to_string(),
210 notes: vec![],
211 fragment: Fragment::None,
212 ..Default::default()
213 }));
214 assert_eq!(extract_remote_address(&err), None);
215 }
216
217 #[test]
218 fn test_extract_remote_token() {
219 let err = Error(Box::new(Diagnostic {
220 code: "REMOTE_001".to_string(),
221 message: "Remote namespace".to_string(),
222 notes: vec![
223 "Namespace 'test' is configured as a remote namespace".to_string(),
224 "Remote gRPC address: http://localhost:50051".to_string(),
225 "Remote token: my-secret".to_string(),
226 ],
227 fragment: Fragment::None,
228 ..Default::default()
229 }));
230 assert_eq!(extract_remote_token(&err), Some("my-secret".to_string()));
231 }
232
233 #[test]
234 fn test_extract_remote_token_missing() {
235 let err = make_remote_error("http://localhost:50051");
236 assert_eq!(extract_remote_token(&err), None);
237 }
238
239 #[test]
240 fn test_is_transport_error() {
241 let grpc_err = Error(Box::new(Diagnostic {
242 code: "GRPC_Unavailable".to_string(),
243 message: "channel closed".to_string(),
244 ..Default::default()
245 }));
246 assert!(is_transport_error(&grpc_err));
247
248 let app_err = Error(Box::new(Diagnostic {
249 code: "CATALOG_001".to_string(),
250 message: "Table not found".to_string(),
251 ..Default::default()
252 }));
253 assert!(!is_transport_error(&app_err));
254 }
255
256 #[test]
257 fn test_cache_key_distinguishes_tokens() {
258 assert_ne!(cache_key("addr", Some("a")), cache_key("addr", Some("b")));
259 assert_ne!(cache_key("addr", None), cache_key("addr", Some("a")));
260 assert_eq!(cache_key("addr", Some("a")), cache_key("addr", Some("a")));
261 }
262
263 #[test]
264 fn test_connect_failure_does_not_pollute_cache() {
265 let runtime = Runtime::from_config(RuntimeConfig::default(), PoolConfig::default());
266 let registry = RemoteRegistry::new(runtime.tokio());
267
268 let err = registry.forward_query("http://127.0.0.1:1", "FROM x", Params::None, None).unwrap_err();
270 assert!(err.0.code.starts_with("GRPC_") || err.0.code == "REMOTE_002");
271 assert_eq!(registry.cache_len(), 0);
272 }
273
274 #[test]
275 fn test_evict_missing_key_is_noop() {
276 let runtime = Runtime::from_config(RuntimeConfig::default(), PoolConfig::default());
277 let registry = RemoteRegistry::new(runtime.tokio());
278 registry.evict("http://127.0.0.1:1", None);
279 registry.evict("http://127.0.0.1:1", Some("tok"));
280 assert_eq!(registry.cache_len(), 0);
281 }
282}