115 lines
3.7 KiB
Go
115 lines
3.7 KiB
Go
package ratelimits
|
|
|
|
import (
|
|
"time"
|
|
|
|
"github.com/jmhodges/clock"
|
|
)
|
|
|
|
// maybeSpend uses the GCRA algorithm to decide whether to allow a request. It
|
|
// returns a Decision struct with the result of the decision and the updated
|
|
// TAT. The cost must be 0 or greater and <= the burst capacity of the limit.
|
|
func maybeSpend(clk clock.Clock, txn Transaction, tat time.Time) *Decision {
|
|
if txn.cost < 0 || txn.cost > txn.limit.burst {
|
|
// The condition above is the union of the conditions checked in Check
|
|
// and Spend methods of Limiter. If this panic is reached, it means that
|
|
// the caller has introduced a bug.
|
|
panic("invalid cost for maybeSpend")
|
|
}
|
|
nowUnix := clk.Now().UnixNano()
|
|
tatUnix := tat.UnixNano()
|
|
|
|
// If the TAT is in the future, use it as the starting point for the
|
|
// calculation. Otherwise, use the current time. This is to prevent the
|
|
// bucket from being filled with capacity from the past.
|
|
if nowUnix > tatUnix {
|
|
tatUnix = nowUnix
|
|
}
|
|
|
|
// Compute the cost increment.
|
|
costIncrement := txn.limit.emissionInterval * txn.cost
|
|
|
|
// Deduct the cost to find the new TAT and residual capacity.
|
|
newTAT := tatUnix + costIncrement
|
|
difference := nowUnix - (newTAT - txn.limit.burstOffset)
|
|
|
|
if difference < 0 {
|
|
// Too little capacity to satisfy the cost, deny the request.
|
|
residual := (nowUnix - (tatUnix - txn.limit.burstOffset)) / txn.limit.emissionInterval
|
|
return &Decision{
|
|
allowed: false,
|
|
remaining: residual,
|
|
retryIn: -time.Duration(difference),
|
|
resetIn: time.Duration(tatUnix - nowUnix),
|
|
newTAT: time.Unix(0, tatUnix).UTC(),
|
|
transaction: txn,
|
|
}
|
|
}
|
|
|
|
// There is enough capacity to satisfy the cost, allow the request.
|
|
var retryIn time.Duration
|
|
residual := difference / txn.limit.emissionInterval
|
|
if difference < costIncrement {
|
|
retryIn = time.Duration(costIncrement - difference)
|
|
}
|
|
return &Decision{
|
|
allowed: true,
|
|
remaining: residual,
|
|
retryIn: retryIn,
|
|
resetIn: time.Duration(newTAT - nowUnix),
|
|
newTAT: time.Unix(0, newTAT).UTC(),
|
|
transaction: txn,
|
|
}
|
|
}
|
|
|
|
// maybeRefund uses the Generic Cell Rate Algorithm (GCRA) to attempt to refund
|
|
// the cost of a request which was previously spent. The refund cost must be 0
|
|
// or greater. A cost will only be refunded up to the burst capacity of the
|
|
// limit. A partial refund is still considered successful.
|
|
func maybeRefund(clk clock.Clock, txn Transaction, tat time.Time) *Decision {
|
|
if txn.cost < 0 || txn.cost > txn.limit.burst {
|
|
// The condition above is checked in the Refund method of Limiter. If
|
|
// this panic is reached, it means that the caller has introduced a bug.
|
|
panic("invalid cost for maybeRefund")
|
|
}
|
|
nowUnix := clk.Now().UnixNano()
|
|
tatUnix := tat.UnixNano()
|
|
|
|
// The TAT must be in the future to refund capacity.
|
|
if nowUnix > tatUnix {
|
|
// The TAT is in the past, therefore the bucket is full.
|
|
return &Decision{
|
|
allowed: false,
|
|
remaining: txn.limit.burst,
|
|
retryIn: time.Duration(0),
|
|
resetIn: time.Duration(0),
|
|
newTAT: tat,
|
|
transaction: txn,
|
|
}
|
|
}
|
|
|
|
// Compute the refund increment.
|
|
refundIncrement := txn.limit.emissionInterval * txn.cost
|
|
|
|
// Subtract the refund increment from the TAT to find the new TAT.
|
|
newTAT := tatUnix - refundIncrement
|
|
|
|
// Ensure the new TAT is not earlier than now.
|
|
if newTAT < nowUnix {
|
|
newTAT = nowUnix
|
|
}
|
|
|
|
// Calculate the new capacity.
|
|
difference := nowUnix - (newTAT - txn.limit.burstOffset)
|
|
residual := difference / txn.limit.emissionInterval
|
|
|
|
return &Decision{
|
|
allowed: (newTAT != tatUnix),
|
|
remaining: residual,
|
|
retryIn: time.Duration(0),
|
|
resetIn: time.Duration(newTAT - nowUnix),
|
|
newTAT: time.Unix(0, newTAT).UTC(),
|
|
transaction: txn,
|
|
}
|
|
}
|