Introduce Task Processor optimization

tubignat · tubignat · commit c98c02ed3658 · 2025-10-14T15:04:26.000+02:00
diff --git a/common/dynamicconfig/dynamicproperties/constants.go b/common/dynamicconfig/dynamicproperties/constants.go
@@ -2113,6 +2113,7 @@ const (
 	EnableTimerQueueV2
 	EnableTransferQueueV2PendingTaskCountAlert
 	EnableTimerQueueV2PendingTaskCountAlert
+	EnableTimerProcessorInMemoryQueue
 
 	// LastBoolKey must be the last one in this const group
 	LastBoolKey
@@ -4680,6 +4681,12 @@ var BoolKeys = map[BoolKey]DynamicBool{
 		Filters:      []Filter{ShardID},
 		DefaultValue: false,
 	},
+	EnableTimerProcessorInMemoryQueue: {
+		KeyName:      "history.enableTimerProcessorInMemoryQueue",
+		Description:  "EnableTimerProcessorInMemoryQueue is the flag to enable in-memory queue for timer processor",
+		Filters:      []Filter{ShardID},
+		DefaultValue: false,
+	},
 }
 
 var FloatKeys = map[FloatKey]DynamicFloat{
diff --git a/config/development.yaml b/config/development.yaml
@@ -172,4 +172,4 @@ shardDistribution:
       prefix: "store"
   process:
     period: 1s
-    heartbeatTTL: 2s
+    heartbeatTTL: 2s
diff --git a/service/history/common/type.go b/service/history/common/type.go
@@ -36,5 +36,9 @@ type (
 		Activities       map[int64]*persistence.ActivityInfo
 		History          events.PersistedBlobs
 		PersistenceError bool
+
+		// if true, the task will be scheduled in memory for the current execution, otherwise
+		// it will only be scheduled after the next DB scan
+		ScheduleInMemory bool
 	}
 )
diff --git a/service/history/config/config.go b/service/history/config/config.go
@@ -154,6 +154,7 @@ type Config struct {
 	TimerProcessorSplitQueueIntervalJitterCoefficient dynamicproperties.FloatPropertyFn
 	TimerProcessorMaxRedispatchQueueSize              dynamicproperties.IntPropertyFn
 	TimerProcessorMaxTimeShift                        dynamicproperties.DurationPropertyFn
+	EnableTimerProcessorInMemoryQueue                 dynamicproperties.BoolPropertyFnWithShardIDFilter
 	TimerProcessorHistoryArchivalSizeLimit            dynamicproperties.IntPropertyFn
 	TimerProcessorArchivalTimeLimit                   dynamicproperties.DurationPropertyFn
 	DisableTimerFailoverQueue                         dynamicproperties.BoolPropertyFn
@@ -453,6 +454,7 @@ func New(dc *dynamicconfig.Collection, numberOfShards int, maxMessageSize int, i
 		TimerProcessorSplitQueueIntervalJitterCoefficient:    dc.GetFloat64Property(dynamicproperties.TimerProcessorSplitQueueIntervalJitterCoefficient),
 		TimerProcessorMaxRedispatchQueueSize:                 dc.GetIntProperty(dynamicproperties.TimerProcessorMaxRedispatchQueueSize),
 		TimerProcessorMaxTimeShift:                           dc.GetDurationProperty(dynamicproperties.TimerProcessorMaxTimeShift),
+		EnableTimerProcessorInMemoryQueue:                    dc.GetBoolPropertyFilteredByShardID(dynamicproperties.EnableTimerProcessorInMemoryQueue),
 		TimerProcessorHistoryArchivalSizeLimit:               dc.GetIntProperty(dynamicproperties.TimerProcessorHistoryArchivalSizeLimit),
 		TimerProcessorArchivalTimeLimit:                      dc.GetDurationProperty(dynamicproperties.TimerProcessorArchivalTimeLimit),
 		DisableTimerFailoverQueue:                            dc.GetBoolProperty(dynamicproperties.DisableTimerFailoverQueue),
diff --git a/service/history/config/config_test.go b/service/history/config/config_test.go
@@ -278,6 +278,7 @@ func TestNewConfig(t *testing.T) {
 		"QueueMaxVirtualQueueCount":                            {dynamicproperties.QueueMaxVirtualQueueCount, 101},
 		"VirtualSliceForceAppendInterval":                      {dynamicproperties.VirtualSliceForceAppendInterval, time.Second},
 		"ReplicationTaskProcessorLatencyLogThreshold":          {dynamicproperties.ReplicationTaskProcessorLatencyLogThreshold, time.Duration(0)},
+		"EnableTimerProcessorInMemoryQueue":                    {dynamicproperties.EnableTimerProcessorInMemoryQueue, false},
 	}
 	client := dynamicconfig.NewInMemoryClient()
 	for fieldName, expected := range fields {
diff --git a/service/history/execution/context.go b/service/history/execution/context.go
@@ -992,6 +992,7 @@ func notifyTasks(
 		ExecutionInfo:    executionInfo,
 		Tasks:            tasksByCategory[persistence.HistoryTaskCategoryTimer],
 		PersistenceError: persistenceError,
+		ScheduleInMemory: true,
 	}
 	replicationTaskInfo := &hcommon.NotifyTaskInfo{
 		ExecutionInfo:    executionInfo,
diff --git a/service/history/execution/context_test.go b/service/history/execution/context_test.go
@@ -289,6 +289,7 @@ func TestNotifyTasksFromWorkflowSnapshot(t *testing.T) {
 						},
 					},
 					PersistenceError: true,
+					ScheduleInMemory: true,
 				})
 				mockEngine.EXPECT().NotifyNewReplicationTasks(&hcommon.NotifyTaskInfo{
 					ExecutionInfo: &persistence.WorkflowExecutionInfo{
@@ -419,6 +420,7 @@ func TestNotifyTasksFromWorkflowMutation(t *testing.T) {
 						},
 					},
 					PersistenceError: true,
+					ScheduleInMemory: true,
 				})
 				mockEngine.EXPECT().NotifyNewReplicationTasks(&hcommon.NotifyTaskInfo{
 					ExecutionInfo: &persistence.WorkflowExecutionInfo{
diff --git a/service/history/queuev2/queue_base.go b/service/history/queuev2/queue_base.go
@@ -274,6 +274,14 @@ func (q *queueBase) processNewTasks() bool {
 	return true
 }
 
+func (q *queueBase) insertSingleTask(task task.Task) bool {
+	return q.virtualQueueManager.InsertSingleTaskToRootQueue(task)
+}
+
+func (q *queueBase) removeScheduledTasksAfter(t time.Time) {
+	q.virtualQueueManager.RemoveScheduledTasksAfter(t)
+}
+
 func (q *queueBase) updateQueueState(ctx context.Context) {
 	q.metricsScope.IncCounter(metrics.AckLevelUpdateCounter)
 	queueState := &QueueState{
diff --git a/service/history/queuev2/queue_scheduled.go b/service/history/queuev2/queue_scheduled.go
@@ -143,15 +143,42 @@ func (q *scheduledQueue) NotifyNewTask(clusterName string, info *hcommon.NotifyT
 		return
 	}
 
-	nextTime := info.Tasks[0].GetVisibilityTimestamp()
-	for i := 1; i < numTasks; i++ {
-		ts := info.Tasks[i].GetVisibilityTimestamp()
-		if ts.Before(nextTime) {
-			nextTime = ts
+	q.base.logger.Debug(
+		"New timer task notification received",
+		tag.Dynamic("numTasks", numTasks),
+		tag.Dynamic("scheduleInMemory", info.ScheduleInMemory),
+		tag.Dynamic("persistenceError", info.PersistenceError),
+		tag.Dynamic("shardId", q.base.shard.GetShardID()),
+	)
+
+	tasksToBeReadFromDB := make([]persistence.Task, 0)
+
+	if info.ScheduleInMemory && !info.PersistenceError {
+		for _, task := range info.Tasks {
+			ts := task.GetVisibilityTimestamp()
+			q.base.logger.Debug("Submitting task to an in-memory queue", tag.Dynamic("scheduledTime", ts), tag.Dynamic("shardId", q.base.shard.GetShardID()))
+
+			if !q.base.insertSingleTask(q.base.taskInitializer(task)) {
+				tasksToBeReadFromDB = append(tasksToBeReadFromDB, task)
+			}
 		}
+	} else {
+		tasksToBeReadFromDB = info.Tasks
+	}
+
+	var nextReadTime time.Time
+	for _, task := range tasksToBeReadFromDB {
+		ts := task.GetVisibilityTimestamp()
+		if nextReadTime.IsZero() || ts.Before(nextReadTime) {
+			nextReadTime = ts
+		}
+	}
+
+	if !nextReadTime.IsZero() {
+		q.base.removeScheduledTasksAfter(nextReadTime)
+		q.notify(nextReadTime)
 	}
 
-	q.notify(nextTime)
 	q.base.metricsScope.AddCounter(metrics.NewHistoryTaskCounter, int64(numTasks))
 }
 
diff --git a/service/history/queuev2/queue_scheduled_test.go b/service/history/queuev2/queue_scheduled_test.go
@@ -37,6 +37,7 @@ func TestScheduledQueue_LifeCycle(t *testing.T) {
 	mockExecutionManager := persistence.NewMockExecutionManager(ctrl)
 
 	// Setup mock expectations
+	mockShard.EXPECT().GetShardID().Return(1).AnyTimes()
 	mockShard.EXPECT().GetClusterMetadata().Return(cluster.TestActiveClusterMetadata).AnyTimes()
 	mockShard.EXPECT().GetTimeSource().Return(mockTimeSource).AnyTimes()
 	mockShard.EXPECT().GetQueueState(persistence.HistoryTaskCategoryTimer).Return(&types.QueueState{
diff --git a/service/history/queuev2/virtual_queue.go b/service/history/queuev2/virtual_queue.go
@@ -67,6 +67,10 @@ type (
 		SplitSlices(func(VirtualSlice) (remaining []VirtualSlice, split bool))
 		// Pause pauses the virtual queue for a while
 		Pause(time.Duration)
+		// InsertSingleTask inserts a single task to the virtual queue. Return false if the task's timestamp is out of range of the current queue slice..
+		InsertSingleTask(task task.Task) bool
+		// RemoveScheduledTasksAfter removes the scheduled tasks after the given time
+		RemoveScheduledTasksAfter(time.Time)
 	}
 
 	VirtualQueueOptions struct {
@@ -394,35 +398,14 @@ func (q *virtualQueueImpl) loadAndSubmitTasks() {
 
 	now := q.timeSource.Now()
 	for _, task := range tasks {
-		if persistence.IsTaskCorrupted(task) {
-			q.logger.Error("Virtual queue encountered a corrupted task", tag.Dynamic("task", task))
-			q.metricsScope.IncCounter(metrics.CorruptedHistoryTaskCounter)
-			task.Ack()
-			continue
-		}
-
-		scheduledTime := task.GetTaskKey().GetScheduledTime()
-		// if the scheduled time is in the future, we need to reschedule the task
-		if now.Before(scheduledTime) {
-			q.rescheduler.RescheduleTask(task, scheduledTime)
-			continue
-		}
-		// shard level metrics for the duration between a task being written to a queue and being fetched from it
-		q.metricsScope.RecordHistogramDuration(metrics.TaskEnqueueToFetchLatency, now.Sub(task.GetVisibilityTimestamp()))
-		task.SetInitialSubmitTime(now)
-		submitted, err := q.processor.TrySubmit(task)
-		if err != nil {
+		if err := q.submitTask(now, task); err != nil {
 			select {
 			case <-q.ctx.Done():
 				return
 			default:
 				q.logger.Error("Virtual queue failed to submit task", tag.Error(err))
 			}
 		}
-		if !submitted {
-			q.metricsScope.IncCounter(metrics.ProcessingQueueThrottledCounter)
-			q.rescheduler.RescheduleTask(task, q.timeSource.Now().Add(taskSchedulerThrottleBackoffInterval))
-		}
 	}
 
 	if sliceToRead.HasMoreTasks() {
@@ -436,6 +419,86 @@ func (q *virtualQueueImpl) loadAndSubmitTasks() {
 	}
 }
 
+func (q *virtualQueueImpl) InsertSingleTask(task task.Task) bool {
+	q.Lock()
+	defer q.Unlock()
+
+	taskKey := task.GetTaskKey()
+	var slice VirtualSlice
+
+	for e := q.virtualSlices.Front(); e != nil; e = e.Next() {
+		s := e.Value.(VirtualSlice)
+		r := s.GetState().Range
+		if taskKey.Compare(r.InclusiveMinTaskKey) >= 0 && taskKey.Compare(r.ExclusiveMaxTaskKey) < 0 {
+			slice = s
+			break
+		}
+	}
+
+	if slice == nil {
+		// the new task is outside of the current range, it will be read from the DB on the next poll
+		return false
+	}
+
+	slice.InsertTask(task)
+	q.monitor.SetSlicePendingTaskCount(slice, slice.GetPendingTaskCount())
+
+	now := q.timeSource.Now()
+	if err := q.submitTask(now, task); err != nil {
+		q.logger.Error("Virtual queue failed to submit task", tag.Error(err))
+		return false
+	}
+
+	return true
+}
+
+func (q *virtualQueueImpl) RemoveScheduledTasksAfter(t time.Time) {
+	q.Lock()
+	defer q.Unlock()
+
+	for e := q.virtualSlices.Front(); e != nil; e = e.Next() {
+		s := e.Value.(VirtualSlice)
+		r := s.GetState().Range
+
+		if t.Before(r.InclusiveMinTaskKey.GetScheduledTime()) {
+			continue
+		}
+
+		// TODO: remove scheduled tasks from virtual slices
+		q.monitor.SetSlicePendingTaskCount(s, s.GetPendingTaskCount())
+	}
+}
+
+func (q *virtualQueueImpl) submitTask(now time.Time, task task.Task) error {
+	if persistence.IsTaskCorrupted(task) {
+		q.logger.Error("Virtual queue encountered a corrupted task", tag.Dynamic("task", task))
+		q.metricsScope.IncCounter(metrics.CorruptedHistoryTaskCounter)
+		task.Ack()
+		return nil
+	}
+
+	scheduledTime := task.GetTaskKey().GetScheduledTime()
+	// if the scheduled time is in the future, we need to reschedule the task
+	if now.Before(scheduledTime) {
+		q.rescheduler.RescheduleTask(task, scheduledTime)
+		return nil
+	}
+	// shard level metrics for the duration between a task being written to a queue and being fetched from it
+	q.metricsScope.RecordHistogramDuration(metrics.TaskEnqueueToFetchLatency, now.Sub(task.GetVisibilityTimestamp()))
+	task.SetInitialSubmitTime(now)
+	submitted, err := q.processor.TrySubmit(task)
+	if err != nil {
+		return err
+	}
+
+	if !submitted {
+		q.metricsScope.IncCounter(metrics.ProcessingQueueThrottledCounter)
+		q.rescheduler.RescheduleTask(task, q.timeSource.Now().Add(taskSchedulerThrottleBackoffInterval))
+	}
+
+	return nil
+}
+
 func (q *virtualQueueImpl) resetNextReadSliceLocked() {
 	q.sliceToRead = nil
 	for element := q.virtualSlices.Front(); element != nil; element = element.Next() {
diff --git a/service/history/queuev2/virtual_queue_manager.go b/service/history/queuev2/virtual_queue_manager.go
@@ -56,6 +56,9 @@ type (
 		// Add a new virtual slice to the root queue. This is used when new tasks are generated and max read level is updated.
 		// By default, all new tasks belong to the root queue, so we need to add a new virtual slice to the root queue.
 		AddNewVirtualSliceToRootQueue(VirtualSlice)
+		// Insert a single task to the current slice. Return false if the task's timestamp is out of range of the current slice.
+		InsertSingleTaskToRootQueue(task.Task) bool
+		RemoveScheduledTasksAfter(time.Time)
 	}
 
 	virtualQueueManagerImpl struct {
@@ -218,6 +221,24 @@ func (m *virtualQueueManagerImpl) AddNewVirtualSliceToRootQueue(s VirtualSlice)
 	m.virtualQueues[rootQueueID].Start()
 }
 
+func (m *virtualQueueManagerImpl) InsertSingleTaskToRootQueue(t task.Task) bool {
+	m.Lock()
+	defer m.Unlock()
+	if vq, ok := m.virtualQueues[rootQueueID]; ok {
+		return vq.InsertSingleTask(t)
+	}
+
+	// if a root queue is not created yet, no need to schedule an incoming task, it will be read from the slice
+	return false
+}
+
+func (m *virtualQueueManagerImpl) RemoveScheduledTasksAfter(t time.Time) {
+	m.Lock()
+	defer m.Unlock()
+	for _, vq := range m.virtualQueues {
+		vq.RemoveScheduledTasksAfter(t)
+	}
+}
 func (m *virtualQueueManagerImpl) appendOrMergeSlice(vq VirtualQueue, s VirtualSlice) {
 	now := m.timeSource.Now()
 	newVirtualSliceState := s.GetState()
diff --git a/service/history/queuev2/virtual_queue_manager_mock.go b/service/history/queuev2/virtual_queue_manager_mock.go
diff --git a/service/history/queuev2/virtual_queue_mock.go b/service/history/queuev2/virtual_queue_mock.go
diff --git a/service/history/queuev2/virtual_slice.go b/service/history/queuev2/virtual_slice.go
diff --git a/service/history/queuev2/virtual_slice_mock.go b/service/history/queuev2/virtual_slice_mock.go
diff --git a/service/history/shard/context.go b/service/history/shard/context.go
diff --git a/service/sharddistributor/store/etcd/go.mod b/service/sharddistributor/store/etcd/go.mod

Original file line number	Diff line number	Diff line change
`@@ -36,5 +36,9 @@ type (`
`36`	`36`	`Activities map[int64]*persistence.ActivityInfo`
`37`	`37`	`History events.PersistedBlobs`
`38`	`38`	`PersistenceError bool`
	`39`	`+`
	`40`	`+ // if true, the task will be scheduled in memory for the current execution, otherwise`
	`41`	`+ // it will only be scheduled after the next DB scan`
	`42`	`+ ScheduleInMemory bool`
`39`	`43`	`}`
`40`	`44`	`)`
Original file line number	Diff line number	Diff line change
`@@ -278,6 +278,7 @@ func TestNewConfig(t *testing.T) {`
`278`	`278`	`"QueueMaxVirtualQueueCount": {dynamicproperties.QueueMaxVirtualQueueCount, 101},`
`279`	`279`	`"VirtualSliceForceAppendInterval": {dynamicproperties.VirtualSliceForceAppendInterval, time.Second},`
`280`	`280`	`"ReplicationTaskProcessorLatencyLogThreshold": {dynamicproperties.ReplicationTaskProcessorLatencyLogThreshold, time.Duration(0)},`
	`281`	`+ "EnableTimerProcessorInMemoryQueue": {dynamicproperties.EnableTimerProcessorInMemoryQueue, false},`
`281`	`282`	`}`
`282`	`283`	`client := dynamicconfig.NewInMemoryClient()`
`283`	`284`	`for fieldName, expected := range fields {`
Original file line number	Diff line number	Diff line change
`@@ -992,6 +992,7 @@ func notifyTasks(`
`992`	`992`	`ExecutionInfo: executionInfo,`
`993`	`993`	`Tasks: tasksByCategory[persistence.HistoryTaskCategoryTimer],`
`994`	`994`	`PersistenceError: persistenceError,`
	`995`	`+ ScheduleInMemory: true,`
`995`	`996`	`}`
`996`	`997`	`replicationTaskInfo := &hcommon.NotifyTaskInfo{`
`997`	`998`	`ExecutionInfo: executionInfo,`