use std::io::{Error, ErrorKind, Result};
use std::sync::Mutex;
use std::time::Duration;
use redis::{cmd, pipe, Commands, Connection, ConnectionLike, Pipeline, RedisResult};
use crate::runtime::checkpoint_codec;
use crate::runtime::state::{
check_claim, claim_matches, clear_claim, validate_claim, validate_renew, Checkpoint,
LeaseOperationClock, StateStore, StateStoreSpec,
};
const KEY_PREFIX: &str = "vv_agent:checkpoint:";
const IO_TIMEOUT: Duration = Duration::from_secs(1);
const TRANSACTION_MAX_ATTEMPTS: usize = 8;
const RENEW_CLAIM_SCRIPT: &str = r#"
local redis_time = redis.call("TIME")
local server_now_ms = tonumber(redis_time[1]) * 1000 + math.floor(tonumber(redis_time[2]) / 1000)
local client_now_ms = tonumber(ARGV[5])
local current_now_ms = math.max(server_now_ms, client_now_ms)
local previous_expiry_ms = tonumber(ARGV[3])
local requested_expiry_ms = tonumber(ARGV[4])
if previous_expiry_ms <= current_now_ms or requested_expiry_ms <= current_now_ms then
return 2
end
local current = redis.call("GET", KEYS[1])
if current ~= ARGV[1] then
return 0
end
redis.call("SET", KEYS[1], ARGV[2])
return 1
"#;
pub struct RedisStateStore {
connection: Mutex<Connection>,
redis_url: String,
}
impl std::fmt::Debug for RedisStateStore {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
formatter
.debug_struct("RedisStateStore")
.finish_non_exhaustive()
}
}
impl RedisStateStore {
pub fn new(redis_url: impl AsRef<str>) -> Result<Self> {
let redis_url = redis_url.as_ref().to_string();
let client = redis::Client::open(redis_url.as_str()).map_err(redis_to_io)?;
let connection = client
.get_connection_with_timeout(IO_TIMEOUT)
.map_err(redis_to_io)?;
connection
.set_read_timeout(Some(IO_TIMEOUT))
.map_err(redis_to_io)?;
connection
.set_write_timeout(Some(IO_TIMEOUT))
.map_err(redis_to_io)?;
Ok(Self {
connection: Mutex::new(connection),
redis_url,
})
}
pub fn checkpoint_key(task_id: &str) -> String {
format!("{KEY_PREFIX}{task_id}")
}
pub fn checkpoint_to_json(checkpoint: &Checkpoint) -> Result<String> {
checkpoint_codec::checkpoint_to_json(checkpoint)
}
pub fn checkpoint_from_json(raw: &str) -> Result<Checkpoint> {
checkpoint_codec::checkpoint_from_json(raw)
}
fn transaction<T>(
&self,
key: &str,
operation: impl FnMut(&mut Connection, &mut redis::Pipeline) -> redis::RedisResult<Option<T>>,
) -> Result<T> {
let mut connection = self
.connection
.lock()
.map_err(|_| Error::other("redis state store lock is poisoned"))?;
transaction_with_connection(&mut *connection, key, operation)
}
}
fn transaction_with_connection<C, T>(
connection: &mut C,
key: &str,
mut operation: impl FnMut(&mut C, &mut Pipeline) -> RedisResult<Option<T>>,
) -> Result<T>
where
C: ConnectionLike,
{
for _attempt in 0..TRANSACTION_MAX_ATTEMPTS {
cmd("WATCH")
.arg(key)
.exec(&mut *connection)
.map_err(redis_to_io)?;
let mut pipeline = pipe();
if let Some(result) = operation(connection, pipeline.atomic()).map_err(redis_to_io)? {
cmd("UNWATCH").exec(&mut *connection).map_err(redis_to_io)?;
return Ok(result);
}
}
cmd("UNWATCH").exec(&mut *connection).map_err(redis_to_io)?;
Err(Error::new(
ErrorKind::TimedOut,
"redis checkpoint transaction retry limit exceeded",
))
}
fn query_committed_transaction<C, T>(
connection: &mut C,
pipeline: &Pipeline,
result: T,
) -> RedisResult<Option<T>>
where
C: ConnectionLike,
{
pipeline
.query::<Option<()>>(connection)
.map(|committed| committed.map(|()| result))
}
fn renew_checkpoint_claim_with_connection<C>(
connection: &mut C,
key: &str,
claim_token: &str,
expected_revision: u64,
lease_expires_at_ms: u64,
clock: &LeaseOperationClock,
) -> Result<bool>
where
C: ConnectionLike,
{
let Some(raw) = connection
.get::<_, Option<String>>(key)
.map_err(redis_to_io)?
else {
return Ok(false);
};
let mut checkpoint = RedisStateStore::checkpoint_from_json(&raw)?;
let current_now_ms = clock.now_ms();
if checkpoint.revision != expected_revision
|| checkpoint.claim_token.as_deref() != Some(claim_token)
|| checkpoint.lease_expires_at_ms.unwrap_or(0) <= current_now_ms
|| lease_expires_at_ms <= current_now_ms
{
return Ok(false);
}
let previous_lease_expires_at_ms = checkpoint
.lease_expires_at_ms
.expect("validated lease expiry must be present");
checkpoint.lease_expires_at_ms = Some(lease_expires_at_ms);
let payload = RedisStateStore::checkpoint_to_json(&checkpoint)?;
let result = cmd("EVAL")
.arg(RENEW_CLAIM_SCRIPT)
.arg(1)
.arg(key)
.arg(&raw)
.arg(payload)
.arg(previous_lease_expires_at_ms)
.arg(lease_expires_at_ms)
.arg(clock.now_ms())
.query::<i64>(connection)
.map_err(redis_to_io)?;
match result {
0 => Ok(false),
1 => Ok(true),
2 => Err(Error::new(ErrorKind::TimedOut, "claim lease expired")),
unexpected => Err(Error::new(
ErrorKind::InvalidData,
format!("redis checkpoint renewal returned unexpected result: {unexpected}"),
)),
}
}
impl StateStore for RedisStateStore {
fn create_checkpoint(&self, checkpoint: Checkpoint) -> Result<bool> {
let key = Self::checkpoint_key(&checkpoint.task_id);
let payload = Self::checkpoint_to_json(&checkpoint)?;
self.connection
.lock()
.map_err(|_| Error::other("redis state store lock is poisoned"))?
.set_nx(key, payload)
.map_err(redis_to_io)
}
fn save_checkpoint(&self, checkpoint: Checkpoint) -> Result<()> {
let key = Self::checkpoint_key(&checkpoint.task_id);
let payload = Self::checkpoint_to_json(&checkpoint)?;
self.connection
.lock()
.map_err(|_| Error::other("redis state store lock is poisoned"))?
.set::<_, _, ()>(key, payload)
.map_err(redis_to_io)
}
fn load_checkpoint(&self, task_id: &str) -> Result<Option<Checkpoint>> {
let key = Self::checkpoint_key(task_id);
let raw = self
.connection
.lock()
.map_err(|_| Error::other("redis state store lock is poisoned"))?
.get::<_, Option<String>>(key)
.map_err(redis_to_io)?;
raw.as_deref().map(Self::checkpoint_from_json).transpose()
}
fn claim_checkpoint(
&self,
task_id: &str,
cycle_index: u32,
claim_token: &str,
lease_expires_at_ms: u64,
now_ms: u64,
) -> Result<Option<Checkpoint>> {
validate_claim(cycle_index, claim_token, lease_expires_at_ms, now_ms)?;
let key = Self::checkpoint_key(task_id);
self.transaction(&key, |connection, pipe| {
let Some(raw) = connection.get::<_, Option<String>>(&key)? else {
return Ok(Some(None));
};
let mut checkpoint = Self::checkpoint_from_json(&raw).map_err(io_to_redis)?;
check_claim(&checkpoint, cycle_index, now_ms).map_err(io_to_redis)?;
checkpoint.revision = checkpoint.revision.checked_add(1).ok_or_else(|| {
io_to_redis(Error::new(
ErrorKind::InvalidData,
"checkpoint revision overflow",
))
})?;
checkpoint.claim_token = Some(claim_token.to_string());
checkpoint.claimed_cycle = Some(cycle_index);
checkpoint.lease_expires_at_ms = Some(lease_expires_at_ms);
let payload = Self::checkpoint_to_json(&checkpoint).map_err(io_to_redis)?;
pipe.set(&key, payload).ignore();
query_committed_transaction(connection, pipe, Some(checkpoint))
})
}
fn commit_checkpoint(
&self,
mut checkpoint: Checkpoint,
claim_token: &str,
expected_revision: u64,
) -> Result<bool> {
let key = Self::checkpoint_key(&checkpoint.task_id);
self.transaction(&key, |connection, pipe| {
let current = connection.get::<_, Option<String>>(&key)?;
let current = current
.as_deref()
.map(Self::checkpoint_from_json)
.transpose()
.map_err(io_to_redis)?;
if !claim_matches(
current.as_ref(),
&checkpoint,
claim_token,
expected_revision,
) {
return Ok(Some(false));
}
checkpoint.revision = expected_revision.checked_add(1).ok_or_else(|| {
io_to_redis(Error::new(
ErrorKind::InvalidData,
"checkpoint revision overflow",
))
})?;
clear_claim(&mut checkpoint);
let payload = Self::checkpoint_to_json(&checkpoint).map_err(io_to_redis)?;
pipe.set(&key, payload).ignore();
query_committed_transaction(connection, pipe, true)
})
}
fn renew_checkpoint_claim(
&self,
task_id: &str,
claim_token: &str,
expected_revision: u64,
lease_expires_at_ms: u64,
now_ms: u64,
) -> Result<bool> {
validate_renew(claim_token, expected_revision, lease_expires_at_ms, now_ms)?;
let clock = LeaseOperationClock::new(now_ms);
let key = Self::checkpoint_key(task_id);
let mut connection = self
.connection
.lock()
.map_err(|_| Error::other("redis state store lock is poisoned"))?;
renew_checkpoint_claim_with_connection(
&mut *connection,
&key,
claim_token,
expected_revision,
lease_expires_at_ms,
&clock,
)
}
fn finalize_checkpoint(
&self,
mut checkpoint: Checkpoint,
expected_revision: u64,
) -> Result<bool> {
if checkpoint.terminal_result.is_none() {
return Err(Error::new(
ErrorKind::InvalidInput,
"finalized checkpoint must include terminal_result",
));
}
let key = Self::checkpoint_key(&checkpoint.task_id);
self.transaction(&key, |connection, pipe| {
let current = connection.get::<_, Option<String>>(&key)?;
let current = current
.as_deref()
.map(Self::checkpoint_from_json)
.transpose()
.map_err(io_to_redis)?;
if current.as_ref().is_none_or(|current| {
current.revision != expected_revision
|| current.claim_token.is_some()
|| current.terminal_result.is_some()
}) {
return Ok(Some(false));
}
checkpoint.revision = expected_revision.checked_add(1).ok_or_else(|| {
io_to_redis(Error::new(
ErrorKind::InvalidData,
"checkpoint revision overflow",
))
})?;
clear_claim(&mut checkpoint);
let payload = Self::checkpoint_to_json(&checkpoint).map_err(io_to_redis)?;
pipe.set(&key, payload).ignore();
query_committed_transaction(connection, pipe, true)
})
}
fn delete_checkpoint(&self, task_id: &str) -> Result<()> {
let key = Self::checkpoint_key(task_id);
self.connection
.lock()
.map_err(|_| Error::other("redis state store lock is poisoned"))?
.del::<_, ()>(key)
.map_err(redis_to_io)
}
fn acknowledge_terminal(&self, task_id: &str, expected_revision: u64) -> Result<bool> {
let key = Self::checkpoint_key(task_id);
self.transaction(&key, |connection, pipe| {
let current = connection.get::<_, Option<String>>(&key)?;
let matches = current
.as_deref()
.map(Self::checkpoint_from_json)
.transpose()
.map_err(io_to_redis)?
.is_some_and(|checkpoint| {
checkpoint.revision == expected_revision && checkpoint.terminal_result.is_some()
});
if !matches {
return Ok(Some(false));
}
pipe.del(&key).ignore();
query_committed_transaction(connection, pipe, true)
})
}
fn list_checkpoints(&self) -> Result<Vec<String>> {
let pattern = format!("{KEY_PREFIX}*");
let mut connection = self
.connection
.lock()
.map_err(|_| Error::other("redis state store lock is poisoned"))?;
let iter = connection
.scan_match::<_, String>(pattern)
.map_err(redis_to_io)?;
let mut keys = iter
.map(|key| key.strip_prefix(KEY_PREFIX).unwrap_or(&key).to_string())
.collect::<Vec<_>>();
keys.sort();
Ok(keys)
}
fn state_store_spec(&self) -> Option<StateStoreSpec> {
StateStoreSpec::redis(&self.redis_url).ok()
}
}
fn redis_to_io(error: redis::RedisError) -> Error {
Error::other(error.to_string())
}
fn io_to_redis(error: Error) -> redis::RedisError {
redis::RedisError::from((
redis::ErrorKind::TypeError,
"checkpoint state error",
error.to_string(),
))
}
#[cfg(test)]
mod tests {
use super::*;
use std::collections::VecDeque;
struct ScriptedConnection {
exec_responses: VecDeque<redis::Value>,
watch_calls: usize,
unwatch_calls: usize,
exec_calls: usize,
}
impl ScriptedConnection {
fn new(exec_responses: impl IntoIterator<Item = redis::Value>) -> Self {
Self {
exec_responses: exec_responses.into_iter().collect(),
watch_calls: 0,
unwatch_calls: 0,
exec_calls: 0,
}
}
}
impl ConnectionLike for ScriptedConnection {
fn req_packed_command(&mut self, command: &[u8]) -> RedisResult<redis::Value> {
if command
.windows(b"UNWATCH".len())
.any(|window| window == b"UNWATCH")
{
self.unwatch_calls += 1;
} else if command
.windows(b"WATCH".len())
.any(|window| window == b"WATCH")
{
self.watch_calls += 1;
} else {
panic!("unexpected standalone Redis command");
}
Ok(redis::Value::Okay)
}
fn req_packed_commands(
&mut self,
_command: &[u8],
_offset: usize,
_count: usize,
) -> RedisResult<Vec<redis::Value>> {
self.exec_calls += 1;
let response = self
.exec_responses
.pop_front()
.expect("scripted EXEC response");
Ok(vec![response])
}
fn get_db(&self) -> i64 {
0
}
fn check_connection(&mut self) -> bool {
true
}
fn is_open(&self) -> bool {
true
}
}
struct AtomicRenewConnection {
value: Option<String>,
replace_before_eval: Option<String>,
server_now_ms: u64,
get_calls: usize,
eval_calls: usize,
}
impl AtomicRenewConnection {
fn new(value: String, server_now_ms: u64) -> Self {
Self {
value: Some(value),
replace_before_eval: None,
server_now_ms,
get_calls: 0,
eval_calls: 0,
}
}
fn replace_before_eval(mut self, value: String) -> Self {
self.replace_before_eval = Some(value);
self
}
}
impl ConnectionLike for AtomicRenewConnection {
fn req_packed_command(&mut self, command: &[u8]) -> RedisResult<redis::Value> {
let arguments = packed_command_arguments(command);
match arguments.first().map(Vec::as_slice) {
Some(b"GET") => {
self.get_calls += 1;
Ok(self
.value
.as_ref()
.map(|value| redis::Value::BulkString(value.as_bytes().to_vec()))
.unwrap_or(redis::Value::Nil))
}
Some(b"EVAL") => {
self.eval_calls += 1;
assert_eq!(arguments[1], RENEW_CLAIM_SCRIPT.as_bytes());
assert_eq!(arguments[2], b"1");
if let Some(replacement) = self.replace_before_eval.take() {
self.value = Some(replacement);
}
let expected = std::str::from_utf8(&arguments[4]).expect("expected payload");
let updated = std::str::from_utf8(&arguments[5]).expect("updated payload");
let previous_expiry = parse_u64_argument(&arguments[6]);
let requested_expiry = parse_u64_argument(&arguments[7]);
let client_now = parse_u64_argument(&arguments[8]);
let current_now = self.server_now_ms.max(client_now);
if previous_expiry <= current_now || requested_expiry <= current_now {
return Ok(redis::Value::Int(2));
}
if self.value.as_deref() != Some(expected) {
return Ok(redis::Value::Int(0));
}
self.value = Some(updated.to_string());
Ok(redis::Value::Int(1))
}
command => panic!("unexpected standalone Redis command: {command:?}"),
}
}
fn req_packed_commands(
&mut self,
_command: &[u8],
_offset: usize,
_count: usize,
) -> RedisResult<Vec<redis::Value>> {
panic!("atomic renewal does not use a pipeline")
}
fn get_db(&self) -> i64 {
0
}
fn check_connection(&mut self) -> bool {
true
}
fn is_open(&self) -> bool {
true
}
}
fn packed_command_arguments(command: &[u8]) -> Vec<Vec<u8>> {
let mut cursor = 0;
assert_eq!(command.get(cursor), Some(&b'*'));
cursor += 1;
let count = read_resp_number(command, &mut cursor);
(0..count)
.map(|_| {
assert_eq!(command.get(cursor), Some(&b'$'));
cursor += 1;
let length = read_resp_number(command, &mut cursor);
let value = command[cursor..cursor + length].to_vec();
cursor += length;
assert_eq!(&command[cursor..cursor + 2], b"\r\n");
cursor += 2;
value
})
.collect()
}
fn read_resp_number(command: &[u8], cursor: &mut usize) -> usize {
let line_start = *cursor;
while &command[*cursor..*cursor + 2] != b"\r\n" {
*cursor += 1;
}
let number = std::str::from_utf8(&command[line_start..*cursor])
.expect("RESP number")
.parse()
.expect("valid RESP number");
*cursor += 2;
number
}
fn parse_u64_argument(value: &[u8]) -> u64 {
std::str::from_utf8(value)
.expect("integer argument")
.parse()
.expect("valid u64 argument")
}
#[test]
fn exec_nil_retries_until_a_transaction_commits() {
let mut connection = ScriptedConnection::new([
redis::Value::Nil,
redis::Value::Nil,
redis::Value::Array(vec![redis::Value::Okay]),
]);
let mut operation_calls = 0;
let committed =
transaction_with_connection(&mut connection, "checkpoint", |connection, pipe| {
operation_calls += 1;
pipe.set("checkpoint", "value").ignore();
query_committed_transaction(connection, pipe, true)
})
.expect("third transaction attempt commits");
assert!(committed);
assert_eq!(operation_calls, 3);
assert_eq!(connection.watch_calls, 3);
assert_eq!(connection.exec_calls, 3);
assert_eq!(connection.unwatch_calls, 1);
}
#[test]
fn exec_nil_stops_after_the_transaction_retry_limit() {
let mut connection = ScriptedConnection::new(std::iter::repeat_n(
redis::Value::Nil,
TRANSACTION_MAX_ATTEMPTS,
));
let mut operation_calls = 0;
let error =
transaction_with_connection(&mut connection, "checkpoint", |connection, pipe| {
operation_calls += 1;
pipe.set("checkpoint", "value").ignore();
query_committed_transaction(connection, pipe, true)
})
.expect_err("transaction conflicts must be bounded");
assert_eq!(error.kind(), ErrorKind::TimedOut);
assert_eq!(
error.to_string(),
"redis checkpoint transaction retry limit exceeded"
);
assert_eq!(operation_calls, TRANSACTION_MAX_ATTEMPTS);
assert_eq!(connection.watch_calls, TRANSACTION_MAX_ATTEMPTS);
assert_eq!(connection.exec_calls, TRANSACTION_MAX_ATTEMPTS);
assert_eq!(connection.unwatch_calls, 1);
}
#[test]
fn renewal_atomic_script_updates_an_active_owner() {
let checkpoint = renewal_checkpoint("renewal-active", 200);
let raw = RedisStateStore::checkpoint_to_json(&checkpoint).expect("checkpoint json");
let mut connection = AtomicRenewConnection::new(raw, 150);
let clock = LeaseOperationClock::new(150);
let renewed = renew_checkpoint_claim_with_connection(
&mut connection,
"checkpoint",
"owner",
1,
300,
&clock,
)
.expect("renewal outcome");
assert!(renewed);
let persisted = RedisStateStore::checkpoint_from_json(
connection.value.as_deref().expect("persisted checkpoint"),
)
.expect("persisted checkpoint json");
assert_eq!(persisted.lease_expires_at_ms, Some(300));
assert_eq!(connection.get_calls, 1);
assert_eq!(connection.eval_calls, 1);
}
#[test]
fn renewal_atomic_script_rejects_owner_expired_at_write() {
let checkpoint = renewal_checkpoint("renewal-expired-at-write", 110);
let raw = RedisStateStore::checkpoint_to_json(&checkpoint).expect("checkpoint json");
let mut connection = AtomicRenewConnection::new(raw.clone(), 110);
let clock = LeaseOperationClock::new(100);
let error = renew_checkpoint_claim_with_connection(
&mut connection,
"checkpoint",
"owner",
1,
1_000,
&clock,
)
.expect_err("an owner expired at the atomic write boundary must fail explicitly");
assert_eq!(error.kind(), ErrorKind::TimedOut);
assert_eq!(error.to_string(), "claim lease expired");
assert_eq!(connection.value.as_deref(), Some(raw.as_str()));
let persisted = RedisStateStore::checkpoint_from_json(&raw).expect("checkpoint json");
check_claim(&persisted, 1, 110).expect("an expired claim is immediately reclaimable");
assert_eq!(connection.get_calls, 1);
assert_eq!(connection.eval_calls, 1);
}
#[test]
fn renewal_atomic_script_does_not_overwrite_a_new_owner() {
let checkpoint = renewal_checkpoint("renewal-cas-mismatch", 200);
let replacement = Checkpoint {
revision: 2,
claim_token: Some("contender".to_string()),
lease_expires_at_ms: Some(500),
..checkpoint.clone()
};
let raw = RedisStateStore::checkpoint_to_json(&checkpoint).expect("checkpoint json");
let replacement_raw =
RedisStateStore::checkpoint_to_json(&replacement).expect("replacement checkpoint json");
let mut connection =
AtomicRenewConnection::new(raw, 150).replace_before_eval(replacement_raw.clone());
let clock = LeaseOperationClock::new(150);
let renewed = renew_checkpoint_claim_with_connection(
&mut connection,
"checkpoint",
"owner",
1,
300,
&clock,
)
.expect("renewal outcome");
assert!(!renewed);
assert_eq!(connection.value.as_deref(), Some(replacement_raw.as_str()));
}
#[test]
fn renewal_atomic_script_prioritizes_expiry_over_cas_mismatch() {
let checkpoint = renewal_checkpoint("renewal-expiry-before-cas", 200);
let replacement = Checkpoint {
revision: 2,
claim_token: None,
claimed_cycle: None,
lease_expires_at_ms: None,
..checkpoint.clone()
};
let raw = RedisStateStore::checkpoint_to_json(&checkpoint).expect("checkpoint json");
let replacement_raw =
RedisStateStore::checkpoint_to_json(&replacement).expect("replacement checkpoint json");
let mut connection =
AtomicRenewConnection::new(raw, 200).replace_before_eval(replacement_raw.clone());
let clock = LeaseOperationClock::new(150);
let error = renew_checkpoint_claim_with_connection(
&mut connection,
"checkpoint",
"owner",
1,
300,
&clock,
)
.expect_err("authoritative expiry must take precedence over CAS loss");
assert_eq!(error.kind(), ErrorKind::TimedOut);
assert_eq!(error.to_string(), "claim lease expired");
assert_eq!(connection.value.as_deref(), Some(replacement_raw.as_str()));
}
fn renewal_checkpoint(task_id: &str, lease_expires_at_ms: u64) -> Checkpoint {
Checkpoint {
task_id: task_id.to_string(),
cycle_index: 0,
status: crate::types::AgentStatus::Running,
messages: Vec::new(),
cycles: Vec::new(),
shared_state: Default::default(),
revision: 1,
claim_token: Some("owner".to_string()),
claimed_cycle: Some(1),
lease_expires_at_ms: Some(lease_expires_at_ms),
terminal_result: None,
budget_usage: None,
}
}
}