vv_agent/runtime/stores/
redis.rs

1use std::io::{Error, Result};
2use std::sync::Mutex;
3
4use redis::{Commands, Connection};
5
6use crate::runtime::checkpoint_codec;
7use crate::runtime::state::{Checkpoint, StateStore};
8
9const KEY_PREFIX: &str = "vv_agent:checkpoint:";
10
11pub struct RedisStateStore {
12    connection: Mutex<Connection>,
13}
14
15impl std::fmt::Debug for RedisStateStore {
16    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
17        formatter
18            .debug_struct("RedisStateStore")
19            .finish_non_exhaustive()
20    }
21}
22
23impl RedisStateStore {
24    pub fn new(redis_url: impl AsRef<str>) -> Result<Self> {
25        let client = redis::Client::open(redis_url.as_ref()).map_err(redis_to_io)?;
26        let connection = client.get_connection().map_err(redis_to_io)?;
27        Ok(Self {
28            connection: Mutex::new(connection),
29        })
30    }
31
32    pub fn checkpoint_key(task_id: &str) -> String {
33        format!("{KEY_PREFIX}{task_id}")
34    }
35
36    pub fn checkpoint_to_json(checkpoint: &Checkpoint) -> Result<String> {
37        checkpoint_codec::checkpoint_to_json(checkpoint)
38    }
39
40    pub fn checkpoint_from_json(raw: &str) -> Result<Checkpoint> {
41        checkpoint_codec::checkpoint_from_json(raw)
42    }
43}
44
45impl StateStore for RedisStateStore {
46    fn save_checkpoint(&self, checkpoint: Checkpoint) -> Result<()> {
47        let key = Self::checkpoint_key(&checkpoint.task_id);
48        let payload = Self::checkpoint_to_json(&checkpoint)?;
49        self.connection
50            .lock()
51            .map_err(|_| Error::other("redis state store lock is poisoned"))?
52            .set::<_, _, ()>(key, payload)
53            .map_err(redis_to_io)
54    }
55
56    fn load_checkpoint(&self, task_id: &str) -> Result<Option<Checkpoint>> {
57        let key = Self::checkpoint_key(task_id);
58        let raw = self
59            .connection
60            .lock()
61            .map_err(|_| Error::other("redis state store lock is poisoned"))?
62            .get::<_, Option<String>>(key)
63            .map_err(redis_to_io)?;
64        raw.as_deref().map(Self::checkpoint_from_json).transpose()
65    }
66
67    fn delete_checkpoint(&self, task_id: &str) -> Result<()> {
68        let key = Self::checkpoint_key(task_id);
69        self.connection
70            .lock()
71            .map_err(|_| Error::other("redis state store lock is poisoned"))?
72            .del::<_, ()>(key)
73            .map_err(redis_to_io)
74    }
75
76    fn list_checkpoints(&self) -> Result<Vec<String>> {
77        let pattern = format!("{KEY_PREFIX}*");
78        let mut connection = self
79            .connection
80            .lock()
81            .map_err(|_| Error::other("redis state store lock is poisoned"))?;
82        let iter = connection
83            .scan_match::<_, String>(pattern)
84            .map_err(redis_to_io)?;
85        let mut keys = iter
86            .map(|key| key.strip_prefix(KEY_PREFIX).unwrap_or(&key).to_string())
87            .collect::<Vec<_>>();
88        keys.sort();
89        Ok(keys)
90    }
91}
92
93fn redis_to_io(error: redis::RedisError) -> Error {
94    Error::other(error.to_string())
95}
vv_agent/runtime/stores/redis.rs

vv_agent/runtime/stores/
redis.rs
