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_003".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_003".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::Rbcf).await;
120 let _ = tx.send(result);
121 });
122
123 let mut client = rx.recv().map_err(|_| {
124 Error(Box::new(Diagnostic {
125 code: "REMOTE_003".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 matches!(err.0.code.as_str(), "CONNECTION_LOST" | "TRANSPORT")
145}
146
147pub fn is_remote_query(err: &Error) -> bool {
148 err.0.code == "REMOTE_001"
149}
150
151#[cfg(test)]
152mod tests {
153 use reifydb_runtime::{Runtime, RuntimeConfig, pool::PoolConfig};
154 use reifydb_value::{error::Diagnostic, fragment::Fragment};
155
156 use super::*;
157
158 fn make_remote_error(address: &str) -> Error {
159 Error(Box::new(Diagnostic {
160 code: "REMOTE_001".to_string(),
161 message: format!(
162 "Remote namespace 'remote_ns': source 'users' is on remote instance at {}",
163 address
164 ),
165 notes: vec![
166 "Namespace 'remote_ns' is configured as a remote namespace".to_string(),
167 format!("Remote gRPC address: {}", address),
168 ],
169 fragment: Fragment::None,
170 ..Default::default()
171 }))
172 }
173
174 #[test]
175 fn test_is_remote_query_true() {
176 let err = make_remote_error("http://localhost:50051");
177 assert!(is_remote_query(&err));
178 }
179
180 #[test]
181 fn test_is_remote_query_false() {
182 let err = Error(Box::new(Diagnostic {
183 code: "CATALOG_001".to_string(),
184 message: "Table not found".to_string(),
185 fragment: Fragment::None,
186 ..Default::default()
187 }));
188 assert!(!is_remote_query(&err));
189 }
190
191 #[test]
192 fn test_is_transport_error() {
193 let transport_err = Error(Box::new(Diagnostic {
194 code: "CONNECTION_LOST".to_string(),
195 message: "channel closed".to_string(),
196 ..Default::default()
197 }));
198 assert!(is_transport_error(&transport_err));
199
200 let app_err = Error(Box::new(Diagnostic {
201 code: "CATALOG_001".to_string(),
202 message: "Table not found".to_string(),
203 ..Default::default()
204 }));
205 assert!(!is_transport_error(&app_err));
206 }
207
208 #[test]
209 fn test_cache_key_distinguishes_tokens() {
210 assert_ne!(cache_key("addr", Some("a")), cache_key("addr", Some("b")));
211 assert_ne!(cache_key("addr", None), cache_key("addr", Some("a")));
212 assert_eq!(cache_key("addr", Some("a")), cache_key("addr", Some("a")));
213 }
214
215 #[test]
216 fn test_connect_failure_does_not_pollute_cache() {
217 let runtime = Runtime::from_config(RuntimeConfig::default(), PoolConfig::default());
218 let registry = RemoteRegistry::new(runtime.tokio());
219
220 let err = registry.forward_query("http://127.0.0.1:1", "FROM x", Params::None, None).unwrap_err();
222 assert!(err.0.code == "TRANSPORT" || err.0.code == "REMOTE_003");
223 assert_eq!(registry.cache_len(), 0);
224 }
225
226 #[test]
227 fn test_evict_missing_key_is_noop() {
228 let runtime = Runtime::from_config(RuntimeConfig::default(), PoolConfig::default());
229 let registry = RemoteRegistry::new(runtime.tokio());
230 registry.evict("http://127.0.0.1:1", None);
231 registry.evict("http://127.0.0.1:1", Some("tok"));
232 assert_eq!(registry.cache_len(), 0);
233 }
234}