WIP generic-download-job

crates/lib/database_queries/src/mediators/badge_granter.rs

@@ -12,7 +12,7 @@ pub enum UserBadge {
   // Granted for early vocodes users.
   EarlyUser,
 
-  // Granted for uploading models
+  // Granted for uploading tts models
   TtsModelUploader1,
   TtsModelUploader5,
   TtsModelUploader10,
@@ -23,6 +23,17 @@ pub enum UserBadge {
   TtsModelUploader200,
   TtsModelUploader250,
 
+  // Granted for vocoder models
+  VocoderModelUploader1,
+  VocoderModelUploader5,
+  VocoderModelUploader10,
+  VocoderModelUploader20,
+  VocoderModelUploader50,
+  VocoderModelUploader100,
+  VocoderModelUploader150,
+  VocoderModelUploader200,
+  VocoderModelUploader250,
+
   // Granted for uploading templates
   W2lTemplateUploader1,
   W2lTemplateUploader10,
@@ -63,6 +74,15 @@ impl UserBadge {
       UserBadge::TtsModelUploader150 => "tts_model_uploader_150",
       UserBadge::TtsModelUploader200 => "tts_model_uploader_200",
       UserBadge::TtsModelUploader250 => "tts_model_uploader_250",
+      UserBadge::VocoderModelUploader1 => "vocoder_model_uploader_1",
+      UserBadge::VocoderModelUploader5 => "vocoder_model_uploader_5",
+      UserBadge::VocoderModelUploader10 => "vocoder_model_uploader_10",
+      UserBadge::VocoderModelUploader20 => "vocoder_model_uploader_20",
+      UserBadge::VocoderModelUploader50 => "vocoder_model_uploader_50",
+      UserBadge::VocoderModelUploader100 => "vocoder_model_uploader_100",
+      UserBadge::VocoderModelUploader150 => "vocoder_model_uploader_150",
+      UserBadge::VocoderModelUploader200 => "vocoder_model_uploader_200",
+      UserBadge::VocoderModelUploader250 => "vocoder_model_uploader_250",
       UserBadge::W2lTemplateUploader1 => "w2l_template_uploader_1",
       UserBadge::W2lTemplateUploader10 => "w2l_template_uploader_10",
       UserBadge::W2lTemplateUploader50 => "w2l_template_uploader_50",
@@ -206,6 +226,51 @@ impl BadgeGranter {
     Ok(())
   }
 
+  /// This needs to be called *after* successful upload.
+  pub async fn maybe_grant_vocoder_model_uploads_badge(&self, user_token: &str) -> AnyhowResult<()> {
+    let count = self.count_vocoder_models_uploaded(user_token).await?;
+
+    let mut maybe_badge = None;
+
+    if count >= 250 {
+      maybe_badge = Some(UserBadge::VocoderModelUploader250);
+    } else if count >= 200 {
+      maybe_badge = Some(UserBadge::VocoderModelUploader200);
+    } else if count >= 150 {
+      maybe_badge = Some(UserBadge::VocoderModelUploader150);
+    } else if count >= 100 {
+      maybe_badge = Some(UserBadge::VocoderModelUploader100);
+    } else if count >= 50 {
+      maybe_badge = Some(UserBadge::VocoderModelUploader50);
+    } else if count >= 20 {
+      maybe_badge = Some(UserBadge::VocoderModelUploader20);
+    } else if count >= 10 {
+      maybe_badge = Some(UserBadge::VocoderModelUploader10);
+    } else if count >= 5 {
+      maybe_badge = Some(UserBadge::VocoderModelUploader5);
+    } else if count >= 1 {
+      maybe_badge = Some(UserBadge::VocoderModelUploader1);
+    }
+
+    let badge = match maybe_badge {
+      Some(badge) => badge,
+      None => return Ok(()),
+    };
+
+    if self.has_badge(user_token, badge).await? {
+      return Ok(())
+    }
+
+    let _record_id = self.insert(
+      badge,
+      user_token,
+    ).await?;
+
+    self.firehose_publisher.publish_user_badge_granted(user_token, badge.to_db_value())
+        .await?;
+
+    Ok(())
+  }
 
   // =======================================================================
 
@@ -267,6 +332,28 @@ LIMIT 1
     self.handle_count_query(maybe_result)
   }
 
+  async fn count_vocoder_models_uploaded(
+    &self,
+    user_token: &str,
+  ) -> AnyhowResult<u64> {
+    // NB: This could get expensive!
+    let maybe_result = sqlx::query_as!(
+      CountRecord,
+        r#"
+SELECT count(*) as count
+FROM vocoder_models
+WHERE
+  creator_user_token = ?
+LIMIT 1
+        "#,
+      user_token
+    )
+        .fetch_one(&self.mysql_pool)
+        .await;
+
+    self.handle_count_query(maybe_result)
+  }
+
   async fn count_w2l_templates_uploaded(
     &self,
     user_token: &str,

crates/lib/database_queries/src/queries/generic_download/job/_keys.rs

@@ -0,0 +1,5 @@
+//! Primary and foreign keys
+
+/// Since jobs are internal, we'll use IDs as primary keys.
+#[derive(Clone, Copy, PartialOrd, PartialEq, Eq, Ord, Debug)]
+pub struct GenericDownloadJobId(pub i64);

crates/lib/database_queries/src/queries/generic_download/job/list_available_generic_download_jobs.rs

@@ -0,0 +1,123 @@
+use anyhow::anyhow;
+use chrono::Utc;
+use container_common::anyhow_result::AnyhowResult;
+use crate::column_types::job_status::JobStatus;
+use crate::queries::generic_download::job::_keys::GenericDownloadJobId;
+use crate::tokens::Tokens;
+use reusable_types::entity_visibility::EntityVisibility;
+use reusable_types::generic_download_type::GenericDownloadType;
+use sqlx::MySqlPool;
+use std::path::Path;
+
+/// table: generic_download_jobs
+#[derive(Debug)]
+pub struct AvailableDownloadJob {
+  pub id: GenericDownloadJobId,
+  pub download_job_token: String,
+
+  pub creator_user_token: String,
+  pub creator_ip_address: String,
+  pub creator_set_visibility: EntityVisibility,
+
+  pub download_type: GenericDownloadType,
+  pub download_url: Option<String>,
+  pub title: String,
+
+  pub status: JobStatus,
+  pub attempt_count: i32,
+  pub failure_reason: Option<String>,
+
+  pub created_at: chrono::DateTime<Utc>,
+  pub updated_at: chrono::DateTime<Utc>,
+  pub retry_at: Option<chrono::DateTime<Utc>>,
+}
+
+pub async fn list_available_generic_download_jobs(pool: &MySqlPool, num_records: u32)
+  -> AnyhowResult<Vec<AvailableDownloadJob>>
+{
+  let job_records = sqlx::query_as!(
+      AvailableDownloadJobRawInternal,
+        r#"
+SELECT
+  id,
+  token AS download_job_token,
+
+  creator_user_token,
+  creator_ip_address,
+  creator_set_visibility as `creator_set_visibility: reusable_types::entity_visibility::EntityVisibility`,
+
+  download_type as `download_type: reusable_types::generic_download_type::GenericDownloadType`,
+  download_url,
+  title,
+
+  status as `status: crate::column_types::job_status::JobStatus`,
+  attempt_count,
+  failure_reason,
+
+  created_at,
+  updated_at,
+  retry_at
+FROM generic_download_jobs
+WHERE
+  (
+    status IN ("pending", "attempt_failed")
+  )
+  AND
+  (
+    retry_at IS NULL
+    OR
+    retry_at < CURRENT_TIMESTAMP
+  )
+  ORDER BY id ASC
+  LIMIT ?
+        "#,
+      num_records,
+    )
+      .fetch_all(pool)
+      .await?;
+
+  let job_records = job_records.into_iter()
+      .map(|record : AvailableDownloadJobRawInternal| {
+        AvailableDownloadJob {
+          id: GenericDownloadJobId(record.id),
+          download_job_token: record.download_job_token,
+          creator_ip_address: record.creator_ip_address,
+          creator_user_token: record.creator_user_token,
+          creator_set_visibility: record.creator_set_visibility,
+          download_type: record.download_type,
+          download_url: record.download_url,
+          title: record.title,
+          status: record.status,
+          attempt_count: record.attempt_count,
+          failure_reason: record.failure_reason,
+          created_at: record.created_at,
+          updated_at: record.updated_at,
+          retry_at: record.retry_at,
+        }
+      })
+      .collect::<Vec<AvailableDownloadJob>>();
+
+  Ok(job_records)
+}
+
+#[derive(Debug)]
+struct AvailableDownloadJobRawInternal {
+  pub id: i64,
+  pub download_job_token: String,
+
+  pub creator_user_token: String,
+  pub creator_ip_address: String,
+  pub creator_set_visibility: EntityVisibility,
+
+  pub download_type: GenericDownloadType,
+  pub download_url: Option<String>,
+  pub title: String,
+
+  pub status: JobStatus,
+  pub attempt_count: i32,
+  pub failure_reason: Option<String>,
+
+  pub created_at: chrono::DateTime<Utc>,
+  pub updated_at: chrono::DateTime<Utc>,
+  pub retry_at: Option<chrono::DateTime<Utc>>,
+}

crates/lib/database_queries/src/queries/generic_download/job/mark_generic_download_job_done.rs

@@ -0,0 +1,39 @@
+use anyhow::anyhow;
+use container_common::anyhow_result::AnyhowResult;
+use crate::queries::generic_download::job::list_available_generic_download_jobs::AvailableDownloadJob;
+use sqlx::MySqlPool;
+
+pub async fn mark_generic_download_job_done(
+  pool: &MySqlPool,
+  job: &AvailableDownloadJob,
+  success: bool,
+  maybe_entity_token: Option<&str>,
+  maybe_entity_type: Option<&str>,
+) -> AnyhowResult<()>
+{
+  let status = if success { "complete_success" } else { "complete_failure" };
+
+  let query_result = sqlx::query!(
+        r#"
+UPDATE generic_download_jobs
+SET
+  status = ?,
+  on_success_downloaded_entity_token = ?,
+  on_success_downloaded_entity_type = ?,
+  failure_reason = NULL,
+  retry_at = NULL
+WHERE id = ?
+        "#,
+        status,
+        maybe_entity_token,
+        maybe_entity_type,
+        job.id.0,
+    )
+      .execute(pool)
+      .await;
+
+  match query_result {
+    Err(err) => Err(anyhow!("error with query: {:?}", err)),
+    Ok(_r) => Ok(()),
+  }
+}

crates/lib/database_queries/src/queries/generic_download/job/mark_generic_download_job_failure.rs

@@ -0,0 +1,45 @@
+use anyhow::anyhow;
+use container_common::anyhow_result::AnyhowResult;
+use crate::queries::generic_download::job::list_available_generic_download_jobs::AvailableDownloadJob;
+use sqlx::MySqlPool;
+use sqlx;
+
+/// Mark a single download job failure. The job may be re-run.
+pub async fn mark_generic_download_job_failure(
+  pool: &MySqlPool,
+  job: &AvailableDownloadJob,
+  failure_reason: &str,
+  max_attempts: i32
+) -> AnyhowResult<()> {
+
+  // statuses: "attempt_failed", "complete_failure", "dead"
+  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#"
+UPDATE generic_download_jobs
+SET
+  status = ?,
+  failure_reason = ?,
+  retry_at = NOW() + interval 2 minute
+WHERE id = ?
+        "#,
+        next_status,
+        failure_reason.to_string(),
+        job.id.0,
+    )
+      .execute(pool)
+      .await;
+
+  match query_result {
+    Err(err) => {
+      Err(anyhow!("error with query: {:?}", err))
+    },
+    Ok(_r) => Ok(()),
+  }
+}

crates/lib/database_queries/src/queries/generic_download/job/mark_generic_download_job_pending_and_grab_lock.rs

@@ -0,0 +1,84 @@
+use anyhow::anyhow;
+use container_common::anyhow_result::AnyhowResult;
+use sqlx::MySqlPool;
+use sqlx;
+use crate::queries::generic_download::job::_keys::GenericDownloadJobId;
+
+// TODO(2022-02-09): It's unclear how stuck "pending" jobs don't become perma-dead
+//  I think this logic is flawed and the job system needs to pick up stuck "pending" jobs.
+
+pub struct GenericDownloadJobLockRecord {
+  id: i64,
+  status: String,
+}
+
+pub async fn mark_generic_download_job_pending_and_grab_lock(
+  pool: &MySqlPool,
+  job_id: GenericDownloadJobId
+) -> AnyhowResult<bool> {
+
+  // NB: We use transactions and "SELECT ... FOR UPDATE" to simulate mutexes.
+  let mut transaction = pool.begin().await?;
+
+  let maybe_record = sqlx::query_as!(
+    TtsInferenceLockRecord,
+        r#"
+SELECT
+  id,
+  status
+FROM generic_download_jobs
+WHERE id = ?
+FOR UPDATE
+        "#,
+        job_id.0,
+    )
+      .fetch_one(&mut transaction)
+      .await;
+
+  let record : GenericDownloadJobLockRecord = match maybe_record {
+    Ok(record) => record,
+    Err(err) => {
+      match err {
+        sqlx::Error::RowNotFound => {
+          return Err(anyhow!("could not job"));
+        },
+        _ => {
+          return Err(anyhow!("query error"));
+        }
+      }
+    }
+  };
+
+  let can_transact = match record.status.as_ref() {
+    "pending" => true, // It's okay for us to take the lock.
+    "attempt_failed" => true, // We can retry.
+    "started" => false, // Job in progress (another job beat us, and we can't take the lock)
+    "complete_success" => false, // Job already complete
+    "complete_failure" => false, // Job already complete (permanently dead; no need to retry)
+    "dead" => false, // Job already complete (permanently dead; retries exhausted)
+    _ => false, // Future-proof
+  };
+
+  if !can_transact {
+    transaction.rollback().await?;
+    return Ok(false);
+  }
+
+  let _acquire_lock = sqlx::query!(
+        r#"
+UPDATE generic_download_jobs
+SET
+  status = 'started',
+  attempt_count = attempt_count + 1,
+  retry_at = NOW() + interval 2 minute
+WHERE id = ?
+        "#,
+        job_id.0,
+    )
+      .execute(&mut transaction)
+      .await?;
+
+  transaction.commit().await?;
+
+  Ok(true)
+}

crates/lib/database_queries/src/queries/generic_download/job/mod.rs

@@ -0,0 +1,5 @@
+pub mod _keys;
+pub mod list_available_generic_download_jobs;
+pub mod mark_generic_download_job_done;
+pub mod mark_generic_download_job_failure;
+pub mod mark_generic_download_job_pending_and_grab_lock;

crates/lib/database_queries/src/queries/generic_download/mod.rs

@@ -1,2 +1,3 @@
 pub mod get_generic_download_job_status;
 pub mod insert_generic_download_job;
+pub mod job;

crates/lib/database_queries/src/queries/tts/tts_download_jobs/tts_download_job_queries.rs

@@ -1,7 +1,3 @@
-//! NB: This seems required for sqlx to generate the cached queries.
-//! Sqlx's prepare needs a *single* binary to work against, so we need to
-//! include these in the main binary to generate all the queries.
-
 use anyhow::anyhow;
 use chrono::Utc;
 use container_common::anyhow_result::AnyhowResult;

crates/service/generic_download_job/src/job_state.rs

@@ -0,0 +1,41 @@
+use crate::HifiGanModelCheckCommand;
+use database_queries::mediators::badge_granter::BadgeGranter;
+use database_queries::mediators::firehose_publisher::FirehosePublisher;
+use google_drive_common::google_drive_download_command::GoogleDriveDownloadCommand;
+use r2d2_redis::RedisConnectionManager;
+use r2d2_redis::r2d2;
+use sqlx::MySqlPool;
+use std::path::PathBuf;
+use storage_buckets_common::bucket_client::BucketClient;
+use storage_buckets_common::bucket_path_unifier::BucketPathUnifier;
+
+pub struct JobState {
+  pub download_temp_directory: PathBuf,
+  pub mysql_pool: MySqlPool,
+
+  pub redis_pool: r2d2::Pool<RedisConnectionManager>,
+
+  pub bucket_client: BucketClient,
+  pub firehose_publisher: FirehosePublisher,
+  pub badge_granter: BadgeGranter,
+  pub google_drive_downloader: GoogleDriveDownloadCommand,
+
+  pub bucket_path_unifier: BucketPathUnifier,
+
+  pub hifigan_model_check_command: HifiGanModelCheckCommand,
+
+  // Command to run
+  pub download_script: String,
+  // Root to store TTS results
+  pub bucket_root_tts_model_uploads: String,
+
+  // Sleep between batches
+  pub job_batch_wait_millis: u64,
+
+  // How long to wait between log lines
+  pub no_op_logger_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,
+}

crates/service/generic_download_job/src/job_steps/mod.rs

@@ -0,0 +1,2 @@
+pub mod process_single_job;
+pub mod process_hifigan_vocoder;

crates/service/generic_download_job/src/job_steps/process_hifigan_vocoder.rs

@@ -0,0 +1,118 @@
+use std::path::PathBuf;
+use log::info;
+use tempdir::TempDir;
+use container_common::anyhow_result::AnyhowResult;
+use database_queries::queries::generic_download::job::list_available_generic_download_jobs::AvailableDownloadJob;
+use crate::JobState;
+
+pub async fn process_hifigan_vocoder(
+  job_state: &JobState,
+  job: &AvailableDownloadJob,
+  temp_dir: &TempDir
+) -> AnyhowResult<()> {
+
+  // ==================== RUN MODEL CHECK ==================== //
+
+  info!("Checking that model is valid...");
+
+  redis_logger.log_status("checking model")?;
+
+  let file_path = PathBuf::from(download_filename.clone());
+
+  let output_metadata_fs_path = temp_dir.path().join("metadata.json");
+
+  match model_type {
+    TtsModelType::Tacotron2 => {
+      let result = job_state.tacotron_tts_check.execute(
+        &file_path,
+        &output_metadata_fs_path,
+        false,
+      );
+
+      if let Err(e) = result {
+        safe_delete_temp_file(&file_path);
+        safe_delete_temp_directory(&temp_dir);
+      }
+    },
+    TtsModelType::Talknet => {
+      let result = job_state.talknet_tts_check.execute(
+        &file_path,
+        &output_metadata_fs_path,
+        false,
+      );
+
+      if let Err(e) = result {
+        safe_delete_temp_file(&file_path);
+        safe_delete_temp_directory(&temp_dir);
+      }
+    },
+  }
+
+  // ==================== CHECK ALL FILES EXIST AND GET METADATA ==================== //
+
+  info!("Checking that metadata output file exists...");
+
+  check_file_exists(&output_metadata_fs_path)?;
+
+  let file_metadata = match read_metadata_file(&output_metadata_fs_path) {
+    Ok(metadata) => metadata,
+    Err(e) => {
+      safe_delete_temp_file(&file_path);
+      safe_delete_temp_file(&output_metadata_fs_path);
+      safe_delete_temp_directory(&temp_dir);
+      return Err(e);
+    }
+  };
+
+  // ==================== UPLOAD MODEL FILE ==================== //
+
+  info!("Uploading model to GCS...");
+
+  let private_bucket_hash = hash_file_sha2(&download_filename)?;
+
+  info!("File hash: {}", private_bucket_hash);
+
+  let synthesizer_model_bucket_path = match model_type {
+    TtsModelType::Tacotron2 => job_state.bucket_path_unifier.tts_synthesizer_path(&private_bucket_hash),
+    TtsModelType::Talknet => job_state.bucket_path_unifier.tts_zipped_synthesizer_path(&private_bucket_hash),
+  };
+
+  info!("Destination bucket path: {:?}", &synthesizer_model_bucket_path);
+
+  redis_logger.log_status("uploading model")?;
+
+  if let Err(e) = job_state.bucket_client.upload_filename(&synthesizer_model_bucket_path, &file_path).await {
+    safe_delete_temp_file(&output_metadata_fs_path);
+    safe_delete_temp_file(&file_path);
+    safe_delete_temp_directory(&temp_dir);
+    return Err(e);
+  }
+
+  // ==================== DELETE DOWNLOADED FILE ==================== //
+
+  // NB: We should be using a tempdir, but to make absolutely certain we don't overflow the disk...
+  safe_delete_temp_file(&output_metadata_fs_path);
+  safe_delete_temp_file(&file_path);
+  safe_delete_temp_directory(&temp_dir);
+
+  // ==================== SAVE RECORDS ==================== //
+
+  info!("Saving model record...");
+  let (id, model_token) = insert_tts_model(
+    &job_state.mysql_pool,
+    job,
+    &private_bucket_hash,
+    synthesizer_model_bucket_path,
+    file_metadata.file_size_bytes)
+      .await?;
+
+  job_state.badge_granter.maybe_grant_tts_model_uploads_badge(&job.creator_user_token)
+      .await
+      .map_err(|e| {
+        warn!("error maybe awarding badge: {:?}", e);
+        anyhow!("error maybe awarding badge")
+      })?;
+
+
+  Ok(())
+}

crates/service/generic_download_job/src/job_steps/process_single_job.rs

@@ -0,0 +1,94 @@
+use anyhow::anyhow;
+use config::bad_urls::is_bad_tts_model_download_url;
+use container_common::anyhow_result::AnyhowResult;
+use container_common::filesystem::check_file_exists::check_file_exists;
+use container_common::filesystem::safe_delete_temp_directory::safe_delete_temp_directory;
+use container_common::filesystem::safe_delete_temp_file::safe_delete_temp_file;
+use container_common::hashing::hash_file_sha2::hash_file_sha2;
+use crate::JobState;
+use database_queries::queries::generic_download::job::list_available_generic_download_jobs::AvailableDownloadJob;
+use database_queries::queries::generic_download::job::mark_generic_download_job_pending_and_grab_lock::mark_generic_download_job_pending_and_grab_lock;
+use jobs_common::redis_job_status_logger::RedisJobStatusLogger;
+use log::{info, warn};
+use std::path::PathBuf;
+use tempdir::TempDir;
+use database_queries::queries::generic_download::job::mark_generic_download_job_done::mark_generic_download_job_done;
+use crate::job_steps::process_hifigan_vocoder::process_hifigan_vocoder;
+
+pub async fn process_single_job(job_state: &JobState, job: &AvailableDownloadJob) -> AnyhowResult<()> {
+  let mut redis = job_state.redis_pool.get()?;
+  let mut redis_logger = RedisJobStatusLogger::new_tts_download(
+    &mut redis,
+    &job.token);
+
+  // ==================== ATTEMPT TO GRAB JOB LOCK ==================== //
+
+  let lock_acquired = mark_generic_download_job_pending_and_grab_lock(&job_state.mysql_pool, job.id).await?;
+
+  if !lock_acquired {
+    warn!("Could not acquire job lock for: {}", &job.id.0);
+    return Ok(())
+  }
+
+  // ==================== SETUP TEMP DIRS ==================== //
+
+  let temp_dir = format!("temp_{}", job.id.0);
+  let temp_dir = TempDir::new(&temp_dir)?;
+
+  // ==================== DOWNLOAD MODEL FILE ==================== //
+
+  info!("Calling downloader...");
+
+  redis_logger.log_status("downloading model")?;
+
+  let download_url = job.download_url.as_ref()
+      .map(|c| c.to_string())
+      .unwrap_or("".to_string());
+
+  if is_bad_tts_model_download_url(&download_url)? {
+    warn!("Bad download URL: `{}`", &download_url);
+    return Err(anyhow!("Bad download URL: `{}`", &download_url));
+  }
+
+  let download_filename = match job_state.google_drive_downloader.download_file(&download_url, &temp_dir).await {
+    Ok(filename) => filename,
+    Err(e) => {
+      safe_delete_temp_directory(&temp_dir);
+      return Err(e);
+    }
+  };
+
+  // ==================== HANDLE DIFFERENT DOWNLOAD TYPES ==================== //
+
+  match job.download_type {
+    GenericDownloadType::HifiGan => {
+      process_hifigan_vocoder(job_state, job, &temp_dir).await?;
+    }
+  }
+
+  // =====================================================
+
+  info!("Marking job complete...");
+  mark_generic_download_job_done(
+    &job_state.mysql_pool,
+    job,
+    true,
+    Some(&entity_token),
+    Some(&entity_type),
+  ).await?;
+
+  info!("Saved model record: {}", id);
+
+  job_state.firehose_publisher.publish_generic_download_finished(&job.creator_user_token, &model_token)
+      .await
+      .map_err(|e| {
+        warn!("error publishing event: {:?}", e);
+        anyhow!("error publishing event")
+      })?;
+
+  redis_logger.log_status("done")?;
+
+  info!("Job done: {}", job.id.0);
+
+  Ok(())
+}

crates/service/generic_download_job/src/main.rs

@@ -10,6 +10,8 @@
 
 #[macro_use] extern crate serde_derive;
 
+pub mod job_state;
+pub mod job_steps;
 pub mod script_execution;
 
 use anyhow::anyhow;
@@ -24,16 +26,20 @@ use container_common::filesystem::check_file_exists::check_file_exists;
 use container_common::filesystem::safe_delete_temp_directory::safe_delete_temp_directory;
 use container_common::filesystem::safe_delete_temp_file::safe_delete_temp_file;
 use container_common::hashing::hash_file_sha2::hash_file_sha2;
+use crate::job_state::JobState;
+use crate::job_steps::process_single_job::process_single_job;
+use crate::script_execution::hifigan_model_check_command::HifiGanModelCheckCommand;
 use crate::script_execution::tacotron_model_check_command::TacotronModelCheckCommand;
 use crate::script_execution::talknet_model_check_command::TalknetModelCheckCommand;
 use database_queries::column_types::tts_model_type::TtsModelType;
 use database_queries::mediators::badge_granter::BadgeGranter;
 use database_queries::mediators::firehose_publisher::FirehosePublisher;
+use database_queries::queries::generic_download::job::list_available_generic_download_jobs::{AvailableDownloadJob, list_available_generic_download_jobs};
+use database_queries::queries::generic_download::job::mark_generic_download_job_failure::mark_generic_download_job_failure;
 use database_queries::queries::tts::tts_download_jobs::tts_download_job_queries::TtsUploadJobRecord;
 use database_queries::queries::tts::tts_download_jobs::tts_download_job_queries::grab_job_lock_and_mark_pending;
 use database_queries::queries::tts::tts_download_jobs::tts_download_job_queries::insert_tts_model;
 use database_queries::queries::tts::tts_download_jobs::tts_download_job_queries::mark_tts_upload_job_done;
-use database_queries::queries::tts::tts_download_jobs::tts_download_job_queries::mark_tts_upload_job_failure;
 use database_queries::queries::tts::tts_download_jobs::tts_download_job_queries::query_tts_upload_job_records;
 use google_drive_common::google_drive_download_command::GoogleDriveDownloadCommand;
 use jobs_common::noop_logger::NoOpLogger;
@@ -53,7 +59,6 @@ use std::time::Duration;
 use storage_buckets_common::bucket_client::BucketClient;
 use storage_buckets_common::bucket_path_unifier::BucketPathUnifier;
 use tempdir::TempDir;
-use crate::script_execution::hifigan_model_check_command::HifiGanModelCheckCommand;
 
 // Buckets
 const ENV_ACCESS_KEY : &'static str = "ACCESS_KEY";
@@ -64,37 +69,6 @@ const ENV_BUCKET_ROOT : &'static str = "TTS_DOWNLOAD_BUCKET_ROOT";
 
 const DEFAULT_TEMP_DIR: &'static str = "/tmp";
 
-struct Downloader {
-  pub download_temp_directory: PathBuf,
-  pub mysql_pool: MySqlPool,
-
-  pub redis_pool: r2d2::Pool<RedisConnectionManager>,
-
-  pub bucket_client: BucketClient,
-  pub firehose_publisher: FirehosePublisher,
-  pub badge_granter: BadgeGranter,
-  pub google_drive_downloader: GoogleDriveDownloadCommand,
-
-  pub bucket_path_unifier: BucketPathUnifier,
-
-  pub hifigan_model_check_command: HifiGanModelCheckCommand,
-
-  // Command to run
-  pub download_script: String,
-  // Root to store TTS results
-  pub bucket_root_tts_model_uploads: String,
-
-  // Sleep between batches
-  pub job_batch_wait_millis: u64,
-
-  // How long to wait between log lines
-  pub no_op_logger_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]
 async fn main() -> AnyhowResult<()> {
   easyenv::init_all_with_default_logging(Some(DEFAULT_RUST_LOG));
@@ -106,7 +80,7 @@ async fn main() -> AnyhowResult<()> {
   let server_hostname = hostname::get()
     .ok()
     .and_then(|h| h.into_string().ok())
-    .unwrap_or("tts-download-job".to_string());
+    .unwrap_or("generic-download-job".to_string());
 
   info!("Hostname: {}", &server_hostname);
 
@@ -170,22 +144,17 @@ async fn main() -> AnyhowResult<()> {
     firehose_publisher: firehose_publisher.clone(), // NB: Also safe
   };
 
-  let hifigan_root_code_directory = easyenv::get_env_string_required("HIFIGAN_ROOT_CODE_DIRECTORY")?;
-  let hifigan_virtual_env_activation_command = easyenv::get_env_string_or_default(
-    "HIFIGAN_VIRTUAL_ENV_ACTIVATION_COMMAND",
-    "source python-tacotron/bin/activate");
-
-  let hifigan_model_check_script_name = easyenv::get_env_string_or_default(
-    "HIFIGAN_MODEL_CHECK_SCRIPT_NAME",
-    "vocodes_model_check_hifigan.py");
-
   let hifigan_model_check_command= HifiGanModelCheckCommand::new(
-    &hifigan_root_code_directory,
-    &hifigan_virtual_env_activation_command,
-    &hifigan_model_check_script_name,
+    &easyenv::get_env_string_required("HIFIGAN_ROOT_CODE_DIRECTORY")?,
+    &easyenv::get_env_string_or_default(
+    "HIFIGAN_VIRTUAL_ENV_ACTIVATION_COMMAND",
+    "source python-tacotron/bin/activate"),
+  &easyenv::get_env_string_or_default(
+    "HIFIGAN_MODEL_CHECK_SCRIPT_NAME",
+    "vocodes_model_check_hifigan.py"),
   );
 
-  let downloader = Downloader {
+  let job_state = JobState {
     download_temp_directory: temp_directory,
     mysql_pool,
     redis_pool,
@@ -202,7 +171,7 @@ async fn main() -> AnyhowResult<()> {
     no_op_logger_millis: common_env.no_op_logger_millis,
   };
 
-  main_loop(downloader).await;
+  main_loop(job_state).await;
 
   Ok(())
 }
@@ -210,16 +179,16 @@ async fn main() -> AnyhowResult<()> {
 const START_TIMEOUT_MILLIS : u64 = 500;
 const INCREASE_TIMEOUT_MILLIS : u64 = 1000;
 
-async fn main_loop(downloader: Downloader) {
+async fn main_loop(job_state: JobState) {
   let mut error_timeout_millis = START_TIMEOUT_MILLIS;
 
-  let mut noop_logger = NoOpLogger::new(downloader.no_op_logger_millis as i64);
+  let mut noop_logger = NoOpLogger::new(job_state.no_op_logger_millis as i64);
 
   loop {
     let num_records = 1;
-    let query_result = query_tts_upload_job_records(&downloader.mysql_pool, num_records).await;
+    let maybe_available_jobs = list_available_generic_download_jobs(&job_state.mysql_pool, num_records).await;
 
-    let jobs = match query_result {
+    let jobs = match maybe_available_jobs {
       Ok(jobs) => jobs,
       Err(e) => {
         warn!("Error querying jobs: {:?}", e);
@@ -232,11 +201,11 @@ async fn main_loop(downloader: Downloader) {
     if jobs.is_empty() {
       noop_logger.log_after_awhile();
 
-      std::thread::sleep(Duration::from_millis(downloader.job_batch_wait_millis));
+      std::thread::sleep(Duration::from_millis(job_state.job_batch_wait_millis));
       continue;
     }
 
-    let result = process_jobs(&downloader, jobs).await;
+    let result = process_jobs(&job_state, jobs).await;
 
     match result {
       Ok(_) => {},
@@ -250,23 +219,23 @@ async fn main_loop(downloader: Downloader) {
 
     error_timeout_millis = START_TIMEOUT_MILLIS; // reset
 
-    std::thread::sleep(Duration::from_millis(downloader.job_batch_wait_millis));
+    std::thread::sleep(Duration::from_millis(job_state.job_batch_wait_millis));
   }
 }
 
-async fn process_jobs(downloader: &Downloader, jobs: Vec<TtsUploadJobRecord>) -> AnyhowResult<()> {
+async fn process_jobs(job_state: &JobState, jobs: Vec<AvailableDownloadJob>) -> AnyhowResult<()> {
   for job in jobs.into_iter() {
-    let result = process_job(downloader, &job).await;
+    let result = process_single_job(job_state, &job).await;
     match result {
       Ok(_) => {},
       Err(e) => {
         warn!("Failure to process job: {:?}", e);
         let failure_reason = "";
-        let _r = mark_tts_upload_job_failure(
-          &downloader.mysql_pool,
+        let _r = mark_generic_download_job_failure(
+          &job_state.mysql_pool,
           &job,
           failure_reason,
-          downloader.job_max_attempts
+          job_state.job_max_attempts
         ).await;
       }
     }
@@ -287,191 +256,3 @@ fn read_metadata_file(filename: &PathBuf) -> AnyhowResult<FileMetadata> {
   Ok(serde_json::from_str(&buffer)?)
 }
 
-async fn process_job(downloader: &Downloader, job: &TtsUploadJobRecord) -> AnyhowResult<()> {
-  // TODO: 1. Mark processing. (DONE)
-  // TODO: 2. Download. (DONE)
-  // TODO: 3. Upload. (DONE)
-  // TODO: 4. Save record. (DONE)
-  // TODO: 5. Mark job done. (DONE)
-
-  let mut redis = downloader.redis_pool.get()?;
-  let mut redis_logger = RedisJobStatusLogger::new_tts_download(
-    &mut redis,
-    &job.token);
-
-  // ==================== ATTEMPT TO GRAB JOB LOCK ==================== //
-
-  let lock_acquired = grab_job_lock_and_mark_pending(&downloader.mysql_pool, job).await?;
-
-  if !lock_acquired {
-    warn!("Could not acquire job lock for: {}", &job.id);
-    return Ok(())
-  }
-
-  // ==================== SETUP TEMP DIRS ==================== //
-
-  let temp_dir = format!("temp_{}", job.id);
-  let temp_dir = TempDir::new(&temp_dir)?;
-
-  // ==================== DOWNLOAD MODEL FILE ==================== //
-
-  info!("Calling downloader...");
-
-  redis_logger.log_status("downloading model")?;
-
-  let download_url = job.download_url.as_ref()
-    .map(|c| c.to_string())
-    .unwrap_or("".to_string());
-
-  if is_bad_tts_model_download_url(&download_url)? {
-    warn!("Bad download URL: `{}`", &download_url);
-    return Err(anyhow!("Bad download URL: `{}`", &download_url));
-  }
-
-  let download_filename = match downloader.google_drive_downloader.download_file(&download_url, &temp_dir).await {
-    Ok(filename) => filename,
-    Err(e) => {
-      safe_delete_temp_directory(&temp_dir);
-      return Err(e);
-    }
-  };
-
-  let model_type = if download_filename.to_lowercase().ends_with("zip") {
-    TtsModelType::Talknet
-  } else {
-    TtsModelType::Tacotron2
-  };
-
-  info!("Uploaded model type: {:?}", model_type);
-
-  if model_type == TtsModelType::Talknet {
-    // TODO: Finish supporting TalkNet
-    warn!("Unsupported model type");
-    return Err(anyhow!("unsupported model type"));
-  }
-
-  // ==================== RUN MODEL CHECK ==================== //
-
-  info!("Checking that model is valid...");
-
-  redis_logger.log_status("checking model")?;
-
-  let file_path = PathBuf::from(download_filename.clone());
-
-  let output_metadata_fs_path = temp_dir.path().join("metadata.json");
-
-  match model_type {
-    TtsModelType::Tacotron2 => {
-      let result = downloader.tacotron_tts_check.execute(
-        &file_path,
-        &output_metadata_fs_path,
-        false,
-      );
-
-      if let Err(e) = result {
-        safe_delete_temp_file(&file_path);
-        safe_delete_temp_directory(&temp_dir);
-      }
-    },
-    TtsModelType::Talknet => {
-      let result = downloader.talknet_tts_check.execute(
-        &file_path,
-        &output_metadata_fs_path,
-        false,
-      );
-
-      if let Err(e) = result {
-        safe_delete_temp_file(&file_path);
-        safe_delete_temp_directory(&temp_dir);
-      }
-    },
-  }
-
-  // ==================== CHECK ALL FILES EXIST AND GET METADATA ==================== //
-
-  info!("Checking that metadata output file exists...");
-
-  check_file_exists(&output_metadata_fs_path)?;
-
-  let file_metadata = match read_metadata_file(&output_metadata_fs_path) {
-    Ok(metadata) => metadata,
-    Err(e) => {
-      safe_delete_temp_file(&file_path);
-      safe_delete_temp_file(&output_metadata_fs_path);
-      safe_delete_temp_directory(&temp_dir);
-      return Err(e);
-    }
-  };
-
-  // ==================== UPLOAD MODEL FILE ==================== //
-
-  info!("Uploading model to GCS...");
-
-  let private_bucket_hash = hash_file_sha2(&download_filename)?;
-
-  info!("File hash: {}", private_bucket_hash);
-
-  let synthesizer_model_bucket_path = match model_type {
-    TtsModelType::Tacotron2 => downloader.bucket_path_unifier.tts_synthesizer_path(&private_bucket_hash),
-    TtsModelType::Talknet => downloader.bucket_path_unifier.tts_zipped_synthesizer_path(&private_bucket_hash),
-  };
-
-  info!("Destination bucket path: {:?}", &synthesizer_model_bucket_path);
-
-  redis_logger.log_status("uploading model")?;
-
-  if let Err(e) = downloader.bucket_client.upload_filename(&synthesizer_model_bucket_path, &file_path).await {
-    safe_delete_temp_file(&output_metadata_fs_path);
-    safe_delete_temp_file(&file_path);
-    safe_delete_temp_directory(&temp_dir);
-    return Err(e);
-  }
-
-  // ==================== DELETE DOWNLOADED FILE ==================== //
-
-  // NB: We should be using a tempdir, but to make absolutely certain we don't overflow the disk...
-  safe_delete_temp_file(&output_metadata_fs_path);
-  safe_delete_temp_file(&file_path);
-  safe_delete_temp_directory(&temp_dir);
-
-  // ==================== SAVE RECORDS ==================== //
-
-  info!("Saving model record...");
-  let (id, model_token) = insert_tts_model(
-    &downloader.mysql_pool,
-    job,
-    &private_bucket_hash,
-    synthesizer_model_bucket_path,
-    file_metadata.file_size_bytes)
-    .await?;
-
-  info!("Marking job complete...");
-  mark_tts_upload_job_done(
-    &downloader.mysql_pool,
-    job,
-    true,
-    Some(&model_token)
-  ).await?;
-
-  info!("Saved model record: {}", id);
-
-  downloader.firehose_publisher.publish_tts_model_upload_finished(&job.creator_user_token, &model_token)
-      .await
-      .map_err(|e| {
-        warn!("error publishing event: {:?}", e);
-        anyhow!("error publishing event")
-      })?;
-
-  downloader.badge_granter.maybe_grant_tts_model_uploads_badge(&job.creator_user_token)
-      .await
-      .map_err(|e| {
-        warn!("error maybe awarding badge: {:?}", e);
-        anyhow!("error maybe awarding badge")
-      })?;
-
-  redis_logger.log_status("done")?;
-
-  info!("Job done: {}", job.id);
-
-  Ok(())
-}