Merge pull request #2 from buildbuddy-io/jpcoenen-go-redis-v8

Add store for go-redis.v8 and introduce context

.github/workflows/ci.yml

@@ -15,7 +15,6 @@ jobs:
     strategy:
       matrix:
         go:
-          - 1.12.x
           - 1.13.x
           - 1.14.x
           - 1.15.x

README.md

@@ -65,7 +65,7 @@ func myHandlerFunc(w http.ResponseWriter, r *http.Request) {
 }
 
 func main() {
-	store, err := memstore.New(65536)
+	store, err := memstore.NewCtx(65536)
 	if err != nil {
 		log.Fatal(err)
 	}
@@ -74,12 +74,12 @@ func main() {
 		MaxRate:  throttled.PerMin(20),
 		MaxBurst: 5,
 	}
-	rateLimiter, err := throttled.NewGCRARateLimiter(store, quota)
+	rateLimiter, err := throttled.NewGCRARateLimiterCtx(store, quota)
 	if err != nil {
 		log.Fatal(err)
 	}
 
-	httpRateLimiter := throttled.HTTPRateLimiter{
+	httpRateLimiter := throttled.HTTPRateLimiterCtx{
 		RateLimiter: rateLimiter,
 		VaryBy:      &throttled.VaryBy{Path: true},
 	}
@@ -89,6 +89,16 @@ func main() {
 }
 ```
 
+### Upgrading to `context.Context` aware version of `throttled`
+To upgrade to the new `context.Context` aware version of `throttled`, update the package to the latest version and replace the following function with their context-aware equivalent:
+- `memstore.New` => `memstore.NewCtx`
+- `goredisstore.New` => `goredisstore.NewCtx`
+- `redigostore.New` => `redigostore.NewCtx`
+- `throttled.NewGCRARateLimiter` => `throttled.NewGCRARateLimiterCtx`
+- `throttled.HTTPRateLimiter` => `throttled.HTTPRateLimiterCtx`
+
+Please note that not all stores make use of the passed `context.Context` yet.
+
 ## Related Projects
 
 See [throttled/gcra][throttled-gcra] for a list of other projects related to

deprecated.go

@@ -1,6 +1,7 @@
 package throttled
 
 import (
+	"context"
 	"net/http"
 	"time"
 )
@@ -65,7 +66,7 @@ func RateLimit(q Quota, vary *VaryBy, store GCRAStore) *Throttler {
 	}
 
 	rate := Rate{period: period / time.Duration(count)}
-	limiter, err := NewGCRARateLimiter(store, RateQuota{rate, count - 1})
+	limiter, err := NewGCRARateLimiterCtx(WrapStoreWithContext(store), RateQuota{rate, count - 1})
 
 	// This panic in unavoidable because the original interface does
 	// not support returning an error.
@@ -87,3 +88,122 @@ func RateLimit(q Quota, vary *VaryBy, store GCRAStore) *Throttler {
 type Store interface {
 	GCRAStore
 }
+
+// HTTPRateLimiter is an adapter for HTTPRateLimiterCtx to provide backwards
+// compatibility.
+//
+// Deprecated: Use HTTPRateLimiterCtx instead. If the used RateLimiter does
+// not implement RateLimiterCtx, wrap it with WrapRateLimiterWithContext().
+type HTTPRateLimiter struct {
+	// DeniedHandler is called if the request is disallowed. If it is
+	// nil, the DefaultDeniedHandler variable is used.
+	DeniedHandler http.Handler
+
+	// Error is called if the RateLimiter returns an error. If it is
+	// nil, the DefaultErrorFunc is used.
+	Error func(w http.ResponseWriter, r *http.Request, err error)
+
+	// Limiter is call for each request to determine whether the
+	// request is permitted and update internal state. It must be set.
+	RateLimiter RateLimiter
+
+	// VaryBy is called for each request to generate a key for the
+	// limiter. If it is nil, all requests use an empty string key.
+	VaryBy interface {
+		Key(*http.Request) string
+	}
+}
+
+// RateLimit provides an adapter for HTTPRateLimiterCtx.RateLimit.
+//
+// Deprecated: Use HTTPRateLimiterCtx instead
+func (t *HTTPRateLimiter) RateLimit(h http.Handler) http.Handler {
+	l := HTTPRateLimiterCtx{
+		DeniedHandler: t.DeniedHandler,
+		Error:         t.Error,
+		RateLimiter:   WrapRateLimiterWithContext(t.RateLimiter),
+		VaryBy:        t.VaryBy,
+	}
+	return l.RateLimit(h)
+}
+
+// GCRAStore is the version of GCRAStoreCtx that is not aware of context.
+//
+// Deprecated: Implement GCRAStoreCtx instead.
+type GCRAStore interface {
+	GetWithTime(key string) (int64, time.Time, error)
+	SetIfNotExistsWithTTL(key string, value int64, ttl time.Duration) (bool, error)
+	CompareAndSwapWithTTL(key string, old, new int64, ttl time.Duration) (bool, error)
+}
+
+// NewGCRARateLimiter is a backwards compatible adapter for NewGCRARateLimiterCtx.
+//
+// Deprecated: Use NewGCRARateLimiterCtx instead. If the used store does
+// not implement GCRAStoreCtx, wrap it with WrapStoreWithContext().
+func NewGCRARateLimiter(st GCRAStore, quota RateQuota) (*GCRARateLimiterCtx, error) {
+	return NewGCRARateLimiterCtx(WrapStoreWithContext(st), quota)
+}
+
+// A RateLimiter manages limiting the rate of actions by key.
+//
+// Deprecated: Use RateLimiterCtx instead.
+type RateLimiter interface {
+	// RateLimit checks whether a particular key has exceeded a rate
+	// limit. It also returns a RateLimitResult to provide additional
+	// information about the state of the RateLimiter.
+	//
+	// If the rate limit has not been exceeded, the underlying storage
+	// is updated by the supplied quantity. For example, a quantity of
+	// 1 might be used to rate limit a single request while a greater
+	// quantity could rate limit based on the size of a file upload in
+	// megabytes. If quantity is 0, no update is performed allowing
+	// you to "peek" at the state of the RateLimiter for a given key.
+	RateLimit(key string, quantity int) (bool, RateLimitResult, error)
+}
+
+// RateLimit is provided as a backwards compatible variant of RateLimitCtx.
+//
+// Deprecated: Use RateLimitCtx instead.
+func (g *GCRARateLimiterCtx) RateLimit(key string, quantity int) (bool, RateLimitResult, error) {
+	return g.RateLimitCtx(context.Background(), key, quantity)
+}
+
+// WrapStoreWithContext can be used to use GCRAStore in a place where a GCRAStoreCtx is required.
+func WrapStoreWithContext(store GCRAStore) GCRAStoreCtx {
+	return gcraStoreCtxAdapter{
+		gcraStore: store,
+	}
+}
+
+// WrapRateLimiterWithContext can be used to use RateLimiter in a place where a RateLimiterCtx is required.
+func WrapRateLimiterWithContext(rateLimier RateLimiter) RateLimiterCtx {
+	return rateLimiterCtxAdapter{
+		rateLimiter: rateLimier,
+	}
+}
+
+// gcraStoreCtxAdapter is an adapter that is used to use a GCRAStore where a GCRAStoreCtx is required.
+type gcraStoreCtxAdapter struct {
+	gcraStore GCRAStore
+}
+
+func (g gcraStoreCtxAdapter) GetWithTime(_ context.Context, key string) (int64, time.Time, error) {
+	return g.gcraStore.GetWithTime(key)
+}
+
+func (g gcraStoreCtxAdapter) SetIfNotExistsWithTTL(_ context.Context, key string, value int64, ttl time.Duration) (bool, error) {
+	return g.gcraStore.SetIfNotExistsWithTTL(key, value, ttl)
+}
+
+func (g gcraStoreCtxAdapter) CompareAndSwapWithTTL(_ context.Context, key string, old, new int64, ttl time.Duration) (bool, error) {
+	return g.gcraStore.CompareAndSwapWithTTL(key, old, new, ttl)
+}
+
+// rateLimiterCtxAdapter is an adapter that is used to use a RateLimiter where a RateLimiterCtx is required.
+type rateLimiterCtxAdapter struct {
+	rateLimiter RateLimiter
+}
+
+func (r rateLimiterCtxAdapter) RateLimitCtx(_ context.Context, key string, quantity int) (bool, RateLimitResult, error) {
+	return r.rateLimiter.RateLimit(key, quantity)
+}

example_test.go

@@ -1,6 +1,7 @@
 package throttled_test
 
 import (
+	"context"
 	"fmt"
 	"log"
 	"net/http"
@@ -13,44 +14,44 @@ var myHandler = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 	w.Write([]byte("hi there!"))
 })
 
-// ExampleHTTPRateLimiter demonstrates the usage of HTTPRateLimiter
+// ExampleHTTPRateLimiter demonstrates the usage of HTTPRateLimiterCtx
 // for rate-limiting access to an http.Handler to 20 requests per path
 // per minute with a maximum burst of 5 requests.
-func ExampleHTTPRateLimiter() {
-	store, err := memstore.New(65536)
+func ExampleHTTPRateLimiterCtx() {
+	store, err := memstore.NewCtx(65536)
 	if err != nil {
 		log.Fatal(err)
 	}
 
 	// Maximum burst of 5 which refills at 20 tokens per minute.
 	quota := throttled.RateQuota{MaxRate: throttled.PerMin(20), MaxBurst: 5}
 
-	rateLimiter, err := throttled.NewGCRARateLimiter(store, quota)
+	rateLimiter, err := throttled.NewGCRARateLimiterCtx(store, quota)
 	if err != nil {
 		log.Fatal(err)
 	}
 
-	httpRateLimiter := throttled.HTTPRateLimiter{
+	httpRateLimiter := throttled.HTTPRateLimiterCtx{
 		RateLimiter: rateLimiter,
 		VaryBy:      &throttled.VaryBy{Path: true},
 	}
 
 	http.ListenAndServe(":8080", httpRateLimiter.RateLimit(myHandler))
 }
 
-// Demonstrates direct use of GCRARateLimiter's RateLimit function (and the
+// Demonstrates direct use of GCRARateLimiterCtx's RateLimit function (and the
 // more general RateLimiter interface). This should be used anywhere where
 // granular control over rate limiting is required.
-func ExampleGCRARateLimiter() {
-	store, err := memstore.New(65536)
+func ExampleGCRARateLimiterCtx() {
+	store, err := memstore.NewCtx(65536)
 	if err != nil {
 		log.Fatal(err)
 	}
 
 	// Maximum burst of 5 which refills at 1 token per hour.
 	quota := throttled.RateQuota{MaxRate: throttled.PerHour(1), MaxBurst: 5}
 
-	rateLimiter, err := throttled.NewGCRARateLimiter(store, quota)
+	rateLimiter, err := throttled.NewGCRARateLimiterCtx(store, quota)
 	if err != nil {
 		log.Fatal(err)
 	}
@@ -65,7 +66,7 @@ func ExampleGCRARateLimiter() {
 	for i := 0; i < 20; i++ {
 		bucket := fmt.Sprintf("by-order:%v", i/10)
 
-		limited, result, err := rateLimiter.RateLimit(bucket, 1)
+		limited, result, err := rateLimiter.RateLimitCtx(context.Background(), bucket, 1)
 		if err != nil {
 			log.Fatal(err)
 		}

go.mod

@@ -4,17 +4,11 @@ go 1.13
 
 require (
 	github.com/go-redis/redis v6.15.8+incompatible
-	github.com/golang/protobuf v1.4.2 // indirect
-	github.com/gomodule/redigo v1.8.4
-	github.com/google/go-cmp v0.5.0 // indirect
+	github.com/go-redis/redis/v8 v8.4.2
+	github.com/gomodule/redigo v2.0.0+incompatible
 	github.com/hashicorp/golang-lru v0.5.4
 	github.com/kr/pretty v0.1.0 // indirect
-	github.com/onsi/ginkgo v1.10.1 // indirect
-	github.com/onsi/gomega v1.7.0 // indirect
-	github.com/stretchr/testify v1.5.1
-	golang.org/x/net v0.0.0-20190923162816-aa69164e4478 // indirect
-	golang.org/x/sys v0.0.0-20191010194322-b09406accb47 // indirect
+	github.com/stretchr/testify v1.6.1
 	golang.org/x/text v0.3.7 // indirect
 	gopkg.in/check.v1 v1.0.0-20190902080502-41f04d3bba15 // indirect
-	gopkg.in/yaml.v2 v2.2.7 // indirect
 )