vv_agent/runtime/stores/
redis_v2.rs

1//! Redis checkpoint v2 store.
2
3use std::sync::Mutex;
4use std::time::Duration;
5
6use redis::{Commands, Connection, Pipeline};
7use serde_json::Value;
8use sha2::{Digest, Sha256};
9
10use crate::checkpoint::{CheckpointError, CheckpointResult, ClaimMode, EventCursor};
11use crate::runtime::checkpoint_codec_v2::{checkpoint_v2_from_json, checkpoint_v2_to_json};
12use crate::runtime::state_v2::{
13    apply_claim, claim_candidate, prepare_ack, prepare_commit, prepare_event_delivery,
14    prepare_finalize, prepare_finalize_claimed, prepare_progress, prepare_suspend,
15    CheckpointStoreV2, CheckpointV2,
16};
17
18const KEY_PREFIX: &str = "vv-agent:checkpoint:v2:";
19const LEASE_SUFFIX: &str = ":lease";
20const IO_TIMEOUT: Duration = Duration::from_secs(1);
21const TRANSACTION_MAX_ATTEMPTS: usize = 8;
22const MAX_EXTENSION_STATE_BYTES: u64 = crate::checkpoint::MAX_WIRE_INTEGER;
23
24pub struct RedisCheckpointStoreV2 {
25    connection: Mutex<Connection>,
26    redis_url: String,
27}
28
29impl std::fmt::Debug for RedisCheckpointStoreV2 {
30    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
31        formatter
32            .debug_struct("RedisCheckpointStoreV2")
33            .field("redis_url", &self.redis_url)
34            .finish_non_exhaustive()
35    }
36}
37
38impl RedisCheckpointStoreV2 {
39    pub fn new(redis_url: impl AsRef<str>) -> CheckpointResult<Self> {
40        let redis_url = redis_url.as_ref().to_string();
41        let client = redis::Client::open(redis_url.as_str()).map_err(redis_error)?;
42        let connection = client
43            .get_connection_with_timeout(IO_TIMEOUT)
44            .map_err(redis_error)?;
45        connection
46            .set_read_timeout(Some(IO_TIMEOUT))
47            .map_err(redis_error)?;
48        connection
49            .set_write_timeout(Some(IO_TIMEOUT))
50            .map_err(redis_error)?;
51        Ok(Self {
52            connection: Mutex::new(connection),
53            redis_url,
54        })
55    }
56
57    pub fn data_key(checkpoint_key: &str) -> String {
58        let digest = Sha256::digest(checkpoint_key.as_bytes());
59        format!("{KEY_PREFIX}{digest:x}")
60    }
61
62    pub fn lease_key(checkpoint_key: &str) -> String {
63        format!("{}{LEASE_SUFFIX}", Self::data_key(checkpoint_key))
64    }
65
66    pub fn checkpoint_v2_key(checkpoint_key: &str) -> String {
67        Self::data_key(checkpoint_key)
68    }
69
70    pub fn checkpoint_v2_lease_key(checkpoint_key: &str) -> String {
71        Self::lease_key(checkpoint_key)
72    }
73
74    pub fn redis_url(&self) -> &str {
75        &self.redis_url
76    }
77
78    fn lock(&self) -> CheckpointResult<std::sync::MutexGuard<'_, Connection>> {
79        self.connection.lock().map_err(|_| {
80            CheckpointError::new(
81                "checkpoint_store_lock_poisoned",
82                "Redis store lock poisoned",
83            )
84        })
85    }
86
87    fn load_from_connection(
88        connection: &mut Connection,
89        data_key: &str,
90        lease_key: &str,
91    ) -> CheckpointResult<Option<CheckpointV2>> {
92        for _ in 0..TRANSACTION_MAX_ATTEMPTS {
93            let Some(raw) = connection
94                .get::<_, Option<String>>(data_key)
95                .map_err(redis_error)?
96            else {
97                return Ok(None);
98            };
99            let lease = connection
100                .get::<_, Option<u64>>(lease_key)
101                .map_err(redis_error)?;
102            let raw_again = connection
103                .get::<_, Option<String>>(data_key)
104                .map_err(redis_error)?;
105            if raw_again.as_deref() != Some(raw.as_str()) {
106                continue;
107            }
108            return decode_storage(&raw, lease).map(Some);
109        }
110        Err(CheckpointError::new(
111            "checkpoint_store_read_conflict",
112            "Redis checkpoint load could not obtain a stable snapshot",
113        ))
114    }
115
116    fn transaction<T>(
117        &self,
118        data_key: &str,
119        lease_key: &str,
120        operation: impl Fn(&mut Connection, &mut Pipeline) -> CheckpointResult<Option<T>>,
121    ) -> CheckpointResult<T> {
122        let mut connection = self.lock()?;
123        for _ in 0..TRANSACTION_MAX_ATTEMPTS {
124            redis::cmd("WATCH")
125                .arg(data_key)
126                .arg(lease_key)
127                .query::<()>(&mut *connection)
128                .map_err(redis_error)?;
129            let mut pipeline = redis::pipe();
130            pipeline.atomic();
131            match operation(&mut connection, &mut pipeline)? {
132                None => {
133                    redis::cmd("UNWATCH")
134                        .query::<()>(&mut *connection)
135                        .map_err(redis_error)?;
136                    return Err(CheckpointError::new(
137                        "checkpoint_store_conflict",
138                        "checkpoint operation did not match its compare-and-set precondition",
139                    ));
140                }
141                Some(value) => match pipeline.query::<Option<()>>(&mut *connection) {
142                    Ok(Some(())) => {
143                        redis::cmd("UNWATCH")
144                            .query::<()>(&mut *connection)
145                            .map_err(redis_error)?;
146                        return Ok(value);
147                    }
148                    Ok(None) => continue,
149                    Err(error) => return Err(redis_error(error)),
150                },
151            }
152        }
153        Err(CheckpointError::new(
154            "checkpoint_store_transaction_retry_exhausted",
155            "Redis checkpoint transaction retry limit exceeded",
156        ))
157    }
158}
159
160impl CheckpointStoreV2 for RedisCheckpointStoreV2 {
161    fn create_checkpoint_v2(&self, checkpoint: CheckpointV2) -> CheckpointResult<bool> {
162        checkpoint.validate()?;
163        let data_key = Self::data_key(&checkpoint.checkpoint_key);
164        let lease_key = Self::lease_key(&checkpoint.checkpoint_key);
165        let payload = checkpoint_v2_to_json(&checkpoint, MAX_EXTENSION_STATE_BYTES)?;
166        let mut connection = self.lock()?;
167        let created: bool = connection.set_nx(&data_key, payload).map_err(redis_error)?;
168        if created {
169            connection.del::<_, ()>(&lease_key).map_err(redis_error)?;
170        }
171        Ok(created)
172    }
173
174    fn load_checkpoint_v2(&self, checkpoint_key: &str) -> CheckpointResult<Option<CheckpointV2>> {
175        let data_key = Self::data_key(checkpoint_key);
176        let lease_key = Self::lease_key(checkpoint_key);
177        let mut connection = self.lock()?;
178        Self::load_from_connection(&mut connection, &data_key, &lease_key)
179    }
180
181    fn claim_checkpoint_v2(
182        &self,
183        checkpoint_key: &str,
184        cycle_index: u64,
185        claim_token: &str,
186        lease_expires_at_ms: u64,
187        now_ms: u64,
188        claim_mode: ClaimMode,
189    ) -> CheckpointResult<Option<CheckpointV2>> {
190        if claim_token.trim().is_empty() || lease_expires_at_ms <= now_ms {
191            return Err(CheckpointError::new(
192                "checkpoint_claim_invalid",
193                "claim token must be non-empty and lease must be in the future",
194            ));
195        }
196        let data_key = Self::data_key(checkpoint_key);
197        let lease_key = Self::lease_key(checkpoint_key);
198        let result = self.transaction(&data_key, &lease_key, |connection, pipeline| {
199            let Some(raw) = connection
200                .get::<_, Option<String>>(&data_key)
201                .map_err(redis_error)?
202            else {
203                return Ok(None);
204            };
205            let lease = connection
206                .get::<_, Option<u64>>(&lease_key)
207                .map_err(redis_error)?;
208            let current = decode_storage(&raw, lease)?;
209            if !claim_candidate(&current, cycle_index, now_ms, claim_mode)? {
210                return Ok(None);
211            }
212            let mut claimed = current;
213            apply_claim(
214                &mut claimed,
215                cycle_index,
216                claim_token,
217                lease_expires_at_ms,
218                claim_mode,
219            )?;
220            let payload = checkpoint_v2_to_json(&claimed, MAX_EXTENSION_STATE_BYTES)?;
221            pipeline.set(&data_key, payload).ignore();
222            pipeline.set(&lease_key, lease_expires_at_ms).ignore();
223            Ok(Some(claimed))
224        });
225        match result {
226            Ok(value) => Ok(Some(value)),
227            Err(error) if error.code() == "checkpoint_store_conflict" => Ok(None),
228            Err(error) => Err(error),
229        }
230    }
231
232    fn progress_checkpoint_v2(
233        &self,
234        checkpoint: CheckpointV2,
235        claim_token: &str,
236        expected_revision: u64,
237    ) -> CheckpointResult<bool> {
238        self.replace_claimed(
239            checkpoint,
240            claim_token,
241            expected_revision,
242            ReplaceKind::Progress,
243        )
244    }
245
246    fn suspend_checkpoint_v2(
247        &self,
248        checkpoint: CheckpointV2,
249        claim_token: &str,
250        expected_revision: u64,
251    ) -> CheckpointResult<bool> {
252        self.replace_claimed(
253            checkpoint,
254            claim_token,
255            expected_revision,
256            ReplaceKind::Suspend,
257        )
258    }
259
260    fn commit_checkpoint_v2(
261        &self,
262        checkpoint: CheckpointV2,
263        claim_token: &str,
264        expected_revision: u64,
265    ) -> CheckpointResult<bool> {
266        self.replace_claimed(
267            checkpoint,
268            claim_token,
269            expected_revision,
270            ReplaceKind::Commit,
271        )
272    }
273
274    fn finalize_claimed_v2(
275        &self,
276        checkpoint: CheckpointV2,
277        claim_token: &str,
278        expected_revision: u64,
279    ) -> CheckpointResult<bool> {
280        self.replace_claimed(
281            checkpoint,
282            claim_token,
283            expected_revision,
284            ReplaceKind::FinalizeClaimed,
285        )
286    }
287
288    fn finalize_checkpoint_v2(
289        &self,
290        checkpoint: CheckpointV2,
291        expected_revision: u64,
292    ) -> CheckpointResult<bool> {
293        let data_key = Self::data_key(&checkpoint.checkpoint_key);
294        let lease_key = Self::lease_key(&checkpoint.checkpoint_key);
295        let result = self.transaction(&data_key, &lease_key, |connection, pipeline| {
296            let Some(raw) = connection
297                .get::<_, Option<String>>(&data_key)
298                .map_err(redis_error)?
299            else {
300                return Ok(None);
301            };
302            let current = decode_storage(
303                &raw,
304                connection
305                    .get::<_, Option<u64>>(&lease_key)
306                    .map_err(redis_error)?,
307            )?;
308            let Some(updated) = prepare_finalize(&current, checkpoint.clone(), expected_revision)?
309            else {
310                return Ok(None);
311            };
312            let payload = checkpoint_v2_to_json(&updated, MAX_EXTENSION_STATE_BYTES)?;
313            pipeline.set(&data_key, payload).ignore();
314            pipeline.del(&lease_key).ignore();
315            Ok(Some(true))
316        });
317        match result {
318            Ok(value) => Ok(value),
319            Err(error) if error.code() == "checkpoint_store_conflict" => Ok(false),
320            Err(error) => Err(error),
321        }
322    }
323
324    fn renew_checkpoint_claim_v2(
325        &self,
326        checkpoint_key: &str,
327        claim_token: &str,
328        lease_expires_at_ms: u64,
329        now_ms: u64,
330    ) -> CheckpointResult<bool> {
331        if claim_token.trim().is_empty() || lease_expires_at_ms <= now_ms {
332            return Err(CheckpointError::new(
333                "checkpoint_claim_invalid",
334                "claim token must be non-empty and lease must be in the future",
335            ));
336        }
337        let data_key = Self::data_key(checkpoint_key);
338        let lease_key = Self::lease_key(checkpoint_key);
339        let result = self.transaction(&data_key, &lease_key, |connection, pipeline| {
340            let Some(raw) = connection
341                .get::<_, Option<String>>(&data_key)
342                .map_err(redis_error)?
343            else {
344                return Ok(None);
345            };
346            let current_lease = connection
347                .get::<_, Option<u64>>(&lease_key)
348                .map_err(redis_error)?;
349            let current = decode_storage(&raw, current_lease)?;
350            if current.claim_token.as_deref() != Some(claim_token)
351                || current
352                    .lease_expires_at_ms
353                    .is_none_or(|expiry| expiry <= now_ms)
354            {
355                return Ok(None);
356            }
357            pipeline.set(&lease_key, lease_expires_at_ms).ignore();
358            Ok(Some(true))
359        });
360        match result {
361            Ok(value) => Ok(value),
362            Err(error) if error.code() == "checkpoint_store_conflict" => Ok(false),
363            Err(error) => Err(error),
364        }
365    }
366
367    fn acknowledge_terminal_v2(
368        &self,
369        checkpoint_key: &str,
370        expected_revision: u64,
371    ) -> CheckpointResult<bool> {
372        let data_key = Self::data_key(checkpoint_key);
373        let lease_key = Self::lease_key(checkpoint_key);
374        let result = self.transaction(&data_key, &lease_key, |connection, pipeline| {
375            let Some(raw) = connection
376                .get::<_, Option<String>>(&data_key)
377                .map_err(redis_error)?
378            else {
379                return Ok(None);
380            };
381            let current = decode_storage(
382                &raw,
383                connection
384                    .get::<_, Option<u64>>(&lease_key)
385                    .map_err(redis_error)?,
386            )?;
387            let Some(updated) = prepare_ack(&current, expected_revision)? else {
388                return Ok(None);
389            };
390            let payload = checkpoint_v2_to_json(&updated, MAX_EXTENSION_STATE_BYTES)?;
391            pipeline.set(&data_key, payload).ignore();
392            pipeline.del(&lease_key).ignore();
393            Ok(Some(true))
394        });
395        match result {
396            Ok(value) => Ok(value),
397            Err(error) if error.code() == "checkpoint_store_conflict" => Ok(false),
398            Err(error) => Err(error),
399        }
400    }
401
402    fn record_event_delivery_v2(
403        &self,
404        checkpoint_key: &str,
405        claim_token: Option<&str>,
406        expected_revision: u64,
407        event_id: &str,
408        payload_digest: &str,
409        cursor: EventCursor,
410    ) -> CheckpointResult<bool> {
411        let data_key = Self::data_key(checkpoint_key);
412        let lease_key = Self::lease_key(checkpoint_key);
413        let result = self.transaction(&data_key, &lease_key, |connection, pipeline| {
414            let Some(raw) = connection
415                .get::<_, Option<String>>(&data_key)
416                .map_err(redis_error)?
417            else {
418                return Ok(None);
419            };
420            let current = decode_storage(
421                &raw,
422                connection
423                    .get::<_, Option<u64>>(&lease_key)
424                    .map_err(redis_error)?,
425            )?;
426            let Some(updated) = prepare_event_delivery(
427                &current,
428                claim_token,
429                expected_revision,
430                event_id,
431                payload_digest,
432                cursor.clone(),
433            )?
434            else {
435                return Ok(None);
436            };
437            let payload = checkpoint_v2_to_json(&updated, MAX_EXTENSION_STATE_BYTES)?;
438            pipeline.set(&data_key, payload).ignore();
439            if updated.claim_token.is_none() {
440                pipeline.del(&lease_key).ignore();
441            }
442            Ok(Some(true))
443        });
444        match result {
445            Ok(value) => Ok(value),
446            Err(error) if error.code() == "checkpoint_store_conflict" => Ok(false),
447            Err(error) => Err(error),
448        }
449    }
450
451    fn delete_checkpoint_v2(&self, checkpoint_key: &str) -> CheckpointResult<()> {
452        let mut connection = self.lock()?;
453        let data_key = Self::data_key(checkpoint_key);
454        let lease_key = Self::lease_key(checkpoint_key);
455        let keys = [data_key.as_str(), lease_key.as_str()];
456        let _: usize = connection.del(&keys).map_err(redis_error)?;
457        Ok(())
458    }
459
460    fn list_checkpoints_v2(&self) -> CheckpointResult<Vec<String>> {
461        let mut connection = self.lock()?;
462        let keys = connection
463            .scan_match::<_, String>(format!("{KEY_PREFIX}*"))
464            .map_err(redis_error)?
465            .filter(|key| !key.ends_with(LEASE_SUFFIX))
466            .collect::<Vec<_>>();
467        let mut checkpoint_keys = Vec::new();
468        for key in keys {
469            let Some(raw) = connection
470                .get::<_, Option<String>>(&key)
471                .map_err(redis_error)?
472            else {
473                continue;
474            };
475            let checkpoint = decode_storage(
476                &raw,
477                connection
478                    .get::<_, Option<u64>>(format!("{key}{LEASE_SUFFIX}"))
479                    .map_err(redis_error)?,
480            )?;
481            checkpoint_keys.push(checkpoint.checkpoint_key);
482        }
483        checkpoint_keys.sort();
484        Ok(checkpoint_keys)
485    }
486}
487
488impl RedisCheckpointStoreV2 {
489    fn replace_claimed(
490        &self,
491        checkpoint: CheckpointV2,
492        claim_token: &str,
493        expected_revision: u64,
494        kind: ReplaceKind,
495    ) -> CheckpointResult<bool> {
496        let data_key = Self::data_key(&checkpoint.checkpoint_key);
497        let lease_key = Self::lease_key(&checkpoint.checkpoint_key);
498        let result = self.transaction(&data_key, &lease_key, |connection, pipeline| {
499            let Some(raw) = connection
500                .get::<_, Option<String>>(&data_key)
501                .map_err(redis_error)?
502            else {
503                return Ok(None);
504            };
505            let current = decode_storage(
506                &raw,
507                connection
508                    .get::<_, Option<u64>>(&lease_key)
509                    .map_err(redis_error)?,
510            )?;
511            let updated = match kind {
512                ReplaceKind::Progress => {
513                    prepare_progress(&current, checkpoint.clone(), claim_token, expected_revision)?
514                }
515                ReplaceKind::Suspend => {
516                    prepare_suspend(&current, checkpoint.clone(), claim_token, expected_revision)?
517                }
518                ReplaceKind::Commit => {
519                    prepare_commit(&current, checkpoint.clone(), claim_token, expected_revision)?
520                }
521                ReplaceKind::FinalizeClaimed => prepare_finalize_claimed(
522                    &current,
523                    checkpoint.clone(),
524                    claim_token,
525                    expected_revision,
526                )?,
527            };
528            let Some(updated) = updated else {
529                return Ok(None);
530            };
531            let payload = checkpoint_v2_to_json(&updated, MAX_EXTENSION_STATE_BYTES)?;
532            pipeline.set(&data_key, payload).ignore();
533            if updated.claim_token.is_none() {
534                pipeline.del(&lease_key).ignore();
535            }
536            Ok(Some(true))
537        });
538        match result {
539            Ok(value) => Ok(value),
540            Err(error) if error.code() == "checkpoint_store_conflict" => Ok(false),
541            Err(error) => Err(error),
542        }
543    }
544}
545
546#[derive(Clone, Copy)]
547enum ReplaceKind {
548    Progress,
549    Suspend,
550    Commit,
551    FinalizeClaimed,
552}
553
554fn decode_storage(raw: &str, lease: Option<u64>) -> CheckpointResult<CheckpointV2> {
555    let mut value: Value = serde_json::from_str(raw)
556        .map_err(|error| CheckpointError::new("checkpoint_json_invalid", error.to_string()))?;
557    if let Some(object) = value.as_object_mut() {
558        object.insert(
559            "lease_expires_at_ms".to_string(),
560            lease.map_or(Value::Null, Value::from),
561        );
562    }
563    let payload = serde_json::to_string(&value)
564        .map_err(|error| CheckpointError::new("checkpoint_json_invalid", error.to_string()))?;
565    checkpoint_v2_from_json(&payload, MAX_EXTENSION_STATE_BYTES)
566}
567
568fn redis_error(error: redis::RedisError) -> CheckpointError {
569    CheckpointError::new("checkpoint_store_redis", error.to_string())
570}
571
572pub type RedisStateStoreV2 = RedisCheckpointStoreV2;
vv_agent/runtime/stores/redis_v2.rs

vv_agent/runtime/stores/
redis_v2.rs
