report on w2l template download status

src/http_server/endpoints/w2l/get_w2l_upload_template_job_status.rs

@@ -15,10 +15,12 @@ use crate::http_server::web_utils::response_error_helpers::to_simple_json_error;
 use crate::http_server::web_utils::session_checker::SessionRecord;
 use crate::server_state::ServerState;
 use crate::util::random_crockford_token::random_crockford_token;
+use crate::util::redis::redis_keys::RedisKeys;
 use crate::validations::passwords::validate_passwords;
 use crate::validations::username::validate_username;
 use derive_more::{Display, Error};
 use log::{info, warn, log};
+use r2d2_redis::redis::Commands;
 use regex::Regex;
 use sqlx::MySqlPool;
 use sqlx::error::DatabaseError;
@@ -36,7 +38,13 @@ pub struct GetW2lUploadTemplateStatusPathInfo {
 pub struct W2lUploadTemplateJobStatusForResponse {
   pub job_token: String,
 
+  /// Primary status from the database (a state machine).
   pub status: String,
+
+  /// Extra, temporary status from Redis.
+  /// This can denote inference progress, and the Python code can write to it.
+  pub maybe_extra_status_description: Option<String>,
+
   pub attempt_count: u8,
   pub maybe_template_token: Option<String>,
 
@@ -127,9 +135,29 @@ WHERE jobs.token = ?
     }
   };
 
+  let mut redis = server_state.redis_pool
+      .get()
+      .map_err(|e| {
+        warn!("redis error: {:?}", e);
+        GetW2lUploadTemplateStatusError::ServerError
+      })?;
+
+  let extra_status_key = RedisKeys::w2l_download_extra_status_info(&path.token);
+  let maybe_extra_status_description : Option<String> = match redis.get(&extra_status_key) {
+    Ok(Some(status)) => {
+      Some(status)
+    },
+    Ok(None) => None,
+    Err(e) => {
+      warn!("redis error: {:?}", e);
+      None // Fail open
+    },
+  };
+
   let template_for_response = W2lUploadTemplateJobStatusForResponse {
     job_token: record.job_token.clone(),
     status: record.status.clone(),
+    maybe_extra_status_description,
     attempt_count: record.attempt_count as u8,
     maybe_template_token: record.maybe_template_token.clone(),
     created_at: record.created_at.clone(),

src/util/redis/redis_job_status_logger.rs

@@ -19,6 +19,11 @@ impl <'a> RedisJobStatusLogger <'a> {
     Self { redis, status_key }
   }
 
+  pub fn new_w2l_download(redis: &'a mut RedisPool, w2l_job_token: &str) -> Self {
+    let status_key = RedisKeys::w2l_download_extra_status_info(w2l_job_token);
+    Self { redis, status_key }
+  }
+
   pub fn new_w2l_inference(redis: &'a mut RedisPool, w2l_job_token: &str) -> Self {
     let status_key = RedisKeys::w2l_inference_extra_status_info(w2l_job_token);
     Self { redis, status_key }

src/util/redis/redis_keys.rs

@@ -22,6 +22,12 @@ impl RedisKeys {
     format!("ttsInferenceExtraStatus:{}", inference_job_token)
   }
 
+  /// We write extra in-progress status information to keys.
+  /// These keys should have a TTL.
+  pub fn w2l_download_extra_status_info(inference_job_token: &str) -> String {
+    format!("w2lDownloadExtraStatus:{}", inference_job_token)
+  }
+
   /// We write extra in-progress status information to keys.
   /// These keys should have a TTL.
   pub fn w2l_inference_extra_status_info(inference_job_token: &str) -> String {

src/w2l_download_job.rs

@@ -43,13 +43,17 @@ use crate::util::filesystem::check_file_exists;
 use crate::util::hashing::hash_file_sha2::hash_file_sha2;
 use crate::util::noop_logger::NoOpLogger;
 use crate::util::random_crockford_token::random_crockford_token;
+use crate::util::redis::redis_job_status_logger::RedisJobStatusLogger;
 use data_encoding::{HEXUPPER, HEXLOWER, HEXLOWER_PERMISSIVE};
 use log::{warn, info};
+use r2d2_redis::RedisConnectionManager;
+use r2d2_redis::r2d2;
 use ring::digest::{Context, Digest, SHA256};
 use sqlx::MySqlPool;
 use sqlx::mysql::MySqlPoolOptions;
 use std::fs::{File, metadata};
 use std::io::{BufReader, Read};
+use std::ops::Deref;
 use std::path::{PathBuf, Path};
 use std::process::Command;
 use std::thread;
@@ -80,6 +84,8 @@ struct Downloader {
   pub download_temp_directory: PathBuf,
   pub mysql_pool: MySqlPool,
 
+  pub redis_pool: r2d2::Pool<RedisConnectionManager>,
+
   pub private_bucket_client: BucketClient,
   pub public_bucket_client: BucketClient,
 
@@ -195,7 +201,17 @@ async fn main() -> AnyhowResult<()> {
     .max_connections(5)
     .connect(&db_connection_string)
     .await?;
-  
+
+  let common_env = CommonEnv::read_from_env()?;
+
+  info!("Connecting to redis...");
+
+  let redis_manager =
+      RedisConnectionManager::new(common_env.redis_connection_string.deref())?;
+
+  let redis_pool = r2d2::Pool::builder()
+      .build(redis_manager)?;
+
   let firehose_publisher = FirehosePublisher {
     mysql_pool: mysql_pool.clone(), // NB: Pool is sync/send/clone-safe
   };
@@ -205,11 +221,10 @@ async fn main() -> AnyhowResult<()> {
     firehose_publisher: firehose_publisher.clone(), // NB: Also safe
   };
 
-  let common_env = CommonEnv::read_from_env()?;
-
   let downloader = Downloader {
     download_temp_directory: temp_directory,
     mysql_pool,
+    redis_pool,
     public_bucket_client,
     private_bucket_client,
     download_script,
@@ -335,6 +350,11 @@ async fn process_job(downloader: &Downloader, job: &W2lTemplateUploadJobRecord)
   // TODO: 7. Save record. (DONE)
   // TODO: 8. Mark job done. (DONE)
 
+  let mut redis = downloader.redis_pool.get()?;
+  let mut redis_logger = RedisJobStatusLogger::new_w2l_download(
+    &mut redis,
+    &job.token);
+
   // ==================== ATTEMPT TO GRAB JOB LOCK ==================== //
 
   let lock_acquired = grab_job_lock_and_mark_pending(&downloader.mysql_pool, job).await?;
@@ -351,6 +371,8 @@ async fn process_job(downloader: &Downloader, job: &W2lTemplateUploadJobRecord)
 
   // ==================== DOWNLOAD FILE ==================== //
 
+  redis_logger.log_status("downloading media")?;
+
   let download_url = job.download_url.as_ref()
       .map(|c| c.to_string())
       .unwrap_or("".to_string());
@@ -366,6 +388,8 @@ async fn process_job(downloader: &Downloader, job: &W2lTemplateUploadJobRecord)
 
   // ==================== PROCESS FACES ==================== //
 
+  redis_logger.log_status("detecting faces")?;
+
   // This is the Python Pickle file with all the face frames.
   // We'll retain it forever since it's expensive to compute.
   let cached_faces_filename = format!("{}_detected_faces.pickle", &download_filename);
@@ -421,6 +445,8 @@ async fn process_job(downloader: &Downloader, job: &W2lTemplateUploadJobRecord)
   let mut maybe_video_preview_filename : Option<PathBuf> = None;
   let mut maybe_video_preview_object_name : Option<String> = None;
 
+  redis_logger.log_status("generating preview")?;
+
   if file_metadata.is_video {
     let preview_filename = format!("{}_preview.webp", &download_filename);
 
@@ -461,6 +487,8 @@ async fn process_job(downloader: &Downloader, job: &W2lTemplateUploadJobRecord)
 
   // ==================== UPLOAD TO BUCKETS ==================== //
 
+  redis_logger.log_status("uploading results")?;
+
   info!("Image/video destination bucket path: {}", full_object_path);
 
   // Full path to cached faces
@@ -575,6 +603,8 @@ async fn process_job(downloader: &Downloader, job: &W2lTemplateUploadJobRecord)
     thread::sleep(Duration::from_millis(downloader.debug_job_end_sleep_millis));
   }
 
+  redis_logger.log_status("done")?;
+
   info!("Job {} complete success! Downloaded, processed, and uploaded. Saved model record: {}",
         job.id, id);