功能说明
该示例按“权重”而非任务个数限流:每个任务声明自身成本,只有已用容量加上该成本不超过总容量时才被接纳。这适合内存、带宽或下游配额这类资源——十个轻任务与一个重任务占用的额度不同。提供直接尝试与包裹执行两种用法,后者确保无论成功失败都释放权重。
关键代码
TryAcquire 先做非正权重与超容量权重的快速拒绝,避免出现永远无法满足的请求。在互斥保护下检查已用加权重是否超过容量,满足则累加并返回,否则返回耗尽错误。Release 递减已用并在异常时为负数归零,防止调用方多释放导致容量被虚增。WithContext 用 defer 保证释放,使执行函数抛错也不泄漏容量。
使用注意
本实现是非阻塞的:容量不足立即返回错误,调用方需自行决定重试、排队或降级;如需等待可叠加带超时的重试或改用会阻塞的信号量。权重必须与真实资源成比例,估值过低会让限制失效、过高则浪费容量。Release 必须与 TryAcquire 一一对应,任何遗漏或重复都会破坏容量账本,应尽量通过包裹函数使用而非手工配对。
代码示例
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()
}