Fix Backend.GetWorkflowTask lock bug in PSQL backend

.github/workflows/ci.yaml

@@ -0,0 +1,18 @@
+name: ci
+
+on: [ push ]
+
+jobs:
+  build:
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        go-version: [ '1.22.x' ]
+    steps:
+      - uses: actions/checkout@v4
+      - name: Setup Go ${{ matrix.go-version }}
+        uses: actions/setup-go@v5
+        with:
+          go-version: ${{ matrix.go-version }}
+      - name: Build & Test
+        run: ./build.sh

build.sh

@@ -1,10 +1,13 @@
-#!/bin/sh
+#!/bin/bash
 
-set -o pipefail
+set -eo pipefail
 
 go clean
 go clean -testcache
 go build ./...
 go test ./...
 go test -race ./...
-go test -tags e2e -p 1 ./test/e2e/...
+
+if [[ $ENABLE_E2E_TEST ]]; then
+  go test -tags e2e -p 1 ./test/e2e/...
+fi

pkg/backend/psql/backend.go

@@ -234,7 +234,7 @@ func (b *be) GetWorkflowTask(ctx context.Context) (result *task.WorkflowTask, er
 		return nil, HandleSQLError(err)
 	}
 	currentTimestampUTC := b.getCurrentTimestampLocal()
-	t, err := b.taskRepo.GetAndLockAvailableTask(uowCtx, task.TaskTypeWorkflow, b.lockedBy, b.lockExpirationDuration)
+	t, previouslyLockedBy, err := b.taskRepo.GetAndLockAvailableTask(uowCtx, task.TaskTypeWorkflow, b.lockedBy, b.lockExpirationDuration)
 	if err != nil {
 		if errors.Is(err, gorm.ErrRecordNotFound) {
 			return nil, worker.ErrNoTask
@@ -245,7 +245,7 @@ func (b *be) GetWorkflowTask(ctx context.Context) (result *task.WorkflowTask, er
 	if err != nil {
 		return nil, err
 	}
-	pEvents, err := b.eventRepo.GetAvailableWorkflowEventsAndLock(uowCtx, t.WorkflowID, b.lockedBy)
+	pEvents, err := b.eventRepo.GetAvailableWorkflowEventsAndLock(uowCtx, t.WorkflowID, b.lockedBy, previouslyLockedBy)
 	if err != nil {
 		return nil, err
 	}
@@ -481,7 +481,7 @@ func (b *be) GetActivityTask(ctx context.Context) (result *task.ActivityTask, er
 	if err != nil {
 		return nil, HandleSQLError(err)
 	}
-	t, err := b.taskRepo.GetAndLockAvailableTask(uowCtx, task.TaskTypeActivity, b.lockedBy, b.lockExpirationDuration)
+	t, _, err := b.taskRepo.GetAndLockAvailableTask(uowCtx, task.TaskTypeActivity, b.lockedBy, b.lockExpirationDuration)
 	if err != nil {
 		if errors.Is(err, gorm.ErrRecordNotFound) {
 			return nil, worker.ErrNoTask

pkg/backend/psql/persistent/event.go

@@ -21,7 +21,7 @@ type EventRepository interface {
 	DeleteEventsByWorkflowID(ctx context.Context, workflowID string) (int64, error)
 	DeleteEventsByWorkflowIDAndHeldBy(ctx context.Context, workflowID string, heldBy string) (int64, error)
 	ReleaseEventsByWorkflowIDAndHeldBy(ctx context.Context, workflowID string, heldBy string) (int64, error)
-	GetAvailableWorkflowEventsAndLock(ctx context.Context, workflowID string, heldBy string) ([]*Event, error)
+	GetAvailableWorkflowEventsAndLock(ctx context.Context, workflowID string, heldBy string, previouslyHeldBy *string) ([]*Event, error)
 }
 
 type eventRepository struct {
@@ -61,13 +61,13 @@ func (r *eventRepository) ReleaseEventsByWorkflowIDAndHeldBy(ctx context.Context
 	return result.RowsAffected, result.Error
 }
 
-func (r *eventRepository) GetAvailableWorkflowEventsAndLock(ctx context.Context, workflowID string, heldBy string) ([]*Event, error) {
+func (r *eventRepository) GetAvailableWorkflowEventsAndLock(ctx context.Context, workflowID string, heldBy string, previouslyHeldBy *string) ([]*Event, error) {
 	uow := r.UnitOfWork(ctx)
 	now := time.Now().UnixMilli()
 	result := uow.Tx.Model(&Event{}).Where(
 		"workflow_id = ? AND (held_by IS NULL OR held_by = ?) AND visible_at < ?",
 		workflowID,
-		heldBy,
+		previouslyHeldBy,
 		now,
 	).Clauses().Updates(map[string]interface{}{
 		"held_by": heldBy,

pkg/backend/psql/persistent/task.go

@@ -36,7 +36,7 @@ type TaskRepository interface {
 	ReleaseTask(ctx context.Context, workflowID string, taskID string, taskType task.TaskType, lockedBy string, reason *string, nextScheduleTimestamp *int64) error
 	DeleteTask(ctx context.Context, workflowID string, taskID string, taskType task.TaskType, lockedBy string) error
 	DeleteTaskUnsafe(ctx context.Context, workflowID string, taskID string, taskType task.TaskType) error
-	GetAndLockAvailableTask(ctx context.Context, taskType task.TaskType, lockedBy string, lockExpirationDuration time.Duration) (*Task, error)
+	GetAndLockAvailableTask(ctx context.Context, taskType task.TaskType, lockedBy string, lockExpirationDuration time.Duration) (*Task, *string, error)
 	ResetTaskLastTouchTimestamp(ctx context.Context, workflowID string, taskID string) error
 }
 
@@ -123,7 +123,7 @@ func (r *taskRepository) DeleteTaskUnsafe(ctx context.Context, workflowID string
 	return nil
 }
 
-func (r *taskRepository) GetAndLockAvailableTask(ctx context.Context, taskType task.TaskType, lockedBy string, lockExpirationDuration time.Duration) (*Task, error) {
+func (r *taskRepository) GetAndLockAvailableTask(ctx context.Context, taskType task.TaskType, lockedBy string, lockExpirationDuration time.Duration) (*Task, *string, error) {
 	uow := r.UnitOfWork(ctx)
 	now := time.Now().UnixMilli()
 	t := &Task{}
@@ -136,8 +136,9 @@ func (r *taskRepository) GetAndLockAvailableTask(ctx context.Context, taskType t
 		).
 		Order("last_touch ASC").First(&t)
 	if result.Error != nil {
-		return nil, result.Error
+		return nil, nil, result.Error
 	}
+	previousLockedBy := t.LockedBy
 	result = uow.Tx.
 		Model(&Task{}).
 		Where("workflow_id = ? AND task_id = ?", t.WorkflowID, t.TaskID).
@@ -147,9 +148,9 @@ func (r *taskRepository) GetAndLockAvailableTask(ctx context.Context, taskType t
 			"last_touch": now,
 		}).First(&t)
 	if result.Error != nil {
-		return nil, result.Error
+		return nil, nil, result.Error
 	}
-	return t, nil
+	return t, previousLockedBy, nil
 }
 
 func (r *taskRepository) ResetTaskLastTouchTimestamp(ctx context.Context, workflowID string, taskID string) error {