chronon_backend_redis/
queue.rs

1//! Redis sorted-set ready queue for run claims.
2
3use std::sync::Arc;
4
5use chrono::{DateTime, Utc};
6use chronon_core::error::ChrononError;
7use chronon_core::Result;
8use redis::aio::ConnectionManager;
9use redis::AsyncCommands;
10use tokio::sync::Mutex;
11
12/// Redis ZSET layer for queued runs (`{prefix}:ready:{pool}` or hash-tagged for cluster).
13///
14/// Used only as the claim hot path inside [`crate::PostgresRedisSchedulerStore`] — not a
15/// standalone [`SchedulerStore`](chronon_core::SchedulerStore). Connect with [`Self::connect`]:
16///
17/// | Knob | Effect |
18/// |------|--------|
19/// | `url` / `CHRONON_REDIS_URL` | Standalone Redis |
20/// | `CHRONON_REDIS_CLUSTER_URLS` | Comma-separated cluster nodes |
21/// | `CHRONON_REDIS_HASH_TAGS=1` | Hash-tagged pool keys for slot affinity |
22/// | `key_prefix` (default `chronon`) | Key namespace |
23///
24/// # Examples
25///
26/// ```no_run
27/// use chronon_backend_redis::RedisQueueLayer;
28///
29/// # async fn demo() -> chronon_core::Result<()> {
30/// let redis = RedisQueueLayer::connect("redis://127.0.0.1:6379", Some("myapp")).await?;
31/// # let _ = redis;
32/// # Ok(())
33/// # }
34/// ```
35pub struct RedisQueueLayer {
36    single: Option<ConnectionManager>,
37    cluster: Option<Arc<Mutex<redis::cluster_async::ClusterConnection>>>,
38    key_prefix: String,
39    hash_tags: bool,
40}
41
42impl std::fmt::Debug for RedisQueueLayer {
43    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
44        f.debug_struct("RedisQueueLayer")
45            .field("key_prefix", &self.key_prefix)
46            .field("hash_tags", &self.hash_tags)
47            .field("cluster", &self.cluster.is_some())
48            .finish_non_exhaustive()
49    }
50}
51
52impl RedisQueueLayer {
53    /// Connect to Redis at `url` with optional key prefix (default `chronon`).
54    ///
55    /// When `CHRONON_REDIS_CLUSTER_URLS` is set (comma-separated), uses Redis Cluster.
56    /// When `CHRONON_REDIS_HASH_TAGS=1`, pool keys use hash tags for cluster slot affinity.
57    ///
58    /// # Errors
59    ///
60    /// Returns a storage error when the connection cannot be established.
61    pub async fn connect(url: &str, key_prefix: Option<&str>) -> Result<Self> {
62        let hash_tags = std::env::var("CHRONON_REDIS_HASH_TAGS")
63            .is_ok_and(|v| matches!(v.as_str(), "1" | "true" | "yes"));
64        let prefix = key_prefix.unwrap_or("chronon").to_string();
65
66        if let Ok(urls) = std::env::var("CHRONON_REDIS_CLUSTER_URLS") {
67            let label = chronon_core::redact_endpoint(&urls);
68            let nodes: Vec<String> = urls
69                .split(',')
70                .map(str::trim)
71                .filter(|s| !s.is_empty())
72                .map(String::from)
73                .collect();
74            if !nodes.is_empty() {
75                let client = redis::cluster::ClusterClient::new(nodes)
76                    .map_err(|e| map_connect_err(&label, e))?;
77                let conn = client
78                    .get_async_connection()
79                    .await
80                    .map_err(|e| map_connect_err(&label, e))?;
81                return Ok(Self {
82                    single: None,
83                    cluster: Some(Arc::new(Mutex::new(conn))),
84                    key_prefix: prefix,
85                    hash_tags: true,
86                });
87            }
88        }
89
90        let client = redis::Client::open(url).map_err(|e| map_connect_err(url, e))?;
91        let conn = ConnectionManager::new(client)
92            .await
93            .map_err(|e| map_connect_err(url, e))?;
94        Ok(Self {
95            single: Some(conn),
96            cluster: None,
97            key_prefix: prefix,
98            hash_tags,
99        })
100    }
101
102    /// Redis URL for tests (`CHRONON_TEST_REDIS_URL`, then `CHRONON_REDIS_URL`, or local default).
103    #[must_use]
104    pub fn test_url() -> String {
105        std::env::var("CHRONON_TEST_REDIS_URL")
106            .or_else(|_| std::env::var("CHRONON_REDIS_URL"))
107            .unwrap_or_else(|_| "redis://127.0.0.1:6379".into())
108    }
109
110    fn ready_key(&self, pool_id: &str) -> String {
111        format_ready_key(&self.key_prefix, pool_id, self.hash_tags)
112    }
113
114    /// Enqueue a run id ordered by `scheduled_for` (ZADD score = epoch millis).
115    ///
116    /// # Errors
117    ///
118    /// Returns a storage error when Redis commands fail.
119    pub async fn enqueue_run(
120        &self,
121        pool_id: &str,
122        run_id: &str,
123        scheduled_for: DateTime<Utc>,
124    ) -> Result<()> {
125        let score = scheduled_for.timestamp_millis() as f64;
126        let key = self.ready_key(pool_id);
127        if let Some(conn) = &self.single {
128            let mut conn = conn.clone();
129            let _: () = conn.zadd(key, run_id, score).await.map_err(map_err)?;
130        } else if let Some(conn) = &self.cluster {
131            let mut conn = conn.lock().await;
132            let _: () = conn.zadd(key, run_id, score).await.map_err(map_err)?;
133        }
134        Ok(())
135    }
136
137    /// Atomically pop the earliest **due** run id (`score <= now`) from the pool queue.
138    ///
139    /// # Errors
140    ///
141    /// Returns a storage error when Redis commands fail.
142    pub async fn claim_next_run_id(
143        &self,
144        pool_id: &str,
145        now: DateTime<Utc>,
146    ) -> Result<Option<String>> {
147        let ids = self.claim_next_run_ids(pool_id, 1, now).await?;
148        Ok(ids.into_iter().next())
149    }
150
151    /// Pop up to `count` earliest due run ids (`scheduled_for` score ≤ `now`).
152    ///
153    /// # Errors
154    ///
155    /// Returns a storage error when Redis commands fail.
156    pub async fn claim_next_run_ids(
157        &self,
158        pool_id: &str,
159        count: usize,
160        now: DateTime<Utc>,
161    ) -> Result<Vec<String>> {
162        if count == 0 {
163            return Ok(Vec::new());
164        }
165        let key = self.ready_key(pool_id);
166        let max_score = now.timestamp_millis() as f64;
167        // Atomic: take due members then remove them from the ZSET.
168        let script = redis::Script::new(
169            r"
170            local ids = redis.call('ZRANGEBYSCORE', KEYS[1], '-inf', ARGV[1], 'LIMIT', 0, tonumber(ARGV[2]))
171            for _, id in ipairs(ids) do
172              redis.call('ZREM', KEYS[1], id)
173            end
174            return ids
175            ",
176        );
177        let ids: Vec<String> = if let Some(conn) = &self.single {
178            let mut conn = conn.clone();
179            script
180                .key(&key)
181                .arg(max_score)
182                .arg(count)
183                .invoke_async(&mut conn)
184                .await
185                .map_err(map_err)?
186        } else if let Some(conn) = &self.cluster {
187            let mut conn = conn.lock().await;
188            script
189                .key(&key)
190                .arg(max_score)
191                .arg(count)
192                .invoke_async(&mut *conn)
193                .await
194                .map_err(map_err)?
195        } else {
196            Vec::new()
197        };
198        Ok(ids)
199    }
200
201    /// Remove a run from the ready queue (e.g. after cancellation).
202    ///
203    /// # Errors
204    ///
205    /// Returns a storage error when Redis commands fail.
206    pub async fn remove_run(&self, pool_id: &str, run_id: &str) -> Result<()> {
207        let key = self.ready_key(pool_id);
208        if let Some(conn) = &self.single {
209            let mut conn = conn.clone();
210            let _: () = conn.zrem(key, run_id).await.map_err(map_err)?;
211        } else if let Some(conn) = &self.cluster {
212            let mut conn = conn.lock().await;
213            let _: () = conn.zrem(key, run_id).await.map_err(map_err)?;
214        }
215        Ok(())
216    }
217
218    /// Delete all keys with this layer's prefix (test isolation).
219    ///
220    /// # Errors
221    ///
222    /// Returns a storage error when Redis commands fail.
223    pub async fn flush_keys(&self) -> Result<()> {
224        if self.cluster.is_some() {
225            return Ok(());
226        }
227        let Some(single) = &self.single else {
228            return Ok(());
229        };
230        let pattern = format!("{}:*", self.key_prefix);
231        let mut conn = single.clone();
232        let mut cursor = 0_u64;
233        loop {
234            let (next, batch): (u64, Vec<String>) = redis::cmd("SCAN")
235                .arg(cursor)
236                .arg("MATCH")
237                .arg(&pattern)
238                .arg("COUNT")
239                .arg(500)
240                .query_async(&mut conn)
241                .await
242                .map_err(map_err)?;
243            if !batch.is_empty() {
244                let _: () = conn.del(batch).await.map_err(map_err)?;
245            }
246            cursor = next;
247            if cursor == 0 {
248                break;
249            }
250        }
251        Ok(())
252    }
253}
254
255fn format_ready_key(prefix: &str, pool_id: &str, hash_tags: bool) -> String {
256    if hash_tags {
257        format!("{prefix}:{{{pool_id}}}:ready")
258    } else {
259        format!("{prefix}:ready:{pool_id}")
260    }
261}
262
263fn map_err(e: impl std::error::Error + Send + Sync + 'static) -> ChrononError {
264    let message = chronon_core::redact_credentials_in_text(&e.to_string());
265    ChrononError::storage_source(message, e)
266}
267
268fn map_connect_err(url: &str, e: impl std::error::Error + Send + Sync + 'static) -> ChrononError {
269    let detail = chronon_core::redact_credentials_in_text(&e.to_string());
270    ChrononError::storage_source(
271        format!(
272            "redis connect {}: {detail}",
273            chronon_core::redact_endpoint(url)
274        ),
275        e,
276    )
277}
278
279#[cfg(test)]
280mod tests {
281    use chrono::Utc;
282
283    use super::{format_ready_key, RedisQueueLayer};
284
285    #[test]
286    fn ready_key_plain_and_hash_tagged() {
287        assert_eq!(
288            format_ready_key("chronon_test", "workers", false),
289            "chronon_test:ready:workers"
290        );
291        assert_eq!(
292            format_ready_key("chronon_test", "general-0", true),
293            "chronon_test:{general-0}:ready"
294        );
295    }
296
297    async fn layer() -> Option<RedisQueueLayer> {
298        let url = RedisQueueLayer::test_url();
299        let prefix = format!("chronon_test_{}", uuid_like());
300        let connect = RedisQueueLayer::connect(&url, Some(&prefix));
301        let layer = tokio::time::timeout(std::time::Duration::from_secs(2), connect)
302            .await
303            .ok()?
304            .ok()?;
305        tokio::time::timeout(std::time::Duration::from_secs(2), layer.flush_keys())
306            .await
307            .ok()?
308            .ok()?;
309        Some(layer)
310    }
311
312    fn uuid_like() -> String {
313        std::time::SystemTime::now()
314            .duration_since(std::time::UNIX_EPOCH)
315            .map_or_else(|_| "0".into(), |d| d.as_nanos().to_string())
316    }
317
318    #[tokio::test]
319    async fn enqueue_and_claim_orders_by_scheduled_for() {
320        let Some(layer) = layer().await else {
321            return;
322        };
323        let pool = "workers-order";
324        let now = Utc::now();
325        layer
326            .enqueue_run(pool, "run-late", now)
327            .await
328            .expect("enqueue");
329        layer
330            .enqueue_run(pool, "run-early", now - chrono::Duration::minutes(1))
331            .await
332            .expect("enqueue");
333
334        let first = layer.claim_next_run_id(pool, now).await.expect("claim");
335        assert_eq!(first.as_deref(), Some("run-early"));
336
337        let second = layer.claim_next_run_id(pool, now).await.expect("claim");
338        assert_eq!(second.as_deref(), Some("run-late"));
339
340        let empty = layer.claim_next_run_id(pool, now).await.expect("claim");
341        assert!(empty.is_none());
342    }
343
344    #[tokio::test]
345    async fn claim_skips_future_scheduled_for() {
346        let Some(layer) = layer().await else {
347            return;
348        };
349        let pool = "workers-future";
350        let now = Utc::now();
351        layer
352            .enqueue_run(pool, "run-future", now + chrono::Duration::minutes(5))
353            .await
354            .expect("enqueue");
355        let empty = layer.claim_next_run_id(pool, now).await.expect("claim");
356        assert!(empty.is_none());
357    }
358
359    #[tokio::test]
360    async fn claim_next_run_ids_batch() {
361        let Some(layer) = layer().await else {
362            return;
363        };
364        let pool = "workers-batch";
365        let now = Utc::now();
366        for i in 0..3_u64 {
367            layer
368                .enqueue_run(pool, &format!("run-{i}"), now)
369                .await
370                .expect("enqueue");
371        }
372        let batch = layer.claim_next_run_ids(pool, 2, now).await.expect("batch");
373        assert_eq!(batch.len(), 2);
374    }
375}
chronon_backend_redis/queue.rs

chronon_backend_redis/
queue.rs
