add a flag to select target types to deregister instance from (#214)

* add deregister-target-types flag to select target types to deregister Signed-off-by: Yuxuan Zhu <[email protected]> * upgrade golang image to 1.21 Signed-off-by: Yuxuan Zhu <[email protected]> * update readme Signed-off-by: Yuxuan Zhu <[email protected]> * add debug Signed-off-by: Yuxuan Zhu <[email protected]> * remove debug logging Signed-off-by: Yuxuan Zhu <[email protected]> * revert golang version to 1.19 Signed-off-by: Yuxuan Zhu <[email protected]> --------- Signed-off-by: Yuxuan Zhu <[email protected]> Co-authored-by: Yuxuan Zhu <[email protected]> Co-authored-by: Todd Ekenstam <[email protected]>
keikoproj · Sep 9, 2024 · f3b1488 · f3b1488
1 parent ad9b7d7
commit f3b1488
Show file tree

Hide file tree

Showing 5 changed files with 28 additions and 12 deletions.
diff --git a/README.md b/README.md
@@ -134,6 +134,7 @@ time="2020-03-11T07:24:49Z" level=info msg="event ce25c321-ec67-3f0b-c156-a7c1f7
 | polling-interval | 10 | Int | interval in seconds for which to poll SQS |
 | with-deregister | true | Bool | try to deregister deleting instance from target groups |
 | refresh-expired-credentials | false | Bool | refreshes expired credentials (requires shared credentials file) |
+| deregister-target-types | "classic-elb,target-group" | String | comma separated list of target types to deregister instance from (classic-elb, target-group) |
 
 
 ## Release History

diff --git a/cmd/serve.go b/cmd/serve.go
@@ -1,6 +1,7 @@
 package cmd
 
 import (
+	"fmt"
 	"os"
 	"time"
 
@@ -30,6 +31,7 @@ var (
 	nodeName                   string
 	logLevel                   string
 	deregisterTargetGroups     bool
+	deregisterTargetTypes      []string
 	refreshExpiredCredentials  bool
 	drainRetryIntervalSeconds  int
 	maxDrainConcurrency        int64
@@ -83,6 +85,7 @@ var serveCmd = &cobra.Command{
 			DrainRetryAttempts:         uint(drainRetryAttempts),
 			Region:                     region,
 			WithDeregister:             deregisterTargetGroups,
+			DeregisterTargetTypes:      deregisterTargetTypes,
 		}
 
 		s := service.New(auth, context)
@@ -105,6 +108,8 @@ func init() {
 	serveCmd.Flags().IntVar(&drainRetryAttempts, "drain-retries", 3, "number of times to retry the node drain operation")
 	serveCmd.Flags().IntVar(&pollingIntervalSeconds, "polling-interval", 10, "interval in seconds for which to poll SQS")
 	serveCmd.Flags().BoolVar(&deregisterTargetGroups, "with-deregister", true, "try to deregister deleting instance from target groups")
+	serveCmd.Flags().StringSliceVar(&deregisterTargetTypes, "deregister-target-types", []string{service.TargetTypeClassicELB.String(), service.TargetTypeTargetGroup.String()},
+		fmt.Sprintf("comma separated list of target types to deregister instance from (%s, %s)", service.TargetTypeClassicELB.String(), service.TargetTypeTargetGroup.String()))
 	serveCmd.Flags().BoolVar(&refreshExpiredCredentials, "refresh-expired-credentials", false, "refreshes expired credentials (requires shared credentials file)")
 }
 

diff --git a/pkg/service/manager.go b/pkg/service/manager.go
@@ -48,6 +48,7 @@ type ManagerContext struct {
 	DrainRetryAttempts         uint
 	PollingIntervalSeconds     int64
 	WithDeregister             bool
+	DeregisterTargetTypes      []string
 	MaxDrainConcurrency        *semaphore.Weighted
 	MaxTimeToProcessSeconds    int64
 }

diff --git a/pkg/service/server.go b/pkg/service/server.go
@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"math/rand"
 	"runtime"
+	"slices"
 	"strings"
 	"time"
 
@@ -73,6 +74,7 @@ func (mgr *Manager) Start() {
 	log.Infof("node drain retry interval seconds = %v", ctx.DrainRetryIntervalSeconds)
 	log.Infof("node drain retry attempts = %v", ctx.DrainRetryAttempts)
 	log.Infof("with alb deregister = %v", ctx.WithDeregister)
+	log.Infof("deregister target types = %v", ctx.DeregisterTargetTypes)
 
 	// start metrics server
 	log.Infof("starting metrics server on %v%v", MetricsEndpoint, MetricsPort)
@@ -305,6 +307,7 @@ func (mgr *Manager) deleteNodeTarget(event *LifecycleEvent) error {
 
 func (mgr *Manager) scanMembership(event *LifecycleEvent) (*ScanResult, error) {
 	var (
+		ctx                 = &mgr.context
 		elbv2Client         = mgr.authenticator.ELBv2Client
 		elbClient           = mgr.authenticator.ELBClient
 		instanceID          = event.EC2InstanceID
@@ -315,22 +318,26 @@ func (mgr *Manager) scanMembership(event *LifecycleEvent) (*ScanResult, error) {
 
 	// get all target groups
 	targetGroups := []*elbv2.TargetGroup{}
-	err := elbv2Client.DescribeTargetGroupsPages(&elbv2.DescribeTargetGroupsInput{}, func(page *elbv2.DescribeTargetGroupsOutput, lastPage bool) bool {
-		targetGroups = append(targetGroups, page.TargetGroups...)
-		return page.NextMarker != nil
-	})
-	if err != nil {
-		return scanResult, err
+	if slices.Contains(ctx.DeregisterTargetTypes, TargetTypeTargetGroup.String()) {
+		err := elbv2Client.DescribeTargetGroupsPages(&elbv2.DescribeTargetGroupsInput{}, func(page *elbv2.DescribeTargetGroupsOutput, lastPage bool) bool {
+			targetGroups = append(targetGroups, page.TargetGroups...)
+			return page.NextMarker != nil
+		})
+		if err != nil {
+			return scanResult, err
+		}
 	}
 
 	// get all classic elbs
 	elbDescriptions := []*elb.LoadBalancerDescription{}
-	err = elbClient.DescribeLoadBalancersPages(&elb.DescribeLoadBalancersInput{}, func(page *elb.DescribeLoadBalancersOutput, lastPage bool) bool {
-		elbDescriptions = append(elbDescriptions, page.LoadBalancerDescriptions...)
-		return page.NextMarker != nil
-	})
-	if err != nil {
-		return scanResult, err
+	if slices.Contains(ctx.DeregisterTargetTypes, TargetTypeClassicELB.String()) {
+		err := elbClient.DescribeLoadBalancersPages(&elb.DescribeLoadBalancersInput{}, func(page *elb.DescribeLoadBalancersOutput, lastPage bool) bool {
+			elbDescriptions = append(elbDescriptions, page.LoadBalancerDescriptions...)
+			return page.NextMarker != nil
+		})
+		if err != nil {
+			return scanResult, err
+		}
 	}
 
 	log.Infof("%v> checking targetgroup/elb membership", instanceID)

diff --git a/pkg/service/server_test.go b/pkg/service/server_test.go
@@ -297,6 +297,7 @@ func Test_HandleEventWithDeregister(t *testing.T) {
 
 	ctx := _newBasicContext()
 	ctx.WithDeregister = true
+	ctx.DeregisterTargetTypes = []string{TargetTypeClassicELB.String(), TargetTypeTargetGroup.String()}
 
 	fakeNodes := []v1.Node{
 		{
@@ -398,6 +399,7 @@ func Test_HandleEventWithDeregisterError(t *testing.T) {
 
 	ctx := _newBasicContext()
 	ctx.WithDeregister = true
+	ctx.DeregisterTargetTypes = []string{TargetTypeClassicELB.String(), TargetTypeTargetGroup.String()}
 
 	fakeNodes := []v1.Node{
 		{