klukai_client/
lib.rs

1pub mod sub;
2
3use bytes::Bytes;
4use hickory_resolver::{ResolveError, Resolver, name_server::TokioConnectionProvider};
5use http_body_util::{BodyExt, Full};
6use hyper::http::uri::PathAndQuery;
7use hyper::{StatusCode, http::HeaderName};
8use hyper_util::client::legacy::{Client, connect::HttpConnector};
9use klukai_types::api::{ChangeId, ExecResponse, ExecResult, SqliteValue, Statement};
10use serde::de::DeserializeOwned;
11use std::{
12    net::SocketAddr,
13    ops::Deref,
14    path::Path,
15    sync::Arc,
16    time::{self, Duration, Instant},
17};
18use sub::{QueryStream, SubscriptionStream, UpdatesStream};
19use tokio::{
20    sync::{RwLock, RwLockReadGuard},
21    time::timeout,
22};
23use tracing::{debug, info};
24use uuid::Uuid;
25
26pub mod util;
27
28const HTTP2_CONNECT_TIMEOUT: Duration = Duration::from_secs(3);
29const HTTP2_KEEP_ALIVE_INTERVAL: Duration = Duration::from_secs(10);
30const DNS_RESOLVE_TIMEOUT: Duration = Duration::from_secs(3);
31
32#[derive(Clone)]
33pub struct CorrosionApiClient {
34    api_addr: SocketAddr,
35    api_client: Client<HttpConnector, Full<Bytes>>,
36}
37
38impl CorrosionApiClient {
39    pub fn new(api_addr: SocketAddr) -> Self {
40        let mut connector = HttpConnector::new();
41        connector.set_connect_timeout(Some(HTTP2_CONNECT_TIMEOUT));
42        Self {
43            api_addr,
44            api_client: Client::builder(hyper_util::rt::TokioExecutor::new())
45                .timer(hyper_util::rt::TokioTimer::new())
46                .http2_only(true)
47                .http2_keep_alive_interval(Some(HTTP2_KEEP_ALIVE_INTERVAL))
48                .http2_keep_alive_timeout(HTTP2_KEEP_ALIVE_INTERVAL / 2)
49                .build(connector),
50        }
51    }
52
53    pub async fn query_typed<T: DeserializeOwned + Unpin>(
54        &self,
55        statement: &Statement,
56        timeout: Option<u64>,
57    ) -> Result<QueryStream<T>, Error> {
58        let params = timeout
59            .map(|t| format!("?timeout={}", t))
60            .unwrap_or_default();
61        let req = hyper::Request::builder()
62            .method(hyper::Method::POST)
63            .uri(format!("http://{}/v1/queries{}", self.api_addr, params))
64            .header(hyper::header::CONTENT_TYPE, "application/json")
65            .header(hyper::header::ACCEPT, "application/json")
66            .body(Full::new(Bytes::from(serde_json::to_vec(statement)?)))?;
67
68        let res = self.api_client.request(req).await?;
69
70        if !res.status().is_success() {
71            let status = res.status();
72            match res.into_body().collect().await {
73                Ok(collected) => match serde_json::from_slice(&collected.to_bytes()) {
74                    Ok(res) => match res {
75                        ExecResult::Error { error } => return Err(Error::ResponseError(error)),
76                        res => return Err(Error::UnexpectedResult(res)),
77                    },
78                    Err(e) => {
79                        debug!(
80                            error = %e,
81                            "could not deserialize response body, sending generic error..."
82                        );
83                        return Err(Error::UnexpectedStatusCode(status));
84                    }
85                },
86                Err(e) => {
87                    debug!(
88                        error = %e,
89                        "could not aggregate response body bytes, sending generic error..."
90                    );
91                    return Err(Error::UnexpectedStatusCode(status));
92                }
93            }
94        }
95
96        Ok(QueryStream::new(res.into_body()))
97    }
98
99    pub async fn query(
100        &self,
101        statement: &Statement,
102        timeout: Option<u64>,
103    ) -> Result<QueryStream<Vec<SqliteValue>>, Error> {
104        self.query_typed(statement, timeout).await
105    }
106
107    pub async fn subscribe_typed<T: DeserializeOwned + Unpin>(
108        &self,
109        statement: &Statement,
110        skip_rows: bool,
111        from: Option<ChangeId>,
112    ) -> Result<SubscriptionStream<T>, Error> {
113        let p_and_q: PathAndQuery = if let Some(change_id) = from {
114            format!(
115                "/v1/subscriptions?skip_rows={}&from={}",
116                skip_rows, change_id.0
117            )
118            .try_into()?
119        } else {
120            format!("/v1/subscriptions?skip_rows={}", skip_rows).try_into()?
121        };
122        let url = hyper::Uri::builder()
123            .scheme("http")
124            .authority(self.api_addr.to_string())
125            .path_and_query(p_and_q)
126            .build()?;
127
128        let req = hyper::Request::builder()
129            .method(hyper::Method::POST)
130            .uri(url)
131            .header(hyper::header::CONTENT_TYPE, "application/json")
132            .header(hyper::header::ACCEPT, "application/json")
133            .body(Full::new(Bytes::from(serde_json::to_vec(statement)?)))?;
134
135        let res = self.api_client.request(req).await?;
136
137        if !res.status().is_success() {
138            return Err(Error::UnexpectedStatusCode(res.status()));
139        }
140
141        // TODO: make that header name a const in corro-types
142        let id = res
143            .headers()
144            .get(HeaderName::from_static("corro-query-id"))
145            .and_then(|v| v.to_str().ok().and_then(|v| v.parse().ok()))
146            .ok_or(Error::ExpectedQueryId)?;
147        let hash = res
148            .headers()
149            .get(HeaderName::from_static("corro-query-hash"))
150            .and_then(|v| v.to_str().map(ToOwned::to_owned).ok());
151
152        Ok(SubscriptionStream::new(
153            id,
154            hash,
155            self.api_client.clone(),
156            self.api_addr,
157            res.into_body(),
158            from,
159        ))
160    }
161
162    pub async fn subscribe(
163        &self,
164        statement: &Statement,
165        skip_rows: bool,
166        from: Option<ChangeId>,
167    ) -> Result<SubscriptionStream<Vec<SqliteValue>>, Error> {
168        self.subscribe_typed(statement, skip_rows, from).await
169    }
170
171    pub async fn subscription_typed<T: DeserializeOwned + Unpin>(
172        &self,
173        id: Uuid,
174        skip_rows: bool,
175        from: Option<ChangeId>,
176    ) -> Result<SubscriptionStream<T>, Error> {
177        let p_and_q: PathAndQuery = if let Some(change_id) = from {
178            format!(
179                "/v1/subscriptions/{id}?skip_rows={}&from={}",
180                skip_rows, change_id.0
181            )
182            .try_into()?
183        } else {
184            format!("/v1/subscriptions/{id}?skip_rows={}", skip_rows).try_into()?
185        };
186        let url = hyper::Uri::builder()
187            .scheme("http")
188            .authority(self.api_addr.to_string())
189            .path_and_query(p_and_q)
190            .build()?;
191
192        let req = hyper::Request::builder()
193            .method(hyper::Method::GET)
194            .uri(url)
195            .header(hyper::header::ACCEPT, "application/json")
196            .body(Full::new(Bytes::new()))?;
197
198        let res = self.api_client.request(req).await?;
199
200        if !res.status().is_success() {
201            return Err(Error::UnexpectedStatusCode(res.status()));
202        }
203
204        let hash = res
205            .headers()
206            .get(HeaderName::from_static("corro-query-hash"))
207            .and_then(|v| v.to_str().map(ToOwned::to_owned).ok());
208
209        Ok(SubscriptionStream::new(
210            id,
211            hash,
212            self.api_client.clone(),
213            self.api_addr,
214            res.into_body(),
215            from,
216        ))
217    }
218
219    pub async fn subscription(
220        &self,
221        id: Uuid,
222        skip_rows: bool,
223        from: Option<ChangeId>,
224    ) -> Result<SubscriptionStream<Vec<SqliteValue>>, Error> {
225        self.subscription_typed(id, skip_rows, from).await
226    }
227
228    pub async fn updates_typed<T: DeserializeOwned + Unpin>(
229        &self,
230        table: &str,
231    ) -> Result<UpdatesStream<T>, Error> {
232        let p_and_q: PathAndQuery = format!("/v1/updates/{}", table).try_into()?;
233
234        let url = hyper::Uri::builder()
235            .scheme("http")
236            .authority(self.api_addr.to_string())
237            .path_and_query(p_and_q)
238            .build()?;
239
240        let req = hyper::Request::builder()
241            .method(hyper::Method::POST)
242            .uri(url)
243            .header(hyper::header::CONTENT_TYPE, "application/json")
244            .header(hyper::header::ACCEPT, "application/json")
245            .body(Full::new(Bytes::new()))?;
246
247        let res = self.api_client.request(req).await?;
248
249        if !res.status().is_success() {
250            return Err(Error::UnexpectedStatusCode(res.status()));
251        }
252
253        // TODO: make that header name a const in corro-types
254        let id = res
255            .headers()
256            .get(HeaderName::from_static("corro-query-id"))
257            .and_then(|v| v.to_str().ok().and_then(|v| v.parse().ok()))
258            .ok_or(Error::ExpectedQueryId)?;
259
260        Ok(UpdatesStream::new(id, res.into_body()))
261    }
262
263    pub async fn updates(&self, table: &str) -> Result<UpdatesStream<Vec<SqliteValue>>, Error> {
264        self.updates_typed(table).await
265    }
266
267    pub async fn execute(
268        &self,
269        statements: &[Statement],
270        timeout: Option<u64>,
271    ) -> Result<ExecResponse, Error> {
272        let uri = if let Some(timeout) = timeout {
273            format!(
274                "http://{}/v1/transactions?timeout={}",
275                self.api_addr, timeout
276            )
277        } else {
278            format!("http://{}/v1/transactions", self.api_addr)
279        };
280        // println!("uri: {:?}", uri);
281        let req = hyper::Request::builder()
282            .method(hyper::Method::POST)
283            .uri(uri)
284            .header(hyper::header::CONTENT_TYPE, "application/json")
285            .header(hyper::header::ACCEPT, "application/json")
286            .body(Full::new(Bytes::from(serde_json::to_vec(statements)?)))?;
287
288        let res = self.api_client.request(req).await?;
289
290        let status = res.status();
291        if !status.is_success() {
292            match res.into_body().collect().await {
293                Ok(collected) => match serde_json::from_slice(&collected.to_bytes()) {
294                    Ok(ExecResponse { results, .. }) => {
295                        if let Some(ExecResult::Error { error }) = results
296                            .into_iter()
297                            .find(|r| matches!(r, ExecResult::Error { .. }))
298                        {
299                            return Err(Error::ResponseError(error));
300                        }
301                        return Err(Error::UnexpectedStatusCode(status));
302                    }
303                    Err(e) => {
304                        debug!(
305                            error = %e,
306                            "could not deserialize response body, sending generic error..."
307                        );
308                        return Err(Error::UnexpectedStatusCode(status));
309                    }
310                },
311                Err(e) => {
312                    debug!(
313                        error = %e,
314                        "could not aggregate response body bytes, sending generic error..."
315                    );
316                    return Err(Error::UnexpectedStatusCode(status));
317                }
318            }
319        }
320        let collected = res.into_body().collect().await?;
321        let bytes = collected.to_bytes();
322        Ok(serde_json::from_slice(&bytes)?)
323    }
324
325    pub async fn schema(&self, statements: &[Statement]) -> Result<ExecResponse, Error> {
326        let req = hyper::Request::builder()
327            .method(hyper::Method::POST)
328            .uri(format!("http://{}/v1/migrations", self.api_addr))
329            .header(hyper::header::CONTENT_TYPE, "application/json")
330            .header(hyper::header::ACCEPT, "application/json")
331            .body(Full::new(Bytes::from(serde_json::to_vec(statements)?)))?;
332
333        let res = self.api_client.request(req).await?;
334
335        if !res.status().is_success() {
336            return Err(Error::UnexpectedStatusCode(res.status()));
337        }
338
339        let collected = res.into_body().collect().await?;
340        let bytes = collected.to_bytes();
341
342        Ok(serde_json::from_slice(&bytes)?)
343    }
344
345    pub async fn schema_from_paths<P: AsRef<Path>>(
346        &self,
347        schema_paths: &[P],
348    ) -> Result<Option<ExecResponse>, Error> {
349        let statements: Vec<Statement> = crate::util::read_files_from_paths(schema_paths)
350            .await
351            .map_err(|e| Error::ResponseError(e.to_string()))?
352            .into_iter()
353            .map(Statement::Simple)
354            .collect();
355
356        if statements.is_empty() {
357            return Ok(None);
358        }
359
360        Ok(Some(self.schema(&statements).await?))
361    }
362}
363
364#[derive(Clone)]
365pub struct CorrosionClient {
366    api_client: CorrosionApiClient,
367    pool: klukai_types::sqlite_pool::RusqlitePool,
368}
369
370impl CorrosionClient {
371    pub fn new<P: AsRef<Path>>(api_addr: SocketAddr, db_path: P) -> Self {
372        Self {
373            api_client: CorrosionApiClient::new(api_addr),
374            pool: klukai_types::sqlite_pool::Config::new(db_path.as_ref())
375                .max_size(5)
376                .create_pool()
377                .expect("could not build pool, this can't fail because we specified a runtime"),
378        }
379    }
380
381    pub fn with_sqlite_pool(
382        api_addr: SocketAddr,
383        pool: klukai_types::sqlite_pool::RusqlitePool,
384    ) -> Self {
385        Self {
386            api_client: CorrosionApiClient::new(api_addr),
387            pool,
388        }
389    }
390
391    pub fn pool(&self) -> &klukai_types::sqlite_pool::RusqlitePool {
392        &self.pool
393    }
394}
395
396impl Deref for CorrosionClient {
397    type Target = CorrosionApiClient;
398
399    fn deref(&self) -> &Self::Target {
400        &self.api_client
401    }
402}
403
404#[derive(Clone)]
405pub struct CorrosionPooledClient {
406    inner: Arc<RwLock<PooledClientInner>>,
407}
408
409struct PooledClientInner {
410    picker: AddrPicker,
411
412    // For how long to stick with a chosen server
413    stickiness_timeout: time::Duration,
414    // Currently chosen client
415    client: Option<CorrosionApiClient>,
416    // Whether or not the chosen client has made a successful request
417    had_success: bool,
418    // Time when the first fail occured after at least one successful call
419    first_fail_at: Option<Instant>,
420    // Current client generation, incremented after each client change
421    generation: u64,
422}
423
424impl CorrosionPooledClient {
425    pub fn new(
426        addrs: Vec<String>,
427        stickiness_timeout: time::Duration,
428        resolver: Resolver<TokioConnectionProvider>,
429    ) -> Self {
430        Self {
431            inner: Arc::new(RwLock::new(PooledClientInner {
432                picker: AddrPicker::new(addrs, resolver),
433
434                stickiness_timeout,
435                client: None,
436                had_success: false,
437                first_fail_at: None,
438                generation: 0,
439            })),
440        }
441    }
442
443    pub async fn query_typed<T: DeserializeOwned + Unpin>(
444        &self,
445        statement: &Statement,
446        timeout: Option<u64>,
447    ) -> Result<QueryStream<T>, Error> {
448        let (response, generation) = {
449            let (client, generation) = self.get_client().await?;
450            let response = client.query_typed(statement, timeout).await;
451
452            (response, generation)
453        };
454
455        if matches!(response, Err(Error::Hyper(_)) | Err(Error::HyperUtil(_))) {
456            // We only care about I/O related errors
457            self.handle_error(generation).await;
458        } else {
459            // The rest are considered a success
460            self.handle_success(generation).await;
461        }
462
463        response
464    }
465
466    pub async fn subscribe_typed<T: DeserializeOwned + Unpin>(
467        &self,
468        statement: &Statement,
469        skip_rows: bool,
470        from: Option<ChangeId>,
471    ) -> Result<SubscriptionStream<T>, Error> {
472        let (response, generation) = {
473            let (client, generation) = self.get_client().await?;
474            let response = client.subscribe_typed(statement, skip_rows, from).await;
475
476            (response, generation)
477        };
478
479        if matches!(response, Err(Error::Hyper(_)) | Err(Error::HyperUtil(_))) {
480            // We only care about I/O related errors
481            self.handle_error(generation).await;
482        } else {
483            // The rest are considered a success
484            self.handle_success(generation).await;
485        }
486
487        response
488    }
489
490    pub async fn subscription_typed<T: DeserializeOwned + Unpin>(
491        &self,
492        id: Uuid,
493        skip_rows: bool,
494        from: Option<ChangeId>,
495    ) -> Result<SubscriptionStream<T>, Error> {
496        let (response, generation) = {
497            let (client, generation) = self.get_client().await?;
498            let response = client.subscription_typed(id, skip_rows, from).await;
499
500            (response, generation)
501        };
502
503        if matches!(response, Err(Error::Hyper(_)) | Err(Error::HyperUtil(_))) {
504            // We only care about I/O related errors
505            self.handle_error(generation).await;
506        } else {
507            // The rest are considered a success
508            self.handle_success(generation).await;
509        }
510
511        response
512    }
513
514    async fn get_client(&self) -> Result<(RwLockReadGuard<'_, CorrosionApiClient>, u64), Error> {
515        let mut inner = self.inner.write().await;
516        let generation = inner.generation;
517
518        if inner.client.is_none() {
519            let addr = inner.picker.next().await?;
520            info!(
521                "next Corrosion server to attempt: {}, generation: {}",
522                addr, generation
523            );
524            inner.client = Some(CorrosionApiClient::new(addr))
525        }
526
527        Ok((
528            RwLockReadGuard::map(inner.downgrade(), |inner| inner.client.as_ref().unwrap()),
529            generation,
530        ))
531    }
532
533    async fn handle_success(&self, generation: u64) {
534        let mut inner = self.inner.write().await;
535
536        // Even though the call was successul, another failed call has advanced the client already.
537        if inner.generation != generation {
538            return;
539        }
540
541        // Mark that this client was able to perform a successful call to Corrosion
542        // and we should stick with it.
543        inner.had_success = true;
544        // And reset the time of the first fail.
545        inner.first_fail_at = None;
546    }
547
548    async fn handle_error(&self, generation: u64) {
549        let mut inner = self.inner.write().await;
550
551        // Somebody else has already handled an error with this client
552        if generation != inner.generation {
553            return;
554        }
555
556        match inner.first_fail_at {
557            // First fail after success
558            None if inner.had_success => {
559                inner.first_fail_at = Some(Instant::now());
560            }
561
562            // Still within stickiness timeout, try the same server again
563            Some(first) if Instant::now().duration_since(first) < inner.stickiness_timeout => {}
564
565            // Otherwise, pick a new server for the next attempt
566            _ => {
567                // If we had a successful call before, try to fallback to the first server, so the first
568                // one is always preferred by all the clients.
569                // Otherwise, continue iterating over the rest of the servers.
570                if inner.had_success {
571                    inner.picker.reset()
572                }
573
574                inner.client = None;
575                inner.first_fail_at = None;
576                inner.had_success = false;
577                inner.generation += 1;
578            }
579        }
580    }
581}
582
583struct AddrPicker {
584    // Resolver used to resolve the addresses
585    resolver: Resolver<TokioConnectionProvider>,
586    // List of addresses/hostname to try in order
587    addrs: Vec<String>,
588    // Next address/hostname to try
589    next_addr: usize,
590
591    // List of addresses returned by the last resolve attempt
592    last_resolved_addrs: Option<Vec<SocketAddr>>,
593    // Next address to return
594    next_resolved_addr: usize,
595}
596
597impl AddrPicker {
598    fn new(addrs: Vec<String>, resolver: Resolver<TokioConnectionProvider>) -> AddrPicker {
599        Self {
600            resolver,
601            addrs,
602            next_addr: 0,
603
604            last_resolved_addrs: None,
605            next_resolved_addr: 0,
606        }
607    }
608
609    async fn next(&mut self) -> Result<SocketAddr, Error> {
610        // Either we don't have any address or we tried them all, resolve again.
611        if self.next_resolved_addr
612            >= self
613                .last_resolved_addrs
614                .as_ref()
615                .map(|v| v.len())
616                .unwrap_or_default()
617        {
618            let host_port = self
619                .addrs
620                .get(self.next_addr)
621                .ok_or(ResolveError::from("No addresses available"))?;
622            self.next_addr = (self.next_addr + 1) % self.addrs.len();
623
624            let mut addrs = if let Ok(addr) = host_port.parse() {
625                vec![addr]
626            } else {
627                // split host port
628                let (host, port) = host_port
629                    .rsplit_once(':')
630                    .and_then(|(host, port)| Some((host, port.parse().ok()?)))
631                    .ok_or(ResolveError::from("Invalid Corrosion server address"))?;
632
633                timeout(DNS_RESOLVE_TIMEOUT, self.resolver.lookup_ip(host))
634                    .await
635                    .map_err(|_| ResolveError::from("DNS lookup timeout"))??
636                    .iter()
637                    .map(|addr| (addr, port).into())
638                    .collect::<Vec<_>>()
639            };
640            // Sort so all the nodes try the addresses in the same order
641            addrs.sort();
642
643            debug!("got the following Corrosion servers: {:?}", addrs);
644
645            self.last_resolved_addrs = Some(addrs);
646            self.next_resolved_addr = 0;
647        }
648
649        if let Some(addr) = self
650            .last_resolved_addrs
651            .as_ref()
652            .and_then(|a| a.get(self.next_resolved_addr).copied())
653        {
654            self.next_resolved_addr += 1;
655
656            Ok(addr)
657        } else {
658            Err(ResolveError::from("DNS didn't return any addresses").into())
659        }
660    }
661
662    fn reset(&mut self) {
663        self.next_addr = 0;
664        self.last_resolved_addrs = None;
665        self.next_resolved_addr = 0;
666    }
667}
668
669#[derive(Debug, thiserror::Error)]
670pub enum Error {
671    #[error(transparent)]
672    Dns(#[from] ResolveError),
673    #[error(transparent)]
674    Hyper(#[from] hyper::Error),
675    #[error(transparent)]
676    HyperUtil(#[from] hyper_util::client::legacy::Error),
677    #[error(transparent)]
678    InvalidUri(#[from] hyper::http::uri::InvalidUri),
679    #[error(transparent)]
680    Http(#[from] hyper::http::Error),
681    #[error(transparent)]
682    Serde(#[from] serde_json::Error),
683
684    #[error("received unexpected response code: {0}")]
685    UnexpectedStatusCode(StatusCode),
686
687    #[error("{0}")]
688    ResponseError(String),
689
690    #[error("unexpected result: {0:?}")]
691    UnexpectedResult(ExecResult),
692
693    #[error("could not retrieve subscription id from headers")]
694    ExpectedQueryId,
695}
696
697#[cfg(test)]
698mod tests {
699    use crate::{CorrosionPooledClient, Error};
700    use bytes::Bytes;
701    use hickory_resolver::Resolver;
702    use hyper::body::Incoming;
703    use hyper::{Request, Response, header::HeaderValue, service::service_fn};
704    use hyper_util::rt::TokioIo;
705    use klukai_types::api::SqliteValue;
706    use std::{
707        convert::Infallible,
708        net::SocketAddr,
709        sync::{
710            Arc,
711            atomic::{AtomicBool, Ordering},
712        },
713        time::Duration,
714    };
715    use tokio::{net::TcpListener, pin, sync::broadcast};
716    use uuid::Uuid;
717
718    struct Server {
719        id: Uuid,
720        addr: SocketAddr,
721        refuse: Arc<AtomicBool>,
722        drop_conn_tx: broadcast::Sender<()>,
723    }
724
725    impl Server {
726        async fn new(id: Uuid) -> Self {
727            let refuse = Arc::new(AtomicBool::new(false));
728            let (drop_conn_tx, drop_conn_rx) = broadcast::channel::<()>(1);
729            let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
730            let addr = listener.local_addr().unwrap();
731
732            tokio::spawn({
733                let refuse = refuse.clone();
734
735                async move {
736                    loop {
737                        let (stream, _) = listener.accept().await.unwrap();
738                        if refuse.load(Ordering::Relaxed) {
739                            drop(stream);
740                            continue;
741                        }
742
743                        let mut drop_conn_rx = drop_conn_rx.resubscribe();
744                        tokio::spawn(async move {
745                            let http = hyper::server::conn::http2::Builder::new(
746                                hyper_util::rt::TokioExecutor::new(),
747                            );
748                            let conn = http.serve_connection(
749                                TokioIo::new(stream),
750                                service_fn(move |_: Request<Incoming>| async move {
751                                    let mut res =
752                                        Response::new(http_body_util::Empty::<Bytes>::new());
753                                    res.headers_mut().insert(
754                                        "corro-query-id",
755                                        HeaderValue::from_str(&id.to_string()).unwrap(),
756                                    );
757                                    Ok::<_, Infallible>(res)
758                                }),
759                            );
760                            pin!(conn);
761
762                            tokio::select! {
763                                _ = conn.as_mut() => (),
764                                _ = drop_conn_rx.recv() => (),
765                            }
766                        });
767                    }
768                }
769            });
770
771            Server {
772                id,
773                addr,
774                refuse,
775                drop_conn_tx,
776            }
777        }
778
779        fn refuse_new_conns(&self, refuse: bool) {
780            self.refuse.store(refuse, Ordering::Relaxed)
781        }
782
783        fn kill_existing_conns(&self) {
784            _ = self.drop_conn_tx.send(())
785        }
786    }
787
788    async fn gen_servers(num: usize) -> (Vec<Server>, Vec<String>) {
789        let mut servers = Vec::new();
790
791        for _ in 0..num {
792            servers.push(Server::new(Uuid::new_v4()).await);
793        }
794
795        // sort the way the client is supposed to try them
796        servers.sort_by(|a, b| a.addr.partial_cmp(&b.addr).unwrap());
797        let addrs = servers.iter().map(|s| s.addr.to_string()).collect();
798
799        (servers, addrs)
800    }
801
802    #[tokio::test]
803    async fn test_single_address() {
804        let statement = "".into();
805        let (servers, addresses) = gen_servers(1).await;
806
807        let resolver = Resolver::builder_tokio().unwrap().build();
808        let client = CorrosionPooledClient::new(addresses, Duration::from_nanos(1), resolver);
809
810        let sub = client
811            .subscribe_typed::<SqliteValue>(&statement, false, None)
812            .await
813            .unwrap();
814        assert_eq!(sub.id(), servers[0].id);
815
816        // Drop the connection, next attempt should error
817        servers[0].kill_existing_conns();
818
819        let res = client
820            .subscribe_typed::<SqliteValue>(&statement, false, None)
821            .await;
822        assert!(matches!(
823            res,
824            Result::Err(Error::Hyper(_)) | Result::Err(Error::HyperUtil(_))
825        ));
826
827        // But the new one should succeed
828        let sub = client
829            .subscribe_typed::<SqliteValue>(&statement, false, None)
830            .await
831            .unwrap();
832        assert_eq!(sub.id(), servers[0].id);
833    }
834
835    #[tokio::test]
836    async fn test_multiple_addresses() {
837        let statement = "".into();
838        let (servers, addresses) = gen_servers(3).await;
839
840        let resolver = Resolver::builder_tokio().unwrap().build();
841        let client = CorrosionPooledClient::new(addresses, Duration::from_nanos(1), resolver);
842
843        // Refuse connections on the first server
844        servers[0].refuse_new_conns(true);
845
846        // First one should error
847        let res = client
848            .subscribe_typed::<SqliteValue>(&statement, false, None)
849            .await;
850        assert!(matches!(
851            res,
852            Result::Err(Error::Hyper(_)) | Result::Err(Error::HyperUtil(_))
853        ));
854
855        // Second one should succeed
856        let sub = client
857            .subscribe_typed::<SqliteValue>(&statement, false, None)
858            .await
859            .unwrap();
860        assert_eq!(sub.id(), servers[1].id);
861
862        // Abort the second server, the client should fallback back to the first one after the first two attempts
863        servers[1].kill_existing_conns();
864        servers[1].refuse_new_conns(true);
865        servers[0].refuse_new_conns(false);
866
867        // First and second one should error
868        for _ in 0..2 {
869            let res = client
870                .subscribe_typed::<SqliteValue>(&statement, false, None)
871                .await;
872            assert!(matches!(
873                res,
874                Result::Err(Error::Hyper(_)) | Result::Err(Error::HyperUtil(_))
875            ));
876        }
877
878        // The next one should succeed
879        let sub = client
880            .subscribe_typed::<SqliteValue>(&statement, false, None)
881            .await
882            .unwrap();
883        assert_eq!(sub.id(), servers[0].id);
884    }
885
886    #[tokio::test]
887    async fn test_multiple_addresses_sticky() {
888        let statement = "".into();
889        let (servers, addresses) = gen_servers(3).await;
890
891        let resolver = Resolver::builder_tokio().unwrap().build();
892        let client = CorrosionPooledClient::new(addresses, Duration::from_millis(50), resolver);
893
894        // Refuse connections on the first server
895        servers[0].refuse_new_conns(true);
896
897        // First one should error
898        let res = client
899            .subscribe_typed::<SqliteValue>(&statement, false, None)
900            .await;
901        assert!(matches!(
902            res,
903            Result::Err(Error::Hyper(_)) | Result::Err(Error::HyperUtil(_))
904        ));
905
906        // Second one should succeed
907        let sub = client
908            .subscribe_typed::<SqliteValue>(&statement, false, None)
909            .await
910            .unwrap();
911        assert_eq!(sub.id(), servers[1].id);
912
913        // Abort the second server, the client should continue trying it until the timeout expires
914        servers[1].kill_existing_conns();
915        servers[1].refuse_new_conns(true);
916        servers[0].refuse_new_conns(false);
917
918        let mut attempts = 0;
919        loop {
920            let res = client
921                .subscribe_typed::<SqliteValue>(&statement, false, None)
922                .await;
923
924            match res {
925                Ok(sub) => {
926                    assert_eq!(sub.id(), servers[0].id);
927                    break;
928                }
929                Err(_) => attempts += 1,
930            };
931        }
932        assert!(attempts > 2);
933    }
934
935    #[tokio::test]
936    async fn test_more_servers() {
937        let statement = "".into();
938        let (pool1_servers, pool1_addresses) = gen_servers(2).await;
939        let (pool2_servers, pool2_addresses) = gen_servers(2).await;
940
941        let mut addresses = pool1_addresses;
942        addresses.extend_from_slice(&pool2_addresses);
943
944        let resolver = Resolver::builder_tokio().unwrap().build();
945        let client = CorrosionPooledClient::new(addresses, Duration::from_nanos(1), resolver);
946
947        // Refuse connections on all servers
948        for i in 0..2 {
949            pool1_servers[i].refuse_new_conns(true);
950            pool2_servers[i].refuse_new_conns(true);
951        }
952
953        // Try to connect multiple times, all should fail
954        for _ in 0..15 {
955            let res = client
956                .subscribe_typed::<SqliteValue>(&statement, false, None)
957                .await;
958            assert!(matches!(
959                res,
960                Result::Err(Error::Hyper(_)) | Result::Err(Error::HyperUtil(_))
961            ));
962        }
963
964        // Accept connections on first server in the backup pool
965        pool2_servers[0].refuse_new_conns(false);
966        let mut succeeded = false;
967        for _i in 0..10 {
968            let res = client
969                .subscribe_typed::<SqliteValue>(&statement, false, None)
970                .await;
971            match res {
972                Result::Err(_) => (),
973                Ok(sub) => {
974                    assert_eq!(sub.id(), pool2_servers[0].id);
975                    succeeded = true;
976                    break;
977                }
978            }
979        }
980        assert!(
981            succeeded,
982            "Failed to connect to backup server after multiple attempts"
983        );
984
985        // Kill the connection, it should fallback to the first pool
986        pool2_servers[0].kill_existing_conns();
987        pool2_servers[0].refuse_new_conns(true);
988        pool1_servers[0].refuse_new_conns(false);
989        pool1_servers[1].refuse_new_conns(false);
990
991        // First and second one should error
992        for _ in 0..2 {
993            let res = client
994                .subscribe_typed::<SqliteValue>(&statement, false, None)
995                .await;
996            assert!(matches!(
997                res,
998                Result::Err(Error::Hyper(_)) | Result::Err(Error::HyperUtil(_))
999            ));
1000        }
1001
1002        // Thirst one should succeed
1003        let sub = client
1004            .subscribe_typed::<SqliteValue>(&statement, false, None)
1005            .await
1006            .unwrap();
1007        assert_eq!(sub.id(), pool1_servers[0].id);
1008    }
1009}
klukai_client/lib.rs

klukai_client/
lib.rs
