diff --git a/src/http_server/endpoints/w2l/get_w2l_upload_template_job_status.rs b/src/http_server/endpoints/w2l/get_w2l_upload_template_job_status.rs index 0674d2ccdc..6424bd7cb7 100644 --- a/src/http_server/endpoints/w2l/get_w2l_upload_template_job_status.rs +++ b/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, + pub attempt_count: u8, pub maybe_template_token: Option, @@ -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 = 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(), diff --git a/src/util/redis/redis_job_status_logger.rs b/src/util/redis/redis_job_status_logger.rs index e03d99d8ad..a407e80339 100644 --- a/src/util/redis/redis_job_status_logger.rs +++ b/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 } diff --git a/src/util/redis/redis_keys.rs b/src/util/redis/redis_keys.rs index afadb4f7e1..3455c60183 100644 --- a/src/util/redis/redis_keys.rs +++ b/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 { diff --git a/src/w2l_download_job.rs b/src/w2l_download_job.rs index 681318341f..b9044dea7d 100644 --- a/src/w2l_download_job.rs +++ b/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, + 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 = None; let mut maybe_video_preview_object_name : Option = 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);