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/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 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..e39abac43 --- /dev/null +++ b/crates/provisioner-task-libsql/migrations/01_provisioner_task.sql @@ -0,0 +1,26 @@ +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, + 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)), + 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..d85a93c94 --- /dev/null +++ b/crates/provisioner-task-libsql/src/lib.rs @@ -0,0 +1,1016 @@ +//! 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::{ + AttemptStart, CreateOutcome, NewProvisionerTaskRecord, ProvisionerTaskRecord, 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, job_token"; + +/// 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 fail_if_overdue( + conn: &Connection, + task: ProvisionerTaskRecord, + 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`. + /// + /// A Task past its deadline is failed as timed out first, and then not ended again. + async fn finish( + &self, + id: Uuid, + from: &[ProvisionerTaskState], + state: ProvisionerTaskState, + 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 finish")?; + + let Some(task) = Self::find(&tx, id).await? else { + tx.rollback().await.context("end task finish")?; + return Ok(false); + }; + let task = Self::fail_if_overdue(&tx, task, now).await?; + + 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(true) + } +} + +#[async_trait] +impl ProvisionerTaskStore for LibSqlProvisionerTaskStore { + async fn create(&self, task: NewProvisionerTaskRecord, 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::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")?; + + 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::fail_if_overdue(&tx, unfinished, now).await?; + + if !unfinished.state.is_finished() { + tx.rollback().await.context("end task creation")?; + return Ok(CreateOutcome::TargetBusy(unfinished)); + } + } + + tx.execute( + "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(), + task.target.clone(), + task.params.to_string(), + now.unix_timestamp(), + task.deadline_at.unix_timestamp(), + task.job_token.to_string(), + ], + ) + .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 { + tx.rollback().await.context("end task lookup")?; + return Ok(None); + }; + let task = Self::fail_if_overdue(&tx, task, now).await?; + tx.commit().await.context("commit task lookup")?; + + Ok(Some(task)) + } + + 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 { + tx.rollback().await.context("end task attempt")?; + return Ok(AttemptStart::Unknown); + }; + + if task.job_token != job_token { + tx.rollback().await.context("end task attempt")?; + return Ok(AttemptStart::Superseded); + } + + let task = Self::fail_if_overdue(&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 { + tx.rollback().await.context("end job token replacement")?; + return Ok(false); + }; + let task = Self::fail_if_overdue(&tx, task, now).await?; + + let replaced = !task.state.is_finished(); + + if replaced { + tx.execute( + "UPDATE provisioner_task_records SET job_token = ?2 WHERE id = ?1", + params![id.to_string(), job_token.to_string()], + ) + .await + .context("replace task job token")?; + } + + tx.commit().await.context("commit job token replacement")?; + + 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::fail_if_overdue(&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, now: OffsetDateTime) -> anyhow::Result { + let conn = self.conn.lock().await; + 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::fail_if_overdue(&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")?; + } + + 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, + &[ProvisionerTaskState::Running], + ProvisionerTaskState::Succeeded, + payload, + now, + ) + .await + } + + async fn fail(&self, id: Uuid, payload: serde_json::Value, now: OffsetDateTime) -> anyhow::Result { + 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:?}"), + } +} + +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(ProvisionerTaskRecord { + 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: 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")?, + 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")? + .parse() + .context("parse task job token")?, + }) +} + +/// 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) -> NewProvisionerTaskRecord { + NewProvisionerTaskRecord { + 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), + job_token: Uuid::new_v4(), + } + } + + async fn create(store: &LibSqlProvisionerTaskStore, task: NewProvisionerTaskRecord) -> 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; + + 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] + 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_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"); + + for other in [other_params, other_kind, other_target] { + 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(_) + )); + 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()) + .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, + overdue.job_token, + 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!(matches!( + store + .start_attempt(task.id, task.job_token, serde_json::json!({ "done": 0 }), now()) + .await + .expect("start"), + AttemptStart::Started(_) + )); + assert!( + store + .update_running(task.id, serde_json::json!({ "done": 1 }), now()) + .await + .expect("update") + ); + assert!(matches!( + store + .start_attempt(task.id, task.job_token, serde_json::json!({ "done": 1 }), now()) + .await + .expect("retry"), + AttemptStart::Started(_) + )); + + 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!(matches!( + store + .start_attempt(task.id, task.job_token, serde_json::json!({}), now()) + .await + .expect("start after finish"), + AttemptStart::Finished(_) + )); + assert!( + !store + .update_running(task.id, serde_json::json!({}), now()) + .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 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, task.job_token, 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, overdue.job_token, serde_json::json!({ "done": 3 }), now()) + .await + .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) + .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; + let mut overdue = new_task("session-1"); + overdue.deadline_at = now() - Duration::minutes(1); + create(&store, overdue.clone()).await; + + assert!(matches!( + store + .start_attempt(overdue.id, overdue.job_token, serde_json::json!({}), now()) + .await + .expect("start attempt"), + AttemptStart::Finished(_) + )); + 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 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 + .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 == new_token && started.attempts == 1 + )); + + store + .fail(task.id, serde_json::json!({ "reason": "boom" }), now()) + .await + .expect("fail"); + assert!( + !store + .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"), + 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] + 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/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. + diff --git a/crates/provisioner-task/src/lib.rs b/crates/provisioner-task/src/lib.rs new file mode 100644 index 000000000..be3d8431d --- /dev/null +++ b/crates/provisioner-task/src/lib.rs @@ -0,0 +1,132 @@ +//! 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. +/// +/// 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, + Running, + Succeeded, + Failed, +} + +impl ProvisionerTaskState { + pub fn is_finished(self) -> bool { + matches!(self, Self::Succeeded | Self::Failed) + } +} + +/// What is kept of a Task: what was asked, how it goes and how it ended. +#[derive(Debug, Clone, PartialEq)] +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`. + 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, + /// 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: Uuid, +} + +/// The record of a Task about to be created. +#[derive(Debug, Clone, PartialEq)] +pub struct NewProvisionerTaskRecord { + pub id: Uuid, + pub kind: String, + pub target: String, + pub params: serde_json::Value, + pub deadline_at: OffsetDateTime, + /// Token of the job that will run the new Task. + pub job_token: Uuid, +} + +#[derive(Debug, Clone, PartialEq)] +pub enum CreateOutcome { + Created(ProvisionerTaskRecord), + /// The same request came in again; nothing new was created. + Existing(ProvisionerTaskRecord), + /// The ID is already used by a Task with another kind, target or parameters. + IdConflict(ProvisionerTaskRecord), + /// Another unfinished Task of the same kind already works on this target. + 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(ProvisionerTaskRecord), + /// The Task is finished, possibly just now because it is past its deadline. + Finished(ProvisionerTaskRecord), + /// 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. + /// + /// 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. A job that no longer holds the token gets + /// [`AttemptStart::Superseded`], and nothing changes. + 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. + /// + /// 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, 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, 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, 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; +}