diff --git a/pkg/scheduling/v1/scheduler.go b/pkg/scheduling/v1/scheduler.go index 33604986d3..afac8ec9a6 100644 --- a/pkg/scheduling/v1/scheduler.go +++ b/pkg/scheduling/v1/scheduler.go @@ -7,6 +7,7 @@ import ( "slices" "strings" "sync" + "sync/atomic" "time" "github.com/google/uuid" @@ -36,8 +37,7 @@ type Scheduler struct { workersMu mutex workers map[uuid.UUID]*worker - assignedCount int - assignedCountMu mutex + assignedCount atomic.Int64 // unackedSlots are slots which have been assigned to a worker, but have not been flushed // to the database yet. They negatively count towards a worker's available slot count. @@ -52,19 +52,18 @@ func newScheduler(cf *sharedConfig, tenantId uuid.UUID, rl *rateLimiter, exts *E l := cf.l.With().Str("tenant_id", tenantId.String()).Logger() return &Scheduler{ - repo: cf.repo.Assignment(), - tenantId: tenantId, - l: &l, - actions: make(map[string]*action), - unackedSlots: make(map[int]*assignedSlots), - rl: rl, - actionsMu: newRWMu(cf.l), - replenishMu: newMu(cf.l), - workers: map[uuid.UUID]*worker{}, - workersMu: newMu(cf.l), - assignedCountMu: newMu(cf.l), - unackedMu: newMu(cf.l), - exts: exts, + repo: cf.repo.Assignment(), + tenantId: tenantId, + l: &l, + actions: make(map[string]*action), + unackedSlots: make(map[int]*assignedSlots), + rl: rl, + actionsMu: newRWMu(cf.l), + replenishMu: newMu(cf.l), + workers: map[uuid.UUID]*worker{}, + workersMu: newMu(cf.l), + unackedMu: newMu(cf.l), + exts: exts, } } @@ -954,10 +953,7 @@ func (s *Scheduler) tryAssignSingleton( return res, nil } - s.assignedCountMu.Lock() - s.assignedCount++ - res.ackId = s.assignedCount - s.assignedCountMu.Unlock() + res.ackId = int(s.assignedCount.Add(1)) s.unackedMu.Lock() s.unackedSlots[res.ackId] = assignedSlot