perf(engine): replace assignedCountMu with atomic.Int64

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