Fix TestWorkflowTaskAndActivityTaskTimeoutsWorkflow test timed out after 35m0s

.github/workflows/repro-flaky-test.yml

@@ -0,0 +1,40 @@
+name: Reproduce Flaky Test
+
+on:
+  pull_request:
+  push:
+    branches:
+      - fix-flaky-task-timeouts-test
+
+env:
+  DOCKER_COMPOSE_FILE: ./develop/github/docker-compose.yml
+  TEMPORAL_VERSION_CHECK_DISABLED: 1
+
+jobs:
+  repro-flaky-test:
+    name: Repro flaky test (run ${{ matrix.run }})
+    runs-on: ubuntu-latest
+    strategy:
+      fail-fast: false
+      matrix:
+        run: [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20]
+    steps:
+      - uses: actions/checkout@v4
+
+      - uses: actions/setup-go@v5
+        with:
+          go-version-file: go.mod
+          cache: true
+
+      - run: make pre-build-functional-test-coverage
+
+      - name: Start dependencies
+        run: docker compose -f ${{ env.DOCKER_COMPOSE_FILE }} up -d
+
+      - name: Run flaky test
+        timeout-minutes: 40
+        run: |
+          make functional-test-coverage
+        env:
+          PERSISTENCE_DRIVER: sqlite
+          TEST_ARGS: "-run TestWorkflowTestSuite/TestWorkflowTaskAndActivityTaskTimeoutsWorkflow -count 1"

chasm/lib/activity/activity.go

@@ -1,6 +1,7 @@
 package activity
 
 import (
+	"errors"
 	"fmt"
 	"slices"
 	"time"
@@ -22,6 +23,7 @@ import (
 	"go.temporal.io/server/common/dynamicconfig"
 	"go.temporal.io/server/common/metrics"
 	"go.temporal.io/server/common/payload"
+	serviceerrors "go.temporal.io/server/common/serviceerror"
 	"go.temporal.io/server/common/tqid"
 	"google.golang.org/protobuf/types/known/durationpb"
 	"google.golang.org/protobuf/types/known/timestamppb"
@@ -41,9 +43,6 @@ var (
 var _ chasm.VisibilitySearchAttributesProvider = (*Activity)(nil)
 
 type ActivityStore interface {
-	// PopulateRecordStartedResponse populates the response for RecordActivityTaskStarted
-	PopulateRecordStartedResponse(ctx chasm.Context, key chasm.ExecutionKey, namespace string, response *historyservice.RecordActivityTaskStartedResponse) error
-
 	// RecordCompleted applies the provided function to record activity completion
 	RecordCompleted(ctx chasm.MutableContext, applyFn func(ctx chasm.MutableContext) error) error
 }
@@ -187,58 +186,66 @@ func (a *Activity) createAddActivityTaskRequest(ctx chasm.Context, namespaceID s
 		TaskQueue:              a.GetTaskQueue(),
 		Priority:               a.GetPriority(),
 		ComponentRef:           componentRef,
+		Stamp:                  a.LastAttempt.Get(ctx).GetStamp(),
 	}, nil
 }
 
 // HandleStarted updates the activity on recording activity task started and populates the response.
 func (a *Activity) HandleStarted(ctx chasm.MutableContext, request *historyservice.RecordActivityTaskStartedRequest) (
 	*historyservice.RecordActivityTaskStartedResponse, error,
 ) {
+	lastAttempt := a.LastAttempt.Get(ctx)
+	// If already started, return existing response if request ID matches to make retry idempotent, else error.
+	if a.StateMachineState() == activitypb.ACTIVITY_EXECUTION_STATUS_STARTED && request.GetRequestId() == lastAttempt.GetStartRequestId() {
+		return a.GenerateRecordActivityTaskStartedResponse(ctx, request.GetPollRequest().GetNamespace())
+	}
+	if lastAttempt.GetStamp() != request.GetStamp() {
+		return nil, serviceerrors.NewObsoleteMatchingTask("activity attempt stamp mismatch")
+	}
 	if err := TransitionStarted.Apply(a, ctx, request); err != nil {
+		if errors.Is(err, chasm.ErrInvalidTransition) {
+			return nil, serviceerrors.NewObsoleteMatchingTask(err.Error())
+		}
 		return nil, err
 	}
-	response := &historyservice.RecordActivityTaskStartedResponse{}
-	err := a.StoreOrSelf(ctx).PopulateRecordStartedResponse(ctx, ctx.ExecutionKey(), request.GetPollRequest().GetNamespace(), response)
-	return response, err
+	return a.GenerateRecordActivityTaskStartedResponse(ctx, request.GetPollRequest().GetNamespace())
 }
 
-// PopulateRecordStartedResponse populates the response for HandleStarted.
-func (a *Activity) PopulateRecordStartedResponse(
+// GenerateRecordActivityTaskStartedResponse generates the response for HandleStarted.
+func (a *Activity) GenerateRecordActivityTaskStartedResponse(
 	ctx chasm.Context,
-	key chasm.ExecutionKey,
 	namespace string,
-	response *historyservice.RecordActivityTaskStartedResponse,
-) error {
+) (*historyservice.RecordActivityTaskStartedResponse, error) {
+	key := ctx.ExecutionKey()
 	lastHeartbeat, _ := a.LastHeartbeat.TryGet(ctx)
-	if lastHeartbeat != nil {
-		response.HeartbeatDetails = lastHeartbeat.GetDetails()
-	}
 	requestData := a.RequestData.Get(ctx)
 	attempt := a.LastAttempt.Get(ctx)
-	response.StartedTime = attempt.GetStartedTime()
-	response.Attempt = attempt.GetCount()
-	response.Priority = a.GetPriority()
-	response.RetryPolicy = a.GetRetryPolicy()
-	response.ActivityRunId = key.RunID
-	response.WorkflowNamespace = namespace
-	response.ScheduledEvent = &historypb.HistoryEvent{
-		EventType: enumspb.EVENT_TYPE_ACTIVITY_TASK_SCHEDULED,
-		EventTime: a.GetScheduleTime(),
-		Attributes: &historypb.HistoryEvent_ActivityTaskScheduledEventAttributes{
-			ActivityTaskScheduledEventAttributes: &historypb.ActivityTaskScheduledEventAttributes{
-				ActivityId:             key.BusinessID,
-				ActivityType:           a.GetActivityType(),
-				Input:                  requestData.GetInput(),
-				Header:                 requestData.GetHeader(),
-				TaskQueue:              a.GetTaskQueue(),
-				ScheduleToCloseTimeout: a.GetScheduleToCloseTimeout(),
-				ScheduleToStartTimeout: a.GetScheduleToStartTimeout(),
-				StartToCloseTimeout:    a.GetStartToCloseTimeout(),
-				HeartbeatTimeout:       a.GetHeartbeatTimeout(),
+	return &historyservice.RecordActivityTaskStartedResponse{
+		StartedTime:       attempt.GetStartedTime(),
+		Attempt:           attempt.GetCount(),
+		Priority:          a.GetPriority(),
+		RetryPolicy:       a.GetRetryPolicy(),
+		ActivityRunId:     key.RunID,
+		WorkflowNamespace: namespace,
+		HeartbeatDetails:  lastHeartbeat.GetDetails(),
+		ScheduledEvent: &historypb.HistoryEvent{
+			EventType: enumspb.EVENT_TYPE_ACTIVITY_TASK_SCHEDULED,
+			EventTime: a.GetScheduleTime(),
+			Attributes: &historypb.HistoryEvent_ActivityTaskScheduledEventAttributes{
+				ActivityTaskScheduledEventAttributes: &historypb.ActivityTaskScheduledEventAttributes{
+					ActivityId:             key.BusinessID,
+					ActivityType:           a.GetActivityType(),
+					Input:                  requestData.GetInput(),
+					Header:                 requestData.GetHeader(),
+					TaskQueue:              a.GetTaskQueue(),
+					ScheduleToCloseTimeout: a.GetScheduleToCloseTimeout(),
+					ScheduleToStartTimeout: a.GetScheduleToStartTimeout(),
+					StartToCloseTimeout:    a.GetStartToCloseTimeout(),
+					HeartbeatTimeout:       a.GetHeartbeatTimeout(),
+				},
 			},
 		},
-	}
-	return nil
+	}, nil
 }
 
 // RecordCompleted applies the provided function to record activity completion.
@@ -571,7 +578,7 @@ func (a *Activity) RecordHeartbeat(
 			ScheduledTime: ctx.Now(a).Add(a.GetHeartbeatTimeout().AsDuration()),
 		},
 		&activitypb.HeartbeatTimeoutTask{
-			Attempt: a.LastAttempt.Get(ctx).GetCount(),
+			Stamp: a.LastAttempt.Get(ctx).GetStamp(),
 		},
 	)
 	return &historyservice.RecordActivityTaskHeartbeatResponse{

chasm/lib/activity/activity_tasks.go

@@ -37,7 +37,7 @@ func (e *activityDispatchTaskExecutor) Validate(
 ) (bool, error) {
 	// TODO(saa-preview): make sure we handle resets when we support them, as they will reset the attempt count
 	return (TransitionStarted.Possible(activity) &&
-		task.Attempt == activity.LastAttempt.Get(ctx).GetCount()), nil
+		task.Stamp == activity.LastAttempt.Get(ctx).GetStamp()), nil
 }
 
 func (e *activityDispatchTaskExecutor) Execute(
@@ -86,7 +86,7 @@ func (e *scheduleToStartTimeoutTaskExecutor) Validate(
 	task *activitypb.ScheduleToStartTimeoutTask,
 ) (bool, error) {
 	return (activity.Status == activitypb.ACTIVITY_EXECUTION_STATUS_SCHEDULED &&
-		task.Attempt == activity.LastAttempt.Get(ctx).GetCount()), nil
+		task.Stamp == activity.LastAttempt.Get(ctx).GetStamp()), nil
 }
 
 func (e *scheduleToStartTimeoutTaskExecutor) Execute(
@@ -181,7 +181,7 @@ func (e *startToCloseTimeoutTaskExecutor) Validate(
 	task *activitypb.StartToCloseTimeoutTask,
 ) (bool, error) {
 	valid := (activity.Status == activitypb.ACTIVITY_EXECUTION_STATUS_STARTED &&
-		task.Attempt == activity.LastAttempt.Get(ctx).GetCount())
+		task.Stamp == activity.LastAttempt.Get(ctx).GetStamp())
 	return valid, nil
 }
 
@@ -258,7 +258,7 @@ func (e *heartbeatTimeoutTaskExecutor) Validate(
 	}
 	// Task attempt must still match current attempt.
 	attempt := activity.LastAttempt.Get(ctx)
-	if attempt.GetCount() != task.Attempt {
+	if attempt.GetStamp() != task.Stamp {
 		return false, nil
 	}
 

chasm/lib/activity/activity_test.go

@@ -0,0 +1,156 @@
+package activity
+
+import (
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+	commonpb "go.temporal.io/api/common/v1"
+	taskqueuepb "go.temporal.io/api/taskqueue/v1"
+	"go.temporal.io/server/api/historyservice/v1"
+	"go.temporal.io/server/chasm"
+	"go.temporal.io/server/chasm/lib/activity/gen/activitypb/v1"
+	serviceerrors "go.temporal.io/server/common/serviceerror"
+	"google.golang.org/protobuf/types/known/durationpb"
+	"google.golang.org/protobuf/types/known/timestamppb"
+)
+
+func TestHandleStarted(t *testing.T) {
+	testTime := time.Date(2000, 1, 1, 0, 0, 0, 0, time.UTC)
+	testRequestID := "test-request-id"
+	testStamp := int32(1)
+
+	testCases := []struct {
+		name           string
+		activityStatus activitypb.ActivityExecutionStatus
+		attemptStamp   int32
+		requestStamp   int32
+		startRequestID string
+		requestID      string
+		checkOutcome   func(t *testing.T, response *historyservice.RecordActivityTaskStartedResponse, err error)
+	}{
+		{
+			name:           "successful transition from scheduled",
+			activityStatus: activitypb.ACTIVITY_EXECUTION_STATUS_SCHEDULED,
+			attemptStamp:   testStamp,
+			requestStamp:   testStamp,
+			requestID:      testRequestID,
+			checkOutcome: func(t *testing.T, response *historyservice.RecordActivityTaskStartedResponse, err error) {
+				require.Equal(t, int32(1), response.Attempt)
+				require.NoError(t, err)
+			},
+		},
+		{
+			name:           "idempotent retry - same request ID",
+			activityStatus: activitypb.ACTIVITY_EXECUTION_STATUS_STARTED,
+			attemptStamp:   testStamp,
+			requestStamp:   testStamp,
+			startRequestID: testRequestID,
+			requestID:      testRequestID,
+			checkOutcome: func(t *testing.T, response *historyservice.RecordActivityTaskStartedResponse, err error) {
+				require.Equal(t, int32(1), response.Attempt)
+				require.NoError(t, err)
+			},
+		},
+		{
+			name:           "error - already started with different request ID",
+			activityStatus: activitypb.ACTIVITY_EXECUTION_STATUS_STARTED,
+			attemptStamp:   testStamp,
+			requestStamp:   testStamp,
+			startRequestID: "different-request-id",
+			requestID:      testRequestID,
+			checkOutcome: func(t *testing.T, response *historyservice.RecordActivityTaskStartedResponse, err error) {
+				require.ErrorAs(t, err, new(*serviceerrors.ObsoleteMatchingTask))
+			},
+		},
+		{
+			name:           "error - stamp mismatch",
+			activityStatus: activitypb.ACTIVITY_EXECUTION_STATUS_SCHEDULED,
+			attemptStamp:   testStamp,
+			requestStamp:   testStamp + 1,
+			requestID:      testRequestID,
+			checkOutcome: func(t *testing.T, response *historyservice.RecordActivityTaskStartedResponse, err error) {
+				require.ErrorAs(t, err, new(*serviceerrors.ObsoleteMatchingTask))
+			},
+		},
+		{
+			name:           "error - invalid transition from completed",
+			activityStatus: activitypb.ACTIVITY_EXECUTION_STATUS_COMPLETED,
+			attemptStamp:   testStamp,
+			requestStamp:   testStamp,
+			requestID:      testRequestID,
+			checkOutcome: func(t *testing.T, response *historyservice.RecordActivityTaskStartedResponse, err error) {
+				require.ErrorAs(t, err, new(*serviceerrors.ObsoleteMatchingTask))
+			},
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			// Setup mock context
+			ctx := &chasm.MockMutableContext{
+				MockContext: chasm.MockContext{
+					HandleNow: func(chasm.Component) time.Time { return testTime },
+					HandleExecutionKey: func() chasm.ExecutionKey {
+						return chasm.ExecutionKey{
+							BusinessID: "test-activity-id",
+							RunID:      "test-run-id",
+						}
+					},
+				},
+			}
+
+			// Setup activity state
+			attemptState := &activitypb.ActivityAttemptState{
+				Count:          1,
+				Stamp:          tc.attemptStamp,
+				StartRequestId: tc.startRequestID,
+			}
+			if tc.activityStatus == activitypb.ACTIVITY_EXECUTION_STATUS_STARTED {
+				attemptState.StartedTime = timestamppb.New(testTime.Add(-1 * time.Minute))
+			}
+
+			// Determine heartbeat timeout based on test case
+			heartbeatTimeout := 1 * time.Minute
+			if tc.name == "successful transition without heartbeat timeout" {
+				heartbeatTimeout = 0
+			}
+
+			activity := &Activity{
+				ActivityState: &activitypb.ActivityState{
+					ActivityType:           &commonpb.ActivityType{Name: "test-activity-type"},
+					Status:                 tc.activityStatus,
+					TaskQueue:              &taskqueuepb.TaskQueue{Name: "test-task-queue"},
+					ScheduleToCloseTimeout: durationpb.New(10 * time.Minute),
+					ScheduleToStartTimeout: durationpb.New(2 * time.Minute),
+					StartToCloseTimeout:    durationpb.New(3 * time.Minute),
+					HeartbeatTimeout:       durationpb.New(heartbeatTimeout),
+					ScheduleTime:           timestamppb.New(testTime.Add(-30 * time.Second)),
+				},
+				LastAttempt: chasm.NewDataField(ctx, attemptState),
+				RequestData: chasm.NewDataField(ctx, &activitypb.ActivityRequestData{
+					Input: &commonpb.Payloads{
+						Payloads: []*commonpb.Payload{{Data: []byte("test-input")}},
+					},
+					Header: &commonpb.Header{
+						Fields: map[string]*commonpb.Payload{
+							"test-header": {Data: []byte("test-value")},
+						},
+					},
+				}),
+				Outcome: chasm.NewDataField(ctx, &activitypb.ActivityOutcome{}),
+			}
+
+			// Create request
+			request := &historyservice.RecordActivityTaskStartedRequest{
+				Stamp:     tc.requestStamp,
+				RequestId: tc.requestID,
+			}
+
+			// Execute HandleStarted
+			response, err := activity.HandleStarted(ctx, request)
+
+			tc.checkOutcome(t, response, err)
+		})
+	}
+}