IWF-138: Add stateId and waitForKey to WorkflowWaitForStateCompletion (…

…#440)
indeedeng · Oct 8, 2024 · 3db7524 · 3db7524
1 parent 575d52d
commit 3db7524
Show file tree

Hide file tree

Showing 4 changed files with 65 additions and 33 deletions.
diff --git a/integ/wait_for_state_completion_test.go b/integ/wait_for_state_completion_test.go
@@ -74,37 +74,61 @@ func doTestWaitForStateCompletion(
 			WorkflowConfigOverride: config,
 		},
 	}
+
+	assertions := assert.New(t)
+
 	if useStateId {
-		startReq.WaitForCompletionStateIds = []string{"S1"}
+		startReq.WaitForCompletionStateIds = []string{"S2"}
+
+		_, httpResp, err := req.WorkflowStartRequest(startReq).Execute()
+		panicAtHttpError(err, httpResp)
+
+		req := apiClient.DefaultApi.ApiV1WorkflowWaitForStateCompletionPost(context.Background())
+		_, httpResp, err = req.WorkflowWaitForStateCompletionRequest(
+			iwfidl.WorkflowWaitForStateCompletionRequest{
+				WorkflowId:      wfId,
+				WaitForKey:      ptr.Any("testKey"),
+				StateId:         ptr.Any("S2"),
+				WaitTimeSeconds: iwfidl.PtrInt32(30),
+			}).Execute()
+		panicAtHttpError(err, httpResp)
+
+		assertions.Equal(200, httpResp.StatusCode)
+		// read httpResp body
+		var output iwfidl.WorkflowWaitForStateCompletionResponse
+		defer httpResp.Body.Close()
+		err = json.NewDecoder(httpResp.Body).Decode(&output)
+		if err != nil {
+			log.Fatalf("Failed to decode the response: %v", err)
+		}
 	} else {
 		startReq.WaitForCompletionStateExecutionIds = []string{"S1-1"}
-	}
-	_, httpResp, err := req.WorkflowStartRequest(startReq).Execute()
-	panicAtHttpError(err, httpResp)
 
-	req1 := apiClient.DefaultApi.ApiV1WorkflowWaitForStateCompletionPost(context.Background())
-	_, httpResp, err = req1.WorkflowWaitForStateCompletionRequest(
-		iwfidl.WorkflowWaitForStateCompletionRequest{
-			WorkflowId:       wfId,
-			StateExecutionId: ptr.Any("S1-1"),
-			WaitTimeSeconds:  iwfidl.PtrInt32(30),
-		}).Execute()
-	panicAtHttpError(err, httpResp)
+		_, httpResp, err := req.WorkflowStartRequest(startReq).Execute()
+		panicAtHttpError(err, httpResp)
 
-	assertions := assert.New(t)
-	assertions.Equal(200, httpResp.StatusCode)
-	// read httpResp body
-	var output iwfidl.WorkflowWaitForStateCompletionResponse
-	defer httpResp.Body.Close()
-	err = json.NewDecoder(httpResp.Body).Decode(&output)
-	if err != nil {
-		log.Fatalf("Failed to decode the response: %v", err)
+		req := apiClient.DefaultApi.ApiV1WorkflowWaitForStateCompletionPost(context.Background())
+		_, httpResp, err = req.WorkflowWaitForStateCompletionRequest(
+			iwfidl.WorkflowWaitForStateCompletionRequest{
+				WorkflowId:       wfId,
+				StateExecutionId: ptr.Any("S1-1"),
+				WaitTimeSeconds:  iwfidl.PtrInt32(30),
+			}).Execute()
+		panicAtHttpError(err, httpResp)
+
+		assertions.Equal(200, httpResp.StatusCode)
+		// read httpResp body
+		var output iwfidl.WorkflowWaitForStateCompletionResponse
+		defer httpResp.Body.Close()
+		err = json.NewDecoder(httpResp.Body).Decode(&output)
+		if err != nil {
+			log.Fatalf("Failed to decode the response: %v", err)
+		}
 	}
-	assertions.Equal("S1-1", output.StateCompletionOutput.CompletedStateExecutionId)
 
 	// wait for the workflow
 	req2 := apiClient.DefaultApi.ApiV1WorkflowGetWithWaitPost(context.Background())
-	_, httpResp, err = req2.WorkflowGetRequest(iwfidl.WorkflowGetRequest{
+	_, httpResp, err := req2.WorkflowGetRequest(iwfidl.WorkflowGetRequest{
 		WorkflowId: wfId,
 	}).Execute()
 	panicAtHttpError(err, httpResp)

diff --git a/integ/workflow/wait_for_state_completion/routers.go b/integ/workflow/wait_for_state_completion/routers.go
@@ -95,7 +95,8 @@ func (h *handler) ApiV1WorkflowStateDecide(c *gin.Context) {
 				StateDecision: &iwfidl.StateDecision{
 					NextStates: []iwfidl.StateMovement{
 						{
-							StateId: State2,
+							StateId:    State2,
+							WaitForKey: ptr.Any("testKey"),
 						},
 					},
 				},

diff --git a/service/api/service.go b/service/api/service.go
@@ -158,7 +158,7 @@ func (s *serviceImpl) ApiV1WorkflowWaitForStateCompletion(
 
 	var workflowId string
 	if req.WaitForKey != nil {
-		workflowId = service.IwfSystemConstPrefix + req.WorkflowId + "_" + *req.WaitForKey
+		workflowId = service.IwfSystemConstPrefix + req.WorkflowId + "_" + *req.StateId + "_" + *req.WaitForKey
 	} else {
 		workflowId = service.IwfSystemConstPrefix + req.WorkflowId + "_" + *req.StateExecutionId
 	}

diff --git a/service/interpreter/workflowImpl.go b/service/interpreter/workflowImpl.go
@@ -484,9 +484,9 @@ func executeState(
 	isResumeFromContinueAsNew := stateReq.IsResumeRequest()
 
 	options := state.GetStateOptions()
-	skipStart := compatibility.GetSkipStartApi(&options)
-	if skipStart {
-		return executeStateDecide(ctx, provider, basicInfo, state, stateExeId, persistenceManager, interStateChannel, executionContext,
+	skipWaitUntil := compatibility.GetSkipStartApi(&options)
+	if skipWaitUntil {
+		return invokeStateExecute(ctx, provider, basicInfo, state, stateExeId, persistenceManager, interStateChannel, executionContext,
 			nil, continueAsNewer, configer, executeApi, stateExecutionLocal, shouldSendSignalOnCompletion)
 	}
 
@@ -706,11 +706,11 @@ func executeState(
 		commandRes.SetInterStateChannelResults(interStateChannelResults)
 	}
 
-	return executeStateDecide(ctx, provider, basicInfo, state, stateExeId, persistenceManager, interStateChannel, executionContext,
+	return invokeStateExecute(ctx, provider, basicInfo, state, stateExeId, persistenceManager, interStateChannel, executionContext,
 		commandRes, continueAsNewer, configer, executeApi, stateExecutionLocal, shouldSendSignalOnCompletion)
 }
 
-func executeStateDecide(
+func invokeStateExecute(
 	ctx UnifiedContext,
 	provider WorkflowProvider,
 	basicInfo service.BasicInfo,
@@ -762,16 +762,23 @@ func executeStateDecide(
 		// NOTE: here uses NOT IsReplaying to signalWithStart, to save an activity for this operation
 		// this is not a problem because the signalWithStart will be very fast and highly available
 		unifiedClient := env.GetUnifiedClient()
+
+		var workflowId string
+
+		if state.WaitForKey != nil {
+			workflowId = service.IwfSystemConstPrefix + executionContext.WorkflowId + "_" + state.StateId + "_" + *state.WaitForKey
+		} else {
+			workflowId = service.IwfSystemConstPrefix + executionContext.WorkflowId + "_" + *executionContext.StateExecutionId
+		}
+
 		err := unifiedClient.SignalWithStartWaitForStateCompletionWorkflow(
 			context.Background(),
 			uclient.StartWorkflowOptions{
-				ID:                       service.IwfSystemConstPrefix + executionContext.WorkflowId + "_" + *executionContext.StateExecutionId,
+				ID:                       workflowId,
 				TaskQueue:                env.GetTaskQueue(),
 				WorkflowExecutionTimeout: 60 * time.Second, // timeout doesn't matter here as it will complete immediate with the signal
 			},
-			iwfidl.StateCompletionOutput{
-				CompletedStateExecutionId: *executionContext.StateExecutionId,
-			})
+			iwfidl.StateCompletionOutput{})
 		if err != nil && !unifiedClient.IsWorkflowAlreadyStartedError(err) {
 			// WorkflowAlreadyStartedError is returned when the started workflow is closed and the signal is not sent
 			// panic will let the workflow task will retry until the signal is sent