various: fix evaluator cancellation and retry lifecycles

Signed-off-by: NotAShelf <raf@notashelf.dev>
Change-Id: I061874ad8c52cd82a427d598887494ef6a6a6964

Author: NotAShelf

crates/common/src/models.rs

@@ -67,8 +67,8 @@ pub struct Evaluation {
 }
 
 #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, sqlx::Type)]
-#[serde(rename_all = "lowercase")]
-#[sqlx(type_name = "text", rename_all = "lowercase")]
+#[serde(rename_all = "snake_case")]
+#[sqlx(type_name = "text", rename_all = "snake_case")]
 pub enum EvaluationStatus {
   Pending,
   Running,
@@ -513,7 +513,7 @@ impl<'r> sqlx::Decode<'r, sqlx::Postgres> for BuildStatus {
 
 #[cfg(test)]
 mod build_status_tests {
-  use super::BuildStatus;
+  use super::{BuildStatus, EvaluationStatus};
 
   #[test]
   fn build_status_i32_round_trips() {
@@ -544,6 +544,14 @@ mod build_status_tests {
     assert_eq!(BuildStatus::from_i32(-1), None);
     assert_eq!(BuildStatus::from_i32(15), None);
   }
+
+  #[test]
+  fn timed_out_status_uses_database_spelling() {
+    assert_eq!(
+      serde_json::to_string(&EvaluationStatus::TimedOut).unwrap(),
+      "\"timed_out\""
+    );
+  }
 }
 
 #[derive(Debug, Clone, Serialize, Deserialize, FromRow)]

crates/common/src/repo/build_dependencies.rs

@@ -1,4 +1,4 @@
-use sqlx::PgPool;
+use sqlx::{Executor, PgPool, Postgres, Transaction};
 use uuid::Uuid;
 
 use crate::{
@@ -17,13 +17,37 @@ pub async fn create(
   build_id: Uuid,
   dependency_build_id: Uuid,
 ) -> Result<BuildDependency> {
+  create_with(pool, build_id, dependency_build_id).await
+}
+
+/// Create a build dependency relationship within an existing transaction.
+///
+/// # Errors
+///
+/// Returns an error if database insert fails or dependency already exists.
+pub async fn create_in_transaction(
+  tx: &mut Transaction<'_, Postgres>,
+  build_id: Uuid,
+  dependency_build_id: Uuid,
+) -> Result<BuildDependency> {
+  create_with(&mut **tx, build_id, dependency_build_id).await
+}
+
+async fn create_with<'e, E>(
+  executor: E,
+  build_id: Uuid,
+  dependency_build_id: Uuid,
+) -> Result<BuildDependency>
+where
+  E: Executor<'e, Database = Postgres>,
+{
   sqlx::query_as::<_, BuildDependency>(
     "INSERT INTO build_dependencies (build_id, dependency_build_id) VALUES \
      ($1, $2) RETURNING *",
   )
   .bind(build_id)
   .bind(dependency_build_id)
-  .fetch_one(pool)
+  .fetch_one(executor)
   .await
   .on_unique_violation(|| {
     format!(

crates/common/src/repo/builds.rs

@@ -1,6 +1,6 @@
 use std::collections::HashSet;
 
-use sqlx::PgPool;
+use sqlx::{Executor, PgPool, Postgres, Transaction};
 use uuid::Uuid;
 
 use crate::{
@@ -14,6 +14,25 @@ use crate::{
 ///
 /// Returns error if database insert fails or job already exists.
 pub async fn create(pool: &PgPool, input: CreateBuild) -> Result<Build> {
+  create_with(pool, input).await
+}
+
+/// Create a new build record within an existing transaction.
+///
+/// # Errors
+///
+/// Returns an error if database insert fails or job already exists.
+pub async fn create_in_transaction(
+  tx: &mut Transaction<'_, Postgres>,
+  input: CreateBuild,
+) -> Result<Build> {
+  create_with(&mut **tx, input).await
+}
+
+async fn create_with<'e, E>(executor: E, input: CreateBuild) -> Result<Build>
+where
+  E: Executor<'e, Database = Postgres>,
+{
   let is_aggregate = input.is_aggregate.unwrap_or(false);
   let is_fod = input.is_fod.unwrap_or(false);
   sqlx::query_as::<_, Build>(
@@ -37,7 +56,7 @@ pub async fn create(pool: &PgPool, input: CreateBuild) -> Result<Build> {
   .bind(&input.meta_homepage)
   .bind(&input.meta_maintainers)
   .bind(&input.required_features)
-  .fetch_one(pool)
+  .fetch_one(executor)
   .await
   .on_unique_violation(|| {
     format!(

crates/common/src/repo/evaluations.rs

@@ -1,4 +1,4 @@
-use sqlx::PgPool;
+use sqlx::{PgPool, Postgres, Transaction};
 use uuid::Uuid;
 
 use crate::{
@@ -365,17 +365,18 @@ pub async fn cancel(pool: &PgPool, id: Uuid) -> Result<Option<Evaluation>> {
   )
 }
 
-/// Requeue a cancelled or failed evaluation after discarding its stale builds.
+/// Requeue a cancelled, failed, or timed-out evaluation after discarding its
+/// stale builds. One-shot jobsets are re-enabled for the new attempt.
 ///
 /// # Errors
 ///
 /// Returns an error if the database transaction fails.
 pub async fn restart(pool: &PgPool, id: Uuid) -> Result<Option<Evaluation>> {
   let mut tx = pool.begin().await?;
   let evaluation = sqlx::query_as::<_, Evaluation>(
-    "UPDATE evaluations SET status = 'pending', error_message = NULL, \
-     inputs_hash = NULL WHERE id = $1 AND status IN ('cancelled', 'failed') \
-     RETURNING *",
+    "UPDATE evaluations SET status = 'pending', evaluation_time = NOW(), \
+     error_message = NULL, inputs_hash = NULL WHERE id = $1 AND status IN \
+     ('cancelled', 'failed', 'timed_out') RETURNING *",
   )
   .bind(id)
   .fetch_optional(&mut *tx)
@@ -386,11 +387,64 @@ pub async fn restart(pool: &PgPool, id: Uuid) -> Result<Option<Evaluation>> {
       .bind(id)
       .execute(&mut *tx)
       .await?;
+    sqlx::query(
+      "UPDATE jobsets SET enabled = true WHERE id = (SELECT jobset_id FROM \
+       evaluations WHERE id = $1) AND state = 'one_shot'",
+    )
+    .bind(id)
+    .execute(&mut *tx)
+    .await?;
   }
   tx.commit().await?;
   Ok(evaluation)
 }
 
+/// Lock a running evaluation before atomically persisting its result.
+///
+/// # Errors
+///
+/// Returns an error if the database query fails.
+pub async fn lock_running(
+  tx: &mut Transaction<'_, Postgres>,
+  id: Uuid,
+) -> Result<bool> {
+  Ok(
+    sqlx::query_scalar::<_, Uuid>(
+      "SELECT id FROM evaluations WHERE id = $1 AND status = 'running' FOR \
+       UPDATE",
+    )
+    .bind(id)
+    .fetch_optional(&mut **tx)
+    .await?
+    .is_some(),
+  )
+}
+
+/// Finish a locked running evaluation within its result transaction.
+///
+/// # Errors
+///
+/// Returns an error if the database query fails.
+pub async fn finish_running_in_transaction(
+  tx: &mut Transaction<'_, Postgres>,
+  id: Uuid,
+  status: EvaluationStatus,
+  error_message: Option<&str>,
+) -> Result<bool> {
+  Ok(
+    sqlx::query_scalar::<_, Uuid>(
+      "UPDATE evaluations SET status = $1, error_message = $2 WHERE id = $3 \
+       AND status = 'running' RETURNING id",
+    )
+    .bind(status)
+    .bind(error_message)
+    .bind(id)
+    .fetch_optional(&mut **tx)
+    .await?
+    .is_some(),
+  )
+}
+
 /// Return whether an evaluator should cancel its currently-running work.
 ///
 /// # Errors

crates/common/src/repo/jobsets.rs

@@ -269,15 +269,14 @@ pub async fn list_active(pool: &PgPool) -> Result<Vec<ActiveJobset>> {
   )
 }
 
-/// Mark a one-shot jobset as complete (set state to disabled).
+/// Mark a one-shot jobset as complete without losing its one-shot state.
 ///
 /// # Errors
 ///
 /// Returns error if database update fails.
 pub async fn mark_one_shot_complete(pool: &PgPool, id: Uuid) -> Result<()> {
   sqlx::query(
-    "UPDATE jobsets SET state = 'disabled', enabled = false WHERE id = $1 AND \
-     state = 'one_shot'",
+    "UPDATE jobsets SET enabled = false WHERE id = $1 AND state = 'one_shot'",
   )
   .bind(id)
   .execute(pool)