scheduler/
valkey_coordinated_store.rs

1use crate::coordinated_store::{
2    CoordinatedClaim, CoordinatedLeaseConfig, CoordinatedPendingTrigger, CoordinatedRuntimeState,
3    CoordinatedStateStore,
4};
5use crate::error::{ExecutionGuardErrorKind, StoreErrorKind};
6use crate::execution_guard::{ExecutionGuardRenewal, ExecutionGuardScope, ExecutionLease};
7use crate::model::JobState;
8use crate::valkey_execution_support::{
9    next_token, now_millis, occurrence_index_key, occurrence_lease_key, resource_lock_key,
10};
11use crate::valkey_runtime::{ValkeyCommandOutcome, ValkeyRecoveryConfig, ValkeyRuntime};
12use crate::valkey_scripts;
13use crate::valkey_store::ValkeyStoreError;
14use chrono::SecondsFormat;
15use chrono::{DateTime, Utc};
16use redis::{AsyncCommands, Client, cmd};
17use std::collections::HashMap;
18use std::sync::atomic::AtomicU64;
19
20const DEFAULT_STATE_KEY_PREFIX: &str = "scheduler:valkey:job-state:";
21const LEGACY_DEFAULT_STATE_KEY_PREFIX: &str = "scheduler:job-state:";
22const DEFAULT_EXECUTION_KEY_PREFIX: &str = "scheduler:valkey:execution-lease:";
23
24const FIELD_VERSION: &str = "version";
25const FIELD_STATE: &str = "state";
26const FIELD_PAUSED: &str = "paused";
27const FIELD_INFLIGHT_SCHEDULED_AT: &str = "inflight_scheduled_at";
28const FIELD_INFLIGHT_CATCH_UP: &str = "inflight_catch_up";
29const FIELD_INFLIGHT_TRIGGER_COUNT: &str = "inflight_trigger_count";
30const FIELD_INFLIGHT_RESOURCE_ID: &str = "inflight_resource_id";
31const FIELD_INFLIGHT_SCOPE: &str = "inflight_scope";
32const FIELD_INFLIGHT_TOKEN: &str = "inflight_token";
33const FIELD_INFLIGHT_LEASE_KEY: &str = "inflight_lease_key";
34const FIELD_INFLIGHT_LEASE_EXPIRES_AT: &str = "inflight_lease_expires_at";
35
36static COORDINATED_TOKEN_COUNTER: AtomicU64 = AtomicU64::new(1);
37
38#[derive(Debug, Clone)]
39pub struct ValkeyCoordinatedStateStore {
40    runtime: ValkeyRuntime,
41    state_key_prefix: String,
42    execution_key_prefix: String,
43}
44
45impl ValkeyCoordinatedStateStore {
46    pub async fn new(url: impl AsRef<str>) -> Result<Self, redis::RedisError> {
47        Self::with_prefixes(url, DEFAULT_STATE_KEY_PREFIX, DEFAULT_EXECUTION_KEY_PREFIX).await
48    }
49
50    pub async fn with_prefixes(
51        url: impl AsRef<str>,
52        state_key_prefix: impl Into<String>,
53        execution_key_prefix: impl Into<String>,
54    ) -> Result<Self, redis::RedisError> {
55        Self::with_prefixes_and_recovery_config(
56            url,
57            state_key_prefix,
58            execution_key_prefix,
59            ValkeyRecoveryConfig::default(),
60        )
61        .await
62    }
63
64    pub async fn with_recovery_config(
65        url: impl AsRef<str>,
66        config: ValkeyRecoveryConfig,
67    ) -> Result<Self, redis::RedisError> {
68        Self::with_prefixes_and_recovery_config(
69            url,
70            DEFAULT_STATE_KEY_PREFIX,
71            DEFAULT_EXECUTION_KEY_PREFIX,
72            config,
73        )
74        .await
75    }
76
77    pub async fn with_prefixes_and_recovery_config(
78        url: impl AsRef<str>,
79        state_key_prefix: impl Into<String>,
80        execution_key_prefix: impl Into<String>,
81        config: ValkeyRecoveryConfig,
82    ) -> Result<Self, redis::RedisError> {
83        let client = Client::open(url.as_ref())?;
84        let runtime = ValkeyRuntime::from_client(client, config).await?;
85        Ok(Self {
86            runtime,
87            state_key_prefix: state_key_prefix.into(),
88            execution_key_prefix: execution_key_prefix.into(),
89        })
90    }
91
92    fn state_key(&self, job_id: &str) -> String {
93        format!("{}{}", self.state_key_prefix, job_id)
94    }
95
96    fn legacy_state_key(&self, job_id: &str) -> Option<String> {
97        if self.state_key_prefix == DEFAULT_STATE_KEY_PREFIX {
98            Some(format!("{LEGACY_DEFAULT_STATE_KEY_PREFIX}{job_id}"))
99        } else {
100            None
101        }
102    }
103
104    fn resource_lock_key(&self, resource_id: &str) -> String {
105        resource_lock_key(&self.execution_key_prefix, resource_id)
106    }
107
108    fn occurrence_index_key(&self, resource_id: &str) -> String {
109        occurrence_index_key(&self.execution_key_prefix, resource_id)
110    }
111
112    fn occurrence_lease_key(&self, resource_id: &str, scheduled_at: DateTime<Utc>) -> String {
113        occurrence_lease_key(&self.execution_key_prefix, resource_id, scheduled_at)
114    }
115
116    async fn key_type(&self, key: &str) -> Result<String, ValkeyStoreError> {
117        let key = key.to_string();
118        let result: ValkeyCommandOutcome<String> = self
119            .runtime
120            .execute(move |mut connection| {
121                let key = key.clone();
122                async move { cmd("TYPE").arg(key).query_async(&mut connection).await }
123            })
124            .await
125            .map_err(ValkeyStoreError::from)?;
126        match result {
127            ValkeyCommandOutcome::Available(value) => Ok(value),
128            ValkeyCommandOutcome::Degraded => Err(ValkeyStoreError::Unavailable),
129        }
130    }
131
132    async fn load_hash(
133        &self,
134        key: &str,
135    ) -> Result<Option<CoordinatedRuntimeState>, ValkeyStoreError> {
136        let key = key.to_string();
137        let result: ValkeyCommandOutcome<HashMap<String, String>> = self
138            .runtime
139            .execute(move |mut connection| {
140                let key = key.clone();
141                async move { connection.hgetall(key).await }
142            })
143            .await
144            .map_err(ValkeyStoreError::from)?;
145        let ValkeyCommandOutcome::Available(fields) = result else {
146            return Err(ValkeyStoreError::Unavailable);
147        };
148        if fields.is_empty() {
149            return Ok(None);
150        }
151
152        Ok(Some(parse_runtime_state(&fields)?))
153    }
154
155    async fn migrate_string_state(
156        &self,
157        key: &str,
158        payload: String,
159    ) -> Result<CoordinatedRuntimeState, ValkeyStoreError> {
160        let state: JobState = serde_json::from_str(&payload).map_err(ValkeyStoreError::from)?;
161        let runtime = CoordinatedRuntimeState {
162            state,
163            revision: 0,
164            paused: false,
165        };
166        self.write_runtime(key, &runtime).await?;
167        Ok(runtime)
168    }
169
170    async fn write_runtime(
171        &self,
172        key: &str,
173        runtime: &CoordinatedRuntimeState,
174    ) -> Result<(), ValkeyStoreError> {
175        let payload = serde_json::to_string(&runtime.state).map_err(ValkeyStoreError::from)?;
176        let key = key.to_string();
177        let revision = runtime.revision;
178        let paused = runtime.paused;
179        let result = self
180            .runtime
181            .execute(move |mut connection| {
182                let key = key.clone();
183                let payload = payload.clone();
184                async move {
185                    let _: () = cmd("DEL").arg(&key).query_async(&mut connection).await?;
186                    let _: () = cmd("HSET")
187                        .arg(key)
188                        .arg(FIELD_VERSION)
189                        .arg(revision)
190                        .arg(FIELD_STATE)
191                        .arg(payload)
192                        .arg(FIELD_PAUSED)
193                        .arg(if paused { "1" } else { "0" })
194                        .query_async(&mut connection)
195                        .await?;
196                    Ok(())
197                }
198            })
199            .await
200            .map_err(ValkeyStoreError::from)?;
201        match result {
202            ValkeyCommandOutcome::Available(()) => Ok(()),
203            ValkeyCommandOutcome::Degraded => Err(ValkeyStoreError::Unavailable),
204        }
205    }
206
207    async fn load_payload_state(&self, key: &str) -> Result<Option<String>, ValkeyStoreError> {
208        let key = key.to_string();
209        let result = self
210            .runtime
211            .execute(move |mut connection| {
212                let key = key.clone();
213                async move { connection.get(key).await }
214            })
215            .await
216            .map_err(ValkeyStoreError::from)?;
217        match result {
218            ValkeyCommandOutcome::Available(value) => Ok(value),
219            ValkeyCommandOutcome::Degraded => Err(ValkeyStoreError::Unavailable),
220        }
221    }
222}
223
224impl CoordinatedStateStore for ValkeyCoordinatedStateStore {
225    type Error = ValkeyStoreError;
226
227    async fn load_or_initialize(
228        &self,
229        job_id: &str,
230        initial_state: JobState,
231    ) -> Result<CoordinatedRuntimeState, Self::Error> {
232        let key = self.state_key(job_id);
233        match self.key_type(&key).await?.as_str() {
234            "hash" => {
235                if let Some(runtime) = self.load_hash(&key).await? {
236                    return Ok(runtime);
237                }
238            }
239            "string" => {
240                if let Some(payload) = self.load_payload_state(&key).await? {
241                    return self.migrate_string_state(&key, payload).await;
242                }
243            }
244            "none" => {}
245            _ => {}
246        }
247
248        if let Some(legacy_key) = self.legacy_state_key(job_id) {
249            if self.key_type(&legacy_key).await?.as_str() == "string" {
250                if let Some(payload) = self.load_payload_state(&legacy_key).await? {
251                    let runtime = self.migrate_string_state(&key, payload).await?;
252                    let result: ValkeyCommandOutcome<()> = self
253                        .runtime
254                        .execute(move |mut connection| {
255                            let legacy_key = legacy_key.clone();
256                            async move {
257                                cmd("DEL")
258                                    .arg(legacy_key)
259                                    .query_async(&mut connection)
260                                    .await
261                            }
262                        })
263                        .await
264                        .map_err(ValkeyStoreError::from)?;
265                    if matches!(result, ValkeyCommandOutcome::Degraded) {
266                        return Err(ValkeyStoreError::Unavailable);
267                    }
268                    return Ok(runtime);
269                }
270            }
271        }
272
273        let runtime = CoordinatedRuntimeState {
274            state: initial_state,
275            revision: 0,
276            paused: false,
277        };
278        self.write_runtime(&key, &runtime).await?;
279        Ok(runtime)
280    }
281
282    async fn save_state(
283        &self,
284        job_id: &str,
285        revision: u64,
286        state: &JobState,
287    ) -> Result<bool, Self::Error> {
288        let key = self.state_key(job_id);
289        let payload = serde_json::to_string(state).map_err(ValkeyStoreError::from)?;
290        let result: ValkeyCommandOutcome<i32> = self
291            .runtime
292            .execute(move |mut connection| {
293                let key = key.clone();
294                let payload = payload.clone();
295                async move {
296                    valkey_scripts::script(valkey_scripts::coordinated::SAVE_STATE)
297                        .key(key)
298                        .arg(FIELD_VERSION)
299                        .arg(revision)
300                        .arg(FIELD_INFLIGHT_TOKEN)
301                        .arg(FIELD_STATE)
302                        .arg(payload)
303                        .invoke_async(&mut connection)
304                        .await
305                }
306            })
307            .await
308            .map_err(ValkeyStoreError::from)?;
309        match result {
310            ValkeyCommandOutcome::Available(updated) => Ok(updated == 1),
311            ValkeyCommandOutcome::Degraded => Err(ValkeyStoreError::Unavailable),
312        }
313    }
314
315    async fn reclaim_inflight(
316        &self,
317        job_id: &str,
318        resource_id: &str,
319        lease_config: CoordinatedLeaseConfig,
320    ) -> Result<Option<CoordinatedClaim>, Self::Error> {
321        let key = self.state_key(job_id);
322        let lease_key = self.occurrence_lease_key(resource_id, Utc::now());
323        let token = next_token(&COORDINATED_TOKEN_COUNTER, "coord");
324        let ttl_millis = u64::try_from(lease_config.ttl.as_millis()).unwrap_or(u64::MAX);
325        let now_millis = now_millis();
326        let expires_at_millis = now_millis.saturating_add(ttl_millis);
327        let resource_lock_key = self.resource_lock_key(resource_id);
328        let occurrence_index_key = self.occurrence_index_key(resource_id);
329        let occurrence_prefix = format!("{}{}:occurrence:", self.execution_key_prefix, resource_id);
330        let result: ValkeyCommandOutcome<Option<Vec<String>>> = self
331            .runtime
332            .execute({
333                let token = token.clone();
334                let lease_key = lease_key.clone();
335                move |mut connection| {
336                    let key = key.clone();
337                    let resource_lock_key = resource_lock_key.clone();
338                    let lease_key = lease_key.clone();
339                    let occurrence_index_key = occurrence_index_key.clone();
340                    let occurrence_prefix = occurrence_prefix.clone();
341                    let token = token.clone();
342                    async move {
343                        valkey_scripts::script(valkey_scripts::coordinated::RECLAIM_INFLIGHT)
344                            .key(key)
345                            .key(resource_lock_key)
346                            .key(lease_key)
347                            .key(occurrence_index_key)
348                            // ARGV layout:
349                            //  1..9  = hash field names read from KEYS[1]
350                            // 10..14 = reclaim timing + new lease payload
351                            // 15..16 = hash field names updated on successful reclaim
352                            .arg(FIELD_INFLIGHT_SCHEDULED_AT)
353                            .arg(FIELD_INFLIGHT_CATCH_UP)
354                            .arg(FIELD_INFLIGHT_TRIGGER_COUNT)
355                            .arg(FIELD_INFLIGHT_RESOURCE_ID)
356                            .arg(FIELD_INFLIGHT_SCOPE)
357                            .arg(FIELD_INFLIGHT_LEASE_EXPIRES_AT)
358                            .arg(FIELD_STATE)
359                            .arg(FIELD_VERSION)
360                            .arg(FIELD_PAUSED)
361                            .arg(now_millis)
362                            .arg(occurrence_prefix)
363                            .arg(&token)
364                            .arg(ttl_millis)
365                            .arg(expires_at_millis)
366                            .arg(FIELD_INFLIGHT_TOKEN)
367                            .arg(FIELD_INFLIGHT_LEASE_KEY)
368                            .invoke_async(&mut connection)
369                            .await
370                    }
371                }
372            })
373            .await
374            .map_err(ValkeyStoreError::from)?;
375
376        let ValkeyCommandOutcome::Available(result) = result else {
377            return Err(ValkeyStoreError::Unavailable);
378        };
379
380        let Some(values) = result else {
381            return Ok(None);
382        };
383        if values.len() != 8 {
384            return Ok(None);
385        }
386        let revision = values[0].parse::<u64>().unwrap_or(0);
387        let state: JobState = serde_json::from_str(&values[1]).map_err(ValkeyStoreError::from)?;
388        let scheduled_at = DateTime::parse_from_rfc3339(&values[2])
389            .map_err(|error| {
390                ValkeyStoreError::Codec(serde_json::Error::io(std::io::Error::other(
391                    error.to_string(),
392                )))
393            })?
394            .with_timezone(&Utc);
395        let catch_up = values[3].parse::<bool>().unwrap_or(false);
396        let trigger_count = values[4].parse::<u32>().unwrap_or(0);
397        let scope = parse_scope(&values[5]);
398        Ok(Some(CoordinatedClaim {
399            state: CoordinatedRuntimeState {
400                state,
401                revision,
402                paused: false,
403            },
404            trigger: CoordinatedPendingTrigger {
405                scheduled_at,
406                catch_up,
407                trigger_count,
408            },
409            lease: ExecutionLease::new(
410                job_id.to_string(),
411                resource_id.to_string(),
412                scope,
413                Some(scheduled_at),
414                values[7].clone(),
415                values[6].clone(),
416            ),
417            replayed: true,
418        }))
419    }
420
421    async fn claim_trigger(
422        &self,
423        job_id: &str,
424        resource_id: &str,
425        revision: u64,
426        trigger: CoordinatedPendingTrigger,
427        next_state: &JobState,
428        lease_config: CoordinatedLeaseConfig,
429    ) -> Result<Option<CoordinatedClaim>, Self::Error> {
430        let key = self.state_key(job_id);
431        let lease_key = self.occurrence_lease_key(resource_id, trigger.scheduled_at);
432        let token = next_token(&COORDINATED_TOKEN_COUNTER, "coord");
433        let ttl_millis = u64::try_from(lease_config.ttl.as_millis()).unwrap_or(u64::MAX);
434        let now_millis = now_millis();
435        let expires_at_millis = now_millis.saturating_add(ttl_millis);
436        let next_state_payload =
437            serde_json::to_string(next_state).map_err(ValkeyStoreError::from)?;
438        let resource_lock_key = self.resource_lock_key(resource_id);
439        let occurrence_index_key = self.occurrence_index_key(resource_id);
440        let scheduled_at = trigger
441            .scheduled_at
442            .to_rfc3339_opts(SecondsFormat::Nanos, true);
443        let resource_id_arg = resource_id.to_string();
444        let command_trigger = trigger.clone();
445        let result: ValkeyCommandOutcome<i64> = self
446            .runtime
447            .execute({
448                let token = token.clone();
449                let lease_key = lease_key.clone();
450                let command_trigger = command_trigger.clone();
451                move |mut connection| {
452                    let key = key.clone();
453                    let resource_lock_key = resource_lock_key.clone();
454                    let lease_key = lease_key.clone();
455                    let occurrence_index_key = occurrence_index_key.clone();
456                    let token = token.clone();
457                    let next_state_payload = next_state_payload.clone();
458                    let scheduled_at = scheduled_at.clone();
459                    let resource_id_arg = resource_id_arg.clone();
460                    let command_trigger = command_trigger.clone();
461                    async move {
462                        valkey_scripts::script(valkey_scripts::coordinated::CLAIM_TRIGGER)
463                            .key(key)
464                            .key(resource_lock_key)
465                            .key(&lease_key)
466                            .key(occurrence_index_key)
467                            .arg(FIELD_VERSION)
468                            .arg(FIELD_INFLIGHT_TOKEN)
469                            .arg(FIELD_INFLIGHT_LEASE_EXPIRES_AT)
470                            .arg(FIELD_PAUSED)
471                            .arg(now_millis)
472                            .arg(revision)
473                            .arg(&token)
474                            .arg(ttl_millis)
475                            .arg(expires_at_millis)
476                            .arg(FIELD_STATE)
477                            .arg(next_state_payload)
478                            .arg(FIELD_INFLIGHT_SCHEDULED_AT)
479                            .arg(scheduled_at)
480                            .arg(FIELD_INFLIGHT_CATCH_UP)
481                            .arg(command_trigger.catch_up)
482                            .arg(FIELD_INFLIGHT_TRIGGER_COUNT)
483                            .arg(command_trigger.trigger_count)
484                            .arg(FIELD_INFLIGHT_RESOURCE_ID)
485                            .arg(resource_id_arg)
486                            .arg(FIELD_INFLIGHT_SCOPE)
487                            .arg("occurrence")
488                            .arg(FIELD_INFLIGHT_TOKEN)
489                            .arg(FIELD_INFLIGHT_LEASE_KEY)
490                            .invoke_async(&mut connection)
491                            .await
492                    }
493                }
494            })
495            .await
496            .map_err(ValkeyStoreError::from)?;
497        let ValkeyCommandOutcome::Available(new_revision) = result else {
498            return Err(ValkeyStoreError::Unavailable);
499        };
500
501        if new_revision <= 0 {
502            return Ok(None);
503        }
504
505        Ok(Some(CoordinatedClaim {
506            state: CoordinatedRuntimeState {
507                state: next_state.clone(),
508                revision: new_revision as u64,
509                paused: false,
510            },
511            trigger: trigger.clone(),
512            lease: ExecutionLease::new(
513                job_id.to_string(),
514                resource_id.to_string(),
515                ExecutionGuardScope::Occurrence,
516                Some(trigger.scheduled_at),
517                token,
518                lease_key,
519            ),
520            replayed: false,
521        }))
522    }
523
524    async fn renew(
525        &self,
526        lease: &ExecutionLease,
527        lease_config: CoordinatedLeaseConfig,
528    ) -> Result<ExecutionGuardRenewal, Self::Error> {
529        let ttl_millis = u64::try_from(lease_config.ttl.as_millis()).unwrap_or(u64::MAX);
530        let expires_at_millis = now_millis().saturating_add(ttl_millis);
531        let lease = lease.clone();
532        let occurrence_index_key = self.occurrence_index_key(&lease.resource_id);
533        let state_key = self.state_key(&lease.job_id);
534        let result: ValkeyCommandOutcome<i32> = self
535            .runtime
536            .execute(move |mut connection| {
537                let lease = lease.clone();
538                let occurrence_index_key = occurrence_index_key.clone();
539                let state_key = state_key.clone();
540                async move {
541                    valkey_scripts::script(valkey_scripts::coordinated::RENEW_LEASE)
542                        .key(&lease.lease_key)
543                        .key(occurrence_index_key)
544                        .key(state_key)
545                        .arg(&lease.token)
546                        .arg(ttl_millis)
547                        .arg(expires_at_millis)
548                        .arg(FIELD_INFLIGHT_LEASE_EXPIRES_AT)
549                        .invoke_async(&mut connection)
550                        .await
551                }
552            })
553            .await
554            .map_err(ValkeyStoreError::from)?;
555        let ValkeyCommandOutcome::Available(renewed) = result else {
556            return Ok(ExecutionGuardRenewal::Lost);
557        };
558        Ok(if renewed == 1 {
559            ExecutionGuardRenewal::Renewed
560        } else {
561            ExecutionGuardRenewal::Lost
562        })
563    }
564
565    async fn complete(
566        &self,
567        job_id: &str,
568        revision: u64,
569        lease: &ExecutionLease,
570        state: &JobState,
571    ) -> Result<bool, Self::Error> {
572        let key = self.state_key(job_id);
573        let payload = serde_json::to_string(state).map_err(ValkeyStoreError::from)?;
574        let lease = lease.clone();
575        let occurrence_index_key = self.occurrence_index_key(&lease.resource_id);
576        let result: ValkeyCommandOutcome<i32> = self
577            .runtime
578            .execute(move |mut connection| {
579                let key = key.clone();
580                let lease = lease.clone();
581                let occurrence_index_key = occurrence_index_key.clone();
582                let payload = payload.clone();
583                async move {
584                    valkey_scripts::script(valkey_scripts::coordinated::COMPLETE)
585                        .key(key)
586                        .key(&lease.lease_key)
587                        .key(occurrence_index_key)
588                        .arg(FIELD_VERSION)
589                        .arg(FIELD_INFLIGHT_TOKEN)
590                        .arg(revision)
591                        .arg(&lease.token)
592                        .arg(FIELD_STATE)
593                        .arg(payload)
594                        .arg(FIELD_INFLIGHT_SCHEDULED_AT)
595                        .arg(FIELD_INFLIGHT_CATCH_UP)
596                        .arg(FIELD_INFLIGHT_TRIGGER_COUNT)
597                        .arg(FIELD_INFLIGHT_RESOURCE_ID)
598                        .arg(FIELD_INFLIGHT_SCOPE)
599                        .arg(FIELD_INFLIGHT_LEASE_KEY)
600                        .invoke_async(&mut connection)
601                        .await
602                }
603            })
604            .await
605            .map_err(ValkeyStoreError::from)?;
606        match result {
607            ValkeyCommandOutcome::Available(completed) => Ok(completed == 1),
608            ValkeyCommandOutcome::Degraded => Err(ValkeyStoreError::Unavailable),
609        }
610    }
611
612    async fn delete(&self, job_id: &str) -> Result<(), Self::Error> {
613        let key = self.state_key(job_id);
614        let result = self
615            .runtime
616            .execute(move |mut connection| {
617                let key = key.clone();
618                async move { cmd("DEL").arg(key).query_async(&mut connection).await }
619            })
620            .await
621            .map_err(ValkeyStoreError::from)?;
622        match result {
623            ValkeyCommandOutcome::Available(()) => Ok(()),
624            ValkeyCommandOutcome::Degraded => Err(ValkeyStoreError::Unavailable),
625        }
626    }
627
628    async fn pause(&self, job_id: &str) -> Result<bool, Self::Error> {
629        let key = self.state_key(job_id);
630        let result: ValkeyCommandOutcome<i32> = self
631            .runtime
632            .execute(move |mut connection| {
633                let key = key.clone();
634                async move {
635                    valkey_scripts::script(valkey_scripts::coordinated::PAUSE)
636                        .key(key)
637                        .arg(FIELD_PAUSED)
638                        .invoke_async(&mut connection)
639                        .await
640                }
641            })
642            .await
643            .map_err(ValkeyStoreError::from)?;
644        match result {
645            ValkeyCommandOutcome::Available(changed) => Ok(changed == 1),
646            ValkeyCommandOutcome::Degraded => Err(ValkeyStoreError::Unavailable),
647        }
648    }
649
650    async fn resume(&self, job_id: &str) -> Result<bool, Self::Error> {
651        let key = self.state_key(job_id);
652        let result: ValkeyCommandOutcome<i32> = self
653            .runtime
654            .execute(move |mut connection| {
655                let key = key.clone();
656                async move {
657                    valkey_scripts::script(valkey_scripts::coordinated::RESUME)
658                        .key(key)
659                        .arg(FIELD_PAUSED)
660                        .invoke_async(&mut connection)
661                        .await
662                }
663            })
664            .await
665            .map_err(ValkeyStoreError::from)?;
666        match result {
667            ValkeyCommandOutcome::Available(changed) => Ok(changed == 1),
668            ValkeyCommandOutcome::Degraded => Err(ValkeyStoreError::Unavailable),
669        }
670    }
671
672    fn classify_store_error(error: &Self::Error) -> StoreErrorKind
673    where
674        Self: Sized,
675    {
676        if matches!(error, ValkeyStoreError::Codec(_)) {
677            StoreErrorKind::Data
678        } else if error.is_connection_issue() {
679            StoreErrorKind::Connection
680        } else {
681            StoreErrorKind::Unknown
682        }
683    }
684
685    fn classify_guard_error(error: &Self::Error) -> ExecutionGuardErrorKind
686    where
687        Self: Sized,
688    {
689        if matches!(error, ValkeyStoreError::Codec(_)) {
690            ExecutionGuardErrorKind::Data
691        } else if error.is_connection_issue() {
692            ExecutionGuardErrorKind::Connection
693        } else {
694            ExecutionGuardErrorKind::Unknown
695        }
696    }
697}
698
699fn parse_runtime_state(
700    fields: &HashMap<String, String>,
701) -> Result<CoordinatedRuntimeState, ValkeyStoreError> {
702    let revision = fields
703        .get(FIELD_VERSION)
704        .and_then(|value| value.parse::<u64>().ok())
705        .unwrap_or(0);
706    let paused = fields
707        .get(FIELD_PAUSED)
708        .map(|value| value == "1" || value.eq_ignore_ascii_case("true"))
709        .unwrap_or(false);
710    let state = serde_json::from_str(fields.get(FIELD_STATE).map(String::as_str).unwrap_or("{}"))
711        .map_err(ValkeyStoreError::from)?;
712    Ok(CoordinatedRuntimeState {
713        state,
714        revision,
715        paused,
716    })
717}
718
719fn parse_scope(raw: &str) -> ExecutionGuardScope {
720    match raw {
721        "resource" => ExecutionGuardScope::Resource,
722        _ => ExecutionGuardScope::Occurrence,
723    }
724}
scheduler/valkey_coordinated_store.rs

scheduler/
valkey_coordinated_store.rs
