Revert "Remove timeouts during activity retry (#8509)" (#8526)

This reverts commit 051fb08.

## What changed?
Revert activity timeout change

## Why?
Reverting activity timeout timer improvment as it might have caused lost
heartbeat timer.

## How did you test it?
- [x] built
- [ ] run locally and tested manually
- [ ] covered by existing tests
- [ ] added new unit test(s)
- [ ] added new functional test(s)

Author: prathyushpv

service/history/workflow/mutable_state_impl.go

@@ -5816,11 +5816,6 @@ func (ms *MutableStateImpl) RetryActivity(
 		ai.Attempt+1,
 		activityFailure)
 
-	// Delete old per-attempt timeout tasks since they're invalidated by the retry.
-	ms.deleteActivityTimeoutTask(ai.ScheduledEventId, enumspb.TIMEOUT_TYPE_SCHEDULE_TO_START, ai)
-	ms.deleteActivityTimeoutTask(ai.ScheduledEventId, enumspb.TIMEOUT_TYPE_START_TO_CLOSE, ai)
-	ms.deleteActivityTimeoutTask(ai.ScheduledEventId, enumspb.TIMEOUT_TYPE_HEARTBEAT, ai)
-
 	if err := ms.taskGenerator.GenerateActivityRetryTasks(ai); err != nil {
 		return enumspb.RETRY_STATE_INTERNAL_SERVER_ERROR, err
 	}
@@ -6261,12 +6256,6 @@ func (ms *MutableStateImpl) deleteActivityTimeoutTask(
 	case enumspb.TIMEOUT_TYPE_SCHEDULE_TO_START:
 		task = timeoutTasks.ScheduleToStartTimeoutTask
 		timeoutTasks.ScheduleToStartTimeoutTask = nil
-	case enumspb.TIMEOUT_TYPE_SCHEDULE_TO_CLOSE:
-		task = timeoutTasks.ScheduleToCloseTimeoutTask
-		timeoutTasks.ScheduleToCloseTimeoutTask = nil
-	case enumspb.TIMEOUT_TYPE_START_TO_CLOSE:
-		task = timeoutTasks.StartToCloseTimeoutTask
-		timeoutTasks.StartToCloseTimeoutTask = nil
 	case enumspb.TIMEOUT_TYPE_HEARTBEAT:
 		task = timeoutTasks.HeartbeatTimeoutTask
 		timeoutTasks.HeartbeatTimeoutTask = nil
@@ -6295,11 +6284,39 @@ func (ms *MutableStateImpl) recordActivityTimeoutTasksForDeletion(
 		return
 	}
 
+	timeoutTasks, exists := ms.activityTimeoutTasks[scheduledEventID]
+	if !exists {
+		return
+	}
+
 	// Delete all timeout tasks for this activity
-	ms.deleteActivityTimeoutTask(scheduledEventID, enumspb.TIMEOUT_TYPE_SCHEDULE_TO_START, activityInfo)
-	ms.deleteActivityTimeoutTask(scheduledEventID, enumspb.TIMEOUT_TYPE_SCHEDULE_TO_CLOSE, activityInfo)
-	ms.deleteActivityTimeoutTask(scheduledEventID, enumspb.TIMEOUT_TYPE_START_TO_CLOSE, activityInfo)
-	ms.deleteActivityTimeoutTask(scheduledEventID, enumspb.TIMEOUT_TYPE_HEARTBEAT, activityInfo)
+	if task := timeoutTasks.ScheduleToStartTimeoutTask; task != nil {
+		ms.BestEffortDeleteTasks[tasks.CategoryTimer] = append(
+			ms.BestEffortDeleteTasks[tasks.CategoryTimer],
+			task.GetKey(),
+		)
+	}
+
+	if task := timeoutTasks.ScheduleToCloseTimeoutTask; task != nil {
+		ms.BestEffortDeleteTasks[tasks.CategoryTimer] = append(
+			ms.BestEffortDeleteTasks[tasks.CategoryTimer],
+			task.GetKey(),
+		)
+	}
+
+	if task := timeoutTasks.StartToCloseTimeoutTask; task != nil {
+		ms.BestEffortDeleteTasks[tasks.CategoryTimer] = append(
+			ms.BestEffortDeleteTasks[tasks.CategoryTimer],
+			task.GetKey(),
+		)
+	}
+
+	if task := timeoutTasks.HeartbeatTimeoutTask; task != nil {
+		ms.BestEffortDeleteTasks[tasks.CategoryTimer] = append(
+			ms.BestEffortDeleteTasks[tasks.CategoryTimer],
+			task.GetKey(),
+		)
+	}
 
 	delete(ms.activityTimeoutTasks, scheduledEventID)
 }

service/history/workflow/mutable_state_impl_test.go

@@ -1159,116 +1159,6 @@ func (s *mutableStateSuite) TestActivityTimeoutDeletion_AllTimersDeletedOnComple
 	s.True(foundHeartbeat, "Heartbeat should be deleted on completion (if still exists)")
 }
 
-func (s *mutableStateSuite) TestActivityTimeoutDeletion_PerAttemptTimersDeletedOnRetry() {
-	// Test that when activity is retried, per-attempt timeout tasks (ScheduleToStart, StartToClose, Heartbeat)
-	// are deleted but ScheduleToClose is preserved since it spans all retry attempts.
-	s.mockEventsCache.EXPECT().PutEvent(gomock.Any(), gomock.Any()).AnyTimes()
-
-	// Schedule and start an activity with retry policy
-	_, activityInfo, err := s.mutableState.AddActivityTaskScheduledEvent(
-		int64(2), // workflowTaskCompletedEventID
-		&commandpb.ScheduleActivityTaskCommandAttributes{
-			ActivityId:   "activity",
-			ActivityType: &commonpb.ActivityType{Name: "type"},
-			TaskQueue:    &taskqueuepb.TaskQueue{Name: "tq"},
-			RetryPolicy:  &commonpb.RetryPolicy{InitialInterval: durationpb.New(time.Second)},
-		},
-		false,
-	)
-	s.NoError(err)
-
-	_, err = s.mutableState.AddActivityTaskStartedEvent(activityInfo, activityInfo.ScheduledEventId, "", "", nil, nil, nil)
-	s.NoError(err)
-
-	// Create mock timeout tasks for all types
-	workflowKey := s.mutableState.GetWorkflowKey()
-	scheduledEventID := activityInfo.ScheduledEventId
-
-	mockScheduleToCloseTask := &tasks.ActivityTimeoutTask{
-		WorkflowKey: workflowKey,
-		TaskID:      200,
-		TimeoutType: enumspb.TIMEOUT_TYPE_SCHEDULE_TO_CLOSE,
-		EventID:     scheduledEventID,
-		Stamp:       activityInfo.Stamp,
-	}
-	mockScheduleToStartTask := &tasks.ActivityTimeoutTask{
-		WorkflowKey: workflowKey,
-		TaskID:      201,
-		TimeoutType: enumspb.TIMEOUT_TYPE_SCHEDULE_TO_START,
-		EventID:     scheduledEventID,
-		Stamp:       activityInfo.Stamp,
-	}
-	mockStartToCloseTask := &tasks.ActivityTimeoutTask{
-		WorkflowKey: workflowKey,
-		TaskID:      202,
-		TimeoutType: enumspb.TIMEOUT_TYPE_START_TO_CLOSE,
-		EventID:     scheduledEventID,
-		Stamp:       activityInfo.Stamp,
-	}
-	mockHeartbeatTask := &tasks.ActivityTimeoutTask{
-		WorkflowKey: workflowKey,
-		TaskID:      203,
-		TimeoutType: enumspb.TIMEOUT_TYPE_HEARTBEAT,
-		EventID:     scheduledEventID,
-		Stamp:       activityInfo.Stamp,
-	}
-
-	// Store all tasks
-	if s.mutableState.activityTimeoutTasks == nil {
-		s.mutableState.activityTimeoutTasks = make(map[int64]*activityTimeoutTasks)
-	}
-	s.mutableState.activityTimeoutTasks[scheduledEventID] = &activityTimeoutTasks{
-		ScheduleToStartTimeoutTask: mockScheduleToStartTask,
-		ScheduleToCloseTimeoutTask: mockScheduleToCloseTask,
-		StartToCloseTimeoutTask:    mockStartToCloseTask,
-		HeartbeatTimeoutTask:       mockHeartbeatTask,
-	}
-
-	// Retry the activity
-	retryState, err := s.mutableState.RetryActivity(activityInfo, &failurepb.Failure{})
-	s.NoError(err)
-	s.Equal(enumspb.RETRY_STATE_IN_PROGRESS, retryState)
-
-	// Verify that per-attempt timeout tasks are in the deletion list
-	del := s.mutableState.BestEffortDeleteTasks
-	s.Contains(del, tasks.CategoryTimer)
-
-	foundScheduleToClose := false
-	foundStartToClose := false
-	foundScheduleToStart := false
-	foundHeartbeat := false
-	for _, key := range del[tasks.CategoryTimer] {
-		if key == mockScheduleToCloseTask.GetKey() {
-			foundScheduleToClose = true
-		}
-		if key == mockStartToCloseTask.GetKey() {
-			foundStartToClose = true
-		}
-		if key == mockScheduleToStartTask.GetKey() {
-			foundScheduleToStart = true
-		}
-		if key == mockHeartbeatTask.GetKey() {
-			foundHeartbeat = true
-		}
-	}
-
-	// Per-attempt timers should be deleted
-	s.True(foundStartToClose, "StartToClose should be deleted on retry")
-	s.True(foundScheduleToStart, "ScheduleToStart should be deleted on retry")
-	s.True(foundHeartbeat, "Heartbeat should be deleted on retry")
-
-	// ScheduleToClose should NOT be deleted (it spans all retry attempts)
-	s.False(foundScheduleToClose, "ScheduleToClose should NOT be deleted on retry (spans all attempts)")
-
-	// Verify the in-memory storage reflects the deletions
-	timeoutTasks := s.mutableState.activityTimeoutTasks[scheduledEventID]
-	s.NotNil(timeoutTasks, "Activity timeout tasks map entry should still exist after retry")
-	s.Nil(timeoutTasks.ScheduleToStartTimeoutTask, "ScheduleToStart task reference should be nil")
-	s.Nil(timeoutTasks.StartToCloseTimeoutTask, "StartToClose task reference should be nil")
-	s.Nil(timeoutTasks.HeartbeatTimeoutTask, "Heartbeat task reference should be nil")
-	s.NotNil(timeoutTasks.ScheduleToCloseTimeoutTask, "ScheduleToClose task reference should be preserved")
-}
-
 func (s *mutableStateSuite) TestDeleteScheduleToStartTimeout_BothPathsWhenActivityStarts() {
 	// Test that ScheduleToStart is deleted for both transient and non-transient activity starts
 	s.mockEventsCache.EXPECT().PutEvent(gomock.Any(), gomock.Any()).AnyTimes()