From ac5b8f20fbf481382e830a065358831d71ecccb3 Mon Sep 17 00:00:00 2001
From: salarkhannn <salarkhanlala@gmail.com>
Date: Mon, 15 Jun 2026 15:23:05 +0500
Subject: [PATCH 1/2] scheduler: replace assignedCountMu with atomic.Int64

assignedCount is only ever incremented and read under a single mutex
critical section. Replace the int+mutex pair with atomic.Int64 so the
increment-and-read sequence is a single Add(1) call with no lock
contention, eliminating one mutex acquisition per task assignment
on the scheduler hot path.
---
 pkg/scheduling/v1/scheduler.go | 14 +++++---------
 1 file changed, 5 insertions(+), 9 deletions(-)

diff --git a/pkg/scheduling/v1/scheduler.go b/pkg/scheduling/v1/scheduler.go
index 33604986d3..f1f33a65f8 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.
@@ -61,9 +61,8 @@ func newScheduler(cf *sharedConfig, tenantId uuid.UUID, rl *rateLimiter, exts *E
 		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),
+		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

From a10363c10148e6f01d830345233bffd1b7f8bb68 Mon Sep 17 00:00:00 2001
From: salarkhannn <salarkhanlala@gmail.com>
Date: Mon, 15 Jun 2026 18:27:27 +0500
Subject: [PATCH 2/2] style(engine): gofmt alignment in scheduler.go

---
 pkg/scheduling/v1/scheduler.go | 24 ++++++++++++------------
 1 file changed, 12 insertions(+), 12 deletions(-)

diff --git a/pkg/scheduling/v1/scheduler.go b/pkg/scheduling/v1/scheduler.go
index f1f33a65f8..afac8ec9a6 100644
--- a/pkg/scheduling/v1/scheduler.go
+++ b/pkg/scheduling/v1/scheduler.go
@@ -52,18 +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),
-		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,
 	}
 }