1use crate::error::{ExecutionGuardError, ExecutionGuardErrorKind};
2use crate::valkey_execution_support::{
3 lease_key, next_token, now_millis, occurrence_index_key, resource_lock_key,
4};
5use crate::valkey_runtime::{
6 ValkeyCommandOutcome, ValkeyRecoveryConfig, ValkeyRuntime, ValkeyRuntimeEvent,
7 is_connection_issue,
8};
9use crate::valkey_scripts;
10use crate::{
11 ExecutionGuard, ExecutionGuardAcquire, ExecutionGuardEvent, ExecutionGuardRenewal,
12 ExecutionGuardScope, ExecutionLease, ExecutionSlot,
13};
14use redis::Client;
15use std::fmt::{self, Display, Formatter};
16use std::num::TryFromIntError;
17use std::sync::atomic::AtomicU64;
18use std::time::Duration;
19
20const DEFAULT_KEY_PREFIX: &str = "scheduler:valkey:execution-lease:";
21
22static TOKEN_COUNTER: AtomicU64 = AtomicU64::new(1);
23
24#[derive(Debug, Clone, Copy, PartialEq, Eq)]
25pub struct ValkeyLeaseConfig {
26 pub ttl: Duration,
27 pub renew_interval: Duration,
28}
29
30#[derive(Debug, Clone)]
31pub struct ValkeyExecutionGuard {
32 runtime: ValkeyRuntime,
33 key_prefix: String,
34 lease_config: ValkeyLeaseConfig,
35}
36
37impl ValkeyExecutionGuard {
38 pub async fn new(
39 url: impl AsRef<str>,
40 lease_config: ValkeyLeaseConfig,
41 ) -> Result<Self, ExecutionGuardError> {
42 Self::with_prefix(url, DEFAULT_KEY_PREFIX, lease_config).await
43 }
44
45 pub async fn with_prefix(
46 url: impl AsRef<str>,
47 key_prefix: impl Into<String>,
48 lease_config: ValkeyLeaseConfig,
49 ) -> Result<Self, ExecutionGuardError> {
50 Self::with_prefix_and_recovery_config(
51 url,
52 key_prefix,
53 lease_config,
54 ValkeyRecoveryConfig::default(),
55 )
56 .await
57 }
58
59 pub async fn with_recovery_config(
60 url: impl AsRef<str>,
61 lease_config: ValkeyLeaseConfig,
62 recovery_config: ValkeyRecoveryConfig,
63 ) -> Result<Self, ExecutionGuardError> {
64 Self::with_prefix_and_recovery_config(
65 url,
66 DEFAULT_KEY_PREFIX,
67 lease_config,
68 recovery_config,
69 )
70 .await
71 }
72
73 pub async fn with_prefix_and_recovery_config(
74 url: impl AsRef<str>,
75 key_prefix: impl Into<String>,
76 lease_config: ValkeyLeaseConfig,
77 recovery_config: ValkeyRecoveryConfig,
78 ) -> Result<Self, ExecutionGuardError> {
79 lease_config.validate()?;
80 let client = Client::open(url.as_ref()).map_err(|error| {
81 let kind = classify_redis_error(&error);
82 ExecutionGuardError::new(error, kind)
83 })?;
84 let runtime = ValkeyRuntime::from_client(client, recovery_config)
85 .await
86 .map_err(|error| {
87 let kind = classify_redis_error(&error);
88 ExecutionGuardError::new(error, kind)
89 })?;
90
91 Ok(Self {
92 runtime,
93 key_prefix: key_prefix.into(),
94 lease_config,
95 })
96 }
97
98 fn lease_key(&self, slot: &ExecutionSlot) -> String {
99 lease_key(&self.key_prefix, slot)
100 }
101
102 fn resource_lock_key(&self, resource_id: &str) -> String {
103 resource_lock_key(&self.key_prefix, resource_id)
104 }
105
106 fn occurrence_index_key(&self, resource_id: &str) -> String {
107 occurrence_index_key(&self.key_prefix, resource_id)
108 }
109
110 fn ttl_millis(&self) -> Result<u64, ValkeyExecutionGuardError> {
111 u64::try_from(self.lease_config.ttl.as_millis())
112 .map_err(ValkeyExecutionGuardError::DurationOutOfRange)
113 }
114}
115
116impl ExecutionGuard for ValkeyExecutionGuard {
117 type Error = ValkeyExecutionGuardError;
118
119 async fn acquire(&self, slot: ExecutionSlot) -> Result<ExecutionGuardAcquire, Self::Error> {
120 let lease_key = self.lease_key(&slot);
121 let token = next_token(&TOKEN_COUNTER, "lease");
122 let ttl_millis = self.ttl_millis()?;
123 let now_millis = now_millis();
124 let expires_at_millis = now_millis.saturating_add(ttl_millis);
125 let resource_lock_key = self.resource_lock_key(&slot.resource_id);
126 let occurrence_index_key = self.occurrence_index_key(&slot.resource_id);
127 let scope = slot.scope;
128 let command_lease_key = lease_key.clone();
129 let command_token = token.clone();
130 let result: ValkeyCommandOutcome<bool> = self
131 .runtime
132 .execute(move |mut connection| {
133 let lease_key = command_lease_key.clone();
134 let token = command_token.clone();
135 let resource_lock_key = resource_lock_key.clone();
136 let occurrence_index_key = occurrence_index_key.clone();
137 async move {
138 match scope {
139 ExecutionGuardScope::Occurrence => {
140 let acquired: i32 =
141 valkey_scripts::script(valkey_scripts::guard::ACQUIRE_OCCURRENCE)
142 .key(resource_lock_key)
143 .key(&lease_key)
144 .key(occurrence_index_key)
145 .arg(now_millis)
146 .arg(&token)
147 .arg(ttl_millis)
148 .arg(expires_at_millis)
149 .invoke_async(&mut connection)
150 .await?;
151 Ok(acquired == 1)
152 }
153 ExecutionGuardScope::Resource => {
154 let acquired: i32 =
155 valkey_scripts::script(valkey_scripts::guard::ACQUIRE_RESOURCE)
156 .key(&resource_lock_key)
157 .key(occurrence_index_key)
158 .arg(now_millis)
159 .arg(&token)
160 .arg(ttl_millis)
161 .invoke_async(&mut connection)
162 .await?;
163 Ok(acquired == 1)
164 }
165 }
166 }
167 })
168 .await
169 .map_err(ValkeyExecutionGuardError::Redis)?;
170
171 let ValkeyCommandOutcome::Available(acquired) = result else {
172 return Ok(ExecutionGuardAcquire::Contended);
173 };
174
175 Ok(if acquired {
176 ExecutionGuardAcquire::Acquired(ExecutionLease::new(
177 slot.job_id,
178 slot.resource_id,
179 slot.scope,
180 slot.scheduled_at,
181 token,
182 lease_key,
183 ))
184 } else {
185 ExecutionGuardAcquire::Contended
186 })
187 }
188
189 async fn renew(&self, lease: &ExecutionLease) -> Result<ExecutionGuardRenewal, Self::Error> {
190 let ttl_millis = self.ttl_millis()?;
191 let expires_at_millis = now_millis().saturating_add(ttl_millis);
192 let occurrence_index_key = self.occurrence_index_key(&lease.resource_id);
193 let lease = lease.clone();
194 let result: ValkeyCommandOutcome<i32> = self
195 .runtime
196 .execute(move |mut connection| {
197 let lease = lease.clone();
198 let occurrence_index_key = occurrence_index_key.clone();
199 async move {
200 match lease.scope {
201 ExecutionGuardScope::Occurrence => {
202 valkey_scripts::script(valkey_scripts::guard::RENEW_OCCURRENCE)
203 .key(&lease.lease_key)
204 .key(occurrence_index_key)
205 .arg(&lease.token)
206 .arg(ttl_millis)
207 .arg(expires_at_millis)
208 .invoke_async(&mut connection)
209 .await
210 }
211 ExecutionGuardScope::Resource => {
212 valkey_scripts::script(valkey_scripts::guard::RENEW_RESOURCE)
213 .key(&lease.lease_key)
214 .arg(&lease.token)
215 .arg(ttl_millis)
216 .invoke_async(&mut connection)
217 .await
218 }
219 }
220 }
221 })
222 .await
223 .map_err(ValkeyExecutionGuardError::Redis)?;
224
225 let ValkeyCommandOutcome::Available(renewed) = result else {
226 return Ok(ExecutionGuardRenewal::Lost);
227 };
228
229 Ok(if renewed == 1 {
230 ExecutionGuardRenewal::Renewed
231 } else {
232 ExecutionGuardRenewal::Lost
233 })
234 }
235
236 async fn release(&self, lease: &ExecutionLease) -> Result<(), Self::Error> {
237 let occurrence_index_key = self.occurrence_index_key(&lease.resource_id);
238 let lease = lease.clone();
239 let _: ValkeyCommandOutcome<i32> = self
240 .runtime
241 .execute(move |mut connection| {
242 let lease = lease.clone();
243 let occurrence_index_key = occurrence_index_key.clone();
244 async move {
245 match lease.scope {
246 ExecutionGuardScope::Occurrence => {
247 valkey_scripts::script(valkey_scripts::guard::RELEASE_OCCURRENCE)
248 .key(&lease.lease_key)
249 .key(occurrence_index_key)
250 .arg(&lease.token)
251 .invoke_async(&mut connection)
252 .await
253 }
254 ExecutionGuardScope::Resource => {
255 valkey_scripts::script(valkey_scripts::guard::RELEASE_RESOURCE)
256 .key(&lease.lease_key)
257 .arg(&lease.token)
258 .invoke_async(&mut connection)
259 .await
260 }
261 }
262 }
263 })
264 .await
265 .map_err(ValkeyExecutionGuardError::Redis)?;
266
267 Ok(())
268 }
269
270 async fn drain_events(&self) -> Result<Vec<ExecutionGuardEvent>, Self::Error> {
271 Ok(self
272 .runtime
273 .drain_events()
274 .await
275 .into_iter()
276 .map(|event| match event {
277 ValkeyRuntimeEvent::Degraded { error } => ExecutionGuardEvent::Degraded { error },
278 ValkeyRuntimeEvent::Recovered => ExecutionGuardEvent::Recovered,
279 })
280 .collect())
281 }
282
283 fn classify_error(error: &Self::Error) -> ExecutionGuardErrorKind
284 where
285 Self: Sized,
286 {
287 match error {
288 ValkeyExecutionGuardError::Config(_)
289 | ValkeyExecutionGuardError::DurationOutOfRange(_) => ExecutionGuardErrorKind::Data,
290 ValkeyExecutionGuardError::Redis(error) => classify_redis_error(error),
291 }
292 }
293
294 fn renew_interval(&self, _lease: &ExecutionLease) -> Option<Duration> {
295 Some(self.lease_config.renew_interval)
296 }
297}
298
299#[derive(Debug)]
300pub enum ValkeyExecutionGuardError {
301 Redis(redis::RedisError),
302 Config(LeaseConfigError),
303 DurationOutOfRange(TryFromIntError),
304}
305
306impl Display for ValkeyExecutionGuardError {
307 fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
308 match self {
309 Self::Redis(error) => write!(f, "{error}"),
310 Self::Config(error) => write!(f, "{error}"),
311 Self::DurationOutOfRange(error) => write!(f, "{error}"),
312 }
313 }
314}
315
316impl std::error::Error for ValkeyExecutionGuardError {
317 fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
318 match self {
319 Self::Redis(error) => Some(error),
320 Self::Config(error) => Some(error),
321 Self::DurationOutOfRange(error) => Some(error),
322 }
323 }
324}
325
326#[derive(Debug)]
327pub struct LeaseConfigError {
328 message: &'static str,
329}
330
331impl Display for LeaseConfigError {
332 fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
333 f.write_str(self.message)
334 }
335}
336
337impl std::error::Error for LeaseConfigError {}
338
339impl ValkeyLeaseConfig {
340 fn validate(self) -> Result<Self, ExecutionGuardError> {
341 if self.ttl.is_zero() {
342 return Err(ExecutionGuardError::new(
343 LeaseConfigError {
344 message: "execution guard ttl must be greater than zero",
345 },
346 ExecutionGuardErrorKind::Data,
347 ));
348 }
349
350 if self.renew_interval.is_zero() {
351 return Err(ExecutionGuardError::new(
352 LeaseConfigError {
353 message: "execution guard renew_interval must be greater than zero",
354 },
355 ExecutionGuardErrorKind::Data,
356 ));
357 }
358
359 if self.renew_interval >= self.ttl {
360 return Err(ExecutionGuardError::new(
361 LeaseConfigError {
362 message: "execution guard renew_interval must be less than ttl",
363 },
364 ExecutionGuardErrorKind::Data,
365 ));
366 }
367
368 if u64::try_from(self.ttl.as_millis()).is_err() {
369 return Err(ExecutionGuardError::new(
370 LeaseConfigError {
371 message: "execution guard ttl is too large to encode in milliseconds",
372 },
373 ExecutionGuardErrorKind::Data,
374 ));
375 }
376
377 Ok(self)
378 }
379}
380
381fn classify_redis_error(error: &redis::RedisError) -> ExecutionGuardErrorKind {
382 if is_connection_issue(error) {
383 ExecutionGuardErrorKind::Connection
384 } else {
385 ExecutionGuardErrorKind::Unknown
386 }
387}
388
389#[cfg(test)]
390mod tests {
391 use super::{DEFAULT_KEY_PREFIX, lease_key, next_token};
392 use crate::{ExecutionGuardScope, ExecutionSlot};
393 use chrono::{TimeZone, Utc};
394
395 #[test]
396 fn generated_tokens_are_not_empty() {
397 assert_ne!(next_token(&super::TOKEN_COUNTER, "lease"), "");
398 assert_ne!(next_token(&super::TOKEN_COUNTER, "lease"), "");
399 }
400
401 #[test]
402 fn lease_key_uses_default_prefix_and_occurrence_time() {
403 let slot = ExecutionSlot::for_occurrence(
404 "job-1",
405 "resource-1",
406 Utc.with_ymd_and_hms(2026, 4, 3, 1, 2, 3).unwrap(),
407 );
408
409 assert_eq!(
410 lease_key(DEFAULT_KEY_PREFIX, &slot),
411 "scheduler:valkey:execution-lease:resource-1:occurrence:2026-04-03T01:02:03.000000000Z"
412 );
413 }
414
415 #[test]
416 fn resource_scope_lease_key_uses_resource_lock() {
417 let slot = ExecutionSlot::for_resource("job-1", "resource-1");
418
419 assert_eq!(slot.scope, ExecutionGuardScope::Resource);
420 assert_eq!(
421 lease_key(DEFAULT_KEY_PREFIX, &slot),
422 "scheduler:valkey:execution-lease:resource-1:resource"
423 );
424 }
425}