From b01b83cccf620b83798d746e377fe5b64e098b73 Mon Sep 17 00:00:00 2001
From: Brandon Thomas <echelon@gmail.com>
Date: Tue, 8 Jun 2021 03:05:31 -0400
Subject: [PATCH] job max attempts

---
 src/common_env.rs                            |  8 +++++++-
 src/job_queries/tts_inference_job_queries.rs | 13 ++++++++++---
 src/tts_inference_job.rs                     |  8 +++++++-
 3 files changed, 24 insertions(+), 5 deletions(-)

diff --git a/src/common_env.rs b/src/common_env.rs
index e04938a702..36499cf53a 100644
--- a/src/common_env.rs
+++ b/src/common_env.rs
@@ -3,13 +3,18 @@ use crate::util::anyhow_result::AnyhowResult;
 // TODO: Move more shared configs here.
 
 /// Environment variables shared between one or more jobs or the main server.
-/// These are
+/// These do not have to have the same value for each instance, but they behave
+/// similarly across all usages.
 pub struct CommonEnv {
 
   /// The amount of time to wait between job batches (not individual jobs).
   /// This prevents the outer loop of querying batches from flooding the DB.
   /// (In theory, there's work within jobs that prevents rapidly pegging the DB.)
   pub job_batch_wait_millis: u64,
+
+  /// The maximum number of retries a job will get.
+  /// After a job exhausts these attempts, it will become "dead".
+  pub job_max_attempts: u8,
 }
 
 impl CommonEnv {
@@ -17,6 +22,7 @@ impl CommonEnv {
   pub fn read_from_env() -> AnyhowResult<Self> {
     Ok(Self {
       job_batch_wait_millis: easyenv::get_env_num("JOB_BATCH_WAIT_MILLIS", 100)?,
+      job_max_attempts: easyenv::get_env_num("JOB_MAX_ATTEMPTS", 3)?,
     })
   }
 }
\ No newline at end of file
diff --git a/src/job_queries/tts_inference_job_queries.rs b/src/job_queries/tts_inference_job_queries.rs
index ef3ed7b7ea..38f0106ecc 100644
--- a/src/job_queries/tts_inference_job_queries.rs
+++ b/src/job_queries/tts_inference_job_queries.rs
@@ -142,6 +142,7 @@ FOR UPDATE
 UPDATE tts_inference_jobs
 SET
   status = 'started',
+  attempt_count = attempt_count + 1,
   retry_at = NOW() + interval 2 minute
 WHERE id = ?
         "#,
@@ -158,11 +159,17 @@ WHERE id = ?
 pub async fn mark_tts_inference_job_failure(
   pool: &MySqlPool,
   job: &TtsInferenceJobRecord,
-  failure_reason: &str
+  failure_reason: &str,
+  max_attempts: i32
 ) -> AnyhowResult<()> {
 
   // statuses: "attempt_failed", "complete_failure", "dead"
-  let status = "attempt_failed";
+  let mut next_status = "attempt_failed";
+
+  if job.attempt_count >= max_attempts {
+    // NB: Job attempt count is incremented at start
+    next_status = "dead";
+  }
 
   let query_result = sqlx::query!(
         r#"
@@ -173,7 +180,7 @@ SET
   retry_at = NOW() + interval 1 minute
 WHERE id = ?
         "#,
-        status,
+        next_status,
         failure_reason.to_string(),
         job.id,
     )
diff --git a/src/tts_inference_job.rs b/src/tts_inference_job.rs
index a76e6760b3..3d4778999c 100644
--- a/src/tts_inference_job.rs
+++ b/src/tts_inference_job.rs
@@ -91,6 +91,10 @@ struct Inferencer {
 
   // Sleep between batches
   pub job_batch_wait_millis: u64,
+
+  // Max job attempts before failure.
+  // NB: This is an i32 so we don't need to convert to db column type.
+  pub job_max_attempts: i32,
 }
 
 #[tokio::main]
@@ -198,6 +202,7 @@ async fn main() -> AnyhowResult<()> {
     firehose_publisher,
     tts_vocoder_model_filename,
     job_batch_wait_millis: common_env.job_batch_wait_millis,
+    job_max_attempts: common_env.job_max_attempts as i32,
   };
 
   main_loop(inferencer).await;
@@ -264,7 +269,8 @@ async fn process_jobs(inferencer: &Inferencer, jobs: Vec<TtsInferenceJobRecord>)
         let _r = mark_tts_inference_job_failure(
           &inferencer.mysql_pool,
           &job,
-          failure_reason)
+          failure_reason,
+          inferencer.job_max_attempts)
             .await;
       }
     }