From 0a62d8760ac558f42f84c9105894ac55e6e2a6eb Mon Sep 17 00:00:00 2001
From: Yuchen Wang <yuchen.wang@databricks.com>
Date: Sun, 13 Oct 2024 03:53:54 -0700
Subject: [PATCH] update downscale

---
 pkg/controller/delay.go | 18 +++++++++++++-----
 1 file changed, 13 insertions(+), 5 deletions(-)

diff --git a/pkg/controller/delay.go b/pkg/controller/delay.go
index da671a0f0..19afb5512 100644
--- a/pkg/controller/delay.go
+++ b/pkg/controller/delay.go
@@ -20,6 +20,13 @@ import (
 	"github.com/grafana/rollout-operator/pkg/config"
 )
 
+func getStsSvcName(sts *v1.StatefulSet) string {
+	if sts.Spec.ServiceName != "" {
+		return sts.Spec.ServiceName
+	}
+	return sts.GetName()
+}
+
 func cancelDelayedDownscaleIfConfigured(ctx context.Context, logger log.Logger, sts *v1.StatefulSet, httpClient httpClient, replicas int32) {
 	delay, prepareURL, err := parseDelayedDownscaleAnnotations(sts.GetAnnotations())
 	if delay == 0 || prepareURL == nil {
@@ -31,7 +38,7 @@ func cancelDelayedDownscaleIfConfigured(ctx context.Context, logger log.Logger,
 		return
 	}
 
-	endpoints := createPrepareDownscaleEndpoints(sts.Namespace, sts.GetName(), 0, int(replicas), prepareURL)
+	endpoints := createPrepareDownscaleEndpoints(sts.Namespace, getStsSvcName(sts), 0, int(replicas), prepareURL)
 
 	callCancelDelayedDownscale(ctx, logger, httpClient, endpoints)
 }
@@ -54,19 +61,19 @@ func checkScalingDelay(ctx context.Context, logger log.Logger, sts *v1.StatefulS
 	}
 
 	if desiredReplicas >= currentReplicas {
-		callCancelDelayedDownscale(ctx, logger, httpClient, createPrepareDownscaleEndpoints(sts.Namespace, sts.GetName(), 0, int(currentReplicas), prepareURL))
+		callCancelDelayedDownscale(ctx, logger, httpClient, createPrepareDownscaleEndpoints(sts.Namespace, getStsSvcName(sts), 0, int(currentReplicas), prepareURL))
 		// Proceed even if calling cancel of delayed downscale fails. We call cancellation repeatedly, so it will happen during next reconcile.
 		return desiredReplicas, nil
 	}
 
 	{
 		// Replicas in [0, desired) interval should cancel any delayed downscale, if they have any.
-		cancelEndpoints := createPrepareDownscaleEndpoints(sts.Namespace, sts.GetName(), 0, int(desiredReplicas), prepareURL)
+		cancelEndpoints := createPrepareDownscaleEndpoints(sts.Namespace, getStsSvcName(sts), 0, int(desiredReplicas), prepareURL)
 		callCancelDelayedDownscale(ctx, logger, httpClient, cancelEndpoints)
 	}
 
 	// Replicas in [desired, current) interval are going to be stopped.
-	downscaleEndpoints := createPrepareDownscaleEndpoints(sts.Namespace, sts.GetName(), int(desiredReplicas), int(currentReplicas), prepareURL)
+	downscaleEndpoints := createPrepareDownscaleEndpoints(sts.Namespace, getStsSvcName(sts), int(desiredReplicas), int(currentReplicas), prepareURL)
 	elapsedTimeSinceDownscaleInitiated, err := callPrepareDownscaleAndReturnElapsedDurationsSinceInitiatedDownscale(ctx, logger, httpClient, downscaleEndpoints)
 	if err != nil {
 		return currentReplicas, fmt.Errorf("failed prepare pods for delayed downscale: %v", err)
@@ -147,7 +154,8 @@ func createPrepareDownscaleEndpoints(namespace, serviceName string, from, to int
 		}
 
 		ep.url = *url
-		ep.url.Host = fmt.Sprintf("%s.%v.%v.svc.cluster.local.", ep.podName, serviceName, ep.namespace)
+		port := ep.url.Port()
+		ep.url.Host = fmt.Sprintf("%s.%v.%v.svc.cluster.local:%s", ep.podName, serviceName, ep.namespace, port)
 
 		eps = append(eps, ep)
 	}