Adding fix for Issue 503

Author: saad1998

service/history/api/respondworkflowtaskcompleted/workflow_task_completed_handler.go

@@ -725,6 +725,10 @@ func (handler *workflowTaskCompletedHandler) handleCommandCompleteWorkflow(
 		newExecutionRunID = uuid.NewString()
 	}
 
+	if err := handler.mutableState.FlushPendingActivityEventsForCompletion(); err != nil {
+		return nil, err
+	}
+
 	// Always add workflow completed event to this one
 	event, err := handler.mutableState.AddCompletedWorkflowEvent(handler.workflowTaskCompletedID, attr, newExecutionRunID)
 	if err != nil {
@@ -1024,6 +1028,10 @@ func (handler *workflowTaskCompletedHandler) handleCommandContinueAsNewWorkflow(
 		}
 	}
 
+	if err := handler.mutableState.FlushPendingActivityEventsForCompletion(); err != nil {
+		return nil, err
+	}
+
 	event, newMutableState, err := handler.mutableState.AddContinueAsNewEvent(
 		ctx,
 		handler.workflowTaskCompletedID,

service/history/api/respondworkflowtaskcompleted/workflow_task_completed_handler_test.go

@@ -186,8 +186,9 @@ func TestCommandProtocolMessage(t *testing.T) {
 			CompleteWorkflowExecutionCommandAttributes: completeWorkflowExecutionCommandAttributes,
 		}
 
-		// mock a failed event creation.
+		// mock a failed event creation (handler calls FlushPendingActivityEventsForCompletion before AddCompletedWorkflowEvent).
 		event := &historypb.HistoryEvent{}
+		tc.ms.EXPECT().FlushPendingActivityEventsForCompletion().Return(nil)
 		tc.ms.EXPECT().AddCompletedWorkflowEvent(tc.handler.workflowTaskCompletedID, completeWorkflowExecutionCommandAttributes, "").MaxTimes(1).Return(event, fmt.Errorf("FAIL"))
 		tc.ms.EXPECT().GetExecutionInfo().AnyTimes().Return(&persistencespb.WorkflowExecutionInfo{})
 		tc.ms.EXPECT().GetExecutionState().AnyTimes().Return(&persistencespb.WorkflowExecutionState{})

service/history/interfaces/mutable_state.go

@@ -67,6 +67,7 @@ type (
 		AddChildWorkflowExecutionTimedOutEvent(int64, *commonpb.WorkflowExecution, *historypb.WorkflowExecutionTimedOutEventAttributes) (*historypb.HistoryEvent, error)
 		AddCompletedWorkflowEvent(int64, *commandpb.CompleteWorkflowExecutionCommandAttributes, string) (*historypb.HistoryEvent, error)
 		AddContinueAsNewEvent(context.Context, int64, int64, namespace.Name, *commandpb.ContinueAsNewWorkflowExecutionCommandAttributes, worker_versioning.IsWFTaskQueueInVersionDetector) (*historypb.HistoryEvent, MutableState, error)
+		FlushPendingActivityEventsForCompletion() error
 		AddWorkflowTaskCompletedEvent(*WorkflowTaskInfo, *workflowservice.RespondWorkflowTaskCompletedRequest, WorkflowTaskCompletionLimits) (*historypb.HistoryEvent, error)
 		AddWorkflowTaskFailedEvent(workflowTask *WorkflowTaskInfo, cause enumspb.WorkflowTaskFailedCause, failure *failurepb.Failure, identity string, versioningStamp *commonpb.WorkerVersionStamp, binChecksum, baseRunID, newRunID string, forkEventVersion int64) (*historypb.HistoryEvent, error)
 		AddWorkflowTaskScheduleToStartTimeoutEvent(workflowTask *WorkflowTaskInfo) (*historypb.HistoryEvent, error)

service/history/interfaces/mutable_state_mock.go

@@ -4048,6 +4048,71 @@ func (ms *MutableStateImpl) addStartedEventForTransientActivity(
 	return ms.ApplyActivityTaskStartedEvent(event)
 }
 
+// FlushPendingActivityEventsForCompletion records ActivityTaskStarted and (where applicable)
+// ActivityTaskFailed for pending activities that are in retry, so history is complete before
+// the workflow completes or continues-as-new. See issue #503.
+func (ms *MutableStateImpl) FlushPendingActivityEventsForCompletion() error {
+	opTag := tag.WorkflowActionWorkflowCompleted
+	if err := ms.checkMutability(opTag); err != nil {
+		return err
+	}
+
+	pending := ms.GetPendingActivityInfos()
+	if len(pending) == 0 {
+		return nil
+	}
+
+	scheduledEventIDs := make([]int64, 0, len(pending))
+	for id := range pending {
+		scheduledEventIDs = append(scheduledEventIDs, id)
+	}
+
+	for _, scheduledEventID := range scheduledEventIDs {
+		ai, ok := ms.GetActivityInfo(scheduledEventID)
+		if !ok || !ai.HasRetryPolicy {
+			continue
+		}
+
+		if ai.RetryLastFailure != nil && ai.StartedEventId != common.TransientEventID {
+			failedAttempt := ai.Attempt - 1
+			if failedAttempt < 1 {
+				failedAttempt = 1
+			}
+			startedEvent := ms.hBuilder.AddActivityTaskStartedEvent(
+				scheduledEventID,
+				failedAttempt,
+				"",
+				ai.RetryLastWorkerIdentity,
+				ai.RetryLastFailure,
+				nil,
+				0,
+			)
+			if err := ms.ApplyActivityTaskStartedEvent(startedEvent); err != nil {
+				return err
+			}
+			if _, err := ms.AddActivityTaskFailedEvent(
+				scheduledEventID,
+				startedEvent.GetEventId(),
+				ai.RetryLastFailure,
+				enumspb.RETRY_STATE_IN_PROGRESS,
+				ai.RetryLastWorkerIdentity,
+				nil,
+			); err != nil {
+				return err
+			}
+			continue
+		}
+
+		if ai.StartedEventId == common.TransientEventID {
+			if err := ms.addStartedEventForTransientActivity(scheduledEventID, nil); err != nil {
+				return err
+			}
+		}
+	}
+
+	return nil
+}
+
 func (ms *MutableStateImpl) AddActivityTaskStartedEvent(
 	ai *persistencespb.ActivityInfo,
 	scheduledEventID int64,

service/history/workflow/mutable_state_impl.go

@@ -2959,6 +2959,114 @@ func (s *mutableStateSuite) TestRetryActivity_PausedIncrementsStamp() {
 	s.Equal(int32(2), updatedActivityInfo.Attempt)
 }
 
+func (s *mutableStateSuite) TestFlushPendingActivityEventsForCompletion() {
+	s.mockEventsCache.EXPECT().PutEvent(gomock.Any(), gomock.Any()).AnyTimes()
+
+	tq := &taskqueuepb.TaskQueue{Name: "tq"}
+	retryPolicy := &commonpb.RetryPolicy{
+		InitialInterval: durationpb.New(time.Second),
+	}
+
+	s.Run("no_pending_activities", func() {
+		s.createVersionedMutableStateWithCompletedWFT(tq)
+		nextBefore := s.mutableState.GetNextEventID()
+		err := s.mutableState.FlushPendingActivityEventsForCompletion()
+		s.NoError(err)
+		s.Equal(nextBefore, s.mutableState.GetNextEventID())
+	})
+
+	s.Run("pending_activity_with_transient_started_flushes_started", func() {
+		s.createVersionedMutableStateWithCompletedWFT(tq)
+		workflowTaskCompletedEventID := s.mutableState.GetNextEventID() - 1
+		_, activityInfo, err := s.mutableState.AddActivityTaskScheduledEvent(
+			workflowTaskCompletedEventID,
+			&commandpb.ScheduleActivityTaskCommandAttributes{
+				ActivityId:   "flush-started",
+				ActivityType: &commonpb.ActivityType{Name: "activity-type"},
+				TaskQueue:    tq,
+				RetryPolicy:  retryPolicy,
+			},
+			false,
+		)
+		s.NoError(err)
+		_, err = s.mutableState.AddActivityTaskStartedEvent(
+			activityInfo,
+			activityInfo.ScheduledEventId,
+			uuid.NewString(),
+			"worker",
+			nil,
+			nil,
+			nil,
+		)
+		s.NoError(err)
+		ai, ok := s.mutableState.GetActivityInfo(activityInfo.ScheduledEventId)
+		s.True(ok)
+		s.Equal(common.TransientEventID, ai.StartedEventId)
+
+		err = s.mutableState.FlushPendingActivityEventsForCompletion()
+		s.NoError(err)
+		ai, ok = s.mutableState.GetActivityInfo(activityInfo.ScheduledEventId)
+		s.True(ok)
+		s.NotEqual(common.TransientEventID, ai.StartedEventId, "flush should have written ActivityTaskStarted so StartedEventId is no longer transient")
+	})
+
+	s.Run("pending_activity_between_retries_flushes_started_and_failed", func() {
+		s.createVersionedMutableStateWithCompletedWFT(tq)
+		workflowTaskCompletedEventID := s.mutableState.GetNextEventID() - 1
+		_, activityInfo, err := s.mutableState.AddActivityTaskScheduledEvent(
+			workflowTaskCompletedEventID,
+			&commandpb.ScheduleActivityTaskCommandAttributes{
+				ActivityId:   "flush-failed",
+				ActivityType: &commonpb.ActivityType{Name: "activity-type"},
+				TaskQueue:    tq,
+				RetryPolicy:  retryPolicy,
+			},
+			false,
+		)
+		s.NoError(err)
+		_, err = s.mutableState.AddActivityTaskStartedEvent(
+			activityInfo,
+			activityInfo.ScheduledEventId,
+			uuid.NewString(),
+			"worker",
+			nil,
+			nil,
+			nil,
+		)
+		s.NoError(err)
+		_, err = s.mutableState.RetryActivity(activityInfo, &failurepb.Failure{Message: "attempt failed"})
+		s.NoError(err)
+		s.Len(s.mutableState.GetPendingActivityInfos(), 1)
+
+		err = s.mutableState.FlushPendingActivityEventsForCompletion()
+		s.NoError(err)
+		s.Len(s.mutableState.GetPendingActivityInfos(), 0)
+		_, ok := s.mutableState.GetActivityInfo(activityInfo.ScheduledEventId)
+		s.False(ok)
+	})
+
+	s.Run("pending_activity_without_retry_skipped", func() {
+		s.createVersionedMutableStateWithCompletedWFT(tq)
+		workflowTaskCompletedEventID := s.mutableState.GetNextEventID() - 1
+		_, _, err := s.mutableState.AddActivityTaskScheduledEvent(
+			workflowTaskCompletedEventID,
+			&commandpb.ScheduleActivityTaskCommandAttributes{
+				ActivityId:   "no-retry",
+				ActivityType: &commonpb.ActivityType{Name: "activity-type"},
+				TaskQueue:    tq,
+			},
+			false,
+		)
+		s.NoError(err)
+		s.Len(s.mutableState.GetPendingActivityInfos(), 1)
+		nextBefore := s.mutableState.GetNextEventID()
+		err = s.mutableState.FlushPendingActivityEventsForCompletion()
+		s.NoError(err)
+		s.Equal(nextBefore, s.mutableState.GetNextEventID())
+		s.Len(s.mutableState.GetPendingActivityInfos(), 1)
+	})
+}
+
 func (s *mutableStateSuite) TestupdateBuildIdsAndDeploymentSearchAttributes() {
 	versioned := func(buildId string) *commonpb.WorkerVersionStamp {
 		return &commonpb.WorkerVersionStamp{BuildId: buildId, UseVersioning: true}