fix redis tx over database tx pooling issues

crates/service/web/storyteller_web/src/http_server/endpoints/inference_job/get/batch_get_inference_job_status_handler.rs

@@ -10,6 +10,7 @@ use crate::http_server::endpoints::inference_job::utils::extractors::extract_liv
 use crate::http_server::endpoints::inference_job::utils::extractors::extract_polymorphic_inference_args::extract_polymorphic_inference_args;
 use crate::http_server::endpoints::media_files::helpers::get_media_domain::get_media_domain;
 use crate::http_server::web_utils::filter_model_name::maybe_filter_model_name;
+use crate::http_server::web_utils::job_keepalives::write_job_keepalives;
 use crate::state::server_state::ServerState;
 use actix_web::web::Json;
 use actix_web::{web, HttpRequest};
@@ -23,22 +24,14 @@ use enums::by_table::generic_inference_jobs::frontend_failure_category::Frontend
 use enums::by_table::generic_inference_jobs::inference_category::InferenceCategory;
 use enums::common::job_status_plus::JobStatusPlus;
 use enums::no_table::style_transfer::style_transfer_name::StyleTransferName;
-use log::{error, warn};
+use log::warn;
 use mysql_queries::queries::generic_inference::web::batch_get_inference_job_status::batch_get_inference_job_status;
 use mysql_queries::queries::generic_inference::web::job_status::GenericInferenceJobStatus;
-use redis::Commands;
-use redis_common::redis_keys::RedisKeys;
 use server_environment::ServerEnvironment;
 use tokens::tokens::generic_inference_jobs::InferenceJobToken;
 use tokens::tokens::prompts::PromptToken;
 use utoipa::{IntoParams, ToSchema};
 
-/// For certain jobs or job classes (eg. non-premium), we kill the jobs if the user hasn't
-/// maintained a keepalive. This prevents wasted work when users who are unlikely to return
-/// navigate away. Premium users have accounts and can always return to the site, so they
-/// typically do not require keepalive.
-const JOB_KEEPALIVE_TTL_SECONDS : u64 = 60 * 3;
-
 #[derive(Deserialize, ToSchema, IntoParams)]
 pub struct BatchGetInferenceJobStatusQueryParams {
   /// A grab bag of tokens of various types.
@@ -191,32 +184,14 @@ pub async fn batch_get_inference_job_status_handler(
         CommonWebError::from_anyhow_error(err)
       })?;
 
-  let mut redis = server_state.redis_pool
-      .get()
-      .map_err(|e| {
-        warn!("Redis pool error: {:?}", e);
-        CommonWebError::from_error(e)
-      })?;
-
   // TODO(bt,2024-04-22): Look up the extra redis statuses per item.
 
-  let keepalive_keys = records.iter()
+  let keepalive_job_tokens = records.iter()
       .filter(|record| record.is_keepalive_required)
-      .map(|record| RedisKeys::generic_inference_keepalive(record.job_token.as_str()))
+      .map(|record| record.job_token.as_str())
       .collect::<Vec<_>>();
 
-  for key in keepalive_keys.iter() {
-    // TODO(bt,2024-04-22): There is no msetex. We'll need to run a Redis pipeline here.
-    //  https://stackoverflow.com/questions/16423342/redis-multi-set-with-a-ttl
-    let _: Option<String> = match redis.set_ex(key, "1", JOB_KEEPALIVE_TTL_SECONDS) {
-      Ok(Some(status)) => Some(status),
-      Ok(None) => None,
-      Err(e) => {
-        error!("redis error setting job keepalive: {:?}", e);
-        None // Fail open (which in this case is bad! it will kill jobs if cluster has many jobs / is slow!)
-      },
-    };
-  }
+  write_job_keepalives(&server_state.redis_pool, &keepalive_job_tokens);
 
   let media_domain = get_media_domain(&http_request);
 

crates/service/web/storyteller_web/src/http_server/endpoints/inference_job/get/get_inference_job_status_handler.rs

@@ -10,6 +10,7 @@ use crate::http_server::endpoints::inference_job::utils::extractors::extract_liv
 use crate::http_server::endpoints::inference_job::utils::extractors::extract_polymorphic_inference_args::extract_polymorphic_inference_args;
 use crate::http_server::endpoints::media_files::helpers::get_media_domain::get_media_domain;
 use crate::http_server::web_utils::filter_model_name::maybe_filter_model_name;
+use crate::http_server::web_utils::job_keepalives::write_job_keepalives;
 use crate::state::server_state::ServerState;
 use actix_web::web::{Json, Path};
 use actix_web::{web, HttpRequest};
@@ -32,12 +33,6 @@ use tokens::tokens::generic_inference_jobs::InferenceJobToken;
 use tokens::tokens::prompts::PromptToken;
 use utoipa::ToSchema;
 
-/// For certain jobs or job classes (eg. non-premium), we kill the jobs if the user hasn't
-/// maintained a keepalive. This prevents wasted work when users who are unlikely to return
-/// navigate away. Premium users have accounts and can always return to the site, so they
-/// typically do not require keepalive.
-const JOB_KEEPALIVE_TTL_SECONDS : u64 = 60 * 3;
-
 /// For the URL PathInfo
 #[derive(Deserialize, ToSchema)]
 pub struct GetInferenceJobStatusPathInfo {
@@ -202,18 +197,10 @@ pub async fn get_inference_job_status_handler(
   };
 
   if record.is_keepalive_required {
-    // TODO(bt,2023-05-21): Make async.
-    let keepalive_key = RedisKeys::generic_inference_keepalive(path.token.as_str());
-    let _: Option<String> = match redis.set_ex(&extra_status_key, "1", JOB_KEEPALIVE_TTL_SECONDS) {
-      Ok(Some(status)) => {
-        Some(status)
-      },
-      Ok(None) => None,
-      Err(e) => {
-        error!("redis error setting job keepalive: {:?}", e);
-        None // Fail open (which in this case is bad! it will kill jobs if cluster has many jobs / is slow!)
-      },
-    };
+    // NB: This used to SET the *extra status* key instead of the keepalive key — which both
+    // clobbered the extra status with "1" (see the workaround above) and let keepalive-required
+    // jobs get killed while being actively polled.
+    write_job_keepalives(&server_state.redis_pool, &[path.token.as_str()]);
   }
 
   let media_domain = get_media_domain(&http_request);

crates/service/web/storyteller_web/src/http_server/endpoints/inference_job/list/list_session_jobs_handler.rs

@@ -10,6 +10,7 @@ use crate::http_server::endpoints::inference_job::utils::extractors::extract_liv
 use crate::http_server::endpoints::inference_job::utils::extractors::extract_polymorphic_inference_args::extract_polymorphic_inference_args;
 use crate::http_server::endpoints::media_files::helpers::get_media_domain::get_media_domain;
 use crate::http_server::web_utils::filter_model_name::maybe_filter_model_name;
+use crate::http_server::web_utils::job_keepalives::write_job_keepalives;
 use crate::state::server_state::ServerState;
 use actix_web::web::Json;
 use actix_web::{web, HttpRequest};
@@ -22,23 +23,14 @@ use enums::api_safe::by_table::generic_inference_jobs::frontend_failure_category
 use enums::api_safe::by_table::generic_inference_jobs::frontend_failure_category_for_old_clients::FrontendFailureCategoryForOldClients;
 use enums::by_table::generic_inference_jobs::inference_category::InferenceCategory;
 use enums::common::job_status_plus::JobStatusPlus;
-use log::{error, warn};
 use mysql_queries::queries::generic_inference::web::job_status::GenericInferenceJobStatus;
 use mysql_queries::queries::generic_inference::web::list_session_jobs::{list_session_jobs_from_connection, ListSessionJobsForUserArgs, SessionUser};
 use primitives::numerics::i64_to_u64_zero_clamped::i64_to_u64_zero_clamped;
-use redis::Commands;
-use redis_common::redis_keys::RedisKeys;
 use server_environment::ServerEnvironment;
 use tokens::tokens::generic_inference_jobs::InferenceJobToken;
 use tokens::tokens::media_files::MediaFileToken;
 use utoipa::IntoParams;
 
-/// For certain jobs or job classes (eg. non-premium), we kill the jobs if the user hasn't
-/// maintained a keepalive. This prevents wasted work when users who are unlikely to return
-/// navigate away. Premium users have accounts and can always return to the site, so they
-/// typically do not require keepalive.
-const JOB_KEEPALIVE_TTL_SECONDS : u64 = 60 * 3;
-
 /// List job statuses for jobs that are associated with the user's session.
 ///
 /// The user must be logged in. This endpoint will only show the last 36 hours
@@ -109,29 +101,18 @@ pub async fn list_session_jobs_handler(
   // NB: Since this is publicly exposed, we don't query sensitive data.
   let records = list_session_jobs_from_connection(args, &mut mysql_connection).await?;
 
-  let mut redis = server_state.redis_pool
-      .get()
-      .map_err(CommonWebError::from_error)?;
+  // The Redis keepalive writes below are blocking I/O; release the pooled MySQL connection
+  // first so it isn't held hostage across Redis round trips (this starves the pool).
+  drop(mysql_connection);
 
   // TODO(bt,2024-04-22): Look up the extra redis statuses per item.
 
-  let keepalive_keys = records.iter()
+  let keepalive_job_tokens = records.iter()
       .filter(|record| record.is_keepalive_required)
-      .map(|record| RedisKeys::generic_inference_keepalive(record.job_token.as_str()))
+      .map(|record| record.job_token.as_str())
       .collect::<Vec<_>>();
 
-  for key in keepalive_keys.iter() {
-    // TODO(bt,2024-04-22): There is no msetex. We'll need to run a Redis pipeline here.
-    //  https://stackoverflow.com/questions/16423342/redis-multi-set-with-a-ttl
-    let _: Option<String> = match redis.set_ex(key, "1", JOB_KEEPALIVE_TTL_SECONDS) {
-      Ok(Some(status)) => Some(status),
-      Ok(None) => None,
-      Err(e) => {
-        error!("redis error setting job keepalive: {:?}", e);
-        None // Fail open (which in this case is bad! it will kill jobs if cluster has many jobs / is slow!)
-      },
-    };
-  }
+  write_job_keepalives(&server_state.redis_pool, &keepalive_job_tokens);
 
   let media_domain = get_media_domain(&http_request);
 

crates/service/web/storyteller_web/src/http_server/endpoints/omni_api/job_status/batch_get_job_status_handler.rs

@@ -3,13 +3,11 @@ use std::sync::Arc;
 use actix_web::web::Json;
 use actix_web::{web, HttpRequest};
 use actix_web_lab::extract::Query;
-use log::{error, warn};
-use redis::Commands;
+use log::warn;
 
 use artcraft_api_defs::omni_api::job_status::omni_api_batch_get_job_status::{OmniApiBatchGetJobStatusQueryParams, OmniApiBatchGetJobStatusSuccessResponse};
 use mysql_queries::queries::generic_inference::web::batch_get_inference_job_status::batch_get_inference_job_status;
 use mysql_queries::queries::generic_inference::web::job_status::GenericInferenceJobStatus;
-use redis_common::redis_keys::RedisKeys;
 use server_environment::ServerEnvironment;
 use tokens::tokens::generic_inference_jobs::InferenceJobToken;
 
@@ -18,14 +16,9 @@ use crate::http_server::common_responses::media::media_domain::MediaDomain;
 use crate::http_server::endpoints::media_files::helpers::get_media_domain::get_media_domain;
 use crate::http_server::endpoints::omni_api::job_status::record_to_payload::record_to_payload;
 use crate::http_server::user_lookup::api_keys::require_api_key_user::require_api_key_user;
+use crate::http_server::web_utils::job_keepalives::write_job_keepalives;
 use crate::state::server_state::ServerState;
 
-/// For certain jobs or job classes (eg. non-premium), we kill the jobs if the user hasn't
-/// maintained a keepalive. This prevents wasted work when users who are unlikely to return
-/// navigate away. Premium users have accounts and can always return to the site, so they
-/// typically do not require keepalive.
-const JOB_KEEPALIVE_TTL_SECONDS: u64 = 60 * 3;
-
 /// Get job statuses for a batch of multiple job tokens (Omni API, API-key authenticated).
 ///
 /// Mirrors `/v1/jobs/batch` but reads the caller's identity from the
@@ -82,31 +75,14 @@ pub async fn omni_api_batch_get_job_status_handler(
         CommonWebError::from_anyhow_error(err)
       })?;
 
-  let mut redis = server_state.redis_pool
-      .get()
-      .map_err(|e| {
-        warn!("Redis pool error: {:?}", e);
-        CommonWebError::from_error(e)
-      })?;
-
   // TODO(bt,2024-04-22): Look up the extra redis statuses per item.
 
-  let keepalive_keys = records.iter()
+  let keepalive_job_tokens = records.iter()
       .filter(|record| record.is_keepalive_required)
-      .map(|record| RedisKeys::generic_inference_keepalive(record.job_token.as_str()))
+      .map(|record| record.job_token.as_str())
       .collect::<Vec<_>>();
 
-  for key in keepalive_keys.iter() {
-    // TODO(bt,2024-04-22): There is no msetex. We'll need to run a Redis pipeline here.
-    let _: Option<String> = match redis.set_ex(key, "1", JOB_KEEPALIVE_TTL_SECONDS) {
-      Ok(Some(status)) => Some(status),
-      Ok(None) => None,
-      Err(e) => {
-        error!("redis error setting job keepalive: {:?}", e);
-        None // Fail open (which in this case is bad! it will kill jobs if cluster has many jobs / is slow!)
-      }
-    };
-  }
+  write_job_keepalives(&server_state.redis_pool, &keepalive_job_tokens);
 
   let media_domain = get_media_domain(&http_request);
 

crates/service/web/storyteller_web/src/http_server/endpoints/omni_api/job_status/get_job_status_handler.rs

@@ -2,25 +2,18 @@ use std::sync::Arc;
 
 use actix_web::web::{Json, Path};
 use actix_web::{web, HttpRequest};
-use log::{error, warn};
-use redis::Commands;
+use log::warn;
 
 use artcraft_api_defs::omni_api::job_status::omni_api_get_job_status::{OmniApiGetJobStatusPathInfo, OmniApiGetJobStatusSuccessResponse};
 use mysql_queries::queries::generic_inference::web::get_inference_job_status::get_inference_job_status;
-use redis_common::redis_keys::RedisKeys;
 
 use crate::http_server::common_responses::common_web_error::CommonWebError;
 use crate::http_server::endpoints::media_files::helpers::get_media_domain::get_media_domain;
 use crate::http_server::endpoints::omni_api::job_status::record_to_payload::record_to_payload;
 use crate::http_server::user_lookup::api_keys::require_api_key_user::require_api_key_user;
+use crate::http_server::web_utils::job_keepalives::write_job_keepalives;
 use crate::state::server_state::ServerState;
 
-/// For certain jobs or job classes (eg. non-premium), we kill the jobs if the user hasn't
-/// maintained a keepalive. This prevents wasted work when users who are unlikely to return
-/// navigate away. Premium users have accounts and can always return to the site, so they
-/// typically do not require keepalive.
-const JOB_KEEPALIVE_TTL_SECONDS: u64 = 60 * 3;
-
 /// Get the status for a single job (Omni API, API-key authenticated).
 ///
 /// Mirrors `/v1/jobs/job/{token}` but reads the caller's identity from the
@@ -72,24 +65,10 @@ pub async fn omni_api_get_job_status_handler(
     }
   };
 
-  let mut redis = server_state.redis_pool
-      .get()
-      .map_err(|e| {
-        warn!("Redis pool error: {:?}", e);
-        CommonWebError::from_error(e)
-      })?;
-
   if record.is_keepalive_required {
-    // TODO(bt,2023-05-21): Make async.
-    let extra_status_key = RedisKeys::generic_inference_extra_status_info(path.token.as_str());
-    let _: Option<String> = match redis.set_ex(&extra_status_key, "1", JOB_KEEPALIVE_TTL_SECONDS) {
-      Ok(Some(status)) => Some(status),
-      Ok(None) => None,
-      Err(e) => {
-        error!("redis error setting job keepalive: {:?}", e);
-        None // Fail open (which in this case is bad! it will kill jobs if cluster has many jobs / is slow!)
-      }
-    };
+    // NB: This used to SET the job's *extra status* key instead of the keepalive key, which
+    // clobbered the extra status with "1" and never actually refreshed the keepalive.
+    write_job_keepalives(&server_state.redis_pool, &[path.token.as_str()]);
   }
 
   let media_domain = get_media_domain(&http_request);

crates/service/web/storyteller_web/src/http_server/web_utils/job_keepalives.rs

@@ -0,0 +1,45 @@
+use log::error;
+use r2d2::Pool;
+use redis::Client;
+
+use redis_common::redis_keys::RedisKeys;
+
+/// For certain jobs or job classes (eg. non-premium), we kill the jobs if the user hasn't
+/// maintained a keepalive. This prevents wasted work when users who are unlikely to return
+/// navigate away. Premium users have accounts and can always return to the site, so they
+/// typically do not require keepalive.
+pub const JOB_KEEPALIVE_TTL_SECONDS: u64 = 60 * 3;
+
+/// Refresh the keepalive keys for the given jobs in a single pipelined Redis round trip.
+///
+/// Fails open: Redis errors are logged and swallowed rather than failing the request (a failed
+/// request wouldn't have refreshed the keepalive either, and the job data is already in hand).
+///
+/// The Redis I/O here is blocking, so callers MUST release any pooled MySQL connection before
+/// calling this — holding one across Redis round trips starves the MySQL pool (this crate has
+/// had pool-timeout incidents).
+pub fn write_job_keepalives(redis_pool: &Pool<Client>, job_tokens: &[&str]) {
+  if job_tokens.is_empty() {
+    return;
+  }
+
+  let mut redis = match redis_pool.get() {
+    Ok(connection) => connection,
+    Err(err) => {
+      // NB: Failing open can kill keepalive-required jobs if Redis is down for a while.
+      error!("redis pool error writing job keepalives: {:?}", err);
+      return;
+    }
+  };
+
+  let mut pipe = redis::pipe();
+  for job_token in job_tokens {
+    let key = RedisKeys::generic_inference_keepalive(job_token);
+    pipe.set_ex(key, "1", JOB_KEEPALIVE_TTL_SECONDS).ignore();
+  }
+
+  if let Err(err) = pipe.query::<()>(&mut *redis) {
+    // NB: Failing open can kill keepalive-required jobs if Redis is down for a while.
+    error!("redis error writing job keepalives: {:?}", err);
+  }
+}

crates/service/web/storyteller_web/src/http_server/web_utils/mod.rs

@@ -4,6 +4,7 @@ pub mod get_authorization_header_api_key;
 pub mod get_host_header;
 pub mod get_request_platform_type;
 pub mod handle_multipart_error;
+pub mod job_keepalives;
 pub mod open_zip_archive;
 pub mod read_multipart_field_bytes;
 pub mod redis_rate_limiter;