feat: scheduling extensions

pkg/repository/prisma/scheduler_lease.go

@@ -146,8 +146,9 @@ func (d *leaseRepository) ListActiveWorkers(ctx context.Context, tenantId pgtype
 	for _, worker := range activeWorkers {
 		wId := sqlchelpers.UUIDToStr(worker.ID)
 		res = append(res, &repository.ListActiveWorkersResult{
-			ID:     worker.ID,
-			Labels: workerIdsToLabels[wId],
+			ID:      worker.ID,
+			MaxRuns: int(worker.MaxRuns),
+			Labels:  workerIdsToLabels[wId],
 		})
 	}
 

pkg/repository/scheduler.go

@@ -16,8 +16,9 @@ type SchedulerRepository interface {
 }
 
 type ListActiveWorkersResult struct {
-	ID     pgtype.UUID
-	Labels []*dbsqlc.ListManyWorkerLabelsRow
+	ID      pgtype.UUID
+	MaxRuns int
+	Labels  []*dbsqlc.ListManyWorkerLabelsRow
 }
 
 type LeaseRepository interface {

pkg/scheduling/v2/extension.go

@@ -0,0 +1,86 @@
+package v2
+
+import (
+	"sync"
+
+	"golang.org/x/sync/errgroup"
+
+	"github.com/hatchet-dev/hatchet/pkg/repository/prisma/dbsqlc"
+)
+
+type PostScheduleInput struct {
+	Workers map[string]*WorkerCp
+
+	Slots []*SlotCp
+
+	Unassigned []*dbsqlc.QueueItem
+
+	ActionsToSlots map[string][]*SlotCp
+}
+
+type WorkerCp struct {
+	WorkerId string
+	MaxRuns  int
+	Labels   []*dbsqlc.ListManyWorkerLabelsRow
+}
+
+type SlotCp struct {
+	WorkerId string
+	Used     bool
+}
+
+type SchedulerExtension interface {
+	SetTenants(tenants []*dbsqlc.Tenant)
+	PostSchedule(tenantId string, input *PostScheduleInput)
+	Cleanup() error
+}
+
+type Extensions struct {
+	mu   sync.RWMutex
+	exts []SchedulerExtension
+}
+
+func (e *Extensions) Add(ext SchedulerExtension) {
+	e.mu.Lock()
+	defer e.mu.Unlock()
+
+	if e.exts == nil {
+		e.exts = make([]SchedulerExtension, 0)
+	}
+
+	e.exts = append(e.exts, ext)
+}
+
+func (e *Extensions) PostSchedule(tenantId string, input *PostScheduleInput) {
+	e.mu.RLock()
+	defer e.mu.RUnlock()
+
+	for _, ext := range e.exts {
+		f := ext.PostSchedule
+		go f(tenantId, input)
+	}
+}
+
+func (e *Extensions) Cleanup() error {
+	e.mu.RLock()
+	defer e.mu.RUnlock()
+
+	eg := errgroup.Group{}
+
+	for _, ext := range e.exts {
+		f := ext.Cleanup
+		eg.Go(f)
+	}
+
+	return eg.Wait()
+}
+
+func (e *Extensions) SetTenants(tenants []*dbsqlc.Tenant) {
+	e.mu.RLock()
+	defer e.mu.RUnlock()
+
+	for _, ext := range e.exts {
+		f := ext.SetTenants
+		go f(tenants)
+	}
+}

pkg/scheduling/v2/pool.go

@@ -21,6 +21,8 @@ type sharedConfig struct {
 
 // SchedulingPool is responsible for managing a pool of tenantManagers.
 type SchedulingPool struct {
+	Extensions *Extensions
+
 	tenants sync.Map
 	setMu   mutex
 
@@ -33,6 +35,7 @@ func NewSchedulingPool(repo repository.SchedulerRepository, l *zerolog.Logger, s
 	resultsCh := make(chan *QueueResults, 1000)
 
 	s := &SchedulingPool{
+		Extensions: &Extensions{},
 		cf: &sharedConfig{
 			repo:             repo,
 			l:                l,
@@ -62,6 +65,12 @@ func (p *SchedulingPool) cleanup() {
 	})
 
 	p.cleanupTenants(toCleanup)
+
+	err := p.Extensions.Cleanup()
+
+	if err != nil {
+		p.cf.l.Error().Err(err).Msg("failed to cleanup extensions")
+	}
 }
 
 func (p *SchedulingPool) SetTenants(tenants []*dbsqlc.Tenant) {
@@ -103,6 +112,8 @@ func (p *SchedulingPool) SetTenants(tenants []*dbsqlc.Tenant) {
 		// any cleaned up tenants in the map
 		p.cleanupTenants(toCleanup)
 	}()
+
+	go p.Extensions.SetTenants(tenants)
 }
 
 func (p *SchedulingPool) cleanupTenants(toCleanup []*tenantManager) {
@@ -148,7 +159,7 @@ func (p *SchedulingPool) getTenantManager(tenantId string, storeIfNotFound bool)
 
 	if !ok {
 		if storeIfNotFound {
-			tm = newTenantManager(p.cf, tenantId, p.resultsCh)
+			tm = newTenantManager(p.cf, tenantId, p.resultsCh, p.Extensions)
 			p.tenants.Store(tenantId, tm)
 		} else {
 			return nil

pkg/scheduling/v2/scheduler.go

@@ -39,10 +39,11 @@ type Scheduler struct {
 	unackedSlots map[int]*slot
 	unackedMu    mutex
 
-	rl *rateLimiter
+	rl   *rateLimiter
+	exts *Extensions
 }
 
-func newScheduler(cf *sharedConfig, tenantId pgtype.UUID, rl *rateLimiter) *Scheduler {
+func newScheduler(cf *sharedConfig, tenantId pgtype.UUID, rl *rateLimiter, exts *Extensions) *Scheduler {
 	l := cf.l.With().Str("tenant_id", sqlchelpers.UUIDToStr(tenantId)).Logger()
 
 	return &Scheduler{
@@ -57,6 +58,7 @@ func newScheduler(cf *sharedConfig, tenantId pgtype.UUID, rl *rateLimiter) *Sche
 		workersMu:       newMu(cf.l),
 		assignedCountMu: newMu(cf.l),
 		unackedMu:       newMu(cf.l),
+		exts:            exts,
 	}
 }
 
@@ -665,6 +667,9 @@ func (s *Scheduler) tryAssign(
 		wg := sync.WaitGroup{}
 		startTotal := time.Now()
 
+		extensionResults := make([]*assignResults, 0)
+		extensionResultsMu := sync.Mutex{}
+
 		// process each action id in parallel
 		for actionId, qis := range actionIdToQueueItems {
 			wg.Add(1)
@@ -733,12 +738,18 @@ func (s *Scheduler) tryAssign(
 						s.l.Warn().Dur("duration", sinceStart).Msgf("processing batch of %d queue items took longer than 100ms", len(batchQis))
 					}
 
-					resultsCh <- &assignResults{
+					r := &assignResults{
 						assigned:    batchAssigned,
 						rateLimited: batchRateLimited,
 						unassigned:  batchUnassigned,
 					}
 
+					extensionResultsMu.Lock()
+					extensionResults = append(extensionResults, r)
+					extensionResultsMu.Unlock()
+
+					resultsCh <- r
+
 					return nil
 				})
 
@@ -752,6 +763,10 @@ func (s *Scheduler) tryAssign(
 		span.End()
 		close(resultsCh)
 
+		extInput := s.getExtensionInput(extensionResults)
+
+		s.exts.PostSchedule(sqlchelpers.UUIDToStr(s.tenantId), extInput)
+
 		if sinceStart := time.Since(startTotal); sinceStart > 100*time.Millisecond {
 			s.l.Warn().Dur("duration", sinceStart).Msgf("assigning queue items took longer than 100ms")
 		}
@@ -760,6 +775,80 @@ func (s *Scheduler) tryAssign(
 	return resultsCh
 }
 
+func (s *Scheduler) getExtensionInput(results []*assignResults) *PostScheduleInput {
+	unassigned := make([]*dbsqlc.QueueItem, 0)
+
+	for _, res := range results {
+		unassigned = append(unassigned, res.unassigned...)
+	}
+
+	workers := s.getWorkers()
+
+	res := &PostScheduleInput{
+		Workers:    make(map[string]*WorkerCp),
+		Slots:      make([]*SlotCp, 0),
+		Unassigned: unassigned,
+	}
+
+	for workerId, worker := range workers {
+		res.Workers[workerId] = &WorkerCp{
+			WorkerId: workerId,
+			MaxRuns:  worker.MaxRuns,
+			Labels:   worker.Labels,
+		}
+	}
+
+	// NOTE: these locks are important because we must acquire locks in the same order as the replenish and tryAssignBatch
+	// functions. we always acquire actionsMu first and then the specific action's lock.
+	actionKeys := make([]string, 0, len(s.actions))
+
+	s.actionsMu.RLock()
+
+	for actionId := range s.actions {
+		actionKeys = append(actionKeys, actionId)
+	}
+
+	s.actionsMu.RUnlock()
+
+	uniqueSlots := make(map[*slot]*SlotCp)
+	actionsToSlots := make(map[string][]*SlotCp)
+
+	for _, actionId := range actionKeys {
+		s.actionsMu.RLock()
+		action, ok := s.actions[actionId]
+		s.actionsMu.RUnlock()
+
+		if !ok || action == nil {
+			continue
+		}
+
+		action.mu.RLock()
+		actionsToSlots[action.actionId] = make([]*SlotCp, 0, len(action.slots))
+
+		for _, slot := range action.slots {
+			if _, ok := uniqueSlots[slot]; ok {
+				continue
+			}
+
+			uniqueSlots[slot] = &SlotCp{
+				WorkerId: slot.getWorkerId(),
+				Used:     slot.used,
+			}
+
+			actionsToSlots[action.actionId] = append(actionsToSlots[action.actionId], uniqueSlots[slot])
+		}
+		action.mu.RUnlock()
+	}
+
+	for _, slot := range uniqueSlots {
+		res.Slots = append(res.Slots, slot)
+	}
+
+	res.ActionsToSlots = actionsToSlots
+
+	return res
+}
+
 func isTimedOut(qi *dbsqlc.QueueItem) bool {
 	// if the current time is after the scheduleTimeoutAt, then mark this as timed out
 	now := time.Now().UTC().UTC()

pkg/scheduling/v2/tenant_manager.go

@@ -34,11 +34,11 @@ type tenantManager struct {
 	cleanup func()
 }
 
-func newTenantManager(cf *sharedConfig, tenantId string, resultsCh chan *QueueResults) *tenantManager {
+func newTenantManager(cf *sharedConfig, tenantId string, resultsCh chan *QueueResults, exts *Extensions) *tenantManager {
 	tenantIdUUID := sqlchelpers.UUIDFromStr(tenantId)
 
 	rl := newRateLimiter(cf, tenantIdUUID)
-	s := newScheduler(cf, tenantIdUUID, rl)
+	s := newScheduler(cf, tenantIdUUID, rl, exts)
 	leaseManager, workersCh, queuesCh := newLeaseManager(cf, tenantIdUUID)
 
 	t := &tenantManager{