Last active
November 29, 2023 09:56
-
-
Save jicowan/ad5e13d12577b41a22f83ed91a3e61bf to your computer and use it in GitHub Desktop.
ECS Fargate tasks that are stopped by Spot interruptions are not deregistered from load balancers automatically. This function deregister an ECS Fargate Spot task from an AWS load balancer when it is interrupted..
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
package main | |
import ( | |
"context" | |
"github.com/aws/aws-lambda-go/lambda" | |
"github.com/aws/aws-sdk-go-v2/aws" | |
"github.com/aws/aws-sdk-go-v2/config" | |
"github.com/aws/aws-sdk-go-v2/service/ec2" | |
"github.com/aws/aws-sdk-go-v2/service/ecs" | |
"github.com/aws/aws-sdk-go-v2/service/elasticloadbalancingv2" | |
"github.com/aws/aws-sdk-go-v2/service/elasticloadbalancingv2/types" | |
"log" | |
"strings" | |
"time" | |
) | |
type ECSEvent struct { | |
Version string `json:"version,omitempty"` | |
ID string `json:"id,omitempty"` | |
DetailType string `json:"detail-type,omitempty"` | |
Source string `json:"source,omitempty"` | |
Account string `json:"account,omitempty"` | |
Time time.Time `json:"time,omitempty"` | |
Region string `json:"region,omitempty"` | |
Resources []string `json:"resources,omitempty"` | |
Detail struct { | |
ClusterArn string `json:"clusterArn,omitempty"` | |
Containers []struct { | |
ContainerArn string `json:"containerArn,omitempty"` | |
LastStatus string `json:"lastStatus,omitempty"` | |
Name string `json:"name,omitempty"` | |
TaskArn string `json:"taskArn,omitempty"` | |
NetworkInterfaces []struct { | |
AttachmentID string `json:"attachmentId,omitempty"` | |
PrivateIpv4Address string `json:"privateIpv4Address,omitempty"` | |
} `json:"networkInterfaces,omitempty"` | |
CPU string `json:"cpu,omitempty"` | |
Memory string `json:"memory,omitempty"` | |
} `json:"containers,omitempty"` | |
CreatedAt time.Time `json:"createdAt,omitempty"` | |
LaunchType string `json:"launchType,omitempty"` | |
CPU string `json:"cpu,omitempty"` | |
Memory string `json:"memory,omitempty"` | |
DesiredStatus string `json:"desiredStatus,omitempty"` | |
Group string `json:"group,omitempty"` | |
LastStatus string `json:"lastStatus,omitempty"` | |
Overrides struct { | |
ContainerOverrides []struct { | |
Name string `json:"name,omitempty"` | |
} `json:"containerOverrides,omitempty"` | |
} `json:"overrides,omitempty"` | |
Attachments []struct { | |
ID string `json:"id,omitempty"` | |
Type string `json:"type,omitempty"` | |
Status string `json:"status,omitempty"` | |
Details []struct { | |
Name string `json:"name,omitempty"` | |
Value string `json:"value,omitempty"` | |
} `json:"details,omitempty"` | |
} `json:"attachments,omitempty"` | |
Connectivity string `json:"connectivity,omitempty"` | |
ConnectivityAt time.Time `json:"connectivityAt,omitempty"` | |
PullStartedAt time.Time `json:"pullStartedAt,omitempty"` | |
StartedAt time.Time `json:"startedAt,omitempty"` | |
StartedBy string `json:"startedBy,omitempty"` | |
StoppingAt time.Time `json:"stoppingAt,omitempty"` | |
PullStoppedAt time.Time `json:"pullStoppedAt,omitempty"` | |
StoppedReason string `json:"stoppedReason,omitempty"` | |
StopCode string `json:"stopCode,omitempty"` | |
UpdatedAt time.Time `json:"updatedAt,omitempty"` | |
TaskArn string `json:"taskArn,omitempty"` | |
TaskDefinitionArn string `json:"taskDefinitionArn,omitempty"` | |
Version int `json:"version,omitempty"` | |
PlatformVersion string `json:"platformVersion,omitempty"` | |
} `json:"detail,omitempty"` | |
} | |
func main() { | |
lambda.Start(HandleRequest) | |
} | |
func HandleRequest(e ECSEvent) error { | |
var privateIPv4Address string | |
var subnetId []string | |
var service []string | |
for _, attachment := range e.Detail.Attachments { | |
if e.Detail.StopCode == "TerminationNotice" && strings.Contains(e.Detail.Group, "service:") { | |
for _, detail := range attachment.Details { | |
if detail.Name == "privateIPv4Address" { | |
privateIPv4Address = detail.Value | |
} | |
if detail.Name == "subnetId" { | |
subnetId = append(subnetId, detail.Value) | |
} | |
} | |
cluster := e.Detail.ClusterArn | |
service = append(service, strings.Split(e.Detail.Group, ":")[1]) | |
targetGroup := getTargetGroup(service, cluster) | |
az := getAvailabilityZone(subnetId) | |
deregisterTask(&privateIPv4Address, az, targetGroup, nil) | |
} | |
} | |
return nil | |
} | |
func getAvailabilityZone(subnetId []string) *string { | |
ctx := context.Background() | |
config, err := config.LoadDefaultConfig(context.TODO()) | |
if err != nil { | |
log.Fatal(err) | |
} | |
var az *string | |
client := ec2.NewFromConfig(config) | |
output, err := client.DescribeSubnets(ctx, &ec2.DescribeSubnetsInput{SubnetIds: subnetId}) | |
if err != nil { | |
log.Println(err) | |
} | |
for _, subnet := range output.Subnets { | |
az = subnet.AvailabilityZone | |
} | |
return az | |
} | |
func deregisterTask(ip *string, az *string, tg *string, port *int32) { | |
ctx := context.Background() | |
config, err := config.LoadDefaultConfig(context.TODO()) | |
if err != nil { | |
log.Fatal(err) | |
} | |
client := elasticloadbalancingv2.NewFromConfig(config) | |
params := &elasticloadbalancingv2.DeregisterTargetsInput{ | |
TargetGroupArn: tg, | |
Targets: []types.TargetDescription{ | |
{ | |
Id: ip, | |
AvailabilityZone: az, | |
Port: port, | |
}, | |
}, | |
} | |
_, err = client.DeregisterTargets(ctx, params) | |
if err != nil { | |
log.Println(err) | |
} else { | |
log.Printf("The target %v was deregistered\n", aws.ToString(ip)) | |
} | |
} | |
func getTargetGroup(svc []string, cluster string) *string { | |
log.Printf("The service name is: %v", svc[0]) | |
log.Printf("The clusterArn is: %v\n", cluster) | |
ctx := context.Background() | |
config, err := config.LoadDefaultConfig(context.TODO()) | |
if err != nil { | |
log.Fatal(err) | |
} | |
var tg *string | |
client := ecs.NewFromConfig(config) | |
log.Printf("Finding target group\n") | |
output, err := client.DescribeServices(ctx, &ecs.DescribeServicesInput{ | |
Services: svc, | |
Cluster: aws.String(cluster), | |
}) | |
if err != nil { | |
log.Println(err) | |
} | |
for _, service := range output.Services { | |
for _, lb := range service.LoadBalancers { | |
tg = lb.TargetGroupArn | |
} | |
} | |
log.Printf("The target group is: %v\n", aws.ToString(tg)) | |
return tg | |
} |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Lambda function needs the following permissions: DescribeServices (ECS), DeregisterTargets (ELBv2), and DescribeSubnets (EC2)