Refactor transient workflow task state machine (#3740)

alexshtin · web-flow · commit 661880f37a40 · 2022-12-21T00:02:21.000-08:00
diff --git a/service/history/ndc/branch_manager.go b/service/history/ndc/branch_manager.go
@@ -170,7 +170,7 @@ func (r *BranchMgrImpl) flushBufferedEvents(
 	// check whether there are buffered events, if so, flush it
 	// NOTE: buffered events does not show in version history or next event id
 	if !r.mutableState.HasBufferedEvents() {
-		if r.mutableState.HasTransientWorkflowTask() {
+		if r.mutableState.HasInFlightWorkflowTask() && r.mutableState.IsTransientWorkflowTask() {
 			if err := r.mutableState.ClearTransientWorkflowTask(); err != nil {
 				return nil, 0, err
 			}
diff --git a/service/history/ndc/branch_manager_test.go b/service/history/ndc/branch_manager_test.go
@@ -193,7 +193,8 @@ func (s *branchMgrSuite) TestClearTransientWorkflowTask() {
 
 	s.mockMutableState.EXPECT().GetLastWriteVersion().Return(lastWriteVersion, nil).AnyTimes()
 	s.mockMutableState.EXPECT().HasBufferedEvents().Return(false).AnyTimes()
-	s.mockMutableState.EXPECT().HasTransientWorkflowTask().Return(true).AnyTimes()
+	s.mockMutableState.EXPECT().HasInFlightWorkflowTask().Return(true).AnyTimes()
+	s.mockMutableState.EXPECT().IsTransientWorkflowTask().Return(true).AnyTimes()
 	s.mockMutableState.EXPECT().ClearTransientWorkflowTask().Return(nil).AnyTimes()
 
 	s.mockMutableState.EXPECT().GetExecutionInfo().Return(&persistencespb.WorkflowExecutionInfo{
@@ -285,7 +286,8 @@ func (s *branchMgrSuite) TestPrepareVersionHistory_BranchAppendable_NoMissingEve
 
 	s.mockMutableState.EXPECT().GetExecutionInfo().Return(&persistencespb.WorkflowExecutionInfo{VersionHistories: versionHistories}).AnyTimes()
 	s.mockMutableState.EXPECT().HasBufferedEvents().Return(false).AnyTimes()
-	s.mockMutableState.EXPECT().HasTransientWorkflowTask().Return(false).AnyTimes()
+	s.mockMutableState.EXPECT().HasInFlightWorkflowTask().Return(true).AnyTimes()
+	s.mockMutableState.EXPECT().IsTransientWorkflowTask().Return(false).AnyTimes()
 
 	doContinue, index, err := s.nDCBranchMgr.prepareVersionHistory(
 		context.Background(),
@@ -316,7 +318,8 @@ func (s *branchMgrSuite) TestPrepareVersionHistory_BranchAppendable_MissingEvent
 	s.NoError(err)
 
 	s.mockMutableState.EXPECT().HasBufferedEvents().Return(false).AnyTimes()
-	s.mockMutableState.EXPECT().HasTransientWorkflowTask().Return(false).AnyTimes()
+	s.mockMutableState.EXPECT().HasInFlightWorkflowTask().Return(true).AnyTimes()
+	s.mockMutableState.EXPECT().IsTransientWorkflowTask().Return(false).AnyTimes()
 	s.mockMutableState.EXPECT().GetExecutionInfo().Return(&persistencespb.WorkflowExecutionInfo{
 		NamespaceId:      s.namespaceID,
 		WorkflowId:       s.workflowID,
@@ -357,7 +360,8 @@ func (s *branchMgrSuite) TestPrepareVersionHistory_BranchNotAppendable_NoMissing
 	newBranchToken := []byte("some random new branch token")
 
 	s.mockMutableState.EXPECT().HasBufferedEvents().Return(false).AnyTimes()
-	s.mockMutableState.EXPECT().HasTransientWorkflowTask().Return(false).AnyTimes()
+	s.mockMutableState.EXPECT().HasInFlightWorkflowTask().Return(true).AnyTimes()
+	s.mockMutableState.EXPECT().IsTransientWorkflowTask().Return(false).AnyTimes()
 	s.mockMutableState.EXPECT().GetExecutionInfo().Return(&persistencespb.WorkflowExecutionInfo{
 		NamespaceId:      s.namespaceID,
 		WorkflowId:       s.workflowID,
@@ -416,7 +420,8 @@ func (s *branchMgrSuite) TestPrepareVersionHistory_BranchNotAppendable_MissingEv
 	})
 
 	s.mockMutableState.EXPECT().HasBufferedEvents().Return(false).AnyTimes()
-	s.mockMutableState.EXPECT().HasTransientWorkflowTask().Return(false).AnyTimes()
+	s.mockMutableState.EXPECT().HasInFlightWorkflowTask().Return(true).AnyTimes()
+	s.mockMutableState.EXPECT().IsTransientWorkflowTask().Return(false).AnyTimes()
 	s.mockMutableState.EXPECT().GetExecutionInfo().Return(&persistencespb.WorkflowExecutionInfo{
 		NamespaceId:      s.namespaceID,
 		WorkflowId:       s.workflowID,
diff --git a/service/history/workflow/mutable_state.go b/service/history/workflow/mutable_state.go
@@ -182,7 +182,7 @@ type (
 		GetWorkflowType() *commonpb.WorkflowType
 		GetWorkflowStateStatus() (enumsspb.WorkflowExecutionState, enumspb.WorkflowExecutionStatus)
 		GetQueryRegistry() QueryRegistry
-		HasTransientWorkflowTask() bool
+		IsTransientWorkflowTask() bool
 		ClearTransientWorkflowTask() error
 		HasBufferedEvents() bool
 		HasInFlightWorkflowTask() bool
diff --git a/service/history/workflow/mutable_state_impl.go b/service/history/workflow/mutable_state_impl.go
@@ -607,15 +607,15 @@ func (ms *MutableStateImpl) IsStickyTaskQueueEnabled() bool {
 	return ms.executionInfo.StickyTaskQueue != ""
 }
 
-func (e *MutableStateImpl) TaskQueue() *taskqueuepb.TaskQueue {
-	if e.IsStickyTaskQueueEnabled() {
+func (ms *MutableStateImpl) TaskQueue() *taskqueuepb.TaskQueue {
+	if ms.IsStickyTaskQueueEnabled() {
 		return &taskqueuepb.TaskQueue{
-			Name: e.executionInfo.StickyTaskQueue,
+			Name: ms.executionInfo.StickyTaskQueue,
 			Kind: enumspb.TASK_QUEUE_KIND_STICKY,
 		}
 	}
 	return &taskqueuepb.TaskQueue{
-		Name: e.executionInfo.TaskQueue,
+		Name: ms.executionInfo.TaskQueue,
 		Kind: enumspb.TASK_QUEUE_KIND_NORMAL,
 	}
 }
@@ -1294,25 +1294,18 @@ func (ms *MutableStateImpl) GetInFlightWorkflowTask() (*WorkflowTaskInfo, bool)
 	return ms.workflowTaskManager.GetInFlightWorkflowTask()
 }
 
-func (ms *MutableStateImpl) HasTransientWorkflowTask() bool {
-	workflowTask, ok := ms.GetInFlightWorkflowTask()
-	if !ok {
-		return false
-	}
-	return workflowTask.ScheduledEventID >= ms.GetNextEventID()
+func (ms *MutableStateImpl) IsTransientWorkflowTask() bool {
+	return ms.executionInfo.WorkflowTaskAttempt > 1
 }
 
 func (ms *MutableStateImpl) ClearTransientWorkflowTask() error {
-	workflowTask, ok := ms.GetInFlightWorkflowTask()
-	if !ok {
+	if !ms.HasInFlightWorkflowTask() {
 		return serviceerror.NewInternal("cannot clear transient workflow task when task is missing")
 	}
-
-	if workflowTask.ScheduledEventID < ms.GetNextEventID() {
+	if !ms.IsTransientWorkflowTask() {
 		return serviceerror.NewInternal("cannot clear transient workflow task when task is not transient")
 	}
-	// workflowTask.ScheduledEventID >= ms.GetNextEventID()
-	// this is transient workflow
+	// this is transient workflow task
 	if ms.HasBufferedEvents() {
 		return serviceerror.NewInternal("cannot clear transient workflow task when there are buffered events")
 	}
@@ -1791,6 +1784,9 @@ func (ms *MutableStateImpl) GetTransientWorkflowTaskInfo(
 	workflowTask *WorkflowTaskInfo,
 	identity string,
 ) *historyspb.TransientWorkflowTaskInfo {
+	if !ms.IsTransientWorkflowTask() {
+		return nil
+	}
 	return ms.workflowTaskManager.GetTransientWorkflowTaskInfo(workflowTask, identity)
 }
 
diff --git a/service/history/workflow/mutable_state_mock.go b/service/history/workflow/mutable_state_mock.go
diff --git a/service/history/workflow/workflow_task_state_machine.go b/service/history/workflow/workflow_task_state_machine.go

Original file line number	Diff line number	Diff line change
`@@ -170,7 +170,7 @@ func (r *BranchMgrImpl) flushBufferedEvents(`
`170`	`170`	`// check whether there are buffered events, if so, flush it`
`171`	`171`	`// NOTE: buffered events does not show in version history or next event id`
`172`	`172`	`if !r.mutableState.HasBufferedEvents() {`
`173`		`- if r.mutableState.HasTransientWorkflowTask() {`
	`173`	`+ if r.mutableState.HasInFlightWorkflowTask() && r.mutableState.IsTransientWorkflowTask() {`
`174`	`174`	`if err := r.mutableState.ClearTransientWorkflowTask(); err != nil {`
`175`	`175`	`return nil, 0, err`
`176`	`176`	`}`
Original file line number	Diff line number	Diff line change
`@@ -607,15 +607,15 @@ func (ms *MutableStateImpl) IsStickyTaskQueueEnabled() bool {`
`607`	`607`	`return ms.executionInfo.StickyTaskQueue != ""`
`608`	`608`	`}`
`609`	`609`
`610`		`-func (e MutableStateImpl) TaskQueue() taskqueuepb.TaskQueue {`
`611`		`- if e.IsStickyTaskQueueEnabled() {`
	`610`	`+func (ms MutableStateImpl) TaskQueue() taskqueuepb.TaskQueue {`
	`611`	`+ if ms.IsStickyTaskQueueEnabled() {`
`612`	`612`	`return &taskqueuepb.TaskQueue{`
`613`		`- Name: e.executionInfo.StickyTaskQueue,`
	`613`	`+ Name: ms.executionInfo.StickyTaskQueue,`
`614`	`614`	`Kind: enumspb.TASK_QUEUE_KIND_STICKY,`
`615`	`615`	`}`
`616`	`616`	`}`
`617`	`617`	`return &taskqueuepb.TaskQueue{`
`618`		`- Name: e.executionInfo.TaskQueue,`
	`618`	`+ Name: ms.executionInfo.TaskQueue,`
`619`	`619`	`Kind: enumspb.TASK_QUEUE_KIND_NORMAL,`
`620`	`620`	`}`
`621`	`621`	`}`
`@@ -1294,25 +1294,18 @@ func (ms MutableStateImpl) GetInFlightWorkflowTask() (WorkflowTaskInfo, bool)`
`1294`	`1294`	`return ms.workflowTaskManager.GetInFlightWorkflowTask()`
`1295`	`1295`	`}`
`1296`	`1296`
`1297`		`-func (ms *MutableStateImpl) HasTransientWorkflowTask() bool {`
`1298`		`- workflowTask, ok := ms.GetInFlightWorkflowTask()`
`1299`		`- if !ok {`
`1300`		`- return false`
`1301`		`- }`
`1302`		`- return workflowTask.ScheduledEventID >= ms.GetNextEventID()`
	`1297`	`+func (ms *MutableStateImpl) IsTransientWorkflowTask() bool {`
	`1298`	`+ return ms.executionInfo.WorkflowTaskAttempt > 1`
`1303`	`1299`	`}`
`1304`	`1300`
`1305`	`1301`	`func (ms *MutableStateImpl) ClearTransientWorkflowTask() error {`
`1306`		`- workflowTask, ok := ms.GetInFlightWorkflowTask()`
`1307`		`- if !ok {`
	`1302`	`+ if !ms.HasInFlightWorkflowTask() {`
`1308`	`1303`	`return serviceerror.NewInternal("cannot clear transient workflow task when task is missing")`
`1309`	`1304`	`}`
`1310`		`-`
`1311`		`- if workflowTask.ScheduledEventID < ms.GetNextEventID() {`
	`1305`	`+ if !ms.IsTransientWorkflowTask() {`
`1312`	`1306`	`return serviceerror.NewInternal("cannot clear transient workflow task when task is not transient")`
`1313`	`1307`	`}`
`1314`		`- // workflowTask.ScheduledEventID >= ms.GetNextEventID()`
`1315`		`- // this is transient workflow`
	`1308`	`+ // this is transient workflow task`
`1316`	`1309`	`if ms.HasBufferedEvents() {`
`1317`	`1310`	`return serviceerror.NewInternal("cannot clear transient workflow task when there are buffered events")`
`1318`	`1311`	`}`
`@@ -1791,6 +1784,9 @@ func (ms *MutableStateImpl) GetTransientWorkflowTaskInfo(`
`1791`	`1784`	`workflowTask *WorkflowTaskInfo,`
`1792`	`1785`	`identity string,`
`1793`	`1786`	`) *historyspb.TransientWorkflowTaskInfo {`
	`1787`	`+ if !ms.IsTransientWorkflowTask() {`
	`1788`	`+ return nil`
	`1789`	`+ }`
`1794`	`1790`	`return ms.workflowTaskManager.GetTransientWorkflowTaskInfo(workflowTask, identity)`
`1795`	`1791`	`}`
`1796`	`1792`