Fix heartbeats recording on activity failure (#578)

Fixes #521

internal/internal_task_handlers.go

@@ -1771,10 +1771,12 @@ func (ath *activityTaskHandlerImpl) Execute(taskQueue string, t *workflowservice
 	activityMetricsScope := metrics.GetMetricsScopeForActivity(ath.metricsScope, workflowType, activityType, ath.taskQueueName)
 	ctx := WithActivityTask(canCtx, t, taskQueue, invoker, ath.logger, activityMetricsScope, ath.dataConverter, ath.workerStopCh, ath.contextPropagators, ath.tracer)
 
-	defer func() {
+	// We must capture the context here because it is changed later to one that is
+	// cancelled when the activity is done
+	defer func(ctx context.Context) {
 		_, activityCompleted := result.(*workflowservice.RespondActivityTaskCompletedRequest)
 		invoker.Close(ctx, !activityCompleted) // flush buffered heartbeat if activity was not successfully completed.
-	}()
+	}(ctx)
 
 	activityImplementation := ath.getActivity(activityType)
 	if activityImplementation == nil {

test/activity_test.go

@@ -102,6 +102,32 @@ func (a *Activities) LongRunningHeartbeat(ctx context.Context, delay time.Durati
 	return nil
 }
 
+func (a *Activities) HeartbeatTwiceAndFailNTimes(
+	ctx context.Context,
+	times int,
+	id string,
+) (heartbeatCounts int, err error) {
+	// Get details
+	if activity.HasHeartbeatDetails(ctx) {
+		if err = activity.GetHeartbeatDetails(ctx, &heartbeatCounts); err != nil {
+			return
+		}
+	}
+
+	// Heartbeat twice, incrementing before each
+	heartbeatCounts++
+	activity.RecordHeartbeat(ctx, heartbeatCounts)
+	heartbeatCounts++
+	activity.RecordHeartbeat(ctx, heartbeatCounts)
+
+	// Set error if haven't reached enough times
+	a.append(id)
+	if a.invokedCount(id) <= times {
+		err = errFailOnPurpose
+	}
+	return
+}
+
 func (a *Activities) fail(_ context.Context) error {
 	a.append("fail")
 	return errFailOnPurpose

test/integration_test.go

@@ -418,6 +418,18 @@ func (ts *IntegrationTestSuite) TestLongRunningActivityWithHBAndGrpcRetries() {
 	ts.assertReportedOperationCount("temporal_request_attempt", "RecordActivityTaskHeartbeat", int(totalHeartbeats+3))
 }
 
+func (ts *IntegrationTestSuite) TestHeartbeatOnActivityFailure() {
+	var heartbeatCounts int
+	err := ts.executeWorkflow("test-heartbeat-on-activity-failure",
+		ts.workflows.ActivityHeartbeatWithRetry, &heartbeatCounts)
+	ts.NoError(err)
+	// Final count should be 6 because the activity is called 3 times (first 2
+	// fail) and each activity heartbeats twice. Before fixing a bug where the
+	// gRPC call wasn't made on activity failure, this was 4 because the first 2
+	// failing activities didn't have their second heartbeats recorded.
+	ts.Equal(6, heartbeatCounts)
+}
+
 func (ts *IntegrationTestSuite) TestContinueAsNew() {
 	var result int
 	err := ts.executeWorkflow("test-continueasnew", ts.workflows.ContinueAsNew, &result, 4, ts.taskQueueName)

test/workflow_test.go

@@ -268,6 +268,18 @@ func (w *Workflows) ActivityRetryOnHBTimeout(ctx workflow.Context) ([]string, er
 	return []string{"heartbeatAndSleep", "heartbeatAndSleep", "heartbeatAndSleep"}, nil
 }
 
+func (w *Workflows) ActivityHeartbeatWithRetry(ctx workflow.Context) (heartbeatCounts int, err error) {
+	// Make retries fast
+	opts := w.defaultActivityOptions()
+	opts.RetryPolicy = &temporal.RetryPolicy{InitialInterval: 5 * time.Millisecond, BackoffCoefficient: 1}
+	ctx = workflow.WithActivityOptions(ctx, opts)
+
+	// Fail twice then succeed
+	err = workflow.ExecuteActivity(ctx, "HeartbeatTwiceAndFailNTimes", 2,
+		"activity-heartbeat-"+workflow.GetInfo(ctx).WorkflowExecution.ID).Get(ctx, &heartbeatCounts)
+	return
+}
+
 func (w *Workflows) ContinueAsNew(ctx workflow.Context, count int, taskQueue string) (int, error) {
 	tq := workflow.GetInfo(ctx).TaskQueueName
 	if tq != taskQueue {
@@ -1265,6 +1277,7 @@ func (w *Workflows) WaitSignalReturnParam(ctx workflow.Context, v interface{}) (
 func (w *Workflows) register(worker worker.Worker) {
 	worker.RegisterWorkflow(w.ActivityCancelRepro)
 	worker.RegisterWorkflow(w.ActivityCompletionUsingID)
+	worker.RegisterWorkflow(w.ActivityHeartbeatWithRetry)
 	worker.RegisterWorkflow(w.ActivityRetryOnError)
 	worker.RegisterWorkflow(w.CallUnregisteredActivityRetry)
 	worker.RegisterWorkflow(w.ActivityRetryOnHBTimeout)