Skip to content

Commit

Permalink
only do alb deregister for nodes with specific label
Browse files Browse the repository at this point in the history 
Signed-off-by: Zihan Jiang <[email protected]>
  • Loading branch information
@ZihanJiang96
ZihanJiang96 committed Jul 1, 2024
1 parent 87c603d commit 298eeec
Show file tree
Hide file tree
Showing 2 changed files with 133 additions and 1 deletion.
6 changes: 5 additions & 1 deletion pkg/service/server.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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)
Expand Down
128 changes: 128 additions & 0 deletions pkg/service/server_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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",
},
Expand All @@ -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",
Expand All @@ -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)
Expand All @@ -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 (
Expand Down Expand Up @@ -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",
},
Expand All @@ -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",
Expand All @@ -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)
Expand Down

0 comments on commit 298eeec

Please sign in to comment.