feat: add prepare downscale handler

Author: grobinson-grafana

pkg/dataobj/consumer/http.go

@@ -10,6 +10,13 @@ import (
 	"github.com/grafana/loki/v3/pkg/util"
 )
 
+// PrepareDownscaleHandler is a special handler called by the rollout operator
+// immediately before the pod is downscaled. It can stop a downscale by
+// responding with a non 2xx status code.
+func (s *Service) PrepareDownscaleHandler(_ http.ResponseWriter, _ *http.Request) {
+	s.partitionInstanceLifecycler.SetRemoveOwnerOnShutdown(true)
+}
+
 // PrepareDelayedDownscaleHandler is a special handler called by the rollout
 // operator to prepare for a delayed downscale. This allows the service to
 // perform any number of actions in preparation of being scaled down at the

pkg/loki/modules.go

@@ -2404,7 +2404,11 @@ func (t *Loki) initDataObjConsumer() (services.Service, error) {
 		serverutil.RecoveryHTTPMiddleware,
 	)
 	t.Server.HTTP.
-		Methods(http.MethodGet, http.MethodPost, http.MethoDelete).
+		Methods(http.MethodGet, http.MethodPost, http.MethodDelete).
+		Path("/dataobj-consumer/prepare-downscale").
+		Handler(httpMiddleware.Wrap(http.HandlerFunc(t.dataObjConsumer.PrepareDownscaleHandler)))
+	t.Server.HTTP.
+		Methods(http.MethodGet, http.MethodPost, http.MethodDelete).
 		Path("/dataobj-consumer/prepare-delayed-downscale").
 		Handler(httpMiddleware.Wrap(http.HandlerFunc(t.dataObjConsumer.PrepareDelayedDownscaleHandler)))