Enforce context timeout for retry policies in execution context and c…

…ache (cadence-workflow#4303)

common/util.go

@@ -140,6 +140,23 @@ func CreatePersistenceRetryPolicy() backoff.RetryPolicy {
 	return policy
 }
 
+// CreatePersistenceRetryPolicyWithContext create a retry policy for persistence layer operations
+// which has an expiration interval computed based on the context's deadline
+func CreatePersistenceRetryPolicyWithContext(ctx context.Context) backoff.RetryPolicy {
+	if ctx == nil {
+		return CreatePersistenceRetryPolicy()
+	}
+	deadline, ok := ctx.Deadline()
+	if !ok {
+		return CreatePersistenceRetryPolicy()
+	}
+
+	policy := backoff.NewExponentialRetryPolicy(retryPersistenceOperationInitialInterval)
+	policy.SetMaximumInterval(retryPersistenceOperationMaxInterval)
+	policy.SetExpirationInterval(deadline.Sub(time.Now()))
+	return policy
+}
+
 // CreateHistoryServiceRetryPolicy creates a retry policy for calls to history service
 func CreateHistoryServiceRetryPolicy() backoff.RetryPolicy {
 	policy := backoff.NewExponentialRetryPolicy(historyServiceOperationInitialInterval)

service/history/execution/cache.go

@@ -27,6 +27,7 @@ import (
 
 	"github.com/pborman/uuid"
 
+	"github.com/uber/cadence/common"
 	"github.com/uber/cadence/common/backoff"
 	"github.com/uber/cadence/common/cache"
 	"github.com/uber/cadence/common/definition"
@@ -315,7 +316,11 @@ func (c *Cache) getCurrentExecutionWithRetry(
 		return err
 	}
 
-	err := backoff.Retry(op, persistenceOperationRetryPolicy, persistence.IsTransientError)
+	err := backoff.Retry(
+		op,
+		common.CreatePersistenceRetryPolicyWithContext(ctx),
+		persistence.IsTransientError,
+	)
 	if err != nil {
 		c.metricsClient.IncCounter(metrics.HistoryCacheGetCurrentExecutionScope, metrics.CacheFailures)
 		return nil, err

service/history/execution/context.go

@@ -153,10 +153,6 @@ type (
 
 var _ Context = (*contextImpl)(nil)
 
-var (
-	persistenceOperationRetryPolicy = common.CreatePersistenceRetryPolicy()
-)
-
 // NewContext creates a new workflow execution context
 func NewContext(
 	domainID string,
@@ -1003,7 +999,7 @@ func (c *contextImpl) appendHistoryV2EventsWithRetry(
 
 	err := backoff.Retry(
 		op,
-		persistenceOperationRetryPolicy,
+		common.CreatePersistenceRetryPolicyWithContext(ctx),
 		persistence.IsTransientError,
 	)
 	return int64(resp), err
@@ -1033,7 +1029,7 @@ func (c *contextImpl) createWorkflowExecutionWithRetry(
 
 	err := backoff.Retry(
 		op,
-		persistenceOperationRetryPolicy,
+		common.CreatePersistenceRetryPolicyWithContext(ctx),
 		isRetryable,
 	)
 	switch err.(type) {
@@ -1071,7 +1067,7 @@ func (c *contextImpl) getWorkflowExecutionWithRetry(
 
 	err := backoff.Retry(
 		op,
-		persistenceOperationRetryPolicy,
+		common.CreatePersistenceRetryPolicyWithContext(ctx),
 		persistence.IsTransientError,
 	)
 	switch err.(type) {
@@ -1114,7 +1110,8 @@ func (c *contextImpl) updateWorkflowExecutionWithRetry(
 	}
 
 	err := backoff.Retry(
-		op, persistenceOperationRetryPolicy,
+		op,
+		common.CreatePersistenceRetryPolicyWithContext(ctx),
 		isRetryable,
 	)
 	switch err.(type) {