Skip to content
This repository has been archived by the owner on Apr 19, 2024. It is now read-only.

Commit

Permalink
Drain Over Limit behavior (#209)
Browse files Browse the repository at this point in the history
* Fix protobuf compile process.

* Add `DRAIN_OVERLIMIT` behavior to protobuf.

* Implement `DRAIN_OVERLIMIT` behavior.

* Update golangci-lint.

* Move protobuf compile to `buf` tool.

* Update .gitignore.

* Refactor `DRAIN_OVERLIMIT` to `DRAIN_OVER_LIMIT`.

* Remove unused proto code.

* Update README.
  • Loading branch information
Baliedge authored Feb 9, 2024
1 parent 885519d commit 7ffdfef
Show file tree
Hide file tree
Showing 16 changed files with 237 additions and 124 deletions.
2 changes: 1 addition & 1 deletion .github/workflows/lint.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@ on:
pull_request:

env:
GOLANGCI_LINT_VERSION: v1.54.2
GOLANGCI_LINT_VERSION: v1.55.2

jobs:
lint:
Expand Down
2 changes: 0 additions & 2 deletions .gitignore
Original file line number Diff line number Diff line change
Expand Up @@ -2,10 +2,8 @@
.vscode/
__pycache__
*.pyc
gubernator.egg-info/
.DS_Store
*.iml
googleapis/
coverage.out
coverage.html
/gubernator
Expand Down
5 changes: 3 additions & 2 deletions Makefile
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@
VERSION=$(shell cat version)
LDFLAGS="-X main.Version=$(VERSION)"
GOLANGCI_LINT = $(GOPATH)/bin/golangci-lint
GOLANGCI_LINT_VERSION = 1.54.2
GOLANGCI_LINT_VERSION = 1.55.2

$(GOLANGCI_LINT):
curl -sfL https://raw.githubusercontent.com/golangci/golangci-lint/master/install.sh | sh -s -- -b $(GOPATH)/bin $(GOLANGCI_LINT_VERSION)
Expand Down Expand Up @@ -37,7 +37,8 @@ clean:

.PHONY: proto
proto:
scripts/proto.sh
# Install buf: https://buf.build/docs/installation
buf generate

.PHONY: certs
certs:
Expand Down
31 changes: 28 additions & 3 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -72,7 +72,7 @@ rate_limits:
limit: 100
# The duration of the rate limit in milliseconds
duration: 1000
# The algorithm used to calculate the rate limit
# The algorithm used to calculate the rate limit
# 0 = Token Bucket
# 1 = Leaky Bucket
algorithm: 0
Expand Down Expand Up @@ -166,7 +166,7 @@ Given the following `Duration` values
* 3 = Weeks
* 4 = Months
* 5 = Years

Examples when using `Behavior = DURATION_IS_GREGORIAN`
* If `Duration = 2` (Days) then the rate limit will reset to `Current = 0` at the end of the current day the rate limit was created.
* If `Duration = 0` (Minutes) then the rate limit will reset to `Current = 0` at the end of the minute the rate limit was created.
Expand All @@ -178,6 +178,31 @@ This will reset the rate limit as if created new on first use.

When using Reset Remaining, the `Hits` field should be 0.

## Drain Over Limit Behavior
Users may add behavior `Behavior_DRAIN_OVER_LIMIT` to the rate check request.
A `GetRateLimits` call drains the remaining counter on first over limit event.
Then, successive `GetRateLimits` calls will return zero remaining counter and
not any residual value. This behavior works best with token bucket algorithm
because the `Remaining` counter will stay zero after an over limit until reset
time, whereas leaky bucket algorithm will immediately update `Remaining` to a
non-zero value.

This facilitates scenarios that require an over limit event to stay over limit
until the rate limit resets. This approach is necessary if a process must make
two rate checks, before and after a process, and the `Hit` amount is not known
until after the process.

- Before process: Call `GetRateLimits` with `Hits=0` to check the value of
`Remaining` counter. If `Remaining` is zero, it's known
that the rate limit is depleted and the process can be aborted.
- After process: Call `GetRateLimits` with a user specified `Hits` value. If
the call returns over limit, the process cannot be aborted because it had
already completed. Using `DRAIN_OVER_LIMIT` behavior, the `Remaining` count
will be drained to zero.

Once an over limit occurs in the "After" step, successive processes will detect
the over limit state in the "Before" step.

## Gubernator as a library
If you are using golang, you can use Gubernator as a library. This is useful if
you wish to implement a rate limit service with your own company specific model
Expand Down Expand Up @@ -346,4 +371,4 @@ Gubernator publishes Prometheus metrics for realtime monitoring. See
[prometheus.md](docs/prometheus.md) for details.

## OpenTelemetry Tracing (OTEL)
Gubernator supports OpenTelemetry. See [tracing.md](docs/tracing.md) for details.
Gubernator supports OpenTelemetry. See [tracing.md](docs/tracing.md) for details.
10 changes: 10 additions & 0 deletions algorithms.go
Original file line number Diff line number Diff line change
Expand Up @@ -181,6 +181,11 @@ func tokenBucket(ctx context.Context, s Store, c Cache, r *RateLimitReq) (resp *
trace.SpanFromContext(ctx).AddEvent("Over the limit")
metricOverLimitCounter.Add(1)
rl.Status = Status_OVER_LIMIT
if HasBehavior(r.Behavior, Behavior_DRAIN_OVER_LIMIT) {
// DRAIN_OVER_LIMIT behavior drains the remaining counter.
t.Remaining = 0
rl.Remaining = 0
}
return rl, nil
}

Expand Down Expand Up @@ -394,6 +399,11 @@ func leakyBucket(ctx context.Context, s Store, c Cache, r *RateLimitReq) (resp *
if r.Hits > int64(b.Remaining) {
metricOverLimitCounter.Add(1)
rl.Status = Status_OVER_LIMIT
if HasBehavior(r.Behavior, Behavior_DRAIN_OVER_LIMIT) {
// DRAIN_OVER_LIMIT behavior drains the remaining counter.
b.Remaining = 0
rl.Remaining = 0
}
return rl, nil
}

Expand Down
2 changes: 1 addition & 1 deletion config.go
Original file line number Diff line number Diff line change
Expand Up @@ -62,7 +62,7 @@ type BehaviorConfig struct {
GlobalTimeout time.Duration
// The max number of global updates we can batch into a single peer request
GlobalBatchLimit int
// ForceGlobal forces global mode on all rate limit checks.
// ForceGlobal forces global behavior on all rate limit checks.
ForceGlobal bool

// Number of concurrent requests that will be made to peers. Defaults to 100
Expand Down
74 changes: 74 additions & 0 deletions functional_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -37,6 +37,14 @@ import (
json "google.golang.org/protobuf/encoding/protojson"
)

var algos = []struct {
Name string
Algorithm guber.Algorithm
}{
{Name: "Token bucket", Algorithm: guber.Algorithm_TOKEN_BUCKET},
{Name: "Leaky bucket", Algorithm: guber.Algorithm_LEAKY_BUCKET},
}

// Setup and shutdown the mock gubernator cluster for the entire test suite
func TestMain(m *testing.M) {
if err := cluster.StartWith([]guber.PeerInfo{
Expand Down Expand Up @@ -363,6 +371,72 @@ func TestTokenBucketNegativeHits(t *testing.T) {
}
}

func TestDrainOverLimit(t *testing.T) {
defer clock.Freeze(clock.Now()).Unfreeze()
client, errs := guber.DialV1Server(cluster.PeerAt(0).GRPCAddress, nil)
require.Nil(t, errs)

tests := []struct {
Name string
Hits int64
Remaining int64
Status guber.Status
}{
{
Name: "check remaining before hit",
Hits: 0,
Remaining: 10,
Status: guber.Status_UNDER_LIMIT,
}, {
Name: "first hit",
Hits: 1,
Remaining: 9,
Status: guber.Status_UNDER_LIMIT,
}, {
Name: "over limit hit",
Hits: 100,
Remaining: 0,
Status: guber.Status_OVER_LIMIT,
}, {
Name: "check remaining",
Hits: 0,
Remaining: 0,
Status: guber.Status_UNDER_LIMIT,
},
}

for idx, algoCase := range algos {
t.Run(algoCase.Name, func(t *testing.T) {
for _, test := range tests {
ctx := context.Background()
t.Run(test.Name, func(t *testing.T) {
resp, err := client.GetRateLimits(ctx, &guber.GetRateLimitsReq{
Requests: []*guber.RateLimitReq{
{
Name: "test_drain_over_limit",
UniqueKey: fmt.Sprintf("account:1234:%d", idx),
Algorithm: algoCase.Algorithm,
Behavior: guber.Behavior_DRAIN_OVER_LIMIT,
Duration: guber.Second * 30,
Hits: test.Hits,
Limit: 10,
},
},
})
require.NoError(t, err)
require.Len(t, resp.Responses, 1)

rl := resp.Responses[0]
assert.Equal(t, test.Status, rl.Status)
assert.Equal(t, test.Remaining, rl.Remaining)
assert.Equal(t, int64(10), rl.Limit)
assert.NotZero(t, rl.ResetTime)
})
}
})
}
}

func TestLeakyBucket(t *testing.T) {
defer clock.Freeze(clock.Now()).Unfreeze()

Expand Down
4 changes: 2 additions & 2 deletions go.mod
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,7 @@ require (
github.com/grpc-ecosystem/grpc-gateway/v2 v2.18.0
github.com/hashicorp/memberlist v0.5.0
github.com/mailgun/errors v0.1.5
github.com/mailgun/holster/v4 v4.16.2-0.20231121154636-69040cb71a3b
github.com/mailgun/holster/v4 v4.16.3
github.com/miekg/dns v1.1.50
github.com/pkg/errors v0.9.1
github.com/prometheus/client_golang v1.13.0
Expand All @@ -26,7 +26,7 @@ require (
golang.org/x/time v0.3.0
google.golang.org/genproto/googleapis/api v0.0.0-20231016165738-49dd2c1f3d0b
google.golang.org/grpc v1.59.0
google.golang.org/protobuf v1.31.0
google.golang.org/protobuf v1.32.0
k8s.io/api v0.23.3
k8s.io/apimachinery v0.23.3
k8s.io/client-go v0.23.3
Expand Down
8 changes: 4 additions & 4 deletions go.sum
Original file line number Diff line number Diff line change
Expand Up @@ -291,8 +291,8 @@ github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY=
github.com/kr/text v0.2.0/go.mod h1:eLer722TekiGuMkidMxC/pM04lWEeraHUUmBw8l2grE=
github.com/mailgun/errors v0.1.5 h1:riRpZqfUKTdc8saXvoEg2tYkbRyZESU1KvQ3UxPbdus=
github.com/mailgun/errors v0.1.5/go.mod h1:lw+Nh4r/aoUTz6uK915FdfZJo3yq60gPiflFHNpK4NQ=
github.com/mailgun/holster/v4 v4.16.2-0.20231121154636-69040cb71a3b h1:ohMhrwmmA4JbXNukFpriztFWEVLlMuL90Cssg2Vl2TU=
github.com/mailgun/holster/v4 v4.16.2-0.20231121154636-69040cb71a3b/go.mod h1:phAg61z7LZ1PBfedyt2GXkGSlHhuVKK9AcVJO+Cm0/U=
github.com/mailgun/holster/v4 v4.16.3 h1:YMTkDoaFV83ViSaFuAfiyIvzrHJD1UNw7RjNv6J3Kfg=
github.com/mailgun/holster/v4 v4.16.3/go.mod h1:phAg61z7LZ1PBfedyt2GXkGSlHhuVKK9AcVJO+Cm0/U=
github.com/mailru/easyjson v0.0.0-20190614124828-94de47d64c63/go.mod h1:C1wdFJiN94OJF2b5HbByQZoLdCWB1Yqtg26g4irojpc=
github.com/mailru/easyjson v0.0.0-20190626092158-b2ccc519800e/go.mod h1:C1wdFJiN94OJF2b5HbByQZoLdCWB1Yqtg26g4irojpc=
github.com/mattn/go-colorable v0.0.9/go.mod h1:9vuHe8Xs5qXnSaW/c/ABM9alt+Vo+STaOChaDxuIBZU=
Expand Down Expand Up @@ -839,8 +839,8 @@ google.golang.org/protobuf v1.25.0/go.mod h1:9JNX74DMeImyA3h4bdi1ymwjUzf21/xIlba
google.golang.org/protobuf v1.26.0-rc.1/go.mod h1:jlhhOSvTdKEhbULTjvd4ARK9grFBp09yW+WbY/TyQbw=
google.golang.org/protobuf v1.26.0/go.mod h1:9q0QmTI4eRPtz6boOQmLYwt+qCgq0jsYwAQnmE0givc=
google.golang.org/protobuf v1.27.1/go.mod h1:9q0QmTI4eRPtz6boOQmLYwt+qCgq0jsYwAQnmE0givc=
google.golang.org/protobuf v1.31.0 h1:g0LDEJHgrBl9N9r17Ru3sqWhkIx2NB67okBHPwC7hs8=
google.golang.org/protobuf v1.31.0/go.mod h1:HV8QOd/L58Z+nl8r43ehVNZIU/HEI6OcFqwMG9pJV4I=
google.golang.org/protobuf v1.32.0 h1:pPC6BG5ex8PDFnkbrGU3EixyhKcQ2aDuBS36lqK/C7I=
google.golang.org/protobuf v1.32.0/go.mod h1:c6P6GXX6sHbq/GpV6MGZEdwhWPcYBgnhAHhKbcUYpos=
gopkg.in/alecthomas/kingpin.v2 v2.2.6/go.mod h1:FMv+mEhP44yOT+4EoQTLFTRgOQ1FBLkstjWtayDeSgw=
gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
Expand Down
72 changes: 39 additions & 33 deletions gubernator.pb.go

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

6 changes: 5 additions & 1 deletion gubernator.proto
Original file line number Diff line number Diff line change
Expand Up @@ -34,7 +34,6 @@ service V1 {
};
}


// This method is for round trip benchmarking and can be used by
// the client to determine connectivity to the server
rpc HealthCheck (HealthCheckReq) returns (HealthCheckResp) {
Expand Down Expand Up @@ -127,6 +126,11 @@ enum Behavior {
// least 2 instances of Gubernator.
MULTI_REGION = 16;

// A GetRateLimits call drains the remaining counter on first over limit
// event. Then, successive GetRateLimits calls will return zero remaining
// counter and not any residual value.
DRAIN_OVER_LIMIT = 32;

// TODO: Add support for LOCAL. Which would force the rate limit to be handled by the local instance
}

Expand Down
Loading

0 comments on commit 7ffdfef

Please sign in to comment.