hatchet-dev · abelanger5 · Aug 13, 2024 · Aug 12, 2024 · Aug 13, 2024 · Aug 13, 2024
diff --git a/api-contracts/openapi/components/schemas/workflow_run.yaml b/api-contracts/openapi/components/schemas/workflow_run.yaml
@@ -133,6 +133,7 @@ StepRunStatus:
     - SUCCEEDED
     - FAILED
     - CANCELLED
+    - CANCELLING
 
 JobRunStatus:
   type: string

diff --git a/api/v1/server/oas/gen/openapi.gen.go b/api/v1/server/oas/gen/openapi.gen.go
diff --git a/frontend/app/src/lib/api/generated/data-contracts.ts b/frontend/app/src/lib/api/generated/data-contracts.ts
@@ -714,6 +714,7 @@ export enum StepRunStatus {
   SUCCEEDED = 'SUCCEEDED',
   FAILED = 'FAILED',
   CANCELLED = 'CANCELLED',
+  CANCELLING = 'CANCELLING',
 }
 
 export interface JobRun {

diff --git a/frontend/app/src/pages/main/workers/$worker/components/slot-grid.tsx b/frontend/app/src/pages/main/workers/$worker/components/slot-grid.tsx
@@ -15,6 +15,7 @@ const statusColors = {
   RUNNING: 'bg-green-500',
   SUCCEEDED: 'bg-green-700',
   FAILED: 'bg-red-500',
+  CANCELLING: 'bg-yellow-500',
   CANCELLED: 'bg-gray-500',
   UNDEFINED: 'bg-gray-300', // Default color for undefined status
 };

diff --git a/frontend/app/src/pages/main/workflow-runs/components/run-statuses.tsx b/frontend/app/src/pages/main/workflow-runs/components/run-statuses.tsx
@@ -28,6 +28,10 @@ const RUN_STATUS_VARIANTS: Record<RunStatusType, RunStatusVariant> = {
     text: 'Cancelled',
     variant: 'failed',
   },
+  CANCELLING: {
+    text: 'Cancelling',
+    variant: 'inProgress',
+  },
   RUNNING: {
     text: 'Running',
     variant: 'inProgress',

diff --git a/pkg/client/rest/gen.go b/pkg/client/rest/gen.go
diff --git a/pkg/repository/prisma/dbsqlc/step_runs.sql b/pkg/repository/prisma/dbsqlc/step_runs.sql
@@ -556,7 +556,30 @@ WHERE
 FOR UPDATE SKIP LOCKED;
 
 -- name: BulkAssignStepRunsToWorkers :many
-WITH updated_step_runs AS (
+WITH already_assigned_step_runs AS (
+    SELECT
+        input."id",
+        wss."id" AS "slotId"
+    FROM
+        (
+            SELECT
+                unnest(@stepRunIds::uuid[]) AS "id"
+        ) AS input
+    JOIN
+        "WorkerSemaphoreSlot" wss ON input."id" = wss."stepRunId"
+), already_assigned_slots AS (
+    SELECT
+        wss."id"
+    FROM
+        (
+            SELECT
+                unnest(@slotIds::uuid[]) AS "id"
+        ) AS input
+    JOIN
+        "WorkerSemaphoreSlot" wss ON input."id" = wss."id"
+    WHERE
+        wss."stepRunId" IS NOT NULL
+), updated_step_runs AS (
     UPDATE
         "StepRun" sr
     SET
@@ -567,24 +590,33 @@ WITH updated_step_runs AS (
         "timeoutAt" = CURRENT_TIMESTAMP + convert_duration_to_interval(input."stepTimeout")
     FROM (
         SELECT
-            unnest(@stepRunIds::uuid[]) AS "id",
-            unnest(@stepRunTimeouts::text[]) AS "stepTimeout",
-            unnest(@workerIds::uuid[]) AS "workerId"
+            "id",
+            "stepTimeout",
+            "workerId",
+            "slotId"
+        FROM
+            (
+                SELECT
+                    unnest(@stepRunIds::uuid[]) AS "id",
+                    unnest(@stepRunTimeouts::text[]) AS "stepTimeout",
+                    unnest(@workerIds::uuid[]) AS "workerId",
+                    unnest(@slotIds::uuid[]) AS "slotId"
+            ) AS subquery
+        WHERE
+            "id" NOT IN (SELECT "id" FROM already_assigned_step_runs)
+            AND "slotId" NOT IN (SELECT "id" FROM already_assigned_slots)
     ) AS input
     WHERE
         sr."id" = input."id"
+    RETURNING input."id", input."slotId"
 )
 UPDATE
     "WorkerSemaphoreSlot" wss
 SET
-    "stepRunId" = input."stepRunId"
-FROM (
-    SELECT
-        unnest(@slotIds::uuid[]) AS "id",
-        unnest(@stepRunIds::uuid[]) AS "stepRunId"
-) AS input
+    "stepRunId" = updated_step_runs."id"
+FROM updated_step_runs
 WHERE
-    wss."id" = input."id"
+    wss."id" = updated_step_runs."slotId"
 RETURNING wss."id";
 
 -- name: BulkMarkStepRunsAsCancelling :many

diff --git a/pkg/repository/prisma/dbsqlc/step_runs.sql.go b/pkg/repository/prisma/dbsqlc/step_runs.sql.go
diff --git a/pkg/repository/prisma/step_run.go b/pkg/repository/prisma/step_run.go
@@ -965,6 +965,10 @@ func (s *stepRunEngineRepository) QueueStepRuns(ctx context.Context, tenantId st
 		return emptyRes, nil
 	}
 
+	var duplicates []*scheduling.QueueItemWithOrder
+
+	queueItems, duplicates = removeDuplicates(queueItems)
+
 	// sort the queue items by Order from least to greatest, then by queue id
 	sort.Slice(queueItems, func(i, j int) bool {
 		// sort by priority, then by order, then by id
@@ -1163,7 +1167,17 @@ func (s *stepRunEngineRepository) QueueStepRuns(ctx context.Context, tenantId st
 		return emptyRes, fmt.Errorf("could not bulk assign step runs to workers: %w", err)
 	}
 
-	err = s.queries.BulkQueueItems(ctx, tx, plan.QueuedItems)
+	popItems := plan.QueuedItems
+
+	// we'd like to remove duplicates from the queue items as well
+	for _, item := range duplicates {
+		// print a warning for duplicates
+		s.l.Warn().Msgf("duplicate queue item: %d for step run %s", item.QueueItem.ID, sqlchelpers.UUIDToStr(item.QueueItem.StepRunId))
+
+		popItems = append(popItems, item.QueueItem.ID)
+	}
+
+	err = s.queries.BulkQueueItems(ctx, tx, popItems)
 
 	if err != nil {
 		return emptyRes, fmt.Errorf("could not bulk queue items: %w", err)
@@ -2087,3 +2101,23 @@ func (s *stepRunEngineRepository) ClearStepRunPayloadData(ctx context.Context, t
 func getCacheName(tenantId, queue string) string {
 	return fmt.Sprintf("%s:%s", tenantId, queue)
 }
+
+// removes duplicates from a slice of queue items by step run id
+func removeDuplicates(qis []*scheduling.QueueItemWithOrder) ([]*scheduling.QueueItemWithOrder, []*scheduling.QueueItemWithOrder) {
+	encountered := map[string]bool{}
+	result := []*scheduling.QueueItemWithOrder{}
+	duplicates := []*scheduling.QueueItemWithOrder{}
+
+	for _, v := range qis {
+		stepRunId := sqlchelpers.UUIDToStr(v.StepRunId)
+		if encountered[stepRunId] {
+			duplicates = append(duplicates, v)
+			continue
+		}
+
+		encountered[stepRunId] = true
+		result = append(result, v)
+	}
+
+	return result, duplicates
+}