Skip to content

Commit

Permalink
Narrow scope of used slots lock for tracking slot supplier (#1617)
Browse files Browse the repository at this point in the history
  • Loading branch information
Sushisource authored Sep 3, 2024
1 parent eb41d13 commit 2af8c72
Showing 1 changed file with 16 additions and 13 deletions.
29 changes: 16 additions & 13 deletions internal/tuning.go
Original file line number Diff line number Diff line change
Expand Up @@ -417,7 +417,10 @@ func (t *trackingSlotSupplier) ReserveSlot(
return nil, fmt.Errorf("slot supplier returned nil permit")
}
t.issuedSlotsAtomic.Add(1)
t.publishMetrics(false)
t.slotsMutex.Lock()
usedSlots := len(t.usedSlots)
t.slotsMutex.Unlock()
t.publishMetrics(usedSlots)
return permit, nil
}

Expand All @@ -432,7 +435,10 @@ func (t *trackingSlotSupplier) TryReserveSlot(data *slotReservationData) *SlotPe
})
if permit != nil {
t.issuedSlotsAtomic.Add(1)
t.publishMetrics(false)
t.slotsMutex.Lock()
usedSlots := len(t.usedSlots)
t.slotsMutex.Unlock()
t.publishMetrics(usedSlots)
}
return permit
}
Expand All @@ -442,42 +448,39 @@ func (t *trackingSlotSupplier) MarkSlotUsed(permit *SlotPermit) {
panic("Cannot mark nil permit as used")
}
t.slotsMutex.Lock()
defer t.slotsMutex.Unlock()
t.usedSlots[permit] = struct{}{}
usedSlots := len(t.usedSlots)
t.slotsMutex.Unlock()
t.inner.MarkSlotUsed(&slotMarkUsedContextImpl{
permit: permit,
logger: t.logger,
metrics: t.metrics,
})
t.publishMetrics(true)
t.publishMetrics(usedSlots)
}

func (t *trackingSlotSupplier) ReleaseSlot(permit *SlotPermit, reason SlotReleaseReason) {
if permit == nil {
panic("Cannot release with nil permit")
}
t.slotsMutex.Lock()
defer t.slotsMutex.Unlock()
delete(t.usedSlots, permit)
usedSlots := len(t.usedSlots)
t.slotsMutex.Unlock()
t.inner.ReleaseSlot(&slotReleaseContextImpl{
permit: permit,
reason: reason,
logger: t.logger,
metrics: t.metrics,
})
t.issuedSlotsAtomic.Add(-1)
delete(t.usedSlots, permit)
if permit.extraReleaseCallback != nil {
permit.extraReleaseCallback()
}
t.publishMetrics(true)
t.publishMetrics(usedSlots)
}

func (t *trackingSlotSupplier) publishMetrics(lockAlreadyHeld bool) {
if !lockAlreadyHeld {
t.slotsMutex.Lock()
defer t.slotsMutex.Unlock()
}
usedSlots := len(t.usedSlots)
func (t *trackingSlotSupplier) publishMetrics(usedSlots int) {
if t.inner.MaxSlots() != 0 {
t.taskSlotsAvailableGauge.Update(float64(t.inner.MaxSlots() - usedSlots))
}
Expand Down

0 comments on commit 2af8c72

Please sign in to comment.