compactor: allow grace period for final status update (#15113)

#### What this PR does

This adds a grace period for compactor workers to send final status
updates when the parent context of the scheduler executor is canceled
(e.g. during shutdown). The purpose is for compactors to be able to
request job reassignment before shutting down. This avoids having to
wait for lease expiration if restarts occur mid-job and during a rollout
would help avoid the confusing situation of having more active jobs than
running compactors.
 
<img width="2420" height="800" alt="image"
src="https://github.com/user-attachments/assets/26a68ffd-0b33-4dba-80c4-40fd90d96bb7"
/>



My original attempt at fixing this was to check `if ctx.Err() != nil`,
then pass a different context, but that subtly never provided the grace
period if the shutdown happened during the final update itself. This
approach watches the context and starts a timer only if the context
cancellation occurs. This restricts when the deadline is imposed to not
get in front of the normal retry behavior.

#### Which issue(s) this PR fixes or relates to

Fixes N/A

#### Checklist

- [x] Tests updated.
- [ ] Documentation added.
- [ ] `CHANGELOG.md` updated - the order of entries should be
`[CHANGE]`, `[FEATURE]`, `[ENHANCEMENT]`, `[BUGFIX]`. If changelog entry
is not needed, please add the `changelog-not-needed` label to the PR.
- [ ]
[`about-versioning.md`](https://github.com/grafana/mimir/blob/main/docs/sources/mimir/configure/about-versioning.md)
updated with experimental features.

Author: andyasp

pkg/compactor/executor.go

@@ -40,9 +40,10 @@ import (
 )
 
 var (
-	errCompactionJobHasNoBlocks = errors.New("compaction job has no blocks")
-	errNoBlockMetadataProvided  = errors.New("no block metadata provided")
-	errJobCanceledByScheduler   = errors.New("job canceled by scheduler")
+	errCompactionJobHasNoBlocks      = errors.New("compaction job has no blocks")
+	errNoBlockMetadataProvided       = errors.New("no block metadata provided")
+	errJobCanceledByScheduler        = errors.New("job canceled by scheduler")
+	errFinalStatusGracePeriodTimeout = errors.New("final status grace period timed out")
 )
 
 // compactionExecutor defines how compaction work is executed.
@@ -82,25 +83,27 @@ func (e *standaloneExecutor) stop() error {
 }
 
 var (
-	errInvalidSchedulerEndpoint          = fmt.Errorf("invalid compactor.scheduler-client.scheduler-endpoint, required when compactor.scheduler-client.enabled is true")
-	errInvalidSchedulerUpdateInterval    = fmt.Errorf("invalid compactor.scheduler-client.update-interval, interval must be positive")
-	errInvalidSchedulerLeasingMinBackoff = fmt.Errorf("invalid compactor.scheduler-client.leasing-min-backoff, must be positive")
-	errInvalidSchedulerLeasingMaxBackoff = fmt.Errorf("invalid compactor.scheduler-client.leasing-max-backoff, must be greater than min backoff")
-	errInvalidSchedulerUpdateMinBackoff  = fmt.Errorf("invalid compactor.scheduler-client.update-min-backoff, must be positive")
-	errInvalidSchedulerUpdateMaxBackoff  = fmt.Errorf("invalid compactor.scheduler-client.update-max-backoff, must be greater than min backoff")
+	errInvalidSchedulerEndpoint                      = fmt.Errorf("invalid compactor.scheduler-client.scheduler-endpoint, required when compactor.scheduler-client.enabled is true")
+	errInvalidSchedulerUpdateInterval                = fmt.Errorf("invalid compactor.scheduler-client.update-interval, interval must be positive")
+	errInvalidSchedulerLeasingMinBackoff             = fmt.Errorf("invalid compactor.scheduler-client.leasing-min-backoff, must be positive")
+	errInvalidSchedulerLeasingMaxBackoff             = fmt.Errorf("invalid compactor.scheduler-client.leasing-max-backoff, must be greater than min backoff")
+	errInvalidSchedulerUpdateMinBackoff              = fmt.Errorf("invalid compactor.scheduler-client.update-min-backoff, must be positive")
+	errInvalidSchedulerUpdateMaxBackoff              = fmt.Errorf("invalid compactor.scheduler-client.update-max-backoff, must be greater than min backoff")
+	errInvalidSchedulerTerminatingFinalStatusTimeout = fmt.Errorf("invalid compactor.scheduler-client.terminating-final-status-timeout, must be positive")
 )
 
 type SchedulerClientConfig struct {
-	Enabled                      bool                           `yaml:"enabled" category:"experimental"`
-	SchedulerEndpoint            string                         `yaml:"scheduler_endpoint" category:"experimental"`
-	GRPCClientConfig             grpcclient.Config              `yaml:"grpc_client_config" category:"experimental"`
-	LeasingMinBackoff            time.Duration                  `yaml:"leasing_min_backoff" category:"experimental"`
-	LeasingMaxBackoff            time.Duration                  `yaml:"leasing_max_backoff" category:"experimental"`
-	UpdateInterval               time.Duration                  `yaml:"update_interval" category:"experimental"`
-	UpdateMinBackoff             time.Duration                  `yaml:"update_min_backoff" category:"experimental"`
-	UpdateMaxBackoff             time.Duration                  `yaml:"update_max_backoff" category:"experimental"`
-	CompactionDirCleanupInterval time.Duration                  `yaml:"compaction_dir_cleanup_interval" category:"experimental"`
-	MetadataCacheConfig          mimir_tsdb.MetadataCacheConfig `yaml:"metadata_cache" category:"experimental"`
+	Enabled                       bool                           `yaml:"enabled" category:"experimental"`
+	SchedulerEndpoint             string                         `yaml:"scheduler_endpoint" category:"experimental"`
+	GRPCClientConfig              grpcclient.Config              `yaml:"grpc_client_config" category:"experimental"`
+	LeasingMinBackoff             time.Duration                  `yaml:"leasing_min_backoff" category:"experimental"`
+	LeasingMaxBackoff             time.Duration                  `yaml:"leasing_max_backoff" category:"experimental"`
+	UpdateInterval                time.Duration                  `yaml:"update_interval" category:"experimental"`
+	UpdateMinBackoff              time.Duration                  `yaml:"update_min_backoff" category:"experimental"`
+	UpdateMaxBackoff              time.Duration                  `yaml:"update_max_backoff" category:"experimental"`
+	CompactionDirCleanupInterval  time.Duration                  `yaml:"compaction_dir_cleanup_interval" category:"experimental"`
+	MetadataCacheConfig           mimir_tsdb.MetadataCacheConfig `yaml:"metadata_cache" category:"experimental"`
+	TerminatingFinalStatusTimeout time.Duration                  `yaml:"terminating_final_status_timeout" category:"experimental"`
 }
 
 func (cfg *SchedulerClientConfig) RegisterFlags(f *flag.FlagSet) {
@@ -113,6 +116,7 @@ func (cfg *SchedulerClientConfig) RegisterFlags(f *flag.FlagSet) {
 	f.DurationVar(&cfg.UpdateMinBackoff, flagPrefix+"update-min-backoff", 1*time.Second, "Minimum backoff time for compaction executor retries when sending scheduler status updates.")
 	f.DurationVar(&cfg.UpdateMaxBackoff, flagPrefix+"update-max-backoff", 32*time.Second, "Maximum backoff time for compaction executor retries when sending scheduler status updates.")
 	f.DurationVar(&cfg.CompactionDirCleanupInterval, flagPrefix+"compaction-dir-cleanup-interval", 30*time.Minute, "Defines how frequently to clean up the compaction working directory. The directory is cleaned on startup and then only when this interval has elapsed since the last cleanup. Set to 0 to disable periodic cleanup.")
+	f.DurationVar(&cfg.TerminatingFinalStatusTimeout, flagPrefix+"terminating-final-status-timeout", 30*time.Second, "Timeout for sending a final job status update to the scheduler when the parent context is canceled (e.g. during shutdown).")
 	cfg.GRPCClientConfig.RegisterFlagsWithPrefix(flagPrefix+"grpc-client-config", f)
 	cfg.MetadataCacheConfig.RegisterFlagsWithPrefix(f, flagPrefix+"metadata-cache.")
 }
@@ -142,6 +146,9 @@ func (cfg *SchedulerClientConfig) Validate() error {
 	if err := cfg.MetadataCacheConfig.Validate(); err != nil {
 		return err
 	}
+	if cfg.TerminatingFinalStatusTimeout <= 0 {
+		return errInvalidSchedulerTerminatingFinalStatusTimeout
+	}
 	return nil
 }
 
@@ -332,24 +339,38 @@ func (e *schedulerExecutor) startJobStatusUpdater(ctx context.Context, c *Multit
 	}
 }
 
-// sendFinalJobStatus sends a final status update to the scheduler with retry policy.
+// sendFinalJobStatus sends a final status update to the scheduler with a retry policy.
 // Compaction jobs send final statuses on completion, planning jobs only on failure for reassignment.
+// If ctx is canceled (e.g. during shutdown), attempting to send a final status can continue for up to TerminatingFinalStatusTimeout.
 func (e *schedulerExecutor) sendFinalJobStatus(ctx context.Context, key *compactorschedulerpb.JobKey, spec *compactorschedulerpb.JobSpec, status compactorschedulerpb.UpdateType) {
+	graceCtx, cancel := context.WithCancelCause(context.WithoutCancel(ctx))
+	defer cancel(nil)
+	stop := context.AfterFunc(ctx, func() {
+		timer := time.NewTimer(e.cfg.TerminatingFinalStatusTimeout)
+		defer timer.Stop()
+		select {
+		case <-timer.C:
+			cancel(errFinalStatusGracePeriodTimeout)
+		case <-graceCtx.Done():
+		}
+	})
+	defer stop()
+
 	jobId := key.Id
 	jobTenant := spec.Tenant
 
 	var err error
 	switch spec.JobType {
 	case compactorschedulerpb.JOB_TYPE_COMPACTION:
 		req := &compactorschedulerpb.UpdateCompactionJobRequest{Key: key, Tenant: spec.Tenant, Update: status}
-		err = e.retryable.WithContext(ctx).Run(func() error {
-			_, err := e.schedulerClient.UpdateCompactionJob(ctx, req)
+		err = e.retryable.WithContext(graceCtx).Run(func() error {
+			_, err := e.schedulerClient.UpdateCompactionJob(graceCtx, req)
 			return err
 		})
 	case compactorschedulerpb.JOB_TYPE_PLANNING:
 		req := &compactorschedulerpb.UpdatePlanJobRequest{Key: key, Tenant: spec.Tenant, Update: status}
-		err = e.retryable.WithContext(ctx).Run(func() error {
-			_, err := e.schedulerClient.UpdatePlanJob(ctx, req)
+		err = e.retryable.WithContext(graceCtx).Run(func() error {
+			_, err := e.schedulerClient.UpdatePlanJob(graceCtx, req)
 			return err
 		})
 	default:

pkg/compactor/executor_test.go

@@ -568,6 +568,63 @@ func TestSchedulerExecutor_JobCancellationOn_NotFoundResponse(t *testing.T) {
 	require.Equal(t, 2, mockSchedulerClient.GetUpdateJobCallCount(), "should have sent exactly 2 updates: one successful, then one NOT_FOUND")
 }
 
+func TestSchedulerExecutor_TerminatingFinalJobStatus(t *testing.T) {
+	key := &compactorschedulerpb.JobKey{Id: "test-job"}
+	spec := &compactorschedulerpb.JobSpec{
+		Tenant:  "tenant",
+		JobType: compactorschedulerpb.JOB_TYPE_PLANNING,
+	}
+
+	t.Run("canceled_context_uses_fresh_context_for_update", func(t *testing.T) {
+		mock := &mockCompactorSchedulerClient{
+			UpdatePlanJobFunc: func(ctx context.Context, _ *compactorschedulerpb.UpdatePlanJobRequest) (*compactorschedulerpb.UpdateJobResponse, error) {
+				assert.NoError(t, ctx.Err())
+				return &compactorschedulerpb.UpdateJobResponse{}, nil
+			},
+		}
+		exec := newTestSchedulerExecutor(t, makeTestCompactorConfig(), mock)
+
+		canceledCtx, cancel := context.WithCancel(context.Background())
+		cancel()
+
+		exec.sendFinalJobStatus(canceledCtx, key, spec, compactorschedulerpb.UPDATE_TYPE_REASSIGN)
+		require.Equal(t, 1, mock.GetUpdateJobCallCount())
+	})
+
+	t.Run("canceled_context_returns_after_final_status_timeout", func(t *testing.T) {
+		mock := &mockCompactorSchedulerClient{
+			UpdatePlanJobFunc: func(ctx context.Context, _ *compactorschedulerpb.UpdatePlanJobRequest) (*compactorschedulerpb.UpdateJobResponse, error) {
+				<-ctx.Done()
+				return nil, ctx.Err()
+			},
+		}
+		cfg := makeTestCompactorConfig()
+		cfg.SchedulerClientConfig.TerminatingFinalStatusTimeout = time.Millisecond // arbitrary
+
+		exec := newTestSchedulerExecutor(t, cfg, mock)
+		cancelledCtx, cancel := context.WithCancel(context.Background())
+		cancel()
+
+		synctest.Test(t, func(t *testing.T) {
+			done := make(chan struct{})
+			go func() {
+				exec.sendFinalJobStatus(cancelledCtx, key, spec, compactorschedulerpb.UPDATE_TYPE_REASSIGN)
+				close(done) // unblock select
+			}()
+
+			// Advance synctest time past timeout
+			time.Sleep(cfg.SchedulerClientConfig.TerminatingFinalStatusTimeout + time.Millisecond)
+			synctest.Wait()
+
+			select {
+			case <-done:
+			default:
+				t.Fatal("sendFinalJobStatus should have returned after FinalStatusTimeout elapsed")
+			}
+		})
+	})
+}
+
 func TestSchedulerExecutor_ExecuteCompactionJob_InvalidInput(t *testing.T) {
 	tests := map[string]struct {
 		spec           *compactorschedulerpb.JobSpec