From 530bb314a69e245c2b7dcfa60bc9cfa72fc46c6a Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Tue, 6 Oct 2026 14:20:27 -0400 Subject: [PATCH 1/9] feat(dgw): record provisioner tasks in gateway.db A Task is the permanent record of work the provisioner asked Gateway to do: its kind, target and parameters, its state and a payload the kind defines. The work itself will run as a job on the job queue, which forgets it once done; the Task stays for status and auditing. The provisioner chooses the Task ID, so asking again with the same request returns the existing Task. At most one unfinished Task works on a target per kind, and a Task past its deadline is failed when another one asks for its target, so a lost job never blocks it. The provisioner-task crate defines the operations and provisioner-task-libsql stores them in the provisioner_task_records table of gateway.db, with its own migrations and a schema_version row in provisioner_task_metadata. Co-Authored-By: Claude Opus 5.5 (1M context) --- Cargo.lock | 28 + crates/provisioner-task-libsql/Cargo.toml | 26 + .../migrations/01_provisioner_task.sql | 25 + crates/provisioner-task-libsql/src/lib.rs | 660 ++++++++++++++++++ crates/provisioner-task/Cargo.toml | 16 + crates/provisioner-task/src/lib.rs | 114 +++ 6 files changed, 869 insertions(+) create mode 100644 crates/provisioner-task-libsql/Cargo.toml create mode 100644 crates/provisioner-task-libsql/migrations/01_provisioner_task.sql create mode 100644 crates/provisioner-task-libsql/src/lib.rs create mode 100644 crates/provisioner-task/Cargo.toml create mode 100644 crates/provisioner-task/src/lib.rs diff --git a/Cargo.lock b/Cargo.lock index 1fd7f9dc0..21ed6c292 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -6009,6 +6009,34 @@ version = "3.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "95067976aca6421a523e491fce939a3e65249bac4b977adee0ee9771568e8aa3" +[[package]] +name = "provisioner-task" +version = "0.0.0" +dependencies = [ + "anyhow", + "async-trait", + "serde_json", + "time", + "uuid", +] + +[[package]] +name = "provisioner-task-libsql" +version = "0.0.0" +dependencies = [ + "anyhow", + "async-trait", + "gateway-db", + "libsql", + "provisioner-task", + "serde_json", + "tempfile", + "time", + "tokio 1.52.3", + "tracing", + "uuid", +] + [[package]] name = "proxy-generators" version = "0.0.0" diff --git a/crates/provisioner-task-libsql/Cargo.toml b/crates/provisioner-task-libsql/Cargo.toml new file mode 100644 index 000000000..ec683efc3 --- /dev/null +++ b/crates/provisioner-task-libsql/Cargo.toml @@ -0,0 +1,26 @@ +[package] +name = "provisioner-task-libsql" +version = "0.0.0" +edition = "2024" +authors = ["Devolutions Inc. "] +publish = false + +[lints] +workspace = true + +[dependencies] +anyhow = "1" +async-trait = "0.1" +libsql = { version = "0.9", default-features = false, features = ["core"] } +provisioner-task.path = "../provisioner-task" +serde_json = "1" +time = { version = "0.3", default-features = false } +tokio = { version = "1.52", features = ["sync"] } +tracing = "0.1" +uuid = "1.23" + +[dev-dependencies] +gateway-db.path = "../gateway-db" +tempfile = "3.24" +tokio = { version = "1.52", features = ["macros", "rt-multi-thread"] } +uuid = { version = "1.23", features = ["v4"] } diff --git a/crates/provisioner-task-libsql/migrations/01_provisioner_task.sql b/crates/provisioner-task-libsql/migrations/01_provisioner_task.sql new file mode 100644 index 000000000..f37e3b4dc --- /dev/null +++ b/crates/provisioner-task-libsql/migrations/01_provisioner_task.sql @@ -0,0 +1,25 @@ +CREATE TABLE provisioner_task_records ( + id TEXT PRIMARY KEY, + kind TEXT NOT NULL, + target TEXT NOT NULL, + params TEXT NOT NULL, + state TEXT NOT NULL, + payload TEXT NOT NULL, + attempts INTEGER NOT NULL DEFAULT 0, + created_at INTEGER NOT NULL, + started_at INTEGER, + finished_at INTEGER, + deadline_at INTEGER NOT NULL, + CHECK (state IN ('queued', 'running', 'succeeded', 'failed')), + CHECK ((state IN ('succeeded', 'failed')) = (finished_at IS NOT NULL)), + CHECK (json_valid(params) AND json_valid(payload)), + CHECK (attempts >= 0) +); + +CREATE UNIQUE INDEX provisioner_task_records_unfinished ON provisioner_task_records (kind, target) + WHERE state IN ('queued', 'running'); + +CREATE TABLE provisioner_task_metadata ( + key TEXT PRIMARY KEY, + value BLOB NOT NULL +); diff --git a/crates/provisioner-task-libsql/src/lib.rs b/crates/provisioner-task-libsql/src/lib.rs new file mode 100644 index 000000000..05cedf5dc --- /dev/null +++ b/crates/provisioner-task-libsql/src/lib.rs @@ -0,0 +1,660 @@ +//! The provisioner task table in `gateway.db`: one row per Task DVLS asked Gateway to run. +//! +//! The tables are named `provisioner_task_*`, and this crate alone creates and migrates them. Their schema version is +//! the `schema_version` row of `provisioner_task_metadata`, not the file's `PRAGMA user_version`, because other +//! features keep their own tables in the same file. + +#[macro_use] +extern crate tracing; + +use anyhow::{Context as _, bail}; +use async_trait::async_trait; +use libsql::{Connection, Row, TransactionBehavior, params}; +use provisioner_task::{ + CreateOutcome, NewProvisionerTask, ProvisionerTask, ProvisionerTaskState, ProvisionerTaskStore, +}; +use time::OffsetDateTime; +use tokio::sync::Mutex; +use uuid::Uuid; + +/// Every schema change ever made to the `provisioner_task_*` tables, oldest first. Only ever append to it. +const MIGRATIONS: &[&str] = &[include_str!("../migrations/01_provisioner_task.sql")]; + +const SCHEMA_VERSION_KEY: &str = "schema_version"; + +const COLUMNS: &str = + "id, kind, target, params, state, payload, attempts, created_at, started_at, finished_at, deadline_at"; + +/// Tasks kept in the `provisioner_task_records` table of `gateway.db`. +pub struct LibSqlProvisionerTaskStore { + conn: Mutex, +} + +impl LibSqlProvisionerTaskStore { + /// Opens the store on a `gateway.db` connection of its own, from `GatewayDb::connect`. + /// + /// Opening creates the `provisioner_task_*` tables or brings them up to date. + pub async fn open(conn: Connection) -> anyhow::Result { + migrate(&conn).await?; + + Ok(Self { conn: Mutex::new(conn) }) + } + + async fn find(conn: &Connection, id: Uuid) -> anyhow::Result> { + conn.query( + &format!("SELECT {COLUMNS} FROM provisioner_task_records WHERE id = ?1"), + params![id.to_string()], + ) + .await + .context("query task")? + .next() + .await + .context("read task")? + .as_ref() + .map(task_from_row) + .transpose() + } + + async fn find_unfinished(conn: &Connection, kind: &str, target: &str) -> anyhow::Result> { + conn.query( + &format!( + "SELECT {COLUMNS} FROM provisioner_task_records + WHERE kind = ?1 AND target = ?2 AND state IN ('queued', 'running')" + ), + params![kind, target], + ) + .await + .context("query unfinished task")? + .next() + .await + .context("read unfinished task")? + .as_ref() + .map(task_from_row) + .transpose() + } + + /// Fails `task` if it is unfinished and past its deadline, then returns it as stored. + async fn expire(conn: &Connection, task: ProvisionerTask, now: OffsetDateTime) -> anyhow::Result { + if task.state.is_finished() || now < task.deadline_at { + return Ok(task); + } + + let payload = serde_json::json!({ + "reason": "timed out", + "details": { "lastPayload": task.payload }, + "attempts": task.attempts, + }); + + conn.execute( + "UPDATE provisioner_task_records SET state = 'failed', payload = ?2, finished_at = ?3 + WHERE id = ?1 AND state IN ('queued', 'running')", + params![task.id.to_string(), payload.to_string(), now.unix_timestamp()], + ) + .await + .context("fail overdue task")?; + + warn!(task.id = %task.id, task.kind = %task.kind, attempts = task.attempts, "Task timed out"); + + Self::find(conn, task.id).await?.context("overdue task is missing") + } + + /// Ends the Task as `state` if it is currently in one of `from`, given as an SQL list such as `'running'`. + async fn finish( + &self, + id: Uuid, + from: &'static str, + state: ProvisionerTaskState, + payload: serde_json::Value, + now: OffsetDateTime, + ) -> anyhow::Result { + let conn = self.conn.lock().await; + let changed = conn + .execute( + &format!( + "UPDATE provisioner_task_records SET state = ?2, payload = ?3, finished_at = ?4 + WHERE id = ?1 AND state IN ({from})" + ), + params![ + id.to_string(), + state.as_str(), + payload.to_string(), + now.unix_timestamp() + ], + ) + .await + .context("finish task")?; + + Ok(changed == 1) + } +} + +#[async_trait] +impl ProvisionerTaskStore for LibSqlProvisionerTaskStore { + async fn create(&self, task: NewProvisionerTask, now: OffsetDateTime) -> anyhow::Result { + let conn = self.conn.lock().await; + let tx = conn + .transaction_with_behavior(TransactionBehavior::Immediate) + .await + .context("begin task creation")?; + + if let Some(existing) = Self::find(&tx, task.id).await? { + let existing = Self::expire(&tx, existing, now).await?; + let same_request = + existing.kind == task.kind && existing.target == task.target && existing.params == task.params; + tx.commit().await.context("commit task lookup")?; + + return Ok(if same_request { + CreateOutcome::Existing(existing) + } else { + CreateOutcome::IdConflict(existing) + }); + } + + if let Some(unfinished) = Self::find_unfinished(&tx, &task.kind, &task.target).await? { + let unfinished = Self::expire(&tx, unfinished, now).await?; + + if !unfinished.state.is_finished() { + return Ok(CreateOutcome::TargetBusy(unfinished)); + } + } + + tx.execute( + "INSERT INTO provisioner_task_records (id, kind, target, params, state, payload, created_at, deadline_at) + VALUES (?1, ?2, ?3, ?4, 'queued', 'null', ?5, ?6)", + params![ + task.id.to_string(), + task.kind.clone(), + task.target.clone(), + task.params.to_string(), + now.unix_timestamp(), + task.deadline_at.unix_timestamp(), + ], + ) + .await + .context("insert task")?; + + let created = Self::find(&tx, task.id).await?.context("created task is missing")?; + tx.commit().await.context("commit task creation")?; + + Ok(CreateOutcome::Created(created)) + } + + async fn get(&self, id: Uuid, now: OffsetDateTime) -> anyhow::Result> { + let conn = self.conn.lock().await; + let tx = conn + .transaction_with_behavior(TransactionBehavior::Immediate) + .await + .context("begin task lookup")?; + + let Some(task) = Self::find(&tx, id).await? else { + return Ok(None); + }; + let task = Self::expire(&tx, task, now).await?; + tx.commit().await.context("commit task lookup")?; + + Ok(Some(task)) + } + + async fn start_attempt(&self, id: Uuid, payload: serde_json::Value, now: OffsetDateTime) -> anyhow::Result { + let conn = self.conn.lock().await; + let tx = conn + .transaction_with_behavior(TransactionBehavior::Immediate) + .await + .context("begin task attempt")?; + + let Some(task) = Self::find(&tx, id).await? else { + return Ok(false); + }; + let task = Self::expire(&tx, task, now).await?; + let runs = !task.state.is_finished(); + + if runs { + tx.execute( + "UPDATE provisioner_task_records + SET state = 'running', payload = ?2, attempts = attempts + 1, started_at = COALESCE(started_at, ?3) + WHERE id = ?1", + params![id.to_string(), payload.to_string(), now.unix_timestamp()], + ) + .await + .context("start task attempt")?; + } + + tx.commit().await.context("commit task attempt")?; + + Ok(runs) + } + + async fn update_running(&self, id: Uuid, payload: serde_json::Value) -> anyhow::Result { + let conn = self.conn.lock().await; + let changed = conn + .execute( + "UPDATE provisioner_task_records SET payload = ?2 WHERE id = ?1 AND state = 'running'", + params![id.to_string(), payload.to_string()], + ) + .await + .context("update running task")?; + + Ok(changed == 1) + } + + async fn succeed(&self, id: Uuid, payload: serde_json::Value, now: OffsetDateTime) -> anyhow::Result { + self.finish(id, "'running'", ProvisionerTaskState::Succeeded, payload, now) + .await + } + + async fn fail(&self, id: Uuid, payload: serde_json::Value, now: OffsetDateTime) -> anyhow::Result { + self.finish(id, "'queued', 'running'", ProvisionerTaskState::Failed, payload, now) + .await + } +} + +fn task_from_row(row: &Row) -> anyhow::Result { + let timestamp = |index: i32, name: &str| -> anyhow::Result { + let seconds = row.get::(index).with_context(|| format!("decode task {name}"))?; + OffsetDateTime::from_unix_timestamp(seconds).with_context(|| format!("invalid task {name}")) + }; + let optional_timestamp = |index: i32, name: &str| -> anyhow::Result> { + row.get::>(index) + .with_context(|| format!("decode task {name}"))? + .map(|seconds| OffsetDateTime::from_unix_timestamp(seconds).with_context(|| format!("invalid task {name}"))) + .transpose() + }; + let json = |index: i32, name: &str| -> anyhow::Result { + let text = row + .get::(index) + .with_context(|| format!("decode task {name}"))?; + serde_json::from_str(&text).with_context(|| format!("parse task {name}")) + }; + + Ok(ProvisionerTask { + id: row + .get::(0) + .context("decode task id")? + .parse() + .context("parse task id")?, + kind: row.get::(1).context("decode task kind")?, + target: row.get::(2).context("decode task target")?, + params: json(3, "params")?, + state: row + .get::(4) + .context("decode task state")? + .parse() + .context("parse task state")?, + payload: json(5, "payload")?, + attempts: u32::try_from(row.get::(6).context("decode task attempts")?).context("invalid task attempts")?, + created_at: timestamp(7, "creation time")?, + started_at: optional_timestamp(8, "start time")?, + finished_at: optional_timestamp(9, "finish time")?, + deadline_at: timestamp(10, "deadline")?, + }) +} + +/// Creates the `provisioner_task_*` tables or brings them up to date. +async fn migrate(conn: &Connection) -> anyhow::Result<()> { + loop { + // Another connection may migrate at the same time, so the version is read under the write lock. + let tx = conn + .transaction_with_behavior(TransactionBehavior::Immediate) + .await + .context("begin provisioner task migration")?; + let version = schema_version(&tx).await?; + + let Some(migration) = MIGRATIONS.get(version) else { + tx.rollback().await.context("end provisioner task migration")?; + + if MIGRATIONS.len() < version { + bail!( + "provisioner task schema version {version} is newer than supported version {}", + MIGRATIONS.len() + ); + } + + return Ok(()); + }; + + let version = version + 1; + tx.execute_batch(migration) + .await + .with_context(|| format!("apply provisioner task migration {version}"))?; + tx.execute( + "INSERT OR REPLACE INTO provisioner_task_metadata (key, value) VALUES (?1, ?2)", + params![SCHEMA_VERSION_KEY, i64::try_from(version)?], + ) + .await + .with_context(|| format!("record provisioner task migration {version}"))?; + tx.commit() + .await + .with_context(|| format!("commit provisioner task migration {version}"))?; + } +} + +/// Reads the version of the `provisioner_task_*` tables; 0 when there are none yet. +async fn schema_version(conn: &Connection) -> anyhow::Result { + let metadata_exists = conn + .query( + "SELECT 1 FROM sqlite_master WHERE type = 'table' AND name = 'provisioner_task_metadata'", + (), + ) + .await + .context("look for the provisioner task tables")? + .next() + .await + .context("read the provisioner task tables")? + .is_some(); + + if !metadata_exists { + return Ok(0); + } + + let version = conn + .query( + "SELECT value FROM provisioner_task_metadata WHERE key = ?1", + params![SCHEMA_VERSION_KEY], + ) + .await + .context("query provisioner task schema version")? + .next() + .await + .context("read provisioner task schema version")? + .context("provisioner task tables have no schema version")? + .get::(0) + .context("decode provisioner task schema version")?; + + usize::try_from(version).context("provisioner task schema version is out of range") +} +#[cfg(test)] +mod tests { + use time::Duration; + + use super::*; + + async fn store() -> LibSqlProvisionerTaskStore { + let conn = gateway_db::GatewayDb::open_path(":memory:") + .await + .expect("open gateway database") + .connect() + .await + .expect("connect to gateway database"); + LibSqlProvisionerTaskStore::open(conn) + .await + .expect("open provisioner task store") + } + + fn now() -> OffsetDateTime { + OffsetDateTime::from_unix_timestamp(1_800_000_000).expect("valid timestamp") + } + + fn new_task(target: &str) -> NewProvisionerTask { + NewProvisionerTask { + id: Uuid::new_v4(), + kind: String::from("recording.ai-analysis"), + target: String::from(target), + params: serde_json::json!({ "provider": "openai", "model": "gpt-6-luna" }), + deadline_at: now() + Duration::hours(2), + } + } + + async fn create(store: &LibSqlProvisionerTaskStore, task: NewProvisionerTask) -> CreateOutcome { + store.create(task, now()).await.expect("create task") + } + + #[tokio::test] + async fn created_task_is_queued() { + let store = store().await; + let task = new_task("session-1"); + + let CreateOutcome::Created(created) = create(&store, task.clone()).await else { + panic!("task must be created"); + }; + + assert_eq!(created.state, ProvisionerTaskState::Queued); + assert_eq!(created.attempts, 0); + assert_eq!(created.params, task.params); + assert_eq!(store.get(task.id, now()).await.expect("get task"), Some(created)); + } + + #[tokio::test] + async fn same_request_again_returns_the_existing_task() { + let store = store().await; + let task = new_task("session-1"); + create(&store, task.clone()).await; + + assert!(matches!(create(&store, task).await, CreateOutcome::Existing(_))); + } + + #[tokio::test] + async fn same_id_with_other_parameters_conflicts() { + let store = store().await; + let task = new_task("session-1"); + create(&store, task.clone()).await; + + let mut other = task; + other.params = serde_json::json!({ "provider": "anthropic" }); + + assert!(matches!(create(&store, other).await, CreateOutcome::IdConflict(_))); + } + + #[tokio::test] + async fn one_unfinished_task_per_kind_and_target() { + let store = store().await; + let first = new_task("session-1"); + create(&store, first.clone()).await; + + let CreateOutcome::TargetBusy(busy) = create(&store, new_task("session-1")).await else { + panic!("target must be busy"); + }; + assert_eq!(busy.id, first.id); + assert!(matches!( + create(&store, new_task("session-2")).await, + CreateOutcome::Created(_) + )); + + store + .fail(first.id, serde_json::json!({}), now()) + .await + .expect("finish task"); + assert!(matches!( + create(&store, new_task("session-1")).await, + CreateOutcome::Created(_) + )); + } + + #[tokio::test] + async fn only_a_running_task_succeeds() { + let store = store().await; + let task = new_task("session-1"); + create(&store, task.clone()).await; + + assert!( + !store + .succeed(task.id, serde_json::json!({}), now()) + .await + .expect("succeed queued task") + ); + assert_eq!( + store.get(task.id, now()).await.expect("get task").expect("task").state, + ProvisionerTaskState::Queued + ); + + assert!( + store + .fail(task.id, serde_json::json!({ "reason": "not queued" }), now()) + .await + .expect("fail queued task") + ); + assert_eq!( + store.get(task.id, now()).await.expect("get task").expect("task").state, + ProvisionerTaskState::Failed + ); + } + + #[tokio::test] + async fn overdue_task_is_failed_and_frees_its_target() { + let store = store().await; + let mut overdue = new_task("session-1"); + overdue.deadline_at = now() - Duration::minutes(1); + create(&store, overdue.clone()).await; + store + .start_attempt( + overdue.id, + serde_json::json!({ "done": 3, "total": 10 }), + now() - Duration::minutes(2), + ) + .await + .expect("start attempt"); + + assert!(matches!( + create(&store, new_task("session-1")).await, + CreateOutcome::Created(_) + )); + + let failed = store + .get(overdue.id, now()) + .await + .expect("get task") + .expect("task exists"); + assert_eq!(failed.state, ProvisionerTaskState::Failed); + assert_eq!(failed.payload["reason"], "timed out"); + assert_eq!(failed.payload["attempts"], 1); + assert_eq!(failed.payload["details"]["lastPayload"]["done"], 3); + } + + #[tokio::test] + async fn attempts_running_payload_and_finish() { + let store = store().await; + let task = new_task("session-1"); + create(&store, task.clone()).await; + + assert!( + store + .start_attempt(task.id, serde_json::json!({ "done": 0 }), now()) + .await + .expect("start") + ); + assert!( + store + .update_running(task.id, serde_json::json!({ "done": 1 })) + .await + .expect("update") + ); + assert!( + store + .start_attempt(task.id, serde_json::json!({ "done": 1 }), now()) + .await + .expect("retry") + ); + + let running = store.get(task.id, now()).await.expect("get task").expect("task exists"); + assert_eq!(running.state, ProvisionerTaskState::Running); + assert_eq!(running.attempts, 2); + assert_eq!(running.started_at, Some(now())); + + assert!( + store + .fail(task.id, serde_json::json!({ "reason": "boom" }), now()) + .await + .expect("finish") + ); + assert!( + !store + .succeed(task.id, serde_json::json!({}), now()) + .await + .expect("finish again") + ); + assert!( + !store + .start_attempt(task.id, serde_json::json!({}), now()) + .await + .expect("start after finish") + ); + assert!( + !store + .update_running(task.id, serde_json::json!({})) + .await + .expect("update after finish") + ); + + let finished = store.get(task.id, now()).await.expect("get task").expect("task exists"); + assert_eq!(finished.state, ProvisionerTaskState::Failed); + assert_eq!(finished.payload["reason"], "boom"); + assert_eq!(finished.finished_at, Some(now())); + } + + #[tokio::test] + async fn overdue_task_is_failed_when_read_or_started() { + let store = store().await; + let mut overdue = new_task("session-1"); + overdue.deadline_at = now() - Duration::minutes(1); + create(&store, overdue.clone()).await; + + assert!( + !store + .start_attempt(overdue.id, serde_json::json!({}), now()) + .await + .expect("start attempt") + ); + let read = store + .get(overdue.id, now()) + .await + .expect("get task") + .expect("task exists"); + assert_eq!(read.state, ProvisionerTaskState::Failed); + assert_eq!(read.payload["reason"], "timed out"); + + let CreateOutcome::Existing(again) = create(&store, overdue).await else { + panic!("same request must return the existing task"); + }; + assert_eq!(again.state, ProvisionerTaskState::Failed); + } + + #[tokio::test] + async fn unknown_task() { + let store = store().await; + + assert_eq!(store.get(Uuid::new_v4(), now()).await.expect("get task"), None); + assert!( + !store + .start_attempt(Uuid::new_v4(), serde_json::json!({}), now()) + .await + .expect("start attempt") + ); + } + + #[tokio::test] + async fn new_tables_record_their_schema_version_in_metadata() { + let store = store().await; + let conn = store.conn.lock().await; + + assert_eq!( + schema_version(&conn).await.expect("read schema version"), + MIGRATIONS.len() + ); + } + + #[tokio::test] + async fn newer_schema_is_rejected() { + let temp_dir = tempfile::tempdir().expect("create temporary directory"); + let path = temp_dir.path().join("gateway.db"); + let db = gateway_db::GatewayDb::open_path(path.to_str().expect("temporary path is UTF-8")) + .await + .expect("open gateway database"); + drop( + LibSqlProvisionerTaskStore::open(db.connect().await.expect("connect")) + .await + .expect("open provisioner task store"), + ); + let conn = db.connect().await.expect("connect"); + conn.execute( + "UPDATE provisioner_task_metadata SET value = 99 WHERE key = ?1", + params![SCHEMA_VERSION_KEY], + ) + .await + .expect("set unsupported schema version"); + + let error = LibSqlProvisionerTaskStore::open(conn) + .await + .err() + .expect("newer schema must be rejected"); + assert!(error.to_string().contains("newer than supported")); + } +} diff --git a/crates/provisioner-task/Cargo.toml b/crates/provisioner-task/Cargo.toml new file mode 100644 index 000000000..e47db4a96 --- /dev/null +++ b/crates/provisioner-task/Cargo.toml @@ -0,0 +1,16 @@ +[package] +name = "provisioner-task" +version = "0.0.0" +edition = "2024" +authors = ["Devolutions Inc. "] +publish = false + +[lints] +workspace = true + +[dependencies] +anyhow = "1" +async-trait = "0.1" +serde_json = "1" +time = { version = "0.3", default-features = false } +uuid = "1.23" diff --git a/crates/provisioner-task/src/lib.rs b/crates/provisioner-task/src/lib.rs new file mode 100644 index 000000000..1f44295d8 --- /dev/null +++ b/crates/provisioner-task/src/lib.rs @@ -0,0 +1,114 @@ +//! Tasks: the permanent record of work the provisioner (DVLS, Hub) asked Gateway to do. +//! +//! A Task remembers what was asked, how it is going and how it ended, and it stays around for auditing. The work +//! itself runs as a job on the job queue; that job is gone once it finishes, but its Task is not. + +use std::sync::Arc; + +use async_trait::async_trait; +use time::OffsetDateTime; +use uuid::Uuid; + +pub type DynProvisionerTaskStore = Arc; + +/// Where a Task is in its life. Every state carries a payload whose shape belongs to the Task kind. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum ProvisionerTaskState { + Queued, + Running, + Succeeded, + Failed, +} + +impl ProvisionerTaskState { + pub fn as_str(self) -> &'static str { + match self { + Self::Queued => "queued", + Self::Running => "running", + Self::Succeeded => "succeeded", + Self::Failed => "failed", + } + } + + pub fn is_finished(self) -> bool { + matches!(self, Self::Succeeded | Self::Failed) + } +} + +impl std::str::FromStr for ProvisionerTaskState { + type Err = anyhow::Error; + + fn from_str(value: &str) -> anyhow::Result { + match value { + "queued" => Ok(Self::Queued), + "running" => Ok(Self::Running), + "succeeded" => Ok(Self::Succeeded), + "failed" => Ok(Self::Failed), + _ => anyhow::bail!("unknown task state {value:?}"), + } + } +} + +#[derive(Debug, Clone, PartialEq)] +pub struct ProvisionerTask { + /// Chosen by the provisioner, so it can ask again without knowing whether the first request went through. + pub id: Uuid, + /// What kind of work this is, e.g. `recording.ai-analysis`. + pub kind: String, + /// What the work is about, e.g. a session ID. At most one unfinished Task per kind and target. + pub target: String, + /// What was asked, without secrets. + pub params: serde_json::Value, + pub state: ProvisionerTaskState, + pub payload: serde_json::Value, + pub attempts: u32, + pub created_at: OffsetDateTime, + pub started_at: Option, + pub finished_at: Option, + /// After this, an unfinished Task is failed, so a lost job can't block its target forever. + pub deadline_at: OffsetDateTime, +} + +#[derive(Debug, Clone, PartialEq)] +pub struct NewProvisionerTask { + pub id: Uuid, + pub kind: String, + pub target: String, + pub params: serde_json::Value, + pub deadline_at: OffsetDateTime, +} + +#[derive(Debug, Clone, PartialEq)] +pub enum CreateOutcome { + Created(ProvisionerTask), + /// The same request came in again; nothing new was created. + Existing(ProvisionerTask), + /// The ID is already used by a Task with another kind, target or parameters. + IdConflict(ProvisionerTask), + /// Another unfinished Task of the same kind already works on this target. + TargetBusy(ProvisionerTask), +} + +#[async_trait] +pub trait ProvisionerTaskStore: Send + Sync { + /// Records a new `queued` Task, unless the ID is known or its target is busy. + /// + /// Like every method that takes `now`, it first fails the Tasks it looks at that are past their deadline. + async fn create(&self, task: NewProvisionerTask, now: OffsetDateTime) -> anyhow::Result; + + async fn get(&self, id: Uuid, now: OffsetDateTime) -> anyhow::Result>; + + /// Marks the start of one more attempt: `running`, with one more attempt counted. + /// + /// Returns `false` if the Task is unknown, finished or past its deadline; the attempt must not run then. + async fn start_attempt(&self, id: Uuid, payload: serde_json::Value, now: OffsetDateTime) -> anyhow::Result; + + /// Replaces the payload of a `running` Task, e.g. to report progress. Returns `false` if it is not running. + async fn update_running(&self, id: Uuid, payload: serde_json::Value) -> anyhow::Result; + + /// Ends a `running` Task as `succeeded`. Returns `false` if it is not running, and changes nothing. + async fn succeed(&self, id: Uuid, payload: serde_json::Value, now: OffsetDateTime) -> anyhow::Result; + + /// Ends a `queued` or `running` Task as `failed`. Returns `false` if it was already finished, and changes nothing. + async fn fail(&self, id: Uuid, payload: serde_json::Value, now: OffsetDateTime) -> anyhow::Result; +} From 3d4fc7519946703bdc9a68793e73ed9c39112725 Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Wed, 7 Oct 2026 16:09:54 -0400 Subject: [PATCH 2/9] docs(dgw): add the provisioner task intent --- crates/provisioner-task/INTENT.md | 23 +++++++++++++++++++++++ 1 file changed, 23 insertions(+) create mode 100644 crates/provisioner-task/INTENT.md diff --git a/crates/provisioner-task/INTENT.md b/crates/provisioner-task/INTENT.md new file mode 100644 index 000000000..08888362a --- /dev/null +++ b/crates/provisioner-task/INTENT.md @@ -0,0 +1,23 @@ +# Background +Provisioner tasks are tasks that are assigned to the Gateway. + + +# Invariants +Each provisioner task will have a task id. +Each provisioner task will have the following state: +- Queued +- Running {task specific payload} +- Succeeded {task specific payload} +- Failed {task specific payload} +The state transitions are as follows: +- Queued -> Running {task specific payload} +- Running {task specific payload} -> Succeeded {task specific payload} +- Running {task specific payload} -> Failed {task specific payload} +- Running {task specific payload} -> Running {task specific payload} (update task specific payload) +- Queued -> Failed {task specific payload} + +Tasks will have a deadline which is relative to the time it is received by the Gateway. + +# Footnote +Provisioner tasks and jobs are not the same thing. Task has its own state, the job queue is merely for scheduling and retrying. + From 7e2bec0fec4f0da4eba12de6937817d9b2e4aec3 Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Thu, 8 Oct 2026 12:40:01 -0400 Subject: [PATCH 3/9] fix(dgw): expire overdue provisioner tasks before finishing them succeed and fail ended a Task without looking at its deadline, so an overdue running Task could still succeed, or have its timeout replaced by the caller's failure, although every method that takes now is meant to fail overdue Tasks first. Both now expire the Task and end it in the same transaction, and return false when the timeout wins. Also covers the main path, a running Task that succeeds with its result. --- crates/provisioner-task-libsql/src/lib.rs | 73 ++++++++++++++++++++++- 1 file changed, 72 insertions(+), 1 deletion(-) diff --git a/crates/provisioner-task-libsql/src/lib.rs b/crates/provisioner-task-libsql/src/lib.rs index 05cedf5dc..45fa93b8b 100644 --- a/crates/provisioner-task-libsql/src/lib.rs +++ b/crates/provisioner-task-libsql/src/lib.rs @@ -99,6 +99,8 @@ impl LibSqlProvisionerTaskStore { } /// Ends the Task as `state` if it is currently in one of `from`, given as an SQL list such as `'running'`. + /// + /// A Task past its deadline is failed as timed out first, and then not ended again. async fn finish( &self, id: Uuid, @@ -108,7 +110,17 @@ impl LibSqlProvisionerTaskStore { now: OffsetDateTime, ) -> anyhow::Result { let conn = self.conn.lock().await; - let changed = conn + let tx = conn + .transaction_with_behavior(TransactionBehavior::Immediate) + .await + .context("begin task finish")?; + + let Some(task) = Self::find(&tx, id).await? else { + return Ok(false); + }; + Self::expire(&tx, task, now).await?; + + let changed = tx .execute( &format!( "UPDATE provisioner_task_records SET state = ?2, payload = ?3, finished_at = ?4 @@ -123,6 +135,7 @@ impl LibSqlProvisionerTaskStore { ) .await .context("finish task")?; + tx.commit().await.context("commit task finish")?; Ok(changed == 1) } @@ -580,6 +593,64 @@ mod tests { assert_eq!(finished.finished_at, Some(now())); } + #[tokio::test] + async fn running_task_succeeds_with_its_result() { + let store = store().await; + let task = new_task("session-1"); + create(&store, task.clone()).await; + store + .start_attempt(task.id, serde_json::json!({ "done": 0 }), now()) + .await + .expect("start"); + + assert!( + store + .succeed(task.id, serde_json::json!({ "log": "ai-analysis-0.slog" }), now()) + .await + .expect("succeed") + ); + + let succeeded = store.get(task.id, now()).await.expect("get task").expect("task exists"); + assert_eq!(succeeded.state, ProvisionerTaskState::Succeeded); + assert_eq!(succeeded.payload["log"], "ai-analysis-0.slog"); + assert_eq!(succeeded.finished_at, Some(now())); + } + + #[tokio::test] + async fn overdue_task_cannot_succeed_or_fail_with_another_reason() { + let store = store().await; + let mut overdue = new_task("session-1"); + overdue.deadline_at = now() + Duration::minutes(1); + create(&store, overdue.clone()).await; + store + .start_attempt(overdue.id, serde_json::json!({ "done": 3 }), now()) + .await + .expect("start"); + let late = now() + Duration::minutes(2); + + assert!( + !store + .succeed(overdue.id, serde_json::json!({ "log": "late" }), late) + .await + .expect("succeed late") + ); + assert!( + !store + .fail(overdue.id, serde_json::json!({ "reason": "boom" }), late) + .await + .expect("fail late") + ); + + let failed = store + .get(overdue.id, late) + .await + .expect("get task") + .expect("task exists"); + assert_eq!(failed.state, ProvisionerTaskState::Failed); + assert_eq!(failed.payload["reason"], "timed out"); + assert_eq!(failed.payload["details"]["lastPayload"]["done"], 3); + } + #[tokio::test] async fn overdue_task_is_failed_when_read_or_started() { let store = store().await; From 26134983e986143328278dcf6e42c64583103585 Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Thu, 8 Oct 2026 13:54:51 -0400 Subject: [PATCH 4/9] feat(dgw): tie each provisioner task to the job that runs it A Task now records the token of the job allowed to run it. An attempt only starts with that token, so a stale or duplicate job finds the Task superseded and stops. start_attempt says why an attempt cannot run: finished, superseded or unknown, and returns the started Task. replace_job_token hands an unfinished Task to a new job, and list_unfinished returns the Tasks still to run, so Gateway can queue them again at startup without ever running one twice. --- .../migrations/01_provisioner_task.sql | 1 + crates/provisioner-task-libsql/src/lib.rs | 266 +++++++++++++++--- crates/provisioner-task/src/lib.rs | 37 ++- 3 files changed, 263 insertions(+), 41 deletions(-) diff --git a/crates/provisioner-task-libsql/migrations/01_provisioner_task.sql b/crates/provisioner-task-libsql/migrations/01_provisioner_task.sql index f37e3b4dc..09149315b 100644 --- a/crates/provisioner-task-libsql/migrations/01_provisioner_task.sql +++ b/crates/provisioner-task-libsql/migrations/01_provisioner_task.sql @@ -10,6 +10,7 @@ CREATE TABLE provisioner_task_records ( started_at INTEGER, finished_at INTEGER, deadline_at INTEGER NOT NULL, + job_token TEXT, CHECK (state IN ('queued', 'running', 'succeeded', 'failed')), CHECK ((state IN ('succeeded', 'failed')) = (finished_at IS NOT NULL)), CHECK (json_valid(params) AND json_valid(payload)), diff --git a/crates/provisioner-task-libsql/src/lib.rs b/crates/provisioner-task-libsql/src/lib.rs index 45fa93b8b..24f19a1e0 100644 --- a/crates/provisioner-task-libsql/src/lib.rs +++ b/crates/provisioner-task-libsql/src/lib.rs @@ -11,7 +11,7 @@ use anyhow::{Context as _, bail}; use async_trait::async_trait; use libsql::{Connection, Row, TransactionBehavior, params}; use provisioner_task::{ - CreateOutcome, NewProvisionerTask, ProvisionerTask, ProvisionerTaskState, ProvisionerTaskStore, + AttemptStart, CreateOutcome, NewProvisionerTask, ProvisionerTask, ProvisionerTaskState, ProvisionerTaskStore, }; use time::OffsetDateTime; use tokio::sync::Mutex; @@ -23,7 +23,7 @@ const MIGRATIONS: &[&str] = &[include_str!("../migrations/01_provisioner_task.sq const SCHEMA_VERSION_KEY: &str = "schema_version"; const COLUMNS: &str = - "id, kind, target, params, state, payload, attempts, created_at, started_at, finished_at, deadline_at"; + "id, kind, target, params, state, payload, attempts, created_at, started_at, finished_at, deadline_at, job_token"; /// Tasks kept in the `provisioner_task_records` table of `gateway.db`. pub struct LibSqlProvisionerTaskStore { @@ -172,8 +172,9 @@ impl ProvisionerTaskStore for LibSqlProvisionerTaskStore { } tx.execute( - "INSERT INTO provisioner_task_records (id, kind, target, params, state, payload, created_at, deadline_at) - VALUES (?1, ?2, ?3, ?4, 'queued', 'null', ?5, ?6)", + "INSERT INTO provisioner_task_records + (id, kind, target, params, state, payload, created_at, deadline_at, job_token) + VALUES (?1, ?2, ?3, ?4, 'queued', 'null', ?5, ?6, ?7)", params![ task.id.to_string(), task.kind.clone(), @@ -181,6 +182,7 @@ impl ProvisionerTaskStore for LibSqlProvisionerTaskStore { task.params.to_string(), now.unix_timestamp(), task.deadline_at.unix_timestamp(), + task.job_token.to_string(), ], ) .await @@ -208,33 +210,111 @@ impl ProvisionerTaskStore for LibSqlProvisionerTaskStore { Ok(Some(task)) } - async fn start_attempt(&self, id: Uuid, payload: serde_json::Value, now: OffsetDateTime) -> anyhow::Result { + async fn start_attempt( + &self, + id: Uuid, + job_token: Uuid, + payload: serde_json::Value, + now: OffsetDateTime, + ) -> anyhow::Result { let conn = self.conn.lock().await; let tx = conn .transaction_with_behavior(TransactionBehavior::Immediate) .await .context("begin task attempt")?; + let Some(task) = Self::find(&tx, id).await? else { + return Ok(AttemptStart::Unknown); + }; + + if task.job_token != Some(job_token) { + return Ok(AttemptStart::Superseded); + } + + let task = Self::expire(&tx, task, now).await?; + + if task.state.is_finished() { + tx.commit().await.context("commit task attempt")?; + return Ok(AttemptStart::Finished(task)); + } + + tx.execute( + "UPDATE provisioner_task_records + SET state = 'running', payload = ?2, attempts = attempts + 1, started_at = COALESCE(started_at, ?3) + WHERE id = ?1", + params![id.to_string(), payload.to_string(), now.unix_timestamp()], + ) + .await + .context("start task attempt")?; + + let started = Self::find(&tx, id).await?.context("started task is missing")?; + tx.commit().await.context("commit task attempt")?; + + Ok(AttemptStart::Started(started)) + } + + async fn replace_job_token(&self, id: Uuid, job_token: Uuid, now: OffsetDateTime) -> anyhow::Result { + let conn = self.conn.lock().await; + let tx = conn + .transaction_with_behavior(TransactionBehavior::Immediate) + .await + .context("begin job token replacement")?; + let Some(task) = Self::find(&tx, id).await? else { return Ok(false); }; let task = Self::expire(&tx, task, now).await?; - let runs = !task.state.is_finished(); - if runs { + let replaced = !task.state.is_finished(); + + if replaced { tx.execute( - "UPDATE provisioner_task_records - SET state = 'running', payload = ?2, attempts = attempts + 1, started_at = COALESCE(started_at, ?3) - WHERE id = ?1", - params![id.to_string(), payload.to_string(), now.unix_timestamp()], + "UPDATE provisioner_task_records SET job_token = ?2 WHERE id = ?1", + params![id.to_string(), job_token.to_string()], ) .await - .context("start task attempt")?; + .context("replace task job token")?; } - tx.commit().await.context("commit task attempt")?; + tx.commit().await.context("commit job token replacement")?; - Ok(runs) + Ok(replaced) + } + + async fn list_unfinished(&self, now: OffsetDateTime) -> anyhow::Result> { + let conn = self.conn.lock().await; + let tx = conn + .transaction_with_behavior(TransactionBehavior::Immediate) + .await + .context("begin unfinished task listing")?; + + let mut rows = tx + .query( + &format!( + "SELECT {COLUMNS} FROM provisioner_task_records + WHERE state IN ('queued', 'running') ORDER BY created_at" + ), + (), + ) + .await + .context("query unfinished tasks")?; + let mut tasks = Vec::new(); + while let Some(row) = rows.next().await.context("read unfinished task")? { + tasks.push(task_from_row(&row)?); + } + drop(rows); + + let mut unfinished = Vec::with_capacity(tasks.len()); + for task in tasks { + let task = Self::expire(&tx, task, now).await?; + if !task.state.is_finished() { + unfinished.push(task); + } + } + + tx.commit().await.context("commit unfinished task listing")?; + + Ok(unfinished) } async fn update_running(&self, id: Uuid, payload: serde_json::Value) -> anyhow::Result { @@ -299,6 +379,11 @@ fn task_from_row(row: &Row) -> anyhow::Result { started_at: optional_timestamp(8, "start time")?, finished_at: optional_timestamp(9, "finish time")?, deadline_at: timestamp(10, "deadline")?, + job_token: row + .get::>(11) + .context("decode task job token")? + .map(|token| token.parse().context("parse task job token")) + .transpose()?, }) } @@ -404,6 +489,7 @@ mod tests { target: String::from(target), params: serde_json::json!({ "provider": "openai", "model": "gpt-6-luna" }), deadline_at: now() + Duration::hours(2), + job_token: Uuid::new_v4(), } } @@ -510,6 +596,7 @@ mod tests { store .start_attempt( overdue.id, + overdue.job_token, serde_json::json!({ "done": 3, "total": 10 }), now() - Duration::minutes(2), ) @@ -538,24 +625,26 @@ mod tests { let task = new_task("session-1"); create(&store, task.clone()).await; - assert!( + assert!(matches!( store - .start_attempt(task.id, serde_json::json!({ "done": 0 }), now()) + .start_attempt(task.id, task.job_token, serde_json::json!({ "done": 0 }), now()) .await - .expect("start") - ); + .expect("start"), + AttemptStart::Started(_) + )); assert!( store .update_running(task.id, serde_json::json!({ "done": 1 })) .await .expect("update") ); - assert!( + assert!(matches!( store - .start_attempt(task.id, serde_json::json!({ "done": 1 }), now()) + .start_attempt(task.id, task.job_token, serde_json::json!({ "done": 1 }), now()) .await - .expect("retry") - ); + .expect("retry"), + AttemptStart::Started(_) + )); let running = store.get(task.id, now()).await.expect("get task").expect("task exists"); assert_eq!(running.state, ProvisionerTaskState::Running); @@ -574,12 +663,13 @@ mod tests { .await .expect("finish again") ); - assert!( - !store - .start_attempt(task.id, serde_json::json!({}), now()) + assert!(matches!( + store + .start_attempt(task.id, task.job_token, serde_json::json!({}), now()) .await - .expect("start after finish") - ); + .expect("start after finish"), + AttemptStart::Finished(_) + )); assert!( !store .update_running(task.id, serde_json::json!({})) @@ -599,7 +689,7 @@ mod tests { let task = new_task("session-1"); create(&store, task.clone()).await; store - .start_attempt(task.id, serde_json::json!({ "done": 0 }), now()) + .start_attempt(task.id, task.job_token, serde_json::json!({ "done": 0 }), now()) .await .expect("start"); @@ -623,7 +713,7 @@ mod tests { overdue.deadline_at = now() + Duration::minutes(1); create(&store, overdue.clone()).await; store - .start_attempt(overdue.id, serde_json::json!({ "done": 3 }), now()) + .start_attempt(overdue.id, overdue.job_token, serde_json::json!({ "done": 3 }), now()) .await .expect("start"); let late = now() + Duration::minutes(2); @@ -658,12 +748,13 @@ mod tests { overdue.deadline_at = now() - Duration::minutes(1); create(&store, overdue.clone()).await; - assert!( - !store - .start_attempt(overdue.id, serde_json::json!({}), now()) + assert!(matches!( + store + .start_attempt(overdue.id, overdue.job_token, serde_json::json!({}), now()) .await - .expect("start attempt") - ); + .expect("start attempt"), + AttemptStart::Finished(_) + )); let read = store .get(overdue.id, now()) .await @@ -679,15 +770,114 @@ mod tests { } #[tokio::test] - async fn unknown_task() { + async fn only_the_job_holding_the_token_runs_the_task() { let store = store().await; + let task = new_task("session-1"); + create(&store, task.clone()).await; - assert_eq!(store.get(Uuid::new_v4(), now()).await.expect("get task"), None); + assert_eq!( + store + .start_attempt(task.id, Uuid::new_v4(), serde_json::json!({}), now()) + .await + .expect("start with another token"), + AttemptStart::Superseded + ); + assert_eq!( + store + .get(task.id, now()) + .await + .expect("get task") + .expect("task") + .attempts, + 0, + "a superseded job counts no attempt" + ); + + let new_token = Uuid::new_v4(); + assert!( + store + .replace_job_token(task.id, new_token, now()) + .await + .expect("replace token") + ); + assert_eq!( + store + .start_attempt(task.id, task.job_token, serde_json::json!({}), now()) + .await + .expect("start with the old token"), + AttemptStart::Superseded + ); + assert!(matches!( + store + .start_attempt(task.id, new_token, serde_json::json!({}), now()) + .await + .expect("start with the new token"), + AttemptStart::Started(started) if started.job_token == Some(new_token) && started.attempts == 1 + )); + + store + .fail(task.id, serde_json::json!({ "reason": "boom" }), now()) + .await + .expect("fail"); assert!( !store - .start_attempt(Uuid::new_v4(), serde_json::json!({}), now()) + .replace_job_token(task.id, Uuid::new_v4(), now()) + .await + .expect("replace token of a finished task"), + "a finished Task gets no new job" + ); + } + + #[tokio::test] + async fn unfinished_tasks_are_listed_and_overdue_ones_failed() { + let store = store().await; + let queued = new_task("session-1"); + let running = new_task("session-2"); + let finished = new_task("session-3"); + let mut overdue = new_task("session-4"); + overdue.deadline_at = now() - Duration::minutes(1); + for task in [&queued, &running, &finished, &overdue] { + create(&store, task.clone()).await; + } + store + .start_attempt(running.id, running.job_token, serde_json::json!({}), now()) + .await + .expect("start"); + store + .fail(finished.id, serde_json::json!({ "reason": "boom" }), now()) + .await + .expect("fail"); + + let unfinished = store.list_unfinished(now()).await.expect("list"); + + let mut ids = unfinished.iter().map(|task| task.id).collect::>(); + ids.sort(); + let mut expected = vec![queued.id, running.id]; + expected.sort(); + assert_eq!(ids, expected); + assert_eq!( + store + // Read as of before its deadline, so only the listing could have failed it. + .get(overdue.id, now() - Duration::hours(1)) + .await + .expect("get task") + .expect("task") + .payload["reason"], + "timed out" + ); + } + + #[tokio::test] + async fn unknown_task() { + let store = store().await; + + assert_eq!(store.get(Uuid::new_v4(), now()).await.expect("get task"), None); + assert_eq!( + store + .start_attempt(Uuid::new_v4(), Uuid::new_v4(), serde_json::json!({}), now()) .await - .expect("start attempt") + .expect("start attempt"), + AttemptStart::Unknown ); } diff --git a/crates/provisioner-task/src/lib.rs b/crates/provisioner-task/src/lib.rs index 1f44295d8..251ae7f3d 100644 --- a/crates/provisioner-task/src/lib.rs +++ b/crates/provisioner-task/src/lib.rs @@ -67,6 +67,8 @@ pub struct ProvisionerTask { pub finished_at: Option, /// After this, an unfinished Task is failed, so a lost job can't block its target forever. pub deadline_at: OffsetDateTime, + /// The job that may run this Task: a job with another token is stale or a duplicate, and must not run it. + pub job_token: Option, } #[derive(Debug, Clone, PartialEq)] @@ -76,6 +78,8 @@ pub struct NewProvisionerTask { pub target: String, pub params: serde_json::Value, pub deadline_at: OffsetDateTime, + /// Token of the job queued for the new Task. + pub job_token: Uuid, } #[derive(Debug, Clone, PartialEq)] @@ -89,6 +93,19 @@ pub enum CreateOutcome { TargetBusy(ProvisionerTask), } +/// What a job finds when it comes to run an attempt of its Task. +#[derive(Debug, Clone, PartialEq)] +pub enum AttemptStart { + /// The attempt is recorded: the Task is `running` with one more attempt counted. + Started(ProvisionerTask), + /// The Task is finished, possibly just now because it is past its deadline. + Finished(ProvisionerTask), + /// Another job now runs the Task, so this one is stale or a duplicate. + Superseded, + /// No Task has this ID. + Unknown, +} + #[async_trait] pub trait ProvisionerTaskStore: Send + Sync { /// Records a new `queued` Task, unless the ID is known or its target is busy. @@ -98,10 +115,24 @@ pub trait ProvisionerTaskStore: Send + Sync { async fn get(&self, id: Uuid, now: OffsetDateTime) -> anyhow::Result>; - /// Marks the start of one more attempt: `running`, with one more attempt counted. + /// Marks the start of one more attempt by the job holding `job_token`: `running`, with one more attempt counted. /// - /// Returns `false` if the Task is unknown, finished or past its deadline; the attempt must not run then. - async fn start_attempt(&self, id: Uuid, payload: serde_json::Value, now: OffsetDateTime) -> anyhow::Result; + /// Only [`AttemptStart::Started`] lets the attempt run. + async fn start_attempt( + &self, + id: Uuid, + job_token: Uuid, + payload: serde_json::Value, + now: OffsetDateTime, + ) -> anyhow::Result; + + /// Hands an unfinished Task to a new job, so the job that had it stops at its next attempt. + /// + /// Returns `false` if the Task is unknown or finished, and changes nothing. + async fn replace_job_token(&self, id: Uuid, job_token: Uuid, now: OffsetDateTime) -> anyhow::Result; + + /// Every unfinished Task, after failing those past their deadline. + async fn list_unfinished(&self, now: OffsetDateTime) -> anyhow::Result>; /// Replaces the payload of a `running` Task, e.g. to report progress. Returns `false` if it is not running. async fn update_running(&self, id: Uuid, payload: serde_json::Value) -> anyhow::Result; From ac47396b4c6e5e52e871793a002f9164fad8f67d Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Thu, 8 Oct 2026 15:07:31 -0400 Subject: [PATCH 5/9] refactor(dgw): name the stored provisioner task a record ProvisionerTask becomes ProvisionerTaskRecord, like the provisioner_task_records table, and NewProvisionerTask becomes NewProvisionerTaskRecord. The name ProvisionerTask is left for the work itself, which Gateway defines next to the code that runs it. --- crates/provisioner-task-libsql/src/lib.rs | 33 ++++++++++++++--------- crates/provisioner-task/src/lib.rs | 24 +++++++++-------- 2 files changed, 34 insertions(+), 23 deletions(-) diff --git a/crates/provisioner-task-libsql/src/lib.rs b/crates/provisioner-task-libsql/src/lib.rs index 24f19a1e0..8e04b33e1 100644 --- a/crates/provisioner-task-libsql/src/lib.rs +++ b/crates/provisioner-task-libsql/src/lib.rs @@ -11,7 +11,8 @@ use anyhow::{Context as _, bail}; use async_trait::async_trait; use libsql::{Connection, Row, TransactionBehavior, params}; use provisioner_task::{ - AttemptStart, CreateOutcome, NewProvisionerTask, ProvisionerTask, ProvisionerTaskState, ProvisionerTaskStore, + AttemptStart, CreateOutcome, NewProvisionerTaskRecord, ProvisionerTaskRecord, ProvisionerTaskState, + ProvisionerTaskStore, }; use time::OffsetDateTime; use tokio::sync::Mutex; @@ -40,7 +41,7 @@ impl LibSqlProvisionerTaskStore { Ok(Self { conn: Mutex::new(conn) }) } - async fn find(conn: &Connection, id: Uuid) -> anyhow::Result> { + async fn find(conn: &Connection, id: Uuid) -> anyhow::Result> { conn.query( &format!("SELECT {COLUMNS} FROM provisioner_task_records WHERE id = ?1"), params![id.to_string()], @@ -55,7 +56,11 @@ impl LibSqlProvisionerTaskStore { .transpose() } - async fn find_unfinished(conn: &Connection, kind: &str, target: &str) -> anyhow::Result> { + async fn find_unfinished( + conn: &Connection, + kind: &str, + target: &str, + ) -> anyhow::Result> { conn.query( &format!( "SELECT {COLUMNS} FROM provisioner_task_records @@ -74,7 +79,11 @@ impl LibSqlProvisionerTaskStore { } /// Fails `task` if it is unfinished and past its deadline, then returns it as stored. - async fn expire(conn: &Connection, task: ProvisionerTask, now: OffsetDateTime) -> anyhow::Result { + async fn expire( + conn: &Connection, + task: ProvisionerTaskRecord, + now: OffsetDateTime, + ) -> anyhow::Result { if task.state.is_finished() || now < task.deadline_at { return Ok(task); } @@ -143,7 +152,7 @@ impl LibSqlProvisionerTaskStore { #[async_trait] impl ProvisionerTaskStore for LibSqlProvisionerTaskStore { - async fn create(&self, task: NewProvisionerTask, now: OffsetDateTime) -> anyhow::Result { + async fn create(&self, task: NewProvisionerTaskRecord, now: OffsetDateTime) -> anyhow::Result { let conn = self.conn.lock().await; let tx = conn .transaction_with_behavior(TransactionBehavior::Immediate) @@ -194,7 +203,7 @@ impl ProvisionerTaskStore for LibSqlProvisionerTaskStore { Ok(CreateOutcome::Created(created)) } - async fn get(&self, id: Uuid, now: OffsetDateTime) -> anyhow::Result> { + async fn get(&self, id: Uuid, now: OffsetDateTime) -> anyhow::Result> { let conn = self.conn.lock().await; let tx = conn .transaction_with_behavior(TransactionBehavior::Immediate) @@ -281,7 +290,7 @@ impl ProvisionerTaskStore for LibSqlProvisionerTaskStore { Ok(replaced) } - async fn list_unfinished(&self, now: OffsetDateTime) -> anyhow::Result> { + async fn list_unfinished(&self, now: OffsetDateTime) -> anyhow::Result> { let conn = self.conn.lock().await; let tx = conn .transaction_with_behavior(TransactionBehavior::Immediate) @@ -341,7 +350,7 @@ impl ProvisionerTaskStore for LibSqlProvisionerTaskStore { } } -fn task_from_row(row: &Row) -> anyhow::Result { +fn task_from_row(row: &Row) -> anyhow::Result { let timestamp = |index: i32, name: &str| -> anyhow::Result { let seconds = row.get::(index).with_context(|| format!("decode task {name}"))?; OffsetDateTime::from_unix_timestamp(seconds).with_context(|| format!("invalid task {name}")) @@ -359,7 +368,7 @@ fn task_from_row(row: &Row) -> anyhow::Result { serde_json::from_str(&text).with_context(|| format!("parse task {name}")) }; - Ok(ProvisionerTask { + Ok(ProvisionerTaskRecord { id: row .get::(0) .context("decode task id")? @@ -482,8 +491,8 @@ mod tests { OffsetDateTime::from_unix_timestamp(1_800_000_000).expect("valid timestamp") } - fn new_task(target: &str) -> NewProvisionerTask { - NewProvisionerTask { + fn new_task(target: &str) -> NewProvisionerTaskRecord { + NewProvisionerTaskRecord { id: Uuid::new_v4(), kind: String::from("recording.ai-analysis"), target: String::from(target), @@ -493,7 +502,7 @@ mod tests { } } - async fn create(store: &LibSqlProvisionerTaskStore, task: NewProvisionerTask) -> CreateOutcome { + async fn create(store: &LibSqlProvisionerTaskStore, task: NewProvisionerTaskRecord) -> CreateOutcome { store.create(task, now()).await.expect("create task") } diff --git a/crates/provisioner-task/src/lib.rs b/crates/provisioner-task/src/lib.rs index 251ae7f3d..d982ffc98 100644 --- a/crates/provisioner-task/src/lib.rs +++ b/crates/provisioner-task/src/lib.rs @@ -49,8 +49,9 @@ impl std::str::FromStr for ProvisionerTaskState { } } +/// What is kept of a Task: what was asked, how it goes and how it ended. #[derive(Debug, Clone, PartialEq)] -pub struct ProvisionerTask { +pub struct ProvisionerTaskRecord { /// Chosen by the provisioner, so it can ask again without knowing whether the first request went through. pub id: Uuid, /// What kind of work this is, e.g. `recording.ai-analysis`. @@ -71,8 +72,9 @@ pub struct ProvisionerTask { pub job_token: Option, } +/// The record of a Task about to be created. #[derive(Debug, Clone, PartialEq)] -pub struct NewProvisionerTask { +pub struct NewProvisionerTaskRecord { pub id: Uuid, pub kind: String, pub target: String, @@ -84,22 +86,22 @@ pub struct NewProvisionerTask { #[derive(Debug, Clone, PartialEq)] pub enum CreateOutcome { - Created(ProvisionerTask), + Created(ProvisionerTaskRecord), /// The same request came in again; nothing new was created. - Existing(ProvisionerTask), + Existing(ProvisionerTaskRecord), /// The ID is already used by a Task with another kind, target or parameters. - IdConflict(ProvisionerTask), + IdConflict(ProvisionerTaskRecord), /// Another unfinished Task of the same kind already works on this target. - TargetBusy(ProvisionerTask), + TargetBusy(ProvisionerTaskRecord), } /// What a job finds when it comes to run an attempt of its Task. #[derive(Debug, Clone, PartialEq)] pub enum AttemptStart { /// The attempt is recorded: the Task is `running` with one more attempt counted. - Started(ProvisionerTask), + Started(ProvisionerTaskRecord), /// The Task is finished, possibly just now because it is past its deadline. - Finished(ProvisionerTask), + Finished(ProvisionerTaskRecord), /// Another job now runs the Task, so this one is stale or a duplicate. Superseded, /// No Task has this ID. @@ -111,9 +113,9 @@ pub trait ProvisionerTaskStore: Send + Sync { /// Records a new `queued` Task, unless the ID is known or its target is busy. /// /// Like every method that takes `now`, it first fails the Tasks it looks at that are past their deadline. - async fn create(&self, task: NewProvisionerTask, now: OffsetDateTime) -> anyhow::Result; + async fn create(&self, task: NewProvisionerTaskRecord, now: OffsetDateTime) -> anyhow::Result; - async fn get(&self, id: Uuid, now: OffsetDateTime) -> anyhow::Result>; + async fn get(&self, id: Uuid, now: OffsetDateTime) -> anyhow::Result>; /// Marks the start of one more attempt by the job holding `job_token`: `running`, with one more attempt counted. /// @@ -132,7 +134,7 @@ pub trait ProvisionerTaskStore: Send + Sync { async fn replace_job_token(&self, id: Uuid, job_token: Uuid, now: OffsetDateTime) -> anyhow::Result; /// Every unfinished Task, after failing those past their deadline. - async fn list_unfinished(&self, now: OffsetDateTime) -> anyhow::Result>; + async fn list_unfinished(&self, now: OffsetDateTime) -> anyhow::Result>; /// Replaces the payload of a `running` Task, e.g. to report progress. Returns `false` if it is not running. async fn update_running(&self, id: Uuid, payload: serde_json::Value) -> anyhow::Result; From 7b123b4d8c54890b3570afd3817910337f279ed3 Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Thu, 8 Oct 2026 16:01:38 -0400 Subject: [PATCH 6/9] docs(dgw): use the AI analysis task kind in provisioner task examples --- crates/provisioner-task-libsql/src/lib.rs | 2 +- crates/provisioner-task/src/lib.rs | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/crates/provisioner-task-libsql/src/lib.rs b/crates/provisioner-task-libsql/src/lib.rs index 8e04b33e1..fef9fdc41 100644 --- a/crates/provisioner-task-libsql/src/lib.rs +++ b/crates/provisioner-task-libsql/src/lib.rs @@ -494,7 +494,7 @@ mod tests { fn new_task(target: &str) -> NewProvisionerTaskRecord { NewProvisionerTaskRecord { id: Uuid::new_v4(), - kind: String::from("recording.ai-analysis"), + kind: String::from("recording-ai-analysis"), target: String::from(target), params: serde_json::json!({ "provider": "openai", "model": "gpt-6-luna" }), deadline_at: now() + Duration::hours(2), diff --git a/crates/provisioner-task/src/lib.rs b/crates/provisioner-task/src/lib.rs index d982ffc98..9f69cd906 100644 --- a/crates/provisioner-task/src/lib.rs +++ b/crates/provisioner-task/src/lib.rs @@ -54,7 +54,7 @@ impl std::str::FromStr for ProvisionerTaskState { pub struct ProvisionerTaskRecord { /// Chosen by the provisioner, so it can ask again without knowing whether the first request went through. pub id: Uuid, - /// What kind of work this is, e.g. `recording.ai-analysis`. + /// What kind of work this is, e.g. `recording-ai-analysis`. pub kind: String, /// What the work is about, e.g. a session ID. At most one unfinished Task per kind and target. pub target: String, From 08f0823249e14f17970c6c4fc6ec4ad6c69aedf3 Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Thu, 8 Oct 2026 18:57:41 -0400 Subject: [PATCH 7/9] docs(dgw): add the provisioner task store intent --- crates/provisioner-task-libsql/INTENT.md | 21 +++++++++++++++++++++ 1 file changed, 21 insertions(+) create mode 100644 crates/provisioner-task-libsql/INTENT.md diff --git a/crates/provisioner-task-libsql/INTENT.md b/crates/provisioner-task-libsql/INTENT.md new file mode 100644 index 000000000..2abc5b13a --- /dev/null +++ b/crates/provisioner-task-libsql/INTENT.md @@ -0,0 +1,21 @@ +# Background +This crate stores provisioner Task records in `gateway.db`, implementing `ProvisionerTaskStore` from `provisioner-task`. +The states and transitions of a Task are defined in `../provisioner-task/INTENT.md`. + +# Invariants + +## Tables and migrations +The tables of this crate are named with the `provisioner_task_` prefix; the prefix is the compatibility boundary. +This crate alone creates and migrates its tables, and its migrations are append-only. +Its schema version is the `schema_version` row of `provisioner_task_metadata`, never the file's `PRAGMA user_version`. +This crate may depend on `gateway-db`, never the other way around. + +## Records +For a given kind and target, at most one record can be unfinished. + +A record past its deadline is failed. + +Each change of state is atomic: a record is never seen half moved from one state to another. + +## Secrets +Records never hold secrets: parameters and payloads are stored as given, in clear. \ No newline at end of file From d583a045e5facdda5c46cc01c99c8f208266132d Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Mon, 12 Oct 2026 20:49:54 -0400 Subject: [PATCH 8/9] fix(dgw): fail overdue tasks on progress updates and keep SQL in the adapter - `update_running` takes `now` and fails a Task past its deadline instead of writing progress after it, as the store intent requires. - `job_token` is `NOT NULL`: every record is created with one, so `Option` allowed a state the store never produces. - `finish` takes the allowed states as `ProvisionerTaskState`s instead of an SQL list, and the SQL encoding of a state moves into the libsql adapter. - Transactions that end early are rolled back explicitly, so a failed rollback on those paths is an error instead of a panic in the drop guard (errors that leave through `?` still rely on the drop guard). - Docs: a queued Task has no payload; a superseded job changes nothing; `get` may fail an overdue Task. Co-Authored-By: Claude Opus 5.5 (1M context) --- .../migrations/01_provisioner_task.sql | 2 +- crates/provisioner-task-libsql/src/lib.rs | 136 ++++++++++++------ crates/provisioner-task/src/lib.rs | 35 ++--- 3 files changed, 105 insertions(+), 68 deletions(-) diff --git a/crates/provisioner-task-libsql/migrations/01_provisioner_task.sql b/crates/provisioner-task-libsql/migrations/01_provisioner_task.sql index 09149315b..e39abac43 100644 --- a/crates/provisioner-task-libsql/migrations/01_provisioner_task.sql +++ b/crates/provisioner-task-libsql/migrations/01_provisioner_task.sql @@ -10,7 +10,7 @@ CREATE TABLE provisioner_task_records ( started_at INTEGER, finished_at INTEGER, deadline_at INTEGER NOT NULL, - job_token TEXT, + job_token TEXT NOT NULL, CHECK (state IN ('queued', 'running', 'succeeded', 'failed')), CHECK ((state IN ('succeeded', 'failed')) = (finished_at IS NOT NULL)), CHECK (json_valid(params) AND json_valid(payload)), diff --git a/crates/provisioner-task-libsql/src/lib.rs b/crates/provisioner-task-libsql/src/lib.rs index fef9fdc41..45c7be628 100644 --- a/crates/provisioner-task-libsql/src/lib.rs +++ b/crates/provisioner-task-libsql/src/lib.rs @@ -107,13 +107,13 @@ impl LibSqlProvisionerTaskStore { Self::find(conn, task.id).await?.context("overdue task is missing") } - /// Ends the Task as `state` if it is currently in one of `from`, given as an SQL list such as `'running'`. + /// Ends the Task as `state` if it is currently in one of `from`. /// /// A Task past its deadline is failed as timed out first, and then not ended again. async fn finish( &self, id: Uuid, - from: &'static str, + from: &[ProvisionerTaskState], state: ProvisionerTaskState, payload: serde_json::Value, now: OffsetDateTime, @@ -125,28 +125,30 @@ impl LibSqlProvisionerTaskStore { .context("begin task finish")?; let Some(task) = Self::find(&tx, id).await? else { + tx.rollback().await.context("end task finish")?; return Ok(false); }; - Self::expire(&tx, task, now).await?; + let task = Self::expire(&tx, task, now).await?; - let changed = tx - .execute( - &format!( - "UPDATE provisioner_task_records SET state = ?2, payload = ?3, finished_at = ?4 - WHERE id = ?1 AND state IN ({from})" - ), - params![ - id.to_string(), - state.as_str(), - payload.to_string(), - now.unix_timestamp() - ], - ) - .await - .context("finish task")?; + if !from.contains(&task.state) { + tx.commit().await.context("commit task finish")?; + return Ok(false); + } + + tx.execute( + "UPDATE provisioner_task_records SET state = ?2, payload = ?3, finished_at = ?4 WHERE id = ?1", + params![ + id.to_string(), + state_to_sql(state), + payload.to_string(), + now.unix_timestamp() + ], + ) + .await + .context("finish task")?; tx.commit().await.context("commit task finish")?; - Ok(changed == 1) + Ok(true) } } @@ -176,6 +178,7 @@ impl ProvisionerTaskStore for LibSqlProvisionerTaskStore { let unfinished = Self::expire(&tx, unfinished, now).await?; if !unfinished.state.is_finished() { + tx.rollback().await.context("end task creation")?; return Ok(CreateOutcome::TargetBusy(unfinished)); } } @@ -211,6 +214,7 @@ impl ProvisionerTaskStore for LibSqlProvisionerTaskStore { .context("begin task lookup")?; let Some(task) = Self::find(&tx, id).await? else { + tx.rollback().await.context("end task lookup")?; return Ok(None); }; let task = Self::expire(&tx, task, now).await?; @@ -233,10 +237,12 @@ impl ProvisionerTaskStore for LibSqlProvisionerTaskStore { .context("begin task attempt")?; let Some(task) = Self::find(&tx, id).await? else { + tx.rollback().await.context("end task attempt")?; return Ok(AttemptStart::Unknown); }; - if task.job_token != Some(job_token) { + if task.job_token != job_token { + tx.rollback().await.context("end task attempt")?; return Ok(AttemptStart::Superseded); } @@ -270,6 +276,7 @@ impl ProvisionerTaskStore for LibSqlProvisionerTaskStore { .context("begin job token replacement")?; let Some(task) = Self::find(&tx, id).await? else { + tx.rollback().await.context("end job token replacement")?; return Ok(false); }; let task = Self::expire(&tx, task, now).await?; @@ -326,27 +333,74 @@ impl ProvisionerTaskStore for LibSqlProvisionerTaskStore { Ok(unfinished) } - async fn update_running(&self, id: Uuid, payload: serde_json::Value) -> anyhow::Result { + async fn update_running(&self, id: Uuid, payload: serde_json::Value, now: OffsetDateTime) -> anyhow::Result { let conn = self.conn.lock().await; - let changed = conn - .execute( - "UPDATE provisioner_task_records SET payload = ?2 WHERE id = ?1 AND state = 'running'", + let tx = conn + .transaction_with_behavior(TransactionBehavior::Immediate) + .await + .context("begin running task update")?; + + let Some(task) = Self::find(&tx, id).await? else { + tx.rollback().await.context("end running task update")?; + return Ok(false); + }; + let task = Self::expire(&tx, task, now).await?; + + let running = task.state == ProvisionerTaskState::Running; + + if running { + tx.execute( + "UPDATE provisioner_task_records SET payload = ?2 WHERE id = ?1", params![id.to_string(), payload.to_string()], ) .await .context("update running task")?; + } - Ok(changed == 1) + tx.commit().await.context("commit running task update")?; + + Ok(running) } async fn succeed(&self, id: Uuid, payload: serde_json::Value, now: OffsetDateTime) -> anyhow::Result { - self.finish(id, "'running'", ProvisionerTaskState::Succeeded, payload, now) - .await + self.finish( + id, + &[ProvisionerTaskState::Running], + ProvisionerTaskState::Succeeded, + payload, + now, + ) + .await } async fn fail(&self, id: Uuid, payload: serde_json::Value, now: OffsetDateTime) -> anyhow::Result { - self.finish(id, "'queued', 'running'", ProvisionerTaskState::Failed, payload, now) - .await + self.finish( + id, + &[ProvisionerTaskState::Queued, ProvisionerTaskState::Running], + ProvisionerTaskState::Failed, + payload, + now, + ) + .await + } +} + +fn state_to_sql(state: ProvisionerTaskState) -> &'static str { + match state { + ProvisionerTaskState::Queued => "queued", + ProvisionerTaskState::Running => "running", + ProvisionerTaskState::Succeeded => "succeeded", + ProvisionerTaskState::Failed => "failed", + } +} + +fn state_from_sql(state: &str) -> anyhow::Result { + match state { + "queued" => Ok(ProvisionerTaskState::Queued), + "running" => Ok(ProvisionerTaskState::Running), + "succeeded" => Ok(ProvisionerTaskState::Succeeded), + "failed" => Ok(ProvisionerTaskState::Failed), + _ => bail!("unknown task state {state:?}"), } } @@ -377,11 +431,7 @@ fn task_from_row(row: &Row) -> anyhow::Result { kind: row.get::(1).context("decode task kind")?, target: row.get::(2).context("decode task target")?, params: json(3, "params")?, - state: row - .get::(4) - .context("decode task state")? - .parse() - .context("parse task state")?, + state: state_from_sql(&row.get::(4).context("decode task state")?)?, payload: json(5, "payload")?, attempts: u32::try_from(row.get::(6).context("decode task attempts")?).context("invalid task attempts")?, created_at: timestamp(7, "creation time")?, @@ -389,10 +439,10 @@ fn task_from_row(row: &Row) -> anyhow::Result { finished_at: optional_timestamp(9, "finish time")?, deadline_at: timestamp(10, "deadline")?, job_token: row - .get::>(11) + .get::(11) .context("decode task job token")? - .map(|token| token.parse().context("parse task job token")) - .transpose()?, + .parse() + .context("parse task job token")?, }) } @@ -643,7 +693,7 @@ mod tests { )); assert!( store - .update_running(task.id, serde_json::json!({ "done": 1 })) + .update_running(task.id, serde_json::json!({ "done": 1 }), now()) .await .expect("update") ); @@ -681,7 +731,7 @@ mod tests { )); assert!( !store - .update_running(task.id, serde_json::json!({})) + .update_running(task.id, serde_json::json!({}), now()) .await .expect("update after finish") ); @@ -727,6 +777,12 @@ mod tests { .expect("start"); let late = now() + Duration::minutes(2); + assert!( + !store + .update_running(overdue.id, serde_json::json!({ "done": 4 }), late) + .await + .expect("update late") + ); assert!( !store .succeed(overdue.id, serde_json::json!({ "log": "late" }), late) @@ -821,7 +877,7 @@ mod tests { .start_attempt(task.id, new_token, serde_json::json!({}), now()) .await .expect("start with the new token"), - AttemptStart::Started(started) if started.job_token == Some(new_token) && started.attempts == 1 + AttemptStart::Started(started) if started.job_token == new_token && started.attempts == 1 )); store diff --git a/crates/provisioner-task/src/lib.rs b/crates/provisioner-task/src/lib.rs index 9f69cd906..355f301ec 100644 --- a/crates/provisioner-task/src/lib.rs +++ b/crates/provisioner-task/src/lib.rs @@ -11,7 +11,9 @@ use uuid::Uuid; pub type DynProvisionerTaskStore = Arc; -/// Where a Task is in its life. Every state carries a payload whose shape belongs to the Task kind. +/// Where a Task is in its life. +/// +/// Running, succeeded and failed Tasks carry a payload whose shape belongs to the Task kind; a queued Task has none. #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum ProvisionerTaskState { Queued, @@ -21,34 +23,11 @@ pub enum ProvisionerTaskState { } impl ProvisionerTaskState { - pub fn as_str(self) -> &'static str { - match self { - Self::Queued => "queued", - Self::Running => "running", - Self::Succeeded => "succeeded", - Self::Failed => "failed", - } - } - pub fn is_finished(self) -> bool { matches!(self, Self::Succeeded | Self::Failed) } } -impl std::str::FromStr for ProvisionerTaskState { - type Err = anyhow::Error; - - fn from_str(value: &str) -> anyhow::Result { - match value { - "queued" => Ok(Self::Queued), - "running" => Ok(Self::Running), - "succeeded" => Ok(Self::Succeeded), - "failed" => Ok(Self::Failed), - _ => anyhow::bail!("unknown task state {value:?}"), - } - } -} - /// What is kept of a Task: what was asked, how it goes and how it ended. #[derive(Debug, Clone, PartialEq)] pub struct ProvisionerTaskRecord { @@ -69,7 +48,7 @@ pub struct ProvisionerTaskRecord { /// After this, an unfinished Task is failed, so a lost job can't block its target forever. pub deadline_at: OffsetDateTime, /// The job that may run this Task: a job with another token is stale or a duplicate, and must not run it. - pub job_token: Option, + pub job_token: Uuid, } /// The record of a Task about to be created. @@ -115,11 +94,13 @@ pub trait ProvisionerTaskStore: Send + Sync { /// Like every method that takes `now`, it first fails the Tasks it looks at that are past their deadline. async fn create(&self, task: NewProvisionerTaskRecord, now: OffsetDateTime) -> anyhow::Result; + /// Reads a Task, after failing it if it is past its deadline. async fn get(&self, id: Uuid, now: OffsetDateTime) -> anyhow::Result>; /// Marks the start of one more attempt by the job holding `job_token`: `running`, with one more attempt counted. /// - /// Only [`AttemptStart::Started`] lets the attempt run. + /// Only [`AttemptStart::Started`] lets the attempt run. A job that no longer holds the token gets + /// [`AttemptStart::Superseded`], and nothing changes. async fn start_attempt( &self, id: Uuid, @@ -137,7 +118,7 @@ pub trait ProvisionerTaskStore: Send + Sync { async fn list_unfinished(&self, now: OffsetDateTime) -> anyhow::Result>; /// Replaces the payload of a `running` Task, e.g. to report progress. Returns `false` if it is not running. - async fn update_running(&self, id: Uuid, payload: serde_json::Value) -> anyhow::Result; + async fn update_running(&self, id: Uuid, payload: serde_json::Value, now: OffsetDateTime) -> anyhow::Result; /// Ends a `running` Task as `succeeded`. Returns `false` if it is not running, and changes nothing. async fn succeed(&self, id: Uuid, payload: serde_json::Value, now: OffsetDateTime) -> anyhow::Result; From ed227ac37e418fdc8d2f3cee2d450601a60a0a2a Mon Sep 17 00:00:00 2001 From: Junyi Ou Date: Tue, 13 Oct 2026 11:03:21 -0400 Subject: [PATCH 9/9] chore(dgw): tighten the provisioner task store docs and tests - `replace_job_token` says that only the start of an attempt checks the token, so the old job must run no attempt meanwhile. - `update_running`, `succeed` and `fail` say that `false` can mean the Task was just failed for being past its deadline. - `expire` is renamed `fail_if_overdue`. - Tests: another kind does not make a target busy; another kind or target is an ID conflict; the same request with a new job token keeps the first token; the remaining methods on an unknown ID. Co-Authored-By: Claude Opus 5.5 (1M context) --- crates/provisioner-task-libsql/src/lib.rs | 56 +++++++++++++++++------ crates/provisioner-task/src/lib.rs | 14 ++++-- 2 files changed, 52 insertions(+), 18 deletions(-) diff --git a/crates/provisioner-task-libsql/src/lib.rs b/crates/provisioner-task-libsql/src/lib.rs index 45c7be628..d85a93c94 100644 --- a/crates/provisioner-task-libsql/src/lib.rs +++ b/crates/provisioner-task-libsql/src/lib.rs @@ -79,7 +79,7 @@ impl LibSqlProvisionerTaskStore { } /// Fails `task` if it is unfinished and past its deadline, then returns it as stored. - async fn expire( + async fn fail_if_overdue( conn: &Connection, task: ProvisionerTaskRecord, now: OffsetDateTime, @@ -128,7 +128,7 @@ impl LibSqlProvisionerTaskStore { tx.rollback().await.context("end task finish")?; return Ok(false); }; - let task = Self::expire(&tx, task, now).await?; + let task = Self::fail_if_overdue(&tx, task, now).await?; if !from.contains(&task.state) { tx.commit().await.context("commit task finish")?; @@ -162,7 +162,7 @@ impl ProvisionerTaskStore for LibSqlProvisionerTaskStore { .context("begin task creation")?; if let Some(existing) = Self::find(&tx, task.id).await? { - let existing = Self::expire(&tx, existing, now).await?; + let existing = Self::fail_if_overdue(&tx, existing, now).await?; let same_request = existing.kind == task.kind && existing.target == task.target && existing.params == task.params; tx.commit().await.context("commit task lookup")?; @@ -175,7 +175,7 @@ impl ProvisionerTaskStore for LibSqlProvisionerTaskStore { } if let Some(unfinished) = Self::find_unfinished(&tx, &task.kind, &task.target).await? { - let unfinished = Self::expire(&tx, unfinished, now).await?; + let unfinished = Self::fail_if_overdue(&tx, unfinished, now).await?; if !unfinished.state.is_finished() { tx.rollback().await.context("end task creation")?; @@ -217,7 +217,7 @@ impl ProvisionerTaskStore for LibSqlProvisionerTaskStore { tx.rollback().await.context("end task lookup")?; return Ok(None); }; - let task = Self::expire(&tx, task, now).await?; + let task = Self::fail_if_overdue(&tx, task, now).await?; tx.commit().await.context("commit task lookup")?; Ok(Some(task)) @@ -246,7 +246,7 @@ impl ProvisionerTaskStore for LibSqlProvisionerTaskStore { return Ok(AttemptStart::Superseded); } - let task = Self::expire(&tx, task, now).await?; + let task = Self::fail_if_overdue(&tx, task, now).await?; if task.state.is_finished() { tx.commit().await.context("commit task attempt")?; @@ -279,7 +279,7 @@ impl ProvisionerTaskStore for LibSqlProvisionerTaskStore { tx.rollback().await.context("end job token replacement")?; return Ok(false); }; - let task = Self::expire(&tx, task, now).await?; + let task = Self::fail_if_overdue(&tx, task, now).await?; let replaced = !task.state.is_finished(); @@ -322,7 +322,7 @@ impl ProvisionerTaskStore for LibSqlProvisionerTaskStore { let mut unfinished = Vec::with_capacity(tasks.len()); for task in tasks { - let task = Self::expire(&tx, task, now).await?; + let task = Self::fail_if_overdue(&tx, task, now).await?; if !task.state.is_finished() { unfinished.push(task); } @@ -344,7 +344,7 @@ impl ProvisionerTaskStore for LibSqlProvisionerTaskStore { tx.rollback().await.context("end running task update")?; return Ok(false); }; - let task = Self::expire(&tx, task, now).await?; + let task = Self::fail_if_overdue(&tx, task, now).await?; let running = task.state == ProvisionerTaskState::Running; @@ -577,7 +577,15 @@ mod tests { let task = new_task("session-1"); create(&store, task.clone()).await; - assert!(matches!(create(&store, task).await, CreateOutcome::Existing(_))); + let mut again = task.clone(); + again.job_token = Uuid::new_v4(); + let CreateOutcome::Existing(existing) = create(&store, again).await else { + panic!("the same request must return the existing task"); + }; + assert_eq!( + existing.job_token, task.job_token, + "the job of the first request keeps the task" + ); } #[tokio::test] @@ -586,10 +594,16 @@ mod tests { let task = new_task("session-1"); create(&store, task.clone()).await; - let mut other = task; - other.params = serde_json::json!({ "provider": "anthropic" }); + let mut other_params = task.clone(); + other_params.params = serde_json::json!({ "provider": "anthropic" }); + let mut other_kind = task.clone(); + other_kind.kind = String::from("other-kind"); + let mut other_target = task; + other_target.target = String::from("session-2"); - assert!(matches!(create(&store, other).await, CreateOutcome::IdConflict(_))); + for other in [other_params, other_kind, other_target] { + assert!(matches!(create(&store, other).await, CreateOutcome::IdConflict(_))); + } } #[tokio::test] @@ -606,6 +620,9 @@ mod tests { create(&store, new_task("session-2")).await, CreateOutcome::Created(_) )); + let mut other_kind = new_task("session-1"); + other_kind.kind = String::from("other-kind"); + assert!(matches!(create(&store, other_kind).await, CreateOutcome::Created(_))); store .fail(first.id, serde_json::json!({}), now()) @@ -944,6 +961,19 @@ mod tests { .expect("start attempt"), AttemptStart::Unknown ); + assert!( + !store + .replace_job_token(Uuid::new_v4(), Uuid::new_v4(), now()) + .await + .expect("replace job token") + ); + for finished in [ + store.update_running(Uuid::new_v4(), serde_json::json!({}), now()).await, + store.succeed(Uuid::new_v4(), serde_json::json!({}), now()).await, + store.fail(Uuid::new_v4(), serde_json::json!({}), now()).await, + ] { + assert!(!finished.expect("update unknown task")); + } } #[tokio::test] diff --git a/crates/provisioner-task/src/lib.rs b/crates/provisioner-task/src/lib.rs index 355f301ec..be3d8431d 100644 --- a/crates/provisioner-task/src/lib.rs +++ b/crates/provisioner-task/src/lib.rs @@ -59,7 +59,7 @@ pub struct NewProvisionerTaskRecord { pub target: String, pub params: serde_json::Value, pub deadline_at: OffsetDateTime, - /// Token of the job queued for the new Task. + /// Token of the job that will run the new Task. pub job_token: Uuid, } @@ -111,18 +111,22 @@ pub trait ProvisionerTaskStore: Send + Sync { /// Hands an unfinished Task to a new job, so the job that had it stops at its next attempt. /// - /// Returns `false` if the Task is unknown or finished, and changes nothing. + /// Only the start of an attempt checks the token, so call it only while the old job runs no attempt, such as at + /// startup before any job runs. Returns `false` if the Task is unknown or finished. async fn replace_job_token(&self, id: Uuid, job_token: Uuid, now: OffsetDateTime) -> anyhow::Result; /// Every unfinished Task, after failing those past their deadline. async fn list_unfinished(&self, now: OffsetDateTime) -> anyhow::Result>; - /// Replaces the payload of a `running` Task, e.g. to report progress. Returns `false` if it is not running. + /// Replaces the payload of a `running` Task, e.g. to report progress. Returns `false` if it is not running, such + /// as when it was just failed for being past its deadline. async fn update_running(&self, id: Uuid, payload: serde_json::Value, now: OffsetDateTime) -> anyhow::Result; - /// Ends a `running` Task as `succeeded`. Returns `false` if it is not running, and changes nothing. + /// Ends a `running` Task as `succeeded`. Returns `false` if it is not running, such as when it was just failed for + /// being past its deadline. async fn succeed(&self, id: Uuid, payload: serde_json::Value, now: OffsetDateTime) -> anyhow::Result; - /// Ends a `queued` or `running` Task as `failed`. Returns `false` if it was already finished, and changes nothing. + /// Ends a `queued` or `running` Task as `failed`. Returns `false` if it was already finished, such as when it was + /// just failed for being past its deadline. async fn fail(&self, id: Uuid, payload: serde_json::Value, now: OffsetDateTime) -> anyhow::Result; }