use crate::{
db::Pool,
tasks::{TaskError, TaskOutcome, TaskRecord, TaskStatus, store::AbstractTaskStore},
};
#[derive(Clone)]
pub struct SqliteTaskStore {
pool: Pool,
batch_size: usize,
lease_duration: chrono::Duration,
}
impl SqliteTaskStore {
pub fn new(pool: Pool, batch_size: usize, lease_duration: std::time::Duration) -> Self {
Self {
pool,
batch_size: batch_size.max(1),
lease_duration: chrono::Duration::from_std(lease_duration).unwrap_or_default(),
}
}
fn map_store_error(err: sqlx::Error) -> TaskError {
if let sqlx::Error::Database(db_err) = &err {
let message = db_err.message();
if message.contains("idx_vyuh_tasks_active_identity")
|| message.contains("UNIQUE constraint failed: vyuh_tasks.identity")
{
return TaskError::IdentityError;
}
}
TaskError::DatabaseError(err)
}
async fn update_running(
&self,
task_id: uuid::Uuid,
runner_id: &str,
status: TaskStatus,
state: Option<String>,
resume_topic: Option<String>,
resume_input: Option<String>,
output: Option<String>,
result: Option<String>,
last_error: Option<String>,
ready_at: Option<chrono::DateTime<chrono::Utc>>,
completed_at: Option<chrono::DateTime<chrono::Utc>>,
increment_attempts: bool,
) -> Result<(), TaskError> {
let rows = sqlx::query(
r#"
UPDATE vyuh_tasks
SET status = ?1,
state = COALESCE(?2, state),
resume_topic = ?3,
resume_input = ?4,
output = ?5,
result = ?6,
last_error = ?7,
ready_at = ?8,
completed_at = ?9,
attempts = attempts + CASE WHEN ?10 THEN 1 ELSE 0 END,
locked_by = NULL,
leased_until = NULL,
updated_at = ?11
WHERE id = ?12
AND locked_by = ?13
AND status = ?14
"#,
)
.bind(status)
.bind(state)
.bind(resume_topic)
.bind(resume_input)
.bind(output)
.bind(result)
.bind(last_error)
.bind(ready_at)
.bind(completed_at)
.bind(increment_attempts)
.bind(chrono::Utc::now())
.bind(task_id)
.bind(runner_id)
.bind(TaskStatus::Running)
.execute(&self.pool)
.await?
.rows_affected();
if rows == 0 {
tracing::warn!(
"Task {} outcome ignored because runner {} no longer owns it",
task_id,
runner_id
);
}
Ok(())
}
}
impl AbstractTaskStore for SqliteTaskStore {
async fn claim_tasks(&self, runner_id: &str) -> Result<Vec<TaskRecord>, TaskError> {
let now = chrono::Utc::now();
let default_lease_ms = self.lease_duration.num_milliseconds();
let tasks = sqlx::query_as::<_, TaskRecord>(
r#"
UPDATE vyuh_tasks
SET status = ?1,
locked_by = ?2,
leased_until = strftime(
'%Y-%m-%dT%H:%M:%fZ',
?3,
printf('+%d milliseconds', COALESCE(lease_duration_ms, ?4))
),
updated_at = ?3
WHERE id IN (
SELECT id
FROM vyuh_tasks
WHERE (
status = ?5
AND (ready_at IS NULL OR ready_at <= ?3)
) OR (
status = ?1
AND leased_until IS NOT NULL
AND leased_until <= ?3
)
ORDER BY
CASE
WHEN status = ?1 THEN leased_until
ELSE COALESCE(ready_at, ?3)
END,
created_at
LIMIT ?6
)
RETURNING
id, name, input, state, resume_topic, resume_input, output, result,
status, attempts, max_attempts, retry_delay_ms, lease_duration_ms,
last_error, identity, locked_by, leased_until, ready_at, created_at,
updated_at, completed_at
"#,
)
.bind(TaskStatus::Running)
.bind(runner_id)
.bind(now)
.bind(default_lease_ms)
.bind(TaskStatus::Pending)
.bind(self.batch_size as i64)
.fetch_all(&self.pool)
.await?;
Ok(tasks)
}
async fn commit_outcome(
&self,
task_id: uuid::Uuid,
runner_id: &str,
outcome: TaskOutcome,
) -> Result<(), TaskError> {
let now = chrono::Utc::now();
match outcome {
TaskOutcome::Complete { result } => {
self.update_running(
task_id,
runner_id,
TaskStatus::Succeeded,
None,
None,
None,
None,
Some(result),
None,
None,
Some(now),
false,
)
.await
}
TaskOutcome::Suspend {
topic,
state,
output,
} => {
self.update_running(
task_id,
runner_id,
TaskStatus::Suspended,
Some(state),
Some(topic),
None,
output,
None,
None,
None,
None,
false,
)
.await
}
TaskOutcome::Sleep { state, delay } => {
self.update_running(
task_id,
runner_id,
TaskStatus::Pending,
Some(state),
None,
None,
None,
None,
None,
Some(now + chrono::Duration::from_std(delay).unwrap_or_default()),
None,
false,
)
.await
}
TaskOutcome::Retry { delay, error } => {
let retry_delay_ms: Option<i64> = sqlx::query_scalar(
r#"
SELECT retry_delay_ms
FROM vyuh_tasks
WHERE id = ?1
AND locked_by = ?2
AND status = ?3
"#,
)
.bind(task_id)
.bind(runner_id)
.bind(TaskStatus::Running)
.fetch_optional(&self.pool)
.await?
.flatten();
let retry_delay = delay
.map(chrono::Duration::from_std)
.transpose()
.unwrap_or_default()
.unwrap_or_else(|| {
retry_delay_ms
.map(chrono::Duration::milliseconds)
.unwrap_or_default()
});
let next_ready = now + retry_delay;
let rows = sqlx::query(
r#"
UPDATE vyuh_tasks
SET status = CASE
WHEN max_attempts IS NOT NULL AND attempts + 1 >= max_attempts
THEN ?1 ELSE ?2 END,
attempts = attempts + 1,
last_error = ?3,
ready_at = CASE
WHEN max_attempts IS NOT NULL AND attempts + 1 >= max_attempts
THEN NULL ELSE ?4 END,
completed_at = CASE
WHEN max_attempts IS NOT NULL AND attempts + 1 >= max_attempts
THEN ?5 ELSE NULL END,
locked_by = NULL,
leased_until = NULL,
updated_at = ?5
WHERE id = ?6
AND locked_by = ?7
AND status = ?8
"#,
)
.bind(TaskStatus::Failed)
.bind(TaskStatus::Pending)
.bind(error)
.bind(next_ready)
.bind(now)
.bind(task_id)
.bind(runner_id)
.bind(TaskStatus::Running)
.execute(&self.pool)
.await?
.rows_affected();
if rows == 0 {
tracing::warn!(
"Task {} retry outcome ignored because runner {} no longer owns it",
task_id,
runner_id
);
}
Ok(())
}
TaskOutcome::Fail { error } => {
self.update_running(
task_id,
runner_id,
TaskStatus::Failed,
None,
None,
None,
None,
None,
Some(error),
None,
Some(now),
false,
)
.await
}
}
}
async fn store_task(&self, record: TaskRecord) -> Result<(), TaskError> {
sqlx::query(
r#"
INSERT INTO vyuh_tasks (
id, name, input, state, resume_topic, resume_input, output, result,
status, attempts, max_attempts, retry_delay_ms, lease_duration_ms, last_error, identity,
locked_by, leased_until, ready_at, created_at, updated_at, completed_at
)
VALUES (
?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8,
?9, ?10, ?11, ?12, ?13, ?14, ?15,
?16, ?17, ?18, ?19, ?20, ?21
)
"#,
)
.bind(record.id)
.bind(&record.name)
.bind(&record.input)
.bind(&record.state)
.bind(&record.resume_topic)
.bind(&record.resume_input)
.bind(&record.output)
.bind(&record.result)
.bind(record.status)
.bind(record.attempts)
.bind(record.max_attempts)
.bind(record.retry_delay_ms)
.bind(record.lease_duration_ms)
.bind(&record.last_error)
.bind(&record.identity)
.bind(&record.locked_by)
.bind(record.leased_until)
.bind(record.ready_at)
.bind(record.created_at)
.bind(record.updated_at)
.bind(record.completed_at)
.execute(&self.pool)
.await
.map_err(Self::map_store_error)?;
Ok(())
}
async fn resume(&self, topic: &str, input: String) -> Result<u64, TaskError> {
let now = chrono::Utc::now();
let rows = sqlx::query(
r#"
UPDATE vyuh_tasks
SET status = ?1,
resume_input = ?2,
resume_topic = NULL,
ready_at = ?3,
updated_at = ?3
WHERE status = ?4
AND resume_topic = ?5
"#,
)
.bind(TaskStatus::Pending)
.bind(input)
.bind(now)
.bind(TaskStatus::Suspended)
.bind(topic)
.execute(&self.pool)
.await?
.rows_affected();
Ok(rows)
}
async fn run_migrations(&self) -> Result<(), TaskError> {
sqlx::query(
r#"
CREATE TABLE IF NOT EXISTS vyuh_tasks (
id BLOB PRIMARY KEY,
name TEXT NOT NULL,
input TEXT NOT NULL,
state TEXT,
resume_topic TEXT,
resume_input TEXT,
output TEXT,
result TEXT,
status INTEGER NOT NULL,
attempts INTEGER NOT NULL DEFAULT 0,
max_attempts INTEGER,
retry_delay_ms INTEGER,
lease_duration_ms INTEGER,
last_error TEXT,
identity TEXT,
locked_by TEXT,
leased_until TEXT,
ready_at TEXT,
created_at TEXT NOT NULL,
updated_at TEXT NOT NULL,
completed_at TEXT
)
"#,
)
.execute(&self.pool)
.await?;
sqlx::query(
r#"
CREATE INDEX IF NOT EXISTS idx_vyuh_tasks_pending_claim
ON vyuh_tasks(status, ready_at, created_at)
WHERE status = 0
"#,
)
.execute(&self.pool)
.await?;
sqlx::query(
r#"
CREATE INDEX IF NOT EXISTS idx_vyuh_tasks_suspended_topic
ON vyuh_tasks(resume_topic)
WHERE status = 2
"#,
)
.execute(&self.pool)
.await?;
sqlx::query(
r#"
CREATE INDEX IF NOT EXISTS idx_vyuh_tasks_running_lease
ON vyuh_tasks(status, leased_until)
WHERE status = 1
"#,
)
.execute(&self.pool)
.await?;
sqlx::query(
r#"
CREATE UNIQUE INDEX IF NOT EXISTS idx_vyuh_tasks_active_identity
ON vyuh_tasks(identity)
WHERE identity IS NOT NULL
AND status IN (0, 1, 2)
"#,
)
.execute(&self.pool)
.await?;
sqlx::query(
r#"
CREATE INDEX IF NOT EXISTS idx_vyuh_tasks_name_status
ON vyuh_tasks(name, status)
"#,
)
.execute(&self.pool)
.await?;
Ok(())
}
}