scheduler/
valkey_coordinated_store.rs

1mod codec;
2mod inflight_commands;
3mod keys;
4mod state_commands;
5
6use crate::coordinated_store::{
7    CoordinatedClaim, CoordinatedCompletion, CoordinatedLeaseConfig, CoordinatedPendingTrigger,
8    CoordinatedRuntimeState, CoordinatedStateStore,
9};
10use crate::error::{ExecutionGuardErrorKind, StoreErrorKind};
11use crate::execution_guard::{ExecutionGuardRenewal, ExecutionGuardScope, ExecutionLease};
12use crate::model::JobState;
13use crate::valkey_runtime::{ValkeyRecoveryConfig, ValkeyRuntime};
14use crate::valkey_store::ValkeyStoreError;
15use keys::{DEFAULT_EXECUTION_KEY_PREFIX, DEFAULT_STATE_KEY_PREFIX};
16use redis::Client;
17
18#[derive(Debug, Clone)]
19pub struct ValkeyCoordinatedStateStore {
20    pub(super) runtime: ValkeyRuntime,
21    pub(super) state_key_prefix: String,
22    pub(super) execution_key_prefix: String,
23}
24
25impl ValkeyCoordinatedStateStore {
26    pub async fn new(url: impl AsRef<str>) -> Result<Self, redis::RedisError> {
27        Self::with_prefixes(url, DEFAULT_STATE_KEY_PREFIX, DEFAULT_EXECUTION_KEY_PREFIX).await
28    }
29
30    pub async fn with_prefixes(
31        url: impl AsRef<str>,
32        state_key_prefix: impl Into<String>,
33        execution_key_prefix: impl Into<String>,
34    ) -> Result<Self, redis::RedisError> {
35        Self::with_prefixes_and_recovery_config(
36            url,
37            state_key_prefix,
38            execution_key_prefix,
39            ValkeyRecoveryConfig::default(),
40        )
41        .await
42    }
43
44    pub async fn with_recovery_config(
45        url: impl AsRef<str>,
46        config: ValkeyRecoveryConfig,
47    ) -> Result<Self, redis::RedisError> {
48        Self::with_prefixes_and_recovery_config(
49            url,
50            DEFAULT_STATE_KEY_PREFIX,
51            DEFAULT_EXECUTION_KEY_PREFIX,
52            config,
53        )
54        .await
55    }
56
57    pub async fn with_prefixes_and_recovery_config(
58        url: impl AsRef<str>,
59        state_key_prefix: impl Into<String>,
60        execution_key_prefix: impl Into<String>,
61        config: ValkeyRecoveryConfig,
62    ) -> Result<Self, redis::RedisError> {
63        let client = Client::open(url.as_ref())?;
64        let runtime = ValkeyRuntime::from_client(client, config).await?;
65        Ok(Self {
66            runtime,
67            state_key_prefix: state_key_prefix.into(),
68            execution_key_prefix: execution_key_prefix.into(),
69        })
70    }
71}
72
73impl CoordinatedStateStore for ValkeyCoordinatedStateStore {
74    type Error = ValkeyStoreError;
75
76    async fn load_or_initialize(
77        &self,
78        job_id: &str,
79        initial_state: JobState,
80    ) -> Result<CoordinatedRuntimeState, Self::Error> {
81        self.load_or_initialize_state(job_id, initial_state).await
82    }
83
84    async fn save_state(
85        &self,
86        job_id: &str,
87        revision: u64,
88        state: &JobState,
89    ) -> Result<bool, Self::Error> {
90        self.save_runtime_state(job_id, revision, state).await
91    }
92
93    async fn reclaim_inflight(
94        &self,
95        job_id: &str,
96        resource_id: &str,
97        lease_config: CoordinatedLeaseConfig,
98    ) -> Result<Option<CoordinatedClaim>, Self::Error> {
99        self.reclaim_inflight_claim(job_id, resource_id, lease_config)
100            .await
101    }
102
103    async fn claim_trigger(
104        &self,
105        job_id: &str,
106        resource_id: &str,
107        revision: u64,
108        trigger: CoordinatedPendingTrigger,
109        next_state: &JobState,
110        lease_config: CoordinatedLeaseConfig,
111        scope: ExecutionGuardScope,
112    ) -> Result<Option<CoordinatedClaim>, Self::Error> {
113        self.claim_trigger_for_scope(
114            job_id,
115            resource_id,
116            revision,
117            trigger,
118            next_state,
119            lease_config,
120            scope,
121        )
122        .await
123    }
124
125    async fn renew(
126        &self,
127        lease: &ExecutionLease,
128        lease_config: CoordinatedLeaseConfig,
129    ) -> Result<ExecutionGuardRenewal, Self::Error> {
130        self.renew_claim_lease(lease, lease_config).await
131    }
132
133    async fn complete(
134        &self,
135        job_id: &str,
136        lease: &ExecutionLease,
137        completion: CoordinatedCompletion,
138    ) -> Result<bool, Self::Error> {
139        self.complete_claim(job_id, lease, completion).await
140    }
141
142    async fn delete(&self, job_id: &str) -> Result<(), Self::Error> {
143        self.delete_state(job_id).await
144    }
145
146    async fn pause(&self, job_id: &str) -> Result<bool, Self::Error> {
147        self.pause_state(job_id).await
148    }
149
150    async fn resume(&self, job_id: &str) -> Result<bool, Self::Error> {
151        self.resume_state(job_id).await
152    }
153
154    fn classify_store_error(error: &Self::Error) -> StoreErrorKind
155    where
156        Self: Sized,
157    {
158        classify_store_error(error)
159    }
160
161    fn classify_guard_error(error: &Self::Error) -> ExecutionGuardErrorKind
162    where
163        Self: Sized,
164    {
165        classify_guard_error(error)
166    }
167}
168
169fn classify_store_error(error: &ValkeyStoreError) -> StoreErrorKind {
170    if matches!(error, ValkeyStoreError::Codec(_)) {
171        StoreErrorKind::Data
172    } else if error.is_connection_issue() {
173        StoreErrorKind::Connection
174    } else {
175        StoreErrorKind::Unknown
176    }
177}
178
179fn classify_guard_error(error: &ValkeyStoreError) -> ExecutionGuardErrorKind {
180    if matches!(error, ValkeyStoreError::Codec(_)) {
181        ExecutionGuardErrorKind::Data
182    } else if error.is_connection_issue() {
183        ExecutionGuardErrorKind::Connection
184    } else {
185        ExecutionGuardErrorKind::Unknown
186    }
187}
scheduler/valkey_coordinated_store.rs

scheduler/
valkey_coordinated_store.rs
