Weighted Semaphore Limiter
Bound concurrent work by cost using a weighted semaphore, admitting tasks only while capacity remains.
Code
package limiter
import (
"context"
"errors"
"sync"
)
var ErrExhausted = errors.New("capacity exhausted")
type Weighted struct {
mu sync.Mutex
capacity int64
inUse int64
}
func New(capacity int64) *Weighted {
return &Weighted{capacity: capacity}
}
func (w *Weighted) TryAcquire(weight int64) error {
if weight <= 0 || weight > w.capacity {
return ErrExhausted
}
w.mu.Lock()
defer w.mu.Unlock()
if w.inUse+weight > w.capacity {
return ErrExhausted
}
w.inUse += weight
return nil
}
func (w *Weighted) Release(weight int64) {
w.mu.Lock()
defer w.mu.Unlock()
w.inUse -= weight
if w.inUse < 0 {
w.inUse = 0
}
}
func (w *Weighted) WithContext(ctx context.Context, weight int64, fn func() error) error {
if err := w.TryAcquire(weight); err != nil {
return err
}
defer w.Release(weight)
return fn()
}Report a problem with this exampleOpens a pre-filled issue on GitHub.