osbuild-worker: use aws sdk v2 for asg scale-in protection

This commit is contained in:
Sanne Raymaekers 2024-08-06 16:41:46 +02:00
parent 990ed6a9ad
commit 2624516f1a
4 changed files with 65 additions and 47 deletions

View file

@ -14,15 +14,11 @@ import (
"time"
"github.com/BurntSushi/toml"
"github.com/aws/aws-sdk-go/aws"
"github.com/aws/aws-sdk-go/aws/awserr"
"github.com/aws/aws-sdk-go/aws/ec2metadata"
"github.com/aws/aws-sdk-go/aws/session"
"github.com/aws/aws-sdk-go/service/autoscaling"
"github.com/sirupsen/logrus"
"github.com/osbuild/images/pkg/arch"
"github.com/osbuild/images/pkg/dnfjson"
"github.com/osbuild/osbuild-composer/internal/cloud/awscloud"
"github.com/osbuild/osbuild-composer/internal/upload/azure"
"github.com/osbuild/osbuild-composer/internal/upload/koji"
"github.com/osbuild/osbuild-composer/internal/upload/oci"
@ -96,61 +92,27 @@ func WatchJob(ctx context.Context, job worker.Job) {
}
}
// protect an AWS instance from scaling and/or terminating.
// Protect an AWS instance from scaling (terminating).
func setProtection(protected bool) {
// create a new session
awsSession, err := session.NewSession()
// This will fail if the worker isn't running in AWS, so just return with a debug message.
region, err := awscloud.RegionFromInstanceMetadata()
if err != nil {
logrus.Debugf("Error getting an AWS session, %s", err)
logrus.Debugf("Error getting the instance region: %v", err)
return
}
// get the identity for the instanceID
identity, err := ec2metadata.New(awsSession).GetInstanceIdentityDocument()
aws, err := awscloud.NewDefault(region)
if err != nil {
logrus.Debugf("Error getting the identity document, %s", err)
logrus.Infof("Unable to get default aws client: %v", err)
return
}
svc := autoscaling.New(awsSession, aws.NewConfig().WithRegion(identity.Region))
// get the autoscaling group info for the auto scaling group name
asInstanceInput := &autoscaling.DescribeAutoScalingInstancesInput{
InstanceIds: []*string{
aws.String(identity.InstanceID),
},
}
asInstanceOutput, err := svc.DescribeAutoScalingInstances(asInstanceInput)
err = aws.ASGSetProtectHost(protected)
if err != nil {
if aerr, ok := err.(awserr.Error); ok {
logrus.Warningf("Error getting the Autoscaling instances: %s %s", aerr.Code(), aerr.Error())
} else {
logrus.Errorf("Error getting the Autoscaling instances: unknown, %s", err)
}
return
}
if len(asInstanceOutput.AutoScalingInstances) == 0 {
logrus.Info("No Autoscaling instace is defined")
logrus.Infof("Unable to protect host, if the host isn't running as part of an autoscaling group, this can safely be ignored: %v", err)
return
}
// make the request to protect (or unprotect) the instance
input := &autoscaling.SetInstanceProtectionInput{
AutoScalingGroupName: asInstanceOutput.AutoScalingInstances[0].AutoScalingGroupName,
InstanceIds: []*string{
aws.String(identity.InstanceID),
},
ProtectedFromScaleIn: aws.Bool(protected),
}
_, err = svc.SetInstanceProtection(input)
if err != nil {
if aerr, ok := err.(awserr.Error); ok {
logrus.Warningf("Error protecting instance: %s %s", aerr.Code(), aerr.Error())
} else {
logrus.Errorf("Error protecting instance: unknown, %s", err)
}
return
}
if protected {
logrus.Info("Instance protected")
} else {

View file

@ -0,0 +1,45 @@
package awscloud
import (
"context"
"github.com/aws/aws-sdk-go-v2/aws"
"github.com/aws/aws-sdk-go-v2/feature/ec2/imds"
"github.com/aws/aws-sdk-go-v2/service/autoscaling"
)
func (a *AWS) ASGSetProtectHost(protect bool) error {
identity, err := a.ec2imds.GetInstanceIdentityDocument(context.Background(), &imds.GetInstanceIdentityDocumentInput{})
if err != nil {
return err
}
descrASG, err := a.asg.DescribeAutoScalingInstances(
context.Background(),
&autoscaling.DescribeAutoScalingInstancesInput{
InstanceIds: []string{
identity.InstanceID,
},
},
)
if err != nil {
return err
}
if len(descrASG.AutoScalingInstances) == 0 {
return nil
}
_, err = a.asg.SetInstanceProtection(
context.Background(),
&autoscaling.SetInstanceProtectionInput{
AutoScalingGroupName: descrASG.AutoScalingInstances[0].AutoScalingGroupName,
InstanceIds: []string{
identity.InstanceID,
},
ProtectedFromScaleIn: aws.Bool(protect),
},
)
return err
}

View file

@ -13,6 +13,7 @@ import (
"github.com/aws/aws-sdk-go-v2/credentials"
"github.com/aws/aws-sdk-go-v2/feature/ec2/imds"
"github.com/aws/aws-sdk-go-v2/feature/s3/manager"
"github.com/aws/aws-sdk-go-v2/service/autoscaling"
"github.com/aws/aws-sdk-go-v2/service/ec2"
ec2types "github.com/aws/aws-sdk-go-v2/service/ec2/types"
"github.com/aws/aws-sdk-go-v2/service/s3"
@ -26,6 +27,7 @@ type AWS struct {
s3 S3
s3uploader S3Manager
s3presign S3Presign
asg ASG
}
func newForTest(ec2cli EC2, ec2imds EC2Imds, s3cli S3, upldr S3Manager, sign S3Presign) *AWS {
@ -35,6 +37,7 @@ func newForTest(ec2cli EC2, ec2imds EC2Imds, s3cli S3, upldr S3Manager, sign S3P
s3: s3cli,
s3uploader: upldr,
s3presign: sign,
asg: nil,
}
}
@ -48,6 +51,7 @@ func newAwsFromConfig(cfg aws.Config) *AWS {
s3: s3cli,
s3uploader: manager.NewUploader(s3cli),
s3presign: s3.NewPresignClient(s3cli),
asg: autoscaling.NewFromConfig(cfg),
}
}
@ -159,6 +163,7 @@ func newAwsFromCredsWithEndpoint(creds config.LoadOptionsFunc, region, endpoint,
s3: s3cli,
s3uploader: manager.NewUploader(s3cli),
s3presign: s3.NewPresignClient(s3cli),
asg: autoscaling.NewFromConfig(cfg),
}, nil
}

View file

@ -6,6 +6,7 @@ import (
"github.com/aws/aws-sdk-go-v2/aws/signer/v4"
"github.com/aws/aws-sdk-go-v2/feature/ec2/imds"
"github.com/aws/aws-sdk-go-v2/feature/s3/manager"
"github.com/aws/aws-sdk-go-v2/service/autoscaling"
"github.com/aws/aws-sdk-go-v2/service/ec2"
"github.com/aws/aws-sdk-go-v2/service/s3"
)
@ -60,6 +61,11 @@ type EC2Imds interface {
GetInstanceIdentityDocument(context.Context, *imds.GetInstanceIdentityDocumentInput, ...func(*imds.Options)) (*imds.GetInstanceIdentityDocumentOutput, error)
}
type ASG interface {
DescribeAutoScalingInstances(context.Context, *autoscaling.DescribeAutoScalingInstancesInput, ...func(*autoscaling.Options)) (*autoscaling.DescribeAutoScalingInstancesOutput, error)
SetInstanceProtection(context.Context, *autoscaling.SetInstanceProtectionInput, ...func(*autoscaling.Options)) (*autoscaling.SetInstanceProtectionOutput, error)
}
type S3 interface {
DeleteObject(context.Context, *s3.DeleteObjectInput, ...func(*s3.Options)) (*s3.DeleteObjectOutput, error)
PutObjectAcl(context.Context, *s3.PutObjectAclInput, ...func(*s3.Options)) (*s3.PutObjectAclOutput, error)