fix: Move getRecords stop condition checking to top of loop

Added tests to validate coordination of the lease renewal with the
record getting and processing routines. This brought to light the value
of checking the loop's "stop conditions" early in the loop instead of at
the end.

Author: vmwjc

clientlibrary/worker/polling-shard-consumer.go

@@ -145,6 +145,16 @@ func (sc *PollingShardConsumer) getRecords() error {
 		leaseRenewalErrChan <- sc.renewLease(ctx)
 	}()
 	for {
+		select {
+		case <-*sc.stop:
+			shutdownInput := &kcl.ShutdownInput{ShutdownReason: kcl.REQUESTED, Checkpointer: recordCheckpointer}
+			sc.recordProcessor.Shutdown(shutdownInput)
+			return nil
+		case leaseRenewalErr := <-leaseRenewalErrChan:
+			return leaseRenewalErr
+		default:
+		}
+
 		getRecordsStartTime := time.Now()
 
 		log.Debugf("Trying to read %d record from iterator: %v", sc.kclConfig.MaxRecords, aws.ToString(shardIterator))
@@ -226,15 +236,6 @@ func (sc *PollingShardConsumer) getRecords() error {
 			time.Sleep(time.Duration(sc.kclConfig.IdleTimeBetweenReadsInMillis) * time.Millisecond)
 		}
 
-		select {
-		case <-*sc.stop:
-			shutdownInput := &kcl.ShutdownInput{ShutdownReason: kcl.REQUESTED, Checkpointer: recordCheckpointer}
-			sc.recordProcessor.Shutdown(shutdownInput)
-			return nil
-		case leaseRenewalErr := <-leaseRenewalErrChan:
-			return leaseRenewalErr
-		default:
-		}
 	}
 }
 

clientlibrary/worker/polling-shard-consumer_test.go

@@ -22,17 +22,21 @@ package worker
 import (
 	"context"
 	"errors"
+	"sync"
 	"testing"
 	"time"
 
 	"github.com/aws/aws-sdk-go-v2/aws"
 	"github.com/aws/aws-sdk-go-v2/service/kinesis"
+	"github.com/aws/aws-sdk-go-v2/service/kinesis/types"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/mock"
 	chk "github.com/vmware/vmware-go-kcl-v2/clientlibrary/checkpoint"
 	"github.com/vmware/vmware-go-kcl-v2/clientlibrary/config"
+	kcl "github.com/vmware/vmware-go-kcl-v2/clientlibrary/interfaces"
 	"github.com/vmware/vmware-go-kcl-v2/clientlibrary/metrics"
 	par "github.com/vmware/vmware-go-kcl-v2/clientlibrary/partition"
+	"github.com/vmware/vmware-go-kcl-v2/logger"
 )
 
 var (
@@ -199,7 +203,8 @@ func (m *MockKinesisSubscriberGetter) GetRecords(ctx context.Context, params *ki
 }
 
 func (m *MockKinesisSubscriberGetter) GetShardIterator(ctx context.Context, params *kinesis.GetShardIteratorInput, optFns ...func(*kinesis.Options)) (*kinesis.GetShardIteratorOutput, error) {
-	return nil, nil
+	ret := m.Called(ctx, params, optFns)
+	return ret.Get(0).(*kinesis.GetShardIteratorOutput), ret.Error(1)
 }
 
 func (m *MockKinesisSubscriberGetter) SubscribeToShard(ctx context.Context, params *kinesis.SubscribeToShardInput, optFns ...func(*kinesis.Options)) (*kinesis.SubscribeToShardOutput, error) {
@@ -455,6 +460,155 @@ func TestPollingShardConsumer_renewLease(t *testing.T) {
 	}
 }
 
+func TestPollingShardConsumer_getRecordsRenewLease(t *testing.T) {
+	log := logger.GetDefaultLogger()
+	type fields struct {
+		checkpointer chk.Checkpointer
+		kclConfig    *config.KinesisClientLibConfiguration
+		mService     metrics.MonitoringService
+	}
+	tests := []struct {
+		name            string
+		fields          fields
+
+		// testMillis must be at least 200ms or you'll trigger the localTPSExceededError
+		testMillis      time.Duration
+		expRenewalCalls int
+		expRenewals     int
+		shardClosed	bool
+		expErr          error
+	}{
+		{
+			"renew once",
+			fields{
+				&mockCheckpointer{},
+				&config.KinesisClientLibConfiguration{
+					LeaseRefreshWaitTime: 200,
+					Logger:               log,
+					InitialPositionInStream: config.LATEST,
+				},
+				&mockMetrics{},
+			},
+			250,
+			1,
+			1,
+			false,
+			nil,
+		},
+		{
+			"renew some",
+			fields{
+				&mockCheckpointer{},
+				&config.KinesisClientLibConfiguration{
+					LeaseRefreshWaitTime: 50,
+					Logger:               log,
+					InitialPositionInStream: config.LATEST,
+				},
+				&mockMetrics{},
+			},
+			50*5 + 10,
+			5,
+			5,
+			false,
+			nil,
+		},
+		{
+			"renew twice every 2.5 seconds",
+			fields{
+				&mockCheckpointer{},
+				&config.KinesisClientLibConfiguration{
+					LeaseRefreshWaitTime: 2500,
+					Logger:               log,
+					InitialPositionInStream: config.LATEST,
+				},
+				&mockMetrics{},
+			},
+			5100,
+			2,
+			2,
+			false,
+			nil,
+		},
+		{
+			"lease error",
+			fields{
+				&mockCheckpointer{fail: true},
+				&config.KinesisClientLibConfiguration{
+					LeaseRefreshWaitTime: 500,
+					Logger:               log,
+					InitialPositionInStream: config.LATEST,
+				},
+				&mockMetrics{},
+			},
+			1100,
+			1,
+			0,
+			false,
+			getLeaseTestFailure,
+		},
+	}
+	iterator := "test-iterator"
+	nextIt := "test-next-iterator"
+	millisBehind := int64(0)
+	stopChan := make(chan struct{})
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			mk := MockKinesisSubscriberGetter{}
+			gro := kinesis.GetRecordsOutput{
+				Records: []types.Record{
+					{
+						Data:                        []byte{},
+						PartitionKey:                new(string),
+						SequenceNumber:              new(string),
+						ApproximateArrivalTimestamp: &time.Time{},
+						EncryptionType:              "",
+					},
+				},
+				MillisBehindLatest: &millisBehind,
+			}
+			if !tt.shardClosed {
+				gro.NextShardIterator = &nextIt
+			}
+			mk.On("GetRecords", mock.Anything, mock.Anything, mock.Anything).Return(&gro, nil)
+			mk.On("GetShardIterator", mock.Anything, mock.Anything, mock.Anything).Return(&kinesis.GetShardIteratorOutput{ShardIterator: &iterator}, nil)
+			rp := mockRecordProcessor{
+				processDurationMillis: tt.testMillis,
+			}
+			sc := &PollingShardConsumer{
+				commonShardConsumer: commonShardConsumer{
+					shard: &par.ShardStatus{
+						ID:  "test-shard-id",
+						Mux: &sync.RWMutex{},
+					},
+					checkpointer:    tt.fields.checkpointer,
+					kclConfig:       tt.fields.kclConfig,
+					kc:              &mk,
+					recordProcessor: &rp,
+					mService: tt.fields.mService,
+				},
+				stop: &stopChan,
+				mService: tt.fields.mService,
+			}
+
+			// Send the stop signal a little before the total time it should
+			// take to get records and process them. This prevents test time
+			// errors due to the threads running longer than the test case
+			// expects.
+			go func() {
+				time.Sleep((tt.testMillis - 1) * time.Millisecond)
+				stopChan <- struct{}{}
+			}()
+
+			err := sc.getRecords()
+
+			assert.Equal(t, tt.expErr, err)
+			assert.Equal(t, tt.expRenewalCalls, sc.checkpointer.(*mockCheckpointer).getLeaseCalledTimes)
+			assert.Equal(t, tt.expRenewals, sc.mService.(*mockMetrics).leaseRenewedCalledTimes)
+			mk.AssertExpectations(t)
+		})
+	}
+}
+
 type mockCheckpointer struct {
 	getLeaseCalledTimes int
 	fail                bool
@@ -478,6 +632,16 @@ func (m mockCheckpointer) ListActiveWorkers(map[string]*par.ShardStatus) (map[st
 }
 func (m mockCheckpointer) ClaimShard(*par.ShardStatus, string) error { return nil }
 
+type mockRecordProcessor struct {
+	processDurationMillis time.Duration
+}
+
+func (m mockRecordProcessor) Initialize(initializationInput *kcl.InitializationInput) {}
+func (m mockRecordProcessor) ProcessRecords(processRecordsInput *kcl.ProcessRecordsInput) {
+	time.Sleep(time.Millisecond * m.processDurationMillis)
+}
+func (m mockRecordProcessor) Shutdown(shutdownInput *kcl.ShutdownInput) {}
+
 type mockMetrics struct {
 	leaseRenewedCalledTimes int
 }