← 中文资源

简体中文版本

加权信号量限流

用加权信号量按成本限定并发,仅在有剩余容量时接纳任务。

阅读英文原文

功能说明

该示例按“权重”而非任务个数限流:每个任务声明自身成本,只有已用容量加上该成本不超过总容量时才被接纳。这适合内存、带宽或下游配额这类资源——十个轻任务与一个重任务占用的额度不同。提供直接尝试与包裹执行两种用法,后者确保无论成功失败都释放权重。

关键代码

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()
}