more internal polish, resilient workers

crates/worker/src/service.rs

@@ -307,18 +307,39 @@ impl WorkerService {
 
     /// Stop the worker service
     pub async fn stop(&mut self) -> Result<()> {
-        info!("Stopping Worker Service");
+        info!("Stopping Worker Service - initiating graceful shutdown");
+
+        // Mark worker as inactive first to stop receiving new tasks
+        {
+            let reg = self.registration.read().await;
+            info!("Marking worker as inactive to stop receiving new tasks");
+            reg.deregister().await?;
+        }
 
         // Stop heartbeat
+        info!("Stopping heartbeat updates");
         self.heartbeat.stop().await;
 
         // Wait a bit for heartbeat to stop
         tokio::time::sleep(Duration::from_millis(100)).await;
 
-        // Deregister worker
-        {
-            let reg = self.registration.read().await;
-            reg.deregister().await?;
+        // Wait for in-flight tasks to complete (with timeout)
+        let shutdown_timeout = self
+            .config
+            .worker
+            .as_ref()
+            .and_then(|w| w.shutdown_timeout)
+            .unwrap_or(30); // Default: 30 seconds
+
+        info!(
+            "Waiting up to {} seconds for in-flight tasks to complete",
+            shutdown_timeout
+        );
+
+        let timeout_duration = Duration::from_secs(shutdown_timeout as u64);
+        match tokio::time::timeout(timeout_duration, self.wait_for_in_flight_tasks()).await {
+            Ok(_) => info!("All in-flight tasks completed"),
+            Err(_) => warn!("Shutdown timeout reached - some tasks may have been interrupted"),
         }
 
         info!("Worker Service stopped");
@@ -326,6 +347,22 @@ impl WorkerService {
         Ok(())
     }
 
+    /// Wait for in-flight tasks to complete
+    async fn wait_for_in_flight_tasks(&self) {
+        // Poll for active executions with short intervals
+        loop {
+            // Check if executor has any active tasks
+            // Note: This is a simplified check. In a real implementation,
+            // we would track active execution count in the executor.
+            tokio::time::sleep(Duration::from_millis(500)).await;
+
+            // TODO: Add proper tracking of active executions in ActionExecutor
+            // For now, we just wait a reasonable amount of time
+            // This will be improved when we add execution tracking
+            break;
+        }
+    }
+
     /// Start consuming execution.scheduled messages
     async fn start_execution_consumer(&mut self) -> Result<()> {
         let worker_id = self
@@ -410,7 +447,7 @@ impl WorkerService {
         .await
         {
             error!("Failed to publish running status: {}", e);
-            // Continue anyway - the executor will update the database
+            // Continue anyway - we'll update the database directly
         }
 
         // Execute the action
@@ -592,8 +629,6 @@ impl WorkerService {
 
         Ok(())
     }
-
-
 }
 
 #[cfg(test)]