diff --git a/lambda/main.go b/lambda/main.go index f11bec3..53cce3a 100644 --- a/lambda/main.go +++ b/lambda/main.go @@ -12,6 +12,7 @@ import ( "time" "github.com/aws/aws-lambda-go/lambda" + "github.com/aws/aws-sdk-go/aws/session" "github.com/buildkite/buildkite-agent-scaler/buildkite" "github.com/buildkite/buildkite-agent-scaler/scaler" "github.com/buildkite/buildkite-agent-scaler/version" @@ -114,6 +115,9 @@ func Handler(ctx context.Context, evt json.RawMessage) (string, error) { return vi } + // establish an AWS session to be re-used + sess := session.New() + for { select { case <-timeout: @@ -124,7 +128,7 @@ func Handler(ctx context.Context, evt json.RawMessage) (string, error) { if ssmTokenKey != "" { var err error - token, err = scaler.RetrieveFromParameterStore(ssmTokenKey) + token, err = scaler.RetrieveFromParameterStore(sess, ssmTokenKey) if err != nil { return "", err } @@ -170,7 +174,7 @@ func Handler(ctx context.Context, evt json.RawMessage) (string, error) { params.ScaleOutParams.Disable = true } - scaler, err := scaler.NewScaler(client, params) + scaler, err := scaler.NewScaler(client, sess, params) if err != nil { log.Fatal(err) } diff --git a/main.go b/main.go index df81545..c35b5ce 100644 --- a/main.go +++ b/main.go @@ -5,6 +5,7 @@ import ( "log" "time" + "github.com/aws/aws-sdk-go/aws/session" "github.com/buildkite/buildkite-agent-scaler/buildkite" "github.com/buildkite/buildkite-agent-scaler/scaler" ) @@ -23,16 +24,21 @@ func main() { includeWaiting = flag.Bool("include-waiting", false, "Whether to include jobs behind a wait step for scaling") // scale in/out params - scaleInFactor = flag.Float64("scale-in-factor", 1.0, "A factor to apply to scale ins") - scaleOutFactor = flag.Float64("scale-out-factor", 1.0, "A factor to apply to scale outs") + scaleInFactor = flag.Float64("scale-in-factor", 1.0, "A factor to apply to scale ins") + scaleOutFactor = flag.Float64("scale-out-factor", 1.0, "A factor to apply to scale outs") // general params dryRun = flag.Bool("dry-run", false, "Whether to just show what would be done") ) flag.Parse() + // establish an AWS session to be re-used + sess := session.Must(session.NewSessionWithOptions(session.Options{ + SharedConfigState: session.SharedConfigEnable, + })) + if *ssmTokenKey != "" { - token, err := scaler.RetrieveFromParameterStore(*ssmTokenKey) + token, err := scaler.RetrieveFromParameterStore(sess, *ssmTokenKey) if err != nil { log.Fatal(err) } @@ -41,7 +47,7 @@ func main() { client := buildkite.NewClient(*buildkiteAgentToken) - scaler, err := scaler.NewScaler(client, scaler.Params{ + scaler, err := scaler.NewScaler(client, sess, scaler.Params{ BuildkiteQueue: *buildkiteQueue, AutoScalingGroupName: *asgName, AgentsPerInstance: *agentsPerInstance, @@ -59,7 +65,7 @@ func main() { log.Printf("Running as a dry-run, no changes will be made") } - var interval time.Duration = 10 * time.Second; + var interval time.Duration = 10 * time.Second for { minPollDuration, err := scaler.Run() diff --git a/scaler/asg.go b/scaler/asg.go index dc421bb..530f03a 100644 --- a/scaler/asg.go +++ b/scaler/asg.go @@ -19,12 +19,13 @@ type AutoscaleGroupDetails struct { type asgDriver struct { name string + sess *session.Session } func (a *asgDriver) Describe() (AutoscaleGroupDetails, error) { log.Printf("Collecting AutoScaling details for ASG %q", a.name) - svc := autoscaling.New(session.New()) + svc := autoscaling.New(a.sess) input := &autoscaling.DescribeAutoScalingGroupsInput{ AutoScalingGroupNames: []*string{ aws.String(a.name), @@ -64,7 +65,7 @@ func (a *asgDriver) Describe() (AutoscaleGroupDetails, error) { } func (a *asgDriver) SetDesiredCapacity(count int64) error { - svc := autoscaling.New(session.New()) + svc := autoscaling.New(a.sess) input := &autoscaling.SetDesiredCapacityInput{ AutoScalingGroupName: aws.String(a.name), DesiredCapacity: aws.Int64(count), diff --git a/scaler/cloudwatch.go b/scaler/cloudwatch.go index 443b240..f418d97 100644 --- a/scaler/cloudwatch.go +++ b/scaler/cloudwatch.go @@ -14,11 +14,12 @@ const ( // cloudWatchMetricsPublisher sends queue metrics to AWS CloudWatch type cloudWatchMetricsPublisher struct { + sess *session.Session } // Publish queue metrics to CloudWatch Metrics func (cp *cloudWatchMetricsPublisher) Publish(orgSlug, queue string, metrics map[string]int64) error { - svc := cloudwatch.New(session.New()) + svc := cloudwatch.New(cp.sess) datum := []*cloudwatch.MetricDatum{} diff --git a/scaler/scaler.go b/scaler/scaler.go index a0802af..21592d2 100644 --- a/scaler/scaler.go +++ b/scaler/scaler.go @@ -5,6 +5,7 @@ import ( "math" "time" + "github.com/aws/aws-sdk-go/aws/session" "github.com/buildkite/buildkite-agent-scaler/buildkite" ) @@ -39,12 +40,12 @@ type Scaler struct { metrics interface { Publish(orgSlug, queue string, metrics map[string]int64) error } - scaling ScalingCalculator + scaling ScalingCalculator scaleInParams ScaleParams scaleOutParams ScaleParams } -func NewScaler(client *buildkite.Client, params Params) (*Scaler, error) { +func NewScaler(client *buildkite.Client, sess *session.Session, params Params) (*Scaler, error) { scaler := &Scaler{ bk: &buildkiteDriver{ client: client, @@ -68,6 +69,7 @@ func NewScaler(client *buildkite.Client, params Params) (*Scaler, error) { } else { scaler.autoscaling = &asgDriver{ name: params.AutoScalingGroupName, + sess: sess, } if params.PublishCloudWatchMetrics { diff --git a/scaler/ssm.go b/scaler/ssm.go index d58c54d..29778c1 100644 --- a/scaler/ssm.go +++ b/scaler/ssm.go @@ -6,8 +6,8 @@ import ( "github.com/aws/aws-sdk-go/service/ssm" ) -func RetrieveFromParameterStore(key string) (string, error) { - ssmClient := ssm.New(session.New()) +func RetrieveFromParameterStore(sess *session.Session, key string) (string, error) { + ssmClient := ssm.New(sess) output, err := ssmClient.GetParameter(&ssm.GetParameterInput{ Name: &key, WithDecryption: aws.Bool(true),