Add lower priority for replication tasks (#3870)

* add lower priority for replication

* test

* linting

* RequestPriorityFn

Author: pdoerner

common/headers/caller_info.go

@@ -31,17 +31,22 @@ import (
 )
 
 const (
-	CallerTypeAPI        = "api"
-	CallerTypeBackground = "background"
+	CallerTypeAPI         = "api"
+	CallerTypeBackground  = "background"
+	CallerTypePreemptable = "preemptable"
 
 	CallerNameSystem = "system"
 )
 
 var (
-	SystemBackgroundCallerInfo CallerInfo = CallerInfo{
+	SystemBackgroundCallerInfo = CallerInfo{
 		CallerName: CallerNameSystem,
 		CallerType: CallerTypeBackground,
 	}
+	SystemPreemptableCallerInfo = CallerInfo{
+		CallerName: CallerNameSystem,
+		CallerType: CallerTypePreemptable,
+	}
 )
 
 type (

common/headers/caller_info_test.go

@@ -83,6 +83,10 @@ func (s *callerInfoSuite) TestSetCallerType() {
 	ctx = SetCallerType(ctx, CallerTypeAPI)
 	info = GetCallerInfo(ctx)
 	s.Equal(CallerTypeAPI, info.CallerType)
+
+	ctx = SetCallerType(ctx, CallerTypePreemptable)
+	info = GetCallerInfo(ctx)
+	s.Equal(CallerTypePreemptable, info.CallerType)
 }
 
 func (s *callerInfoSuite) TestSetCallOrigin() {

common/persistence/client/quotas.go

@@ -33,8 +33,9 @@ import (
 
 var (
 	CallerTypeDefaultPriority = map[string]int{
-		headers.CallerTypeAPI:        1,
-		headers.CallerTypeBackground: 3,
+		headers.CallerTypeAPI:         1,
+		headers.CallerTypeBackground:  3,
+		headers.CallerTypePreemptable: 4,
 	}
 
 	APITypeCallOriginPriorityOverride = map[string]int{
@@ -69,7 +70,7 @@ var (
 		p.ConstructHistoryTaskAPI("GetHistoryTasks", tasks.CategoryVisibility): 2,
 	}
 
-	RequestPrioritiesOrdered = []int{0, 1, 2, 3}
+	RequestPrioritiesOrdered = []int{0, 1, 2, 3, 4}
 )
 
 func NewPriorityRateLimiter(
@@ -150,6 +151,8 @@ func RequestPriorityFn(req quotas.Request) int {
 			return priority
 		}
 		return CallerTypeDefaultPriority[req.CallerType]
+	case headers.CallerTypePreemptable:
+		return CallerTypeDefaultPriority[req.CallerType]
 	default:
 		// default requests to high priority to be consistent with existing behavior
 		return RequestPrioritiesOrdered[0]

service/history/replication/task_fetcher.go

@@ -420,7 +420,7 @@ func (f *replicationTaskFetcherWorker) getMessages() error {
 
 	ctx, cancel := rpc.NewContextWithTimeoutAndVersionHeaders(fetchTaskRequestTimeout)
 	defer cancel()
-	ctx = headers.SetCallerInfo(ctx, headers.SystemBackgroundCallerInfo)
+	ctx = headers.SetCallerInfo(ctx, headers.SystemPreemptableCallerInfo)
 
 	request := &adminservice.GetReplicationMessagesRequest{
 		Tokens:      tokens,

service/history/replication/task_processor.go

@@ -274,7 +274,7 @@ func (p *taskProcessorImpl) applyReplicationTask(
 ) error {
 	ctx := headers.SetCallerInfo(
 		context.Background(),
-		headers.SystemBackgroundCallerInfo,
+		headers.SystemPreemptableCallerInfo,
 	)
 
 	err := p.handleReplicationTask(ctx, replicationTask)
@@ -314,7 +314,7 @@ func (p *taskProcessorImpl) handleSyncShardStatus(
 		metrics.OperationTag(metrics.HistorySyncShardStatusScope))
 	ctx, cancel := context.WithTimeout(context.Background(), replicationTimeout)
 	defer cancel()
-	ctx = headers.SetCallerInfo(ctx, headers.SystemBackgroundCallerInfo)
+	ctx = headers.SetCallerInfo(ctx, headers.SystemPreemptableCallerInfo)
 
 	return p.historyEngine.SyncShardStatus(ctx, &historyservice.SyncShardStatusRequest{
 		SourceCluster: p.sourceCluster,

service/history/replication/task_processor_manager.go

@@ -273,7 +273,7 @@ func (r *taskProcessorManagerImpl) cleanupReplicationTasks() error {
 	)
 
 	ctx, cancel := context.WithTimeout(context.Background(), time.Second*5)
-	ctx = headers.SetCallerInfo(ctx, headers.SystemBackgroundCallerInfo)
+	ctx = headers.SetCallerInfo(ctx, headers.SystemPreemptableCallerInfo)
 	defer cancel()
 
 	err := r.shard.GetExecutionManager().RangeCompleteHistoryTasks(

service/worker/replicator/namespace_replication_message_processor.go

@@ -147,7 +147,7 @@ func (p *namespaceReplicationMessageProcessor) getAndHandleNamespaceReplicationT
 	}
 
 	ctx, cancel := rpc.NewContextWithTimeoutAndVersionHeaders(fetchTaskRequestTimeout)
-	ctx = headers.SetCallerInfo(ctx, headers.SystemBackgroundCallerInfo)
+	ctx = headers.SetCallerInfo(ctx, headers.SystemPreemptableCallerInfo)
 	request := &adminservice.GetNamespaceReplicationMessagesRequest{
 		ClusterName:            p.currentCluster,
 		LastRetrievedMessageId: p.lastRetrievedMessageID,
@@ -164,7 +164,7 @@ func (p *namespaceReplicationMessageProcessor) getAndHandleNamespaceReplicationT
 	p.logger.Debug("Successfully fetched namespace replication tasks", tag.Counter(len(response.Messages.ReplicationTasks)))
 
 	// TODO: specify a timeout for processing namespace replication tasks
-	taskCtx := headers.SetCallerInfo(context.TODO(), headers.SystemBackgroundCallerInfo)
+	taskCtx := headers.SetCallerInfo(context.TODO(), headers.SystemPreemptableCallerInfo)
 	for taskIndex := range response.Messages.ReplicationTasks {
 		task := response.Messages.ReplicationTasks[taskIndex]
 		err := backoff.ThrottleRetry(func() error {