orchestrator: fix queueing bug, retrieve all data for events

Author: ndeodhar

orchestrator/src/db/events.rs

@@ -1,4 +1,5 @@
 // Event-related database operations
+use chrono::{DateTime, Utc};
 use sqlx::{postgres::PgRow, Row};
 use uuid::Uuid;
 
@@ -77,26 +78,30 @@ impl Database {
         let mut sequence_ids = Vec::new();
         let mut last_event_type: Option<String> = None;
         let mut last_data: Option<serde_json::Value> = None;
+        let mut last_event_id: Option<Uuid> = None;
+        let mut last_created_at: Option<DateTime<Utc>> = None;
 
         // Publish all events (all for the same topic)
         for (event_type, data, execution_id, attempt_number) in events {
             // Insert event and return sequence_id
             let row: PgRow = sqlx::query(
-        "INSERT INTO events (id, topic, event_type, data, status, execution_id, attempt_number, project_id)
-         VALUES (gen_random_uuid(), $1, $2, $3, $4, $5, $6, $7)
-         RETURNING sequence_id"
-      )
-      .bind(&topic)
-      .bind(event_type.as_deref())
-      .bind(data.clone())
-      .bind("valid")
-      .bind(execution_id.as_ref())
-      .bind(attempt_number)
-      .bind(project_id)
-      .fetch_one(&mut *tx)
-      .await?;
+                "INSERT INTO events (id, topic, event_type, data, status, execution_id, attempt_number, project_id)
+                VALUES (gen_random_uuid(), $1, $2, $3, $4, $5, $6, $7)
+                RETURNING sequence_id, id, created_at"
+              )
+              .bind(&topic)
+              .bind(event_type.as_deref())
+              .bind(data.clone())
+              .bind("valid")
+              .bind(execution_id.as_ref())
+              .bind(attempt_number)
+              .bind(project_id)
+              .fetch_one(&mut *tx)
+              .await?;
 
             let seq_id: i64 = row.get("sequence_id");
+            last_event_id = Some(row.get("id"));
+            last_created_at = Some(row.get("created_at"));
             sequence_ids.push(seq_id);
             // Store last event data for resuming waiting executions
             last_event_type = event_type.clone();
@@ -106,18 +111,18 @@ impl Database {
         // Check for executions waiting on this topic
         // Use FOR UPDATE SKIP LOCKED to allow parallel processing across multiple orchestrator instances
         let waiting_rows = sqlx::query(
-      "SELECT ws.execution_id, COALESCE(ws.root_execution_id, ws.execution_id) as root_execution_id, ws.step_key
-       FROM wait_steps ws
-       INNER JOIN workflow_executions e ON ws.execution_id = e.id
-       WHERE e.status = 'waiting'
-       AND ws.wait_type = 'event'
-       AND ws.wait_topic = $1
-       AND ws.step_key IS NOT NULL
-       FOR UPDATE SKIP LOCKED"
-    )
-    .bind(&topic)
-    .fetch_all(&mut *tx)
-    .await?;
+            "SELECT ws.execution_id, COALESCE(ws.root_execution_id, ws.execution_id) as root_execution_id, ws.step_key
+            FROM wait_steps ws
+            INNER JOIN workflow_executions e ON ws.execution_id = e.id
+            WHERE e.status = 'waiting'
+            AND ws.wait_type = 'event'
+            AND ws.wait_topic = $1
+            AND ws.step_key IS NOT NULL
+            FOR UPDATE SKIP LOCKED"
+        )
+        .bind(&topic)
+        .fetch_all(&mut *tx)
+        .await?;
 
         // Resume each waiting execution (use last event data)
         for row in waiting_rows {
@@ -130,6 +135,8 @@ impl Database {
               "topic": topic,
               "event_type": last_event_type,
               "data": last_data,
+              "id": last_event_id,
+              "created_at": last_created_at,
             });
 
             // Get project_id from execution
@@ -240,7 +247,9 @@ impl Database {
                     'sequence_id', sequence_id,
                     'topic', wait_topic,
                     'event_type', event_type,
-                    'data', data
+                    'data', data,
+                    'id', event_id,
+                    'created_at', created_at
                 ),
                 NULL::jsonb, -- error
                 true, -- success

orchestrator/src/db/workers.rs

@@ -201,40 +201,53 @@ impl Database {
 
         // Get one queued execution with SELECT FOR UPDATE SKIP LOCKED
         // This allows multiple orchestrators to work on different executions in parallel
-        // Uses CTE to efficiently check queue concurrency limits
+        // Uses CTEs to efficiently check queue concurrency limits AND worker availability
+        // Only selects executions that have at least one available worker for their deployment_id
         let execution_row = sqlx::query(
             "WITH running_counts AS (
-        SELECT 
+        SELECT
           queue_name,
           deployment_id,
           COALESCE(concurrency_key, '') as concurrency_key,
           COUNT(*) as running_count
         FROM workflow_executions
         WHERE status IN ('claimed', 'running')
         GROUP BY queue_name, deployment_id, COALESCE(concurrency_key, '')
+      ),
+      available_deployments AS (
+        -- Get deployment_ids that have at least one available worker
+        SELECT DISTINCT current_deployment_id
+        FROM workers
+        WHERE mode = 'push'
+          AND status = 'online'
+          AND current_execution_count < max_concurrent_executions
+          AND push_failure_count < push_failure_threshold
+          AND last_heartbeat > NOW() - INTERVAL '60 seconds'
       )
-      SELECT e.id, e.workflow_id, e.status, e.payload, e.result, e.error, 
-             e.created_at, e.started_at, e.completed_at, 
-             e.deployment_id, e.parent_execution_id, e.root_execution_id, 
-             e.retry_count, e.step_key, e.queue_name, e.concurrency_key, 
-             e.batch_id, e.session_id, e.user_id, e.output_schema_name, 
-             e.otel_traceparent, e.otel_span_id, e.initial_state, e.final_state, 
+      SELECT e.id, e.workflow_id, e.status, e.payload, e.result, e.error,
+             e.created_at, e.started_at, e.completed_at,
+             e.deployment_id, e.parent_execution_id, e.root_execution_id,
+             e.retry_count, e.step_key, e.queue_name, e.concurrency_key,
+             e.batch_id, e.session_id, e.user_id, e.output_schema_name,
+             e.otel_traceparent, e.otel_span_id, e.initial_state, e.final_state,
              e.claimed_at, e.queued_at, e.run_timeout_seconds, q.concurrency_limit
       FROM workflow_executions e
-      INNER JOIN queues q 
-        ON q.name = e.queue_name 
+      INNER JOIN queues q
+        ON q.name = e.queue_name
         AND q.deployment_id = e.deployment_id
-      LEFT JOIN running_counts rc 
-        ON rc.queue_name = e.queue_name 
+      INNER JOIN available_deployments ad
+        ON ad.current_deployment_id = e.deployment_id
+      LEFT JOIN running_counts rc
+        ON rc.queue_name = e.queue_name
         AND rc.deployment_id = e.deployment_id
         AND rc.concurrency_key = COALESCE(e.concurrency_key, '')
       WHERE e.status = 'queued'
         AND (
           q.concurrency_limit IS NULL
           OR COALESCE(rc.running_count, 0) < q.concurrency_limit
         )
-      ORDER BY COALESCE(e.queued_at, e.created_at) ASC 
-      LIMIT 1 
+      ORDER BY COALESCE(e.queued_at, e.created_at) ASC
+      LIMIT 1
       FOR UPDATE OF e SKIP LOCKED",
         )
         .fetch_optional(&mut *tx)

orchestrator/src/main.rs

@@ -50,8 +50,7 @@ async fn main() -> anyhow::Result<()> {
         // In production: use fmt subscriber with stdout and env filter
         tracing_subscriber::fmt()
             .with_env_filter(
-                EnvFilter::from_default_env()
-                    .add_directive(tracing::Level::INFO.into()),
+                EnvFilter::from_default_env().add_directive(tracing::Level::INFO.into()),
             )
             .init();
     }