use std::time::Duration; use futures::FutureExt; use redis::aio::ConnectionLike; use redis::{Cmd, FromRedisValue, Pipeline, RedisResult}; use super::connection_with_credentials_provider::ConnectionWithCredentialsProvider; pub struct RedisKVClient { client: ConnectionWithCredentialsProvider, } #[allow(async_fn_in_trait)] pub trait Queryable { async fn query(&self, conn: &mut impl ConnectionLike) -> RedisResult; } impl Queryable for Pipeline { async fn query(&self, conn: &mut impl ConnectionLike) -> RedisResult { self.query_async(conn).await } } impl Queryable for Cmd { async fn query(&self, conn: &mut impl ConnectionLike) -> RedisResult { self.query_async(conn).await } } impl RedisKVClient { pub fn new(client: ConnectionWithCredentialsProvider) -> Self { Self { client } } pub async fn try_connect(&mut self) -> anyhow::Result<()> { self.client .connect() .boxed() .await .inspect_err(|e| tracing::error!("failed to connect to redis: {e}")) } pub(crate) async fn query( &mut self, q: &impl Queryable, ) -> anyhow::Result { let e = match q.query(&mut self.client).await { Ok(t) => return Ok(t), Err(e) => e, }; tracing::error!("failed to run query: {e}"); match e.retry_method() { redis::RetryMethod::Reconnect => { tracing::info!("Redis client is disconnected. Reconnecting..."); self.try_connect().await?; } redis::RetryMethod::RetryImmediately => {} redis::RetryMethod::WaitAndRetry => { // somewhat arbitrary. tokio::time::sleep(Duration::from_millis(100)).await; } _ => Err(e)?, } Ok(q.query(&mut self.client).await?) } }