Drain Over Limit behavior (#209)

* 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.

.github/workflows/lint.yaml

@@ -6,7 +6,7 @@ on:
   pull_request:
 
 env:
-  GOLANGCI_LINT_VERSION: v1.54.2
+  GOLANGCI_LINT_VERSION: v1.55.2
 
 jobs:
   lint:

.gitignore

@@ -2,10 +2,8 @@
 .vscode/
 __pycache__
 *.pyc
-gubernator.egg-info/
 .DS_Store
 *.iml
-googleapis/
 coverage.out
 coverage.html
 /gubernator

Makefile

@@ -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)
@@ -37,7 +37,8 @@ clean:
 
 .PHONY: proto
 proto:
-	scripts/proto.sh
+	# Install buf: https://buf.build/docs/installation
+	buf generate
 
 .PHONY: certs
 certs:

README.md

@@ -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
@@ -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.
@@ -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
@@ -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.

algorithms.go

@@ -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
 		}
 
@@ -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
 		}
 

config.go

@@ -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

functional_test.go

@@ -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{
@@ -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()
 

go.mod

@@ -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
@@ -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

go.sum

@@ -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=
@@ -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=

gubernator.proto

@@ -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) {
@@ -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
 }