Redirect replication requests when shard count is different (#3789)

* Redirect replication requests when shard count is different

service/history/replication/dlq_handler.go

@@ -339,15 +339,10 @@ func (r *dlqHandlerImpl) getOrCreateTaskExecutor(ctx context.Context, clusterNam
  if executor, ok := r.taskExecutors[clusterName]; ok {
  return executor, nil
  }
- engine, err := r.shard.GetEngine(ctx)
- if err != nil {
- return nil, err
- }
  taskExecutor := r.taskExecutorProvider(TaskExecutorParams{
  RemoteCluster: clusterName,
  Shard: r.shard,
  HistoryResender: r.resender,
- HistoryEngine: engine,
  DeleteManager: r.deleteManager,
  WorkflowCache: r.workflowCache,
  })

service/history/replication/dlq_handler_test.go

@@ -126,7 +126,6 @@ func (s *dlqHandlerSuite) SetupTest() {
  params.RemoteCluster,
  params.Shard,
  params.HistoryResender,
- params.HistoryEngine,
  params.DeleteManager,
  params.WorkflowCache,
  )

service/history/replication/fx.go

@@ -58,7 +58,6 @@ func ReplicationTaskExecutorProvider() TaskExecutorProvider {
  params.RemoteCluster,
  params.Shard,
  params.HistoryResender,
- params.HistoryEngine,
  params.DeleteManager,
  params.WorkflowCache,
  )

service/history/replication/task_executor.go

@@ -56,48 +56,46 @@ type (
  RemoteCluster string // TODO: Remove this remote cluster from executor then it can use singleton.
  Shard shard.Context
  HistoryResender xdc.NDCHistoryResender
- HistoryEngine shard.Engine
- DeleteManager  workflow.DeleteManager
- WorkflowCache  workflow.Cache
+
+ DeleteManager workflow.DeleteManager
+ WorkflowCache workflow.Cache
  }
 
  TaskExecutorProvider func(params TaskExecutorParams) TaskExecutor
 
  taskExecutorImpl struct {
  currentCluster string
  remoteCluster string
- shard  shard.Context
+ shardContext shard.Context
  namespaceRegistry namespace.Registry
  nDCHistoryResender xdc.NDCHistoryResender
- historyEngine shard.Engine
  deleteManager workflow.DeleteManager
  workflowCache workflow.Cache
  metricsHandler metrics.MetricsHandler
- logger log.Logger
+
+ logger log.Logger
  }
 )
 
 // NewTaskExecutor creates a replication task executor
 // The executor uses by 1) DLQ replication task handler 2) history replication task processor
 func NewTaskExecutor(
  remoteCluster string,
- shard shard.Context,
+ shardContext shard.Context,
  nDCHistoryResender xdc.NDCHistoryResender,
- historyEngine shard.Engine,
  deleteManager workflow.DeleteManager,
  workflowCache workflow.Cache,
 ) TaskExecutor {
  return &taskExecutorImpl{
- currentCluster: shard.GetClusterMetadata().GetCurrentClusterName(),
+ currentCluster: shardContext.GetClusterMetadata().GetCurrentClusterName(),
  remoteCluster: remoteCluster,
- shard:  shard,
- namespaceRegistry: shard.GetNamespaceRegistry(),
+ shardContext: shardContext,
+ namespaceRegistry: shardContext.GetNamespaceRegistry(),
  nDCHistoryResender: nDCHistoryResender,
- historyEngine: historyEngine,
  deleteManager: deleteManager,
  workflowCache: workflowCache,
- metricsHandler: shard.GetMetricsHandler(),
- logger: shard.GetLogger(),
+ metricsHandler: shardContext.GetMetricsHandler(),
+ logger: shardContext.GetLogger(),
  }
 }
 
@@ -165,7 +163,9 @@ func (e *taskExecutorImpl) handleActivityTask(
  ctx, cancel := e.newTaskContext(ctx, attr.NamespaceId)
  defer cancel()
 
- err = e.historyEngine.SyncActivity(ctx, request)
+ // This might be extra cost if the workflow belongs to local shard.
+ // Add a wrapper of the history client to call history engine directly if it becomes an issue.
+ _, err = e.shardContext.GetHistoryClient().SyncActivity(ctx, request)
  switch retryErr := err.(type) {
  case nil:
  return nil
@@ -204,7 +204,10 @@ func (e *taskExecutorImpl) handleActivityTask(
  e.logger.Error("error resend history for history event", tag.Error(resendErr))
  return err
  }
- return e.historyEngine.SyncActivity(ctx, request)
+ // This might be extra cost if the workflow belongs to local shard.
+ // Add a wrapper of the history client to call history engine directly if it becomes an issue.
+ _, err = e.shardContext.GetHistoryClient().SyncActivity(ctx, request)
+ return err
 
  default:
  return err
@@ -245,7 +248,9 @@ func (e *taskExecutorImpl) handleHistoryReplicationTask(
  ctx, cancel := e.newTaskContext(ctx, attr.NamespaceId)
  defer cancel()
 
- err = e.historyEngine.ReplicateEventsV2(ctx, request)
+ // This might be extra cost if the workflow belongs to local shard.
+ // Add a wrapper of the history client to call history engine directly if it becomes an issue.
+ _, err = e.shardContext.GetHistoryClient().ReplicateEventsV2(ctx, request)
  switch retryErr := err.(type) {
  case nil:
  return nil
@@ -285,7 +290,10 @@ func (e *taskExecutorImpl) handleHistoryReplicationTask(
  return err
  }
 
- return e.historyEngine.ReplicateEventsV2(ctx, request)
+ // This might be extra cost if the workflow belongs to local shard.
+ // Add a wrapper of the history client to call history engine directly if it becomes an issue.
+ _, err = e.shardContext.GetHistoryClient().ReplicateEventsV2(ctx, request)
+ return err
 
  default:
  return err
@@ -310,10 +318,13 @@ func (e *taskExecutorImpl) handleSyncWorkflowStateTask(
  ctx, cancel := e.newTaskContext(ctx, executionInfo.NamespaceId)
  defer cancel()
 
- return e.historyEngine.ReplicateWorkflowState(ctx, &historyservice.ReplicateWorkflowStateRequest{
+ // This might be extra cost if the workflow belongs to local shard.
+ // Add a wrapper of the history client to call history engine directly if it becomes an issue.
+ _, err = e.shardContext.GetHistoryClient().ReplicateWorkflowState(ctx, &historyservice.ReplicateWorkflowStateRequest{
  WorkflowState: attr.GetWorkflowState(),
  RemoteCluster: e.remoteCluster,
  })
+ return err
 }
 
 func (e *taskExecutorImpl) filterTask(

service/history/replication/task_executor_test.go

@@ -63,7 +63,6 @@ type (
  remoteCluster string
  mockResource *resourcetest.Test
  mockShard *shard.ContextTest
- mockEngine *shard.MockEngine
  config *configs.Config
  historyClient *historyservicemock.MockHistoryServiceClient
  mockNamespaceCache *namespace.MockRegistry
@@ -106,7 +105,6 @@ func (s *taskExecutorSuite) SetupTest() {
  }},
  s.config,
  )
- s.mockEngine = shard.NewMockEngine(s.controller)
  s.mockResource = s.mockShard.Resource
  s.mockNamespaceCache = s.mockResource.NamespaceCache
  s.clusterMetadata = s.mockResource.ClusterMetadata
@@ -116,12 +114,11 @@ func (s *taskExecutorSuite) SetupTest() {
 
  s.clusterMetadata.EXPECT().GetCurrentClusterName().Return(cluster.TestCurrentClusterName).AnyTimes()
  s.mockNamespaceCache.EXPECT().GetNamespaceName(gomock.Any()).Return(tests.Namespace, nil).AnyTimes()
-
+ s.mockShard.SetHistoryClientForTesting(s.historyClient)
  s.replicationTaskExecutor = NewTaskExecutor(
  s.remoteCluster,
  s.mockShard,
  s.nDCHistoryResender,
- s.mockEngine,
  workflow.NewMockDeleteManager(s.controller),
  s.workflowCache,
  ).(*taskExecutorImpl)
@@ -230,7 +227,7 @@ func (s *taskExecutorSuite) TestProcessTaskOnce_SyncActivityReplicationTask() {
  LastWorkerIdentity: "",
  }
 
- s.mockEngine.EXPECT().SyncActivity(gomock.Any(), request).Return(nil)
+ s.historyClient.EXPECT().SyncActivity(gomock.Any(), request).Return(&historyservice.SyncActivityResponse{}, nil)
  err := s.replicationTaskExecutor.Execute(context.Background(), task, true)
  s.NoError(err)
 }
@@ -283,7 +280,7 @@ func (s *taskExecutorSuite) TestProcessTaskOnce_SyncActivityReplicationTask_Rese
  345,
  456,
  )
- s.mockEngine.EXPECT().SyncActivity(gomock.Any(), request).Return(resendErr)
+ s.historyClient.EXPECT().SyncActivity(gomock.Any(), request).Return(nil, resendErr)
  s.nDCHistoryResender.EXPECT().SendSingleWorkflowHistory(
  gomock.Any(),
  s.remoteCluster,
@@ -295,7 +292,8 @@ func (s *taskExecutorSuite) TestProcessTaskOnce_SyncActivityReplicationTask_Rese
  int64(345),
  int64(456),
  )
- s.mockEngine.EXPECT().SyncActivity(gomock.Any(), request).Return(nil)
+
+ s.historyClient.EXPECT().SyncActivity(gomock.Any(), request).Return(&historyservice.SyncActivityResponse{}, nil)
  err := s.replicationTaskExecutor.Execute(context.Background(), task, true)
  s.NoError(err)
 }
@@ -327,8 +325,7 @@ func (s *taskExecutorSuite) TestProcess_HistoryReplicationTask() {
  Events: nil,
  NewRunEvents: nil,
  }
-
- s.mockEngine.EXPECT().ReplicateEventsV2(gomock.Any(), request).Return(nil)
+ s.historyClient.EXPECT().ReplicateEventsV2(gomock.Any(), request).Return(&historyservice.ReplicateEventsV2Response{}, nil)
  err := s.replicationTaskExecutor.Execute(context.Background(), task, true)
  s.NoError(err)
 }
@@ -371,7 +368,7 @@ func (s *taskExecutorSuite) TestProcess_HistoryReplicationTask_Resend() {
  345,
  456,
  )
- s.mockEngine.EXPECT().ReplicateEventsV2(gomock.Any(), request).Return(resendErr)
+ s.historyClient.EXPECT().ReplicateEventsV2(gomock.Any(), request).Return(nil, resendErr)
  s.nDCHistoryResender.EXPECT().SendSingleWorkflowHistory(
  gomock.Any(),
  s.remoteCluster,
@@ -383,7 +380,8 @@ func (s *taskExecutorSuite) TestProcess_HistoryReplicationTask_Resend() {
  int64(345),
  int64(456),
  )
- s.mockEngine.EXPECT().ReplicateEventsV2(gomock.Any(), request).Return(nil)
+
+ s.historyClient.EXPECT().ReplicateEventsV2(gomock.Any(), request).Return(&historyservice.ReplicateEventsV2Response{}, nil)
  err := s.replicationTaskExecutor.Execute(context.Background(), task, true)
  s.NoError(err)
 }
@@ -402,8 +400,7 @@ func (s *taskExecutorSuite) TestProcessTaskOnce_SyncWorkflowStateTask() {
  },
  },
  }
-
- s.mockEngine.EXPECT().ReplicateWorkflowState(gomock.Any(), gomock.Any()).Return(nil)
+ s.historyClient.EXPECT().ReplicateWorkflowState(gomock.Any(), gomock.Any()).Return(&historyservice.ReplicateWorkflowStateResponse{}, nil)
 
  err := s.replicationTaskExecutor.Execute(context.Background(), task, true)
  s.NoError(err)

service/history/replication/task_processor_manager.go

@@ -189,7 +189,6 @@ func (r *taskProcessorManagerImpl) handleClusterMetadataUpdate(
  RemoteCluster: clusterName,
  Shard: r.shard,
  HistoryResender: r.resender,
- HistoryEngine: r.engine,
  DeleteManager: r.deleteMgr,
  WorkflowCache: r.workflowCache,
  }),

service/history/shard/context_testutil.go

@@ -31,6 +31,7 @@ import (
 
  "github.com/golang/mock/gomock"
 
+ "go.temporal.io/server/api/historyservice/v1"
  persistencespb "go.temporal.io/server/api/persistence/v1"
  "go.temporal.io/server/common/clock"
  "go.temporal.io/server/common/future"
@@ -132,6 +133,11 @@ func (s *ContextTest) SetEventsCacheForTesting(c events.Cache) {
  s.eventsCache = c
 }
 
+// SetHistoryClientForTesting sets history client. Only used by tests.
+func (s *ContextTest) SetHistoryClientForTesting(client historyservice.HistoryServiceClient) {
+ s.historyClient = client
+}
+
 // StopForTest calls private method finishStop(). In general only the controller
 // should call that, but integration tests need to do it also to clean up any
 // background acquireShard goroutines that may exist.