Add resend context timeout for ndc resender (uber#3247)

common/xdc/nDCHistoryResender.go

@@ -36,6 +36,7 @@ import (
 	"github.com/uber/cadence/common/log"
 	"github.com/uber/cadence/common/log/tag"
 	"github.com/uber/cadence/common/persistence"
+	"github.com/uber/cadence/common/service/dynamicconfig"
 )
 
 const (
@@ -67,6 +68,7 @@ type (
 		adminClient          adminClient.Client
 		historyReplicationFn nDCHistoryReplicationFn
 		serializer           persistence.PayloadSerializer
+		rereplicationTimeout dynamicconfig.DurationPropertyFnWithDomainIDFilter
 		logger               log.Logger
 	}
 
@@ -82,6 +84,7 @@ func NewNDCHistoryResender(
 	adminClient adminClient.Client,
 	historyReplicationFn nDCHistoryReplicationFn,
 	serializer persistence.PayloadSerializer,
+	rereplicationTimeout dynamicconfig.DurationPropertyFnWithDomainIDFilter,
 	logger log.Logger,
 ) *NDCHistoryResenderImpl {
 
@@ -90,6 +93,7 @@ func NewNDCHistoryResender(
 		adminClient:          adminClient,
 		historyReplicationFn: historyReplicationFn,
 		serializer:           serializer,
+		rereplicationTimeout: rereplicationTimeout,
 		logger:               logger,
 	}
 }
@@ -105,7 +109,18 @@ func (n *NDCHistoryResenderImpl) SendSingleWorkflowHistory(
 	endEventVersion *int64,
 ) error {
 
+	ctx := context.Background()
+	var cancel context.CancelFunc
+	if n.rereplicationTimeout != nil {
+		resendContextTimeout := n.rereplicationTimeout(domainID)
+		if resendContextTimeout > 0 {
+			ctx, cancel = context.WithTimeout(ctx, resendContextTimeout)
+			defer cancel()
+		}
+	}
+
 	historyIterator := collection.NewPagingIterator(n.getPaginationFn(
+		ctx,
 		domainID,
 		workflowID,
 		runID,
@@ -133,7 +148,7 @@ func (n *NDCHistoryResenderImpl) SendSingleWorkflowHistory(
 			historyBatch.rawEventBatch,
 			historyBatch.versionHistory.GetItems())
 
-		err = n.sendReplicationRawRequest(replicationRequest)
+		err = n.sendReplicationRawRequest(ctx, replicationRequest)
 		if err != nil {
 			n.logger.Error("failed to replicate events",
 				tag.WorkflowDomainID(domainID),
@@ -147,6 +162,7 @@ func (n *NDCHistoryResenderImpl) SendSingleWorkflowHistory(
 }
 
 func (n *NDCHistoryResenderImpl) getPaginationFn(
+	ctx context.Context,
 	domainID string,
 	workflowID string,
 	runID string,
@@ -159,6 +175,7 @@ func (n *NDCHistoryResenderImpl) getPaginationFn(
 	return func(paginationToken []byte) ([]interface{}, []byte, error) {
 
 		response, err := n.getHistory(
+			ctx,
 			domainID,
 			workflowID,
 			runID,
@@ -207,15 +224,17 @@ func (n *NDCHistoryResenderImpl) createReplicationRawRequest(
 }
 
 func (n *NDCHistoryResenderImpl) sendReplicationRawRequest(
+	ctx context.Context,
 	request *history.ReplicateEventsV2Request,
 ) error {
 
-	ctx, cancel := context.WithTimeout(context.Background(), resendContextTimeout)
+	ctx, cancel := context.WithTimeout(ctx, resendContextTimeout)
 	defer cancel()
 	return n.historyReplicationFn(ctx, request)
 }
 
 func (n *NDCHistoryResenderImpl) getHistory(
+	ctx context.Context,
 	domainID string,
 	workflowID string,
 	runID string,
@@ -236,7 +255,7 @@ func (n *NDCHistoryResenderImpl) getHistory(
 	}
 	domainName := domainEntry.GetInfo().Name
 
-	ctx, cancel := context.WithTimeout(context.Background(), resendContextTimeout)
+	ctx, cancel := context.WithTimeout(ctx, resendContextTimeout)
 	defer cancel()
 	response, err := n.adminClient.GetWorkflowExecutionRawHistoryV2(ctx, &admin.GetWorkflowExecutionRawHistoryV2Request{
 		Domain: common.StringPtr(domainName),

common/xdc/nDCHistoryResender_test.go

@@ -116,6 +116,7 @@ func (s *nDCHistoryResenderSuite) SetupTest() {
 			return s.mockHistoryClient.ReplicateEventsV2(ctx, request)
 		},
 		persistence.NewPayloadSerializer(),
+		nil,
 		s.logger,
 	)
 }
@@ -269,7 +270,7 @@ func (s *nDCHistoryResenderSuite) TestSendReplicationRawRequest() {
 	}
 
 	s.mockHistoryClient.EXPECT().ReplicateEventsV2(gomock.Any(), request).Return(nil).Times(1)
-	err := s.rereplicator.sendReplicationRawRequest(request)
+	err := s.rereplicator.sendReplicationRawRequest(context.Background(), request)
 	s.Nil(err)
 }
 
@@ -299,7 +300,7 @@ func (s *nDCHistoryResenderSuite) TestSendReplicationRawRequest_Err() {
 	}
 
 	s.mockHistoryClient.EXPECT().ReplicateEventsV2(gomock.Any(), request).Return(retryErr).Times(1)
-	err := s.rereplicator.sendReplicationRawRequest(request)
+	err := s.rereplicator.sendReplicationRawRequest(context.Background(), request)
 	s.Equal(retryErr, err)
 }
 
@@ -336,6 +337,7 @@ func (s *nDCHistoryResenderSuite) TestGetHistory() {
 	}).Return(response, nil).Times(1)
 
 	out, err := s.rereplicator.getHistory(
+		context.Background(),
 		s.domainID,
 		workflowID,
 		runID,

service/frontend/adminHandler.go

@@ -1008,6 +1008,7 @@ func (adh *AdminHandler) ResendReplicationTasks(
 			return adh.GetHistoryClient().ReplicateEventsV2(ctx, request)
 		},
 		adh.eventSerializder,
+		nil,
 		adh.GetLogger(),
 	)
 	return resender.SendSingleWorkflowHistory(

service/history/historyEngine.go

@@ -258,6 +258,7 @@ func NewEngineWithShardContext(
 			return shard.GetService().GetHistoryClient().ReplicateEventsV2(ctx, request)
 		},
 		shard.GetService().GetPayloadSerializer(),
+		nil,
 		shard.GetLogger(),
 	)
 	historyRereplicator := xdc.NewHistoryRereplicator(

service/history/timerQueueProcessor.go

@@ -111,6 +111,7 @@ func newTimerQueueProcessor(
 					return historyService.ReplicateEventsV2(ctx, request)
 				},
 				shard.GetService().GetPayloadSerializer(),
+				nil,
 				logger,
 			)
 			standbyTimerProcessors[clusterName] = newTimerQueueStandbyProcessor(

service/history/transferQueueProcessor.go

@@ -112,6 +112,7 @@ func newTransferQueueProcessor(
 					return historyService.ReplicateEventsV2(ctx, request)
 				},
 				shard.GetService().GetPayloadSerializer(),
+				nil,
 				logger,
 			)
 			standbyTaskProcessors[clusterName] = newTransferQueueStandbyProcessor(

service/worker/replicator/replicator.go

@@ -187,6 +187,7 @@ func (r *Replicator) createKafkaProcessors(currentClusterName string, clusterNam
 			return historyClient.ReplicateEventsV2(ctx, request)
 		},
 		r.historySerializer,
+		r.config.ReReplicationContextTimeout,
 		logger,
 	)
 	r.processors = append(r.processors, newReplicationTaskProcessor(