fix state-based replication for nexus and enable related test (#6932)

## What changed?
<!-- Describe what has changed in this PR -->
- Use deserialized data when applying hsm node mutation.
- Invalid cache after update hsm node.
- Skip current run check when dispatch workflow update if transition
history is enabled.
- xdc nexus: add test with transition history enabled.

## Why?
<!-- Tell your future self why have you made these changes -->

## 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? -->

## 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/`? -->

## Is hotfix candidate?
<!-- Is this PR a hotfix candidate or does it require a notification to
be sent to the broader community? (Yes/No) -->

Author: hai719

Diff for: service/history/hsm/tree.go

@@ -356,6 +356,10 @@ func (n *Node) Child(path []Key) (*Node, error) {
 	return child.Child(rest)
 }
 
+func (n *Node) InvalidateCache() {
+	n.cache.dataLoaded = false
+}
+
 // AddChild adds an immediate child to a node, serializing the given data.
 // Returns [ErrStateMachineAlreadyExists] if a child with the given key already exists, [ErrNotRegistered] if the key's
 // type is not found in the node's state machine registry and serialization errors.

Diff for: service/history/ndc/transaction_manager_existing_workflow.go

@@ -93,8 +93,9 @@ func (r *nDCTransactionMgrForExistingWorkflowImpl) dispatchForExistingWorkflow(
 		)
 	}
 
-	targetExecutionInfo := targetWorkflow.GetMutableState().GetExecutionInfo()
-	targetExecutionState := targetWorkflow.GetMutableState().GetExecutionState()
+	mutableState := targetWorkflow.GetMutableState()
+	targetExecutionInfo := mutableState.GetExecutionInfo()
+	targetExecutionState := mutableState.GetExecutionState()
 	namespaceID := namespace.ID(targetExecutionInfo.NamespaceId)
 	workflowID := targetExecutionInfo.WorkflowId
 	targetRunID := targetExecutionState.RunId
@@ -115,7 +116,7 @@ func (r *nDCTransactionMgrForExistingWorkflowImpl) dispatchForExistingWorkflow(
 	}
 
 	if currentRunID == targetRunID {
-		if !isWorkflowRebuilt {
+		if !mutableState.IsTransitionHistoryEnabled() && !isWorkflowRebuilt {
 			return serviceerror.NewInternal("transactionMgr: encountered workflow not rebuilt & current workflow not guaranteed")
 		}
 

Diff for: service/history/ndc/transaction_manager_existing_workflow_test.go

@@ -139,6 +139,7 @@ func (s *transactionMgrForExistingWorkflowSuite) TestDispatchForExistingWorkflow
 	targetMutableState.EXPECT().GetExecutionState().Return(&persistencespb.WorkflowExecutionState{
 		RunId: targetRunID,
 	}).AnyTimes()
+	targetMutableState.EXPECT().IsTransitionHistoryEnabled().Return(false)
 	s.mockTransactionMgr.EXPECT().GetCurrentWorkflowRunID(ctx, namespaceID, workflowID).Return(targetRunID, nil)
 
 	err := s.updateMgr.dispatchForExistingWorkflow(ctx, isWorkflowRebuilt, targetWorkflow, newWorkflow)
@@ -482,6 +483,7 @@ func (s *transactionMgrForExistingWorkflowSuite) TestDispatchForExistingWorkflow
 	targetMutableState.EXPECT().GetExecutionState().Return(&persistencespb.WorkflowExecutionState{
 		RunId: targetRunID,
 	}).AnyTimes()
+	targetMutableState.EXPECT().IsTransitionHistoryEnabled().Return(false)
 	s.mockTransactionMgr.EXPECT().GetCurrentWorkflowRunID(ctx, namespaceID, workflowID).Return(targetRunID, nil)
 
 	targetContext.EXPECT().ConflictResolveWorkflowExecution(

Diff for: service/history/shard/task_key_generator.go

@@ -106,6 +106,8 @@ func (a *taskKeyGenerator) setTaskKeys(
 							tag.WorkflowNamespaceID(task.GetNamespaceID()),
 							tag.WorkflowID(task.GetWorkflowID()),
 							tag.WorkflowRunID(task.GetRunID()),
+							tag.TaskType(task.GetType()),
+							tag.TaskID(id),
 							tag.Timestamp(taskScheduledTime),
 							tag.CursorTimestamp(a.taskMinScheduledTime),
 							tag.ValueShardAllocateTimerBeforeRead,

Diff for: service/history/workflow/mutable_state_impl.go

@@ -6929,7 +6929,7 @@ func (ms *MutableStateImpl) applyUpdatesToStateMachineNodes(
 		for _, p := range nodeMutation.Path.Path {
 			incomingPath = append(incomingPath, hsm.Key{Type: p.Type, ID: p.Id})
 		}
-		currentNode, err := root.Child(incomingPath)
+		node, err := root.Child(incomingPath)
 		if err != nil {
 			if !errors.Is(err, hsm.ErrStateMachineNotFound) {
 				return err
@@ -6943,16 +6943,29 @@ func (ms *MutableStateImpl) applyUpdatesToStateMachineNodes(
 				// we don't have enough information to recreate all parents
 				return err
 			}
-			newNode, err := parentNode.AddChild(incomingPath[len(incomingPath)-1], nodeMutation.Data)
-			if err != nil {
-				return err
+
+			key := incomingPath[len(incomingPath)-1]
+			parentInternalNode := parentNode.InternalRepr()
+
+			internalNode = &persistencespb.StateMachineNode{
+				Children: make(map[string]*persistencespb.StateMachineMap),
+			}
+			children, ok := parentInternalNode.Children[key.Type]
+			if !ok {
+				children = &persistencespb.StateMachineMap{MachinesById: make(map[string]*persistencespb.StateMachineNode)}
+				// Children may be nil if the map was empty and the proto message we serialized and deserialized.
+				if parentInternalNode.Children == nil {
+					parentInternalNode.Children = make(map[string]*persistencespb.StateMachineMap, 1)
+				}
+				parentInternalNode.Children[key.Type] = children
 			}
-			internalNode = newNode.InternalRepr()
+			children.MachinesById[key.ID] = internalNode
 		} else {
-			internalNode = currentNode.InternalRepr()
+			internalNode = node.InternalRepr()
 			if CompareVersionedTransition(nodeMutation.LastUpdateVersionedTransition, internalNode.LastUpdateVersionedTransition) == 0 {
 				continue
 			}
+			node.InvalidateCache()
 		}
 		internalNode.Data = nodeMutation.Data
 		internalNode.InitialVersionedTransition = nodeMutation.InitialVersionedTransition

Diff for: tests/xdc/base.go

@@ -72,6 +72,8 @@ type (
 
 		startTime          time.Time
 		onceClusterConnect sync.Once
+
+		enableTransitionHistory bool
 	}
 )
 
@@ -98,6 +100,7 @@ func (s *xdcBaseSuite) setupSuite(clusterNames []string, opts ...testcore.Option
 		s.dynamicConfigOverrides = make(map[dynamicconfig.Key]interface{})
 	}
 	s.dynamicConfigOverrides[dynamicconfig.ClusterMetadataRefreshInterval.Key()] = time.Second * 5
+	s.dynamicConfigOverrides[dynamicconfig.EnableTransitionHistory.Key()] = s.enableTransitionHistory
 
 	fileName := "../testdata/xdc_clusters.yaml"
 	if testcore.TestFlags.TestClusterConfigFile != "" {

Diff for: tests/xdc/nexus_state_replication_test.go

@@ -47,6 +47,7 @@ import (
 	"go.temporal.io/api/workflow/v1"
 	"go.temporal.io/api/workflowservice/v1"
 	sdkclient "go.temporal.io/sdk/client"
+	"go.temporal.io/server/common"
 	"go.temporal.io/server/common/dynamicconfig"
 	commonnexus "go.temporal.io/server/common/nexus"
 	"go.temporal.io/server/common/nexus/nexustest"
@@ -62,7 +63,25 @@ type NexusStateReplicationSuite struct {
 
 func TestNexusStateReplicationTestSuite(t *testing.T) {
 	t.Parallel()
-	suite.Run(t, new(NexusStateReplicationSuite))
+	for _, tc := range []struct {
+		name                    string
+		enableTransitionHistory bool
+	}{
+		{
+			name:                    "DisableTransitionHistory",
+			enableTransitionHistory: false,
+		},
+		{
+			name:                    "EnableTransitionHistory",
+			enableTransitionHistory: true,
+		},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			s := &NexusStateReplicationSuite{}
+			s.enableTransitionHistory = tc.enableTransitionHistory
+			suite.Run(t, s)
+		})
+	}
 }
 
 func (s *NexusStateReplicationSuite) SetupSuite() {
@@ -73,7 +92,8 @@ func (s *NexusStateReplicationSuite) SetupSuite() {
 		dynamicconfig.RefreshNexusEndpointsMinWait.Key(): 1 * time.Millisecond,
 		callbacks.AllowedAddresses.Key():                 []any{map[string]any{"Pattern": "*", "AllowInsecure": true}},
 	}
-	s.setupSuite([]string{"nexus_state_replication_active", "nexus_state_replication_standby"})
+	suffix := "_" + common.GenerateRandomString(5)
+	s.setupSuite([]string{"nexus_state_replication_active" + suffix, "nexus_state_replication_standby" + suffix})
 }
 
 func (s *NexusStateReplicationSuite) SetupTest() {

Diff for: tests/xdc/stream_based_replication_test.go

@@ -65,17 +65,17 @@ import (
 type (
 	streamBasedReplicationTestSuite struct {
 		xdcBaseSuite
-		controller              *gomock.Controller
-		namespaceName           string
-		namespaceID             string
-		serializer              serialization.Serializer
-		generator               test.Generator
-		once                    sync.Once
-		enableTransitionHistory bool
+		controller    *gomock.Controller
+		namespaceName string
+		namespaceID   string
+		serializer    serialization.Serializer
+		generator     test.Generator
+		once          sync.Once
 	}
 )
 
 func TestStreamBasedReplicationTestSuite(t *testing.T) {
+	t.Parallel()
 	for _, tc := range []struct {
 		name                    string
 		enableTransitionHistory bool
@@ -91,9 +91,9 @@ func TestStreamBasedReplicationTestSuite(t *testing.T) {
 	} {
 		t.Run(tc.name, func(t *testing.T) {
 			s := &streamBasedReplicationTestSuite{
-				namespaceName:           "replication-test-" + common.GenerateRandomString(5),
-				enableTransitionHistory: tc.enableTransitionHistory,
+				namespaceName: "replication-test-" + common.GenerateRandomString(5),
 			}
+			s.enableTransitionHistory = tc.enableTransitionHistory
 			suite.Run(t, s)
 		})
 	}
@@ -106,7 +106,6 @@ func (s *streamBasedReplicationTestSuite) SetupSuite() {
 		dynamicconfig.EnableEagerNamespaceRefresher.Key():       true,
 		dynamicconfig.EnableReplicationTaskBatching.Key():       true,
 		dynamicconfig.EnableReplicateLocalGeneratedEvents.Key(): true,
-		dynamicconfig.EnableTransitionHistory.Key():             s.enableTransitionHistory,
 	}
 	s.logger = log.NewTestLogger()
 	s.serializer = serialization.NewSerializer()