Merge branch 'kannan/activity-cancel/persist-worker-key' into kannan/activity-cancel/task-definition

Author: rkannan82

api/persistence/v1/executions.pb.go

@@ -641,7 +641,7 @@ message ActivityInfo {
 
     int64 start_version = 50;
 
-    // Worker instance key of the worker executing this activity.
+    // Unique identifier of the worker that is this activity.
     string worker_instance_key = 51;
 }
 

proto/internal/temporal/server/api/persistence/v1/executions.proto

@@ -88,18 +88,18 @@ func Invoke(
 			// we need to force complete an activity
 			fabricateStartedEvent = ai.StartedEventId == common.EmptyEventID
 			if fabricateStartedEvent {
-			_, err := mutableState.AddActivityTaskStartedEvent(
-				ai,
-				scheduledEventID,
-				"",
-				req.GetCompleteRequest().GetIdentity(),
-				nil,
-				nil,
-				// TODO (shahab): do we need to do anything with wf redirect in this case or any
-				// other case where an activity starts?
-				nil,
-				"", // workerInstanceKey not available for force complete
-			)
+				_, err := mutableState.AddActivityTaskStartedEvent(
+					ai,
+					scheduledEventID,
+					"",
+					req.GetCompleteRequest().GetIdentity(),
+					nil,
+					nil,
+					// TODO (shahab): do we need to do anything with wf redirect in this case or any
+					// other case where an activity starts?
+					nil,
+					"", // workerInstanceKey not available for force complete
+				)
 				if err != nil {
 					return nil, err
 				}

service/history/api/respondactivitytaskcompleted/api.go

@@ -6112,3 +6112,76 @@ func (s *mutableStateSuite) TestCHASMNodeSize() {
 	expectedTotalSize += len(newNodeKey) + newNode.Size()
 	s.Equal(expectedTotalSize, mutableState.GetApproximatePersistedSize())
 }
+
+func (s *mutableStateSuite) TestAddActivityTaskStartedEventStoresWorkerInstanceKey() {
+	s.mockEventsCache.EXPECT().PutEvent(gomock.Any(), gomock.Any()).AnyTimes()
+
+	// Setup workflow execution
+	_, err := s.mutableState.AddWorkflowExecutionStartedEvent(
+		&commonpb.WorkflowExecution{WorkflowId: tests.WorkflowID, RunId: tests.RunID},
+		&historyservice.StartWorkflowExecutionRequest{
+			NamespaceId: tests.NamespaceID.String(),
+			StartRequest: &workflowservice.StartWorkflowExecutionRequest{
+				WorkflowType:        &commonpb.WorkflowType{Name: "workflow-type"},
+				TaskQueue:           &taskqueuepb.TaskQueue{Name: "task-queue"},
+				WorkflowRunTimeout:  durationpb.New(200 * time.Second),
+				WorkflowTaskTimeout: durationpb.New(1 * time.Second),
+			},
+		},
+	)
+	s.NoError(err)
+
+	di, err := s.mutableState.AddWorkflowTaskScheduledEvent(false, enumsspb.WORKFLOW_TASK_TYPE_NORMAL)
+	s.NoError(err)
+	_, _, err = s.mutableState.AddWorkflowTaskStartedEvent(
+		di.ScheduledEventID,
+		di.RequestID,
+		di.TaskQueue,
+		"identity",
+		nil,
+		nil,
+		nil,
+		false,
+		nil,
+	)
+	s.NoError(err)
+	_, err = s.mutableState.AddWorkflowTaskCompletedEvent(
+		di,
+		&workflowservice.RespondWorkflowTaskCompletedRequest{Identity: "identity"},
+		workflowTaskCompletionLimits,
+	)
+	s.NoError(err)
+
+	// Schedule activity
+	workflowTaskCompletedEventID := int64(4)
+	_, activityInfo, err := s.mutableState.AddActivityTaskScheduledEvent(
+		workflowTaskCompletedEventID,
+		&commandpb.ScheduleActivityTaskCommandAttributes{
+			ActivityId:   "test-activity-1",
+			ActivityType: &commonpb.ActivityType{Name: "test-activity-type"},
+			TaskQueue:    &taskqueuepb.TaskQueue{Name: "test-task-queue"},
+		},
+		false,
+	)
+	s.NoError(err)
+	s.Empty(activityInfo.WorkerInstanceKey, "WorkerInstanceKey should be empty before activity starts")
+
+	// Start activity with workerInstanceKey
+	expectedWorkerInstanceKey := "test-worker-instance-key-12345"
+	_, err = s.mutableState.AddActivityTaskStartedEvent(
+		activityInfo,
+		activityInfo.ScheduledEventId,
+		uuid.NewString(),
+		"worker-identity",
+		nil,
+		nil,
+		nil,
+		expectedWorkerInstanceKey,
+	)
+	s.NoError(err)
+
+	// Verify workerInstanceKey is stored
+	updatedActivityInfo, ok := s.mutableState.GetActivityInfo(activityInfo.ScheduledEventId)
+	s.True(ok)
+	s.Equal(expectedWorkerInstanceKey, updatedActivityInfo.WorkerInstanceKey)
+}