From a206b6e9668e44a6df398e0182d4cd34a39450b3 Mon Sep 17 00:00:00 2001
From: krhitesh7 <hiteshrawat@sharechat.co>
Date: Sun, 21 Apr 2024 12:44:22 +0530
Subject: [PATCH] Set endpoint to DRAIN instead of removing it

---
 pkg/cache/v3/simple.go | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/pkg/cache/v3/simple.go b/pkg/cache/v3/simple.go
index 669c14f564..371220429d 100644
--- a/pkg/cache/v3/simple.go
+++ b/pkg/cache/v3/simple.go
@@ -17,6 +17,7 @@ package cache
 import (
 	"context"
 	"fmt"
+	core "github.com/envoyproxy/go-control-plane/envoy/config/core/v3"
 	endpoint "github.com/envoyproxy/go-control-plane/envoy/config/endpoint/v3"
 	"github.com/envoyproxy/go-control-plane/pkg/resource/v3"
 	"strconv"
@@ -388,8 +389,11 @@ func (cache *snapshotCache) DeleteResources(ctx context.Context, _ string, typ s
 					for _, lbEndpoint := range cla.Endpoints[i].LbEndpoints {
 						if resourceToDelete == GetResourceName(lbEndpoint) {
 							didModify = true
-							cache.log.Infof("Removed endpoint %s", resourceToDelete)
-							continue
+							cache.log.Infof("Drain and remove endpoint %s", resourceToDelete)
+
+							// Set to UNHEALTHY/DRAINING and let Envoy gracefully remove them.
+							lbEndpoint.HealthStatus = core.HealthStatus_DRAINING
+							// continue
 						}
 						newEndpoints = append(newEndpoints, lbEndpoint)
 					}