From 78e7e0437778e380381ab5addf7bb1567703d291 Mon Sep 17 00:00:00 2001 From: Chris Lu Date: Wed, 12 Aug 2026 19:38:22 -0700 Subject: [PATCH] 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 --- weed/admin/plugin/plugin_scheduler.go | 71 +++++++---- weed/admin/plugin/plugin_scheduler_test.go | 140 +++++++++++++++++++++ 2 files changed, 186 insertions(+), 25 deletions(-) diff --git a/weed/admin/plugin/plugin_scheduler.go b/weed/admin/plugin/plugin_scheduler.go index 294bfd6f1..94a3aa397 100644 --- a/weed/admin/plugin/plugin_scheduler.go +++ b/weed/admin/plugin/plugin_scheduler.go @@ -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) diff --git a/weed/admin/plugin/plugin_scheduler_test.go b/weed/admin/plugin/plugin_scheduler_test.go index 9b1d5e810..2e9459bb9 100644 --- a/weed/admin/plugin/plugin_scheduler_test.go +++ b/weed/admin/plugin/plugin_scheduler_test.go @@ -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) {