Make fair task gc conditional on rangeid (#12283)

Add ConditionRangeID to CompleteTasksLessThanRequest and make
CompleteTasksLessThan conditional on unchanged rangeid, just like
CreateTasks.

When fair task reader gc's tasks, it's not protected by write fencing
like adding tasks. So there's a small race window: a node advances its
ack level to C, but it's not persisted yet. Another node takes over and
loads the previous ack level, A. It writes a task above A but below C,
call it B. The first node now deletes below C, deleting B. Note that
this only happens with fair queues because fifo queues always write
above the rangeid.

- [ ] built
- [ ] run locally and tested manually
- [x] covered by existing tests
- [ ] added new unit test(s)
- [ ] added new functional test(s)
- [x] validated in tla+ model (not in repo yet)

Author: dnr

common/persistence/cassandra/matching_task_store_queue.go

@@ -76,6 +76,16 @@ const (
 		`AND task_id = ? ` +
 		`IF range_id = ?`
 
+	templateCheckRangeIDQuery = `UPDATE tasks_v2 SET ` +
+		`range_id = ? ` +
+		`WHERE namespace_id = ? ` +
+		`AND task_queue_name = ? ` +
+		`AND task_queue_type = ? ` +
+		`AND type = ? ` +
+		`AND pass = 0 ` +
+		`AND task_id = ? ` +
+		`IF range_id = ?`
+
 	templateUpdateTaskQueueQueryWithTTLPart1 = `INSERT INTO tasks_v2 ` +
 		`(namespace_id, task_queue_name, task_queue_type, type, pass, task_id) ` +
 		`VALUES (?, ?, ?, ?, 0, ?) USING TTL ?`

common/persistence/cassandra/matching_task_store_v1.go

@@ -187,7 +187,7 @@ func (d *matchingTaskStoreV1) CompleteTasksLessThan(
 	ctx context.Context,
 	request *p.CompleteTasksLessThanRequest,
 ) (int, error) {
-	if request.ExclusiveMaxPass != 0 {
+	if request.ExclusiveMaxPass != 0 || request.ConditionRangeID != 0 {
 		return 0, serviceerror.NewInternal("invalid CompleteTasksLessThan request on queue")
 	}
 

common/persistence/cassandra/matching_task_store_v2.go

@@ -203,19 +203,48 @@ func (d *matchingTaskStoreV2) CompleteTasksLessThan(
 		return 0, serviceerror.NewInternal("invalid CompleteTasksLessThan request on fair queue")
 	}
 
-	query := d.Session.Query(
-		templateCompleteTasksLessThanQuery_v2,
+	rowType := rowTypeTaskInSubqueue(request.Subqueue)
+	args := []any{
 		request.NamespaceID,
 		request.TaskQueueName,
 		request.TaskType,
-		rowTypeTaskInSubqueue(request.Subqueue),
+		rowType,
 		request.ExclusiveMaxPass,
 		request.ExclusiveMaxTaskID,
-	).WithContext(ctx)
-	err := query.Exec()
+	}
+
+	if request.ConditionRangeID == 0 {
+		err := d.Session.Query(templateCompleteTasksLessThanQuery_v2, args...).WithContext(ctx).Exec()
+		if err != nil {
+			return 0, gocql.ConvertError("CompleteTasksLessThan", err)
+		}
+		return p.UnknownNumRowsAffected, nil
+	}
+
+	// Tasks and task queue metadata are in the same partition, so we can make the delete
+	// conditional on the range id with a batch.
+	batch := d.Session.NewBatch(gocql.LoggedBatch).WithContext(ctx)
+	batch.Query(templateCompleteTasksLessThanQuery_v2, args...)
+	batch.Query(switchTasksTable(templateCheckRangeIDQuery, matchingTaskVersion2),
+		request.ConditionRangeID,
+		request.NamespaceID,
+		request.TaskQueueName,
+		request.TaskType,
+		rowTypeTaskQueue,
+		taskQueueTaskID,
+		request.ConditionRangeID,
+	)
+	previous := make(map[string]any)
+	applied, _, err := d.Session.MapExecuteBatchCAS(batch, previous)
 	if err != nil {
 		return 0, gocql.ConvertError("CompleteTasksLessThan", err)
 	}
+	if !applied {
+		return 0, &p.ConditionFailedError{
+			Msg: fmt.Sprintf("Failed to complete tasks. TaskQueue: %v, taskQueueType: %v, rangeID: %v, db rangeID: %v",
+				request.TaskQueueName, request.TaskType, request.ConditionRangeID, previous["range_id"]),
+		}
+	}
 	return p.UnknownNumRowsAffected, nil
 }
 

common/persistence/data_interfaces.go

@@ -646,6 +646,9 @@ type (
 		ExclusiveMaxTaskID int64 // Tasks less than this ID will be completed
 		Subqueue           int
 		Limit              int // Limit on the max number of tasks that can be completed. Required param
+		// If set, only delete if the task queue's range id matches this value, otherwise return
+		// ConditionFailedError. Only supported for fair (v2) task queues.
+		ConditionRangeID int64
 	}
 
 	// CreateNamespaceRequest is used to create the namespace

common/persistence/sql/task_v1.go

@@ -157,7 +157,7 @@ func (m *sqlTaskManagerV1) CompleteTasksLessThan(
 	ctx context.Context,
 	request *persistence.CompleteTasksLessThanRequest,
 ) (int, error) {
-	if request.ExclusiveMaxPass != 0 {
+	if request.ExclusiveMaxPass != 0 || request.ConditionRangeID != 0 {
 		return 0, serviceerror.NewInternal("invalid CompleteTasksLessThan request on queue")
 	}
 

common/persistence/sql/task_v2.go

@@ -175,18 +175,47 @@ func (m *sqlTaskManagerV2) CompleteTasksLessThan(
 		TaskPass: request.ExclusiveMaxPass,
 		TaskID:   request.ExclusiveMaxTaskID,
 	}
-	result, err := m.DB.DeleteFromTasksV2(ctx, sqlplugin.TasksFilterV2{
+	filter := sqlplugin.TasksFilterV2{
 		RangeHash:         tqHash,
 		TaskQueueID:       tqId,
 		ExclusiveMaxLevel: &exclusiveMaxLevel,
 		Limit:             &request.Limit,
-	})
-	if err != nil {
-		return 0, serviceerror.NewUnavailable(err.Error())
 	}
-	nRows, err := result.RowsAffected()
+
+	if request.ConditionRangeID == 0 {
+		result, err := m.DB.DeleteFromTasksV2(ctx, filter)
+		if err != nil {
+			return 0, serviceerror.NewUnavailable(err.Error())
+		}
+		nRows, err := result.RowsAffected()
+		if err != nil {
+			return 0, serviceerror.NewUnavailablef("rowsAffected returned error: %v", err)
+		}
+		return int(nRows), nil
+	}
+
+	var nRows int64
+	err = m.SqlStore.txExecute(ctx, "CompleteTasksLessThan", func(tx sqlplugin.Tx) error {
+		// Lock task queue (subqueue zero holds the range id) and check range id first.
+		queueID, queueHash := taskQueueIdAndHash(nidBytes, request.TaskQueueName, request.TaskType, persistence.SubqueueZero)
+		if err := lockTaskQueue(ctx,
+			tx,
+			queueHash,
+			queueID,
+			request.ConditionRangeID,
+			sqlplugin.MatchingTaskVersion2,
+		); err != nil {
+			return err
+		}
+		result, err := tx.DeleteFromTasksV2(ctx, filter)
+		if err != nil {
+			return err
+		}
+		nRows, err = result.RowsAffected()
+		return err
+	})
 	if err != nil {
-		return 0, serviceerror.NewUnavailablef("rowsAffected returned error: %v", err)
+		return 0, err
 	}
 	return int(nRows), nil
 }

common/persistence/tests/task_queue_fair_task.go

@@ -152,6 +152,56 @@ func (s *TaskQueueFairTaskSuite) TestCreateDelete_Range() {
 	s.Nil(resp.NextPageToken)
 }
 
+func (s *TaskQueueFairTaskSuite) TestCreateDeleteWithConditionRangeID() {
+	rangeID := rand.Int63()
+	taskQueue := s.createTaskQueue(rangeID)
+
+	var tasks []*persistencespb.AllocatedTaskInfo
+	for pass := int64(1); pass <= 5; pass++ {
+		tasks = append(tasks, s.randomTask(pass, pass))
+	}
+	_, err := s.taskManager.CreateTasks(s.ctx, &p.CreateTasksRequest{
+		TaskQueueInfo: &p.PersistedTaskQueueInfo{RangeID: rangeID, Data: taskQueue},
+		Tasks:         tasks,
+	})
+	s.NoError(err)
+
+	completeLessThan := func(pass, conditionRangeID int64) error {
+		_, err := s.taskManager.CompleteTasksLessThan(s.ctx, &p.CompleteTasksLessThanRequest{
+			NamespaceID:        s.namespaceID,
+			TaskQueueName:      s.taskQueueName,
+			TaskType:           s.taskQueueType,
+			ExclusiveMaxPass:   pass,
+			ExclusiveMaxTaskID: 0,
+			Limit:              10,
+			ConditionRangeID:   conditionRangeID,
+		})
+		return err
+	}
+	getTasks := func() []*persistencespb.AllocatedTaskInfo {
+		resp, err := s.taskManager.GetTasks(s.ctx, &p.GetTasksRequest{
+			NamespaceID:        s.namespaceID,
+			TaskQueue:          s.taskQueueName,
+			TaskType:           s.taskQueueType,
+			InclusiveMinPass:   1,
+			InclusiveMinTaskID: 0,
+			ExclusiveMaxTaskID: math.MaxInt64,
+			PageSize:           10,
+		})
+		s.NoError(err)
+		return resp.Tasks
+	}
+
+	// matching range id: tasks are deleted
+	s.NoError(completeLessThan(3, rangeID))
+	protorequire.ProtoSliceEqual(s.T(), tasks[2:], getTasks())
+
+	// wrong range id: condition fails and nothing is deleted
+	err = completeLessThan(5, rangeID-1)
+	s.ErrorAs(err, new(*p.ConditionFailedError))
+	protorequire.ProtoSliceEqual(s.T(), tasks[2:], getTasks())
+}
+
 func (s *TaskQueueFairTaskSuite) createTaskQueue(rangeID int64) *persistencespb.TaskQueueInfo {
 	taskQueueKind := enumspb.TaskQueueKind(rand.Int31n(int32(len(enumspb.TaskQueueKind_name)) + 1))
 	taskQueue := s.randomTaskQueueInfo(taskQueueKind)

service/matching/db.go

@@ -2,6 +2,7 @@ package matching
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"math"
 	"slices"
@@ -658,6 +659,11 @@ func (db *taskQueueDB) CompleteFairTasksLessThan(
 	limit int,
 	subqueue subqueueIndex,
 ) (int, error) {
+	// For fair tasks, we do an LWT to avoid issues with lost ownership, therefore we should
+	// serialize this with other LWTs by holding db.lock to avoid timeouts or conflicts.
+	db.Lock()
+	defer db.Unlock()
+
 	n, err := db.store.CompleteTasksLessThan(ctx, &persistence.CompleteTasksLessThanRequest{
 		NamespaceID:        db.queue.NamespaceId(),
 		TaskQueueName:      db.queue.PersistenceName(),
@@ -666,8 +672,14 @@ func (db *taskQueueDB) CompleteFairTasksLessThan(
 		ExclusiveMaxTaskID: exclusiveMaxLevel.id,
 		Subqueue:           int(subqueue),
 		Limit:              limit,
+		// We might have lost ownership without knowing it. A new owner assigns passes
+		// starting from the ack level it loaded, which may be below our in-memory ack level,
+		// so an unconditional delete could delete the new owner's tasks.
+		ConditionRangeID: db.rangeID,
 	})
-	if err != nil {
+	// ConditionFailedError means we lost ownership: not a store failure, caller will unload.
+	var conditionFailed *persistence.ConditionFailedError
+	if err != nil && !errors.As(err, &conditionFailed) {
 		db.logger.Error("Persistent store operation failure",
 			tag.StoreOperationCompleteTasksLessThan,
 			tag.Error(err),

service/matching/fair_task_reader.go

@@ -653,6 +653,7 @@ func (tr *fairTaskReader) doGCAt(ackLevel fairLevel) (int, error) {
 	n, err := tr.backlogMgr.db.CompleteFairTasksLessThan(ctx, ackLevel.inc(), batchSize, tr.subqueue)
 	if err != nil {
 		tr.logger.Warn("failed to gc tasks", tag.Error(err))
+		tr.backlogMgr.signalIfFatal(err)
 	}
 	return n, err
 }

service/matching/matching_engine_test.go

@@ -1826,9 +1826,6 @@ func (s *matchingEngineSuite) TestForceUnloadTaskQueue() {
 }
 
 func (s *matchingEngineSuite) TestMultipleEnginesActivitiesRangeStealing() {
-	if s.newMatcher {
-		s.T().Skip("test is flaky with new matcher")
-	}
 	runID := uuid.NewString()
 	workflowID := "workflow1"
 	workflowExecution := &commonpb.WorkflowExecution{RunId: runID, WorkflowId: workflowID}
@@ -1923,10 +1920,12 @@ func (s *matchingEngineSuite) TestMultipleEnginesActivitiesRangeStealing() {
 					}),
 			}, nil
 		}).AnyTimes()
+	// If tasks are lost, we'd poll forever. Give up after a deadline and check below.
+	pollDeadline := time.Now().Add(30 * time.Second)
 	for j := 0; j < iterations; j++ {
 		for p := 0; p < engineCount; p++ {
 			engine := engines[p]
-			for i := int64(0); i < taskCount; /* incremented explicitly to skip empty polls */ {
+			for i := int64(0); i < taskCount && time.Now().Before(pollDeadline); /* incremented explicitly to skip empty polls */ {
 				result, err := engine.PollActivityTaskQueue(context.Background(), &matchingservice.PollActivityTaskQueueRequest{
 					NamespaceId: namespaceID,
 					PollRequest: &workflowservice.PollActivityTaskQueueRequest{
@@ -1971,8 +1970,12 @@ func (s *matchingEngineSuite) TestMultipleEnginesActivitiesRangeStealing() {
 		e.Stop()
 	}
 
-	s.EqualValues(0, s.taskManager.getTaskCount(tlID))
 	totalTasks := taskCount * engineCount * iterations
+	s.Len(startedTasks, totalTasks, "some tasks were never dispatched")
+	if !s.newMatcher {
+		// new matcher does gc lazily so some acked tasks may remain
+		s.EqualValues(0, s.taskManager.getTaskCount(tlID))
+	}
 	persisted := s.taskManager.getCreateTaskCount(tlID)
 	// No sync matching as all messages are published first
 	s.EqualValues(totalTasks, persisted)
@@ -1982,9 +1985,6 @@ func (s *matchingEngineSuite) TestMultipleEnginesActivitiesRangeStealing() {
 }
 
 func (s *matchingEngineSuite) TestMultipleEnginesWorkflowTasksRangeStealing() {
-	if s.newMatcher {
-		s.T().Skip("test is flaky with new matcher")
-	}
 	runID := uuid.NewString()
 	workflowID := "workflow1"
 	workflowExecution := &commonpb.WorkflowExecution{RunId: runID, WorkflowId: workflowID}
@@ -2068,10 +2068,12 @@ func (s *matchingEngineSuite) TestMultipleEnginesWorkflowTasksRangeStealing() {
 			}, nil
 		}).AnyTimes()
 
+	// If tasks are lost, we'd poll forever. Give up after a deadline and check below.
+	pollDeadline := time.Now().Add(30 * time.Second)
 	for j := 0; j < iterations; j++ {
 		for p := 0; p < engineCount; p++ {
 			engine := engines[p]
-			for i := int64(0); i < taskCount; /* incremented explicitly to skip empty polls */ {
+			for i := int64(0); i < taskCount && time.Now().Before(pollDeadline); /* incremented explicitly to skip empty polls */ {
 				result, err := engine.PollWorkflowTaskQueue(context.Background(), &matchingservice.PollWorkflowTaskQueueRequest{
 					NamespaceId: namespaceID,
 					PollRequest: &workflowservice.PollWorkflowTaskQueueRequest{
@@ -2115,8 +2117,12 @@ func (s *matchingEngineSuite) TestMultipleEnginesWorkflowTasksRangeStealing() {
 		e.Stop()
 	}
 
-	s.EqualValues(0, s.taskManager.getTaskCount(tlID))
 	totalTasks := taskCount * engineCount * iterations
+	s.Len(startedTasks, totalTasks, "some tasks were never dispatched")
+	if !s.newMatcher {
+		// new matcher does gc lazily so some acked tasks may remain
+		s.EqualValues(0, s.taskManager.getTaskCount(tlID))
+	}
 	persisted := s.taskManager.getCreateTaskCount(tlID)
 	// No sync matching as all messages are published first
 	s.EqualValues(totalTasks, persisted)
@@ -3476,8 +3482,18 @@ func (s *matchingEngineSuite) concurrentPublishAndConsumeValidateBacklogCounter(
 	wg.Wait()
 
 	ptqMgr := s.getPhysicalTaskQueueManagerImplFromKey(ptq)
-	dbTasks := int64(s.taskManager.getTaskCount(ptq))
 	backlogCount := totalApproximateBacklogCount(ptqMgr.backlogMgr)
+
+	// Only count tasks above the ack level: acked tasks may not have been gc'ed yet.
+	db := ptqMgr.backlogMgr.getDB()
+	db.Lock()
+	ackLevel := fairLevel{id: db.subqueues[subqueueZero].AckLevel}
+	if s.fairness {
+		ackLevel = fairLevelFromProto(db.subqueues[subqueueZero].FairAckLevel)
+	}
+	db.Unlock()
+	dbTasks := int64(s.taskManager.getTaskCountAbove(ptq, ackLevel))
+
 	if s.fairness {
 		// Relax this condition for fairBacklogManager: it can sometimes reset backlog count on
 		// read, making it more accurate in theory, but breaking this test's assumptions.
@@ -4920,6 +4936,20 @@ func (m *testTaskManager) minTaskID(dbq *PhysicalTaskQueueKey) (int64, bool) {
 	return key.id, ok
 }
 
+// getTaskCountAbove returns the number of tasks in a task queue above the given level.
+func (m *testTaskManager) getTaskCountAbove(q *PhysicalTaskQueueKey, level fairLevel) int {
+	tlm := m.getQueueDataByKey(q)
+	tlm.Lock()
+	defer tlm.Unlock()
+	count := 0
+	for _, k := range tlm.tasks.Keys() {
+		if level.less(k.(fairLevel)) {
+			count++
+		}
+	}
+	return count
+}
+
 // maxTaskID returns the maximum value of the TaskID present in testTaskManager
 func (m *testTaskManager) maxTaskID(dbq *PhysicalTaskQueueKey) (int64, bool) {
 	tlm := m.getQueueDataByKey(dbq)
@@ -4936,7 +4966,7 @@ func (m *testTaskManager) CompleteTasksLessThan(
 ) (int, error) {
 	if m.fairness && request.ExclusiveMaxPass < 1 {
 		return 0, serviceerror.NewInternal("invalid CompleteTasksLessThan request on fair queue")
-	} else if !m.fairness && request.ExclusiveMaxPass != 0 {
+	} else if !m.fairness && (request.ExclusiveMaxPass != 0 || request.ConditionRangeID != 0) {
 		return 0, serviceerror.NewInternal("invalid CompleteTasksLessThan request on queue")
 	}
 
@@ -4946,6 +4976,12 @@ func (m *testTaskManager) CompleteTasksLessThan(
 	tlm := m.getQueueData(request.TaskQueueName, request.NamespaceID, request.TaskType)
 	tlm.Lock()
 	defer tlm.Unlock()
+	if request.ConditionRangeID != 0 && tlm.rangeID != request.ConditionRangeID {
+		return 0, &persistence.ConditionFailedError{
+			Msg: fmt.Sprintf("CompleteTasksLessThan failed, range id mismatch. rangeID: %v, db rangeID: %v",
+				request.ConditionRangeID, tlm.rangeID),
+		}
+	}
 	keys := tlm.tasks.Keys()
 	for _, key := range keys {
 		level := key.(fairLevel)