fix: on-demand lease manager for cold start concurrency queue

cmd/hatchet-cli/cli/quickstart.go

@@ -7,9 +7,10 @@ import (
 
 	"github.com/charmbracelet/huh"
 	"github.com/charmbracelet/lipgloss"
-	quickstarts "github.com/hatchet-dev/hatchet-quickstarts"
 	"github.com/spf13/cobra"
 
+	quickstarts "github.com/hatchet-dev/hatchet-quickstarts"
+
 	"github.com/hatchet-dev/hatchet/cmd/hatchet-cli/cli/internal/config/cli"
 	"github.com/hatchet-dev/hatchet/cmd/hatchet-cli/cli/internal/styles"
 	"github.com/hatchet-dev/hatchet/cmd/hatchet-cli/cli/internal/templater"

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)

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
+				}
+			}
+		}
+	})
+}

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,6 +100,11 @@ func (l *LeaseManager) sendConcurrencyLeases(concurrencyLeases []*sqlcv1.V1StepC
 		}
 	}()
 
+	// 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,
@@ -107,6 +114,16 @@ func (l *LeaseManager) sendConcurrencyLeases(concurrencyLeases []*sqlcv1.V1StepC
 	}
 }
 
+func (l *LeaseManager) drainConcurrencyLeasesCh() {
+	for {
+		select {
+		case <-l.concurrencyLeasesCh:
+		default:
+			return
+		}
+	}
+}
+
 func (l *LeaseManager) acquireWorkerLeases(ctx context.Context) error {
 	l.processMu.RLock()
 	defer l.processMu.RUnlock()
@@ -418,6 +435,9 @@ func (l *LeaseManager) acquireConcurrencyLeases(ctx context.Context) error {
 	return nil
 }
 
+// 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()
@@ -455,12 +475,12 @@ func (l *LeaseManager) notifyNewConcurrencyStrategy(ctx context.Context, strateg
 	}
 
 	if len(lease) == 0 || lease[0].ResourceId == "" {
+		// the lease is owned by another scheduler; it will manage this strategy
 		return nil
 	}
 
 	l.concurrencyLeases = append(l.concurrencyLeases, lease...)
 
-	// send the new concurrency strategy to the channel
 	l.sendConcurrencyLeases([]*sqlcv1.V1StepConcurrency{strategy}, true)
 
 	return nil

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(strategy)
+					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,13 +263,18 @@ 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 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(strategy *sqlcv1.V1StepConcurrency) {
+func (t *tenantManager) addConcurrencyStrategy(ctx context.Context, strategy *sqlcv1.V1StepConcurrency) {
 	t.concurrencyMu.Lock()
 	defer t.concurrencyMu.Unlock()
 
@@ -279,7 +284,11 @@ 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)
+
+	// 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) {
@@ -304,6 +313,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 +366,17 @@ 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
+		leaseCtx := context.WithoutCancel(ctx)
+
+		for id := range strategyIdsMap {
+			go t.notifyNewConcurrencyStrategy(leaseCtx, id)
+		}
+	}
 }
 
 func (t *tenantManager) notifyNewWorker(ctx context.Context, workerId uuid.UUID) {