GoLang-限流器 time/rate

一、引用包

import "golang.org/x/time/rate"

二、构造限流器

// NewLimiter returns a new Limiter that allows events up to rate r and permits
// bursts of at most b tokens.
func NewLimiter(r Limit, b int) *Limiter {
    return &Limiter{
        limit: r,
        burst: b,
    }
}

构造函数如上。其中NewLimiter 的第一个参数 r Limit,是速率 limit,代表了一秒钟可以产生多少 Token;第二个参数是 b int,b 代表 Token 桶的容量大小。
此限流器工作原理大致如下图所示:


image.png

三、Wait/WaitN

// Wait is shorthand for WaitN(ctx, 1).
func (lim *Limiter) Wait(ctx context.Context) (err error) {
    return lim.WaitN(ctx, 1)
}

// WaitN blocks until lim permits n events to happen.
// It returns an error if n exceeds the Limiter's burst size, the Context is
// canceled, or the expected wait time exceeds the Context's Deadline.
// The burst limit is ignored if the rate limit is Inf.
func (lim *Limiter) WaitN(ctx context.Context, n int) (err error) {
    if n > lim.burst && lim.limit != Inf {
        return fmt.Errorf("rate: Wait(n=%d) exceeds limiter's burst %d", n, lim.burst)
    }
    // Check if ctx is already cancelled
    select {
    case <-ctx.Done():
        return ctx.Err()
    default:
    }
    // Determine wait limit
    now := time.Now()
    waitLimit := InfDuration
    if deadline, ok := ctx.Deadline(); ok {
        waitLimit = deadline.Sub(now)
    }
    // Reserve
    r := lim.reserveN(now, n, waitLimit)
    if !r.ok {
        return fmt.Errorf("rate: Wait(n=%d) would exceed context deadline", n)
    }
    // Wait if necessary
    delay := r.DelayFrom(now)
    if delay == 0 {
        return nil
    }
    t := time.NewTimer(delay)
    defer t.Stop()
    select {
    case <-t.C:
        // We can proceed.
        return nil
    case <-ctx.Done():
        // Context was canceled before we could proceed.  Cancel the
        // reservation, which may permit other events to proceed sooner.
        r.Cancel()
        return ctx.Err()
    }
}

代码如上,Wait 实际上就是 WaitN(ctx,1)。
当使用 Wait 方法消费 Token 时,如果此时桶内 Token 数组不足 (小于 N),那么 Wait 方法将会阻塞一段时间,直至 Token 满足条件。如果充足则直接返回。
Wait 方法有一个 context 参数,我们可以设置 context 的 Deadline 或者 Timeout,来决定此次 Wait 的最长时间。

四、Allow/AllowN

// Allow is shorthand for AllowN(time.Now(), 1).
func (lim *Limiter) Allow() bool {
    return lim.AllowN(time.Now(), 1)
}

// AllowN reports whether n events may happen at time now.
// Use this method if you intend to drop / skip events that exceed the rate limit.
// Otherwise use Reserve or Wait.
func (lim *Limiter) AllowN(now time.Time, n int) bool {
    return lim.reserveN(now, n, 0).ok
}

如上代码所示,Allow 实际上就是 AllowN(time.Now(),1)。
AllowN 方法表示,截止到某一时刻,目前桶中数目是否至少为 n 个,满足则返回 true,同时从桶中消费 n 个 token。反之返回不消费 Token,false。

五、Reserve/ReserveN

// Reserve is shorthand for ReserveN(time.Now(), 1).
func (lim *Limiter) Reserve() *Reservation {
    return lim.ReserveN(time.Now(), 1)
}

// ReserveN returns a Reservation that indicates how long the caller must wait before n events happen.
// The Limiter takes this Reservation into account when allowing future events.
// ReserveN returns false if n exceeds the Limiter's burst size.
// Usage example:
//   r := lim.ReserveN(time.Now(), 1)
//   if !r.OK() {
//     // Not allowed to act! Did you remember to set lim.burst to be > 0 ?
//     return
//   }
//   time.Sleep(r.Delay())
//   Act()
// Use this method if you wish to wait and slow down in accordance with the rate limit without dropping events.
// If you need to respect a deadline or cancel the delay, use Wait instead.
// To drop or skip events exceeding rate limit, use Allow instead.
func (lim *Limiter) ReserveN(now time.Time, n int) *Reservation {
    r := lim.reserveN(now, n, InfDuration)
    return &r
}

同理,Reserve 相当于 ReserveN(time.Now(), 1)。
ReserveN 的用法就相对来说复杂一些,当调用完成后,无论 Token 是否充足,都会返回一个 Reservation * 对象。然后,可以调用该对象的 Delay() 方法,该方法返回了需要等待的时间。如果等待时间为 0,则说明不用等待。
否则必须等到等待时间之后,才能进行接下来的工作。
或者,如果不想等待,可以调用 Cancel() 方法,该方法会将 Token 归还。源码中有给出使用的例子。

六、其他

Limiter 支持可以调整速率和桶大小:

  • SetLimit(Limit) 改变放入 Token 的速率
  • SetBurst(int) 改变 Token 桶大小

你可能感兴趣的:(GoLang-限流器 time/rate)