From 8923d128b815bd1ba876ebee8495f1c7fa2aebc6 Mon Sep 17 00:00:00 2001 From: Mohammed Nafees Date: Thu, 25 Jun 2026 16:32:11 +0200 Subject: [PATCH 1/7] ondemand lease for coldstart concurrency queue --- internal/msgqueue/message_ids.go | 1 - .../services/controllers/task/controller.go | 22 ++- internal/services/scheduler/v1/scheduler.go | 15 -- .../services/shared/tasktypes/v1/scheduler.go | 18 -- .../v1/concurrency_integration_test.go | 154 ++++++++++++++++++ pkg/scheduling/v1/lease_manager.go | 51 ++++-- pkg/scheduling/v1/tenant_manager.go | 47 +++++- 7 files changed, 239 insertions(+), 69 deletions(-) diff --git a/internal/msgqueue/message_ids.go b/internal/msgqueue/message_ids.go index cb787210d3..dffabebdec 100644 --- a/internal/msgqueue/message_ids.go +++ b/internal/msgqueue/message_ids.go @@ -9,7 +9,6 @@ const ( MsgIDCheckTenantQueue = "check-tenant-queue" MsgIDNewWorker = "new-worker" MsgIDNewQueue = "new-queue" - MsgIDNewConcurrencyStrategy = "new-concurrency-strategy" MsgIDCreateMonitoringEvent = "create-monitoring-event" MsgIDCreatedDAG = "created-dag" MsgIDCreatedEventTrigger = "created-event-trigger" diff --git a/internal/services/controllers/task/controller.go b/internal/services/controllers/task/controller.go index bd5dc0c17d..64d692a5e9 100644 --- a/internal/services/controllers/task/controller.go +++ b/internal/services/controllers/task/controller.go @@ -60,7 +60,7 @@ type TasksControllerImpl struct { retryTaskOperations *operation.TenantOperationPool emitSleepOperations *operation.TenantOperationPool evictExpiredIdempotencyKeysOperations *operation.TenantOperationPool - // deactivateStaleStepConcurrencyOperations *operation.TenantOperationPool + deactivateStaleStepConcurrencyOperations *operation.TenantOperationPool replayEnabled bool analyzeCronInterval time.Duration signaler *signal.OLAPSignaler @@ -286,16 +286,14 @@ func New(fs ...TasksControllerOpt) (*TasksControllerImpl, error) { opts.repov1.Tasks().DefaultTaskActivityGauge, )) - // FIXME(mnafees): temporarily disabling this operation for the meantime - // - // t.deactivateStaleStepConcurrencyOperations = operation.NewTenantOperationPool(opts.p, opts.l, "deactivate-stale-step-concurrency", timeout, "deactivate stale step concurrency", t.deactivateStaleStepConcurrency, operation.WithPoolInterval( - // opts.repov1.IntervalSettings(), - // jitter, - // 15*time.Minute, - // 30*time.Minute, - // 3, - // opts.repov1.Tasks().DefaultTaskActivityGauge, - // )) + t.deactivateStaleStepConcurrencyOperations = operation.NewTenantOperationPool(opts.p, opts.l, "deactivate-stale-step-concurrency", timeout, "deactivate stale step concurrency", t.deactivateStaleStepConcurrency, operation.WithPoolInterval( + opts.repov1.IntervalSettings(), + jitter, + 15*time.Minute, + 30*time.Minute, + 3, + opts.repov1.Tasks().DefaultTaskActivityGauge, + )) return t, nil } @@ -412,7 +410,7 @@ func (tc *TasksControllerImpl) Start() (func() error, error) { tc.retryTaskOperations.Cleanup() tc.emitSleepOperations.Cleanup() tc.evictExpiredIdempotencyKeysOperations.Cleanup() - // tc.deactivateStaleStepConcurrencyOperations.Cleanup() + tc.deactivateStaleStepConcurrencyOperations.Cleanup() tc.pubBuffer.Stop() diff --git a/internal/services/scheduler/v1/scheduler.go b/internal/services/scheduler/v1/scheduler.go index 28130a94d9..1c4fd6e15f 100644 --- a/internal/services/scheduler/v1/scheduler.go +++ b/internal/services/scheduler/v1/scheduler.go @@ -328,8 +328,6 @@ func (s *Scheduler) handleTask(ctx context.Context, task *msgqueue.Message) (err return s.handleNewWorker(ctx, task) case msgqueue.MsgIDNewQueue: return s.handleNewQueue(ctx, task) - case msgqueue.MsgIDNewConcurrencyStrategy: - return s.handleNewConcurrencyStrategy(ctx, task) } return fmt.Errorf("unknown task: %s", task.ID) @@ -384,19 +382,6 @@ func (s *Scheduler) handleNewQueue(ctx context.Context, msg *msgqueue.Message) e return nil } -func (s *Scheduler) handleNewConcurrencyStrategy(ctx context.Context, msg *msgqueue.Message) error { - ctx, span := telemetry.NewSpanWithCarrier(ctx, "handle-new-concurrency-strategy", msg.OtelCarrier) - defer span.End() - - payloads := msgqueue.JSONConvert[tasktypes.NewConcurrencyStrategyPayload](msg.Payloads) - - for _, payload := range payloads { - s.pool.NotifyNewConcurrencyStrategy(ctx, msg.TenantID, payload.StrategyId) - } - - return nil -} - func (s *Scheduler) runSetTenants(ctx context.Context) func() { return func() { s.l.Debug().Ctx(ctx).Msgf("partition: checking step run requeue") diff --git a/internal/services/shared/tasktypes/v1/scheduler.go b/internal/services/shared/tasktypes/v1/scheduler.go index cd56cc0c92..075c760739 100644 --- a/internal/services/shared/tasktypes/v1/scheduler.go +++ b/internal/services/shared/tasktypes/v1/scheduler.go @@ -122,21 +122,3 @@ func NotifyNewQueue(tenantId uuid.UUID, queueName string) (*msgqueue.Message, er payload, ) } - -type NewConcurrencyStrategyPayload struct { - StrategyId int64 `json:"strategy_id"` -} - -func NotifyNewConcurrencyStrategy(tenantId uuid.UUID, strategyId int64) (*msgqueue.Message, error) { - payload := NewConcurrencyStrategyPayload{ - StrategyId: strategyId, - } - - return msgqueue.NewTenantMessage( - tenantId, - msgqueue.MsgIDNewConcurrencyStrategy, - true, - false, - payload, - ) -} diff --git a/pkg/scheduling/v1/concurrency_integration_test.go b/pkg/scheduling/v1/concurrency_integration_test.go index bea2849442..6a8ffc2f28 100644 --- a/pkg/scheduling/v1/concurrency_integration_test.go +++ b/pkg/scheduling/v1/concurrency_integration_test.go @@ -439,3 +439,157 @@ func TestConcurrency_MultipleStrategiesContention(t *testing.T) { return nil }) } + +// TestConcurrency_ColdStrategyScheduledPromptly is a regression test for cold-start scheduling +// latency on concurrency-keyed tasks. +// +// A concurrency strategy that has been idle long enough to be deactivated (is_active=FALSE, no +// slots) has no ConcurrencyManager running in the scheduler. When the next task arrives, the +// scheduler is notified via NotifyConcurrency (the in-process effect of the CheckTenantQueue / +// NotifyTaskCreated message published on task creation). The scheduler must create a manager for +// that strategy and schedule the waiting task. +// +// Before the fix, NotifyConcurrency only woke *already-running* managers, so a cold strategy was +// not picked up until the periodic lease-acquisition poll (every 5s) happened to discover it -- the +// "first run is slow, then warm" symptom (observed as intermittent 5-11s queued->started spikes). +// After the fix, NotifyConcurrency acquires the lease on-demand and the new manager runs immediately. +// +// This test reproduces the cold path through the full SchedulingPool and asserts the task is queued +// well within the 5s lease-poll window. Run against the pre-fix commit and it waits ~5s and fails the +// deadline; against the fix it completes in milliseconds. +func TestConcurrency_ColdStrategyScheduledPromptly(t *testing.T) { + runWithDatabase(t, func(conf *database.Layer) error { + ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + defer cancel() + requireSchedulerSchema(t, ctx, conf) + + r := conf.V1 + queries := sqlcv1.New() + + tenantId := uuid.New() + tenant, err := r.Tenant().CreateTenant(ctx, &repo.CreateTenantOpts{ + ID: &tenantId, + Name: "concurrency-cold-start-test", + Slug: fmt.Sprintf("concurrency-cold-start-test-%s", tenantId.String()), + }) + require.NoError(t, err) + + desc := "test workflow for cold-start scheduling" + groupRR := "GROUP_ROUND_ROBIN" + var maxRuns int32 = 10 + + wfVersion, err := r.Workflows().PutWorkflowVersion(ctx, tenantId, &repo.CreateWorkflowVersionOpts{ + Name: "concurrency-cold-start-test", + Description: &desc, + Tasks: []repo.CreateStepOpts{ + { + ReadableId: "my-task", + Action: "test:run", + Concurrency: []repo.CreateConcurrencyOpts{ + { + MaxRuns: &maxRuns, + LimitStrategy: &groupRR, + Expression: "input.my_id", + }, + }, + }, + }, + }) + require.NoError(t, err) + + workflowId := wfVersion.WorkflowVersion.WorkflowId + workflowVersionId := wfVersion.WorkflowVersion.ID + + strategies, err := queries.ListActiveConcurrencyStrategies(ctx, conf.Pool, tenantId) + require.NoError(t, err) + require.Len(t, strategies, 1) + strat := strategies[0] + + concurrencyRepo := r.Scheduler().Concurrency() + + // Make the strategy "cold": with no slots yet, the stale-deactivation sweep flips it to + // is_active=FALSE. This mirrors a strategy that has been idle for the deactivation window. + require.NoError(t, concurrencyRepo.DeactivateStaleStepConcurrency(ctx, tenantId)) + + activeAfterDeactivate, err := queries.ListActiveConcurrencyStrategies(ctx, conf.Pool, tenantId) + require.NoError(t, err) + require.Len(t, activeAfterDeactivate, 0, "strategy must be inactive so no manager is leased at pool start") + + // Start the scheduling pool. Its initial lease acquisition runs now and finds no active + // strategy, so no ConcurrencyManager exists for our strategy -- it is genuinely cold. + l := zerolog.Nop() + schedulingPool, cleanup, err := v1.NewSchedulingPool( + r.Scheduler(), + &l, + 100, + 20, + 10*time.Millisecond, + 20*time.Millisecond, + 50*time.Millisecond, + 100*time.Millisecond, + 5*time.Millisecond, + false, + 1, + nil, + ) + require.NoError(t, err) + defer func() { _ = cleanup() }() + + schedulingPool.SetTenants([]*sqlcv1.Tenant{tenant}) + resultsChan := schedulingPool.GetConcurrencyResultsCh() + + // A task arrives for the cold strategy. Creating the task inserts a concurrency slot, whose + // insert trigger reactivates the strategy in the DB (is_active=TRUE) -- but the scheduler does + // not know about it yet. + taskParams := newCreateTasksParams(1) + taskParams.Tenantids[0] = tenantId + taskParams.Queues[0] = "default" + taskParams.Actionids[0] = "test:run" + taskParams.Stepids[0] = strat.StepID + taskParams.Stepreadableids[0] = "my-task" + taskParams.Workflowids[0] = workflowId + taskParams.Scheduletimeouts[0] = "5m" + taskParams.Priorities[0] = 1 + taskParams.Stickies[0] = string(sqlcv1.V1StickyStrategyNONE) + taskParams.Externalids[0] = uuid.New() + taskParams.Displaynames[0] = "cold-task" + taskParams.Inputs[0] = []byte(`{"my_id": "thread-1"}`) + taskParams.Additionalmetadatas[0] = []byte(`{}`) + taskParams.InitialStates[0] = string(sqlcv1.V1TaskInitialStateQUEUED) + taskParams.Concurrencyparentstrategyids[0] = []pgtype.Int8{{}} + taskParams.ConcurrencyStrategyIds[0] = []int64{strat.ID} + taskParams.ConcurrencyKeys[0] = []string{"thread-1"} + taskParams.WorkflowVersionIds[0] = workflowVersionId + taskParams.WorkflowRunIds[0] = uuid.New() + + tasks, err := queries.CreateTasks(ctx, conf.Pool, taskParams) + require.NoError(t, err) + require.Len(t, tasks, 1) + + // This is exactly what the task-creation message does in production. + start := time.Now() + schedulingPool.NotifyConcurrency(ctx, tenantId, []int64{strat.ID}) + + // The fix must schedule the waiting task without waiting for the 5s lease-acquisition poll. + const coldStartDeadline = 2 * time.Second + deadline := time.NewTimer(coldStartDeadline) + defer deadline.Stop() + + for { + select { + case <-ctx.Done(): + t.Fatalf("context cancelled while waiting for cold strategy to be scheduled: %v", ctx.Err()) + case <-deadline.C: + t.Fatalf("cold concurrency strategy was not scheduled within %s; "+ + "the on-demand manager was not created (fell back to the periodic lease poll)", coldStartDeadline) + case res := <-resultsChan: + require.False(t, res.RunConcurrencyResult.FailedAdvisoryLock) + if len(res.RunConcurrencyResult.Queued) > 0 { + require.Len(t, res.RunConcurrencyResult.Queued, 1, "the single waiting task should be queued") + t.Logf("cold strategy scheduled in %s", time.Since(start)) + return nil + } + } + } + }) +} diff --git a/pkg/scheduling/v1/lease_manager.go b/pkg/scheduling/v1/lease_manager.go index d6aaf28896..eef555b053 100644 --- a/pkg/scheduling/v1/lease_manager.go +++ b/pkg/scheduling/v1/lease_manager.go @@ -418,52 +418,67 @@ func (l *LeaseManager) acquireConcurrencyLeases(ctx context.Context) error { return nil } -func (l *LeaseManager) notifyNewConcurrencyStrategy(ctx context.Context, strategyId int64) error { +// acquireConcurrencyStrategyOnDemand acquires a lease for a single concurrency strategy on-demand +// (the cold-start path). It returns the strategy when a ConcurrencyManager should be created for it, +// or nil when the lease is already held here or is owned by another scheduler. +// +// Unlike the periodic lease loop, this takes a BLOCKING lock and returns the strategy to the caller +// rather than delivering it over the (non-blocking, lossy) lease channel. Callers run this off the +// hot path in their own goroutine, so blocking is fine -- and it guarantees the acquisition is not +// silently dropped under contention (the previous TryLock + non-blocking-send approach would no-op +// when several cold strategies raced each other or the periodic poll, falling back to the 5s poll). +func (l *LeaseManager) acquireConcurrencyStrategyOnDemand(ctx context.Context, strategyId int64) (*sqlcv1.V1StepConcurrency, error) { l.processMu.RLock() defer l.processMu.RUnlock() if l.cleanedUp { - return nil - } - - if !l.concurrencyLeasesMu.TryLock() { - return nil + return nil, nil } + l.concurrencyLeasesMu.Lock() defer l.concurrencyLeasesMu.Unlock() - // check that we don't already have a lease for this concurrency strategy + ctx, cancel := context.WithTimeout(ctx, 1*time.Second) + defer cancel() + + // Callers only reach this path when NO manager exists for the strategy, so we must always return + // the strategy and let the caller (idempotently) ensure a manager. Whether we already hold the + // lease only decides if we re-acquire it: a lease can be held WITHOUT a live manager (e.g. the + // bulk periodic lease-channel send was dropped), and in that case returning nil here would leave + // the cold task waiting for the next poll. Holding the lease must never block manager creation. + alreadyLeased := false for _, lease := range l.concurrencyLeases { if lease.ResourceId == fmt.Sprintf("%d", strategyId) { - return nil + alreadyLeased = true + break } } - ctx, cancel := context.WithTimeout(ctx, 1*time.Second) - defer cancel() - strategy, err := l.lr.GetConcurrencyStrategy(ctx, l.tenantId, strategyId) if err != nil { - return err + return nil, err + } + + if alreadyLeased { + // lease already held here; just ensure the manager exists (no re-acquire needed) + return strategy, nil } lease, err := l.lr.AcquireOrExtendLeases(ctx, l.tenantId, sqlcv1.LeaseKindCONCURRENCYSTRATEGY, []string{fmt.Sprintf("%d", strategyId)}, []*sqlcv1.Lease{}) if err != nil { - return err + return nil, err } if len(lease) == 0 || lease[0].ResourceId == "" { - return nil + // the lease is genuinely owned by another scheduler; it will manage this strategy + return nil, nil } l.concurrencyLeases = append(l.concurrencyLeases, lease...) - // send the new concurrency strategy to the channel - l.sendConcurrencyLeases([]*sqlcv1.V1StepConcurrency{strategy}, true) - - return nil + return strategy, nil } func (l *LeaseManager) acquireAllLeases(ctx context.Context) { diff --git a/pkg/scheduling/v1/tenant_manager.go b/pkg/scheduling/v1/tenant_manager.go index 5797412195..5905282127 100644 --- a/pkg/scheduling/v1/tenant_manager.go +++ b/pkg/scheduling/v1/tenant_manager.go @@ -177,7 +177,7 @@ func (t *tenantManager) listenForConcurrencyLeases(ctx context.Context) { case msg := <-t.concurrencyCh: if msg.isIncremental { for _, strategy := range msg.items { - t.addConcurrencyStrategy(strategy) + t.addConcurrencyStrategy(ctx, strategy) } } else { t.setConcurrencyStrategies(msg.items) @@ -269,7 +269,7 @@ func (t *tenantManager) setConcurrencyStrategies(strategies []*sqlcv1.V1StepConc t.concurrencyStrategies = newArr } -func (t *tenantManager) addConcurrencyStrategy(strategy *sqlcv1.V1StepConcurrency) { +func (t *tenantManager) addConcurrencyStrategy(ctx context.Context, strategy *sqlcv1.V1StepConcurrency) { t.concurrencyMu.Lock() defer t.concurrencyMu.Unlock() @@ -279,7 +279,15 @@ func (t *tenantManager) addConcurrencyStrategy(strategy *sqlcv1.V1StepConcurrenc } } - t.concurrencyStrategies = append(t.concurrencyStrategies, newConcurrencyManager(t.cf, t.tenantId, strategy, t.concurrencyResultsCh, t.concurrencyAdvisoryLock, t.concurrencyParentAdvisoryLock)) + c := newConcurrencyManager(t.cf, t.tenantId, strategy, t.concurrencyResultsCh, t.concurrencyAdvisoryLock, t.concurrencyParentAdvisoryLock) + t.concurrencyStrategies = append(t.concurrencyStrategies, c) + + // This manager was created on-demand (incremental lease acquisition), which happens when a task + // arrives for a strategy this scheduler isn't already running (i.e. a "cold" strategy). Notify it + // immediately so it schedules the waiting task on its next loop iteration instead of waiting up to + // maxPollingInterval for the first ticker tick. The bulk setConcurrencyStrategies path deliberately + // does NOT do this, so startup/rebalance keeps the randomticker's load-spreading. + c.notify(ctx) } func (t *tenantManager) replenish(ctx context.Context) { @@ -304,6 +312,9 @@ func (t *tenantManager) notifyConcurrency(ctx context.Context, strategyIds []int continue } + // mark this strategy as already managed so we don't try to acquire a lease for it below + delete(strategyIdsMap, c.strategy.ID) + c.notify(ctx) childStrategyIds := make([]int64, 0) @@ -354,6 +365,22 @@ func (t *tenantManager) notifyConcurrency(ctx context.Context, strategyIds []int } t.concurrencyMu.RUnlock() + + // Any strategy ids still in the map have no manager running on this scheduler. This is the "cold" + // case: a task arrived for a strategy that was deactivated while idle (or has never been leased here). + // Acquire the lease on-demand so a ConcurrencyManager is created now, instead of waiting up to the + // next lease-acquisition poll (every 5s) for it to be discovered. notifyNewConcurrencyStrategy does a + // DB lease acquisition, so run it off the hot notify path; it is a no-op if the lease is already held + // or owned by another scheduler. + if len(strategyIdsMap) > 0 { + // detach from the request context so cancellation when the message handler returns doesn't + // abort the lease-acquisition DB call, while still propagating telemetry/values + leaseCtx := context.WithoutCancel(ctx) + + for id := range strategyIdsMap { + go t.notifyNewConcurrencyStrategy(leaseCtx, id) + } + } } func (t *tenantManager) notifyNewWorker(ctx context.Context, workerId uuid.UUID) { @@ -377,12 +404,22 @@ func (t *tenantManager) notifyNewQueue(ctx context.Context, queueName string) { } func (t *tenantManager) notifyNewConcurrencyStrategy(ctx context.Context, strategyId int64) { - err := t.leaseManager.notifyNewConcurrencyStrategy(ctx, strategyId) + // Acquire the lease and create the manager directly, rather than handing the strategy to the + // lease channel (a non-blocking send that can drop the message under load). addConcurrencyStrategy + // is idempotent and notifies the new manager immediately, so the cold task is scheduled at once. + strategy, err := t.leaseManager.acquireConcurrencyStrategyOnDemand(ctx, strategyId) if err != nil { - t.l.Error().Err(err).Msg("error notifying new concurrency strategy") + t.l.Error().Err(err).Msgf("error acquiring on-demand lease for concurrency strategy %d", strategyId) return } + + if strategy == nil { + // lease genuinely owned by another scheduler; that scheduler will manage it + return + } + + t.addConcurrencyStrategy(ctx, strategy) } func (t *tenantManager) queue(ctx context.Context, queueNames []string) { From 6c4daf660dba838e4899c235a890ee529889b657 Mon Sep 17 00:00:00 2001 From: Mohammed Nafees Date: Fri, 26 Jun 2026 11:21:39 +0200 Subject: [PATCH 2/7] PR comments --- internal/msgqueue/message_ids.go | 1 + internal/services/scheduler/v1/scheduler.go | 15 +++++ .../services/shared/tasktypes/v1/scheduler.go | 18 ++++++ pkg/scheduling/v1/lease_manager.go | 64 ++++++++----------- pkg/scheduling/v1/tenant_manager.go | 24 ++----- 5 files changed, 69 insertions(+), 53 deletions(-) diff --git a/internal/msgqueue/message_ids.go b/internal/msgqueue/message_ids.go index dffabebdec..cb787210d3 100644 --- a/internal/msgqueue/message_ids.go +++ b/internal/msgqueue/message_ids.go @@ -9,6 +9,7 @@ const ( MsgIDCheckTenantQueue = "check-tenant-queue" MsgIDNewWorker = "new-worker" MsgIDNewQueue = "new-queue" + MsgIDNewConcurrencyStrategy = "new-concurrency-strategy" MsgIDCreateMonitoringEvent = "create-monitoring-event" MsgIDCreatedDAG = "created-dag" MsgIDCreatedEventTrigger = "created-event-trigger" diff --git a/internal/services/scheduler/v1/scheduler.go b/internal/services/scheduler/v1/scheduler.go index 1c4fd6e15f..28130a94d9 100644 --- a/internal/services/scheduler/v1/scheduler.go +++ b/internal/services/scheduler/v1/scheduler.go @@ -328,6 +328,8 @@ func (s *Scheduler) handleTask(ctx context.Context, task *msgqueue.Message) (err return s.handleNewWorker(ctx, task) case msgqueue.MsgIDNewQueue: return s.handleNewQueue(ctx, task) + case msgqueue.MsgIDNewConcurrencyStrategy: + return s.handleNewConcurrencyStrategy(ctx, task) } return fmt.Errorf("unknown task: %s", task.ID) @@ -382,6 +384,19 @@ func (s *Scheduler) handleNewQueue(ctx context.Context, msg *msgqueue.Message) e return nil } +func (s *Scheduler) handleNewConcurrencyStrategy(ctx context.Context, msg *msgqueue.Message) error { + ctx, span := telemetry.NewSpanWithCarrier(ctx, "handle-new-concurrency-strategy", msg.OtelCarrier) + defer span.End() + + payloads := msgqueue.JSONConvert[tasktypes.NewConcurrencyStrategyPayload](msg.Payloads) + + for _, payload := range payloads { + s.pool.NotifyNewConcurrencyStrategy(ctx, msg.TenantID, payload.StrategyId) + } + + return nil +} + func (s *Scheduler) runSetTenants(ctx context.Context) func() { return func() { s.l.Debug().Ctx(ctx).Msgf("partition: checking step run requeue") diff --git a/internal/services/shared/tasktypes/v1/scheduler.go b/internal/services/shared/tasktypes/v1/scheduler.go index 075c760739..cd56cc0c92 100644 --- a/internal/services/shared/tasktypes/v1/scheduler.go +++ b/internal/services/shared/tasktypes/v1/scheduler.go @@ -122,3 +122,21 @@ func NotifyNewQueue(tenantId uuid.UUID, queueName string) (*msgqueue.Message, er payload, ) } + +type NewConcurrencyStrategyPayload struct { + StrategyId int64 `json:"strategy_id"` +} + +func NotifyNewConcurrencyStrategy(tenantId uuid.UUID, strategyId int64) (*msgqueue.Message, error) { + payload := NewConcurrencyStrategyPayload{ + StrategyId: strategyId, + } + + return msgqueue.NewTenantMessage( + tenantId, + msgqueue.MsgIDNewConcurrencyStrategy, + true, + false, + payload, + ) +} diff --git a/pkg/scheduling/v1/lease_manager.go b/pkg/scheduling/v1/lease_manager.go index eef555b053..b68a7a02f9 100644 --- a/pkg/scheduling/v1/lease_manager.go +++ b/pkg/scheduling/v1/lease_manager.go @@ -98,12 +98,13 @@ func (l *LeaseManager) sendConcurrencyLeases(concurrencyLeases []*sqlcv1.V1StepC } }() - select { - case l.concurrencyLeasesCh <- notifierMsg[*sqlcv1.V1StepConcurrency]{ + // Blocking send: an acquired concurrency lease must never be dropped, otherwise its manager is + // never instantiated and a cold strategy waits for the next lease poll to schedule its task. The + // listener is always reading except during cleanup, where processMu serializes against this send + // and the recover() above handles a racing send on a closed channel. + l.concurrencyLeasesCh <- notifierMsg[*sqlcv1.V1StepConcurrency]{ items: concurrencyLeases, isIncremental: isIncremental, - }: - default: } } @@ -418,67 +419,58 @@ func (l *LeaseManager) acquireConcurrencyLeases(ctx context.Context) error { return nil } -// acquireConcurrencyStrategyOnDemand acquires a lease for a single concurrency strategy on-demand -// (the cold-start path). It returns the strategy when a ConcurrencyManager should be created for it, -// or nil when the lease is already held here or is owned by another scheduler. +// notifyNewConcurrencyStrategy acquires a lease for a single concurrency strategy on-demand and hands +// it to the lease channel, which instantiates the ConcurrencyManager. This mirrors notifyNewQueue. // -// Unlike the periodic lease loop, this takes a BLOCKING lock and returns the strategy to the caller -// rather than delivering it over the (non-blocking, lossy) lease channel. Callers run this off the -// hot path in their own goroutine, so blocking is fine -- and it guarantees the acquisition is not -// silently dropped under contention (the previous TryLock + non-blocking-send approach would no-op -// when several cold strategies raced each other or the periodic poll, falling back to the 5s poll). -func (l *LeaseManager) acquireConcurrencyStrategyOnDemand(ctx context.Context, strategyId int64) (*sqlcv1.V1StepConcurrency, error) { +// It is a no-op when we already hold the lease for the strategy: because sendConcurrencyLeases is a +// blocking send, holding the lease implies a manager was (or is being) created for it, so there is +// nothing to do. +func (l *LeaseManager) notifyNewConcurrencyStrategy(ctx context.Context, strategyId int64) error { l.processMu.RLock() defer l.processMu.RUnlock() if l.cleanedUp { - return nil, nil + return nil } - l.concurrencyLeasesMu.Lock() - defer l.concurrencyLeasesMu.Unlock() + if !l.concurrencyLeasesMu.TryLock() { + return nil + } - ctx, cancel := context.WithTimeout(ctx, 1*time.Second) - defer cancel() + defer l.concurrencyLeasesMu.Unlock() - // Callers only reach this path when NO manager exists for the strategy, so we must always return - // the strategy and let the caller (idempotently) ensure a manager. Whether we already hold the - // lease only decides if we re-acquire it: a lease can be held WITHOUT a live manager (e.g. the - // bulk periodic lease-channel send was dropped), and in that case returning nil here would leave - // the cold task waiting for the next poll. Holding the lease must never block manager creation. - alreadyLeased := false + // check that we don't already have a lease for this concurrency strategy for _, lease := range l.concurrencyLeases { if lease.ResourceId == fmt.Sprintf("%d", strategyId) { - alreadyLeased = true - break + return nil } } + ctx, cancel := context.WithTimeout(ctx, 1*time.Second) + defer cancel() + strategy, err := l.lr.GetConcurrencyStrategy(ctx, l.tenantId, strategyId) if err != nil { - return nil, err - } - - if alreadyLeased { - // lease already held here; just ensure the manager exists (no re-acquire needed) - return strategy, nil + return err } lease, err := l.lr.AcquireOrExtendLeases(ctx, l.tenantId, sqlcv1.LeaseKindCONCURRENCYSTRATEGY, []string{fmt.Sprintf("%d", strategyId)}, []*sqlcv1.Lease{}) if err != nil { - return nil, err + return err } if len(lease) == 0 || lease[0].ResourceId == "" { - // the lease is genuinely owned by another scheduler; it will manage this strategy - return nil, nil + // the lease is owned by another scheduler; it will manage this strategy + return nil } l.concurrencyLeases = append(l.concurrencyLeases, lease...) - return strategy, nil + l.sendConcurrencyLeases([]*sqlcv1.V1StepConcurrency{strategy}, true) + + return nil } func (l *LeaseManager) acquireAllLeases(ctx context.Context) { diff --git a/pkg/scheduling/v1/tenant_manager.go b/pkg/scheduling/v1/tenant_manager.go index 5905282127..a65630b59d 100644 --- a/pkg/scheduling/v1/tenant_manager.go +++ b/pkg/scheduling/v1/tenant_manager.go @@ -282,11 +282,11 @@ func (t *tenantManager) addConcurrencyStrategy(ctx context.Context, strategy *sq c := newConcurrencyManager(t.cf, t.tenantId, strategy, t.concurrencyResultsCh, t.concurrencyAdvisoryLock, t.concurrencyParentAdvisoryLock) t.concurrencyStrategies = append(t.concurrencyStrategies, c) - // This manager was created on-demand (incremental lease acquisition), which happens when a task - // arrives for a strategy this scheduler isn't already running (i.e. a "cold" strategy). Notify it - // immediately so it schedules the waiting task on its next loop iteration instead of waiting up to - // maxPollingInterval for the first ticker tick. The bulk setConcurrencyStrategies path deliberately - // does NOT do this, so startup/rebalance keeps the randomticker's load-spreading. + // Notify the newly instantiated manager so it evaluates its already-queued slots on the next loop + // iteration, instead of waiting up to maxPollingInterval for its first ticker tick. This is what + // lets a cold strategy schedule its waiting task promptly once its lease is acquired. The bulk + // setConcurrencyStrategies path deliberately does NOT do this, so startup/rebalance keeps the + // randomticker's load-spreading. c.notify(ctx) } @@ -404,22 +404,12 @@ func (t *tenantManager) notifyNewQueue(ctx context.Context, queueName string) { } func (t *tenantManager) notifyNewConcurrencyStrategy(ctx context.Context, strategyId int64) { - // Acquire the lease and create the manager directly, rather than handing the strategy to the - // lease channel (a non-blocking send that can drop the message under load). addConcurrencyStrategy - // is idempotent and notifies the new manager immediately, so the cold task is scheduled at once. - strategy, err := t.leaseManager.acquireConcurrencyStrategyOnDemand(ctx, strategyId) + err := t.leaseManager.notifyNewConcurrencyStrategy(ctx, strategyId) if err != nil { - t.l.Error().Err(err).Msgf("error acquiring on-demand lease for concurrency strategy %d", strategyId) + t.l.Error().Err(err).Msg("error notifying new concurrency strategy") return } - - if strategy == nil { - // lease genuinely owned by another scheduler; that scheduler will manage it - return - } - - t.addConcurrencyStrategy(ctx, strategy) } func (t *tenantManager) queue(ctx context.Context, queueNames []string) { From cffa038c07bbfc2f9947b55d6bd5164678c1023b Mon Sep 17 00:00:00 2001 From: Mohammed Nafees Date: Fri, 26 Jun 2026 11:29:56 +0200 Subject: [PATCH 3/7] also notify from bulk path --- pkg/scheduling/v1/tenant_manager.go | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/pkg/scheduling/v1/tenant_manager.go b/pkg/scheduling/v1/tenant_manager.go index a65630b59d..c4e7edecaa 100644 --- a/pkg/scheduling/v1/tenant_manager.go +++ b/pkg/scheduling/v1/tenant_manager.go @@ -180,7 +180,7 @@ func (t *tenantManager) listenForConcurrencyLeases(ctx context.Context) { t.addConcurrencyStrategy(ctx, strategy) } } else { - t.setConcurrencyStrategies(msg.items) + t.setConcurrencyStrategies(ctx, msg.items) } } } @@ -238,7 +238,7 @@ func (t *tenantManager) addQueuer(queueName string) { t.queue(context.Background(), []string{queueName}) } -func (t *tenantManager) setConcurrencyStrategies(strategies []*sqlcv1.V1StepConcurrency) { +func (t *tenantManager) setConcurrencyStrategies(ctx context.Context, strategies []*sqlcv1.V1StepConcurrency) { t.concurrencyMu.Lock() defer t.concurrencyMu.Unlock() @@ -263,7 +263,14 @@ func (t *tenantManager) setConcurrencyStrategies(strategies []*sqlcv1.V1StepConc } for _, strategy := range strategiesSet { - newArr = append(newArr, newConcurrencyManager(t.cf, t.tenantId, strategy, t.concurrencyResultsCh, t.concurrencyAdvisoryLock, t.concurrencyParentAdvisoryLock)) + c := newConcurrencyManager(t.cf, t.tenantId, strategy, t.concurrencyResultsCh, t.concurrencyAdvisoryLock, t.concurrencyParentAdvisoryLock) + newArr = append(newArr, c) + + // Notify each newly instantiated manager so it evaluates its already-queued slots immediately, + // instead of waiting up to maxPollingInterval for its first ticker tick. This matters on + // startup/rebalance, where managers are (re)created here for strategies that may already have + // queued work. notify() is a non-blocking send, so it is safe to call under the lock. + c.notify(ctx) } t.concurrencyStrategies = newArr From 802003f5e31c21cd4bbb4e23354b1742be2b33a5 Mon Sep 17 00:00:00 2001 From: Mohammed Nafees Date: Fri, 26 Jun 2026 16:46:13 +0200 Subject: [PATCH 4/7] no need to enable to cron for now --- .../services/controllers/task/controller.go | 22 ++++++++++--------- 1 file changed, 12 insertions(+), 10 deletions(-) diff --git a/internal/services/controllers/task/controller.go b/internal/services/controllers/task/controller.go index 64d692a5e9..bd5dc0c17d 100644 --- a/internal/services/controllers/task/controller.go +++ b/internal/services/controllers/task/controller.go @@ -60,7 +60,7 @@ type TasksControllerImpl struct { retryTaskOperations *operation.TenantOperationPool emitSleepOperations *operation.TenantOperationPool evictExpiredIdempotencyKeysOperations *operation.TenantOperationPool - deactivateStaleStepConcurrencyOperations *operation.TenantOperationPool + // deactivateStaleStepConcurrencyOperations *operation.TenantOperationPool replayEnabled bool analyzeCronInterval time.Duration signaler *signal.OLAPSignaler @@ -286,14 +286,16 @@ func New(fs ...TasksControllerOpt) (*TasksControllerImpl, error) { opts.repov1.Tasks().DefaultTaskActivityGauge, )) - t.deactivateStaleStepConcurrencyOperations = operation.NewTenantOperationPool(opts.p, opts.l, "deactivate-stale-step-concurrency", timeout, "deactivate stale step concurrency", t.deactivateStaleStepConcurrency, operation.WithPoolInterval( - opts.repov1.IntervalSettings(), - jitter, - 15*time.Minute, - 30*time.Minute, - 3, - opts.repov1.Tasks().DefaultTaskActivityGauge, - )) + // FIXME(mnafees): temporarily disabling this operation for the meantime + // + // t.deactivateStaleStepConcurrencyOperations = operation.NewTenantOperationPool(opts.p, opts.l, "deactivate-stale-step-concurrency", timeout, "deactivate stale step concurrency", t.deactivateStaleStepConcurrency, operation.WithPoolInterval( + // opts.repov1.IntervalSettings(), + // jitter, + // 15*time.Minute, + // 30*time.Minute, + // 3, + // opts.repov1.Tasks().DefaultTaskActivityGauge, + // )) return t, nil } @@ -410,7 +412,7 @@ func (tc *TasksControllerImpl) Start() (func() error, error) { tc.retryTaskOperations.Cleanup() tc.emitSleepOperations.Cleanup() tc.evictExpiredIdempotencyKeysOperations.Cleanup() - tc.deactivateStaleStepConcurrencyOperations.Cleanup() + // tc.deactivateStaleStepConcurrencyOperations.Cleanup() tc.pubBuffer.Stop() From a171dd24d414e113735760cc279956f5f52fd194 Mon Sep 17 00:00:00 2001 From: Mohammed Nafees Date: Fri, 26 Jun 2026 17:15:45 +0200 Subject: [PATCH 5/7] buffered channel --- pkg/scheduling/v1/lease_manager.go | 27 +++++++++++++++++++++------ 1 file changed, 21 insertions(+), 6 deletions(-) diff --git a/pkg/scheduling/v1/lease_manager.go b/pkg/scheduling/v1/lease_manager.go index b68a7a02f9..30b677263e 100644 --- a/pkg/scheduling/v1/lease_manager.go +++ b/pkg/scheduling/v1/lease_manager.go @@ -17,6 +17,8 @@ import ( "github.com/hatchet-dev/hatchet/pkg/repository/sqlcv1" ) +const concurrencyLeasesChBuffer = 1024 + // LeaseManager is responsible for leases on multiple queues and multiplexing // queue results to callers. It is still tenant-scoped. type LeaseManager struct { @@ -46,7 +48,7 @@ type LeaseManager struct { func newLeaseManager(conf *sharedConfig, tenantId uuid.UUID) (*LeaseManager, notifierCh[*v1.ListActiveWorkersResult], notifierCh[string], notifierCh[*sqlcv1.V1StepConcurrency]) { workersCh := make(notifierCh[*v1.ListActiveWorkersResult]) queuesCh := make(notifierCh[string]) - concurrencyLeasesCh := make(notifierCh[*sqlcv1.V1StepConcurrency]) + concurrencyLeasesCh := make(notifierCh[*sqlcv1.V1StepConcurrency], concurrencyLeasesChBuffer) return &LeaseManager{ lr: conf.repo.Lease(), @@ -98,14 +100,27 @@ func (l *LeaseManager) sendConcurrencyLeases(concurrencyLeases []*sqlcv1.V1StepC } }() - // Blocking send: an acquired concurrency lease must never be dropped, otherwise its manager is - // never instantiated and a cold strategy waits for the next lease poll to schedule its task. The - // listener is always reading except during cleanup, where processMu serializes against this send - // and the recover() above handles a racing send on a closed channel. - l.concurrencyLeasesCh <- notifierMsg[*sqlcv1.V1StepConcurrency]{ + msg := notifierMsg[*sqlcv1.V1StepConcurrency]{ items: concurrencyLeases, isIncremental: isIncremental, } + + if isIncremental { + // On-demand (cold-start) fast path: best-effort, non-blocking. Bursts are absorbed by the + // buffer; if it ever fills, the strategy is still active in the DB and the bulk reconcile + // below (the ~5s lease poll) will instantiate its manager, so dropping here is safe. + select { + case l.concurrencyLeasesCh <- msg: + default: + } + return + } + + // Bulk reconcile from the ~5s lease-acquisition poll: this carries the full active set and creates + // any manager the fast path missed, so it must be delivered. Block until the listener accepts it + // (it drains continuously, so this only waits in the pathological full-buffer case). The recover() + // above handles a racing send on the channel closed during cleanup. + l.concurrencyLeasesCh <- msg } func (l *LeaseManager) acquireWorkerLeases(ctx context.Context) error { From b52d11e1af777a3f6816ce73241610263d6209cb Mon Sep 17 00:00:00 2001 From: Mohammed Nafees Date: Fri, 26 Jun 2026 18:07:58 +0200 Subject: [PATCH 6/7] drain --- pkg/scheduling/v1/lease_manager.go | 36 ++++++++++++++--------------- pkg/scheduling/v1/tenant_manager.go | 23 +++++------------- 2 files changed, 23 insertions(+), 36 deletions(-) diff --git a/pkg/scheduling/v1/lease_manager.go b/pkg/scheduling/v1/lease_manager.go index 30b677263e..b28e0edf71 100644 --- a/pkg/scheduling/v1/lease_manager.go +++ b/pkg/scheduling/v1/lease_manager.go @@ -100,27 +100,28 @@ func (l *LeaseManager) sendConcurrencyLeases(concurrencyLeases []*sqlcv1.V1StepC } }() - msg := notifierMsg[*sqlcv1.V1StepConcurrency]{ + // a full refresh supersedes anything still queued, so drain stale messages first + if !isIncremental { + l.drainConcurrencyLeasesCh() + } + + select { + case l.concurrencyLeasesCh <- notifierMsg[*sqlcv1.V1StepConcurrency]{ items: concurrencyLeases, isIncremental: isIncremental, + }: + default: } +} - if isIncremental { - // On-demand (cold-start) fast path: best-effort, non-blocking. Bursts are absorbed by the - // buffer; if it ever fills, the strategy is still active in the DB and the bulk reconcile - // below (the ~5s lease poll) will instantiate its manager, so dropping here is safe. +func (l *LeaseManager) drainConcurrencyLeasesCh() { + for { select { - case l.concurrencyLeasesCh <- msg: + case <-l.concurrencyLeasesCh: default: + return } - return } - - // Bulk reconcile from the ~5s lease-acquisition poll: this carries the full active set and creates - // any manager the fast path missed, so it must be delivered. Block until the listener accepts it - // (it drains continuously, so this only waits in the pathological full-buffer case). The recover() - // above handles a racing send on the channel closed during cleanup. - l.concurrencyLeasesCh <- msg } func (l *LeaseManager) acquireWorkerLeases(ctx context.Context) error { @@ -434,12 +435,9 @@ func (l *LeaseManager) acquireConcurrencyLeases(ctx context.Context) error { return nil } -// notifyNewConcurrencyStrategy acquires a lease for a single concurrency strategy on-demand and hands -// it to the lease channel, which instantiates the ConcurrencyManager. This mirrors notifyNewQueue. -// -// It is a no-op when we already hold the lease for the strategy: because sendConcurrencyLeases is a -// blocking send, holding the lease implies a manager was (or is being) created for it, so there is -// nothing to do. +// notifyNewConcurrencyStrategy acquires a lease for a single strategy on-demand and hands it to the +// lease channel, which spins up its ConcurrencyManager. Mirrors notifyNewQueue. No-op if we already +// hold the lease. func (l *LeaseManager) notifyNewConcurrencyStrategy(ctx context.Context, strategyId int64) error { l.processMu.RLock() defer l.processMu.RUnlock() diff --git a/pkg/scheduling/v1/tenant_manager.go b/pkg/scheduling/v1/tenant_manager.go index c4e7edecaa..7522970957 100644 --- a/pkg/scheduling/v1/tenant_manager.go +++ b/pkg/scheduling/v1/tenant_manager.go @@ -266,10 +266,8 @@ func (t *tenantManager) setConcurrencyStrategies(ctx context.Context, strategies c := newConcurrencyManager(t.cf, t.tenantId, strategy, t.concurrencyResultsCh, t.concurrencyAdvisoryLock, t.concurrencyParentAdvisoryLock) newArr = append(newArr, c) - // Notify each newly instantiated manager so it evaluates its already-queued slots immediately, - // instead of waiting up to maxPollingInterval for its first ticker tick. This matters on - // startup/rebalance, where managers are (re)created here for strategies that may already have - // queued work. notify() is a non-blocking send, so it is safe to call under the lock. + // notify the new manager so it picks up already-queued work now instead of waiting for its first + // tick (matters on startup/rebalance). notify() is a non-blocking send, so it's fine under the lock. c.notify(ctx) } @@ -289,11 +287,7 @@ func (t *tenantManager) addConcurrencyStrategy(ctx context.Context, strategy *sq c := newConcurrencyManager(t.cf, t.tenantId, strategy, t.concurrencyResultsCh, t.concurrencyAdvisoryLock, t.concurrencyParentAdvisoryLock) t.concurrencyStrategies = append(t.concurrencyStrategies, c) - // Notify the newly instantiated manager so it evaluates its already-queued slots on the next loop - // iteration, instead of waiting up to maxPollingInterval for its first ticker tick. This is what - // lets a cold strategy schedule its waiting task promptly once its lease is acquired. The bulk - // setConcurrencyStrategies path deliberately does NOT do this, so startup/rebalance keeps the - // randomticker's load-spreading. + // notify the new manager so it schedules its waiting task now instead of waiting for its first tick. c.notify(ctx) } @@ -373,15 +367,10 @@ func (t *tenantManager) notifyConcurrency(ctx context.Context, strategyIds []int t.concurrencyMu.RUnlock() - // Any strategy ids still in the map have no manager running on this scheduler. This is the "cold" - // case: a task arrived for a strategy that was deactivated while idle (or has never been leased here). - // Acquire the lease on-demand so a ConcurrencyManager is created now, instead of waiting up to the - // next lease-acquisition poll (every 5s) for it to be discovered. notifyNewConcurrencyStrategy does a - // DB lease acquisition, so run it off the hot notify path; it is a no-op if the lease is already held - // or owned by another scheduler. + // anything still in the map has no manager here (cold strategy): acquire its lease on-demand so a + // manager is created now instead of waiting for the ~5s poll. run off the hot path since it hits the DB. if len(strategyIdsMap) > 0 { - // detach from the request context so cancellation when the message handler returns doesn't - // abort the lease-acquisition DB call, while still propagating telemetry/values + // don't inherit cancellation from the message handler returning, but keep telemetry/values leaseCtx := context.WithoutCancel(ctx) for id := range strategyIdsMap { From b223f5503ca6677e8c9057425bbeee01ec47ce16 Mon Sep 17 00:00:00 2001 From: Mohammed Nafees Date: Fri, 26 Jun 2026 19:43:12 +0200 Subject: [PATCH 7/7] fix PR comments --- .../services/controllers/task/controller.go | 46 +++++++++---------- pkg/scheduling/v1/concurrency.go | 3 ++ pkg/scheduling/v1/tenant_manager.go | 27 ++++------- 3 files changed, 35 insertions(+), 41 deletions(-) diff --git a/internal/services/controllers/task/controller.go b/internal/services/controllers/task/controller.go index bd5dc0c17d..4650a495f7 100644 --- a/internal/services/controllers/task/controller.go +++ b/internal/services/controllers/task/controller.go @@ -42,30 +42,30 @@ type TasksController interface { } type TasksControllerImpl struct { - mq msgqueue.MessageQueue - pubBuffer *msgqueue.MQPubBuffer - l *zerolog.Logger - queueLogger *zerolog.Logger - pgxStatsLogger *zerolog.Logger - repov1 v1.Repository - dv datautils.DataDecoderValidator - s gocron.Scheduler - a *hatcheterrors.Wrapped - p *partition.Partition - celParser *cel.CELParser - opsPoolPollInterval time.Duration - opsPoolJitter time.Duration - timeoutTaskOperations *operation.TenantOperationPool - reassignTaskOperations *operation.TenantOperationPool - retryTaskOperations *operation.TenantOperationPool - emitSleepOperations *operation.TenantOperationPool - evictExpiredIdempotencyKeysOperations *operation.TenantOperationPool + mq msgqueue.MessageQueue + pubBuffer *msgqueue.MQPubBuffer + l *zerolog.Logger + queueLogger *zerolog.Logger + pgxStatsLogger *zerolog.Logger + repov1 v1.Repository + dv datautils.DataDecoderValidator + s gocron.Scheduler + a *hatcheterrors.Wrapped + p *partition.Partition + celParser *cel.CELParser + opsPoolPollInterval time.Duration + opsPoolJitter time.Duration + timeoutTaskOperations *operation.TenantOperationPool + reassignTaskOperations *operation.TenantOperationPool + retryTaskOperations *operation.TenantOperationPool + emitSleepOperations *operation.TenantOperationPool + evictExpiredIdempotencyKeysOperations *operation.TenantOperationPool // deactivateStaleStepConcurrencyOperations *operation.TenantOperationPool - replayEnabled bool - analyzeCronInterval time.Duration - signaler *signal.OLAPSignaler - tw *trigger.TriggerWriter - promGate *prometheus.Gate + replayEnabled bool + analyzeCronInterval time.Duration + signaler *signal.OLAPSignaler + tw *trigger.TriggerWriter + promGate *prometheus.Gate } type TasksControllerOpt func(*TasksControllerOpts) diff --git a/pkg/scheduling/v1/concurrency.go b/pkg/scheduling/v1/concurrency.go index 907d2bf152..66cb44dff3 100644 --- a/pkg/scheduling/v1/concurrency.go +++ b/pkg/scheduling/v1/concurrency.go @@ -97,6 +97,9 @@ func newConcurrencyManager(conf *sharedConfig, tenantId uuid.UUID, strategy *sql go c.loopConcurrency(ctx) go c.loopCheckActive(ctx) + // run once on startup instead of waiting for the first tick + c.notify(context.Background()) + return c } diff --git a/pkg/scheduling/v1/tenant_manager.go b/pkg/scheduling/v1/tenant_manager.go index 7522970957..cbf9a617d9 100644 --- a/pkg/scheduling/v1/tenant_manager.go +++ b/pkg/scheduling/v1/tenant_manager.go @@ -177,10 +177,10 @@ func (t *tenantManager) listenForConcurrencyLeases(ctx context.Context) { case msg := <-t.concurrencyCh: if msg.isIncremental { for _, strategy := range msg.items { - t.addConcurrencyStrategy(ctx, strategy) + t.addConcurrencyStrategy(strategy) } } else { - t.setConcurrencyStrategies(ctx, msg.items) + t.setConcurrencyStrategies(msg.items) } } } @@ -238,7 +238,7 @@ func (t *tenantManager) addQueuer(queueName string) { t.queue(context.Background(), []string{queueName}) } -func (t *tenantManager) setConcurrencyStrategies(ctx context.Context, strategies []*sqlcv1.V1StepConcurrency) { +func (t *tenantManager) setConcurrencyStrategies(strategies []*sqlcv1.V1StepConcurrency) { t.concurrencyMu.Lock() defer t.concurrencyMu.Unlock() @@ -265,16 +265,12 @@ func (t *tenantManager) setConcurrencyStrategies(ctx context.Context, strategies for _, strategy := range strategiesSet { c := newConcurrencyManager(t.cf, t.tenantId, strategy, t.concurrencyResultsCh, t.concurrencyAdvisoryLock, t.concurrencyParentAdvisoryLock) newArr = append(newArr, c) - - // notify the new manager so it picks up already-queued work now instead of waiting for its first - // tick (matters on startup/rebalance). notify() is a non-blocking send, so it's fine under the lock. - c.notify(ctx) } t.concurrencyStrategies = newArr } -func (t *tenantManager) addConcurrencyStrategy(ctx context.Context, strategy *sqlcv1.V1StepConcurrency) { +func (t *tenantManager) addConcurrencyStrategy(strategy *sqlcv1.V1StepConcurrency) { t.concurrencyMu.Lock() defer t.concurrencyMu.Unlock() @@ -286,9 +282,6 @@ func (t *tenantManager) addConcurrencyStrategy(ctx context.Context, strategy *sq c := newConcurrencyManager(t.cf, t.tenantId, strategy, t.concurrencyResultsCh, t.concurrencyAdvisoryLock, t.concurrencyParentAdvisoryLock) t.concurrencyStrategies = append(t.concurrencyStrategies, c) - - // notify the new manager so it schedules its waiting task now instead of waiting for its first tick. - c.notify(ctx) } func (t *tenantManager) replenish(ctx context.Context) { @@ -301,9 +294,11 @@ func (t *tenantManager) replenish(ctx context.Context) { func (t *tenantManager) notifyConcurrency(ctx context.Context, strategyIds []int64) { strategyIdsMap := make(map[int64]struct{}, len(strategyIds)) + unmatchedIds := make(map[int64]struct{}, len(strategyIds)) for _, id := range strategyIds { strategyIdsMap[id] = struct{}{} + unmatchedIds[id] = struct{}{} } t.concurrencyMu.RLock() @@ -313,8 +308,7 @@ func (t *tenantManager) notifyConcurrency(ctx context.Context, strategyIds []int continue } - // mark this strategy as already managed so we don't try to acquire a lease for it below - delete(strategyIdsMap, c.strategy.ID) + delete(unmatchedIds, c.strategy.ID) c.notify(ctx) @@ -367,13 +361,10 @@ func (t *tenantManager) notifyConcurrency(ctx context.Context, strategyIds []int t.concurrencyMu.RUnlock() - // anything still in the map has no manager here (cold strategy): acquire its lease on-demand so a - // manager is created now instead of waiting for the ~5s poll. run off the hot path since it hits the DB. - if len(strategyIdsMap) > 0 { - // don't inherit cancellation from the message handler returning, but keep telemetry/values + if len(unmatchedIds) > 0 { leaseCtx := context.WithoutCancel(ctx) - for id := range strategyIdsMap { + for id := range unmatchedIds { go t.notifyNewConcurrencyStrategy(leaseCtx, id) } }