Special Handling for new workflow replication (#7561)

## What changed?
<!-- Describe what has changed in this PR -->
Add special handling for new workflow replication
## Why?
<!-- Tell your future self why have you made these changes -->
To reduce passive side loadmutablestate not found attempt
## How did you test it?
<!-- How have you verified this change? Tested locally? Added a unit
test? Checked in staging env? -->
unit test
## Potential risks
<!-- Assuming the worst case, what can be broken when deploying this
change to production? -->
no risk
## Documentation
<!-- Have you made sure this change doesn't falsify anything currently
stated in `docs/`? If significant
new behavior is added, have you described that in `docs/`? -->
n/a
## Is hotfix candidate?
<!-- Is this PR a hotfix candidate or does it require a notification to
be sent to the broader community? (Yes/No) -->
no

---------

Co-authored-by: Yichao Yang <[email protected]>

Author: xwduan

api/persistence/v1/executions.pb.go

@@ -1277,6 +1277,7 @@ func (s *TaskSerializer) replicationSyncVersionedTransitionTaskToProto(
 		NextEventId:            syncVersionedTransitionTask.NextEventID,
 		NewRunId:               syncVersionedTransitionTask.NewRunID,
 		LastVersionHistoryItem: syncVersionedTransitionTask.LastVersionHistoryItem,
+		IsFirstTask:            syncVersionedTransitionTask.IsFirstTask,
 		TaskEquivalents:        taskInfoEquivalents,
 	}, nil
 }
@@ -1313,6 +1314,7 @@ func (s *TaskSerializer) replicationSyncVersionedTransitionTaskFromProto(
 		VersionedTransition:    syncVersionedTransitionTask.VersionedTransition,
 		LastVersionHistoryItem: syncVersionedTransitionTask.LastVersionHistoryItem,
 		TaskEquivalents:        taskEquivalents,
+		IsFirstTask:            syncVersionedTransitionTask.IsFirstTask,
 	}, nil
 }
 

common/persistence/serialization/task_serializer.go

@@ -379,6 +379,7 @@ message ReplicationTaskInfo {
     // doesn't need to be disabled.
     repeated ReplicationTaskInfo task_equivalents = 20;
     history.v1.VersionHistoryItem last_version_history_item = 21;
+    bool is_first_task = 22;
 }
 
 // visibility_task_data column

proto/internal/temporal/server/api/persistence/v1/executions.proto

@@ -229,6 +229,17 @@ func (r *WorkflowStateReplicatorImpl) ReplicateVersionedTransition(
 	default:
 		return serviceerror.NewInvalidArgument(fmt.Sprintf("unknown artifact type %T", artifactType))
 	}
+
+	if mutation != nil && mutation.ExclusiveStartVersionedTransition.TransitionCount == 0 {
+		// this is the first replication task for this workflow
+		// TODO: Handle reset case to reduce the amount of history events write
+		err := r.handleFirstReplicationTask(ctx, versionedTransition, sourceClusterName)
+		if !errors.Is(err, consts.ErrDuplicate) {
+			// if ErrDuplicate is returned from creation, it means the workflow is already existed, continue to apply mutation
+			return err
+		}
+	}
+
 	executionState, executionInfo := func() (*persistencespb.WorkflowExecutionState, *persistencespb.WorkflowExecutionInfo) {
 		if snapshot != nil {
 			return snapshot.State.ExecutionState, snapshot.State.ExecutionInfo
@@ -337,6 +348,168 @@ func (r *WorkflowStateReplicatorImpl) ReplicateVersionedTransition(
 	}
 }
 
+//nolint:revive // cognitive complexity 35 (> max enabled 25)
+func (r *WorkflowStateReplicatorImpl) handleFirstReplicationTask(
+	ctx context.Context,
+	versionedTransitionArtifact *replicationspb.VersionedTransitionArtifact,
+	sourceClusterName string,
+) (retErr error) {
+	mutation := versionedTransitionArtifact.GetSyncWorkflowStateMutationAttributes()
+	executionInfo := mutation.StateMutation.ExecutionInfo
+	executionState := mutation.StateMutation.ExecutionState
+	sourceVersionHistories := mutation.StateMutation.ExecutionInfo.VersionHistories
+	currentVersionHistory, err := versionhistory.GetCurrentVersionHistory(sourceVersionHistories)
+	if err != nil {
+		return err
+	}
+	lastVersionHistoryItem, err := versionhistory.GetLastVersionHistoryItem(currentVersionHistory)
+	if err != nil {
+		return err
+	}
+
+	var historyEventBatchs [][]*historypb.HistoryEvent
+	for _, blob := range versionedTransitionArtifact.EventBatches {
+		e, err := r.historySerializer.DeserializeEvents(blob)
+		if err != nil {
+			return err
+		}
+		historyEventBatchs = append(historyEventBatchs, e)
+	}
+	lastBatch := historyEventBatchs[len(historyEventBatchs)-1]
+	lastEvent := lastBatch[len(lastBatch)-1]
+	if lastEvent.EventId < lastVersionHistoryItem.EventId {
+		remoteHistoryIterator := collection.NewPagingIterator(r.getHistoryFromRemotePaginationFn(
+			ctx,
+			sourceClusterName,
+			namespace.ID(executionInfo.NamespaceId),
+			executionInfo.WorkflowId,
+			executionState.RunId,
+			lastEvent.EventId,
+			lastEvent.Version,
+			lastVersionHistoryItem.EventId+1,
+			lastVersionHistoryItem.Version),
+		)
+		for remoteHistoryIterator.HasNext() {
+			batch, err := remoteHistoryIterator.Next()
+			if err != nil {
+				return err
+			}
+			sourceEvents, err := r.historySerializer.DeserializeEvents(batch.rawHistory)
+			if err != nil {
+				return err
+			}
+			historyEventBatchs = append(historyEventBatchs, sourceEvents)
+		}
+	}
+
+	wfCtx, releaseFn, err := r.workflowCache.GetOrCreateWorkflowExecution(
+		ctx,
+		r.shardContext,
+		namespace.ID(executionInfo.NamespaceId),
+		&commonpb.WorkflowExecution{
+			WorkflowId: executionInfo.WorkflowId,
+			RunId:      executionState.RunId,
+		},
+		locks.PriorityLow,
+	)
+	if err != nil {
+		return err
+	}
+	defer func() {
+		if rec := recover(); rec != nil {
+			releaseFn(errPanic)
+			panic(rec)
+		}
+		releaseFn(retErr)
+	}()
+
+	nsEntry, err := r.namespaceRegistry.GetNamespaceByID(namespace.ID(executionInfo.NamespaceId))
+	if err != nil {
+		return err
+	}
+	localMutableState := workflow.NewMutableState(
+		r.shardContext,
+		r.shardContext.GetEventsCache(),
+		r.logger,
+		nsEntry,
+		executionInfo.WorkflowId,
+		executionState.RunId,
+		timestamp.TimeValue(executionState.StartTime),
+	)
+	err = localMutableState.ApplyMutation(mutation.StateMutation)
+	if err != nil {
+		return err
+	}
+
+	err = localMutableState.SetHistoryTree(executionInfo.WorkflowExecutionTimeout, executionInfo.WorkflowRunTimeout, executionState.RunId)
+	if err != nil {
+		return nil
+	}
+	localCurrentVersionHistory, err := versionhistory.GetCurrentVersionHistory(localMutableState.GetExecutionInfo().VersionHistories)
+	if err != nil {
+		return err
+	}
+	defer func() {
+		if retErr != nil {
+			// if we fail to create workflow, we need to clean up the history branch
+			if err := r.shardContext.GetExecutionManager().DeleteHistoryBranch(ctx, &persistence.DeleteHistoryBranchRequest{
+				ShardID:     r.shardContext.GetShardID(),
+				BranchToken: localCurrentVersionHistory.BranchToken,
+			}); err != nil {
+				r.logger.Error("failed to clean up workflow execution", tag.Error(err))
+			}
+		}
+	}()
+	if err != nil {
+		return err
+	}
+
+	localCurrentVersionHistory.Items = versionhistory.CopyVersionHistoryItems(currentVersionHistory.Items)
+	if err != nil {
+		return err
+	}
+
+	localMutableState.SetHistoryBuilder(historybuilder.NewImmutable(historyEventBatchs...))
+	for _, historyEventBatch := range historyEventBatchs {
+		for _, historyEvent := range historyEventBatch {
+			r.addEventToCache(definition.WorkflowKey{
+				NamespaceID: executionInfo.NamespaceId,
+				WorkflowID:  executionInfo.WorkflowId,
+				RunID:       executionState.RunId,
+			}, historyEvent)
+		}
+	}
+	if versionedTransitionArtifact.NewRunInfo != nil {
+		err = r.createNewRunWorkflow(
+			ctx,
+			namespace.ID(executionInfo.NamespaceId),
+			executionInfo.WorkflowId,
+			versionedTransitionArtifact.NewRunInfo,
+			localMutableState,
+			true,
+		)
+		if err != nil {
+			return err
+		}
+	}
+
+	err = r.taskRefresher.Refresh(ctx, localMutableState)
+
+	if err != nil {
+		return err
+	}
+
+	return r.transactionMgr.CreateWorkflow(
+		ctx,
+		NewWorkflow(
+			r.clusterMetadata,
+			wfCtx,
+			localMutableState,
+			releaseFn,
+		),
+	)
+}
+
 func (r *WorkflowStateReplicatorImpl) applyMutation(
 	ctx context.Context,
 	namespaceID namespace.ID,
@@ -370,7 +543,8 @@ func (r *WorkflowStateReplicatorImpl) applyMutation(
 	if workflow.TransitionHistoryStalenessCheck(localTransitionHistory, mutation.ExclusiveStartVersionedTransition) != nil ||
 		workflow.TransitionHistoryStalenessCheck(sourceTransitionHistory, localVersionedTransition) != nil {
 		return serviceerrors.NewSyncState(
-			fmt.Sprintf("Failed to apply mutation due to version check failed. local transition history: %v, source transition history: %v", localTransitionHistory, sourceTransitionHistory),
+			fmt.Sprintf("Failed to apply mutation due to version check failed. local transition history: %v, source transition history: %v, exclusiveStartVersionedTransition: %v",
+				localTransitionHistory, sourceTransitionHistory, mutation.ExclusiveStartVersionedTransition),
 			namespaceID.String(),
 			workflowID,
 			runID,

service/history/ndc/workflow_state_replicator.go

@@ -878,6 +878,103 @@ func (s *workflowReplicatorSuite) Test_ReplicateVersionedTransition_SameBranch_S
 	s.NoError(err)
 }
 
+func (s *workflowReplicatorSuite) Test_ReplicateVersionedTransition_FirstTask_SyncMutation() {
+	workflowStateReplicator := NewWorkflowStateReplicator(
+		s.mockShard,
+		s.mockWorkflowCache,
+		nil,
+		serialization.NewSerializer(),
+		s.logger,
+	)
+	mockTransactionManager := NewMockTransactionManager(s.controller)
+	mockTaskRefresher := workflow.NewMockTaskRefresher(s.controller)
+	workflowStateReplicator.transactionMgr = mockTransactionManager
+	workflowStateReplicator.taskRefresher = mockTaskRefresher
+	namespaceID := uuid.New()
+	s.workflowStateReplicator.transactionMgr = NewMockTransactionManager(s.controller)
+	versionHistories := &historyspb.VersionHistories{
+		CurrentVersionHistoryIndex: 0,
+		Histories: []*historyspb.VersionHistory{
+			{
+				BranchToken: []byte("branchToken"),
+				Items: []*historyspb.VersionHistoryItem{
+					{
+						EventId: int64(4),
+						Version: int64(2),
+					},
+				},
+			},
+		},
+	}
+	eventBatches := []*historypb.HistoryEvent{
+		{EventId: 1, Version: 2}, {EventId: 2, Version: 2},
+		{EventId: 3, Version: 2}, {EventId: 4, Version: 2},
+	}
+	eventBatchBlob, err := serialization.NewSerializer().SerializeEvents(eventBatches, enumspb.ENCODING_TYPE_PROTO3)
+	s.NoError(err)
+	transitionHistory := []*persistencespb.VersionedTransition{
+		{NamespaceFailoverVersion: 2, TransitionCount: 10},
+	}
+	versionedTransitionArtifact := &replicationspb.VersionedTransitionArtifact{
+		StateAttributes: &replicationspb.VersionedTransitionArtifact_SyncWorkflowStateMutationAttributes{
+			SyncWorkflowStateMutationAttributes: &replicationspb.SyncWorkflowStateMutationAttributes{
+				StateMutation: &persistencespb.WorkflowMutableStateMutation{
+					ExecutionInfo: &persistencespb.WorkflowExecutionInfo{
+						WorkflowId:        s.workflowID,
+						NamespaceId:       namespaceID,
+						VersionHistories:  versionHistories,
+						TransitionHistory: transitionHistory,
+					},
+					ExecutionState: &persistencespb.WorkflowExecutionState{
+						RunId:  s.runID,
+						State:  enumsspb.WORKFLOW_EXECUTION_STATE_RUNNING,
+						Status: enumspb.WORKFLOW_EXECUTION_STATUS_RUNNING,
+					},
+				},
+				ExclusiveStartVersionedTransition: &persistencespb.VersionedTransition{
+					NamespaceFailoverVersion: 2, TransitionCount: 0,
+				},
+			},
+		},
+		EventBatches: []*commonpb.DataBlob{eventBatchBlob},
+	}
+	mockWeCtx := historyi.NewMockWorkflowContext(s.controller)
+	s.mockWorkflowCache.EXPECT().GetOrCreateWorkflowExecution(
+		gomock.Any(),
+		s.mockShard,
+		namespace.ID(namespaceID),
+		&commonpb.WorkflowExecution{
+			WorkflowId: s.workflowID,
+			RunId:      s.runID,
+		},
+		locks.PriorityLow,
+	).Return(mockWeCtx, wcache.NoopReleaseFn, nil)
+	s.mockNamespaceCache.EXPECT().GetNamespaceByID(namespace.ID(namespaceID)).Return(namespace.NewNamespaceForTest(
+		&persistencespb.NamespaceInfo{},
+		nil,
+		false,
+		nil,
+		int64(100),
+	), nil).AnyTimes()
+	mockTaskRefresher.EXPECT().Refresh(gomock.Any(), gomock.Any()).Return(nil).Times(1)
+
+	mockTransactionManager.EXPECT().CreateWorkflow(
+		gomock.Any(),
+		gomock.AssignableToTypeOf(&WorkflowImpl{}),
+	).DoAndReturn(func(ctx context.Context, wf Workflow) error {
+		// Capture localMutableState from the workflow
+		localMutableState := wf.GetMutableState()
+
+		// Perform your comparisons here
+		s.Equal(localMutableState.GetExecutionInfo().TransitionHistory, transitionHistory)
+
+		return nil
+	}).Times(1)
+	err = workflowStateReplicator.ReplicateVersionedTransition(context.Background(), versionedTransitionArtifact, "test")
+	s.NoError(err)
+
+}
+
 func (s *workflowReplicatorSuite) Test_ReplicateVersionedTransition_MutationProvidedWithGap_ReturnSyncStateError() {
 	workflowStateReplicator := NewWorkflowStateReplicator(
 		s.mockShard,