diff --git a/pkg/dataobj/consumer/http.go b/pkg/dataobj/consumer/http.go index 730c10040eea4..b0facf0d0ebdc 100644 --- a/pkg/dataobj/consumer/http.go +++ b/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 diff --git a/pkg/loki/modules.go b/pkg/loki/modules.go index 44dfbe1e55736..591e0738f64d7 100644 --- a/pkg/loki/modules.go +++ b/pkg/loki/modules.go @@ -2403,6 +2403,10 @@ func (t *Loki) initDataObjConsumer() (services.Service, error) { httpMiddleware := middleware.Merge( serverutil.RecoveryHTTPMiddleware, ) + t.Server.HTTP. + 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").