Improve operation collector logic (#39)

Encapsulate go routine management during operation collection.
Leverage wait groups and channel lifecycles to simplify collector interface.
Track start time in collector to simplify state management for operation poller.
Add unit tests for operation collector.

Author: astaticvoid

pkg/cloudmap/api.go

@@ -93,7 +93,6 @@ func (sdApi *serviceDiscoveryApi) ListServices(ctx context.Context, nsId string)
 		Name:   types.ServiceFilterNameNamespaceId,
 		Values: []string{nsId},
 	}
-	sdApi.log.Info("paginating", "nsId", nsId)
 
 	pages := sd.NewListServicesPaginator(sdApi.awsFacade, &sd.ListServicesInput{Filters: []types.ServiceFilter{filter}})
 

pkg/cloudmap/client.go

@@ -174,17 +174,15 @@ func (sdc *serviceDiscoveryClient) RegisterEndpoints(ctx context.Context, servic
 		return err
 	}
 
-	startTime := Now()
-	opCollector := NewOperationCollector(len(service.Endpoints))
+	opCollector := NewOperationCollector()
 
 	for _, endpt := range service.Endpoints {
-		go func(endpt *model.Endpoint) {
-			opId, endptErr := sdc.sdApi.RegisterInstance(ctx, svcId, endpt.Id, endpt.GetCloudMapAttributes())
-			opCollector.Add(endpt.Id, opId, endptErr)
-		}(endpt)
+		opCollector.Add(func() (opId string, err error) {
+			return sdc.sdApi.RegisterInstance(ctx, svcId, endpt.Id, endpt.GetCloudMapAttributes())
+		})
 	}
 
-	err = NewRegisterInstancePoller(sdc.sdApi, svcId, opCollector.Collect(), startTime).Poll(ctx)
+	err = NewRegisterInstancePoller(sdc.sdApi, svcId, opCollector.Collect(), opCollector.GetStartTime()).Poll(ctx)
 
 	// Evict cache entry so next list call reflects changes
 	sdc.evictEndpoints(svcId)
@@ -214,17 +212,15 @@ func (sdc *serviceDiscoveryClient) DeleteEndpoints(ctx context.Context, service
 		return err
 	}
 
-	startTime := Now()
-	opCollector := NewOperationCollector(len(service.Endpoints))
+	opCollector := NewOperationCollector()
 
 	for _, endpt := range service.Endpoints {
-		go func(endpt *model.Endpoint) {
-			opId, endptErr := sdc.sdApi.DeregisterInstance(ctx, svcId, endpt.Id)
-			opCollector.Add(endpt.Id, opId, endptErr)
-		}(endpt)
+		opCollector.Add(func() (opId string, err error) {
+			return sdc.sdApi.DeregisterInstance(ctx, svcId, endpt.Id)
+		})
 	}
 
-	err = NewDeregisterInstancePoller(sdc.sdApi, svcId, opCollector.Collect(), startTime).Poll(ctx)
+	err = NewDeregisterInstancePoller(sdc.sdApi, svcId, opCollector.Collect(), opCollector.GetStartTime()).Poll(ctx)
 
 	// Evict cache entry so next list call reflects changes
 	sdc.evictEndpoints(svcId)

pkg/cloudmap/operation_collector.go

@@ -3,52 +3,66 @@ package cloudmap
 import (
 	"github.com/go-logr/logr"
 	ctrl "sigs.k8s.io/controller-runtime"
+	"sync"
 )
 
-// OperationCollector collects a list operations with thread safety.
+// OperationCollector collects a list of operation IDs asynchronously with thread safety.
 type OperationCollector interface {
-	// Add an operation to poll with thread safety, to be called from within a go routine.
-	Add(endpointId string, operationId string, operationError error)
+	// Add calls an operation provider function to asynchronously collect operations to poll.
+	Add(operationProvider func() (operationId string, err error))
 
-	// Collect waits for all operations to be added and returns a list of successfully created operation IDs.
+	// Collect waits for all create operation results to be provided and returns a list of the successfully created operation IDs.
 	Collect() []string
 
+	// GetStartTime returns the start time range to poll the collected operations.
+	GetStartTime() int64
+
 	// IsAllOperationsCreated returns true if all operations were created successfully.
 	IsAllOperationsCreated() bool
 }
 
 type opCollector struct {
 	log              logr.Logger
 	opChan           chan opResult
-	opCount          int
+	wg               sync.WaitGroup
+	startTime        int64
 	createOpsSuccess bool
 }
 
 type opResult struct {
-	instId string
-	opId   string
-	err    error
+	opId string
+	err  error
 }
 
-func NewOperationCollector(opCount int) OperationCollector {
+func NewOperationCollector() OperationCollector {
 	return &opCollector{
 		log:              ctrl.Log.WithName("cloudmap"),
 		opChan:           make(chan opResult),
-		opCount:          opCount,
+		startTime:        Now(),
 		createOpsSuccess: true,
 	}
 }
 
-func (opColl *opCollector) Add(endptId string, opId string, opErr error) {
-	opColl.opChan <- opResult{endptId, opId, opErr}
+func (opColl *opCollector) Add(opProvider func() (opId string, err error)) {
+	opColl.wg.Add(1)
+	go func() {
+		defer opColl.wg.Done()
+
+		opId, opErr := opProvider()
+		opColl.opChan <- opResult{opId, opErr}
+	}()
 }
 
 func (opColl *opCollector) Collect() []string {
 	opIds := make([]string, 0)
 
-	for i := 0; i < opColl.opCount; i++ {
-		op := <-opColl.opChan
+	// Run wait in separate go routine to unblock reading from the channel.
+	go func() {
+		opColl.wg.Wait()
+		close(opColl.opChan)
+	}()
 
+	for op := range opColl.opChan {
 		if op.err != nil {
 			opColl.log.Info("could not create operation", "error", op.err)
 			opColl.createOpsSuccess = false
@@ -61,6 +75,10 @@ func (opColl *opCollector) Collect() []string {
 	return opIds
 }
 
+func (opColl *opCollector) GetStartTime() int64 {
+	return opColl.startTime
+}
+
 func (opColl *opCollector) IsAllOperationsCreated() bool {
 	return opColl.createOpsSuccess
 }

pkg/cloudmap/operation_collector_test.go

@@ -0,0 +1,51 @@
+package cloudmap
+
+import (
+	"errors"
+	"github.com/stretchr/testify/assert"
+	"testing"
+	"time"
+)
+
+func TestOpCollector_HappyCase(t *testing.T) {
+	oc := NewOperationCollector()
+	oc.Add(func() (opId string, err error) { return "one", nil })
+	oc.Add(func() (opId string, err error) { return "two", nil })
+
+	result := oc.Collect()
+	assert.True(t, oc.IsAllOperationsCreated())
+	assert.Equal(t, 2, len(result))
+	assert.Contains(t, result, "one")
+	assert.Contains(t, result, "two")
+}
+
+func TestOpCollector_AllFail(t *testing.T) {
+	oc := NewOperationCollector()
+	oc.Add(func() (opId string, err error) { return "one", errors.New("fail one") })
+	oc.Add(func() (opId string, err error) { return "two", errors.New("fail two") })
+
+	result := oc.Collect()
+	assert.False(t, oc.IsAllOperationsCreated())
+	assert.Equal(t, 0, len(result))
+}
+
+func TestOpCollector_MixedSuccess(t *testing.T) {
+	oc := NewOperationCollector()
+	oc.Add(func() (opId string, err error) { return "one", errors.New("fail one") })
+	oc.Add(func() (opId string, err error) { return "two", nil })
+
+	result := oc.Collect()
+	assert.False(t, oc.IsAllOperationsCreated())
+	assert.Equal(t, []string{"two"}, result)
+}
+
+func TestOpCollector_GetStartTime(t *testing.T) {
+	oc1 := NewOperationCollector()
+	time.Sleep(time.Second)
+	oc2 := NewOperationCollector()
+
+	assert.Equal(t, oc1.GetStartTime(), oc1.GetStartTime(), "Start time should not change")
+	assert.NotEqual(t, oc1.GetStartTime(), oc2.GetStartTime(), "Start time should reflect instantiation")
+	assert.Less(t, oc1.GetStartTime(), oc2.GetStartTime(),
+		"Start time should increase for later instantiations")
+}

pkg/cloudmap/operation_poller.go

@@ -33,10 +33,10 @@ type operationPoller struct {
 
 	svcId  string
 	opType types.OperationType
-	start  int
+	start  int64
 }
 
-func newOperationPoller(sdApi ServiceDiscoveryApi, svcId string, opIds []string, startTime int) operationPoller {
+func newOperationPoller(sdApi ServiceDiscoveryApi, svcId string, opIds []string, startTime int64) operationPoller {
 	return operationPoller{
 		log:   ctrl.Log.WithName("cloudmap"),
 		sdApi: sdApi,
@@ -48,14 +48,14 @@ func newOperationPoller(sdApi ServiceDiscoveryApi, svcId string, opIds []string,
 }
 
 // NewRegisterInstancePoller creates a new operation poller for register instance operations.
-func NewRegisterInstancePoller(sdApi ServiceDiscoveryApi, serviceId string, opIds []string, startTime int) OperationPoller {
+func NewRegisterInstancePoller(sdApi ServiceDiscoveryApi, serviceId string, opIds []string, startTime int64) OperationPoller {
 	poller := newOperationPoller(sdApi, serviceId, opIds, startTime)
 	poller.opType = types.OperationTypeRegisterInstance
 	return &poller
 }
 
 // NewDeregisterInstancePoller creates a new operation poller for de-register instance operations.
-func NewDeregisterInstancePoller(sdApi ServiceDiscoveryApi, serviceId string, opIds []string, startTime int) OperationPoller {
+func NewDeregisterInstancePoller(sdApi ServiceDiscoveryApi, serviceId string, opIds []string, startTime int64) OperationPoller {
 	poller := newOperationPoller(sdApi, serviceId, opIds, startTime)
 	poller.opType = types.OperationTypeDeregisterInstance
 	return &poller
@@ -124,9 +124,9 @@ func (opPoller *operationPoller) buildFilters() []types.OperationFilter {
 		Name:      types.OperationFilterNameUpdateDate,
 		Condition: types.FilterConditionBetween,
 		Values: []string{
-			strconv.Itoa(opPoller.start),
+			Itoa(opPoller.start),
 			// Add one minute to end range in case op updates while list request is in flight
-			strconv.Itoa(Now() + 60000),
+			Itoa(Now() + 60000),
 		},
 	}
 
@@ -143,8 +143,11 @@ func (opPoller *operationPoller) getFailedOpReason(ctx context.Context, opId str
 
 	return aws.ToString(op.ErrorMessage)
 }
+func Itoa(i int64) string {
+	return strconv.FormatInt(i, 10)
+}
 
-// Now returns current time with milliseconds, as used by operation UPDATE_DATE field
-func Now() int {
-	return int(time.Now().UnixNano() / 1000000)
+// Now returns current time with milliseconds, as used by operation filter UPDATE_DATE field
+func Now() int64 {
+	return time.Now().UnixNano() / 1000000
 }