Improve session framework heartbeat mechanism (#965)

Author: yycptt

internal/activity.go

@@ -260,7 +260,7 @@ func RecordActivityHeartbeat(ctx context.Context, details ...interface{}) {
 			panic(err)
 		}
 	}
-	err = env.serviceInvoker.Heartbeat(data)
+	err = env.serviceInvoker.Heartbeat(data, false)
 	if err != nil {
 		log := GetActivityLogger(ctx)
 		log.Debug("RecordActivityHeartbeat With Error:", zap.Error(err))
@@ -271,7 +271,7 @@ func RecordActivityHeartbeat(ctx context.Context, details ...interface{}) {
 // Implement to unit test activities.
 type ServiceInvoker interface {
 	// Returns ActivityTaskCanceledError if activity is cancelled
-	Heartbeat(details []byte) error
+	Heartbeat(details []byte, skipBatching bool) error
 	Close(flushBufferedHeartbeat bool)
 	GetClient(domain string, options *ClientOptions) Client
 }

internal/internal_task_handlers.go

@@ -1621,11 +1621,11 @@ type cadenceInvoker struct {
 	workerStopChannel     <-chan struct{}
 }
 
-func (i *cadenceInvoker) Heartbeat(details []byte) error {
+func (i *cadenceInvoker) Heartbeat(details []byte, skipBatching bool) error {
 	i.Lock()
 	defer i.Unlock()
 
-	if i.hbBatchEndTimer != nil {
+	if i.hbBatchEndTimer != nil && !skipBatching {
 		// If we have started batching window, keep track of last reported progress.
 		i.lastDetailsToReport = &details
 		return nil
@@ -1635,7 +1635,7 @@ func (i *cadenceInvoker) Heartbeat(details []byte) error {
 
 	// If the activity is cancelled, the activity can ignore the cancellation and do its work
 	// and complete. Our cancellation is co-operative, so we will try to heartbeat.
-	if err == nil || isActivityCancelled {
+	if (err == nil || isActivityCancelled) && !skipBatching {
 		// We have successfully sent heartbeat, start next batching window.
 		i.lastDetailsToReport = nil
 
@@ -1671,7 +1671,11 @@ func (i *cadenceInvoker) Heartbeat(details []byte) error {
 			i.Unlock()
 
 			if detailsToReport != nil {
-				i.Heartbeat(*detailsToReport)
+				// TODO: there is a potential race condition here as the lock is released here and
+				// locked again in the Hearbeat() method. This possible that a heartbeat call from
+				// user activity grabs the lock first and calls internalHeartBeat before this
+				// batching goroutine, which means some activity progress will be lost.
+				i.Heartbeat(*detailsToReport, false)
 			}
 		}()
 	}

internal/internal_task_handlers_test.go

@@ -1171,7 +1171,7 @@ func (t *TaskHandlersTestSuite) TestHeartBeat_NoError() {
 		taskToken: nil,
 	}
 
-	heartbeatErr := cadenceInvoker.Heartbeat(nil)
+	heartbeatErr := cadenceInvoker.Heartbeat(nil, false)
 
 	t.Nil(heartbeatErr)
 }
@@ -1191,7 +1191,7 @@ func (t *TaskHandlersTestSuite) TestHeartBeat_NilResponseWithError() {
 		0,
 		make(chan struct{}))
 
-	heartbeatErr := cadenceInvoker.Heartbeat(nil)
+	heartbeatErr := cadenceInvoker.Heartbeat(nil, false)
 	t.NotNil(heartbeatErr)
 	_, ok := (heartbeatErr).(*s.EntityNotExistsError)
 	t.True(ok, "heartbeatErr must be EntityNotExistsError.")
@@ -1215,7 +1215,7 @@ func (t *TaskHandlersTestSuite) TestHeartBeat_NilResponseWithDomainNotActiveErro
 		0,
 		make(chan struct{}))
 
-	heartbeatErr := cadenceInvoker.Heartbeat(nil)
+	heartbeatErr := cadenceInvoker.Heartbeat(nil, false)
 	t.NotNil(heartbeatErr)
 	_, ok := (heartbeatErr).(*s.DomainNotActiveError)
 	t.True(ok, "heartbeatErr must be DomainNotActiveError.")

internal/session.go

@@ -28,6 +28,7 @@ import (
 	"time"
 
 	"github.com/pborman/uuid"
+	"go.uber.org/cadence/internal/common/backoff"
 	"go.uber.org/zap"
 )
 
@@ -121,8 +122,8 @@ const (
 
 	errTooManySessionsMsg string = "too many outstanding sessions"
 
-	defaultSessionHeartBeatTimeout time.Duration = time.Second * 20
-	maxSessionHeartBeatInterval    time.Duration = time.Second * 10
+	defaultSessionHeartbeatTimeout time.Duration = time.Second * 20
+	maxSessionHeartbeatInterval    time.Duration = time.Second * 10
 )
 
 var (
@@ -305,7 +306,7 @@ func createSession(ctx Context, creationTasklist string, options *SessionOptions
 		},
 	}
 
-	heartbeatTimeout := defaultSessionHeartBeatTimeout
+	heartbeatTimeout := defaultSessionHeartbeatTimeout
 	if options.HeartbeatTimeout != time.Duration(0) {
 		heartbeatTimeout = options.HeartbeatTimeout
 	}
@@ -410,22 +411,47 @@ func sessionCreationActivity(ctx context.Context, sessionID string) error {
 
 	activityEnv := getActivityEnv(ctx)
 	heartbeatInterval := activityEnv.heartbeatTimeout / 3
-	if heartbeatInterval > maxSessionHeartBeatInterval {
-		heartbeatInterval = maxSessionHeartBeatInterval
+	if heartbeatInterval > maxSessionHeartbeatInterval {
+		heartbeatInterval = maxSessionHeartbeatInterval
 	}
 	ticker := time.NewTicker(heartbeatInterval)
 	defer ticker.Stop()
 
+	heartbeatRetryPolicy := backoff.NewExponentialRetryPolicy(time.Second)
+	heartbeatRetryPolicy.SetMaximumInterval(time.Second * 2)
+	heartbeatRetryPolicy.SetExpirationInterval(heartbeatInterval)
+
 	for {
 		select {
 		case <-ctx.Done():
 			sessionEnv.CompleteSession(sessionID)
 			return ctx.Err()
 		case <-ticker.C:
-			err := activityEnv.serviceInvoker.Heartbeat([]byte{})
+			heartbeatOp := func() error {
+				// here we skip the internal heartbeat batching, as otherwise the activity has only once chance
+				// for heartbeating and if that failed, the entire session will get fail due to heartbeat timeout.
+				// since the heartbeat interval is controlled by the session framework, we don't need to worry about
+				// calling heartbeat too frequently and causing trouble for the sever. (note the min heartbeat timeout
+				// is 1 sec.)
+				return activityEnv.serviceInvoker.Heartbeat([]byte{}, true)
+			}
+			isRetryable := func(_ error) bool {
+				// there will be two types of error here:
+				// 1. transient errors like timeout, in which case we should not fail the session
+				// 2. non-retryable errors like activity cancelled, activity not found or domain
+				// not active. In those cases, the internal implementation will cancel the context,
+				// so in the next iteration, ctx.Done() will be selected. Here we rely on the heartbeat
+				// internal implementation to tell which error is non-retryable.
+				select {
+				case <-ctx.Done():
+					return false
+				default:
+					return true
+				}
+			}
+			err := backoff.Retry(ctx, heartbeatOp, heartbeatRetryPolicy, isRetryable)
 			if err != nil {
-				sessionEnv.CompleteSession(sessionID)
-				return err
+				GetActivityLogger(ctx).Info("session heartbeat failed", zap.Error(err), zap.String("sessionID", sessionID))
 			}
 		case <-doneCh:
 			return nil