BM-1854: fix(bento): aux worker cleanup/connection handling (#1294)

Various fixes/changes to bento cleanup:
-
1ddbcfb:
Env to allow turning off either cleanup cron job (if any issues in
future or if intending to manually handle db maintenance)
-
afcc428:
Increase aux max client connections default to handle the 4 concurrent
tasks. All other workers are sequentially processing so this should be
sufficient for now
-
2bd9755:
to avoid running on startup to avoid issues with the spinning on startup
where all try to connect at one time

Separated into diff commits if we only want a subset of these changes

Author: austinabell

bento/crates/workflow/src/lib.rs

@@ -161,6 +161,14 @@ pub struct Args {
     /// How often to clean up completed jobs
     #[clap(env, long, default_value_t = 60 * 60)]
     cleanup_poll_interval: u64,
+
+    /// Disable cron to clean up completed jobs in taskdb.
+    #[clap(env = "BENTO_DISABLE_COMPLETED_CLEANUP")]
+    disable_completed_cleanup: bool,
+
+    /// Disable cron to clean up stuck tasks in taskdb.
+    #[clap(env = "BENTO_DISABLE_STUCK_TASK_CLEANUP")]
+    disable_stuck_task_cleanup: bool,
 }
 
 /// Core agent context to hold all optional clients / pools and state
@@ -283,42 +291,46 @@ impl Agent {
 
         // Enable stuck task maintenance for aux workers
         if self.args.task_stream == AUX_WORK_TYPE {
-            let term_sig_copy = term_sig.clone();
-            let db_pool_copy = self.db_pool.clone();
-            let stuck_tasks_interval = self.args.stuck_tasks_poll_interval;
-            tokio::spawn(async move {
-                loop {
-                    if let Err(e) = Self::poll_for_stuck_tasks(
-                        term_sig_copy.clone(),
-                        db_pool_copy.clone(),
-                        stuck_tasks_interval,
-                    )
-                    .await
-                    {
-                        tracing::error!("[BENTO-WF-105] Stuck tasks cleanup failed: {:#}", e);
-                        time::sleep(time::Duration::from_secs(60)).await;
+            if !self.args.disable_stuck_task_cleanup {
+                let term_sig_copy = term_sig.clone();
+                let db_pool_copy = self.db_pool.clone();
+                let stuck_tasks_interval = self.args.stuck_tasks_poll_interval;
+                tokio::spawn(async move {
+                    loop {
+                        if let Err(e) = Self::poll_for_stuck_tasks(
+                            term_sig_copy.clone(),
+                            db_pool_copy.clone(),
+                            stuck_tasks_interval,
+                        )
+                        .await
+                        {
+                            tracing::error!("[BENTO-WF-105] Stuck tasks cleanup failed: {:#}", e);
+                            time::sleep(time::Duration::from_secs(60)).await;
+                        }
                     }
-                }
-            });
+                });
+            }
 
             // Enable completed job cleanup for aux workers
-            let term_sig_copy = term_sig.clone();
-            let db_pool_copy = self.db_pool.clone();
-            let cleanup_interval = self.args.cleanup_poll_interval;
-            tokio::spawn(async move {
-                loop {
-                    if let Err(e) = Self::poll_for_completed_job_cleanup(
-                        term_sig_copy.clone(),
-                        db_pool_copy.clone(),
-                        cleanup_interval,
-                    )
-                    .await
-                    {
-                        tracing::error!("[BENTO-WF-106] Completed job cleanup failed: {:#}", e);
-                        time::sleep(time::Duration::from_secs(cleanup_interval)).await;
+            if !self.args.disable_completed_cleanup {
+                let term_sig_copy = term_sig.clone();
+                let db_pool_copy = self.db_pool.clone();
+                let cleanup_interval = self.args.cleanup_poll_interval;
+                tokio::spawn(async move {
+                    loop {
+                        if let Err(e) = Self::poll_for_completed_job_cleanup(
+                            term_sig_copy.clone(),
+                            db_pool_copy.clone(),
+                            cleanup_interval,
+                        )
+                        .await
+                        {
+                            tracing::error!("[BENTO-WF-106] Completed job cleanup failed: {:#}", e);
+                            time::sleep(time::Duration::from_secs(cleanup_interval)).await;
+                        }
                     }
-                }
-            });
+                });
+            }
         }
 
         while !term_sig.load(Ordering::Relaxed) {
@@ -512,6 +524,9 @@ impl Agent {
         poll_interval: u64,
     ) -> Result<()> {
         while !term_sig.load(Ordering::Relaxed) {
+            // Sleep before each check to avoid running on startup
+            time::sleep(tokio::time::Duration::from_secs(poll_interval)).await;
+
             tracing::debug!("Checking for stuck pending tasks...");
 
             // First check if there are any stuck tasks
@@ -537,9 +552,6 @@ impl Agent {
                     tracing::info!("Fixed {} stuck pending tasks", fixed_count);
                 }
             }
-
-            // Sleep before next check
-            time::sleep(tokio::time::Duration::from_secs(poll_interval)).await;
         }
 
         Ok(())
@@ -555,15 +567,15 @@ impl Agent {
         poll_interval: u64,
     ) -> Result<()> {
         while !term_sig.load(Ordering::Relaxed) {
+            // Sleep before each check to avoid running on startup
+            time::sleep(tokio::time::Duration::from_secs(poll_interval)).await;
+
             tracing::debug!("Cleaning up completed jobs...");
 
             let cleared_count = taskdb::clear_completed_jobs(&db_pool).await?;
             if cleared_count > 0 {
                 tracing::info!("Cleared {} completed jobs", cleared_count);
             }
-
-            // Sleep before next cleanup
-            time::sleep(tokio::time::Duration::from_secs(poll_interval)).await;
         }
 
         Ok(())

bento/crates/workflow/src/redis.rs

@@ -4,7 +4,6 @@
 // as found in the LICENSE-BSL file.
 
 use anyhow::{Context, Result};
-use deadpool_redis::Connection;
 pub use deadpool_redis::{Config, Pool as RedisPool, Runtime, redis::AsyncCommands};
 use redis::{RedisResult, ToRedisArgs};
 

compose.yml

@@ -178,6 +178,8 @@ services:
     cpus: 1
     environment:
       <<: *base-environment
+      # Default to 4 max client connections if not set (work, poll for requeue, completed cleanup, stuck tasks cleanup)
+      DB_MAX_CONNECTIONS: ${DB_MAX_CONNECTIONS:-4}
 
     entrypoint: /app/agent -t aux --monitor-requeue --redis-ttl ${REDIS_TTL:-57600}