Xray initial support

Gopkg.toml

@@ -20,6 +20,7 @@
 #  name = "github.com/x/y"
 #  version = "2.4.0"
 
+ignored = ["github.com/cristim/autospotting/core"]
 
 [[constraint]]
   name = "github.com/aws/aws-sdk-go"
@@ -37,6 +38,10 @@
   name = "github.com/namsral/flag"
   branch = "master"
 
+[[constraint]]
+  name = "github.com/aws/aws-xray-sdk-go"
+  version = "1.0.0-rc.5"
+
 [prune]
   go-tests = true
   non-go = true

autospotting.go

@@ -1,13 +1,15 @@
 package main
 
 import (
+	"context"
 	"fmt"
 	"log"
 	"os"
 
-	"github.com/aws/aws-lambda-go/events"
 	"github.com/aws/aws-lambda-go/lambda"
 	"github.com/aws/aws-sdk-go/aws/endpoints"
+	"github.com/aws/aws-xray-sdk-go/strategy/ctxmissing"
+	"github.com/aws/aws-xray-sdk-go/xray"
 	"github.com/cristim/autospotting/core"
 	"github.com/cristim/ec2-instances-info"
 	"github.com/namsral/flag"
@@ -26,12 +28,11 @@ func main() {
 	if os.Getenv("AWS_LAMBDA_FUNCTION_NAME") != "" {
 		lambda.Start(Handler)
 	} else {
-		run()
+		run(context.Background())
 	}
 }
 
-func run() {
-
+func run(ctx context.Context) {
 	log.Println("Starting autospotting agent, build", Version)
 
 	log.Printf("Parsed command line flags: "+
@@ -58,7 +59,11 @@ func run() {
 		conf.TagFilteringMode,
 		conf.SpotProductDescription)
 
-	autospotting.Run(conf.Config)
+	xray.Configure(xray.Config{
+		ContextMissingStrategy: ctxmissing.NewDefaultLogErrorStrategy(),
+		LogLevel:               "error",
+	})
+	autospotting.Run(ctx, conf.Config)
 	log.Println("Execution completed, nothing left to do")
 }
 
@@ -87,8 +92,8 @@ func init() {
 }
 
 // Handler implements the AWS Lambda handler
-func Handler(request events.APIGatewayProxyRequest) {
-	run()
+func Handler(ctx context.Context) {
+	run(ctx)
 }
 
 // Configuration handling

core/autoscaling.go

@@ -1,8 +1,8 @@
 package autospotting
 
 import (
+	"context"
 	"errors"
-
 	"math"
 	"strconv"
 	"strings"
@@ -249,7 +249,7 @@ func (a *autoScalingGroup) loadDefaultConfig() bool {
 	return done
 }
 
-func (a *autoScalingGroup) loadLaunchConfiguration() error {
+func (a *autoScalingGroup) loadLaunchConfiguration(ctx context.Context) error {
 	//already done
 	if a.launchConfiguration != nil {
 		return nil
@@ -266,7 +266,7 @@ func (a *autoScalingGroup) loadLaunchConfiguration() error {
 	params := &autoscaling.DescribeLaunchConfigurationsInput{
 		LaunchConfigurationNames: []*string{lcName},
 	}
-	resp, err := svc.DescribeLaunchConfigurations(params)
+	resp, err := svc.DescribeLaunchConfigurationsWithContext(ctx, params)
 
 	if err != nil {
 		logger.Println(err.Error())
@@ -279,7 +279,7 @@ func (a *autoScalingGroup) loadLaunchConfiguration() error {
 	return nil
 }
 
-func (a *autoScalingGroup) needReplaceOnDemandInstances() bool {
+func (a *autoScalingGroup) needReplaceOnDemandInstances(ctx context.Context) bool {
 	onDemandRunning, totalRunning := a.alreadyRunningInstanceCount(false, "")
 	if onDemandRunning > a.minOnDemand {
 		logger.Println("Currently more than enough OnDemand instances running")
@@ -298,7 +298,7 @@ func (a *autoScalingGroup) needReplaceOnDemandInstances() bool {
 			} else {
 				logger.Println("Terminating a random spot instance",
 					*randomSpot.Instance.InstanceId)
-				randomSpot.terminate()
+				randomSpot.terminate(ctx)
 			}
 		}
 	}
@@ -310,7 +310,7 @@ func (a *autoScalingGroup) allInstanceRunning() bool {
 	return totalRunning == a.instances.count64()
 }
 
-func (a *autoScalingGroup) process() {
+func (a *autoScalingGroup) process(ctx context.Context) {
 	var spotInstanceID string
 	a.scanInstances()
 	a.loadDefaultConfig()
@@ -331,7 +331,7 @@ func (a *autoScalingGroup) process() {
 				"No running on-demand instances were found, nothing to do here...")
 			return
 		}
-		a.loadLaunchConfiguration()
+		a.loadLaunchConfiguration(context.Background())
 		err := onDemandInstance.launchSpotReplacement()
 		if err != nil {
 			logger.Printf("Could not launch cheapest spot instance: %s", err)
@@ -341,15 +341,15 @@ func (a *autoScalingGroup) process() {
 
 	spotInstanceID = *spotInstance.InstanceId
 
-	if !a.needReplaceOnDemandInstances() || !spotInstance.isReadyToAttach(a) {
+	if !a.needReplaceOnDemandInstances(ctx) || !spotInstance.isReadyToAttach(a) {
 		logger.Println("Waiting for next run while processing", a.name)
 		return
 	}
 
 	logger.Println(a.region.name, "Found spot instance:", spotInstanceID,
 		"Attaching it to", a.name)
 
-	a.replaceOnDemandInstanceWithSpot(spotInstanceID)
+	a.replaceOnDemandInstanceWithSpot(ctx, spotInstanceID)
 
 }
 
@@ -385,7 +385,7 @@ func (a *autoScalingGroup) scanInstances() instances {
 	return a.instances
 }
 
-func (a *autoScalingGroup) replaceOnDemandInstanceWithSpot(
+func (a *autoScalingGroup) replaceOnDemandInstanceWithSpot(ctx context.Context,
 	spotInstanceID string) error {
 
 	minSize, maxSize := *a.MinSize, *a.MaxSize
@@ -394,8 +394,8 @@ func (a *autoScalingGroup) replaceOnDemandInstanceWithSpot(
 	// temporarily increase AutoScaling group in case it's of static size
 	if minSize == maxSize {
 		logger.Println(a.name, "Temporarily increasing MaxSize")
-		a.setAutoScalingMaxSize(maxSize + 1)
-		defer a.setAutoScalingMaxSize(maxSize)
+		a.setAutoScalingMaxSize(ctx, maxSize+1)
+		defer a.setAutoScalingMaxSize(ctx, maxSize)
 	}
 
 	// get the details of our spot instance so we can see its AZ
@@ -416,24 +416,24 @@ func (a *autoScalingGroup) replaceOnDemandInstanceWithSpot(
 		logger.Println(a.name, "found no on-demand instances that could be",
 			"replaced with the new spot instance", *spotInst.InstanceId,
 			"terminating the spot instance.")
-		spotInst.terminate()
+		spotInst.terminate(context.Background())
 		return errors.New("couldn't find ondemand instance to replace")
 	}
 	logger.Println(a.name, "found on-demand instance", *odInst.InstanceId,
 		"replacing with new spot instance", *spotInst.InstanceId)
 	// revert attach/detach order when running on minimum capacity
 	if desiredCapacity == minSize {
-		attachErr := a.attachSpotInstance(spotInstanceID)
+		attachErr := a.attachSpotInstance(ctx, spotInstanceID)
 		if attachErr != nil {
 			logger.Println(a.name, "skipping detaching on-demand due to failure to",
 				"attach the new spot instance", *spotInst.InstanceId)
 			return nil
 		}
 	} else {
-		defer a.attachSpotInstance(spotInstanceID)
+		defer a.attachSpotInstance(ctx, spotInstanceID)
 	}
 
-	return a.detachAndTerminateOnDemandInstance(odInst.InstanceId)
+	return a.detachAndTerminateOnDemandInstance(ctx, odInst.InstanceId)
 }
 
 // Returns the information about the first running instance found in
@@ -558,10 +558,11 @@ func (a *autoScalingGroup) getDisallowedInstanceTypes(baseInstance *instance) []
 	})
 }
 
-func (a *autoScalingGroup) setAutoScalingMaxSize(maxSize int64) error {
+func (a *autoScalingGroup) setAutoScalingMaxSize(ctx context.Context, maxSize int64) error {
 	svc := a.region.services.autoScaling
 
-	_, err := svc.UpdateAutoScalingGroup(
+	_, err := svc.UpdateAutoScalingGroupWithContext(
+		ctx,
 		&autoscaling.UpdateAutoScalingGroupInput{
 			AutoScalingGroupName: aws.String(a.name),
 			MaxSize:              aws.Int64(maxSize),
@@ -576,7 +577,7 @@ func (a *autoScalingGroup) setAutoScalingMaxSize(maxSize int64) error {
 	return nil
 }
 
-func (a *autoScalingGroup) attachSpotInstance(spotInstanceID string) error {
+func (a *autoScalingGroup) attachSpotInstance(ctx context.Context, spotInstanceID string) error {
 
 	svc := a.region.services.autoScaling
 
@@ -587,7 +588,7 @@ func (a *autoScalingGroup) attachSpotInstance(spotInstanceID string) error {
 		},
 	}
 
-	resp, err := svc.AttachInstances(&params)
+	resp, err := svc.AttachInstancesWithContext(ctx, &params)
 
 	if err != nil {
 		logger.Println(err.Error())
@@ -600,7 +601,7 @@ func (a *autoScalingGroup) attachSpotInstance(spotInstanceID string) error {
 
 // Terminates an on-demand instance from the group,
 // but only after it was detached from the autoscaling group
-func (a *autoScalingGroup) detachAndTerminateOnDemandInstance(
+func (a *autoScalingGroup) detachAndTerminateOnDemandInstance(ctx context.Context,
 	instanceID *string) error {
 	logger.Println(a.region.name,
 		a.name,
@@ -617,15 +618,15 @@ func (a *autoScalingGroup) detachAndTerminateOnDemandInstance(
 
 	asSvc := a.region.services.autoScaling
 
-	if _, err := asSvc.DetachInstances(&detachParams); err != nil {
+	if _, err := asSvc.DetachInstancesWithContext(ctx, &detachParams); err != nil {
 		logger.Println(err.Error())
 		return err
 	}
 
 	// Wait till detachment initialize is complete before terminate instance
 	time.Sleep(20 * time.Second * a.region.conf.SleepMultiplier)
 
-	return a.instances.get(*instanceID).terminate()
+	return a.instances.get(*instanceID).terminate(context.Background())
 }
 
 // Counts the number of already running instances on-demand or spot, in any or a specific AZ.

core/autoscaling_test.go

@@ -1,6 +1,7 @@
 package autospotting
 
 import (
+	"context"
 	"errors"
 	"reflect"
 	"testing"
@@ -1380,7 +1381,7 @@ func TestNeedReplaceOnDemandInstances(t *testing.T) {
 			a.DesiredCapacity = tt.desiredCapacity
 			a.instances = tt.asgInstances
 			a.minOnDemand = tt.minOnDemand
-			shouldRun := a.needReplaceOnDemandInstances()
+			shouldRun := a.needReplaceOnDemandInstances(context.Background())
 			if tt.expectedRun != shouldRun {
 				t.Errorf("needReplaceOnDemandInstances returned: %t expected %t",
 					shouldRun, tt.expectedRun)
@@ -1518,7 +1519,7 @@ func TestDetachAndTerminateOnDemandInstance(t *testing.T) {
 				region:    tt.regionASG,
 				instances: tt.instancesASG,
 			}
-			err := a.detachAndTerminateOnDemandInstance(tt.instanceID)
+			err := a.detachAndTerminateOnDemandInstance(context.Background(), tt.instanceID)
 			CheckErrors(t, err, tt.expected)
 		})
 	}
@@ -1558,7 +1559,7 @@ func TestAttachSpotInstance(t *testing.T) {
 				name:   "testASG",
 				region: tt.regionASG,
 			}
-			err := a.attachSpotInstance(tt.instanceID)
+			err := a.attachSpotInstance(context.Background(), tt.instanceID)
 			CheckErrors(t, err, tt.expected)
 		})
 	}
@@ -1629,7 +1630,7 @@ func TestLoadLaunchConfiguration(t *testing.T) {
 					LaunchConfigurationName: tt.nameLC,
 				},
 			}
-			err := a.loadLaunchConfiguration()
+			err := a.loadLaunchConfiguration(context.Background())
 			lc := a.launchConfiguration
 
 			if !reflect.DeepEqual(tt.expectedErr, err) {
@@ -1681,7 +1682,7 @@ func TestSetAutoScalingMaxSize(t *testing.T) {
 				name:   "testASG",
 				region: tt.regionASG,
 			}
-			err := a.setAutoScalingMaxSize(tt.maxSize)
+			err := a.setAutoScalingMaxSize(context.Background(), tt.maxSize)
 			CheckErrors(t, err, tt.expected)
 		})
 	}
@@ -2572,7 +2573,7 @@ func TestReplaceOnDemandInstanceWithSpot(t *testing.T) {
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
-			returned := tt.asg.replaceOnDemandInstanceWithSpot(tt.spotID)
+			returned := tt.asg.replaceOnDemandInstanceWithSpot(context.Background(), tt.spotID)
 			CheckErrors(t, returned, tt.expected)
 		})
 	}

core/connections.go

@@ -10,6 +10,7 @@ import (
 	"github.com/aws/aws-sdk-go/service/autoscaling/autoscalingiface"
 	"github.com/aws/aws-sdk-go/service/ec2"
 	"github.com/aws/aws-sdk-go/service/ec2/ec2iface"
+	"github.com/aws/aws-xray-sdk-go/xray"
 )
 
 type connections struct {
@@ -35,8 +36,16 @@ func (c *connections) connect(region string) {
 	asConn := make(chan *autoscaling.AutoScaling)
 	ec2Conn := make(chan *ec2.EC2)
 
-	go func() { asConn <- autoscaling.New(c.session) }()
-	go func() { ec2Conn <- ec2.New(c.session) }()
+	go func() {
+		c := autoscaling.New(c.session)
+		xray.AWS(c.Client)
+		asConn <- c
+	}()
+	go func() {
+		c := ec2.New(c.session)
+		xray.AWS(c.Client)
+		ec2Conn <- c
+	}()
 
 	c.autoScaling, c.ec2, c.region = <-asConn, <-ec2Conn, region