Ability to generate migration task via frontend (#8078)

## What changed?
- Add capability to generate migration task via frontend.

## Why?
- We don't really have per namespace rate limiting on history client,
going through frontend allow us to rate limiting migration traffic per
namespace and protect the cluster.

## How did you test it?
- [x] built
- [ ] run locally and tested manually
- [ ] covered by existing tests
- [x] added new unit test(s)
- [ ] added new functional test(s)

## Potential risks
- Feature is behind a feature flag, so should be safe.
- The feature flag can only be enabled after the change in the admin
handler (for forwarding the target cluster field) is deployed.

Author: yycptt

api/adminservice/v1/request_response.pb.go

@@ -2703,6 +2703,11 @@ Valid fields: MaxConcurrentActivityExecutionSize, TaskQueueActivitiesPerSecond,
 WorkerActivitiesPerSecond, MaxConcurrentActivityTaskPollers.
 `,
 	)
+	WorkerGenerateMigrationTaskViaFrontend = NewGlobalBoolSetting(
+		"worker.generateMigrationTaskViaFrontend",
+		false,
+		`WorkerGenerateMigrationTaskViaFrontend controls whether to generate migration tasks via frontend admin service.`,
+	)
 	MaxUserMetadataSummarySize = NewNamespaceIntSetting(
 		"limit.userMetadataSummarySize",
 		400,

common/dynamicconfig/constants.go

@@ -563,6 +563,7 @@ message SyncWorkflowStateResponse {
 message GenerateLastHistoryReplicationTasksRequest {
   string namespace = 1;
   temporal.api.common.v1.WorkflowExecution execution = 2;
+  repeated string target_clusters = 3;
 }
 
 message GenerateLastHistoryReplicationTasksResponse {

proto/internal/temporal/server/api/adminservice/v1/request_response.proto

@@ -2151,8 +2151,9 @@ func (adh *AdminHandler) GenerateLastHistoryReplicationTasks(
 	resp, err := adh.historyClient.GenerateLastHistoryReplicationTasks(
 		ctx,
 		&historyservice.GenerateLastHistoryReplicationTasksRequest{
-			NamespaceId: namespaceEntry.ID().String(),
-			Execution:   request.Execution,
+			NamespaceId:    namespaceEntry.ID().String(),
+			Execution:      request.Execution,
+			TargetClusters: request.TargetClusters,
 		},
 	)
 	if err != nil {

service/frontend/admin_handler.go

@@ -15,11 +15,13 @@ import (
 	"go.temporal.io/api/workflowservice/v1"
 	"go.temporal.io/sdk/activity"
 	"go.temporal.io/sdk/temporal"
+	"go.temporal.io/server/api/adminservice/v1"
 	enumsspb "go.temporal.io/server/api/enums/v1"
 	"go.temporal.io/server/api/historyservice/v1"
 	replicationspb "go.temporal.io/server/api/replication/v1"
 	serverClient "go.temporal.io/server/client"
 	"go.temporal.io/server/common/definition"
+	"go.temporal.io/server/common/dynamicconfig"
 	"go.temporal.io/server/common/headers"
 	"go.temporal.io/server/common/log"
 	"go.temporal.io/server/common/log/tag"
@@ -100,18 +102,20 @@ type (
 	}
 
 	activities struct {
-		historyShardCount              int32
-		executionManager               persistence.ExecutionManager
-		taskManager                    persistence.TaskManager
-		namespaceRegistry              namespace.Registry
-		historyClient                  historyservice.HistoryServiceClient
-		frontendClient                 workflowservice.WorkflowServiceClient
-		clientFactory                  serverClient.Factory
-		clientBean                     serverClient.Bean
-		logger                         log.Logger
-		metricsHandler                 metrics.Handler
-		forceReplicationMetricsHandler metrics.Handler
-		namespaceReplicationQueue      persistence.NamespaceReplicationQueue
+		historyShardCount                int32
+		executionManager                 persistence.ExecutionManager
+		taskManager                      persistence.TaskManager
+		namespaceRegistry                namespace.Registry
+		historyClient                    historyservice.HistoryServiceClient
+		frontendClient                   workflowservice.WorkflowServiceClient
+		adminClient                      adminservice.AdminServiceClient
+		clientFactory                    serverClient.Factory
+		clientBean                       serverClient.Bean
+		logger                           log.Logger
+		metricsHandler                   metrics.Handler
+		forceReplicationMetricsHandler   metrics.Handler
+		namespaceReplicationQueue        persistence.NamespaceReplicationQueue
+		generateMigrationTaskViaFrontend dynamicconfig.BoolPropertyFn
 	}
 )
 
@@ -336,7 +340,15 @@ func (a *activities) checkHandoverOnce(ctx context.Context, waitRequest waitHand
 	return readyShardCount == len(resp.Shards), nil
 }
 
-func (a *activities) generateWorkflowReplicationTask(ctx context.Context, rateLimiter quotas.RateLimiter, wKey definition.WorkflowKey, targetClusters []string) error {
+func (a *activities) generateWorkflowReplicationTask(
+	ctx context.Context,
+	rateLimiter quotas.RateLimiter,
+	namespaceName string,
+	namespaceID string,
+	we *commonpb.WorkflowExecution,
+	targetClusters []string,
+	generateViaFrontend bool,
+) error {
 	if err := rateLimiter.WaitN(ctx, 1); err != nil {
 		return err
 	}
@@ -345,24 +357,36 @@ func (a *activities) generateWorkflowReplicationTask(ctx context.Context, rateLi
 	ctx, cancel := context.WithTimeout(ctx, time.Second*10)
 	defer cancel()
 
-	resp, err := a.historyClient.GenerateLastHistoryReplicationTasks(ctx, &historyservice.GenerateLastHistoryReplicationTasksRequest{
-		NamespaceId: wKey.NamespaceID,
-		Execution: &commonpb.WorkflowExecution{
-			WorkflowId: wKey.WorkflowID,
-			RunId:      wKey.RunID,
-		},
-		TargetClusters: targetClusters,
-	})
-
-	if err != nil {
-		return err
+	var stateTransitionCount, historyLength int64
+	if generateViaFrontend {
+		resp, err := a.adminClient.GenerateLastHistoryReplicationTasks(ctx, &adminservice.GenerateLastHistoryReplicationTasksRequest{
+			Namespace:      namespaceName,
+			Execution:      we,
+			TargetClusters: targetClusters,
+		})
+		if err != nil {
+			return err
+		}
+		stateTransitionCount = resp.StateTransitionCount
+		historyLength = resp.HistoryLength
+	} else {
+		resp, err := a.historyClient.GenerateLastHistoryReplicationTasks(ctx, &historyservice.GenerateLastHistoryReplicationTasksRequest{
+			NamespaceId:    namespaceID,
+			Execution:      we,
+			TargetClusters: targetClusters,
+		})
+		if err != nil {
+			return err
+		}
+		stateTransitionCount = resp.StateTransitionCount
+		historyLength = resp.HistoryLength
 	}
 
 	// If workflow has many activity retries (bug in activity code e.g.,), the state transition count can be
 	// large but the number of actual state transition that is applied on target cluster can be very small.
 	// Take the minimum between StateTransitionCount and HistoryLength as heuristic to avoid unnecessary throttling
 	// in such situation.
-	count := min(resp.StateTransitionCount, resp.HistoryLength)
+	count := min(stateTransitionCount, historyLength)
 	for count > 0 {
 		token := min(int(count), rateLimiter.Burst())
 		count -= int64(token)
@@ -475,27 +499,38 @@ func (a *activities) GenerateReplicationTasks(ctx context.Context, request *gene
 		}
 	}
 
+	namespaceName, err := a.namespaceRegistry.GetNamespaceName(namespace.ID(request.NamespaceID))
+	if err != nil {
+		a.logger.Error("force-replication failed to translate namespaceID to name", tag.WorkflowNamespaceID(request.NamespaceID))
+		return err
+	}
+
+	generateViaFrontend := a.generateMigrationTaskViaFrontend()
 	for i := startIndex; i < len(request.Executions); i++ {
-		var executionCandidates []definition.WorkflowKey
-		executionCandidates = []definition.WorkflowKey{definition.NewWorkflowKey(request.NamespaceID, request.Executions[i].GetWorkflowId(), request.Executions[i].GetRunId())}
-
-		for _, we := range executionCandidates {
-			if err := a.generateWorkflowReplicationTask(ctx, rateLimiter, we, request.TargetClusters); err != nil {
-				if !isNotFoundServiceError(err) {
-					a.logger.Error("force-replication failed to generate replication task",
-						tag.WorkflowNamespaceID(we.GetNamespaceID()),
-						tag.WorkflowID(we.GetWorkflowID()),
-						tag.WorkflowRunID(we.GetRunID()),
-						tag.Error(err))
-					return err
-				}
-
-				a.logger.Warn("force-replication ignore replication task due to NotFoundServiceError",
-					tag.WorkflowNamespaceID(we.GetNamespaceID()),
-					tag.WorkflowID(we.GetWorkflowID()),
-					tag.WorkflowRunID(we.GetRunID()),
+		we := request.Executions[i]
+		if err := a.generateWorkflowReplicationTask(
+			ctx,
+			rateLimiter,
+			namespaceName.String(),
+			request.NamespaceID,
+			we,
+			request.TargetClusters,
+			generateViaFrontend,
+		); err != nil {
+			if !isNotFoundServiceError(err) {
+				a.logger.Error("force-replication failed to generate replication task",
+					tag.WorkflowNamespaceID(request.NamespaceID),
+					tag.WorkflowID(we.GetWorkflowId()),
+					tag.WorkflowRunID(we.GetRunId()),
 					tag.Error(err))
+				return err
 			}
+
+			a.logger.Warn("force-replication ignore replication task due to NotFoundServiceError",
+				tag.WorkflowNamespaceID(request.NamespaceID),
+				tag.WorkflowID(we.GetWorkflowId()),
+				tag.WorkflowRunID(we.GetRunId()),
+				tag.Error(err))
 		}
 		activity.RecordHeartbeat(ctx, i)
 	}

service/worker/migration/activities.go

@@ -13,11 +13,14 @@ import (
 	"go.temporal.io/sdk/interceptor"
 	"go.temporal.io/sdk/testsuite"
 	"go.temporal.io/sdk/worker"
+	"go.temporal.io/server/api/adminservice/v1"
+	"go.temporal.io/server/api/adminservicemock/v1"
 	enumsspb "go.temporal.io/server/api/enums/v1"
 	"go.temporal.io/server/api/historyservice/v1"
 	"go.temporal.io/server/api/historyservicemock/v1"
 	persistencespb "go.temporal.io/server/api/persistence/v1"
 	"go.temporal.io/server/client"
+	"go.temporal.io/server/common/dynamicconfig"
 	"go.temporal.io/server/common/log"
 	"go.temporal.io/server/common/metrics"
 	"go.temporal.io/server/common/namespace"
@@ -43,6 +46,7 @@ type activitiesSuite struct {
 	mockClientBean                *client.MockBean
 
 	mockFrontendClient *workflowservicemock.MockWorkflowServiceClient
+	mockAdminClient    *adminservicemock.MockAdminServiceClient
 	mockHistoryClient  *historyservicemock.MockHistoryServiceClient
 	mockRemoteClient   *workflowservicemock.MockWorkflowServiceClient
 
@@ -101,6 +105,7 @@ func (s *activitiesSuite) SetupTest() {
 	s.mockClientBean = client.NewMockBean(s.controller)
 
 	s.mockFrontendClient = workflowservicemock.NewMockWorkflowServiceClient(s.controller)
+	s.mockAdminClient = adminservicemock.NewMockAdminServiceClient(s.controller)
 	s.mockHistoryClient = historyservicemock.NewMockHistoryServiceClient(s.controller)
 	s.mockRemoteClient = workflowservicemock.NewMockWorkflowServiceClient(s.controller)
 
@@ -116,16 +121,18 @@ func (s *activitiesSuite) SetupTest() {
 		Return(&testNamespace, nil).AnyTimes()
 
 	s.a = &activities{
-		namespaceRegistry:              s.mockNamespaceRegistry,
-		namespaceReplicationQueue:      s.mockNamespaceReplicationQueue,
-		clientFactory:                  s.mockClientFactory,
-		clientBean:                     s.mockClientBean,
-		taskManager:                    s.mockTaskManager,
-		frontendClient:                 s.mockFrontendClient,
-		historyClient:                  s.mockHistoryClient,
-		logger:                         log.NewCLILogger(),
-		metricsHandler:                 s.mockMetricsHandler,
-		forceReplicationMetricsHandler: s.mockMetricsHandler,
+		namespaceRegistry:                s.mockNamespaceRegistry,
+		namespaceReplicationQueue:        s.mockNamespaceReplicationQueue,
+		clientFactory:                    s.mockClientFactory,
+		clientBean:                       s.mockClientBean,
+		taskManager:                      s.mockTaskManager,
+		frontendClient:                   s.mockFrontendClient,
+		adminClient:                      s.mockAdminClient,
+		historyClient:                    s.mockHistoryClient,
+		logger:                           log.NewCLILogger(),
+		metricsHandler:                   s.mockMetricsHandler,
+		forceReplicationMetricsHandler:   s.mockMetricsHandler,
+		generateMigrationTaskViaFrontend: dynamicconfig.GetBoolPropertyFn(false),
 	}
 }
 
@@ -651,6 +658,37 @@ func (s *activitiesSuite) TestGenerateReplicationTasks_Failed() {
 	s.Equal(0, lastHeartBeat)
 }
 
+func (s *activitiesSuite) TestGenerateReplicationTasks_Success_ViaFrontend() {
+	env, iceptor := s.initEnv()
+	s.a.generateMigrationTaskViaFrontend = dynamicconfig.GetBoolPropertyFn(true)
+
+	request := generateReplicationTasksRequest{
+		NamespaceID:      mockedNamespaceID,
+		RPS:              10,
+		GetParentInfoRPS: 10,
+		Executions:       []*commonpb.WorkflowExecution{execution1, execution2},
+		TargetClusters:   []string{remoteCluster},
+	}
+
+	// Test startIndex logic, and it should be 1 when running the activity.
+	env.SetHeartbeatDetails(0)
+
+	we := request.Executions[1]
+	s.mockAdminClient.EXPECT().GenerateLastHistoryReplicationTasks(gomock.Any(), protomock.Eq(&adminservice.GenerateLastHistoryReplicationTasksRequest{
+		Namespace:      mockedNamespace,
+		Execution:      we,
+		TargetClusters: []string{remoteCluster},
+	})).Return(&adminservice.GenerateLastHistoryReplicationTasksResponse{}, nil).Times(1)
+
+	_, err := env.ExecuteActivity(s.a.GenerateReplicationTasks, &request)
+	s.NoError(err)
+
+	s.Len(iceptor.generateReplicationRecordedHeartbeats, 1)
+	lastIdx := len(iceptor.generateReplicationRecordedHeartbeats) - 1
+	lastHeartBeat := iceptor.generateReplicationRecordedHeartbeats[lastIdx]
+	s.Equal(1, lastHeartBeat)
+}
+
 func (s *activitiesSuite) TestCountWorkflows() {
 	env, _ := s.initEnv()
 

service/worker/migration/activities_test.go

@@ -8,6 +8,7 @@ import (
 	"go.temporal.io/sdk/workflow"
 	serverClient "go.temporal.io/server/client"
 	"go.temporal.io/server/common/config"
+	"go.temporal.io/server/common/dynamicconfig"
 	"go.temporal.io/server/common/headers"
 	"go.temporal.io/server/common/log"
 	"go.temporal.io/server/common/metrics"
@@ -33,6 +34,7 @@ type (
 		TaskManager               persistence.TaskManager
 		Logger                    log.Logger
 		MetricsHandler            metrics.Handler
+		DynamicCollection         *dynamicconfig.Collection
 	}
 
 	fxResult struct {
@@ -85,17 +87,18 @@ func (wc *replicationWorkerComponent) DedicatedActivityWorkerOptions() *workerco
 
 func (wc *replicationWorkerComponent) activities() *activities {
 	return &activities{
-		historyShardCount:              wc.PersistenceConfig.NumHistoryShards,
-		executionManager:               wc.ExecutionManager,
-		namespaceRegistry:              wc.NamespaceRegistry,
-		historyClient:                  wc.HistoryClient,
-		frontendClient:                 wc.FrontendClient,
-		clientFactory:                  wc.ClientFactory,
-		clientBean:                     wc.ClientBean,
-		namespaceReplicationQueue:      wc.NamespaceReplicationQueue,
-		taskManager:                    wc.TaskManager,
-		logger:                         wc.Logger,
-		metricsHandler:                 wc.MetricsHandler,
-		forceReplicationMetricsHandler: wc.MetricsHandler.WithTags(metrics.WorkflowTypeTag(forceReplicationWorkflowName)),
+		historyShardCount:                wc.PersistenceConfig.NumHistoryShards,
+		executionManager:                 wc.ExecutionManager,
+		namespaceRegistry:                wc.NamespaceRegistry,
+		historyClient:                    wc.HistoryClient,
+		frontendClient:                   wc.FrontendClient,
+		clientFactory:                    wc.ClientFactory,
+		clientBean:                       wc.ClientBean,
+		namespaceReplicationQueue:        wc.NamespaceReplicationQueue,
+		taskManager:                      wc.TaskManager,
+		logger:                           wc.Logger,
+		metricsHandler:                   wc.MetricsHandler,
+		forceReplicationMetricsHandler:   wc.MetricsHandler.WithTags(metrics.WorkflowTypeTag(forceReplicationWorkflowName)),
+		generateMigrationTaskViaFrontend: dynamicconfig.WorkerGenerateMigrationTaskViaFrontend.Get(wc.DynamicCollection),
 	}
 }