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