Validate workflow task start time when complete

temporalio · Jul 21, 2023 · 742f66f · 742f66f
1 parent fdb94a5
commit 742f66f
Show file tree

Hide file tree

Showing 8 changed files with 214 additions and 104 deletions.
diff --git a/api/token/v1/message.pb.go b/api/token/v1/message.pb.go
diff --git a/common/tasktoken/token.go b/common/tasktoken/token.go
@@ -25,6 +25,8 @@
 package tasktoken
 
 import (
+	"time"
+
 	v11 "go.temporal.io/server/api/clock/v1"
 	tokenspb "go.temporal.io/server/api/token/v1"
 )
@@ -35,6 +37,7 @@ func NewWorkflowTaskToken(
 	runID string,
 	scheduledEventID int64,
 	startedEventId int64,
+	startedTime *time.Time,
 	attempt int32,
 	clock *v11.VectorClock,
 	version int64,
@@ -45,6 +48,7 @@ func NewWorkflowTaskToken(
 		RunId:            runID,
 		ScheduledEventId: scheduledEventID,
 		StartedEventId:   startedEventId,
+		StartedTime:      startedTime,
 		Attempt:          attempt,
 		Clock:            clock,
 		Version:          version,

diff --git a/proto/internal/temporal/server/api/token/v1/message.proto b/proto/internal/temporal/server/api/token/v1/message.proto
@@ -24,6 +24,8 @@ package temporal.server.api.token.v1;
 
 option go_package = "go.temporal.io/server/api/token/v1;token";
 
+import "google/protobuf/timestamp.proto";
+import "dependencies/gogoproto/gogo.proto";
 import "temporal/server/api/clock/v1/message.proto";
 import "temporal/server/api/history/v1/message.proto";
 
@@ -63,6 +65,7 @@ message Task {
     temporal.server.api.clock.v1.VectorClock clock = 9;
     int64 started_event_id = 10;
     int64 version = 11;
+    google.protobuf.Timestamp started_time = 12 [(gogoproto.stdtime) = true];
 }
 
 message QueryTask {

diff --git a/service/frontend/workflow_handler.go b/service/frontend/workflow_handler.go
@@ -968,6 +968,7 @@ func (wh *WorkflowHandler) RespondWorkflowTaskCompleted(
 			taskToken.GetRunId(),
 			histResp.StartedResponse.GetScheduledEventId(),
 			histResp.StartedResponse.GetStartedEventId(),
+			histResp.StartedResponse.GetStartedTime(),
 			histResp.StartedResponse.GetAttempt(),
 			histResp.StartedResponse.GetClock(),
 			histResp.StartedResponse.GetVersion(),

diff --git a/service/history/api/startworkflow/api.go b/service/history/api/startworkflow/api.go
@@ -33,10 +33,10 @@ import (
 	historypb "go.temporal.io/api/history/v1"
 	"go.temporal.io/api/serviceerror"
 	"go.temporal.io/api/workflowservice/v1"
+	"go.temporal.io/server/api/historyservice/v1"
 
 	"go.temporal.io/server/common/tasktoken"
 
-	"go.temporal.io/server/api/historyservice/v1"
 	"go.temporal.io/server/common"
 	"go.temporal.io/server/common/definition"
 	"go.temporal.io/server/common/metrics"
@@ -543,6 +543,7 @@ func (s *Starter) generateResponse(
 		runID,
 		workflowTaskInfo.ScheduledEventID,
 		workflowTaskInfo.StartedEventID,
+		workflowTaskInfo.StartedTime,
 		workflowTaskInfo.Attempt,
 		clock,
 		workflowTaskInfo.Version,

diff --git a/service/history/workflow_task_handler_callbacks.go b/service/history/workflow_task_handler_callbacks.go
@@ -325,8 +325,11 @@ func (handler *workflowTaskHandlerCallbacksImpl) handleWorkflowTaskFailed(
 			if workflowTask == nil ||
 				workflowTask.StartedEventID == common.EmptyEventID ||
 				(token.StartedEventId != common.EmptyEventID && token.StartedEventId != workflowTask.StartedEventID) ||
+				(token.StartedTime != nil && workflowTask.StartedTime != nil && !token.StartedTime.Equal(*workflowTask.StartedTime)) ||
 				workflowTask.Attempt != token.Attempt ||
 				(workflowTask.Version != common.EmptyVersion && token.Version != workflowTask.Version) {
+				// we have not alter mutable state yet, so release with it with nil to avoid clear MS.
+				workflowContext.GetReleaseFn()(nil)
 				return nil, serviceerror.NewNotFound("Workflow task not found.")
 			}
 
@@ -393,6 +396,22 @@ func (handler *workflowTaskHandlerCallbacksImpl) handleWorkflowTaskCompleted(
 	if err != nil {
 		return nil, err
 	}
+	weContext := workflowContext.GetContext()
+	ms := workflowContext.GetMutableState()
+
+	currentWorkflowTask := ms.GetWorkflowTaskByID(token.GetScheduledEventId())
+	if !ms.IsWorkflowExecutionRunning() ||
+		currentWorkflowTask == nil ||
+		currentWorkflowTask.StartedEventID == common.EmptyEventID ||
+		(token.StartedEventId != common.EmptyEventID && token.StartedEventId != currentWorkflowTask.StartedEventID) ||
+		(token.StartedTime != nil && currentWorkflowTask.StartedTime != nil && !token.StartedTime.Equal(*currentWorkflowTask.StartedTime)) ||
+		currentWorkflowTask.Attempt != token.Attempt ||
+		(token.Version != common.EmptyVersion && token.Version != currentWorkflowTask.Version) {
+		// we have not alter mutable state yet, so release with it with nil to avoid clear MS.
+		workflowContext.GetReleaseFn()(nil)
+		return nil, serviceerror.NewNotFound("Workflow task not found.")
+	}
+
 	defer func() { workflowContext.GetReleaseFn()(retError) }()
 
 	var effects effect.Buffer
@@ -410,19 +429,6 @@ func (handler *workflowTaskHandlerCallbacksImpl) handleWorkflowTaskCompleted(
 		effects.Apply(ctx)
 	}()
 
-	weContext := workflowContext.GetContext()
-	ms := workflowContext.GetMutableState()
-
-	currentWorkflowTask := ms.GetWorkflowTaskByID(token.GetScheduledEventId())
-	if !ms.IsWorkflowExecutionRunning() ||
-		currentWorkflowTask == nil ||
-		currentWorkflowTask.StartedEventID == common.EmptyEventID ||
-		(token.StartedEventId != common.EmptyEventID && token.StartedEventId != currentWorkflowTask.StartedEventID) ||
-		currentWorkflowTask.Attempt != token.Attempt ||
-		(token.Version != common.EmptyVersion && token.Version != currentWorkflowTask.Version) {
-		return nil, serviceerror.NewNotFound("Workflow task not found.")
-	}
-
 	// It's an error if the workflow has used versioning in the past but this task has no versioning info.
 	if ms.GetWorkerVersionStamp().GetUseVersioning() && !request.GetWorkerVersionStamp().GetUseVersioning() {
 		return nil, serviceerror.NewInvalidArgument("Workflow using versioning must continue to use versioning.")