From cd0c3adfeaea30387bae42ba9e8dafada9f13c51 Mon Sep 17 00:00:00 2001
From: Jakob Haahr Taankvist <jht@uber.com>
Date: Thu, 21 Sep 2023 13:19:13 +0200
Subject: [PATCH] Refactor failed heartbeats log (#1273)

* Refactored the log type check

* Wrong check - reversed

* Added nil check

* Added tests, and fixed nil error
---
 internal/internal_task_handlers.go      | 17 ++++++----
 internal/internal_task_handlers_test.go | 45 +++++++++++++++++++++++++
 2 files changed, 56 insertions(+), 6 deletions(-)

diff --git a/internal/internal_task_handlers.go b/internal/internal_task_handlers.go
index 22d94bd2d..f2b12b01a 100644
--- a/internal/internal_task_handlers.go
+++ b/internal/internal_task_handlers.go
@@ -1495,18 +1495,23 @@ func (i *cadenceInvoker) heartbeatAndScheduleNextRun(details []byte) error {
 			i.Unlock()
 
 			// Log the error outside the lock.
-			if err != nil {
-				// If the error is a canceled error do not log, as this is expected.
-				if _, ok := err.(*CanceledError); !ok {
-					i.logger.Error("Failed to send heartbeat", zap.Error(err), zap.String(tagWorkflowType, i.workflowType), zap.String(tagActivityType, i.activityType))
-				}
-			}
+			i.logFailedHeartBeat(err)
 		}()
 	}
 
 	return err
 }
 
+func (i *cadenceInvoker) logFailedHeartBeat(err error) {
+	// If the error is a canceled error do not log, as this is expected.
+	var canceledErr *CanceledError
+
+	// We need to check for nil as errors.As returns false for nil. Which would cause us to log on nil.
+	if err != nil && !errors.As(err, &canceledErr) {
+		i.logger.Error("Failed to send heartbeat", zap.Error(err), zap.String(tagWorkflowType, i.workflowType), zap.String(tagActivityType, i.activityType))
+	}
+}
+
 func (i *cadenceInvoker) internalHeartBeat(details []byte) (bool, error) {
 	isActivityCancelled := false
 	timeout := time.Duration(i.heartBeatTimeoutInSec) * time.Second
diff --git a/internal/internal_task_handlers_test.go b/internal/internal_task_handlers_test.go
index 751be04fb..358f86310 100644
--- a/internal/internal_task_handlers_test.go
+++ b/internal/internal_task_handlers_test.go
@@ -1381,6 +1381,51 @@ func (t *TaskHandlersTestSuite) TestHeartBeat_Interleaved() {
 	time.Sleep(1 * time.Second)
 }
 
+func (t *TaskHandlersTestSuite) TestHeartBeatLogNil() {
+	core, obs := observer.New(zap.ErrorLevel)
+	logger := zap.New(core)
+
+	cadenceInv := &cadenceInvoker{
+		identity: "Test_Cadence_Invoker",
+		logger:   logger,
+	}
+
+	cadenceInv.logFailedHeartBeat(nil)
+
+	t.Empty(obs.All())
+}
+
+func (t *TaskHandlersTestSuite) TestHeartBeatLogCanceledError() {
+	core, obs := observer.New(zap.ErrorLevel)
+	logger := zap.New(core)
+
+	cadenceInv := &cadenceInvoker{
+		identity: "Test_Cadence_Invoker",
+		logger:   logger,
+	}
+
+	var workflowCompleatedErr CanceledError
+	cadenceInv.logFailedHeartBeat(&workflowCompleatedErr)
+
+	t.Empty(obs.All())
+}
+
+func (t *TaskHandlersTestSuite) TestHeartBeatLogNotCanceled() {
+	core, obs := observer.New(zap.ErrorLevel)
+	logger := zap.New(core)
+
+	cadenceInv := &cadenceInvoker{
+		identity: "Test_Cadence_Invoker",
+		logger:   logger,
+	}
+
+	var workflowCompleatedErr s.WorkflowExecutionAlreadyCompletedError
+	cadenceInv.logFailedHeartBeat(&workflowCompleatedErr)
+
+	t.Len(obs.All(), 1)
+	t.Equal("Failed to send heartbeat", obs.All()[0].Message)
+}
+
 func (t *TaskHandlersTestSuite) TestHeartBeat_NilResponseWithError() {
 	mockCtrl := gomock.NewController(t.T())
 	mockService := workflowservicetest.NewMockClient(mockCtrl)