Fixes 5058: separate set cancel from set cancel_attempted (#908)

pkg/tasks/queue/pgqueue.go

@@ -98,6 +98,10 @@ const (
 		UPDATE tasks
 		SET status = 'canceled', error = (left($2, 4000)), cancel_attempted = true
 		WHERE id = $1 AND finished_at IS NULL`
+	sqlSetCancelAttempted = `
+		UPDATE tasks
+		SET cancel_attempted = true
+		WHERE id = $1`
 	// sqlUpdatePayload
 	sqlUpdatePayload = `
 		UPDATE tasks
@@ -611,6 +615,13 @@ func (p *PgQueue) Cancel(ctx context.Context, taskId uuid.UUID) error {
 		return fmt.Errorf("error canceling task: %w", err)
 	}
 
+	// this query is separated because we must ensure the flag is set regardless of finished_at being null
+	// but we do not want to set the status to canceled if finished_at is not null
+	_, err = tx.Exec(ctx, sqlSetCancelAttempted, taskId)
+	if err != nil {
+		return fmt.Errorf("error setting cancel_attempted: %w", err)
+	}
+
 	dependents, err := p.taskDependents(context.Background(), tx.Conn(), taskId)
 	if err != nil {
 		return fmt.Errorf("error fetching task dependents: %w", err)

pkg/tasks/queue/pgqueue_test.go

@@ -336,6 +336,7 @@ func (s *QueueSuite) TestCannotRequeueCanceledTasks() {
 func (s *QueueSuite) TestCannotRequeueCanceledFailedTasks() {
 	config.Get().Tasking.RetryWaitUpperBound = 0
 
+	// Test when task fails right after cancellation
 	id, err := s.queue.Enqueue(&testTask)
 	require.NoError(s.T(), err)
 	assert.NotEqual(s.T(), uuid.Nil, id)
@@ -361,6 +362,33 @@ func (s *QueueSuite) TestCannotRequeueCanceledFailedTasks() {
 	assert.Equal(s.T(), config.TaskStatusFailed, info.Status)
 	assert.Equal(s.T(), true, info.CancelAttempted)
 	assert.True(s.T(), info.Queued.Equal(*originalQueueTime))
+
+	// Test when task fails right before cancellation
+	id, err = s.queue.Enqueue(&testTask)
+	require.NoError(s.T(), err)
+	assert.NotEqual(s.T(), uuid.Nil, id)
+
+	info, err = s.queue.Status(id)
+	require.NoError(s.T(), err)
+	originalQueueTime = info.Queued
+
+	_, err = s.queue.Dequeue(context.Background(), []string{testTaskType})
+	require.NoError(s.T(), err)
+
+	err = s.queue.Finish(id, fmt.Errorf("something went wrong"))
+	require.NoError(s.T(), err)
+
+	err = s.queue.Cancel(context.Background(), id)
+	require.NoError(s.T(), err)
+
+	err = s.queue.RequeueFailedTasks([]string{testTaskType})
+	assert.NoError(s.T(), err)
+
+	info, err = s.queue.Status(id)
+	require.NoError(s.T(), err)
+	assert.Equal(s.T(), config.TaskStatusFailed, info.Status)
+	assert.Equal(s.T(), true, info.CancelAttempted)
+	assert.True(s.T(), info.Queued.Equal(*originalQueueTime))
 }
 
 func (s *QueueSuite) TestRequeueFailedTasksExceedRetries() {