diff --git a/AGENTS.md b/AGENTS.md index 62cbb46c91..bd912e904e 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -4,6 +4,10 @@ - Comments should explain non-obvious intent, invariants, or constraints in the current code. Do not mention the old implementation/state (for example, "preserve the behavior of the original query"); state the current rule directly. +## Tests + +- Do not include customer names, tenant names, or production namespace names in test names, fixtures, comments, or logs. Describe the scenario by structural properties instead (for example, `dense_high_action_fanout`, not a customer or shard name). + ## Docs MDX - In MDX JSX component bodies, such as ``, avoid Markdown link syntax (`[text](href)`). Prettier can wrap the label across lines and break MDX parsing. Use an explicit JSX link instead: diff --git a/pkg/scheduling/v1/action.go b/pkg/scheduling/v1/action.go index 096bf9950a..633d1557da 100644 --- a/pkg/scheduling/v1/action.go +++ b/pkg/scheduling/v1/action.go @@ -1,83 +1,28 @@ package v1 import ( - "slices" - "sync" + "time" "github.com/google/uuid" ) type action struct { - mu sync.RWMutex - actionId string - lastReplenishedSlotCount int lastReplenishedWorkerCount int - // note that slots can be used across multiple actions, hence the pointer - slots []*slot - - // slotsByTypeAndWorkerId indexes slots by slotType -> workerId -> slots. - // - // NOTE: this index contains pointers to the same slot objects as slots. - // It is built/replaced during replenish under a.mu and read during assignment - // under a.mu (RLock or Lock). - slotsByTypeAndWorkerId map[string]map[uuid.UUID][]*slot + // workerIds is the thin action index into Scheduler.pools. + workerIds []uuid.UUID } -func (a *action) activeCount() int { - a.mu.RLock() - defer a.mu.RUnlock() - +func (a *action) activeCountFromPools(poolsByWorker map[uuid.UUID]map[string]*slotPool, now time.Time) int { count := 0 - - for _, slot := range a.slots { - if slot.active() { - count++ + for _, workerId := range a.workerIds { + for _, pool := range poolsByWorker[workerId] { + if pool.staleAt(now) { + return 0 + } + count += pool.unusedCount() } } - return count } - -// orderedLock acquires the locks in a stable order to prevent deadlocks -func orderedLock(actionsMap map[string]*action) { - actions := sortActions(actionsMap) - - for _, action := range actions { - action.mu.Lock() - } -} - -// orderedUnlock releases the locks in a stable order to prevent deadlocks. it returns -// a function that should be deferred to unlock the locks. -func orderedUnlock(actionsMap map[string]*action) func() { - actions := sortActions(actionsMap) - - return func() { - for _, action := range actions { - action.mu.Unlock() - } - } -} - -func sortActions(actionsMap map[string]*action) []*action { - actions := make([]*action, 0, len(actionsMap)) - - for _, action := range actionsMap { - actions = append(actions, action) - } - - slices.SortStableFunc(actions, func(i, j *action) int { - switch { - case i.actionId < j.actionId: - return -1 - case i.actionId > j.actionId: - return 1 - default: - return 0 - } - }) - - return actions -} diff --git a/pkg/scheduling/v1/scheduler.go b/pkg/scheduling/v1/scheduler.go index e24e6cbb71..b5e8cfaba8 100644 --- a/pkg/scheduling/v1/scheduler.go +++ b/pkg/scheduling/v1/scheduler.go @@ -1,11 +1,10 @@ package v1 import ( + "bytes" "context" "fmt" - "math/rand" "slices" - "strings" "sync" "time" @@ -29,9 +28,11 @@ type Scheduler struct { l *zerolog.Logger - actions map[string]*action - actionsMu rwMutex - replenishMu mutex + actions map[string]*action + pools map[poolKey]*slotPool + poolsByWorker map[uuid.UUID]map[string]*slotPool + actionsMu rwMutex + replenishMu mutex workersMu mutex workers map[uuid.UUID]*worker @@ -63,6 +64,8 @@ func newScheduler(cf *sharedConfig, tenantId uuid.UUID, rl *rateLimiter, exts *E tenantId: tenantId, l: &l, actions: make(map[string]*action), + pools: make(map[poolKey]*slotPool), + poolsByWorker: make(map[uuid.UUID]map[string]*slotPool), unackedSlots: make(map[int]*assignedSlots), warmedSlotTypes: make(map[uuid.UUID]map[string]struct{}), rl: rl, @@ -137,6 +140,12 @@ func (s *Scheduler) copyWorkers() map[uuid.UUID]*worker { return copied } +func (s *Scheduler) workerByID(workerId uuid.UUID) *worker { + s.workersMu.Lock() + defer s.workersMu.Unlock() + return s.workers[workerId] +} + // replenish loads new slots from the database. func (s *Scheduler) replenish(ctx context.Context, mustReplenish bool) error { if ok := s.replenishMu.TryLock(); !ok { @@ -211,7 +220,6 @@ func (s *Scheduler) replenish(ctx context.Context, mustReplenish bool) error { _, computeActionsSpan := telemetry.NewSpan(ctx, "replenish-compute-actions-to-replenish") actionsToWorkerIds := make(map[string][]uuid.UUID) - workerIdsToActions := make(map[uuid.UUID][]string) for _, workerActionTuple := range workersToActiveActions { if !workerActionTuple.ActionId.Valid { @@ -222,7 +230,6 @@ func (s *Scheduler) replenish(ctx context.Context, mustReplenish bool) error { workerId := workerActionTuple.WorkerId actionsToWorkerIds[actionId] = append(actionsToWorkerIds[actionId], workerId) - workerIdsToActions[workerId] = append(workerIdsToActions[workerId], actionId) } // FUNCTION 1: determine which actions should be replenished. Logic is the following: @@ -230,31 +237,26 @@ func (s *Scheduler) replenish(ctx context.Context, mustReplenish bool) error { // - some slots for an action: replenish if 50% of slots have been used, or have expired // - more workers available for an action than previously: fully replenish // - otherwise, do not replenish - actionsToReplenish := make(map[string]*action) + actionsToReplenish := make(map[string]struct{}) // Isolate the activeCount scans: on large tenants this dominates replenish wall time // while actionsMu is held, which blocks tryAssignBatch. _, scanActiveSpan := telemetry.NewSpan(ctx, "replenish-scan-active-counts") actionsScanned := 0 activeSlotsTotal := 0 + scanNow := time.Now() for actionId, workers := range actionsToWorkerIds { // if the action is not in the map, it should be replenished if _, ok := s.actions[actionId]; !ok { - newAction := &action{ - actionId: actionId, - slotsByTypeAndWorkerId: make(map[string]map[uuid.UUID][]*slot), - } - - actionsToReplenish[actionId] = newAction - - s.actions[actionId] = newAction + actionsToReplenish[actionId] = struct{}{} + s.actions[actionId] = new(action) continue } if mustReplenish { - actionsToReplenish[actionId] = s.actions[actionId] + actionsToReplenish[actionId] = struct{}{} continue } @@ -263,7 +265,7 @@ func (s *Scheduler) replenish(ctx context.Context, mustReplenish bool) error { // determine if we match the conditions above var replenish bool - activeCount := storedAction.activeCount() + activeCount := storedAction.activeCountFromPools(s.poolsByWorker, scanNow) actionsScanned++ activeSlotsTotal += activeCount @@ -280,7 +282,7 @@ func (s *Scheduler) replenish(ctx context.Context, mustReplenish bool) error { } if replenish { - actionsToReplenish[actionId] = s.actions[actionId] + actionsToReplenish[actionId] = struct{}{} } } @@ -291,43 +293,6 @@ func (s *Scheduler) replenish(ctx context.Context, mustReplenish bool) error { ) scanActiveSpan.End() - // if there are any workers which have additional actions not in the actionsToReplenish map, we need - // to add them to the actionsToReplenish map. This is a transitive closure over "workers of actions - // being replenished", computed with an explicit worklist so that newly added actions also have their - // workers visited. each worker is visited at most once, since a single visit adds all of its actions; - // this keeps the closure O(workers x actions-per-worker). - _, closureSpan := telemetry.NewSpan(ctx, "replenish-transitive-closure") - actionIdQueue := make([]string, 0, len(actionsToReplenish)) - - for actionId := range actionsToReplenish { - actionIdQueue = append(actionIdQueue, actionId) - } - - visitedWorkers := make(map[uuid.UUID]struct{}) - - for i := 0; i < len(actionIdQueue); i++ { - for _, workerId := range actionsToWorkerIds[actionIdQueue[i]] { - if _, visited := visitedWorkers[workerId]; visited { - continue - } - - visitedWorkers[workerId] = struct{}{} - - for _, otherActionId := range workerIdsToActions[workerId] { - if _, ok := actionsToReplenish[otherActionId]; !ok { - actionsToReplenish[otherActionId] = s.actions[otherActionId] - actionIdQueue = append(actionIdQueue, otherActionId) - } - } - } - } - - telemetry.WithAttributes(closureSpan, - telemetry.AttributeKV{Key: "replenish.actions_to_replenish", Value: len(actionsToReplenish)}, - telemetry.AttributeKV{Key: "replenish.closure_workers_visited", Value: len(visitedWorkers)}, - ) - closureSpan.End() - telemetry.WithAttributes(computeActionsSpan, telemetry.AttributeKV{Key: "replenish.actions_to_replenish", Value: len(actionsToReplenish)}, telemetry.AttributeKV{Key: "replenish.actions_scanned", Value: actionsScanned}, @@ -338,7 +303,12 @@ func (s *Scheduler) replenish(ctx context.Context, mustReplenish bool) error { s.l.Debug().Ctx(ctx).Msgf("determining which actions to replenish took %s", time.Since(checkpoint)) checkpoint = time.Now() - // FUNCTION 2: for each action which should be replenished, load the available slots + if len(actionsToReplenish) == 0 { + telemetry.WithAttributes(span, telemetry.AttributeKV{Key: "replenish.skipped_empty", Value: true}) + return nil + } + + // FUNCTION 2: load the worker-owned pool configuration and capacity. listConfigsCtx, listConfigsSpan := telemetry.NewSpan(ctx, "replenish-list-worker-slot-configs") workerSlotConfigs, err := s.repo.ListWorkerSlotConfigs(listConfigsCtx, s.tenantId, workerIds) @@ -349,73 +319,31 @@ func (s *Scheduler) replenish(ctx context.Context, mustReplenish bool) error { return err } - workerSlotTypes := make(map[uuid.UUID]map[string]bool, len(workerSlotConfigs)) - slotTypeToWorkerIds := make(map[string]map[uuid.UUID]bool) + configuredPools := make(map[poolKey]struct{}, len(workerSlotConfigs)) + slotTypeSet := make(map[string]struct{}) + workerIdSet := make(map[uuid.UUID]struct{}) for _, config := range workerSlotConfigs { - if _, ok := workerSlotTypes[config.WorkerID]; !ok { - workerSlotTypes[config.WorkerID] = make(map[string]bool) - } - - workerSlotTypes[config.WorkerID][config.SlotType] = true - - if _, ok := slotTypeToWorkerIds[config.SlotType]; !ok { - slotTypeToWorkerIds[config.SlotType] = make(map[uuid.UUID]bool) - } - - slotTypeToWorkerIds[config.SlotType][config.WorkerID] = true - } - - // We may update slots for any action that is active on a worker with slot capacity. - // Since tryAssignBatch can hold action.mu without holding actionsMu, we must lock every - // action we might write to here (not just the subset that triggered a replenish). - actionsToLock := make(map[string]*action) - for _, workerSet := range slotTypeToWorkerIds { - for workerId := range workerSet { - for _, actionId := range workerIdsToActions[workerId] { - if a := s.actions[actionId]; a != nil { - actionsToLock[actionId] = a - } - } - } + configuredPools[poolKey{workerId: config.WorkerID, slotType: config.SlotType}] = struct{}{} + slotTypeSet[config.SlotType] = struct{}{} + workerIdSet[config.WorkerID] = struct{}{} } - _, actionLocksSpan := telemetry.NewSpan(ctx, "replenish-acquire-action-locks") - telemetry.WithAttributes(actionLocksSpan, telemetry.AttributeKV{Key: "replenish.actions_to_lock", Value: len(actionsToLock)}) - - orderedLock(actionsToLock) - unlock := orderedUnlock(actionsToLock) - defer unlock() - s.unackedMu.Lock() defer s.unackedMu.Unlock() - actionLocksSpan.End() - - availableSlotsByType := make(map[string]map[uuid.UUID]int, len(slotTypeToWorkerIds)) - - slotTypes := make([]string, 0, len(slotTypeToWorkerIds)) - workerUUIDSet := make(map[uuid.UUID]struct{}) - - for slotType, workerSet := range slotTypeToWorkerIds { + slotTypes := make([]string, 0, len(slotTypeSet)) + for slotType := range slotTypeSet { slotTypes = append(slotTypes, slotType) - - // Preserve the prior behavior of creating a map entry per slot type even if it ends up empty. - if _, ok := availableSlotsByType[slotType]; !ok { - availableSlotsByType[slotType] = make(map[uuid.UUID]int, len(workerSet)) - } - - for workerId := range workerSet { - workerUUIDSet[workerId] = struct{}{} - } } - if len(slotTypes) > 0 && len(workerUUIDSet) > 0 { - workerUUIDs := make([]uuid.UUID, 0, len(workerUUIDSet)) - for workerId := range workerUUIDSet { - workerUUIDs = append(workerUUIDs, workerId) - } + workerUUIDs := make([]uuid.UUID, 0, len(workerIdSet)) + for workerId := range workerIdSet { + workerUUIDs = append(workerUUIDs, workerId) + } + availableByPool := make(map[poolKey]int, len(configuredPools)) + if len(slotTypes) > 0 && len(workerUUIDs) > 0 { listSlotsCtx, listSlotsSpan := telemetry.NewSpan(ctx, "replenish-list-available-slots") availableSlots, err := s.repo.ListAvailableSlotsForWorkersAndTypes(listSlotsCtx, s.tenantId, sqlcv1.ListAvailableSlotsForWorkersAndTypesParams{ @@ -431,229 +359,130 @@ func (s *Scheduler) replenish(ctx context.Context, mustReplenish bool) error { } for _, row := range availableSlots { - if _, ok := availableSlotsByType[row.SlotType]; !ok { - availableSlotsByType[row.SlotType] = make(map[uuid.UUID]int) - } - - availableSlotsByType[row.SlotType][row.ID] = int(row.AvailableSlots) + availableByPool[poolKey{workerId: row.ID, slotType: row.SlotType}] = int(row.AvailableSlots) } } s.l.Debug().Ctx(ctx).Msgf("loading available slots took %s", time.Since(checkpoint)) - // FUNCTION 3: list unacked slots (so they're not counted towards the worker slot count) - workersToUnackedSlots := make(map[uuid.UUID]map[string][]*slot) - - for _, unackedSlot := range s.unackedSlots { - for _, assignedSlot := range unackedSlot.slots { - workerId := assignedSlot.getWorkerId() - + // FUNCTION 3: retain unacked slots in their worker-owned pools. + unackedByPool := make(map[poolKey][]*slot) + for _, assignment := range s.unackedSlots { + for _, assignedSlot := range assignment.slots { slotType, err := assignedSlot.getSlotType() if err != nil { return fmt.Errorf("could not get slot type for unacked slot: %w", err) } - if _, ok := workersToUnackedSlots[workerId]; !ok { - workersToUnackedSlots[workerId] = make(map[string][]*slot) - } - - workersToUnackedSlots[workerId][slotType] = append(workersToUnackedSlots[workerId][slotType], assignedSlot) + key := poolKey{workerId: assignedSlot.getWorkerId(), slotType: slotType} + unackedByPool[key] = append(unackedByPool[key], assignedSlot) + configuredPools[key] = struct{}{} } } - // FUNCTION 4: write the new slots to the scheduler and clean up expired slots + // FUNCTION 4: build each (worker, slot type) pool once, then update the + // action-to-worker index. Slot objects are no longer copied into actions. _, buildSlotsSpan := telemetry.NewSpan(ctx, "replenish-build-slots") - actionsToNewSlots := make(map[string][]*slot) - actionsToTotalSlots := make(map[string]int) - actionsToSlotsByType := make(map[string]map[string]map[uuid.UUID][]*slot) - - // metaCache interns slotMeta so slots with identical metadata share the same pointer. - // Key format: slotType + "\x00" + strings.Join(sortedUniqueActions, "\x00") - metaCache := make(map[string]*slotMeta) - - for slotType, availableSlotsByWorker := range availableSlotsByType { - for workerId, availableSlots := range availableSlotsByWorker { - actions := workerIdsToActions[workerId] - unackedSlots := workersToUnackedSlots[workerId][slotType] - - // create a slot for each available slot - slots := make([]*slot, 0) - availableCount := availableSlots - len(unackedSlots) - if availableCount < 0 { - availableCount = 0 - } - - // Canonicalize actions to increase cache hits across workers. - // Order doesn't matter for correctness anywhere in scheduling. - if len(actions) > 1 { - slices.Sort(actions) - actions = slices.Compact(actions) - } - - metaKey := slotType - if len(actions) > 0 { - metaKey = slotType + "\x00" + strings.Join(actions, "\x00") - } - - meta := metaCache[metaKey] - if meta == nil { - meta = newSlotMeta(actions, slotType) - metaCache[metaKey] = meta - } + nextPools := make(map[poolKey]*slotPool, len(configuredPools)) + nextPoolsByWorker := make(map[uuid.UUID]map[string]*slotPool) + totalSlotsBuilt := 0 + maxSlotsPerPool := 0 - for i := 0; i < availableCount; i++ { - slots = append(slots, newSlot(workers[workerId], meta)) - } + for key := range configuredPools { + w := workers[key.workerId] + if w == nil { + continue + } - // extend expiry of all unacked slots - for _, unackedSlot := range unackedSlots { - unackedSlot.extendExpiry() - } + pool := s.pools[key] + if pool == nil { + pool = &slotPool{worker: w, slotType: key.slotType} + } else { + pool.worker = w + pool.slotType = key.slotType + } - s.l.Debug().Ctx(ctx).Msgf("worker %s has %d total slots (%s), %d unacked slots", workerId, availableSlots, slotType, len(unackedSlots)) + unackedSlots := unackedByPool[key] + availableCount := availableByPool[key] - len(unackedSlots) + if availableCount < 0 { + availableCount = 0 + } - slots = append(slots, unackedSlots...) + // Align pool refreshedAt with every slot's expiresAt so staleAt and + // active() flip together. Otherwise unusedCount can still report + // capacity after individual slots have expired. + refreshedAt := time.Now() + expiresAt := refreshedAt.Add(defaultSlotExpiry) - for _, actionId := range actions { - if s.actions[actionId] == nil { - continue - } + slots := make([]*slot, 0, availableCount+len(unackedSlots)) + meta := newSlotMeta(nil, key.slotType) + for index := 0; index < availableCount; index++ { + slots = append(slots, newSlotWithExpiry(w, meta, expiresAt)) + } + for _, unackedSlot := range unackedSlots { + unackedSlot.setExpiry(expiresAt) + } + slots = append(slots, unackedSlots...) + pool.resetSlotsAt(slots, refreshedAt) - actionsToNewSlots[actionId] = append(actionsToNewSlots[actionId], slots...) - actionsToTotalSlots[actionId] += len(slots) + nextPools[key] = pool + if nextPoolsByWorker[key.workerId] == nil { + nextPoolsByWorker[key.workerId] = make(map[string]*slotPool) + } + nextPoolsByWorker[key.workerId][key.slotType] = pool - if _, ok := actionsToSlotsByType[actionId]; !ok { - actionsToSlotsByType[actionId] = make(map[string]map[uuid.UUID][]*slot) - } - if _, ok := actionsToSlotsByType[actionId][slotType]; !ok { - actionsToSlotsByType[actionId][slotType] = make(map[uuid.UUID][]*slot) - } - // Reuse the per-worker/per-type slice for each action on that worker. - actionsToSlotsByType[actionId][slotType][workerId] = slots - } + totalSlotsBuilt += len(slots) + if len(slots) > maxSlotsPerPool { + maxSlotsPerPool = len(slots) } } - // (we don't need cryptographically secure randomness) - randSource := rand.New(rand.NewSource(time.Now().UnixNano())) // nolint: gosec - - totalSlotsBuilt := 0 - maxSlotsPerAction := 0 + s.pools = nextPools + s.poolsByWorker = nextPoolsByWorker - // first pass: write all actions with new slots to the scheduler - for actionId, newSlots := range actionsToNewSlots { - storedAction := actionsToLock[actionId] - if storedAction == nil { - // Defensive: actionsToNewSlots should only contain actions for workers we locked above. - continue + actionsRemoved := 0 + for actionId, storedAction := range s.actions { + actionWorkerIds := actionsToWorkerIds[actionId] + if len(actionWorkerIds) > 1 { + slices.SortFunc(actionWorkerIds, func(left, right uuid.UUID) int { + return bytes.Compare(left[:], right[:]) + }) + actionWorkerIds = slices.Compact(actionWorkerIds) } - // randomly sort the slots - randSource.Shuffle(len(newSlots), func(i, j int) { newSlots[i], newSlots[j] = newSlots[j], newSlots[i] }) - - // we overwrite the slots for the action. we know that the action is in the map because we checked - // for it in the first pass. - storedAction.slots = newSlots - storedAction.slotsByTypeAndWorkerId = actionsToSlotsByType[actionId] - storedAction.lastReplenishedSlotCount = actionsToTotalSlots[actionId] - storedAction.lastReplenishedWorkerCount = len(actionsToWorkerIds[actionId]) + totalSlots := 0 + for _, workerId := range actionWorkerIds { + for _, pool := range nextPoolsByWorker[workerId] { + totalSlots += len(pool.slots) + } + } - totalSlotsBuilt += len(newSlots) - if len(newSlots) > maxSlotsPerAction { - maxSlotsPerAction = len(newSlots) + if totalSlots == 0 { + delete(s.actions, actionId) + actionsRemoved++ + continue } - s.l.Debug().Ctx(ctx).Msgf("before cleanup, action %s has %d slots", actionId, len(newSlots)) + storedAction.workerIds = actionWorkerIds + storedAction.lastReplenishedSlotCount = totalSlots + storedAction.lastReplenishedWorkerCount = len(actionWorkerIds) } telemetry.WithAttributes(buildSlotsSpan, - telemetry.AttributeKV{Key: "replenish.actions_with_new_slots", Value: len(actionsToNewSlots)}, + telemetry.AttributeKV{Key: "replenish.actions_with_new_slots", Value: len(s.actions)}, telemetry.AttributeKV{Key: "replenish.slots_built", Value: totalSlotsBuilt}, - telemetry.AttributeKV{Key: "replenish.max_slots_per_action", Value: maxSlotsPerAction}, + telemetry.AttributeKV{Key: "replenish.max_slots_per_pool", Value: maxSlotsPerPool}, + telemetry.AttributeKV{Key: "replenish.pool_count", Value: len(nextPools)}, telemetry.AttributeKV{Key: "replenish.unacked_slot_entries", Value: len(s.unackedSlots)}, ) buildSlotsSpan.End() telemetry.WithAttributes(span, - telemetry.AttributeKV{Key: "replenish.actions_with_new_slots", Value: len(actionsToNewSlots)}, + telemetry.AttributeKV{Key: "replenish.actions_with_new_slots", Value: len(s.actions)}, telemetry.AttributeKV{Key: "replenish.slots_built", Value: totalSlotsBuilt}, - telemetry.AttributeKV{Key: "replenish.max_slots_per_action", Value: maxSlotsPerAction}, - telemetry.AttributeKV{Key: "replenish.actions_to_lock", Value: len(actionsToLock)}, - ) - - // second pass: clean up expired slots - _, cleanupSpan := telemetry.NewSpan(ctx, "replenish-cleanup-expired-slots") - defer cleanupSpan.End() - - cleanupNow := time.Now() - actionsCleaned := 0 - slotsRemoved := 0 - actionsRemoved := 0 - - for _, storedAction := range actionsToReplenish { - hasSingleSlotExpired := false - - for i := range storedAction.slots { - if storedAction.slots[i].expiredAt(cleanupNow) { - hasSingleSlotExpired = true - break - } - } - - // NOTE: actions replenished in the first pass were just given brand-new slots, - // so in the common case nothing is expired and we keep the existing slices and - // maps untouched instead of reallocating them on every replenish cycle. - if !hasSingleSlotExpired { - continue - } - - beforeLen := len(storedAction.slots) - newSlots := make([]*slot, 0, beforeLen) - - for i := range storedAction.slots { - slotItem := storedAction.slots[i] - - if !slotItem.expiredAt(cleanupNow) { - newSlots = append(newSlots, slotItem) - } - } - - storedAction.slots = newSlots - storedAction.slotsByTypeAndWorkerId = make(map[string]map[uuid.UUID][]*slot) - actionsCleaned++ - slotsRemoved += beforeLen - len(newSlots) - - for _, slotItem := range newSlots { - slotType, err := slotItem.getSlotType() - if err != nil { - return fmt.Errorf("could not get slot type during cleanup: %w", err) - } - - workerId := slotItem.getWorkerId() - - if _, ok := storedAction.slotsByTypeAndWorkerId[slotType]; !ok { - storedAction.slotsByTypeAndWorkerId[slotType] = make(map[uuid.UUID][]*slot) - } - - storedAction.slotsByTypeAndWorkerId[slotType][workerId] = append(storedAction.slotsByTypeAndWorkerId[slotType][workerId], slotItem) - } - - s.l.Debug().Ctx(ctx).Msgf("after cleanup, action %s has %d slots", storedAction.actionId, len(newSlots)) - } - - // third pass: remove any actions which have no slots - for actionId, storedAction := range actionsToReplenish { - if len(storedAction.slots) == 0 { - s.l.Debug().Ctx(ctx).Msgf("removing action %s because it has no slots", actionId) - delete(s.actions, actionId) - actionsRemoved++ - } - } - - telemetry.WithAttributes(cleanupSpan, - telemetry.AttributeKV{Key: "replenish.actions_cleaned", Value: actionsCleaned}, - telemetry.AttributeKV{Key: "replenish.slots_removed", Value: slotsRemoved}, + telemetry.AttributeKV{Key: "replenish.max_slots_per_pool", Value: maxSlotsPerPool}, + telemetry.AttributeKV{Key: "replenish.pool_count", Value: len(nextPools)}, telemetry.AttributeKV{Key: "replenish.actions_removed", Value: actionsRemoved}, ) @@ -862,15 +691,15 @@ func (s *Scheduler) tryAssignBatch( } rateLimitSpan.End() - // lock the actions map and try to assign the batch of queue items. - // NOTE: if we change the position of this lock, make sure that we are still acquiring locks in the same - // order as the replenish() function, otherwise we may deadlock. - // Spans start before each acquire so lock wait (e.g. behind replenish's actionsMu write - // lock or action.mu holders) is visible in the try-assign-batch parent. + // Hold the read lock for the entire batch so replenish cannot replace the + // action-to-worker index or worker pools while assignment is using them. + // Replenish acquires the write lock before unackedMu; assignment preserves + // that lock order when it records selected slots. _, actionsMuSpan := telemetry.NewSpan(ctx, "try-assign-batch-acquire-actions-mu") s.actionsMu.RLock() + defer s.actionsMu.RUnlock() + action, ok := s.actions[actionId] - s.actionsMu.RUnlock() actionsMuSpan.End() if !ok || action == nil { @@ -889,13 +718,9 @@ func (s *Scheduler) tryAssignBatch( return res, newRingOffset, nil } - _, actionMuRLockSpan := telemetry.NewSpan(ctx, "try-assign-batch-acquire-action-mu-rlock") - action.mu.RLock() - actionMuRLockSpan.End() - slotCount := len(action.slots) - telemetry.WithAttributes(span, telemetry.AttributeKV{Key: "action.slot_count", Value: slotCount}) - if slotCount == 0 { - action.mu.RUnlock() + workerCount := len(action.workerIds) + if workerCount == 0 { + telemetry.WithAttributes(span, telemetry.AttributeKV{Key: "action.worker_count", Value: 0}) s.l.Debug().Ctx(ctx).Msgf("no slots for action %s", actionId) @@ -910,22 +735,18 @@ func (s *Scheduler) tryAssignBatch( return res, newRingOffset, nil } - action.mu.RUnlock() - - _, actionMuLockSpan := telemetry.NewSpan(ctx, "try-assign-batch-acquire-action-mu") - action.mu.Lock() - actionMuLockSpan.End() - defer action.mu.Unlock() - candidateSlots := action.slots - telemetry.WithAttributes(span, telemetry.AttributeKV{Key: "action.slot_count", Value: len(candidateSlots)}) + telemetry.WithAttributes(span, + telemetry.AttributeKV{Key: "action.worker_count", Value: workerCount}, + telemetry.AttributeKV{Key: "action.slot_count", Value: action.lastReplenishedSlotCount}, + ) for i := range res { if res[i].rateLimitResult != nil { continue } - denom := len(candidateSlots) + denom := workerCount if denom == 0 { res[i].noSlots = true @@ -959,11 +780,10 @@ func (s *Scheduler) tryAssignBatch( } } - singleRes, err := s.tryAssignSingleton( + singleRes, err := s.tryAssignSingletonFromPools( ctx, qi, - action, - candidateSlots, + action.workerIds, childRingOffset, labels, requests, @@ -988,30 +808,14 @@ func (s *Scheduler) tryAssignBatch( return res, newRingOffset, nil } -func findAssignableSlots( - candidateSlots []*slot, - action *action, +func (s *Scheduler) findAssignableWorkerPools( + workerIds []uuid.UUID, requests map[string]int32, rateLimitAck func(), rateLimitNack func(), ) *assignedSlots { - // NOTE: the caller must hold action.mu (RLock or Lock) while calling this - // function. We read from action.slots, which is replaced during replenish - // under action.mu. - seenWorkers := make(map[uuid.UUID]struct{}) - - for _, candidateSlot := range candidateSlots { - if !candidateSlot.active() { - continue - } - - workerId := candidateSlot.getWorkerId() - if _, seen := seenWorkers[workerId]; seen { - continue - } - seenWorkers[workerId] = struct{}{} - - selected, ok := selectSlotsForWorker(action.slotsByTypeAndWorkerId, workerId, requests) + for _, workerId := range workerIds { + selected, ok := selectSlotsFromPools(s.poolsByWorker[workerId], requests) if !ok { continue } @@ -1021,9 +825,6 @@ func findAssignableSlots( continue } - // Rate limit callbacks are stored at assignedSlots level, - // not on individual slots. They're called once when the - // entire assignment is acked/nacked. return &assignedSlots{ slots: usedSlots, rateLimitAck: rateLimitAck, @@ -1034,30 +835,7 @@ func findAssignableSlots( return nil } -// useSelectedSlots attempts to reserve each slot in order. If any slot cannot be -// reserved, it rolls back by nacking any slots already reserved in this call. -func useSelectedSlots(selected []*slot) ([]*slot, bool) { - usedSlots := make([]*slot, 0, len(selected)) - - for _, sl := range selected { - if !sl.use(nil, nil) { - for _, used := range usedSlots { - used.nack() - } - return nil, false - } - usedSlots = append(usedSlots, sl) - } - - return usedSlots, true -} - -func selectSlotsForWorker( - slotsByType map[string]map[uuid.UUID][]*slot, - workerId uuid.UUID, - requests map[string]int32, -) ([]*slot, bool) { - // Pre-size the selection slice to the total number of requested units. +func selectSlotsFromPools(poolsByType map[string]*slotPool, requests map[string]int32) ([]*slot, bool) { totalNeeded := 0 for _, units := range requests { if units > 0 { @@ -1066,37 +844,31 @@ func selectSlotsForWorker( } selected := make([]*slot, 0, totalNeeded) - for slotType, units := range requests { if units <= 0 { continue } - slotsByWorker, ok := slotsByType[slotType] - if !ok { - return nil, false - } - - workerSlots := slotsByWorker[workerId] - if len(workerSlots) == 0 { + pool := poolsByType[slotType] + if pool == nil { return nil, false } - needed := int(units) + // Do not gate on unusedCount: expiry does not decrement the counter, so it + // can drift above (or, after concurrent use/nack, below) the number of + // active slots. Selection must follow active() only. found := 0 - - for _, s := range workerSlots { - if !s.active() { + for _, sl := range pool.slots { + if !sl.active() { continue } - selected = append(selected, s) + selected = append(selected, sl) found++ - if found >= needed { + if found == int(units) { break } } - - if found < needed { + if found < int(units) { return nil, false } } @@ -1104,24 +876,68 @@ func selectSlotsForWorker( return selected, true } -// tryAssignSingleton attempts to assign a singleton step to a worker. -func (s *Scheduler) tryAssignSingleton( +func (s *Scheduler) rankWorkerIds( + qi *sqlcv1.V1QueueItem, + labels []*sqlcv1.GetDesiredLabelsRow, + workerIds []uuid.UUID, +) []uuid.UUID { + type rankedWorker struct { + id uuid.UUID + rank int + } + + ranked := make([]rankedWorker, 0, len(workerIds)) + for _, workerId := range workerIds { + rank := 0 + switch qi.Sticky { + case sqlcv1.V1StickyStrategyHARD: + if qi.DesiredWorkerID != nil && workerId != *qi.DesiredWorkerID { + continue + } + case sqlcv1.V1StickyStrategySOFT: + if qi.DesiredWorkerID != nil && workerId == *qi.DesiredWorkerID { + rank = 1 + } + default: + if len(labels) > 0 { + // Label affinity reads worker metadata from s.workers. Do not + // require a poolsByWorker entry — candidates can be listed on + // the action before pools are populated. + worker := s.workerByID(workerId) + if worker == nil { + continue + } + rank = worker.computeWeight(labels) + if rank < 0 { + continue + } + } + } + + ranked = append(ranked, rankedWorker{id: workerId, rank: rank}) + } + + slices.SortStableFunc(ranked, func(left, right rankedWorker) int { + return right.rank - left.rank + }) + + result := make([]uuid.UUID, len(ranked)) + for index := range ranked { + result[index] = ranked[index].id + } + return result +} + +func (s *Scheduler) tryAssignSingletonFromPools( ctx context.Context, qi *sqlcv1.V1QueueItem, - action *action, - candidateSlots []*slot, + workerIds []uuid.UUID, ringOffset int, labels []*sqlcv1.GetDesiredLabelsRow, requests map[string]int32, rateLimitAck func(), rateLimitNack func(), -) ( - res assignSingleResult, err error, -) { - // NOTE: the caller must hold action.mu (RLock or Lock) while calling this - // function. We read from action.slots, which is replaced during replenish - // under action.mu. - +) (res assignSingleResult, err error) { ctx, span := telemetry.NewSpan(ctx, "try-assign-singleton") // nolint: ineffassign defer span.End() @@ -1130,20 +946,22 @@ func (s *Scheduler) tryAssignSingleton( telemetry.AttributeKV{Key: "queue.name", Value: qi.Queue}, ) - ringOffset = ringOffset % len(candidateSlots) - - if (qi.Sticky != sqlcv1.V1StickyStrategyNONE) || len(labels) > 0 { - candidateSlots = getRankedSlots(qi, labels, candidateSlots) + candidates := workerIds + if qi.Sticky != sqlcv1.V1StickyStrategyNONE || len(labels) > 0 { + candidates = s.rankWorkerIds(qi, labels, workerIds) ringOffset = 0 } - - assignedSlot := findAssignableSlots(candidateSlots[ringOffset:], action, requests, rateLimitAck, rateLimitNack) - - if assignedSlot == nil { - assignedSlot = findAssignableSlots(candidateSlots[:ringOffset], action, requests, rateLimitAck, rateLimitNack) + if len(candidates) == 0 { + res.noSlots = true + return res, nil } - if assignedSlot == nil { + ringOffset %= len(candidates) + assigned := s.findAssignableWorkerPools(candidates[ringOffset:], requests, rateLimitAck, rateLimitNack) + if assigned == nil { + assigned = s.findAssignableWorkerPools(candidates[:ringOffset], requests, rateLimitAck, rateLimitNack) + } + if assigned == nil { res.noSlots = true return res, nil } @@ -1154,10 +972,10 @@ func (s *Scheduler) tryAssignSingleton( s.assignedCountMu.Unlock() s.unackedMu.Lock() - s.unackedSlots[res.ackId] = assignedSlot + s.unackedSlots[res.ackId] = assigned s.unackedMu.Unlock() - res.workerId = assignedSlot.workerId() + res.workerId = assigned.workerId() if res.workerId == uuid.Nil { s.l.Error().Ctx(ctx).Msgf("assigned slot %d has no worker id, skipping assignment", res.ackId) res.noSlots = true @@ -1165,10 +983,27 @@ func (s *Scheduler) tryAssignSingleton( } res.succeeded = true - return res, nil } +// useSelectedSlots attempts to reserve each slot in order. If any slot cannot be +// reserved, it rolls back by nacking any slots already reserved in this call. +func useSelectedSlots(selected []*slot) ([]*slot, bool) { + usedSlots := make([]*slot, 0, len(selected)) + + for _, sl := range selected { + if !sl.use(nil, nil) { + for _, used := range usedSlots { + used.nack() + } + return nil, false + } + usedSlots = append(usedSlots, sl) + } + + return usedSlots, true +} + type assignedQueueItem struct { AckId int WorkerId uuid.UUID @@ -1398,16 +1233,6 @@ func (s *Scheduler) getSnapshotInput(ctx context.Context, mustSnapshot bool) (*S } } - // 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)) - - for actionId := range s.actions { - actionKeys = append(actionKeys, actionId) - } - - uniqueSlots := make(map[*slot]bool) - utilizationByType := make(map[uuid.UUID]map[string]*SlotUtilization) for workerId := range workers { @@ -1415,51 +1240,34 @@ func (s *Scheduler) getSnapshotInput(ctx context.Context, mustSnapshot bool) (*S } _, walkSpan := telemetry.NewSpan(ctx, "get-snapshot-input-walk-slots") - for _, actionId := range actionKeys { - action, ok := s.actions[actionId] - - if !ok || action == nil { + uniqueSlotCount := 0 + for key, pool := range s.pools { + if _, active := workers[key.workerId]; !active { continue } + byType := utilizationByType[key.workerId] + if byType == nil { + byType = make(map[string]*SlotUtilization) + utilizationByType[key.workerId] = byType + } + utilization := byType[key.slotType] + if utilization == nil { + utilization = &SlotUtilization{} + byType[key.slotType] = utilization + } - action.mu.RLock() - for _, slot := range action.slots { - if _, ok := uniqueSlots[slot]; ok { - continue - } - - workerId := slot.worker.ID - - slotType, err := slot.getSlotType() - if err != nil { - slotType = "" - } - - byType, ok := utilizationByType[workerId] - if !ok { - byType = make(map[string]*SlotUtilization) - utilizationByType[workerId] = byType - } - - utilization, ok := byType[slotType] - if !ok { - utilization = &SlotUtilization{} - byType[slotType] = utilization - } - - uniqueSlots[slot] = true - + for _, slot := range pool.slots { + uniqueSlotCount++ if slot.isUsed() { utilization.UtilizedSlots++ } else { utilization.NonUtilizedSlots++ } } - action.mu.RUnlock() } telemetry.WithAttributes(walkSpan, - telemetry.AttributeKV{Key: "snapshot.action_count", Value: len(actionKeys)}, - telemetry.AttributeKV{Key: "snapshot.unique_slots", Value: len(uniqueSlots)}, + telemetry.AttributeKV{Key: "snapshot.action_count", Value: len(s.actions)}, + telemetry.AttributeKV{Key: "snapshot.unique_slots", Value: uniqueSlotCount}, ) walkSpan.End() @@ -1531,7 +1339,7 @@ func (s *Scheduler) getSnapshotInput(ctx context.Context, mustSnapshot bool) (*S telemetry.WithAttributes(span, telemetry.AttributeKV{Key: "snapshot.worker_count", Value: len(workers)}, - telemetry.AttributeKV{Key: "snapshot.action_count", Value: len(actionKeys)}, + telemetry.AttributeKV{Key: "snapshot.action_count", Value: len(s.actions)}, ) return res, true diff --git a/pkg/scheduling/v1/scheduler_shape_bench_test.go b/pkg/scheduling/v1/scheduler_shape_bench_test.go new file mode 100644 index 0000000000..f1968a2635 --- /dev/null +++ b/pkg/scheduling/v1/scheduler_shape_bench_test.go @@ -0,0 +1,434 @@ +//go:build !e2e && !load && !rampup && !integration + +package v1 + +import ( + "context" + "fmt" + "testing" + "time" + + "github.com/google/uuid" + "github.com/jackc/pgx/v5/pgtype" + "github.com/rs/zerolog" + "github.com/stretchr/testify/require" + + repo "github.com/hatchet-dev/hatchet/pkg/repository" + "github.com/hatchet-dev/hatchet/pkg/repository/sqlcv1" +) + +// inventoryTopology controls how actions attach to workers. It captures the +// action-to-worker index shapes the scheduler must handle. +type inventoryTopology string + +const ( + // topologyDense: every worker registers every action. + topologyDense inventoryTopology = "dense" + + // topologySparse: each worker registers ActionsPerWorker actions, chosen + // round-robin from the action pool. + topologySparse inventoryTopology = "sparse" + + // topologyPartitioned: workers and actions split into Partitions disjoint + // groups; within a group, every worker has every group action (dense local). + topologyPartitioned inventoryTopology = "partitioned" +) + +// inventoryShape is a synthetic (workers × slots × actions × topology) fixture +// used to baseline the current scheduler inventory before trying alternate shapes. +type inventoryShape struct { + Name string + Workers int + Actions int + SlotsPerWorker int + SlotTypes int + Topology inventoryTopology + ActionsPerWorker int // sparse only + Partitions int // partitioned only +} + +type inventoryFixture struct { + shape inventoryShape + tenantId uuid.UUID + workerIds []uuid.UUID + activeWorkers []*repo.ListActiveWorkersResult + actionIds []string + actionRows []*sqlcv1.ListActionsForWorkersRow + scheduler *Scheduler + uniqueSlots int +} + +func baselineShapes() []inventoryShape { + return []inventoryShape{ + { + Name: "small_dense", + Workers: 10, + Actions: 20, + SlotsPerWorker: 4, + Topology: topologyDense, + }, + // High action fan-out: many shared actions map to the same worker pools. + { + Name: "dense_high_action_fanout", + Workers: 50, + Actions: 1724, + SlotsPerWorker: 20, + Topology: topologyDense, + }, + // High slot capacity: fewer actions with large per-worker pools. + { + Name: "dense_high_slot_capacity", + Workers: 40, + Actions: 30, + SlotsPerWorker: 1400, + Topology: topologyDense, + }, + { + Name: "sparse_low_fanout", + Workers: 200, + Actions: 2000, + SlotsPerWorker: 20, + Topology: topologySparse, + ActionsPerWorker: 5, + }, + { + Name: "partitioned_20", + Workers: 200, + Actions: 2000, + SlotsPerWorker: 20, + Topology: topologyPartitioned, + Partitions: 20, + }, + } +} + +func shapeSlotTypes(shape inventoryShape) []string { + count := shape.SlotTypes + if count <= 0 { + count = 1 + } + + slotTypes := make([]string, count) + slotTypes[0] = repo.SlotTypeDefault + for index := 1; index < count; index++ { + slotTypes[index] = fmt.Sprintf("slot-type-%02d", index) + } + return slotTypes +} + +func buildActionRows(shape inventoryShape, workerIds []uuid.UUID, actionIds []string) []*sqlcv1.ListActionsForWorkersRow { + switch shape.Topology { + case topologyDense: + rows := make([]*sqlcv1.ListActionsForWorkersRow, 0, shape.Workers*shape.Actions) + for _, wid := range workerIds { + for _, aid := range actionIds { + rows = append(rows, &sqlcv1.ListActionsForWorkersRow{ + WorkerId: wid, + ActionId: pgtype.Text{String: aid, Valid: true}, + }) + } + } + return rows + + case topologySparse: + per := shape.ActionsPerWorker + if per <= 0 { + per = 1 + } + rows := make([]*sqlcv1.ListActionsForWorkersRow, 0, shape.Workers*per) + for wi, wid := range workerIds { + for j := 0; j < per; j++ { + aid := actionIds[(wi*per+j)%len(actionIds)] + rows = append(rows, &sqlcv1.ListActionsForWorkersRow{ + WorkerId: wid, + ActionId: pgtype.Text{String: aid, Valid: true}, + }) + } + } + return rows + + case topologyPartitioned: + parts := shape.Partitions + if parts <= 0 { + parts = 1 + } + rows := make([]*sqlcv1.ListActionsForWorkersRow, 0) + for wi, wid := range workerIds { + part := wi % parts + for ai, aid := range actionIds { + if ai%parts != part { + continue + } + rows = append(rows, &sqlcv1.ListActionsForWorkersRow{ + WorkerId: wid, + ActionId: pgtype.Text{String: aid, Valid: true}, + }) + } + } + return rows + + default: + panic(fmt.Sprintf("unknown topology %q", shape.Topology)) + } +} + +func newShapeScheduler(tenantId uuid.UUID, ar repo.AssignmentRepository) *Scheduler { + l := zerolog.Nop() + sr := &mockSchedulerRepo{assignment: ar} + cf := &sharedConfig{repo: sr, l: &l} + return newScheduler(cf, tenantId, nil, &Extensions{}) +} + +func newInventoryFixture(shape inventoryShape) *inventoryFixture { + tenantId := uuid.New() + + workerIds := make([]uuid.UUID, shape.Workers) + activeWorkers := make([]*repo.ListActiveWorkersResult, shape.Workers) + for i := range workerIds { + workerIds[i] = uuid.New() + activeWorkers[i] = testWorker(workerIds[i]) + } + + actionIds := make([]string, shape.Actions) + for i := range actionIds { + actionIds[i] = fmt.Sprintf("action-%05d", i) + } + + actionRows := buildActionRows(shape, workerIds, actionIds) + slotsPerWorker := int32(shape.SlotsPerWorker) + slotTypes := shapeSlotTypes(shape) + + ar := &mockAssignmentRepo{ + listActionsForWorkersFn: func(ctx context.Context, tenantId uuid.UUID, ids []uuid.UUID) ([]*sqlcv1.ListActionsForWorkersRow, error) { + return actionRows, nil + }, + listAvailableSlotsForWorkersFn: func(ctx context.Context, tenantId uuid.UUID, params sqlcv1.ListAvailableSlotsForWorkersParams) ([]*sqlcv1.ListAvailableSlotsForWorkersRow, error) { + rows := make([]*sqlcv1.ListAvailableSlotsForWorkersRow, 0, len(params.Workerids)) + for _, wid := range params.Workerids { + rows = append(rows, &sqlcv1.ListAvailableSlotsForWorkersRow{ + ID: wid, + AvailableSlots: slotsPerWorker, + }) + } + return rows, nil + }, + listWorkerSlotConfigsFn: func(ctx context.Context, tenantId uuid.UUID, workerIds []uuid.UUID) ([]*sqlcv1.ListWorkerSlotConfigsRow, error) { + rows := make([]*sqlcv1.ListWorkerSlotConfigsRow, 0, len(workerIds)*len(slotTypes)) + for _, workerId := range workerIds { + for _, slotType := range slotTypes { + rows = append(rows, &sqlcv1.ListWorkerSlotConfigsRow{ + WorkerID: workerId, + SlotType: slotType, + }) + } + } + return rows, nil + }, + listAvailableSlotsForWorkersAndTypesFn: func(ctx context.Context, tenantId uuid.UUID, params sqlcv1.ListAvailableSlotsForWorkersAndTypesParams) ([]*sqlcv1.ListAvailableSlotsForWorkersAndTypesRow, error) { + rows := make([]*sqlcv1.ListAvailableSlotsForWorkersAndTypesRow, 0, len(params.Workerids)*len(params.Slottypes)) + for _, workerId := range params.Workerids { + for _, slotType := range params.Slottypes { + rows = append(rows, &sqlcv1.ListAvailableSlotsForWorkersAndTypesRow{ + ID: workerId, + SlotType: slotType, + AvailableSlots: slotsPerWorker, + }) + } + } + return rows, nil + }, + } + + s := newShapeScheduler(tenantId, ar) + s.setWorkers(activeWorkers) + + return &inventoryFixture{ + shape: shape, + tenantId: tenantId, + workerIds: workerIds, + activeWorkers: activeWorkers, + actionIds: actionIds, + actionRows: actionRows, + scheduler: s, + } +} + +func (f *inventoryFixture) measureInventory() { + slots := 0 + for _, pool := range f.scheduler.pools { + slots += len(pool.slots) + } + f.uniqueSlots = slots +} + +func (f *inventoryFixture) scanActiveCounts() (total int) { + now := time.Now() + for _, a := range f.scheduler.actions { + total += a.activeCountFromPools(f.scheduler.poolsByWorker, now) + } + return total +} + +func (f *inventoryFixture) firstAssignableAction() string { + for _, aid := range f.actionIds { + if a := f.scheduler.actions[aid]; a != nil && len(a.workerIds) > 0 { + return aid + } + } + return "" +} + +// reportShapeMetrics attaches inventory columns to the go-benchmarks / benchstat +// table. Must be called after the timed loop: ResetTimer clears custom metrics +// (Go 1.24+). +func reportShapeMetrics(b *testing.B, f *inventoryFixture) { + b.Helper() + b.ReportMetric(float64(f.shape.Workers), "workers") + b.ReportMetric(float64(f.shape.Actions), "actions") + b.ReportMetric(float64(f.shape.SlotsPerWorker), "slots_per_worker") + b.ReportMetric(float64(len(shapeSlotTypes(f.shape))), "slot_types") + b.ReportMetric(float64(len(f.actionRows)), "action_rows") + b.ReportMetric(float64(f.uniqueSlots), "unique_slots") +} + +func BenchmarkScheduler_InventoryShape_SlotTypeCardinality(b *testing.B) { + for _, slotTypeCount := range []int{1, 4, 16} { + shape := inventoryShape{ + Name: fmt.Sprintf("slot_types_%02d", slotTypeCount), + Workers: 100, + Actions: 500, + SlotsPerWorker: 20, + SlotTypes: slotTypeCount, + Topology: topologyDense, + } + + b.Run(shape.Name, func(b *testing.B) { + f := newInventoryFixture(shape) + if err := f.scheduler.replenish(context.Background(), true); err != nil { + b.Fatal(err) + } + f.measureInventory() + + b.ResetTimer() + for iteration := 0; iteration < b.N; iteration++ { + if err := f.scheduler.replenish(context.Background(), true); err != nil { + b.Fatal(err) + } + } + b.StopTimer() + reportShapeMetrics(b, f) + }) + } +} + +func TestScheduler_InventoryShape_DenseUniqueSlots(t *testing.T) { + f := newInventoryFixture(inventoryShape{ + Name: "small_dense", + Workers: 10, + Actions: 20, + SlotsPerWorker: 4, + Topology: topologyDense, + }) + require.NoError(t, f.scheduler.replenish(context.Background(), true)) + f.measureInventory() + + require.Equal(t, f.shape.Workers*f.shape.SlotsPerWorker, f.uniqueSlots) +} + +func BenchmarkScheduler_InventoryShape_ReplenishMust(b *testing.B) { + for _, shape := range baselineShapes() { + shape := shape + b.Run(shape.Name, func(b *testing.B) { + f := newInventoryFixture(shape) + // Warm once so subsequent iterations measure rebuild, not first insert. + if err := f.scheduler.replenish(context.Background(), true); err != nil { + b.Fatal(err) + } + f.measureInventory() + + b.ResetTimer() + for i := 0; i < b.N; i++ { + if err := f.scheduler.replenish(context.Background(), true); err != nil { + b.Fatal(err) + } + } + b.StopTimer() + reportShapeMetrics(b, f) + }) + } +} + +func BenchmarkScheduler_InventoryShape_ActiveCountScan(b *testing.B) { + for _, shape := range baselineShapes() { + shape := shape + b.Run(shape.Name, func(b *testing.B) { + f := newInventoryFixture(shape) + if err := f.scheduler.replenish(context.Background(), true); err != nil { + b.Fatal(err) + } + f.measureInventory() + + b.ResetTimer() + for i := 0; i < b.N; i++ { + _ = f.scanActiveCounts() + } + b.StopTimer() + reportShapeMetrics(b, f) + }) + } +} + +func BenchmarkScheduler_InventoryShape_TryAssignBatch(b *testing.B) { + const batchSize = 64 + + for _, shape := range baselineShapes() { + shape := shape + b.Run(shape.Name, func(b *testing.B) { + f := newInventoryFixture(shape) + if err := f.scheduler.replenish(context.Background(), true); err != nil { + b.Fatal(err) + } + f.measureInventory() + + assignAction := f.firstAssignableAction() + if assignAction == "" { + b.Fatal("no assignable action after replenish") + } + + qis := make([]*sqlcv1.V1QueueItem, batchSize) + stepRequests := make(map[uuid.UUID]map[string]int32, batchSize) + for i := range qis { + qi := testQI(f.tenantId, assignAction, int64(i+1)) + qis[i] = qi + stepRequests[qi.StepID] = map[string]int32{repo.SlotTypeDefault: 1} + } + + b.ResetTimer() + for i := 0; i < b.N; i++ { + // Replenish between iterations so the batch always has capacity; + // exclude that cost from ns/op. + b.StopTimer() + if err := f.scheduler.replenish(context.Background(), true); err != nil { + b.Fatal(err) + } + b.StartTimer() + + if _, _, err := f.scheduler.tryAssignBatch( + context.Background(), + assignAction, + qis, + 0, + map[uuid.UUID][]*sqlcv1.GetDesiredLabelsRow{}, + stepRequests, + nil, + nil, + ); err != nil { + b.Fatal(err) + } + } + b.StopTimer() + b.ReportMetric(float64(batchSize), "batch_size") + reportShapeMetrics(b, f) + }) + } +} diff --git a/pkg/scheduling/v1/scheduler_test.go b/pkg/scheduling/v1/scheduler_test.go index f8805437b5..fa02ffce5c 100644 --- a/pkg/scheduling/v1/scheduler_test.go +++ b/pkg/scheduling/v1/scheduler_test.go @@ -146,29 +146,55 @@ func testWorker(id uuid.UUID) *repo.ListActiveWorkersResult { } } -func actionWithSlots(actionId string, slots ...*slot) (*action, error) { - a := &action{ - actionId: actionId, - slots: slots, - // populate index for tests; production code builds it in replenish. - slotsByTypeAndWorkerId: make(map[string]map[uuid.UUID][]*slot), - } - +func seedActionPools(t *testing.T, s *Scheduler, actionId string, slots ...*slot) *action { + t.Helper() + workerSet := make(map[uuid.UUID]struct{}) + workerIds := make([]uuid.UUID, 0) for _, sl := range slots { slotType, err := sl.getSlotType() - if err != nil { - return nil, fmt.Errorf("getSlotType failed: %w", err) + require.NoError(t, err) + workerId := sl.getWorkerId() + if _, ok := workerSet[workerId]; !ok { + workerSet[workerId] = struct{}{} + workerIds = append(workerIds, workerId) } - workerId := sl.getWorkerId() + key := poolKey{workerId: workerId, slotType: slotType} + pool := s.pools[key] + if pool == nil { + pool = &slotPool{worker: sl.worker, slotType: slotType} + s.pools[key] = pool + if s.poolsByWorker[workerId] == nil { + s.poolsByWorker[workerId] = make(map[string]*slotPool) + } + s.poolsByWorker[workerId][slotType] = pool + } - if _, ok := a.slotsByTypeAndWorkerId[slotType]; !ok { - a.slotsByTypeAndWorkerId[slotType] = make(map[uuid.UUID][]*slot) + seen := false + for _, existing := range pool.slots { + if existing == sl { + seen = true + break + } + } + if !seen { + pool.resetSlots(append(pool.slots, sl)) } - a.slotsByTypeAndWorkerId[slotType][workerId] = append(a.slotsByTypeAndWorkerId[slotType][workerId], sl) } - return a, nil + a := &action{workerIds: workerIds} + s.actions[actionId] = a + return a +} + +func slotsForAction(s *Scheduler, a *action) []*slot { + var slots []*slot + for _, workerId := range a.workerIds { + for _, pool := range s.poolsByWorker[workerId] { + slots = append(slots, pool.slots...) + } + } + return slots } func testQI(tenantId uuid.UUID, actionId string, taskId int64) *sqlcv1.V1QueueItem { @@ -278,13 +304,12 @@ func TestSelectSlotsForWorker_SkipsInactive(t *testing.T) { s3 := newSlot(w, newSlotMeta([]string{"A"}, repo.SlotTypeDefault)) - slotsByTypeAndWorkerId := map[string]map[uuid.UUID][]*slot{ - repo.SlotTypeDefault: {workerId: {s1, s2, s3}}, - } + pool := &slotPool{worker: w, slotType: repo.SlotTypeDefault} + pool.resetSlots([]*slot{s1, s2, s3}) + require.Equal(t, 1, pool.unusedCount(), "expired/used slots must not inflate unused") - selected, ok := selectSlotsForWorker( - slotsByTypeAndWorkerId, - workerId, + selected, ok := selectSlotsFromPools( + map[string]*slotPool{repo.SlotTypeDefault: pool}, map[string]int32{repo.SlotTypeDefault: 1}, ) require.True(t, ok) @@ -292,6 +317,39 @@ func TestSelectSlotsForWorker_SkipsInactive(t *testing.T) { require.Same(t, s3, selected[0]) } +func TestSelectSlotsFromPools_IgnoresStaleUnusedCount(t *testing.T) { + workerId := uuid.New() + w := &worker{ListActiveWorkersResult: testWorker(workerId)} + + active := newSlot(w, newSlotMeta([]string{"A"}, repo.SlotTypeDefault)) + expired := newSlot(w, newSlotMeta([]string{"A"}, repo.SlotTypeDefault)) + past := time.Now().Add(-time.Second) + expired.setExpiry(past) + + pool := &slotPool{worker: w, slotType: repo.SlotTypeDefault} + pool.slots = []*slot{expired, active} + active.pool = pool + expired.pool = pool + // Simulate counter drift: expiry never decremented unused, so it still + // reports 2 even though only one slot is schedulable. Also cover the + // opposite drift (undercount) which used to hard-fail selection. + pool.unused.Store(0) + + selected, ok := selectSlotsFromPools( + map[string]*slotPool{repo.SlotTypeDefault: pool}, + map[string]int32{repo.SlotTypeDefault: 1}, + ) + require.True(t, ok, "selection must follow active slots, not unusedCount") + require.Len(t, selected, 1) + require.Same(t, active, selected[0]) + + _, ok = selectSlotsFromPools( + map[string]*slotPool{repo.SlotTypeDefault: pool}, + map[string]int32{repo.SlotTypeDefault: 2}, + ) + require.False(t, ok, "must still fail when too few active slots exist") +} + func TestScheduler_TryAssignSingleton_RingWraparound(t *testing.T) { tenantId := uuid.New() workerId1 := uuid.New() @@ -307,12 +365,11 @@ func TestScheduler_TryAssignSingleton_RingWraparound(t *testing.T) { require.True(t, s1.use(nil, nil)) s2 := newSlot(w2, newSlotMeta([]string{"A"}, repo.SlotTypeDefault)) - a, err := actionWithSlots("A", s1, s2) - require.NoError(t, err) + a := seedActionPools(t, s, "A", s1, s2) req := map[string]int32{repo.SlotTypeDefault: 1} qi := testQI(tenantId, "A", 1) - res, err := s.tryAssignSingleton(context.Background(), qi, a, []*slot{s1, s2}, 1, nil, req, func() {}, func() {}) + res, err := s.tryAssignSingletonFromPools(context.Background(), qi, a.workerIds, 1, nil, req, func() {}, func() {}) require.NoError(t, err) require.True(t, res.succeeded) require.False(t, res.noSlots) @@ -335,12 +392,11 @@ func TestScheduler_TryAssignSingleton_NoSlots(t *testing.T) { s1 := newSlot(w, newSlotMeta([]string{"A"}, repo.SlotTypeDefault)) require.True(t, s1.use(nil, nil)) - a, err := actionWithSlots("A", s1) - require.NoError(t, err) + a := seedActionPools(t, s, "A", s1) req := map[string]int32{repo.SlotTypeDefault: 1} qi := testQI(tenantId, "A", 1) - res, err := s.tryAssignSingleton(context.Background(), qi, a, []*slot{s1}, 0, nil, req, func() {}, func() {}) + res, err := s.tryAssignSingletonFromPools(context.Background(), qi, a.workerIds, 0, nil, req, func() {}, func() {}) require.NoError(t, err) require.False(t, res.succeeded) require.True(t, res.noSlots) @@ -360,20 +416,64 @@ func TestScheduler_TryAssignSingleton_StickyHardForcesRanking(t *testing.T) { otherSlot := newSlot(wOther, newSlotMeta([]string{"A"}, repo.SlotTypeDefault)) desiredSlot := newSlot(wDesired, newSlotMeta([]string{"A"}, repo.SlotTypeDefault)) - a, err := actionWithSlots("A", otherSlot, desiredSlot) - require.NoError(t, err) + a := seedActionPools(t, s, "A", otherSlot, desiredSlot) req := map[string]int32{repo.SlotTypeDefault: 1} qi := testQI(tenantId, "A", 1) qi.Sticky = sqlcv1.V1StickyStrategyHARD qi.DesiredWorkerID = &desiredWorkerId - res, err := s.tryAssignSingleton(context.Background(), qi, a, []*slot{otherSlot, desiredSlot}, 1, nil, req, func() {}, func() {}) + res, err := s.tryAssignSingletonFromPools(context.Background(), qi, a.workerIds, 1, nil, req, func() {}, func() {}) require.NoError(t, err) require.True(t, res.succeeded) require.Equal(t, desiredWorkerId, res.workerId) } +func TestScheduler_RankWorkerIds_StickyDoesNotRequirePoolsByWorker(t *testing.T) { + tenantId := uuid.New() + desiredWorkerId := uuid.New() + + s := newTestScheduler(t, tenantId, &mockAssignmentRepo{}) + + // Candidate is listed for the action, but has no poolsByWorker entry yet. + // Sticky ranking only needs the worker id, so it must not drop the candidate. + qi := testQI(tenantId, "A", 1) + qi.Sticky = sqlcv1.V1StickyStrategyHARD + qi.DesiredWorkerID = &desiredWorkerId + + ranked := s.rankWorkerIds(qi, nil, []uuid.UUID{desiredWorkerId}) + require.Equal(t, []uuid.UUID{desiredWorkerId}, ranked) +} + +func TestScheduler_RankWorkerIds_LabelsUseWorkersMap(t *testing.T) { + tenantId := uuid.New() + workerId := uuid.New() + + s := newTestScheduler(t, tenantId, &mockAssignmentRepo{}) + s.setWorkers([]*repo.ListActiveWorkersResult{{ + ID: workerId, + Name: "w", + Labels: []*sqlcv1.ListManyWorkerLabelsRow{ + { + Key: "region", + StrValue: pgtype.Text{String: "us-east-1", Valid: true}, + }, + }, + }}) + + qi := testQI(tenantId, "A", 1) + labels := []*sqlcv1.GetDesiredLabelsRow{{ + Key: "region", + StrValue: pgtype.Text{String: "us-east-1", Valid: true}, + Comparator: sqlcv1.WorkerLabelComparatorEQUAL, + Weight: 10, + }} + + // No poolsByWorker entry — label ranking must still resolve the worker via s.workers. + ranked := s.rankWorkerIds(qi, labels, []uuid.UUID{workerId}) + require.Equal(t, []uuid.UUID{workerId}, ranked) +} + func TestScheduler_TryAssignSingleton_RateLimitAckIsWiredIntoSlotAck(t *testing.T) { tenantId := uuid.New() workerId := uuid.New() @@ -382,15 +482,14 @@ func TestScheduler_TryAssignSingleton_RateLimitAckIsWiredIntoSlotAck(t *testing. w := &worker{ListActiveWorkersResult: testWorker(workerId)} sl := newSlot(w, newSlotMeta([]string{"A"}, repo.SlotTypeDefault)) - a, err := actionWithSlots("A", sl) - require.NoError(t, err) + a := seedActionPools(t, s, "A", sl) req := map[string]int32{repo.SlotTypeDefault: 1} qi := testQI(tenantId, "A", 1) ackCount := 0 rlAck := func() { ackCount++ } - res, err := s.tryAssignSingleton(context.Background(), qi, a, []*slot{sl}, 0, nil, req, rlAck, func() {}) + res, err := s.tryAssignSingletonFromPools(context.Background(), qi, a.workerIds, 0, nil, req, rlAck, func() {}) require.NoError(t, err) require.True(t, res.succeeded) @@ -455,7 +554,7 @@ func TestScheduler_Replenish_SkipsIfCannotAcquireActionsLock(t *testing.T) { require.NoError(t, s.replenish(context.Background(), false)) } -func TestScheduler_Replenish_DoesNotLockUnackedMuBeforeActionLocks(t *testing.T) { +func TestScheduler_Replenish_CompletesWithoutPerActionLocks(t *testing.T) { tenantId := uuid.New() workerId := uuid.New() @@ -512,14 +611,7 @@ func TestScheduler_Replenish_DoesNotLockUnackedMuBeforeActionLocks(t *testing.T) s := newTestScheduler(t, tenantId, ar) s.setWorkers([]*repo.ListActiveWorkersResult{testWorker(workerId)}) - // Pre-create an action so replenish includes it in orderedLock(actionsToLock). - w := &worker{ListActiveWorkersResult: testWorker(workerId)} - sl := newSlot(w, newSlotMeta([]string{"A"}, repo.SlotTypeDefault)) - a, err := actionWithSlots("A", sl) - require.NoError(t, err) - s.actions["A"] = a - - a.mu.Lock() + s.actions["A"] = new(action) replenishDone := make(chan error, 1) go func() { @@ -531,25 +623,9 @@ func TestScheduler_Replenish_DoesNotLockUnackedMuBeforeActionLocks(t *testing.T) select { case <-workerSlotConfigsCalled: case <-time.After(2 * time.Second): - a.mu.Unlock() t.Fatalf("timed out waiting for replenish to call ListWorkerSlotConfigs") } - // While replenish is blocked trying to acquire action locks, it must not hold unackedMu. - // If lock order ever regresses (unackedMu before action.mu), this will fail. - deadline := time.Now().Add(50 * time.Millisecond) - for time.Now().Before(deadline) { - if ok := s.unackedMu.TryLock(); ok { - s.unackedMu.Unlock() - } else { - a.mu.Unlock() - t.Fatalf("replenish acquired unackedMu while action.mu was held (lock order violation)") - } - time.Sleep(1 * time.Millisecond) - } - - a.mu.Unlock() - select { case err := <-replenishDone: require.NoError(t, err) @@ -638,9 +714,7 @@ func TestScheduler_TryAssignBatch_AssignsUntilExhausted(t *testing.T) { sl1 := newSlot(w, newSlotMeta([]string{"A"}, repo.SlotTypeDefault)) sl2 := newSlot(w, newSlotMeta([]string{"A"}, repo.SlotTypeDefault)) - actA, err := actionWithSlots("A", sl1, sl2) - require.NoError(t, err) - s.actions["A"] = actA + seedActionPools(t, s, "A", sl1, sl2) qis := []*sqlcv1.V1QueueItem{ testQI(tenantId, "A", 1), @@ -682,9 +756,7 @@ func TestScheduler_TryAssignBatch_RateLimitedSkipsAssignment(t *testing.T) { w := &worker{ListActiveWorkersResult: testWorker(workerId)} sl := newSlot(w, newSlotMeta([]string{"A"}, repo.SlotTypeDefault)) - actA, err := actionWithSlots("A", sl) - require.NoError(t, err) - s.actions["A"] = actA + seedActionPools(t, s, "A", sl) qi := testQI(tenantId, "A", 100) qis := []*sqlcv1.V1QueueItem{qi} @@ -703,18 +775,14 @@ func TestScheduler_TryAssignBatch_RateLimitedSkipsAssignment(t *testing.T) { func TestScheduler_TryAssign_GroupsAndFiltersTimedOut(t *testing.T) { tenantId := uuid.New() - workerId := uuid.New() s := newTestScheduler(t, tenantId, &mockAssignmentRepo{}) - w := &worker{ListActiveWorkersResult: testWorker(workerId)} + workerA := &worker{ListActiveWorkersResult: testWorker(uuid.New())} + workerB := &worker{ListActiveWorkersResult: testWorker(uuid.New())} - // A has 1 slot, B has 1 slot - actA, err := actionWithSlots("A", newSlot(w, newSlotMeta([]string{"A"}, repo.SlotTypeDefault))) - require.NoError(t, err) - s.actions["A"] = actA - actB, err := actionWithSlots("B", newSlot(w, newSlotMeta([]string{"B"}, repo.SlotTypeDefault))) - require.NoError(t, err) - s.actions["B"] = actB + // A and B each have one independently owned worker slot. + seedActionPools(t, s, "A", newSlot(workerA, newSlotMeta([]string{"A"}, repo.SlotTypeDefault))) + seedActionPools(t, s, "B", newSlot(workerB, newSlotMeta([]string{"B"}, repo.SlotTypeDefault))) timeoutQI := testQI(tenantId, "A", 1) timeoutQI.ScheduleTimeoutAt = ts(time.Now().UTC().Add(-1 * time.Second)) @@ -856,12 +924,8 @@ func TestScheduler_GetSnapshotInput_DedupSlotsAcrossActions(t *testing.T) { require.True(t, sharedSlot.use(nil, nil)) // used unusedSlot := newSlot(w, newSlotMeta([]string{"A", "B"}, repo.SlotTypeDefault)) - actA, err := actionWithSlots("A", sharedSlot, unusedSlot) - require.NoError(t, err) - s.actions["A"] = actA - actB, err := actionWithSlots("B", sharedSlot, unusedSlot) // duplicate pointers - require.NoError(t, err) - s.actions["B"] = actB + seedActionPools(t, s, "A", sharedSlot, unusedSlot) + seedActionPools(t, s, "B", sharedSlot, unusedSlot) in, ok := s.getSnapshotInput(context.Background(), true) require.True(t, ok) @@ -896,9 +960,7 @@ func TestScheduler_GetSnapshotInput_WarmupAndSaturation(t *testing.T) { // slots appear in the pool: the slot type warms up and utilization is derived from capacity w := &worker{ListActiveWorkersResult: testWorker(workerId)} freeSlot := newSlot(w, newSlotMeta([]string{"A"}, repo.SlotTypeDefault)) - actA, err := actionWithSlots("A", freeSlot) - require.NoError(t, err) - s.actions["A"] = actA + seedActionPools(t, s, "A", freeSlot) in, ok = s.getSnapshotInput(context.Background(), true) require.True(t, ok) @@ -907,6 +969,7 @@ func TestScheduler_GetSnapshotInput_WarmupAndSaturation(t *testing.T) { // the pool empties out (all slots assigned and flushed): the warmed type now reports // full utilization instead of a transient zero delete(s.actions, "A") + s.pools[poolKey{workerId: workerId, slotType: repo.SlotTypeDefault}].resetSlots(nil) in, ok = s.getSnapshotInput(context.Background(), true) require.True(t, ok) @@ -934,9 +997,7 @@ func TestScheduler_GetSnapshotInput_FallsBackToWalkedCountsWithoutCapacity(t *te require.True(t, usedSlot.use(nil, nil)) unusedSlot := newSlot(w, newSlotMeta([]string{"A"}, repo.SlotTypeDefault)) - actA, err := actionWithSlots("A", usedSlot, unusedSlot) - require.NoError(t, err) - s.actions["A"] = actA + seedActionPools(t, s, "A", usedSlot, unusedSlot) in, ok := s.getSnapshotInput(context.Background(), true) require.True(t, ok) @@ -1010,18 +1071,19 @@ func TestSelectSlotsForWorker_MissingTypeOrInsufficientUnitsFails(t *testing.T) one := newSlot(w, newSlotMeta([]string{"A"}, repo.SlotTypeDefault)) - slotsByTypeAndWorkerId := map[string]map[uuid.UUID][]*slot{ - repo.SlotTypeDefault: {workerId: {one}}, - } + pool := &slotPool{worker: w, slotType: repo.SlotTypeDefault} + pool.resetSlots([]*slot{one}) + poolsByType := map[string]*slotPool{repo.SlotTypeDefault: pool} - _, ok := selectSlotsForWorker(slotsByTypeAndWorkerId, workerId, map[string]int32{repo.SlotTypeDurable: 1}) + _, ok := selectSlotsFromPools(poolsByType, map[string]int32{repo.SlotTypeDurable: 1}) require.False(t, ok) - _, ok = selectSlotsForWorker(slotsByTypeAndWorkerId, workerId, map[string]int32{repo.SlotTypeDefault: 2}) + _, ok = selectSlotsFromPools(poolsByType, map[string]int32{repo.SlotTypeDefault: 2}) require.False(t, ok) } func TestFindAssignableSlots_MultiUnitSameType(t *testing.T) { + s := newTestScheduler(t, uuid.New(), &mockAssignmentRepo{}) workerId := uuid.New() w := &worker{ListActiveWorkersResult: testWorker(workerId)} @@ -1030,10 +1092,9 @@ func TestFindAssignableSlots_MultiUnitSameType(t *testing.T) { s3 := newSlot(w, newSlotMeta([]string{"A"}, repo.SlotTypeDefault)) require.True(t, s3.use(nil, nil)) // used; ensure not selected - a, err := actionWithSlots("A", s1, s2, s3) - require.NoError(t, err) + a := seedActionPools(t, s, "A", s1, s2, s3) - assigned := findAssignableSlots(a.slots, a, map[string]int32{repo.SlotTypeDefault: 2}, nil, nil) + assigned := s.findAssignableWorkerPools(a.workerIds, map[string]int32{repo.SlotTypeDefault: 2}, nil, nil) require.NotNil(t, assigned) require.Len(t, assigned.slots, 2) require.Equal(t, workerId, assigned.workerId()) @@ -1045,16 +1106,16 @@ func TestFindAssignableSlots_MultiUnitSameType(t *testing.T) { } func TestFindAssignableSlots_MultiType(t *testing.T) { + s := newTestScheduler(t, uuid.New(), &mockAssignmentRepo{}) workerId := uuid.New() w := &worker{ListActiveWorkersResult: testWorker(workerId)} def := newSlot(w, newSlotMeta([]string{"A"}, repo.SlotTypeDefault)) dur := newSlot(w, newSlotMeta([]string{"A"}, repo.SlotTypeDurable)) - a, err := actionWithSlots("A", def, dur) - require.NoError(t, err) + a := seedActionPools(t, s, "A", def, dur) - assigned := findAssignableSlots(a.slots, a, map[string]int32{repo.SlotTypeDefault: 1, repo.SlotTypeDurable: 1}, nil, nil) + assigned := s.findAssignableWorkerPools(a.workerIds, map[string]int32{repo.SlotTypeDefault: 1, repo.SlotTypeDurable: 1}, nil, nil) require.NotNil(t, assigned) require.Len(t, assigned.slots, 2) @@ -1152,10 +1213,11 @@ func TestScheduler_Replenish_MultipleSlotTypes_CallsRepoPerTypeAndPopulatesSlots a := s.actions["A"] require.NotNil(t, a) - require.Len(t, a.slots, 4) + actionSlots := slotsForAction(s, a) + require.Len(t, actionSlots, 4) countByType := map[string]int{} - for _, sl := range a.slots { + for _, sl := range actionSlots { if sl.getWorkerId() != workerId { continue } @@ -1212,7 +1274,7 @@ func TestScheduler_Replenish_UnackedCountsPerSlotType(t *testing.T) { countDefault := 0 countDurable := 0 foundUnacked := false - for _, sl := range a.slots { + for _, sl := range slotsForAction(s, a) { if sl.getWorkerId() != workerId { continue } @@ -1318,11 +1380,12 @@ func TestScheduler_Replenish_CreatesActionAndSlots(t *testing.T) { a, ok := s.actions["A"] require.True(t, ok) require.NotNil(t, a) - require.Len(t, a.slots, 3) + actionSlots := slotsForAction(s, a) + require.Len(t, actionSlots, 3) require.Equal(t, 3, a.lastReplenishedSlotCount) require.Equal(t, 1, a.lastReplenishedWorkerCount) - for _, sl := range a.slots { + for _, sl := range actionSlots { require.Equal(t, workerId, sl.getWorkerId()) slotType, err := sl.getSlotType() require.NoError(t, err) @@ -1330,7 +1393,7 @@ func TestScheduler_Replenish_CreatesActionAndSlots(t *testing.T) { } } -func TestScheduler_Replenish_CleansExpiredSlotsWhenNoNewSlotsLoaded(t *testing.T) { +func TestScheduler_Replenish_RemovesActionWhenWorkerPoolHasNoCapacity(t *testing.T) { tenantId := uuid.New() workerId := uuid.New() @@ -1360,27 +1423,23 @@ func TestScheduler_Replenish_CleansExpiredSlotsWhenNoNewSlotsLoaded(t *testing.T used := newSlot(w, newSlotMeta([]string{"A"}, repo.SlotTypeDefault)) require.True(t, used.use(nil, nil)) - actA, err := actionWithSlots("A", expired, used) - require.NoError(t, err) - s.actions["A"] = actA + seedActionPools(t, s, "A", expired, used) s.actions["A"].lastReplenishedSlotCount = 2 - err = s.replenish(context.Background(), false) + err := s.replenish(context.Background(), false) require.NoError(t, err) - a := s.actions["A"] - require.NotNil(t, a) - require.Len(t, a.slots, 1) - require.Same(t, used, a.slots[0]) + _, ok := s.actions["A"] + require.False(t, ok) + require.Empty(t, s.pools[poolKey{workerId: workerId, slotType: repo.SlotTypeDefault}].slots) } -func TestScheduler_Replenish_ClosureVisitsWorkersOfNewlyAddedActions(t *testing.T) { +func TestScheduler_Replenish_RefreshesAllWorkerPoolsWhenTriggered(t *testing.T) { tenantId := uuid.New() // Chain topology: w1 registers {A, B}, w2 registers {B, C}, w3 registers {C, D}. - // Only A independently triggers a replenish (it is a new action). D is reachable - // only transitively: A -> w1 -> B -> w2 -> C -> w3 -> D. The closure must include - // D so its expired slots get cleaned up. + // A new action triggers a refresh of the canonical inventory for every active + // worker, including pools that serve existing actions. w1Id := uuid.New() w2Id := uuid.New() w3Id := uuid.New() @@ -1409,7 +1468,7 @@ func TestScheduler_Replenish_ClosureVisitsWorkersOfNewlyAddedActions(t *testing. w2 := &worker{ListActiveWorkersResult: testWorker(w2Id)} w3 := &worker{ListActiveWorkersResult: testWorker(w3Id)} - // seedAction creates an action with enough active slots that FUNCTION 1 does not + // seedAction creates an action with enough active slots that the heuristic does not // independently mark it for replenish (activeCount > lastReplenishedSlotCount/2 // and worker count unchanged). seedAction := func(actionId string, workerCount int, w *worker, extraSlots ...*slot) *action { @@ -1419,8 +1478,7 @@ func TestScheduler_Replenish_ClosureVisitsWorkersOfNewlyAddedActions(t *testing. } slots = append(slots, extraSlots...) - a, err := actionWithSlots(actionId, slots...) - require.NoError(t, err) + a := seedActionPools(t, s, actionId, slots...) a.lastReplenishedSlotCount = 2 a.lastReplenishedWorkerCount = workerCount @@ -1442,11 +1500,14 @@ func TestScheduler_Replenish_ClosureVisitsWorkersOfNewlyAddedActions(t *testing. require.NoError(t, s.replenish(context.Background(), false)) - // D was only reachable through the worklist closure; its expired slot must be gone. - require.Len(t, actD.slots, 2) - for _, sl := range actD.slots { - require.NotSame(t, expired, sl) + // The refreshed worker inventory is authoritative; stale action-owned slots + // are not carried into the new pools. + require.NotNil(t, actD) + for _, pool := range s.pools { + require.Empty(t, pool.slots) } + _, ok := s.actions["D"] + require.False(t, ok) } func TestScheduler_Replenish_DenseSharedActions(t *testing.T) { @@ -1502,10 +1563,11 @@ func TestScheduler_Replenish_DenseSharedActions(t *testing.T) { for _, aid := range actionIds { a := s.actions[aid] require.NotNil(t, a, "action %s missing after replenish", aid) - require.Len(t, a.slots, numWorkers, "action %s should have one slot per worker", aid) + require.Len(t, a.workerIds, numWorkers, "action %s should index every worker", aid) require.Equal(t, numWorkers, a.lastReplenishedSlotCount) require.Equal(t, numWorkers, a.lastReplenishedWorkerCount) } + require.Len(t, s.pools, numWorkers) } func BenchmarkScheduler_Replenish_DenseSharedActions(b *testing.B) { @@ -1592,34 +1654,22 @@ func TestScheduler_Replenish_UpdatesAllWorkerActionsForLockSafety(t *testing.T) usedSlot := newSlot(w, newSlotMeta([]string{"A", "B"}, repo.SlotTypeDefault)) require.True(t, usedSlot.use(nil, nil)) - actA, err := actionWithSlots("A", usedSlot) - require.NoError(t, err) - s.actions["A"] = actA + seedActionPools(t, s, "A", usedSlot) s.actions["A"].lastReplenishedSlotCount = 2 s.actions["A"].lastReplenishedWorkerCount = 1 - actB, err := actionWithSlots("B", newSlot(w, newSlotMeta([]string{"A", "B"}, repo.SlotTypeDefault))) - require.NoError(t, err) - s.actions["B"] = actB + seedActionPools(t, s, "B", newSlot(w, newSlotMeta([]string{"A", "B"}, repo.SlotTypeDefault))) s.actions["B"].lastReplenishedSlotCount = 100 s.actions["B"].lastReplenishedWorkerCount = 1 - err = s.replenish(context.Background(), false) + err := s.replenish(context.Background(), false) require.NoError(t, err) a := s.actions["A"] b := s.actions["B"] require.NotNil(t, a) require.NotNil(t, b) - require.Len(t, a.slots, 2) - require.Len(t, b.slots, 2) - - // Compare as sets (order is randomized per action). - setA := map[*slot]bool{} - for _, sl := range a.slots { - setA[sl] = true - } - for _, sl := range b.slots { - require.True(t, setA[sl], "expected slot pointers shared across actions for same worker capacity") - } + require.Equal(t, []uuid.UUID{workerId}, a.workerIds) + require.Equal(t, []uuid.UUID{workerId}, b.workerIds) + require.Len(t, s.pools[poolKey{workerId: workerId, slotType: repo.SlotTypeDefault}].slots, 2) } diff --git a/pkg/scheduling/v1/scheduler_worker_scale_stress_bench_test.go b/pkg/scheduling/v1/scheduler_worker_scale_stress_bench_test.go new file mode 100644 index 0000000000..9a14a2a296 --- /dev/null +++ b/pkg/scheduling/v1/scheduler_worker_scale_stress_bench_test.go @@ -0,0 +1,33 @@ +//go:build stress && !e2e && !load && !rampup && !integration + +package v1 + +import ( + "context" + "testing" +) + +func BenchmarkScheduler_InventoryShape_Workers10x(b *testing.B) { + for _, base := range baselineShapes() { + shape := base + shape.Name += "_workers_10x" + shape.Workers *= 10 + + b.Run(shape.Name, func(b *testing.B) { + f := newInventoryFixture(shape) + if err := f.scheduler.replenish(context.Background(), true); err != nil { + b.Fatal(err) + } + f.measureInventory() + + b.ResetTimer() + for iteration := 0; iteration < b.N; iteration++ { + if err := f.scheduler.replenish(context.Background(), true); err != nil { + b.Fatal(err) + } + } + b.StopTimer() + reportShapeMetrics(b, f) + }) + } +} diff --git a/pkg/scheduling/v1/slot.go b/pkg/scheduling/v1/slot.go index c6b4f0fba6..263a6b2c4a 100644 --- a/pkg/scheduling/v1/slot.go +++ b/pkg/scheduling/v1/slot.go @@ -32,6 +32,7 @@ func newSlotMeta(actions []string, slotType string) *slotMeta { type slot struct { worker *worker meta *slotMeta + pool *slotPool expiresAt *time.Time additionalAcks []func() additionalNacks []func() @@ -72,9 +73,7 @@ func (a *assignedSlots) nack() { } } -func newSlot(worker *worker, meta *slotMeta) *slot { - expires := time.Now().Add(defaultSlotExpiry) - +func newSlotWithExpiry(worker *worker, meta *slotMeta, expires time.Time) *slot { return &slot{ worker: worker, meta: meta, @@ -94,11 +93,10 @@ func (s *slot) getSlotType() (string, error) { return s.meta.slotType, nil } -func (s *slot) extendExpiry() { +func (s *slot) setExpiry(expires time.Time) { s.mu.Lock() defer s.mu.Unlock() - expires := time.Now().Add(defaultSlotExpiry) s.expiresAt = &expires } @@ -135,6 +133,9 @@ func (s *slot) use(additionalAcks []func(), additionalNacks []func()) bool { } s.used = true + if s.pool != nil { + s.pool.unused.Add(-1) + } s.ackd = false s.additionalAcks = additionalAcks s.additionalNacks = additionalNacks @@ -162,7 +163,11 @@ func (s *slot) nack() { s.mu.Lock() defer s.mu.Unlock() + wasUsed := s.used s.used = false + if wasUsed && s.pool != nil { + s.pool.unused.Add(1) + } s.ackd = true for _, nack := range s.additionalNacks { diff --git a/pkg/scheduling/v1/slot_cost_test.go b/pkg/scheduling/v1/slot_cost_test.go index 730dd882bc..da81a02873 100644 --- a/pkg/scheduling/v1/slot_cost_test.go +++ b/pkg/scheduling/v1/slot_cost_test.go @@ -26,36 +26,36 @@ func defaultSlots(w *worker, n int) []*slot { } func TestSlotCost_MixedHeavyAndLightShareDefaultPool(t *testing.T) { + s := newTestScheduler(t, uuid.New(), &mockAssignmentRepo{}) workerId := uuid.New() w := &worker{ListActiveWorkersResult: testWorker(workerId)} - a, err := actionWithSlots("A", defaultSlots(w, 6)...) - require.NoError(t, err) + a := seedActionPools(t, s, "A", defaultSlots(w, 6)...) - heavy := findAssignableSlots(a.slots, a, map[string]int32{repo.SlotTypeDefault: 5}, nil, nil) + heavy := s.findAssignableWorkerPools(a.workerIds, map[string]int32{repo.SlotTypeDefault: 5}, nil, nil) require.NotNil(t, heavy) require.Len(t, heavy.slots, 5) - light := findAssignableSlots(a.slots, a, map[string]int32{repo.SlotTypeDefault: 1}, nil, nil) + light := s.findAssignableWorkerPools(a.workerIds, map[string]int32{repo.SlotTypeDefault: 1}, nil, nil) require.NotNil(t, light) require.Len(t, light.slots, 1) - none := findAssignableSlots(a.slots, a, map[string]int32{repo.SlotTypeDefault: 1}, nil, nil) + none := s.findAssignableWorkerPools(a.workerIds, map[string]int32{repo.SlotTypeDefault: 1}, nil, nil) require.Nil(t, none) } func TestSlotCost_ReservationMustFitOnOneWorker(t *testing.T) { + s := newTestScheduler(t, uuid.New(), &mockAssignmentRepo{}) w1 := &worker{ListActiveWorkersResult: testWorker(uuid.New())} w2 := &worker{ListActiveWorkersResult: testWorker(uuid.New())} all := append(defaultSlots(w1, 4), defaultSlots(w2, 4)...) - a, err := actionWithSlots("A", all...) - require.NoError(t, err) + a := seedActionPools(t, s, "A", all...) - none := findAssignableSlots(a.slots, a, map[string]int32{repo.SlotTypeDefault: 5}, nil, nil) + none := s.findAssignableWorkerPools(a.workerIds, map[string]int32{repo.SlotTypeDefault: 5}, nil, nil) require.Nil(t, none) - fits := findAssignableSlots(a.slots, a, map[string]int32{repo.SlotTypeDefault: 4}, nil, nil) + fits := s.findAssignableWorkerPools(a.workerIds, map[string]int32{repo.SlotTypeDefault: 4}, nil, nil) require.NotNil(t, fits) require.Len(t, fits.slots, 4) } @@ -70,9 +70,7 @@ func TestSlotCost_OverCapacityWaitsThenSchedulingTimesOut(t *testing.T) { s := newTestScheduler(t, tenantId, &mockAssignmentRepo{}) w := &worker{ListActiveWorkersResult: testWorker(workerId)} - a, err := actionWithSlots("A", defaultSlots(w, 4)...) - require.NoError(t, err) - s.actions["A"] = a + seedActionPools(t, s, "A", defaultSlots(w, 4)...) waiting := testQI(tenantId, "A", 1) waiting.ScheduleTimeoutAt = ts(time.Now().UTC().Add(5 * time.Minute)) @@ -123,9 +121,7 @@ func TestSlotCost_ExplicitDefaultCostBlocksProportionally(t *testing.T) { s := newTestScheduler(t, tenantId, &mockAssignmentRepo{}) w := &worker{ListActiveWorkersResult: testWorker(workerId)} - a, err := actionWithSlots("A", defaultSlots(w, 2)...) - require.NoError(t, err) - s.actions["A"] = a + seedActionPools(t, s, "A", defaultSlots(w, 2)...) qi1 := testQI(tenantId, "A", 1) qi2 := testQI(tenantId, "A", 2) diff --git a/pkg/scheduling/v1/slot_pool.go b/pkg/scheduling/v1/slot_pool.go new file mode 100644 index 0000000000..1346358607 --- /dev/null +++ b/pkg/scheduling/v1/slot_pool.go @@ -0,0 +1,53 @@ +package v1 + +import ( + "sync/atomic" + "time" + + "github.com/google/uuid" +) + +type poolKey struct { + slotType string + workerId uuid.UUID +} + +// slotPool is the single owner of scheduling capacity for one worker and slot +// type. Actions index workers; they do not copy these slot slices. +type slotPool struct { + refreshedAt time.Time + worker *worker + slotType string + slots []*slot + unused atomic.Int64 +} + +func (p *slotPool) unusedCount() int { + if p == nil { + return 0 + } + return int(p.unused.Load()) +} + +// resetSlotsAt replaces the pool's slot list and rebuilds the unused counter from +// currently schedulable slots. Expired unused slots are not counted: expiry does +// not go through use()/nack(), so counting !isUsed() would leave unused inflated +// relative to active(). at is used so replenish can keep pool staleness aligned +// with slot expiry. +func (p *slotPool) resetSlotsAt(slots []*slot, at time.Time) { + p.slots = slots + p.refreshedAt = at + + unused := int64(0) + for _, sl := range slots { + sl.pool = p + if sl.active() { + unused++ + } + } + p.unused.Store(unused) +} + +func (p *slotPool) staleAt(now time.Time) bool { + return p == nil || p.refreshedAt.IsZero() || !p.refreshedAt.Add(defaultSlotExpiry).After(now) +} diff --git a/pkg/scheduling/v1/slot_test.go b/pkg/scheduling/v1/slot_test.go index 9538dff492..a8d5468ea7 100644 --- a/pkg/scheduling/v1/slot_test.go +++ b/pkg/scheduling/v1/slot_test.go @@ -4,6 +4,7 @@ package v1 import ( "testing" + "time" "github.com/google/uuid" "github.com/jackc/pgx/v5/pgtype" @@ -13,6 +14,14 @@ import ( "github.com/hatchet-dev/hatchet/pkg/repository/sqlcv1" ) +func newSlot(worker *worker, meta *slotMeta) *slot { + return newSlotWithExpiry(worker, meta, time.Now().Add(defaultSlotExpiry)) +} + +func (p *slotPool) resetSlots(slots []*slot) { + p.resetSlotsAt(slots, time.Now()) +} + var stableWorkerId1 = uuid.New() var stableWorkerId2 = uuid.New() @@ -159,25 +168,20 @@ func TestSelectSlotsForWorker(t *testing.T) { workerId := uuid.New() worker := &worker{ListActiveWorkersResult: &v1.ListActiveWorkersResult{ID: workerId}} - slotsByType := map[string]map[uuid.UUID][]*slot{ - "cpu": { - workerId: { - newSlot(worker, newSlotMeta([]string{}, "cpu")), - newSlot(worker, newSlotMeta([]string{}, "cpu")), - newSlot(worker, newSlotMeta([]string{}, "cpu")), - }, - }, - "mem": { - workerId: { - newSlot(worker, newSlotMeta([]string{}, "mem")), - }, - }, - } + cpuPool := &slotPool{worker: worker, slotType: "cpu"} + cpuPool.resetSlots([]*slot{ + newSlot(worker, newSlotMeta(nil, "cpu")), + newSlot(worker, newSlotMeta(nil, "cpu")), + newSlot(worker, newSlotMeta(nil, "cpu")), + }) + memPool := &slotPool{worker: worker, slotType: "mem"} + memPool.resetSlots([]*slot{newSlot(worker, newSlotMeta(nil, "mem"))}) + poolsByType := map[string]*slotPool{"cpu": cpuPool, "mem": memPool} - selected, ok := selectSlotsForWorker(slotsByType, workerId, map[string]int32{"cpu": 2, "mem": 1}) + selected, ok := selectSlotsFromPools(poolsByType, map[string]int32{"cpu": 2, "mem": 1}) assert.True(t, ok) assert.Len(t, selected, 3) - _, ok = selectSlotsForWorker(slotsByType, workerId, map[string]int32{"cpu": 4}) + _, ok = selectSlotsFromPools(poolsByType, map[string]int32{"cpu": 4}) assert.False(t, ok) }