feat(engine): improve stability of workers with terminations and disconnects

frontend/app/src/components/ui/json-form.tsx

@@ -15,7 +15,7 @@ import { DynamicSizeInputTemplate } from './form-inputs/dynamic-size-input-templ
 import { createContext, useRef } from 'react';
 
 type JSONPrimitive = string | number | boolean | null;
-type JSONType = { [key: string]: JSONType | JSONPrimitive };
+type JSONType = { [key: string]: JSONType | JSONPrimitive | Array<JSONType> };
 
 export const DEFAULT_COLLAPSED = ['advanced', 'user data'];
 

internal/repository/prisma/dbsqlc/step_runs.sql

@@ -167,6 +167,37 @@ SELECT
 FROM step_run_data
 RETURNING *;
 
+-- name: ListStepRunsToReassign :many
+SELECT
+    sr.*
+FROM
+    "StepRun" sr
+LEFT JOIN
+    "Worker" w ON sr."workerId" = w."id"
+JOIN
+    "Step" s ON sr."stepId" = s."id"
+WHERE
+    sr."tenantId" = @tenantId::uuid
+    AND ((
+        sr."status" = 'RUNNING'
+        AND w."lastHeartbeatAt" < NOW() - INTERVAL '60 seconds'
+        AND s."retries" > sr."retryCount"
+    ) OR (
+        sr."status" = 'ASSIGNED'
+        AND w."lastHeartbeatAt" < NOW() - INTERVAL '5 seconds'
+    ))
+    -- Step run cannot have a failed parent
+    AND NOT EXISTS (
+        SELECT 1
+        FROM "_StepRunOrder" AS order_table
+        JOIN "StepRun" AS prev_sr ON order_table."A" = prev_sr."id"
+        WHERE 
+            order_table."B" = sr."id"
+            AND prev_sr."status" != 'SUCCEEDED'
+    )
+ORDER BY
+    sr."createdAt" ASC;
+
 -- name: ListStepRunsToRequeue :many
 SELECT
     sr.*
@@ -177,22 +208,14 @@ LEFT JOIN
 WHERE
     sr."tenantId" = @tenantId::uuid
     AND sr."requeueAfter" < NOW()
-    AND (
-        (
-            sr."workerId" IS NULL
-            AND (sr."status" = 'PENDING' OR sr."status" = 'PENDING_ASSIGNMENT')
-            AND NOT EXISTS (
-                SELECT 1
-                FROM "_StepRunOrder" AS order_table
-                JOIN "StepRun" AS prev_sr ON order_table."A" = prev_sr."id"
-                WHERE 
-                    order_table."B" = sr."id"
-                    AND prev_sr."status" != 'SUCCEEDED'
-            )
-        ) OR (
-            sr."status" = 'ASSIGNED'
-            AND w."lastHeartbeatAt" < NOW() - INTERVAL '5 seconds'
-        )
+    AND (sr."status" = 'PENDING' OR sr."status" = 'PENDING_ASSIGNMENT')
+    AND NOT EXISTS (
+        SELECT 1
+        FROM "_StepRunOrder" AS order_table
+        JOIN "StepRun" AS prev_sr ON order_table."A" = prev_sr."id"
+        WHERE 
+            order_table."B" = sr."id"
+            AND prev_sr."status" != 'SUCCEEDED'
     )
 ORDER BY
     sr."createdAt" ASC;

internal/repository/prisma/step_run.go

@@ -99,6 +99,33 @@ func (s *stepRunRepository) ListStepRunsToRequeue(tenantId string) ([]*dbsqlc.St
 	return stepRuns, nil
 }
 
+func (s *stepRunRepository) ListStepRunsToReassign(tenantId string) ([]*dbsqlc.StepRun, error) {
+	pgTenantId := sqlchelpers.UUIDFromStr(tenantId)
+
+	tx, err := s.pool.Begin(context.Background())
+
+	if err != nil {
+		return nil, err
+	}
+
+	defer deferRollback(context.Background(), s.l, tx.Rollback)
+
+	// get the step run and make sure it's still in pending
+	stepRuns, err := s.queries.ListStepRunsToReassign(context.Background(), tx, pgTenantId)
+
+	if err != nil {
+		return nil, err
+	}
+
+	err = tx.Commit(context.Background())
+
+	if err != nil {
+		return nil, err
+	}
+
+	return stepRuns, nil
+}
+
 func (s *stepRunRepository) ListStepRuns(tenantId string, opts *repository.ListStepRunsOpts) ([]db.StepRunModel, error) {
 	if err := s.v.Validate(opts); err != nil {
 		return nil, err
@@ -108,27 +135,6 @@ func (s *stepRunRepository) ListStepRuns(tenantId string, opts *repository.ListS
 		db.StepRun.TenantID.Equals(tenantId),
 	}
 
-	if opts.Requeuable != nil {
-		// job runs are requeuable if they are past their requeue after time, don't have a worker assigned, have a pending status,
-		// and their previous step is completed
-		params = append(
-			params,
-			db.StepRun.RequeueAfter.Before(time.Now().UTC()),
-			db.StepRun.WorkerID.IsNull(),
-			db.StepRun.Status.Equals(db.StepRunStatusPendingAssignment),
-			// db.StepRun.Or(
-			// 	db.StepRun.Prev
-			// 	db.StepRun.Step.Where(
-			// 		db.Step.Prev.Where(
-			// 			db.Step.StepRuns.Some(
-			// 				db.StepRun.Status.Equals(db.StepRunStatusSUCCEEDED),
-			// 			),
-			// 		),
-			// 	),
-			// ),
-		)
-	}
-
 	if opts.Status != nil {
 		params = append(params, db.StepRun.Status.Equals(*opts.Status))
 	}

internal/repository/step_run.go

@@ -17,8 +17,6 @@ type ListAllStepRunsOpts struct {
 }
 
 type ListStepRunsOpts struct {
-	Requeuable *bool
-
 	JobRunId *string
 
 	WorkflowRunId *string
@@ -83,6 +81,9 @@ type StepRunRepository interface {
 	// ListStepRunsToRequeue returns a list of step runs which are in a requeueable state.
 	ListStepRunsToRequeue(tenantId string) ([]*dbsqlc.StepRun, error)
 
+	// ListStepRunsToReassign returns a list of step runs which are in a reassignable state.
+	ListStepRunsToReassign(tenantId string) ([]*dbsqlc.StepRun, error)
+
 	UpdateStepRun(tenantId, stepRunId string, opts *UpdateStepRunOpts) (*db.StepRunModel, *StepRunUpdateInfo, error)
 
 	// UpdateStepRunOverridesData updates the overrides data field in the input for a step run. This returns the input