diff --git a/cmd/hatchet-cli/cli/quickstart.go b/cmd/hatchet-cli/cli/quickstart.go index 39637f3650..d4da3aa1e2 100644 --- a/cmd/hatchet-cli/cli/quickstart.go +++ b/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" 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/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..b28e0edf71 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,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 diff --git a/pkg/scheduling/v1/tenant_manager.go b/pkg/scheduling/v1/tenant_manager.go index 5797412195..cbf9a617d9 100644 --- a/pkg/scheduling/v1/tenant_manager.go +++ b/pkg/scheduling/v1/tenant_manager.go @@ -263,7 +263,8 @@ 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) } t.concurrencyStrategies = newArr @@ -279,7 +280,8 @@ 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) } func (t *tenantManager) replenish(ctx context.Context) { @@ -292,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() @@ -304,6 +308,8 @@ func (t *tenantManager) notifyConcurrency(ctx context.Context, strategyIds []int continue } + delete(unmatchedIds, c.strategy.ID) + c.notify(ctx) childStrategyIds := make([]int64, 0) @@ -354,6 +360,14 @@ func (t *tenantManager) notifyConcurrency(ctx context.Context, strategyIds []int } t.concurrencyMu.RUnlock() + + if len(unmatchedIds) > 0 { + leaseCtx := context.WithoutCancel(ctx) + + for id := range unmatchedIds { + go t.notifyNewConcurrencyStrategy(leaseCtx, id) + } + } } func (t *tenantManager) notifyNewWorker(ctx context.Context, workerId uuid.UUID) {