Adds Aws tests

.github/workflows/coverage.yml

@@ -29,7 +29,7 @@ jobs:
         run: go generate ./...
 
       - name: Tests
-        run: go test -coverprofile=coverage.out -covermode=atomic ./...
+        run: go test -race -coverprofile=coverage.out -covermode=atomic ./...
 
       - name: Upload coverage to Codecov
         uses: codecov/codecov-action@v3

cmd/manager/main.go

@@ -28,7 +28,7 @@ func run() error {
 		return err
 	}
 
-	pool := compute.NewPool(logger, aws.NewAWSWorkerFactory(logger, ec2.NewFromConfig(cfg), aws.DefaultAWSInstanceParams, 443))
+	pool := compute.NewPool(logger, aws.NewWorkerFactory(logger, ec2.NewFromConfig(cfg), aws.DefaultInstanceParams, 443))
 	defer func(pool compute.Pool) {
 		_ = pool.Close()
 	}(pool)

go.mod

@@ -10,6 +10,7 @@ require (
 	github.com/aws/aws-sdk-go-v2/service/s3 v1.26.6
 	github.com/golang/mock v1.6.0
 	github.com/jaypipes/ghw v0.9.0
+	github.com/jaypipes/pcidb v1.0.0
 	github.com/launchdarkly/go-test-helpers/v2 v2.3.1
 	github.com/pkg/errors v0.9.1
 	github.com/xfrr/goffmpeg v0.0.0-20210624103149-5ca2d3062daf
@@ -38,7 +39,6 @@ require (
 	github.com/ghodss/yaml v1.0.0 // indirect
 	github.com/go-ole/go-ole v1.2.6 // indirect
 	github.com/golang/protobuf v1.5.2 // indirect
-	github.com/jaypipes/pcidb v1.0.0 // indirect
 	github.com/jmespath/go-jmespath v0.4.0 // indirect
 	github.com/mitchellh/go-homedir v1.1.0 // indirect
 	go.uber.org/atomic v1.7.0 // indirect

internal/compute/options.go

@@ -0,0 +1,22 @@
+package compute
+
+import "time"
+
+type ReadyOptions struct {
+	TickerInterval time.Duration
+	ConnTimeout    time.Duration
+}
+
+type ReadyOptionsFunc func(options *ReadyOptions)
+
+func WithTickerInterval(interval time.Duration) ReadyOptionsFunc {
+	return func(opts *ReadyOptions) {
+		opts.TickerInterval = interval
+	}
+}
+
+func WithConnTimeout(timeout time.Duration) ReadyOptionsFunc {
+	return func(opts *ReadyOptions) {
+		opts.ConnTimeout = timeout
+	}
+}

internal/compute/options_test.go

@@ -0,0 +1,54 @@
+package compute
+
+import (
+	"testing"
+	"time"
+
+	m "github.com/launchdarkly/go-test-helpers/v2/matchers"
+)
+
+func TestWithTickerInterval(t *testing.T) {
+	tests := []struct {
+		name     string
+		interval time.Duration
+	}{
+		{"0", time.Duration(0)},
+		{"seconds", 1 * time.Second},
+		{"ms", 100 * time.Millisecond},
+		{"hours", 2 * time.Hour},
+	}
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			f := WithTickerInterval(test.interval)
+
+			opt := new(ReadyOptions)
+
+			f(opt)
+
+			m.For(t, "val").Assert(opt.TickerInterval, m.Equal(test.interval))
+		})
+	}
+}
+
+func TestWithConnTimeout(t *testing.T) {
+	tests := []struct {
+		name     string
+		interval time.Duration
+	}{
+		{"0", time.Duration(0)},
+		{"seconds", 1 * time.Second},
+		{"ms", 100 * time.Millisecond},
+		{"hours", 2 * time.Hour},
+	}
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			f := WithConnTimeout(test.interval)
+
+			opt := new(ReadyOptions)
+
+			f(opt)
+
+			m.For(t, "val").Assert(opt.ConnTimeout, m.Equal(test.interval))
+		})
+	}
+}

internal/compute/work_queue.go

@@ -46,7 +46,10 @@ func NewQueue(logger *zap.Logger, pool Pool, maxSize int) WorkQueue {
 
 func (q *DefaultWorkQueue) run() {
 	for {
+		q.mtx.Lock()
 		if len(q.workers) < q.maxSize {
+			q.mtx.Unlock()
+
 			// Wait for work
 			work := <-q.workQueue
 
@@ -71,9 +74,11 @@ func (q *DefaultWorkQueue) run() {
 
 			go func() {
 				defer func(worker Worker) {
-					q.wg.Done()
+					q.mtx.Lock()
+					defer q.mtx.Lock()
 					q.workers = removeItem(q.workers, worker)
 					q.pool.ReturnWorker(worker)
+					q.wg.Done()
 				}(worker)
 
 				q.logger.Debug("Waiting for worker ready", zap.Any("req", work.getReq()))
@@ -98,6 +103,8 @@ func (q *DefaultWorkQueue) run() {
 				}
 				q.logger.Info("Work finished", zap.Any("req", work.getReq()))
 			}()
+		} else {
+			q.mtx.Unlock()
 		}
 	}
 }
@@ -117,5 +124,7 @@ func (q *DefaultWorkQueue) GetMaxSize() int {
 }
 
 func (q *DefaultWorkQueue) SetMaxSize(size int) {
+	q.mtx.Lock()
+	defer q.mtx.Unlock()
 	q.maxSize = size
 }

internal/compute/work_queue_test.go

@@ -134,8 +134,8 @@ func TestDefaultWorkQueue_Add(t *testing.T) {
 		AnyTimes()
 
 	mWorker.EXPECT().
-		IsReadyChan(gomock.Any()).
-		DoAndReturn(func(ctx context.Context) <-chan error {
+		IsReadyChan(gomock.Any(), gomock.Any()).
+		DoAndReturn(func(ctx context.Context, opts ...func(options *ReadyOptions)) <-chan error {
 			ch := make(chan error)
 			go func() {
 				ch <- nil

internal/compute/worker.go

@@ -23,8 +23,8 @@ type Worker interface {
 	Worker() proto.WorkerServiceClient
 	Job() proto.JobServiceClient
 
-	IsReady(ctx context.Context) (bool, error)
-	IsReadyChan(ctx context.Context) <-chan error
+	IsReady(ctx context.Context, opts ...ReadyOptionsFunc) (bool, error)
+	IsReadyChan(ctx context.Context, opts ...ReadyOptionsFunc) <-chan error
 }
 
 type WorkerFactory interface {

internal/worker/aws/aws.go

@@ -22,7 +22,7 @@ import (
 //go:embed userdata.sh
 var userData []byte
 
-var DefaultAWSInstanceParams = &ec2.RunInstancesInput{
+var DefaultInstanceParams = &ec2.RunInstancesInput{
 	MinCount:           aws.Int32(1),
 	MaxCount:           aws.Int32(1),
 	IamInstanceProfile: nil,
@@ -37,16 +37,16 @@ var DefaultAWSInstanceParams = &ec2.RunInstancesInput{
 	UserData:         aws.String(base64.StdEncoding.EncodeToString(userData)),
 }
 
-type AWSWorkerEC2Client interface {
+type WorkerEC2Client interface {
 	RunInstances(ctx context.Context, params *ec2.RunInstancesInput, optFns ...func(*ec2.Options)) (*ec2.RunInstancesOutput, error)
 	DescribeInstances(ctx context.Context, params *ec2.DescribeInstancesInput, optFns ...func(*ec2.Options)) (*ec2.DescribeInstancesOutput, error)
 	DescribeInstanceStatus(ctx context.Context, params *ec2.DescribeInstanceStatusInput, optFns ...func(*ec2.Options)) (*ec2.DescribeInstanceStatusOutput, error)
 	TerminateInstances(ctx context.Context, params *ec2.TerminateInstancesInput, optFns ...func(*ec2.Options)) (*ec2.TerminateInstancesOutput, error)
 }
 
-type AWSWorker struct {
+type Worker struct {
 	logger *zap.Logger
-	client AWSWorkerEC2Client
+	client WorkerEC2Client
 
 	id   string
 	port uint16
@@ -60,7 +60,7 @@ type AWSWorker struct {
 	closed bool
 }
 
-func (w *AWSWorker) Close() error {
+func (w *Worker) Close() error {
 	if w.closed {
 		return compute.ErrClosed
 	}
@@ -81,16 +81,16 @@ func (w *AWSWorker) Close() error {
 	return err
 }
 
-func (w *AWSWorker) Equals(other compute.Worker) bool {
+func (w *Worker) Equals(other compute.Worker) bool {
 	switch v := other.(type) {
-	case *AWSWorker:
+	case *Worker:
 		return w.id == v.id
 	default:
 		return false
 	}
 }
 
-func (w *AWSWorker) getIP(ctx context.Context) (netip.Addr, error) {
+func (w *Worker) getIP(ctx context.Context) (netip.Addr, error) {
 	instances, err := w.client.DescribeInstances(ctx, &ec2.DescribeInstancesInput{
 		InstanceIds: []string{w.id},
 	})
@@ -110,7 +110,7 @@ func (w *AWSWorker) getIP(ctx context.Context) (netip.Addr, error) {
 	return netip.ParseAddr(aws.ToString(instance.PublicIpAddress))
 }
 
-func (w *AWSWorker) Connect(ctx context.Context) (err error) {
+func (w *Worker) Connect(ctx context.Context) (err error) {
 	if w.closed {
 		return compute.ErrClosed
 	}
@@ -139,15 +139,15 @@ func (w *AWSWorker) Connect(ctx context.Context) (err error) {
 	return nil
 }
 
-func (w *AWSWorker) Worker() proto.WorkerServiceClient {
+func (w *Worker) Worker() proto.WorkerServiceClient {
 	return w.worker
 }
 
-func (w *AWSWorker) Job() proto.JobServiceClient {
+func (w *Worker) Job() proto.JobServiceClient {
 	return w.job
 }
 
-func (w *AWSWorker) getInstanceStatus(ctx context.Context) (types.InstanceStateName, error) {
+func (w *Worker) getInstanceStatus(ctx context.Context) (types.InstanceStateName, error) {
 	statuses, err := w.client.DescribeInstanceStatus(ctx, &ec2.DescribeInstanceStatusInput{
 		InstanceIds: []string{w.id},
 	})
@@ -167,7 +167,7 @@ func (w *AWSWorker) getInstanceStatus(ctx context.Context) (types.InstanceStateN
 	return types.InstanceStateNamePending, nil
 }
 
-func (w *AWSWorker) IsReady(ctx context.Context) (bool, error) {
+func (w *Worker) IsReady(ctx context.Context, opts ...compute.ReadyOptionsFunc) (bool, error) {
 	if w.closed {
 		return false, compute.ErrClosed
 	}
@@ -181,7 +181,15 @@ func (w *AWSWorker) IsReady(ctx context.Context) (bool, error) {
 		return false, nil
 	}
 
-	connCtx, cancel := context.WithTimeout(ctx, 10*time.Second)
+	options := &compute.ReadyOptions{
+		ConnTimeout: 10 * time.Second,
+	}
+
+	for _, opt := range opts {
+		opt(options)
+	}
+
+	connCtx, cancel := context.WithTimeout(ctx, options.ConnTimeout)
 	defer cancel()
 
 	err = w.Connect(connCtx)
@@ -196,22 +204,30 @@ func (w *AWSWorker) IsReady(ctx context.Context) (bool, error) {
 	return true, nil
 }
 
-func (w *AWSWorker) IsReadyChan(ctx context.Context) <-chan error {
+func (w *Worker) IsReadyChan(ctx context.Context, opts ...compute.ReadyOptionsFunc) <-chan error {
 	ch := make(chan error)
 
-	ticker := time.NewTicker(15 * time.Second)
+	options := &compute.ReadyOptions{
+		TickerInterval: 15 * time.Second,
+		ConnTimeout:    10 * time.Second,
+	}
+
+	for _, opt := range opts {
+		opt(options)
+	}
 
 	go func() {
+		ticker := time.NewTicker(options.TickerInterval)
 		for {
 			select {
 			case <-ctx.Done():
 				ch <- ctx.Err()
 				ticker.Stop()
 				return
 			case <-ticker.C:
-				isReady, err := w.IsReady(ctx)
+				isReady, err := w.IsReady(ctx, opts...)
 				if err != nil {
-					if err.Error() == "no instance statuses returned" {
+					if err.Error() == "instance not found" {
 						continue
 					}
 					ch <- err
@@ -221,6 +237,8 @@ func (w *AWSWorker) IsReadyChan(ctx context.Context) <-chan error {
 
 				if isReady {
 					ch <- nil
+					ticker.Stop()
+					return
 				}
 			}
 		}
@@ -229,23 +247,23 @@ func (w *AWSWorker) IsReadyChan(ctx context.Context) <-chan error {
 	return ch
 }
 
-type AWSWorkerFactory struct {
+type WorkerFactory struct {
 	logger *zap.Logger
-	client AWSWorkerEC2Client
+	client WorkerEC2Client
 	params *ec2.RunInstancesInput
 	port   uint16
 }
 
-func NewAWSWorkerFactory(logger *zap.Logger, client AWSWorkerEC2Client, input *ec2.RunInstancesInput, port uint16) compute.WorkerFactory {
-	return &AWSWorkerFactory{
+func NewWorkerFactory(logger *zap.Logger, client WorkerEC2Client, input *ec2.RunInstancesInput, port uint16) compute.WorkerFactory {
+	return &WorkerFactory{
 		logger: logger,
 		client: client,
 		params: input,
 		port:   port,
 	}
 }
 
-func (f *AWSWorkerFactory) Create(ctx context.Context) (compute.Worker, error) {
+func (f *WorkerFactory) Create(ctx context.Context) (compute.Worker, error) {
 	instances, err := f.client.RunInstances(ctx, f.params)
 	if err != nil {
 		return nil, err
@@ -257,7 +275,7 @@ func (f *AWSWorkerFactory) Create(ctx context.Context) (compute.Worker, error) {
 
 	instance := instances.Instances[0]
 
-	return &AWSWorker{
+	return &Worker{
 		logger: f.logger,
 		client: f.client,
 		id:     aws.ToString(instance.InstanceId),