plugin scheduler: drain started jobs past the window close instead of killing them (#10728)

* plugin scheduler: drain started jobs past the window close instead of killing them * plugin scheduler: never drain-cap an attempt below its declared estimated runtime * plugin scheduler: cap estimated_runtime_seconds before the Duration conversion
2026-08-16 04:06:44 +00:00 · 2026-08-12 19:38:22 -07:00
parent 0799084e98
commit 78e7e04377
2 changed files with 186 additions and 25 deletions
@@ -34,6 +34,8 @@ const (
 	defaultWaitingBacklogFloor                 = 8
 	defaultWaitingBacklogMultiplier            = 4
 	maxEstimatedRuntimeCap                     = 8 * time.Hour
+	// How long started jobs may drain past the run window close.
+	scheduledExecutionDrainGrace = 30 * time.Minute
 )

 type schedulerPolicy struct {
@@ -411,15 +413,10 @@ func (r *Plugin) runJobTypeIteration(jobType string, policy schedulerPolicy) boo
 	if execPolicy.ExecutionTimeout <= 0 {
 		execPolicy.ExecutionTimeout = defaultScheduledExecutionTimeout
 	}
-	if execPolicy.ExecutionTimeout > remaining {
-		execPolicy.ExecutionTimeout = remaining
-	}

-	// jobCtx bounds how long this run keeps STARTING jobs; a started job
-	// runs on its own estimated-runtime deadline (see
-	// executeScheduledJobWithExecutor). Jobs still queued when the window
-	// closes are canceled and re-proposed by a later detection, so one large
-	// backlog cannot monopolise the lane for hours.
+	// jobCtx bounds how long this run keeps STARTING jobs; started jobs
+	// drain past the window close (see scheduledAttemptContext), while
+	// still-queued jobs are canceled and re-proposed by a later detection.
 	successCount, errorCount, canceledCount := r.dispatchScheduledProposals(jobCtx, jobType, filtered, clusterContext, execPolicy)

 	status := "success"
@@ -1201,24 +1198,11 @@ func (r *Plugin) executeScheduledJobWithExecutor(
 		// default execution timeout. This lets handlers like vacuum scale
 		// the timeout based on volume size so large volumes are not killed.
 		timeout := policy.ExecutionTimeout
-		if job.Parameters != nil {
-			if est, ok := job.Parameters["estimated_runtime_seconds"]; ok {
-				if v := est.GetInt64Value(); v > 0 {
-					estimated := time.Duration(v) * time.Second
-					if estimated > maxEstimatedRuntimeCap {
-						estimated = maxEstimatedRuntimeCap
-					}
-					if estimated > timeout {
-						timeout = estimated
-					}
-				}
-			}
+		estimated := scheduledEstimatedRuntime(job.Parameters)
+		if estimated > timeout {
+			timeout = estimated
 		}
-		// The attempt runs against its own deadline, detached from the
-		// dispatch window (ctx): a job started near the end of the window
-		// may run to completion; the window only stops further jobs from
-		// starting.
-		execCtx, cancel := context.WithTimeout(context.Background(), timeout)
+		execCtx, cancel := scheduledAttemptContext(ctx, timeout, estimated)
 		_, err := r.executeJobWithExecutor(execCtx, executor, job, clusterContext, int32(attempt))
 		cancel()
 		if err == nil {
@@ -1253,6 +1237,43 @@ func (r *Plugin) executeScheduledJobWithExecutor(
 	return lastErr
 }

+// scheduledEstimatedRuntime returns the job's declared estimated runtime,
+// capped at maxEstimatedRuntimeCap. The seconds are capped before the
+// Duration conversion so oversized values cannot overflow.
+func scheduledEstimatedRuntime(parameters map[string]*plugin_pb.ConfigValue) time.Duration {
+	est, ok := parameters["estimated_runtime_seconds"]
+	if !ok {
+		return 0
+	}
+	v := est.GetInt64Value()
+	if v <= 0 {
+		return 0
+	}
+	if v > int64(maxEstimatedRuntimeCap/time.Second) {
+		v = int64(maxEstimatedRuntimeCap / time.Second)
+	}
+	return time.Duration(v) * time.Second
+}
+
+// scheduledAttemptContext bounds one execution attempt: its own timeout,
+// capped at the window deadline plus scheduledExecutionDrainGrace so a
+// draining job cannot hold the lane indefinitely, but never below the
+// job's declared estimated runtime.
+func scheduledAttemptContext(window context.Context, timeout, estimated time.Duration) (context.Context, context.CancelFunc) {
+	deadline := time.Now().Add(timeout)
+	if window != nil {
+		if windowEnd, ok := window.Deadline(); ok {
+			if drainDeadline := windowEnd.Add(scheduledExecutionDrainGrace); deadline.After(drainDeadline) {
+				if floor := time.Now().Add(estimated); floor.After(drainDeadline) {
+					drainDeadline = floor
+				}
+				deadline = drainDeadline
+			}
+		}
+	}
+	return context.WithDeadline(context.Background(), deadline)
+}
+
 func (r *Plugin) shouldSkipDetectionForWaitingJobs(jobType string, policy schedulerPolicy) (bool, int, int) {
 	waitingCount := r.countWaitingTrackedJobs(jobType)
 	threshold := waitingBacklogThreshold(policy)
@@ -3,11 +3,13 @@ package plugin
 import (
 	"context"
 	"fmt"
+	"math"
 	"sync"
 	"testing"
 	"time"

 	"github.com/seaweedfs/seaweedfs/weed/pb/plugin_pb"
+	"google.golang.org/protobuf/types/known/timestamppb"
 )

 func TestLoadSchedulerPolicyUsesAdminConfig(t *testing.T) {
@@ -738,6 +740,144 @@ func TestDispatchScheduledProposalsLocksPerJob(t *testing.T) {
 	}
 }

+func TestScheduledAttemptContextDrainGrace(t *testing.T) {
+	t.Parallel()
+
+	assertDeadline := func(t *testing.T, ctx context.Context, want time.Time) {
+		t.Helper()
+		got, ok := ctx.Deadline()
+		if !ok {
+			t.Fatal("expected attempt context to carry a deadline")
+		}
+		if diff := got.Sub(want); diff < -5*time.Second || diff > 5*time.Second {
+			t.Fatalf("unexpected attempt deadline: got=%v want~%v", got, want)
+		}
+	}
+
+	windowEnd := time.Now().Add(time.Minute)
+	window, cancelWindow := context.WithDeadline(context.Background(), windowEnd)
+	defer cancelWindow()
+
+	// Fits inside the window: keeps its own deadline.
+	ctx, cancel := scheduledAttemptContext(window, 30*time.Second, 0)
+	assertDeadline(t, ctx, time.Now().Add(30*time.Second))
+	cancel()
+
+	// Longer: capped at window close plus the grace.
+	ctx, cancel = scheduledAttemptContext(window, 3*time.Hour, 0)
+	assertDeadline(t, ctx, windowEnd.Add(scheduledExecutionDrainGrace))
+	cancel()
+
+	// Estimated runtime: keeps its full deadline.
+	ctx, cancel = scheduledAttemptContext(window, 3*time.Hour, 3*time.Hour)
+	assertDeadline(t, ctx, time.Now().Add(3*time.Hour))
+	cancel()
+
+	// An estimate below the timeout floors the drain cap.
+	ctx, cancel = scheduledAttemptContext(window, 3*time.Hour, 45*time.Minute)
+	assertDeadline(t, ctx, time.Now().Add(45*time.Minute))
+	cancel()
+
+	// No window deadline: the timeout stands alone.
+	ctx, cancel = scheduledAttemptContext(context.Background(), 2*time.Hour, 0)
+	assertDeadline(t, ctx, time.Now().Add(2*time.Hour))
+	cancel()
+}
+
+func TestScheduledEstimatedRuntime(t *testing.T) {
+	t.Parallel()
+
+	estimate := func(v int64) map[string]*plugin_pb.ConfigValue {
+		return map[string]*plugin_pb.ConfigValue{
+			"estimated_runtime_seconds": {Kind: &plugin_pb.ConfigValue_Int64Value{Int64Value: v}},
+		}
+	}
+
+	if got := scheduledEstimatedRuntime(nil); got != 0 {
+		t.Fatalf("nil parameters: got=%v want=0", got)
+	}
+	if got := scheduledEstimatedRuntime(estimate(-5)); got != 0 {
+		t.Fatalf("negative estimate: got=%v want=0", got)
+	}
+	if got := scheduledEstimatedRuntime(estimate(600)); got != 10*time.Minute {
+		t.Fatalf("normal estimate: got=%v want=10m", got)
+	}
+	// Oversized seconds cap before the Duration conversion can overflow.
+	if got := scheduledEstimatedRuntime(estimate(math.MaxInt64)); got != maxEstimatedRuntimeCap {
+		t.Fatalf("oversized estimate: got=%v want=%v", got, maxEstimatedRuntimeCap)
+	}
+}
+
+func TestDispatchScheduledProposalsDrainsStartedJobOnWindowClose(t *testing.T) {
+	t.Parallel()
+
+	pluginSvc, err := New(Options{})
+	if err != nil {
+		t.Fatalf("New: %v", err)
+	}
+	defer pluginSvc.Shutdown()
+
+	const workerID = "worker-drain"
+	const jobType = "s3_lifecycle" // lifecycle lane: no admin lock needed
+	pluginSvc.registry.UpsertFromHello(&plugin_pb.WorkerHello{
+		WorkerId: workerID,
+		Capabilities: []*plugin_pb.JobTypeCapability{
+			{JobType: jobType, CanExecute: true, MaxExecutionConcurrency: 1},
+		},
+	})
+	session := &streamSession{workerID: workerID, outgoing: make(chan *plugin_pb.AdminToWorkerMessage, 8), done: make(chan struct{})}
+	pluginSvc.putSession(session)
+
+	windowCtx, closeWindow := context.WithCancel(context.Background())
+	defer closeWindow()
+
+	policy := schedulerPolicy{
+		ExecutionConcurrency:   1,
+		PerWorkerConcurrency:   1,
+		ExecutionTimeout:       time.Hour,
+		ExecutorReserveBackoff: time.Millisecond,
+	}
+	proposals := []*plugin_pb.JobProposal{
+		{ProposalId: "p1", JobType: jobType, DedupeKey: "k1"},
+		{ProposalId: "p2", JobType: jobType, DedupeKey: "k2"},
+	}
+
+	resultCh := make(chan [3]int, 1)
+	go func() {
+		success, errCount, canceled := pluginSvc.dispatchScheduledProposals(
+			windowCtx, jobType, proposals, &plugin_pb.ClusterContext{}, policy)
+		resultCh <- [3]int{success, errCount, canceled}
+	}()
+
+	first := <-session.outgoing
+	execReq := first.GetExecuteJobRequest()
+	if execReq == nil {
+		t.Fatalf("expected execute_job_request, got %+v", first)
+	}
+
+	// Close the window mid-job: the job must drain to completion while the
+	// still-queued second proposal is canceled.
+	closeWindow()
+	pluginSvc.handleJobCompleted(&plugin_pb.JobCompleted{
+		RequestId:   first.RequestId,
+		JobId:       execReq.Job.JobId,
+		JobType:     jobType,
+		Success:     true,
+		CompletedAt: timestamppb.Now(),
+	})
+
+	result := <-resultCh
+	if result[0] != 1 || result[1] != 0 || result[2] != 1 {
+		t.Fatalf("unexpected dispatch counts: success=%d errors=%d canceled=%d", result[0], result[1], result[2])
+	}
+
+	select {
+	case unexpected := <-session.outgoing:
+		t.Fatalf("expected no further worker messages (no cancel for the draining job), got %+v", unexpected)
+	default:
+	}
+}
+
 // ---------- lane-scoped prune ----------

 func TestPruneSchedulerState_DefaultLaneKeepsForeignLanesAndPrunesOwnStale(t *testing.T) {