From 298eeec65a1589dfc149aa517e1535d5ac32e3b9 Mon Sep 17 00:00:00 2001 From: Zihan Jiang Date: Mon, 1 Jul 2024 12:32:55 -0700 Subject: [PATCH] only do alb deregister for nodes with specific label Signed-off-by: Zihan Jiang --- pkg/service/server.go | 6 +- pkg/service/server_test.go | 128 +++++++++++++++++++++++++++++++++++++ 2 files changed, 133 insertions(+), 1 deletion(-) diff --git a/pkg/service/server.go b/pkg/service/server.go index 23f6bd3..8846ba1 100644 --- a/pkg/service/server.go +++ b/pkg/service/server.go @@ -51,6 +51,8 @@ var ( WaiterMaxDelay time.Duration = 90 * time.Second // WaiterMaxAttempts defines the maximum attempts of the IEB waiter WaiterMaxAttempts uint32 = 120 + // EnableALBDeregisterNodeLabelKey defines whether lifecycle manager need to deregister node from ALB + EnableALBDeregisterNodeLabelKey = "lifecycle-manager.keikoproj.io/enable-alb-deregister" ) // Start starts the lifecycle-manager service @@ -503,7 +505,9 @@ func (mgr *Manager) drainLoadbalancerTarget(event *LifecycleEvent) error { isFinished bool ) - if !ctx.WithDeregister { + log.Info("node", node) + if !ctx.WithDeregister || node.Labels[EnableALBDeregisterNodeLabelKey] != "true" { + log.Infof("%v> skipping load balancer deregistration", instanceID) return nil } log.Infof("%v> starting load balancer drain worker", instanceID) diff --git a/pkg/service/server_test.go b/pkg/service/server_test.go index ec0478c..b0014e6 100644 --- a/pkg/service/server_test.go +++ b/pkg/service/server_test.go @@ -300,11 +300,21 @@ func Test_HandleEventWithDeregister(t *testing.T) { fakeNodes := []v1.Node{ { + ObjectMeta: apimachinery_v1.ObjectMeta{ + Labels: map[string]string{ + EnableALBDeregisterNodeLabelKey: "true", + }, + }, Spec: v1.NodeSpec{ ProviderID: fmt.Sprintf("aws:///us-west-2a/%v", instanceID), }, }, { + ObjectMeta: apimachinery_v1.ObjectMeta{ + Labels: map[string]string{ + EnableALBDeregisterNodeLabelKey: "true", + }, + }, Spec: v1.NodeSpec{ ProviderID: "aws:///us-west-2c/i-22222222222222222", }, @@ -315,6 +325,8 @@ func Test_HandleEventWithDeregister(t *testing.T) { auth.KubernetesClient.CoreV1().Nodes().Create(context.Background(), &node, apimachinery_v1.CreateOptions{}) } + referenceNode, _ := getNodeByInstance(auth.KubernetesClient, instanceID) + event := &LifecycleEvent{ LifecycleHookName: "my-hook", AccountID: "12345689012", @@ -325,6 +337,7 @@ func Test_HandleEventWithDeregister(t *testing.T) { LifecycleActionToken: "cc34960c-1e41-4703-a665-bdb3e5b81ad3", receiptHandle: "MbZj6wDWli+JvwwJaBV+3dcjk2YW2vA3+STFFljTM8tJJg6HRG6PYSasuWXPJB+Cw=", heartbeatInterval: 3, + referencedNode: referenceNode, } g := New(auth, ctx) @@ -342,6 +355,108 @@ func Test_HandleEventWithDeregister(t *testing.T) { } } +func Test_HandleEventWithDeregisterSkipNode(t *testing.T) { + t.Log("Test_HandleEvent: should skip deregistration") + var ( + asgStubber = &stubAutoscaling{} + sqsStubber = &stubSQS{} + arn = "arn:aws:elasticloadbalancing:us-west-2:0000000000:targetgroup/targetgroup-name/some-id" + elbName = "my-classic-elb" + instanceID = "i-123486890234" + port int64 = 122233 + ) + + elbv2Stubber := &stubELBv2{ + targetHealthDescriptions: []*elbv2.TargetHealthDescription{ + { + Target: &elbv2.TargetDescription{ + Id: aws.String(instanceID), + Port: aws.Int64(port), + }, + TargetHealth: &elbv2.TargetHealth{ + State: aws.String(elbv2.TargetHealthStateEnumUnused), + }, + }, + }, + targetGroups: []*elbv2.TargetGroup{ + { + TargetGroupArn: aws.String(arn), + }, + }, + } + + elbStubber := &stubELB{ + loadBalancerDescriptions: []*elb.LoadBalancerDescription{ + { + LoadBalancerName: aws.String(elbName), + }, + }, + instanceStates: []*elb.InstanceState{ + { + InstanceId: aws.String(instanceID), + State: aws.String("OutOfService"), + }, + }, + } + + auth := Authenticator{ + ScalingGroupClient: asgStubber, + SQSClient: sqsStubber, + ELBv2Client: elbv2Stubber, + ELBClient: elbStubber, + KubernetesClient: fake.NewSimpleClientset(), + } + + ctx := _newBasicContext() + ctx.WithDeregister = true + + fakeNodes := []v1.Node{ + { + Spec: v1.NodeSpec{ + ProviderID: fmt.Sprintf("aws:///us-west-2a/%v", instanceID), + }, + }, + { + Spec: v1.NodeSpec{ + ProviderID: "aws:///us-west-2c/i-22222222222222222", + }, + }, + } + + for _, node := range fakeNodes { + auth.KubernetesClient.CoreV1().Nodes().Create(context.Background(), &node, apimachinery_v1.CreateOptions{}) + } + + referenceNode, _ := getNodeByInstance(auth.KubernetesClient, instanceID) + + event := &LifecycleEvent{ + LifecycleHookName: "my-hook", + AccountID: "12345689012", + RequestID: "63f5b5c2-58b3-0574-b7d5-b3162d0268f0", + LifecycleTransition: "autoscaling:EC2_INSTANCE_TERMINATING", + AutoScalingGroupName: "my-asg", + EC2InstanceID: instanceID, + LifecycleActionToken: "cc34960c-1e41-4703-a665-bdb3e5b81ad3", + receiptHandle: "MbZj6wDWli+JvwwJaBV+3dcjk2YW2vA3+STFFljTM8tJJg6HRG6PYSasuWXPJB+Cw=", + heartbeatInterval: 3, + referencedNode: referenceNode, + } + + g := New(auth, ctx) + err := g.handleEvent(event) + if err != nil { + t.Fatalf("handleEvent: expected error not to have occured, %v", err) + } + + if event.drainCompleted != true { + t.Fatal("handleEvent: expected drainCompleted to be true, got: false") + } + + if event.deregisterCompleted != false { + t.Fatal("handleEvent: expected deregisterCompleted to be false, got: true") + } +} + func Test_HandleEventWithDeregisterError(t *testing.T) { t.Log("Test_HandleEvent: should successfully handle events") var ( @@ -401,11 +516,21 @@ func Test_HandleEventWithDeregisterError(t *testing.T) { fakeNodes := []v1.Node{ { + ObjectMeta: apimachinery_v1.ObjectMeta{ + Labels: map[string]string{ + EnableALBDeregisterNodeLabelKey: "true", + }, + }, Spec: v1.NodeSpec{ ProviderID: fmt.Sprintf("aws:///us-west-2a/%v", instanceID), }, }, { + ObjectMeta: apimachinery_v1.ObjectMeta{ + Labels: map[string]string{ + EnableALBDeregisterNodeLabelKey: "true", + }, + }, Spec: v1.NodeSpec{ ProviderID: "aws:///us-west-2c/i-22222222222222222", }, @@ -416,6 +541,8 @@ func Test_HandleEventWithDeregisterError(t *testing.T) { auth.KubernetesClient.CoreV1().Nodes().Create(context.Background(), &node, apimachinery_v1.CreateOptions{}) } + referenceNode, _ := getNodeByInstance(auth.KubernetesClient, instanceID) + event := &LifecycleEvent{ LifecycleHookName: "my-hook", AccountID: "12345689012", @@ -426,6 +553,7 @@ func Test_HandleEventWithDeregisterError(t *testing.T) { LifecycleActionToken: "cc34960c-1e41-4703-a665-bdb3e5b81ad3", receiptHandle: "MbZj6wDWli+JvwwJaBV+3dcjk2YW2vA3+STFFljTM8tJJg6HRG6PYSasuWXPJB+Cw=", heartbeatInterval: 3, + referencedNode: referenceNode, } g := New(auth, ctx)