How can I dynamically increase a rate limiters limit?

Viewed 197

I'm using golang.org/x/time/rate for rate limiting api requests.

lim := rate.NewLimiter(rateLimit, burstLimit)

for range time.NewTicker(time.Second).C {
    fmt.Println(time.Now(), "tick")
    go func() {
        res := lim.Reserve()
        if res.OK() && res.Delay() == 0 {
            fmt.Println(time.Now(), "done")
        } else {
            fmt.Println(time.Now(), "dropped")
            res.Cancel()
        }
    }()
}

For allowed requests that can execute immediately, it may sometimes be necessary to retry the request. Those retried requests should not count against the rate limit. Using res.Cancel() in that case does not return the reserved token to the pool due to https://cs.opensource.google/go/x/time/+/1f47c861:rate/rate.go;l=161.

How can reserved tokens be unreserved or the limiter's status adjusted?

2 Answers

You want rate.Reservation.Cancel() - which does a best effort at "undoing" a reservation - considering the concurrent usage of a rate-limiter.

From the docs:

Cancel is shorthand for CancelAt(time.Now()).

CancelAt indicates that the reservation holder will not perform the reserved action and reverses the effects of this Reservation on the rate limit as much as possible, considering that other reservations may have already been made.

It feels horribly hacky, but it seems one can use ReserveN or AllowN with negative values to return tokens to the pool: https://play.golang.org/p/yt0RS3MJOpi

    lim := rate.NewLimiter(rate.Every(3*time.Second), 1)

    for range time.NewTicker(time.Second).C {
        fmt.Println(time.Now(), "tick --")
        go func() {
            if lim.Allow() {
                fmt.Println(time.Now(), "allow")
                time.Sleep(time.Duration(rand.Intn(20)) * time.Millisecond)

                if rand.Intn(2) == 0 {
                    lim.AllowN(time.Now(), -1)
                    fmt.Println(time.Now(), "revert")
                }
            } else {
                fmt.Println(time.Now(), "drop")
            }
        }()
    }

The example shows that whenever an allowed operation is reverted, the next one is allowed too.

Related