Ensure min number of scan workflow result for shadowing (#1084)

Author: yycptt

internal/internal_workflow_testsuite.go

@@ -529,6 +529,21 @@ func (env *testWorkflowEnvironmentImpl) getWorkflowDefinition(wt WorkflowType) (
 func (env *testWorkflowEnvironmentImpl) executeActivity(
 	activityFn interface{},
 	args ...interface{},
+) (Value, error) {
+	return env.executeActivityWithOptions(
+		activityOptions{
+			ScheduleToCloseTimeoutSeconds: 600,
+			StartToCloseTimeoutSeconds:    600,
+		},
+		activityFn,
+		args...,
+	)
+}
+
+func (env *testWorkflowEnvironmentImpl) executeActivityWithOptions(
+	activityOptions activityOptions,
+	activityFn interface{},
+	args ...interface{},
 ) (Value, error) {
 	activityType, err := getValidatedActivityFunction(activityFn, args, env.registry)
 	if err != nil {
@@ -541,13 +556,10 @@ func (env *testWorkflowEnvironmentImpl) executeActivity(
 	}
 
 	params := executeActivityParams{
-		activityOptions: activityOptions{
-			ScheduleToCloseTimeoutSeconds: 600,
-			StartToCloseTimeoutSeconds:    600,
-		},
-		ActivityType: *activityType,
-		Input:        input,
-		Header:       env.header,
+		activityOptions: activityOptions,
+		ActivityType:    *activityType,
+		Input:           input,
+		Header:          env.header,
 	}
 
 	task := newTestActivityTask(

internal/workflow_shadower_activities.go

@@ -24,6 +24,7 @@ import (
 	"context"
 	"math/rand"
 	"strings"
+	"time"
 
 	"go.uber.org/cadence/.gen/go/cadence/workflowserviceclient"
 	"go.uber.org/cadence/.gen/go/shadower"
@@ -46,6 +47,12 @@ const (
 	workflowReplayerContextKey contextKey = "workflowReplayer"
 )
 
+const (
+	minScanWorkflowResultSize   = 10
+	ratioToCompleteScanWorkflow = 0.8
+	scanWorkflowWaitPeriod      = 100 * time.Millisecond
+)
+
 func scanWorkflowActivity(
 	ctx context.Context,
 	params shadower.ScanWorkflowActivityParams,
@@ -69,6 +76,12 @@ func scanWorkflowExecutionsHelper(
 	params shadower.ScanWorkflowActivityParams,
 	logger *zap.Logger,
 ) (shadower.ScanWorkflowActivityResult, error) {
+	var completionTime time.Time
+	if deadline, ok := ctx.Deadline(); ok {
+		now := time.Now()
+		activityTimeout := deadline.Sub(now)
+		completionTime = now.Add(time.Duration(ratioToCompleteScanWorkflow * float32(activityTimeout)))
+	}
 
 	request := &shared.ListWorkflowExecutionsRequest{
 		Domain:        params.Domain,
@@ -77,39 +90,48 @@ func scanWorkflowExecutionsHelper(
 		PageSize:      params.PageSize,
 	}
 
-	var resp *shared.ListWorkflowExecutionsResponse
-	if err := backoff.Retry(ctx,
-		func() error {
-			tchCtx, cancel, opt := newChannelContext(ctx)
-
-			var err error
-			resp, err = service.ScanWorkflowExecutions(tchCtx, request, opt...)
-			cancel()
-
-			return err
-		},
-		createDynamicServiceRetryPolicy(ctx),
-		isServiceTransientError,
-	); err != nil {
-		logger.Error("Failed to scan workflow executions",
-			zap.String(tagDomain, params.GetDomain()),
-			zap.String(tagVisibilityQuery, params.GetWorkflowQuery()),
-			zap.Error(err),
-		)
-		return shadower.ScanWorkflowActivityResult{}, err
-	}
+	result := shadower.ScanWorkflowActivityResult{}
+	for {
+		var resp *shared.ListWorkflowExecutionsResponse
+		if err := backoff.Retry(ctx,
+			func() error {
+				tchCtx, cancel, opt := newChannelContext(ctx)
+
+				var err error
+				resp, err = service.ScanWorkflowExecutions(tchCtx, request, opt...)
+				cancel()
 
-	executions := make([]*shared.WorkflowExecution, 0, len(resp.Executions))
-	for _, execution := range resp.Executions {
-		if shouldReplay(params.GetSamplingRate()) {
-			executions = append(executions, execution.Execution)
+				return err
+			},
+			createDynamicServiceRetryPolicy(ctx),
+			isServiceTransientError,
+		); err != nil {
+			logger.Error("Failed to scan workflow executions",
+				zap.String(tagDomain, params.GetDomain()),
+				zap.String(tagVisibilityQuery, params.GetWorkflowQuery()),
+				zap.Error(err),
+			)
+			return shadower.ScanWorkflowActivityResult{}, err
 		}
+
+		for _, execution := range resp.Executions {
+			if shouldReplay(params.GetSamplingRate()) {
+				result.Executions = append(result.Executions, execution.Execution)
+			}
+		}
+
+		request.NextPageToken = resp.NextPageToken
+		if len(request.NextPageToken) == 0 ||
+			len(result.Executions) >= minScanWorkflowResultSize ||
+			(!completionTime.IsZero() && time.Now().After(completionTime)) {
+			result.NextPageToken = request.NextPageToken
+			break
+		}
+
+		time.Sleep(scanWorkflowWaitPeriod)
 	}
 
-	return shadower.ScanWorkflowActivityResult{
-		Executions:    executions,
-		NextPageToken: resp.NextPageToken,
-	}, nil
+	return result, nil
 }
 
 func shouldReplay(probability float64) bool {

internal/workflow_shadower_activities_test.go

@@ -22,8 +22,10 @@ package internal
 
 import (
 	"context"
+	"math"
 	"math/rand"
 	"testing"
+	"time"
 
 	"github.com/golang/mock/gomock"
 	"github.com/stretchr/testify/require"
@@ -106,6 +108,57 @@ func (s *workflowShadowerActivitiesSuite) TestScanWorkflowActivity_Succeed() {
 	s.True(len(result.Executions) < numExecutions)
 }
 
+func (s *workflowShadowerActivitiesSuite) TestScanWorkflowActivity_MinResultSize() {
+	numExecutionsPerScan := 3
+	s.mockService.EXPECT().ScanWorkflowExecutions(gomock.Any(), gomock.Any(), callOptions...).Return(&shared.ListWorkflowExecutionsResponse{
+		Executions:    newTestWorkflowExecutions(numExecutionsPerScan),
+		NextPageToken: []byte{1, 2, 3},
+	}, nil).Times(int(math.Ceil(float64(minScanWorkflowResultSize) / float64(numExecutionsPerScan))))
+
+	params := shadower.ScanWorkflowActivityParams{
+		Domain:        common.StringPtr(defaultTestDomain),
+		WorkflowQuery: common.StringPtr("some random workflow visibility query"),
+		SamplingRate:  common.Float64Ptr(1),
+	}
+
+	resultValue, err := s.env.ExecuteActivity(shadower.ScanWorkflowActivityName, params)
+	s.NoError(err)
+
+	var result shadower.ScanWorkflowActivityResult
+	s.NoError(resultValue.Get(&result))
+	s.NotNil(result.NextPageToken)
+	s.True(len(result.Executions) >= minScanWorkflowResultSize)
+}
+
+func (s *workflowShadowerActivitiesSuite) TestScanWorkflowActivity_CompletionTime() {
+	activityTimeoutSeconds := int32(1)
+	s.mockService.EXPECT().ScanWorkflowExecutions(gomock.Any(), gomock.Any(), callOptions...).Return(&shared.ListWorkflowExecutionsResponse{
+		Executions:    newTestWorkflowExecutions(1),
+		NextPageToken: []byte{1, 2, 3},
+	}, nil).MaxTimes(int(time.Duration(activityTimeoutSeconds) * time.Second / scanWorkflowWaitPeriod))
+
+	params := shadower.ScanWorkflowActivityParams{
+		Domain:        common.StringPtr(defaultTestDomain),
+		WorkflowQuery: common.StringPtr("some random workflow visibility query"),
+		SamplingRate:  common.Float64Ptr(0.00000001),
+	}
+
+	resultValue, err := s.env.impl.executeActivityWithOptions(
+		activityOptions{
+			ScheduleToCloseTimeoutSeconds: activityTimeoutSeconds,
+			StartToCloseTimeoutSeconds:    activityTimeoutSeconds,
+		},
+		shadower.ScanWorkflowActivityName,
+		params,
+	)
+	s.NoError(err)
+
+	var result shadower.ScanWorkflowActivityResult
+	s.NoError(resultValue.Get(&result))
+	s.NotNil(result.NextPageToken)
+	s.Empty(result.Executions)
+}
+
 func (s *workflowShadowerActivitiesSuite) TestScanWorkflowActivity_InvalidQuery() {
 	s.mockService.EXPECT().ScanWorkflowExecutions(gomock.Any(), gomock.Any(), callOptions...).Return(nil, &shared.BadRequestError{
 		Message: "invalid query",