From 9739cc13636a4d2cc47d00ca1e9517258186e10f Mon Sep 17 00:00:00 2001
From: lwolczynski <54366429+lwolczynski@users.noreply.github.com>
Date: Fri, 22 Nov 2024 15:07:12 -0600
Subject: [PATCH] NOBUG: Fix emitted metrics (#489)

---
 service/interpreter/activityImpl.go | 22 ++++++++++++++--------
 service/interpreter/workflowImpl.go | 28 +++++++++++++++++-----------
 2 files changed, 31 insertions(+), 19 deletions(-)

diff --git a/service/interpreter/activityImpl.go b/service/interpreter/activityImpl.go
index a1ea6ed9..bba26aa3 100644
--- a/service/interpreter/activityImpl.go
+++ b/service/interpreter/activityImpl.go
@@ -55,10 +55,12 @@ func StateApiWaitUntil(
 	printDebugMsg(logger, err, iwfWorkerBaseUrl)
 	if checkHttpError(err, httpResp) {
 		event.Handle(iwfidl.IwfEvent{
-			EventType:     iwfidl.STATE_WAIT_UNTIL_ATTEMPT_FAIL_EVENT,
-			WorkflowType:  input.Request.WorkflowType,
-			WorkflowId:    activityInfo.WorkflowExecution.ID,
-			WorkflowRunId: activityInfo.WorkflowExecution.RunID,
+			EventType:        iwfidl.STATE_WAIT_UNTIL_ATTEMPT_FAIL_EVENT,
+			WorkflowType:     input.Request.WorkflowType,
+			WorkflowId:       activityInfo.WorkflowExecution.ID,
+			WorkflowRunId:    activityInfo.WorkflowExecution.RunID,
+			StateId:          ptr.Any(input.Request.WorkflowStateId),
+			StateExecutionId: ptr.Any(input.Request.Context.GetStateExecutionId()),
 		})
 		return nil, composeHttpError(
 			activityInfo.IsLocalActivity,
@@ -67,10 +69,12 @@ func StateApiWaitUntil(
 
 	if err := checkCommandRequestFromWaitUntilResponse(resp); err != nil {
 		event.Handle(iwfidl.IwfEvent{
-			EventType:     iwfidl.STATE_WAIT_UNTIL_ATTEMPT_FAIL_EVENT,
-			WorkflowType:  input.Request.WorkflowType,
-			WorkflowId:    activityInfo.WorkflowExecution.ID,
-			WorkflowRunId: activityInfo.WorkflowExecution.RunID,
+			EventType:        iwfidl.STATE_WAIT_UNTIL_ATTEMPT_FAIL_EVENT,
+			WorkflowType:     input.Request.WorkflowType,
+			WorkflowId:       activityInfo.WorkflowExecution.ID,
+			WorkflowRunId:    activityInfo.WorkflowExecution.RunID,
+			StateId:          ptr.Any(input.Request.WorkflowStateId),
+			StateExecutionId: ptr.Any(input.Request.Context.GetStateExecutionId()),
 		})
 		return nil, composeStartApiRespError(provider, err, resp)
 	}
@@ -87,6 +91,8 @@ func StateApiWaitUntil(
 		WorkflowType:       input.Request.WorkflowType,
 		WorkflowId:         activityInfo.WorkflowExecution.ID,
 		WorkflowRunId:      activityInfo.WorkflowExecution.RunID,
+		StateId:            ptr.Any(input.Request.WorkflowStateId),
+		StateExecutionId:   ptr.Any(input.Request.Context.GetStateExecutionId()),
 		StartTimestampInMs: ptr.Any(stateApiWaitUntilStartTime),
 		EndTimestampInMs:   ptr.Any(time.Now().UnixMilli()),
 	})
diff --git a/service/interpreter/workflowImpl.go b/service/interpreter/workflowImpl.go
index c5acdceb..327e7276 100644
--- a/service/interpreter/workflowImpl.go
+++ b/service/interpreter/workflowImpl.go
@@ -562,10 +562,12 @@ func processStateExecution(
 
 		if !provider.IsReplaying(ctx) {
 			event.Handle(iwfidl.IwfEvent{
-				EventType:     iwfidl.STATE_WAIT_UNTIL_EE_START_EVENT,
-				WorkflowType:  basicInfo.IwfWorkflowType,
-				WorkflowId:    provider.GetWorkflowInfo(ctx).WorkflowExecution.ID,
-				WorkflowRunId: provider.GetWorkflowInfo(ctx).WorkflowExecution.RunID,
+				EventType:        iwfidl.STATE_WAIT_UNTIL_EE_START_EVENT,
+				WorkflowType:     basicInfo.IwfWorkflowType,
+				WorkflowId:       provider.GetWorkflowInfo(ctx).WorkflowExecution.ID,
+				WorkflowRunId:    provider.GetWorkflowInfo(ctx).WorkflowExecution.RunID,
+				StateId:          ptr.Any(state.StateId),
+				StateExecutionId: ptr.Any(stateExeId),
 			})
 		}
 		stateWaitUntilApiStartTime := provider.Now(ctx).UnixMilli()
@@ -583,21 +585,25 @@ func processStateExecution(
 			})
 		if !provider.IsReplaying(ctx) {
 			if errStartApi == nil {
-				event.Handle(iwfidl.IwfEvent{
-					EventType:     iwfidl.STATE_WAIT_UNTIL_EE_FAIL_EVENT,
-					WorkflowType:  basicInfo.IwfWorkflowType,
-					WorkflowId:    provider.GetWorkflowInfo(ctx).WorkflowExecution.ID,
-					WorkflowRunId: provider.GetWorkflowInfo(ctx).WorkflowExecution.RunID,
-				})
-			} else {
 				event.Handle(iwfidl.IwfEvent{
 					EventType:          iwfidl.STATE_WAIT_UNTIL_EE_COMPLETE_EVENT,
 					WorkflowType:       basicInfo.IwfWorkflowType,
 					WorkflowId:         provider.GetWorkflowInfo(ctx).WorkflowExecution.ID,
 					WorkflowRunId:      provider.GetWorkflowInfo(ctx).WorkflowExecution.RunID,
+					StateId:            ptr.Any(state.StateId),
+					StateExecutionId:   ptr.Any(stateExeId),
 					StartTimestampInMs: ptr.Any(stateWaitUntilApiStartTime),
 					EndTimestampInMs:   ptr.Any(provider.Now(ctx).UnixMilli()),
 				})
+			} else {
+				event.Handle(iwfidl.IwfEvent{
+					EventType:        iwfidl.STATE_WAIT_UNTIL_EE_FAIL_EVENT,
+					WorkflowType:     basicInfo.IwfWorkflowType,
+					WorkflowId:       provider.GetWorkflowInfo(ctx).WorkflowExecution.ID,
+					WorkflowRunId:    provider.GetWorkflowInfo(ctx).WorkflowExecution.RunID,
+					StateId:          ptr.Any(state.StateId),
+					StateExecutionId: ptr.Any(stateExeId),
+				})
 			}
 		}