Add Redis-compatible coordination substrate

This commit is contained in:
ms2k 2026-07-16 14:30:58 +09:00
commit 91ac6dee85
21 changed files with 2481 additions and 1 deletions

View file

@ -0,0 +1,28 @@
use std::fmt;
#[derive(Debug, thiserror::Error)]
pub enum CoordinationError {
#[error("coordination store error: {0}")]
Redis(#[from] redis::RedisError),
#[error("invalid coordination response: {0}")]
Protocol(String),
#[error("configuration error: {0}")]
Config(String),
#[error("local permit capacity exhausted")]
LocalFull,
/// Value exceeds the IEEE-754 exact integer range used by Lua `tonumber`.
#[error("value {value} exceeds max safe integer {max}")]
UnsafeInteger { value: u64, max: u64 },
#[error("json error: {0}")]
Json(#[from] serde_json::Error),
}
impl CoordinationError {
pub fn protocol(message: impl fmt::Display) -> Self {
Self::Protocol(message.to_string())
}
pub fn config(message: impl fmt::Display) -> Self {
Self::Config(message.to_string())
}
}

View file

@ -0,0 +1,54 @@
//! Namespaced keys for queue, reservation, lease, progress, cancellation, and permits.
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct KeySpace {
prefix: String,
}
impl KeySpace {
pub fn new(prefix: impl Into<String>) -> Self {
Self {
prefix: prefix.into(),
}
}
#[must_use]
pub fn prefix(&self) -> &str {
&self.prefix
}
#[must_use]
pub fn queue(&self) -> String {
format!("{}queue", self.prefix)
}
#[must_use]
pub fn queue_members(&self) -> String {
format!("{}queue:members", self.prefix)
}
#[must_use]
pub fn reserve(&self, job_id: &str) -> String {
format!("{}reserve:{job_id}", self.prefix)
}
#[must_use]
pub fn lease(&self, job_id: &str) -> String {
format!("{}lease:{job_id}", self.prefix)
}
#[must_use]
pub fn progress(&self, job_id: &str) -> String {
format!("{}progress:{job_id}", self.prefix)
}
#[must_use]
pub fn cancel(&self, job_id: &str) -> String {
format!("{}cancel:{job_id}", self.prefix)
}
#[must_use]
pub fn permit_slots(&self) -> String {
format!("{}permit:slots", self.prefix)
}
}

View file

@ -0,0 +1,18 @@
//! Dragonfly/Redis coordination primitives with G-A3 global slots and L-B1
//! process-local permits; feature packages own protocol orchestration.
mod error;
mod keys;
mod local;
mod scripts;
mod store;
mod types;
pub use error::CoordinationError;
pub use keys::KeySpace;
pub use local::LocalPermitGate;
pub use store::{CoordinationStore, StoreConfig};
pub use types::{
ActivateOptions, LeaseOwner, LeaseRecord, LiveLease, MAX_SAFE_INTEGER, OpCode, OpResult,
ProgressHint, Reservation, ensure_safe_integer,
};

View file

@ -0,0 +1,101 @@
//! IF-1 L-B1 permit accounting holds local capacity until cooperative stop or
//! process death, independently of lease expiry.
use std::collections::HashMap;
use std::sync::{Arc, Mutex};
use crate::error::CoordinationError;
#[derive(Clone, Debug)]
pub struct LocalPermitGate {
inner: Arc<Mutex<LocalState>>,
limit: u64,
}
#[derive(Debug, Default)]
struct LocalState {
/// Units held by this incarnation (reserves + executing tasks).
holds: HashMap<String, u64>,
used: u64,
}
impl LocalPermitGate {
#[must_use]
pub fn new(limit: u64) -> Self {
Self {
inner: Arc::new(Mutex::new(LocalState::default())),
limit,
}
}
#[must_use]
pub fn limit(&self) -> u64 {
self.limit
}
pub fn used(&self) -> u64 {
self.inner.lock().expect("local permit lock").used
}
/// Nested acquisition increments a refcount without increasing occupancy;
/// each acquire needs `release`, or cooperative stop needs `release_all`.
pub fn try_acquire(&self, hold_id: impl Into<String>) -> Result<bool, CoordinationError> {
let hold_id = hold_id.into();
let mut state = self.inner.lock().expect("local permit lock");
if let Some(count) = state.holds.get_mut(&hold_id) {
*count = count.saturating_add(1);
return Ok(true);
}
if state.used >= self.limit {
return Ok(false);
}
state.holds.insert(hold_id, 1);
state.used = state.used.saturating_add(1);
Ok(true)
}
/// Acquire or return [`CoordinationError::LocalFull`].
pub fn acquire(&self, hold_id: impl Into<String>) -> Result<(), CoordinationError> {
if self.try_acquire(hold_id)? {
Ok(())
} else {
Err(CoordinationError::LocalFull)
}
}
/// Release one nested hold. When the last nested unit for `hold_id` is
/// released, global local occupancy decreases by one.
pub fn release(&self, hold_id: &str) -> bool {
let mut state = self.inner.lock().expect("local permit lock");
let Some(count) = state.holds.get_mut(hold_id) else {
return false;
};
*count = count.saturating_sub(1);
if *count == 0 {
state.holds.remove(hold_id);
state.used = state.used.saturating_sub(1);
}
true
}
/// Drop every hold for `hold_id` (cooperative stop of one task).
pub fn release_all(&self, hold_id: &str) -> bool {
let mut state = self.inner.lock().expect("local permit lock");
if state.holds.remove(hold_id).is_some() {
state.used = state.used.saturating_sub(1);
true
} else {
false
}
}
/// True while `hold_id` still occupies local capacity (including after
/// lease expiry while the attempt is still executing).
pub fn holds(&self, hold_id: &str) -> bool {
self.inner
.lock()
.expect("local permit lock")
.holds
.contains_key(hold_id)
}
}

View file

@ -0,0 +1,116 @@
-- KEYS: lease, reserve, progress, slots. ARGV: attempt_id, fence, worker_id,
-- lease_ttl_ms, seed_progress, require_reserve, job_id, global_limit, progress_ttl_ms.
local function now_ms()
local t = redis.call('TIME')
return (tonumber(t[1]) * 1000) + math.floor(tonumber(t[2]) / 1000)
end
local function fence_str(f)
if type(f) == 'string' then
return f
end
return string.format('%.0f', tonumber(f))
end
local function owner_eq(rec, attempt_id, fence, worker_id)
return rec.attempt_id == attempt_id
and fence_str(rec.fence) == fence_str(fence)
and rec.worker_id == worker_id
end
local function seed_progress_hint(progress_key, attempt_id, fence, progress_ttl_ms)
local prog = {
attempt_id = attempt_id,
fence = fence_str(fence),
bytes = '0'
}
redis.call('SET', progress_key, cjson.encode(prog))
redis.call('PEXPIRE', progress_key, progress_ttl_ms)
end
local lease_key = KEYS[1]
local reserve_key = KEYS[2]
local progress_key = KEYS[3]
local slots_key = KEYS[4]
local attempt_id = ARGV[1]
local fence = ARGV[2]
local worker_id = ARGV[3]
local lease_ttl_ms = tonumber(ARGV[4])
local seed_progress = ARGV[5] == '1'
local require_reserve = ARGV[6] == '1'
local job_id = ARGV[7]
local global_limit = tonumber(ARGV[8])
local progress_ttl_ms = tonumber(ARGV[9]) or lease_ttl_ms
local now = now_ms()
local raw = redis.call('GET', lease_key)
local rec = nil
if raw then
rec = cjson.decode(raw)
end
if rec and tonumber(rec.expires_at) > now then
if owner_eq(rec, attempt_id, fence, worker_id) then
-- At-least-once activate: refresh lease only; never reset progress bytes.
rec.expires_at = now + lease_ttl_ms
redis.call('SET', lease_key, cjson.encode(rec))
redis.call('PEXPIRE', lease_key, lease_ttl_ms)
redis.call('SADD', slots_key, job_id)
return {
'OK_IDEMPOTENT',
attempt_id,
fence_str(fence),
worker_id,
string.format('%.0f', tonumber(rec.expires_at))
}
end
return {
'FAIL_HELD',
rec.attempt_id,
fence_str(rec.fence),
rec.worker_id,
string.format('%.0f', tonumber(rec.expires_at))
}
end
local res_raw = redis.call('GET', reserve_key)
if res_raw then
local res = cjson.decode(res_raw)
if res and res.worker_id and res.worker_id ~= worker_id then
return {'FAIL_RESERVE', res.worker_id}
end
elseif require_reserve then
return {'FAIL_RESERVE', 'missing'}
end
if redis.call('SISMEMBER', slots_key, job_id) == 0 then
local used = redis.call('SCARD', slots_key)
if used >= global_limit then
return {'FAIL_FULL', tostring(used), tostring(global_limit)}
end
redis.call('SADD', slots_key, job_id)
end
local new_rec = {
attempt_id = attempt_id,
fence = fence_str(fence),
worker_id = worker_id,
expires_at = now + lease_ttl_ms
}
redis.call('SET', lease_key, cjson.encode(new_rec))
redis.call('PEXPIRE', lease_key, lease_ttl_ms)
if res_raw then
redis.call('DEL', reserve_key)
end
if seed_progress then
seed_progress_hint(progress_key, attempt_id, fence, progress_ttl_ms)
else
redis.call('DEL', progress_key)
end
return {
'OK',
attempt_id,
fence_str(fence),
worker_id,
string.format('%.0f', tonumber(new_rec.expires_at))
}

View file

@ -0,0 +1,22 @@
-- KEYS: queue, members, reserve, progress, cancel, slots. ARGV[1]=job_id.
-- Terminal cleanup preserves live leases and frees only reservation occupancy.
local queue_key = KEYS[1]
local members_key = KEYS[2]
local reserve_key = KEYS[3]
local progress_key = KEYS[4]
local cancel_key = KEYS[5]
local slots_key = KEYS[6]
local job_id = ARGV[1]
redis.call('LREM', queue_key, 0, job_id)
redis.call('SREM', members_key, job_id)
local had_reserve = redis.call('DEL', reserve_key)
redis.call('DEL', progress_key)
redis.call('DEL', cancel_key)
if had_reserve > 0 then
redis.call('SREM', slots_key, job_id)
end
return {'OK', tostring(had_reserve)}

View file

@ -0,0 +1,38 @@
-- KEYS[1]=lease KEYS[2]=progress
local function now_ms()
local t = redis.call('TIME')
return (tonumber(t[1]) * 1000) + math.floor(tonumber(t[2]) / 1000)
end
local function fence_str(f)
if type(f) == 'string' then
return f
end
return string.format('%.0f', tonumber(f))
end
local lease_key = KEYS[1]
local progress_key = KEYS[2]
local praw = redis.call('GET', progress_key)
if not praw then
return {'OK_GONE'}
end
local prog = cjson.decode(praw)
local lraw = redis.call('GET', lease_key)
if not lraw then
redis.call('DEL', progress_key)
return {'OK', 'no_lease'}
end
local rec = cjson.decode(lraw)
local now = now_ms()
if tonumber(rec.expires_at) <= now then
redis.call('DEL', progress_key)
return {'OK', 'expired_lease'}
end
if rec.attempt_id == prog.attempt_id
and fence_str(rec.fence) == fence_str(prog.fence)
then
return {'OK_IDEMPOTENT', 'live'}
end
redis.call('DEL', progress_key)
return {'OK', 'mismatched'}

View file

@ -0,0 +1,9 @@
-- KEYS[1]=queue list KEYS[2]=queue members
local queue_key = KEYS[1]
local members_key = KEYS[2]
local job_id = redis.call('LPOP', queue_key)
if not job_id then
return {'NONE'}
end
redis.call('SREM', members_key, job_id)
return {'OK', job_id}

View file

@ -0,0 +1,10 @@
-- KEYS[1]=queue list KEYS[2]=queue members
-- ARGV[1]=job_id — idempotent projection of Queued dispatch.
local queue_key = KEYS[1]
local members_key = KEYS[2]
local job_id = ARGV[1]
if redis.call('SADD', members_key, job_id) == 0 then
return {'OK_IDEMPOTENT', job_id}
end
redis.call('RPUSH', queue_key, job_id)
return {'OK', job_id}

View file

@ -0,0 +1,34 @@
-- KEYS: lease, slots, reserve. ARGV[1]=job_id.
-- C2 expiry preserves G-A3 occupancy while a reservation remains.
local function now_ms()
local t = redis.call('TIME')
return (tonumber(t[1]) * 1000) + math.floor(tonumber(t[2]) / 1000)
end
local function free_slot_if_unreserved(slots_key, reserve_key, job_id)
if redis.call('EXISTS', reserve_key) == 0 then
redis.call('SREM', slots_key, job_id)
end
end
local lease_key = KEYS[1]
local slots_key = KEYS[2]
local reserve_key = KEYS[3]
local job_id = ARGV[1]
local raw = redis.call('GET', lease_key)
if not raw then
free_slot_if_unreserved(slots_key, reserve_key, job_id)
return {'OK_GONE'}
end
local rec = cjson.decode(raw)
local now = now_ms()
if tonumber(rec.expires_at) <= now then
redis.call('DEL', lease_key)
free_slot_if_unreserved(slots_key, reserve_key, job_id)
return {'OK', 'expired'}
end
return {
'FAIL_HELD',
string.format('%.0f', tonumber(rec.expires_at)),
string.format('%.0f', now)
}

View file

@ -0,0 +1,39 @@
-- KEYS: reserve, slots, lease. ARGV: worker_id (empty = any), job_id.
-- Preserve the global slot while a live lease remains.
local function now_ms()
local t = redis.call('TIME')
return (tonumber(t[1]) * 1000) + math.floor(tonumber(t[2]) / 1000)
end
local function live_lease(lease_key)
local lraw = redis.call('GET', lease_key)
if not lraw then
return false
end
local rec = cjson.decode(lraw)
return tonumber(rec.expires_at) > now_ms()
end
local reserve_key = KEYS[1]
local slots_key = KEYS[2]
local lease_key = KEYS[3]
local worker_id = ARGV[1]
local job_id = ARGV[2]
local raw = redis.call('GET', reserve_key)
if not raw then
if not live_lease(lease_key) then
redis.call('SREM', slots_key, job_id)
end
return {'OK_GONE'}
end
if worker_id and worker_id ~= '' then
local res = cjson.decode(raw)
if res.worker_id ~= worker_id then
return {'FAIL_RESERVE', res.worker_id}
end
end
redis.call('DEL', reserve_key)
if not live_lease(lease_key) then
redis.call('SREM', slots_key, job_id)
end
return {'OK'}

View file

@ -0,0 +1,29 @@
-- KEYS[1]=lease
local function now_ms()
local t = redis.call('TIME')
return (tonumber(t[1]) * 1000) + math.floor(tonumber(t[2]) / 1000)
end
local function fence_str(f)
if type(f) == 'string' then
return f
end
return string.format('%.0f', tonumber(f))
end
local raw = redis.call('GET', KEYS[1])
if not raw then
return {'NONE'}
end
local rec = cjson.decode(raw)
local now = now_ms()
local live = tonumber(rec.expires_at) > now and '1' or '0'
return {
'REC',
live,
rec.attempt_id,
fence_str(rec.fence),
rec.worker_id,
string.format('%.0f', tonumber(rec.expires_at)),
string.format('%.0f', now)
}

View file

@ -0,0 +1,14 @@
-- KEYS[1]=cancel ARGV[1]=ttl_ms (0 = no TTL)
local cancel_key = KEYS[1]
local ttl_ms = tonumber(ARGV[1])
if redis.call('EXISTS', cancel_key) == 1 then
if ttl_ms > 0 then
redis.call('PEXPIRE', cancel_key, ttl_ms)
end
return {'OK_IDEMPOTENT'}
end
redis.call('SET', cancel_key, '1')
if ttl_ms > 0 then
redis.call('PEXPIRE', cancel_key, ttl_ms)
end
return {'OK'}

View file

@ -0,0 +1,52 @@
-- KEYS[1]=lease KEYS[2]=slots KEYS[3]=reserve
-- ARGV[1]=attempt_id ARGV[2]=fence ARGV[3]=worker_id ARGV[4]=job_id
local function now_ms()
local t = redis.call('TIME')
return (tonumber(t[1]) * 1000) + math.floor(tonumber(t[2]) / 1000)
end
local function fence_str(f)
if type(f) == 'string' then
return f
end
return string.format('%.0f', tonumber(f))
end
local function owner_eq(rec, attempt_id, fence, worker_id)
return rec.attempt_id == attempt_id
and fence_str(rec.fence) == fence_str(fence)
and rec.worker_id == worker_id
end
local function free_slot_if_unreserved(slots_key, reserve_key, job_id)
if redis.call('EXISTS', reserve_key) == 0 then
redis.call('SREM', slots_key, job_id)
end
end
local lease_key = KEYS[1]
local slots_key = KEYS[2]
local reserve_key = KEYS[3]
local attempt_id = ARGV[1]
local fence = ARGV[2]
local worker_id = ARGV[3]
local job_id = ARGV[4]
local raw = redis.call('GET', lease_key)
if not raw then
free_slot_if_unreserved(slots_key, reserve_key, job_id)
return {'OK_GONE'}
end
local rec = cjson.decode(raw)
local now = now_ms()
if not owner_eq(rec, attempt_id, fence, worker_id) then
if tonumber(rec.expires_at) > now then
return {'FAIL_STOLEN', rec.attempt_id, fence_str(rec.fence), rec.worker_id}
end
redis.call('DEL', lease_key)
free_slot_if_unreserved(slots_key, reserve_key, job_id)
return {'OK_GONE'}
end
redis.call('DEL', lease_key)
free_slot_if_unreserved(slots_key, reserve_key, job_id)
return {'OK'}

View file

@ -0,0 +1,10 @@
-- KEYS[1]=queue list KEYS[2]=queue members ARGV[1]=job_id
local queue_key = KEYS[1]
local members_key = KEYS[2]
local job_id = ARGV[1]
redis.call('LREM', queue_key, 0, job_id)
local removed = redis.call('SREM', members_key, job_id)
if removed == 0 then
return {'OK_GONE', job_id}
end
return {'OK', job_id}

View file

@ -0,0 +1,46 @@
-- KEYS[1]=lease
-- ARGV[1]=attempt_id ARGV[2]=fence ARGV[3]=worker_id ARGV[4]=ttl_ms
local function now_ms()
local t = redis.call('TIME')
return (tonumber(t[1]) * 1000) + math.floor(tonumber(t[2]) / 1000)
end
local function fence_str(f)
if type(f) == 'string' then
return f
end
return string.format('%.0f', tonumber(f))
end
local function owner_eq(rec, attempt_id, fence, worker_id)
return rec.attempt_id == attempt_id
and fence_str(rec.fence) == fence_str(fence)
and rec.worker_id == worker_id
end
local lease_key = KEYS[1]
local attempt_id = ARGV[1]
local fence = ARGV[2]
local worker_id = ARGV[3]
local ttl_ms = tonumber(ARGV[4])
local raw = redis.call('GET', lease_key)
if not raw then
return {'FAIL_MISSING'}
end
local rec = cjson.decode(raw)
local now = now_ms()
if tonumber(rec.expires_at) <= now then
return {
'FAIL_EXPIRED',
string.format('%.0f', tonumber(rec.expires_at)),
string.format('%.0f', now)
}
end
if not owner_eq(rec, attempt_id, fence, worker_id) then
return {'FAIL_STOLEN', rec.attempt_id, fence_str(rec.fence), rec.worker_id}
end
rec.expires_at = now + ttl_ms
redis.call('SET', lease_key, cjson.encode(rec))
redis.call('PEXPIRE', lease_key, ttl_ms)
return {'OK', string.format('%.0f', tonumber(rec.expires_at))}

View file

@ -0,0 +1,48 @@
-- KEYS: lease, progress. ARGV: attempt_id, fence, worker_id, bytes, ttl_ms.
-- Fence and bytes remain decimal strings to preserve integer identity.
local function now_ms()
local t = redis.call('TIME')
return (tonumber(t[1]) * 1000) + math.floor(tonumber(t[2]) / 1000)
end
local function fence_str(f)
if type(f) == 'string' then
return f
end
return string.format('%.0f', tonumber(f))
end
local function owner_eq(rec, attempt_id, fence, worker_id)
return rec.attempt_id == attempt_id
and fence_str(rec.fence) == fence_str(fence)
and rec.worker_id == worker_id
end
local lease_key = KEYS[1]
local progress_key = KEYS[2]
local attempt_id = ARGV[1]
local fence = ARGV[2]
local worker_id = ARGV[3]
local bytes = ARGV[4]
local ttl_ms = tonumber(ARGV[5])
local raw = redis.call('GET', lease_key)
if not raw then
return {'FAIL_MISSING'}
end
local rec = cjson.decode(raw)
local now = now_ms()
if tonumber(rec.expires_at) <= now then
return {'FAIL_EXPIRED'}
end
if not owner_eq(rec, attempt_id, fence, worker_id) then
return {'FAIL_STOLEN'}
end
local prog = {
attempt_id = attempt_id,
fence = fence_str(fence),
bytes = fence_str(bytes)
}
redis.call('SET', progress_key, cjson.encode(prog))
redis.call('PEXPIRE', progress_key, ttl_ms)
return {'OK', fence_str(bytes)}

View file

@ -0,0 +1,47 @@
-- G-A3 global slots: KEYS[1]=reserve KEYS[2]=slots KEYS[3]=lease
-- ARGV[1]=worker_id ARGV[2]=ttl_ms ARGV[3]=global_limit ARGV[4]=job_id
local function now_ms()
local t = redis.call('TIME')
return (tonumber(t[1]) * 1000) + math.floor(tonumber(t[2]) / 1000)
end
local reserve_key = KEYS[1]
local slots_key = KEYS[2]
local lease_key = KEYS[3]
local worker_id = ARGV[1]
local ttl_ms = tonumber(ARGV[2])
local global_limit = tonumber(ARGV[3])
local job_id = ARGV[4]
local raw = redis.call('GET', reserve_key)
if raw then
local res = cjson.decode(raw)
if res.worker_id == worker_id then
redis.call('PEXPIRE', reserve_key, ttl_ms)
redis.call('SADD', slots_key, job_id)
return {'OK_IDEMPOTENT', worker_id}
end
return {'FAIL_RESERVE', res.worker_id}
end
local lraw = redis.call('GET', lease_key)
if lraw then
local rec = cjson.decode(lraw)
if tonumber(rec.expires_at) > now_ms() then
return {'FAIL_HELD', rec.worker_id}
end
end
local already = redis.call('SISMEMBER', slots_key, job_id)
if already == 0 then
local used = redis.call('SCARD', slots_key)
if used >= global_limit then
return {'FAIL_FULL', tostring(used), tostring(global_limit)}
end
end
local rec = { worker_id = worker_id, reserved_at = now_ms() }
redis.call('SET', reserve_key, cjson.encode(rec))
redis.call('PEXPIRE', reserve_key, ttl_ms)
redis.call('SADD', slots_key, job_id)
return {'OK', worker_id, tostring(redis.call('SCARD', slots_key))}

View file

@ -0,0 +1,39 @@
use redis::Script;
pub struct Scripts {
pub reserve: Script,
pub free_reserve: Script,
pub activate: Script,
pub renew: Script,
pub release: Script,
pub expire: Script,
pub get_live: Script,
pub report_progress: Script,
pub clear_stale_progress: Script,
pub enqueue: Script,
pub dequeue: Script,
pub remove_queue: Script,
pub publish_cancel: Script,
pub clean_terminal: Script,
}
impl Scripts {
pub fn load() -> Self {
Self {
reserve: Script::new(include_str!("lua/reserve.lua")),
free_reserve: Script::new(include_str!("lua/free_reserve.lua")),
activate: Script::new(include_str!("lua/activate.lua")),
renew: Script::new(include_str!("lua/renew.lua")),
release: Script::new(include_str!("lua/release.lua")),
expire: Script::new(include_str!("lua/expire.lua")),
get_live: Script::new(include_str!("lua/get_live.lua")),
report_progress: Script::new(include_str!("lua/report_progress.lua")),
clear_stale_progress: Script::new(include_str!("lua/clear_stale_progress.lua")),
enqueue: Script::new(include_str!("lua/enqueue.lua")),
dequeue: Script::new(include_str!("lua/dequeue.lua")),
remove_queue: Script::new(include_str!("lua/remove_queue.lua")),
publish_cancel: Script::new(include_str!("lua/publish_cancel.lua")),
clean_terminal: Script::new(include_str!("lua/clean_terminal.lua")),
}
}
}

View file

@ -0,0 +1,513 @@
use std::time::Duration;
use redis::aio::ConnectionManager;
use redis::{AsyncCommands, Client, FromRedisValue, RedisResult, Value};
use crate::error::CoordinationError;
use crate::keys::KeySpace;
use crate::scripts::Scripts;
use crate::types::{
ActivateOptions, LeaseOwner, LeaseRecord, LiveLease, OpCode, OpResult, ProgressHint,
Reservation, ensure_safe_integer,
};
#[derive(Clone, Debug)]
pub struct StoreConfig {
pub global_limit: u64,
pub lease_ttl: Duration,
pub reserve_ttl: Duration,
pub progress_ttl: Duration,
pub cancel_ttl: Option<Duration>,
}
impl Default for StoreConfig {
fn default() -> Self {
Self {
global_limit: 1024,
lease_ttl: Duration::from_secs(30),
reserve_ttl: Duration::from_secs(30),
progress_ttl: Duration::from_secs(30),
cancel_ttl: None,
}
}
}
/// Dragonfly/Redis substrate using IF-1 G-A3 global job slots; process-local
/// permits live in [`crate::LocalPermitGate`].
#[derive(Clone)]
pub struct CoordinationStore {
conn: ConnectionManager,
keys: KeySpace,
scripts: std::sync::Arc<Scripts>,
config: StoreConfig,
}
impl CoordinationStore {
pub async fn connect(
url: &str,
namespace: impl Into<String>,
config: StoreConfig,
) -> Result<Self, CoordinationError> {
let client = Client::open(url).map_err(|error| {
CoordinationError::config(format!("invalid coordination URL: {error}"))
})?;
let conn = ConnectionManager::new(client).await?;
Ok(Self {
conn,
keys: KeySpace::new(namespace),
scripts: std::sync::Arc::new(Scripts::load()),
config,
})
}
#[must_use]
pub fn keys(&self) -> &KeySpace {
&self.keys
}
#[must_use]
pub fn config(&self) -> &StoreConfig {
&self.config
}
pub async fn global_used(&self) -> Result<u64, CoordinationError> {
let mut conn = self.conn.clone();
let used: u64 = conn.scard(self.keys.permit_slots()).await?;
Ok(used)
}
pub async fn enqueue(&self, job_id: &str) -> Result<OpResult, CoordinationError> {
self.eval_script(
&self.scripts.enqueue,
&[self.keys.queue(), self.keys.queue_members()],
&[job_id],
)
.await
}
pub async fn dequeue(&self) -> Result<Option<String>, CoordinationError> {
let result = self
.eval_script(
&self.scripts.dequeue,
&[self.keys.queue(), self.keys.queue_members()],
&[] as &[&str],
)
.await?;
if result.code == OpCode::Ok {
Ok(result.detail.first().cloned())
} else {
Ok(None)
}
}
pub async fn remove_from_queue(&self, job_id: &str) -> Result<OpResult, CoordinationError> {
self.eval_script(
&self.scripts.remove_queue,
&[self.keys.queue(), self.keys.queue_members()],
&[job_id],
)
.await
}
/// Reserve under global capacity (IF-1 G-A3 + IF-4 reserve key).
pub async fn reserve(
&self,
job_id: &str,
worker_id: &str,
) -> Result<OpResult, CoordinationError> {
let ttl_ms = duration_ms(self.config.reserve_ttl);
self.eval_script(
&self.scripts.reserve,
&[
self.keys.reserve(job_id),
self.keys.permit_slots(),
self.keys.lease(job_id),
],
&[
worker_id,
&ttl_ms.to_string(),
&self.config.global_limit.to_string(),
job_id,
],
)
.await
}
pub async fn free_reserve(
&self,
job_id: &str,
worker_id: Option<&str>,
) -> Result<OpResult, CoordinationError> {
self.eval_script(
&self.scripts.free_reserve,
&[
self.keys.reserve(job_id),
self.keys.permit_slots(),
self.keys.lease(job_id),
],
&[worker_id.unwrap_or(""), job_id],
)
.await
}
pub async fn get_reservation(
&self,
job_id: &str,
) -> Result<Option<Reservation>, CoordinationError> {
let mut conn = self.conn.clone();
let raw: Option<String> = conn.get(self.keys.reserve(job_id)).await?;
match raw {
Some(body) => Ok(Some(serde_json::from_str(&body)?)),
None => Ok(None),
}
}
pub async fn activate(
&self,
job_id: &str,
owner: &LeaseOwner,
options: ActivateOptions,
) -> Result<OpResult, CoordinationError> {
owner.validate()?;
let lease_ttl_ms = duration_ms(self.config.lease_ttl);
let progress_ttl_ms = duration_ms(self.config.progress_ttl);
self.eval_script(
&self.scripts.activate,
&[
self.keys.lease(job_id),
self.keys.reserve(job_id),
self.keys.progress(job_id),
self.keys.permit_slots(),
],
&[
owner.attempt_id.as_str(),
&owner.fence.to_string(),
owner.worker_id.as_str(),
&lease_ttl_ms.to_string(),
if options.seed_progress { "1" } else { "0" },
if options.require_reserve { "1" } else { "0" },
job_id,
&self.config.global_limit.to_string(),
&progress_ttl_ms.to_string(),
],
)
.await
}
pub async fn renew(
&self,
job_id: &str,
owner: &LeaseOwner,
) -> Result<OpResult, CoordinationError> {
owner.validate()?;
let ttl_ms = duration_ms(self.config.lease_ttl);
self.eval_script(
&self.scripts.renew,
&[self.keys.lease(job_id)],
&[
owner.attempt_id.as_str(),
&owner.fence.to_string(),
owner.worker_id.as_str(),
&ttl_ms.to_string(),
],
)
.await
}
pub async fn release(
&self,
job_id: &str,
owner: &LeaseOwner,
) -> Result<OpResult, CoordinationError> {
owner.validate()?;
self.eval_script(
&self.scripts.release,
&[
self.keys.lease(job_id),
self.keys.permit_slots(),
self.keys.reserve(job_id),
],
&[
owner.attempt_id.as_str(),
&owner.fence.to_string(),
owner.worker_id.as_str(),
job_id,
],
)
.await
}
pub async fn expire_lease(&self, job_id: &str) -> Result<OpResult, CoordinationError> {
self.eval_script(
&self.scripts.expire,
&[
self.keys.lease(job_id),
self.keys.permit_slots(),
self.keys.reserve(job_id),
],
&[job_id],
)
.await
}
pub async fn get_live_lease(
&self,
job_id: &str,
) -> Result<Option<LiveLease>, CoordinationError> {
let mut conn = self.conn.clone();
let mut invocation = self.scripts.get_live.prepare_invoke();
invocation.key(self.keys.lease(job_id));
let raw: Value = invocation.invoke_async(&mut conn).await?;
let parts = value_to_strings(&raw)?;
if parts.first().map(String::as_str) != Some("REC") {
return Ok(None);
}
if parts.len() < 7 {
return Err(CoordinationError::protocol("truncated get_live response"));
}
if parts[1] != "1" {
return Ok(None);
}
Ok(Some(LiveLease {
record: LeaseRecord {
attempt_id: parts[2].clone(),
fence: parse_u64(&parts[3])?,
worker_id: parts[4].clone(),
expires_at: parse_u64(&parts[5])?,
},
now_ms: parse_u64(&parts[6])?,
}))
}
pub async fn report_progress(
&self,
job_id: &str,
owner: &LeaseOwner,
bytes: u64,
) -> Result<OpResult, CoordinationError> {
owner.validate()?;
ensure_safe_integer(bytes)?;
let ttl_ms = duration_ms(self.config.progress_ttl);
self.eval_script(
&self.scripts.report_progress,
&[self.keys.lease(job_id), self.keys.progress(job_id)],
&[
owner.attempt_id.as_str(),
&owner.fence.to_string(),
owner.worker_id.as_str(),
&bytes.to_string(),
&ttl_ms.to_string(),
],
)
.await
}
pub async fn get_progress(
&self,
job_id: &str,
) -> Result<Option<ProgressHint>, CoordinationError> {
let mut conn = self.conn.clone();
let raw: Option<String> = conn.get(self.keys.progress(job_id)).await?;
match raw {
Some(body) => Ok(Some(serde_json::from_str(&body)?)),
None => Ok(None),
}
}
pub async fn clear_stale_progress(&self, job_id: &str) -> Result<OpResult, CoordinationError> {
self.eval_script(
&self.scripts.clear_stale_progress,
&[self.keys.lease(job_id), self.keys.progress(job_id)],
&[] as &[&str],
)
.await
}
pub async fn publish_cancel(&self, job_id: &str) -> Result<OpResult, CoordinationError> {
let ttl_ms = self.config.cancel_ttl.map(duration_ms).unwrap_or(0);
self.eval_script(
&self.scripts.publish_cancel,
&[self.keys.cancel(job_id)],
&[&ttl_ms.to_string()],
)
.await
}
pub async fn cancel_notice(&self, job_id: &str) -> Result<bool, CoordinationError> {
let mut conn = self.conn.clone();
let exists: bool = conn.exists(self.keys.cancel(job_id)).await?;
Ok(exists)
}
pub async fn clear_cancel(&self, job_id: &str) -> Result<(), CoordinationError> {
let mut conn = self.conn.clone();
let _: () = conn.del(self.keys.cancel(job_id)).await?;
Ok(())
}
/// Idempotent durable-to-coordination enqueue projection.
pub async fn project_enqueue(&self, job_id: &str) -> Result<OpResult, CoordinationError> {
self.enqueue(job_id).await
}
/// Idempotent cancel-notice projection.
pub async fn project_cancel(&self, job_id: &str) -> Result<OpResult, CoordinationError> {
self.publish_cancel(job_id).await
}
/// Drop queue membership, reservation, cancel notice, and progress for a
/// terminal job. Does not delete a live lease (release/expire first).
pub async fn clean_terminal(&self, job_id: &str) -> Result<OpResult, CoordinationError> {
self.eval_script(
&self.scripts.clean_terminal,
&[
self.keys.queue(),
self.keys.queue_members(),
self.keys.reserve(job_id),
self.keys.progress(job_id),
self.keys.cancel(job_id),
self.keys.permit_slots(),
],
&[job_id],
)
.await
}
/// Dequeue then reserve. On capacity failure the job is re-enqueued.
pub async fn claim_and_reserve(
&self,
worker_id: &str,
) -> Result<Option<String>, CoordinationError> {
for _ in 0..64 {
let Some(job_id) = self.dequeue().await? else {
return Ok(None);
};
let result = self.reserve(&job_id, worker_id).await?;
match result.code {
OpCode::Ok | OpCode::OkIdempotent => return Ok(Some(job_id)),
OpCode::FailFull => {
let _ = self.enqueue(&job_id).await?;
return Ok(None);
}
OpCode::FailHeld | OpCode::FailReserve => continue,
_ => {
let _ = self.enqueue(&job_id).await?;
return Err(CoordinationError::protocol(format!(
"unexpected reserve result {:?}",
result.code
)));
}
}
}
Ok(None)
}
/// Deletes known job and shared keys for tests and reconstruction helpers;
/// this is not a full SCAN wipe.
pub async fn drop_job_state(&self, job_ids: &[&str]) -> Result<(), CoordinationError> {
let mut conn = self.conn.clone();
let mut keys = vec![
self.keys.queue(),
self.keys.queue_members(),
self.keys.permit_slots(),
];
for job_id in job_ids {
keys.push(self.keys.reserve(job_id));
keys.push(self.keys.lease(job_id));
keys.push(self.keys.progress(job_id));
keys.push(self.keys.cancel(job_id));
}
if !keys.is_empty() {
let _: () = conn.del(keys).await?;
}
Ok(())
}
async fn eval_script(
&self,
script: &redis::Script,
keys: &[String],
args: &[&str],
) -> Result<OpResult, CoordinationError> {
let mut conn = self.conn.clone();
let mut invocation = script.prepare_invoke();
for key in keys {
invocation.key(key);
}
for arg in args {
invocation.arg(*arg);
}
let raw: Value = invocation.invoke_async(&mut conn).await?;
parse_op_result(&raw)
}
}
fn duration_ms(duration: Duration) -> u64 {
u64::try_from(duration.as_millis())
.unwrap_or(u64::MAX)
.max(1)
}
fn parse_u64(raw: &str) -> Result<u64, CoordinationError> {
raw.parse()
.map_err(|_| CoordinationError::protocol(format!("invalid u64: {raw}")))
}
fn value_to_strings(value: &Value) -> Result<Vec<String>, CoordinationError> {
match value {
Value::Nil => Ok(Vec::new()),
Value::SimpleString(s) => Ok(vec![s.clone()]),
Value::BulkString(bytes) => Ok(vec![String::from_utf8_lossy(bytes).into_owned()]),
Value::Array(items) => {
let mut out = Vec::with_capacity(items.len());
for item in items {
match item {
Value::Nil => out.push(String::new()),
Value::SimpleString(s) => out.push(s.clone()),
Value::BulkString(bytes) => {
out.push(String::from_utf8_lossy(bytes).into_owned())
}
Value::Int(n) => out.push(n.to_string()),
other => {
// Fall back to Redis display for unexpected tags.
let as_string: RedisResult<String> =
FromRedisValue::from_redis_value(other);
out.push(as_string.unwrap_or_else(|_| format!("{other:?}")));
}
}
}
Ok(out)
}
Value::Int(n) => Ok(vec![n.to_string()]),
other => Err(CoordinationError::protocol(format!(
"unexpected redis value: {other:?}"
))),
}
}
fn parse_op_result(value: &Value) -> Result<OpResult, CoordinationError> {
let parts = value_to_strings(value)?;
if parts.is_empty() {
return Err(CoordinationError::protocol("empty script response"));
}
// get_live / dequeue use non-OpCode tags.
if let Some(code) = OpCode::parse(&parts[0]) {
return Ok(OpResult {
code,
detail: parts[1..].to_vec(),
});
}
match parts[0].as_str() {
"NONE" => Ok(OpResult {
code: OpCode::OkGone,
detail: parts[1..].to_vec(),
}),
"REC" => Ok(OpResult {
code: OpCode::Ok,
detail: parts,
}),
other => Err(CoordinationError::protocol(format!(
"unknown script status: {other}"
))),
}
}

Some files were not shown because too many files have changed in this diff Show more