scheduler/
valkey_runtime.rs

1use redis::{Client, ErrorKind, ServerErrorKind, aio::ConnectionManager};
2use std::collections::VecDeque;
3use std::future::Future;
4use std::sync::Arc;
5use std::time::Duration;
6use tokio::sync::Mutex;
7
8#[derive(Debug, Clone, Copy, PartialEq, Eq)]
9pub struct ValkeyRecoveryConfig {
10    pub retry_attempts: usize,
11    pub initial_backoff: Duration,
12    pub max_backoff: Duration,
13    pub probe_interval: Duration,
14}
15
16impl Default for ValkeyRecoveryConfig {
17    fn default() -> Self {
18        Self {
19            retry_attempts: 3,
20            initial_backoff: Duration::from_millis(25),
21            max_backoff: Duration::from_millis(250),
22            probe_interval: Duration::from_millis(250),
23        }
24    }
25}
26
27#[derive(Debug, Clone, PartialEq, Eq)]
28pub(crate) enum ValkeyRuntimeEvent {
29    Degraded { error: String },
30    Recovered,
31}
32
33#[derive(Debug, Clone, PartialEq, Eq)]
34pub(crate) enum ValkeyCommandOutcome<T> {
35    Available(T),
36    Degraded,
37}
38
39#[derive(Debug, Clone)]
40pub(crate) struct ValkeyRuntime {
41    connection: ConnectionManager,
42    config: ValkeyRecoveryConfig,
43    degraded: Arc<std::sync::atomic::AtomicBool>,
44    events: Arc<Mutex<VecDeque<ValkeyRuntimeEvent>>>,
45}
46
47impl ValkeyRuntime {
48    pub(crate) async fn from_client(
49        client: Client,
50        config: ValkeyRecoveryConfig,
51    ) -> Result<Self, redis::RedisError> {
52        let connection = client.get_connection_manager().await?;
53        Ok(Self {
54            connection,
55            config,
56            degraded: Arc::new(std::sync::atomic::AtomicBool::new(false)),
57            events: Arc::new(Mutex::new(VecDeque::new())),
58        })
59    }
60
61    pub(crate) async fn execute<T, F, Fut>(
62        &self,
63        mut operation: F,
64    ) -> Result<ValkeyCommandOutcome<T>, redis::RedisError>
65    where
66        F: FnMut(ConnectionManager) -> Fut,
67        Fut: Future<Output = Result<T, redis::RedisError>>,
68    {
69        let attempts = self.config.retry_attempts.max(1);
70        let mut backoff = self.config.initial_backoff;
71        let mut last_connection_error = None;
72
73        for attempt in 0..attempts {
74            let result = operation(self.connection.clone()).await;
75            match result {
76                Ok(value) => {
77                    if self
78                        .degraded
79                        .swap(false, std::sync::atomic::Ordering::SeqCst)
80                    {
81                        self.events
82                            .lock()
83                            .await
84                            .push_back(ValkeyRuntimeEvent::Recovered);
85                    }
86                    return Ok(ValkeyCommandOutcome::Available(value));
87                }
88                Err(error) if is_connection_issue(&error) => {
89                    last_connection_error = Some(error);
90                    if attempt + 1 < attempts {
91                        tokio::time::sleep(backoff).await;
92                        backoff = std::cmp::min(backoff.saturating_mul(2), self.config.max_backoff);
93                    }
94                }
95                Err(error) => return Err(error),
96            }
97        }
98
99        if let Some(error) = last_connection_error {
100            if !self
101                .degraded
102                .swap(true, std::sync::atomic::Ordering::SeqCst)
103            {
104                self.events
105                    .lock()
106                    .await
107                    .push_back(ValkeyRuntimeEvent::Degraded {
108                        error: error.to_string(),
109                    });
110            }
111            Ok(ValkeyCommandOutcome::Degraded)
112        } else {
113            Ok(ValkeyCommandOutcome::Degraded)
114        }
115    }
116
117    pub(crate) async fn drain_events(&self) -> Vec<ValkeyRuntimeEvent> {
118        let mut events = self.events.lock().await;
119        events.drain(..).collect()
120    }
121}
122
123pub(crate) fn is_connection_issue(error: &redis::RedisError) -> bool {
124    error.is_connection_dropped()
125        || error.is_connection_refusal()
126        || error.is_timeout()
127        || matches!(
128            error.kind(),
129            ErrorKind::Io
130                | ErrorKind::ClusterConnectionNotFound
131                | ErrorKind::Server(ServerErrorKind::BusyLoading)
132                | ErrorKind::Server(ServerErrorKind::ClusterDown)
133                | ErrorKind::Server(ServerErrorKind::MasterDown)
134                | ErrorKind::Server(ServerErrorKind::TryAgain)
135        )
136}
scheduler/valkey_runtime.rs

scheduler/
valkey_runtime.rs
