scheduler/
valkey_store.rs

1use crate::error::StoreErrorKind;
2use crate::model::JobState;
3use crate::store::{
4    ResilientStateStore, ResilientStoreError, StateStore, StoreEvent, StoreOperation,
5};
6use crate::valkey_runtime::{
7    ValkeyCommandOutcome, ValkeyRecoveryConfig, ValkeyRuntime, ValkeyRuntimeEvent,
8    is_connection_issue,
9};
10use redis::{AsyncCommands, Client};
11use std::error::Error;
12use std::fmt::{self, Display, Formatter};
13
14const DEFAULT_KEY_PREFIX: &str = "scheduler:valkey:job-state:";
15const LEGACY_DEFAULT_KEY_PREFIX: &str = "scheduler:job-state:";
16
17#[derive(Debug, Clone)]
18pub struct ValkeyStateStore {
19    runtime: ValkeyRuntime,
20    key_prefix: String,
21}
22
23impl ValkeyStateStore {
24    pub async fn new(url: impl AsRef<str>) -> Result<Self, redis::RedisError> {
25        Self::with_prefix(url, DEFAULT_KEY_PREFIX).await
26    }
27
28    /// Creates a Valkey-backed store that falls back to an in-process mirror
29    /// after connection-class failures and writes dirty state back once Valkey
30    /// accepts commands again.
31    pub async fn resilient(
32        url: impl AsRef<str>,
33    ) -> Result<ResilientStateStore<Self>, ValkeyStoreError> {
34        Self::with_prefix_resilient(url, DEFAULT_KEY_PREFIX).await
35    }
36
37    pub async fn with_prefix(
38        url: impl AsRef<str>,
39        key_prefix: impl Into<String>,
40    ) -> Result<Self, redis::RedisError> {
41        Self::with_prefix_and_recovery_config(url, key_prefix, ValkeyRecoveryConfig::default())
42            .await
43    }
44
45    pub async fn with_recovery_config(
46        url: impl AsRef<str>,
47        config: ValkeyRecoveryConfig,
48    ) -> Result<Self, redis::RedisError> {
49        Self::with_prefix_and_recovery_config(url, DEFAULT_KEY_PREFIX, config).await
50    }
51
52    pub async fn with_prefix_and_recovery_config(
53        url: impl AsRef<str>,
54        key_prefix: impl Into<String>,
55        config: ValkeyRecoveryConfig,
56    ) -> Result<Self, redis::RedisError> {
57        let client = Client::open(url.as_ref())?;
58        Self::from_client_with_recovery_config(client, key_prefix, config).await
59    }
60
61    pub async fn from_client(
62        client: Client,
63        key_prefix: impl Into<String>,
64    ) -> Result<Self, redis::RedisError> {
65        Self::from_client_with_recovery_config(client, key_prefix, ValkeyRecoveryConfig::default())
66            .await
67    }
68
69    pub async fn from_client_with_recovery_config(
70        client: Client,
71        key_prefix: impl Into<String>,
72        config: ValkeyRecoveryConfig,
73    ) -> Result<Self, redis::RedisError> {
74        let runtime = ValkeyRuntime::from_client(client, config).await?;
75        Ok(Self {
76            runtime,
77            key_prefix: key_prefix.into(),
78        })
79    }
80
81    pub async fn with_prefix_resilient(
82        url: impl AsRef<str>,
83        key_prefix: impl Into<String>,
84    ) -> Result<ResilientStateStore<Self>, ValkeyStoreError> {
85        ResilientStateStore::from_result(
86            Self::with_prefix(url, key_prefix)
87                .await
88                .map_err(ValkeyStoreError::from),
89        )
90    }
91
92    pub async fn from_client_resilient(
93        client: Client,
94        key_prefix: impl Into<String>,
95    ) -> Result<ResilientStateStore<Self>, ValkeyStoreError> {
96        ResilientStateStore::from_result(
97            Self::from_client(client, key_prefix)
98                .await
99                .map_err(ValkeyStoreError::from),
100        )
101    }
102
103    fn state_key(&self, job_id: &str) -> String {
104        state_key(&self.key_prefix, job_id)
105    }
106
107    fn legacy_state_key(&self, job_id: &str) -> Option<String> {
108        if self.key_prefix == DEFAULT_KEY_PREFIX {
109            Some(state_key(LEGACY_DEFAULT_KEY_PREFIX, job_id))
110        } else {
111            None
112        }
113    }
114}
115
116fn state_key(prefix: &str, job_id: &str) -> String {
117    format!("{prefix}{job_id}")
118}
119
120#[derive(Debug)]
121pub enum ValkeyStoreError {
122    Redis(redis::RedisError),
123    Codec(serde_json::Error),
124    Unavailable,
125}
126
127impl Display for ValkeyStoreError {
128    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
129        match self {
130            ValkeyStoreError::Redis(error) => write!(f, "{error}"),
131            ValkeyStoreError::Codec(error) => write!(f, "{error}"),
132            ValkeyStoreError::Unavailable => write!(f, "valkey is temporarily unavailable"),
133        }
134    }
135}
136
137impl Error for ValkeyStoreError {
138    fn source(&self) -> Option<&(dyn Error + 'static)> {
139        match self {
140            ValkeyStoreError::Redis(error) => Some(error),
141            ValkeyStoreError::Codec(error) => Some(error),
142            ValkeyStoreError::Unavailable => None,
143        }
144    }
145}
146
147impl From<redis::RedisError> for ValkeyStoreError {
148    fn from(error: redis::RedisError) -> Self {
149        Self::Redis(error)
150    }
151}
152
153impl From<serde_json::Error> for ValkeyStoreError {
154    fn from(error: serde_json::Error) -> Self {
155        Self::Codec(error)
156    }
157}
158
159impl ValkeyStoreError {
160    pub fn is_connection_issue(&self) -> bool {
161        match self {
162            Self::Redis(error) => is_connection_issue(error),
163            Self::Unavailable => true,
164            Self::Codec(_) => false,
165        }
166    }
167}
168
169impl ResilientStoreError for ValkeyStoreError {
170    fn is_connection_issue(&self) -> bool {
171        self.is_connection_issue()
172    }
173}
174
175impl StateStore for ValkeyStateStore {
176    type Error = ValkeyStoreError;
177
178    async fn load(&self, job_id: &str) -> Result<Option<JobState>, Self::Error> {
179        let state_key = self.state_key(job_id);
180        let legacy_key = self.legacy_state_key(job_id);
181        let payload = self
182            .runtime
183            .execute(move |mut connection| {
184                let state_key = state_key.clone();
185                let legacy_key = legacy_key.clone();
186                async move {
187                    let payload: Option<String> = connection.get(state_key).await?;
188                    match payload {
189                        Some(payload) => Ok(Some(payload)),
190                        None => {
191                            if let Some(legacy_key) = legacy_key {
192                                connection.get(legacy_key).await
193                            } else {
194                                Ok(None)
195                            }
196                        }
197                    }
198                }
199            })
200            .await
201            .map_err(ValkeyStoreError::from)?;
202
203        let ValkeyCommandOutcome::Available(payload) = payload else {
204            return Err(ValkeyStoreError::Unavailable);
205        };
206
207        payload
208            .map(|value| serde_json::from_str(&value).map_err(ValkeyStoreError::from))
209            .transpose()
210    }
211
212    async fn save(&self, state: &JobState) -> Result<(), Self::Error> {
213        let key = self.state_key(&state.job_id);
214        let payload = serde_json::to_string(state).map_err(ValkeyStoreError::from)?;
215        let result = self
216            .runtime
217            .execute(move |mut connection| {
218                let key = key.clone();
219                let payload = payload.clone();
220                async move { connection.set(key, payload).await }
221            })
222            .await
223            .map_err(ValkeyStoreError::from)?;
224        match result {
225            ValkeyCommandOutcome::Available(()) => Ok(()),
226            ValkeyCommandOutcome::Degraded => Err(ValkeyStoreError::Unavailable),
227        }
228    }
229
230    async fn delete(&self, job_id: &str) -> Result<(), Self::Error> {
231        let state_key = self.state_key(job_id);
232        let legacy_key = self.legacy_state_key(job_id);
233        let result = self
234            .runtime
235            .execute(move |mut connection| {
236                let state_key = state_key.clone();
237                let legacy_key = legacy_key.clone();
238                async move {
239                    let _: usize = connection.del(state_key).await?;
240                    if let Some(legacy_key) = legacy_key {
241                        let _: usize = connection.del(legacy_key).await?;
242                    }
243                    Ok(())
244                }
245            })
246            .await
247            .map_err(ValkeyStoreError::from)?;
248
249        match result {
250            ValkeyCommandOutcome::Available(()) => Ok(()),
251            ValkeyCommandOutcome::Degraded => Err(ValkeyStoreError::Unavailable),
252        }
253    }
254
255    async fn drain_events(&self) -> Result<Vec<StoreEvent>, Self::Error> {
256        Ok(self
257            .runtime
258            .drain_events()
259            .await
260            .into_iter()
261            .map(|event| match event {
262                ValkeyRuntimeEvent::Degraded { error } => StoreEvent::Degraded {
263                    operation: StoreOperation::Save,
264                    error,
265                },
266                ValkeyRuntimeEvent::Recovered => StoreEvent::Recovered {
267                    operation: StoreOperation::Save,
268                },
269            })
270            .collect())
271    }
272
273    fn classify_error(error: &Self::Error) -> StoreErrorKind
274    where
275        Self: Sized,
276    {
277        if matches!(error, ValkeyStoreError::Codec(_)) {
278            StoreErrorKind::Data
279        } else if error.is_connection_issue() {
280            StoreErrorKind::Connection
281        } else {
282            StoreErrorKind::Unknown
283        }
284    }
285}
286
287#[cfg(test)]
288mod tests {
289    use super::{DEFAULT_KEY_PREFIX, LEGACY_DEFAULT_KEY_PREFIX, ValkeyStoreError, state_key};
290    use crate::model::JobState;
291    use chrono::{TimeDelta, Utc};
292
293    #[test]
294    fn state_key_uses_custom_prefix() {
295        assert_eq!(state_key("custom:", "job-1"), "custom:job-1");
296        assert_eq!(
297            state_key(DEFAULT_KEY_PREFIX, "job-2"),
298            "scheduler:valkey:job-state:job-2"
299        );
300    }
301
302    #[test]
303    fn legacy_default_prefix_is_stable() {
304        assert_eq!(
305            state_key(LEGACY_DEFAULT_KEY_PREFIX, "job-3"),
306            "scheduler:job-state:job-3"
307        );
308    }
309
310    #[test]
311    fn job_state_json_round_trip() {
312        let state = JobState {
313            job_id: "job-1".to_string(),
314            trigger_count: 2,
315            last_run_at: Some(Utc::now()),
316            last_success_at: Some(Utc::now() + TimeDelta::seconds(1)),
317            next_run_at: Some(Utc::now() + TimeDelta::seconds(5)),
318            last_error: Some("boom".to_string()),
319        };
320
321        let encoded = serde_json::to_string(&state).unwrap();
322        let decoded: JobState = serde_json::from_str(&encoded).unwrap();
323
324        assert_eq!(decoded, state);
325    }
326
327    #[test]
328    fn io_errors_are_classified_as_connection_issues() {
329        let error = ValkeyStoreError::from(redis::RedisError::from(std::io::Error::from(
330            std::io::ErrorKind::BrokenPipe,
331        )));
332
333        assert!(error.is_connection_issue());
334    }
335
336    #[test]
337    fn timeout_errors_are_classified_as_connection_issues() {
338        let error = ValkeyStoreError::from(redis::RedisError::from(std::io::Error::from(
339            std::io::ErrorKind::TimedOut,
340        )));
341
342        assert!(error.is_connection_issue());
343    }
344
345    #[test]
346    fn codec_errors_are_not_classified_as_connection_issues() {
347        let error = ValkeyStoreError::from(serde_json::from_str::<JobState>("{").unwrap_err());
348
349        assert!(!error.is_connection_issue());
350    }
351}
scheduler/valkey_store.rs

scheduler/
valkey_store.rs
