klukai_client/
lib.rs

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

klukai_client/
lib.rs
