From b26647e0421ed566a6e0a07abc252e2732bfc331 Mon Sep 17 00:00:00 2001
From: Will Assis <35489495+gassiss@users.noreply.github.com>
Date: Wed, 1 Oct 2025 16:34:30 -0400
Subject: [PATCH] unified-storage: setup rollout-operator endpoint (#111768)

* implement endpoint so that storage/search api can prepare for downscale if enabled
---
 pkg/server/module_server.go                   |  2 +-
 pkg/storage/unified/sql/service.go            | 22 +++++++++++++++++++
 .../unified/sql/test/integration_test.go      |  2 +-
 pkg/tests/testinfra/testinfra.go              |  2 +-
 4 files changed, 25 insertions(+), 3 deletions(-)

diff --git a/pkg/server/module_server.go b/pkg/server/module_server.go
index 0b477c3aa97..5c1517d354d 100644
--- a/pkg/server/module_server.go
+++ b/pkg/server/module_server.go
@@ -187,7 +187,7 @@ func (s *ModuleServer) Run() error {
 		if err != nil {
 			return nil, err
 		}
-		return sql.ProvideUnifiedStorageGrpcService(s.cfg, s.features, nil, s.log, s.registerer, docBuilders, s.storageMetrics, s.indexMetrics, s.searchServerRing, s.MemberlistKVConfig)
+		return sql.ProvideUnifiedStorageGrpcService(s.cfg, s.features, nil, s.log, s.registerer, docBuilders, s.storageMetrics, s.indexMetrics, s.searchServerRing, s.MemberlistKVConfig, s.httpServerRouter)
 	})
 
 	m.RegisterModule(modules.ZanzanaServer, func() (services.Service, error) {
diff --git a/pkg/storage/unified/sql/service.go b/pkg/storage/unified/sql/service.go
index 6dd851ad74f..79d4615f53c 100644
--- a/pkg/storage/unified/sql/service.go
+++ b/pkg/storage/unified/sql/service.go
@@ -6,10 +6,12 @@ import (
 	"fmt"
 	"hash/fnv"
 	"net"
+	"net/http"
 	"os"
 	"strconv"
 	"time"
 
+	"github.com/gorilla/mux"
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/prometheus/client_golang/prometheus/promauto"
 	"go.opentelemetry.io/otel"
@@ -94,6 +96,7 @@ func ProvideUnifiedStorageGrpcService(
 	indexMetrics *resource.BleveIndexMetrics,
 	searchRing *ring.Ring,
 	memberlistKVConfig kv.Config,
+	httpServerRouter *mux.Router,
 ) (UnifiedStorageGrpcService, error) {
 	var err error
 	tracer := otel.Tracer("unified-storage")
@@ -159,6 +162,10 @@ func ProvideUnifiedStorageGrpcService(
 
 		s.ringLifecycler.SetKeepInstanceInTheRingOnShutdown(true)
 		subservices = append(subservices, s.ringLifecycler)
+
+		if httpServerRouter != nil {
+			httpServerRouter.Path("/prepare-downscale").Methods("GET", "POST", "DELETE").Handler(http.HandlerFunc(s.PrepareDownscale))
+		}
 	}
 
 	if cfg.QOSEnabled {
@@ -194,6 +201,21 @@ func ProvideUnifiedStorageGrpcService(
 	return s, nil
 }
 
+func (s *service) PrepareDownscale(w http.ResponseWriter, r *http.Request) {
+	switch r.Method {
+	case http.MethodPost:
+		s.log.Info("Preparing for downscale. Will not keep instance in ring on shutdown.")
+		s.ringLifecycler.SetKeepInstanceInTheRingOnShutdown(false)
+	case http.MethodDelete:
+		s.log.Info("Downscale canceled. Will keep instance in ring on shutdown.")
+		s.ringLifecycler.SetKeepInstanceInTheRingOnShutdown(true)
+	case http.MethodGet:
+		// used for delayed downscale use case, which we don't support. Leaving here for completion sake
+		s.log.Info("Received GET request for prepare-downscale. Behavior not implemented.")
+	default:
+	}
+}
+
 var (
 	// operation used by the search-servers to check if they own the namespace
 	searchOwnerRead = ring.NewOp([]ring.InstanceState{ring.JOINING, ring.ACTIVE, ring.LEAVING}, nil)
diff --git a/pkg/storage/unified/sql/test/integration_test.go b/pkg/storage/unified/sql/test/integration_test.go
index c5420b71280..347a28c4f11 100644
--- a/pkg/storage/unified/sql/test/integration_test.go
+++ b/pkg/storage/unified/sql/test/integration_test.go
@@ -128,7 +128,7 @@ func TestClientServer(t *testing.T) {
 
 	features := featuremgmt.WithFeatures()
 
-	svc, err := sql.ProvideUnifiedStorageGrpcService(cfg, features, dbstore, nil, prometheus.NewPedanticRegistry(), nil, nil, nil, nil, kv.Config{})
+	svc, err := sql.ProvideUnifiedStorageGrpcService(cfg, features, dbstore, nil, prometheus.NewPedanticRegistry(), nil, nil, nil, nil, kv.Config{}, nil)
 	require.NoError(t, err)
 	var client resourcepb.ResourceStoreClient
 
diff --git a/pkg/tests/testinfra/testinfra.go b/pkg/tests/testinfra/testinfra.go
index a4461902fa6..b9f23fb4d44 100644
--- a/pkg/tests/testinfra/testinfra.go
+++ b/pkg/tests/testinfra/testinfra.go
@@ -129,7 +129,7 @@ func StartGrafanaEnv(t *testing.T, grafDir, cfgPath string) (string, *server.Tes
 	var storage sql.UnifiedStorageGrpcService
 	if runstore {
 		storage, err = sql.ProvideUnifiedStorageGrpcService(env.Cfg, env.FeatureToggles, env.SQLStore,
-			env.Cfg.Logger, prometheus.NewPedanticRegistry(), nil, nil, nil, nil, kv.Config{})
+			env.Cfg.Logger, prometheus.NewPedanticRegistry(), nil, nil, nil, nil, kv.Config{}, nil)
 		require.NoError(t, err)
 		ctx := context.Background()
 		err = storage.StartAsync(ctx)