From ca60baa04cf28742df5936e1ca5cb988fa72fc56 Mon Sep 17 00:00:00 2001 From: Keegan Carruthers-Smith Date: Mon, 10 Jul 2023 16:16:50 +0200 Subject: [PATCH 1/3] graphqlbackend: update throttled to use Ctx stores It seems we were using a weird version of throttled which I don't think was actually based on the v2 series correctly. This updates us to the latest version and moves us away from the removed struct GCRARateLimiter and the deprecated interface GCRAStore. I came across this while attempting to update deps and getting failing compiles. Test Plan: we have unit tests for this, so CI. --- cmd/frontend/graphqlbackend/rate_limit.go | 16 ++++++++-------- cmd/frontend/graphqlbackend/rate_limit_test.go | 4 ++-- cmd/frontend/internal/cli/serve_cmd.go | 6 +++--- deps.bzl | 16 ++++++++-------- go.mod | 2 +- go.sum | 14 ++++++++++---- 6 files changed, 32 insertions(+), 26 deletions(-) diff --git a/cmd/frontend/graphqlbackend/rate_limit.go b/cmd/frontend/graphqlbackend/rate_limit.go index 4d2e2ad209b3..b3bd7bf95dc3 100644 --- a/cmd/frontend/graphqlbackend/rate_limit.go +++ b/cmd/frontend/graphqlbackend/rate_limit.go @@ -376,7 +376,7 @@ type LimitWatcher interface { Get() (Limiter, bool) } -func NewBasicLimitWatcher(logger log.Logger, store throttled.GCRAStore) *BasicLimitWatcher { +func NewBasicLimitWatcher(logger log.Logger, store throttled.GCRAStoreCtx) *BasicLimitWatcher { basic := &BasicLimitWatcher{ store: store, } @@ -392,7 +392,7 @@ func NewBasicLimitWatcher(logger log.Logger, store throttled.GCRAStore) *BasicLi } type BasicLimitWatcher struct { - store throttled.GCRAStore + store throttled.GCRAStoreCtx rl atomic.Value // *RateLimiter } @@ -403,7 +403,7 @@ func (bl *BasicLimitWatcher) updateFromConfig(logger log.Logger, limit int) { return } maxBurstPercentage := 0.2 - l, err := throttled.NewGCRARateLimiter( + l, err := throttled.NewGCRARateLimiterCtx( bl.store, throttled.RateQuota{ MaxRate: throttled.PerHour(limit), @@ -428,23 +428,23 @@ func (bl *BasicLimitWatcher) Get() (Limiter, bool) { } type BasicLimiter struct { - *throttled.GCRARateLimiter + *throttled.GCRARateLimiterCtx enabled bool } // RateLimit limits unauthenticated requests to the GraphQL API with an equal // quantity of 1. func (bl *BasicLimiter) RateLimit(_ string, _ int, args LimiterArgs) (bool, throttled.RateLimitResult, error) { - if args.Anonymous && args.RequestName == "unknown" && args.RequestSource == trace.SourceOther && bl.GCRARateLimiter != nil { - return bl.GCRARateLimiter.RateLimit("basic", 1) + if args.Anonymous && args.RequestName == "unknown" && args.RequestSource == trace.SourceOther && bl.GCRARateLimiterCtx != nil { + return bl.GCRARateLimiterCtx.RateLimit("basic", 1) } return false, throttled.RateLimitResult{}, nil } type RateLimiter struct { enabled bool - ipLimiter *throttled.GCRARateLimiter - userLimiter *throttled.GCRARateLimiter + ipLimiter *throttled.GCRARateLimiterCtx + userLimiter *throttled.GCRARateLimiterCtx overrides map[string]limiter } diff --git a/cmd/frontend/graphqlbackend/rate_limit_test.go b/cmd/frontend/graphqlbackend/rate_limit_test.go index fc3bfbc66f58..fc47d1c6b235 100644 --- a/cmd/frontend/graphqlbackend/rate_limit_test.go +++ b/cmd/frontend/graphqlbackend/rate_limit_test.go @@ -415,7 +415,7 @@ func TestBasicLimiterEnabled(t *testing.T) { for _, tt := range tests { t.Run(fmt.Sprintf("limit:%d", tt.limit), func(t *testing.T) { - store, err := memstore.New(1) + store, err := memstore.NewCtx(1) if err != nil { t.Fatal(err) } @@ -435,7 +435,7 @@ func TestBasicLimiterEnabled(t *testing.T) { } func TestBasicLimiter(t *testing.T) { - store, err := memstore.New(1) + store, err := memstore.NewCtx(1) if err != nil { t.Fatal(err) } diff --git a/cmd/frontend/internal/cli/serve_cmd.go b/cmd/frontend/internal/cli/serve_cmd.go index b0bc7d002315..9389d54d491b 100644 --- a/cmd/frontend/internal/cli/serve_cmd.go +++ b/cmd/frontend/internal/cli/serve_cmd.go @@ -379,14 +379,14 @@ func isAllowedOrigin(origin string, allowedOrigins []string) bool { } func makeRateLimitWatcher() (*graphqlbackend.BasicLimitWatcher, error) { - var store throttled.GCRAStore + var store throttled.GCRAStoreCtx var err error if pool, ok := redispool.Cache.Pool(); ok { - store, err = redigostore.New(pool, "gql:rl:", 0) + store, err = redigostore.NewCtx(pool, "gql:rl:", 0) } else { // If redis is disabled we are in Sourcegraph App and can rely on an // in-memory store. - store, err = memstore.New(0) + store, err = memstore.NewCtx(0) } if err != nil { return nil, err diff --git a/deps.bzl b/deps.bzl index a6d36f285a8e..e986676d4eeb 100644 --- a/deps.bzl +++ b/deps.bzl @@ -899,15 +899,15 @@ def go_dependencies(): name = "com_github_bsm_ginkgo_v2", build_file_proto_mode = "disable_global", importpath = "github.com/bsm/ginkgo/v2", - sum = "h1:aOAnND1T40wEdAtkGSkvSICWeQ8L3UASX7YVCqQx+eQ=", - version = "v2.5.0", + sum = "h1:ItPMPH90RbmZJt5GtkcNvIRuGEdwlBItdNVoyzaNQao=", + version = "v2.7.0", ) go_repository( name = "com_github_bsm_gomega", build_file_proto_mode = "disable_global", importpath = "github.com/bsm/gomega", - sum = "h1:JhAwLmtRzXFTx2AkALSLa8ijZafntmhSoU63Ok18Uq8=", - version = "v1.20.0", + sum = "h1:LhQm+AFcgV2M0WyKroMASzAzCAJVpAxQXv4SaI9a69Y=", + version = "v1.26.0", ) go_repository( @@ -5806,8 +5806,8 @@ def go_dependencies(): name = "com_github_redis_go_redis_v9", build_file_proto_mode = "disable_global", importpath = "github.com/redis/go-redis/v9", - sum = "h1:BA426Zqe/7r56kCcvxYLWe1mkaz71LKF77GwgFzSxfE=", - version = "v9.0.2", + sum = "h1:CuQcn5HIEeK7BgElubPP8CGtE0KakrnbBSTLjathl5o=", + version = "v9.0.5", ) go_repository( @@ -6510,8 +6510,8 @@ def go_dependencies(): name = "com_github_throttled_throttled_v2", build_file_proto_mode = "disable_global", importpath = "github.com/throttled/throttled/v2", - sum = "h1:DOkCb1el7NYzRoPb1pyeHVghsUoonVWEjmo34vrcp/8=", - version = "v2.9.0", + sum = "h1:IezKE1uHlYC/0Al05oZV6Ar+uN/znw3cy9J8banxhEY=", + version = "v2.12.0", ) go_repository( name = "com_github_tidwall_gjson", diff --git a/go.mod b/go.mod index 8ebe6ac0361b..14fd2f190e55 100644 --- a/go.mod +++ b/go.mod @@ -201,7 +201,7 @@ require ( github.com/sourcegraph/sourcegraph/lib v0.0.0-20230613175844-f031949c72f5 github.com/stretchr/testify v1.8.2 github.com/temoto/robotstxt v1.1.2 - github.com/throttled/throttled/v2 v2.9.0 + github.com/throttled/throttled/v2 v2.12.0 github.com/tidwall/gjson v1.14.0 github.com/tj/go-naturaldate v1.3.0 github.com/tomnomnom/linkheader v0.0.0-20180905144013-02ca5825eb80 diff --git a/go.sum b/go.sum index fb7f4887468b..e0bdf885861a 100644 --- a/go.sum +++ b/go.sum @@ -379,7 +379,9 @@ github.com/bradleyjkemp/cupaloy/v2 v2.6.0/go.mod h1:bm7JXdkRd4BHJk9HpwqAI8BoAY1l github.com/bshuster-repo/logrus-logstash-hook v0.4.1/go.mod h1:zsTqEiSzDgAa/8GZR7E1qaXrhYNDKBYy5/dWPTIflbk= github.com/bshuster-repo/logrus-logstash-hook v1.0.0/go.mod h1:zsTqEiSzDgAa/8GZR7E1qaXrhYNDKBYy5/dWPTIflbk= github.com/bsm/ginkgo/v2 v2.5.0/go.mod h1:AiKlXPm7ItEHNc/2+OkrNG4E0ITzojb9/xWzvQ9XZ9w= +github.com/bsm/ginkgo/v2 v2.7.0/go.mod h1:AiKlXPm7ItEHNc/2+OkrNG4E0ITzojb9/xWzvQ9XZ9w= github.com/bsm/gomega v1.20.0/go.mod h1:JifAceMQ4crZIWYUKrlGcmbN3bqHogVTADMD2ATsbwk= +github.com/bsm/gomega v1.26.0/go.mod h1:JyEr/xRbxbtgWNi8tIEVPUYZ5Dzef52k01W3YH0H+O0= github.com/bufbuild/buf v1.4.0 h1:GqE3a8CMmcFvWPzuY3Mahf9Kf3S9XgZ/ORpfYFzO+90= github.com/bufbuild/buf v1.4.0/go.mod h1:mwHG7klTHnX+rM/ym8LXGl7vYpVmnwT96xWoRB4H5QI= github.com/buger/jsonparser v0.0.0-20180808090653-f4dd9f5a6b44/go.mod h1:bbYlZJ7hK1yFx9hf58LP0zeX7UjIGs20ufpu3evjr+s= @@ -963,6 +965,7 @@ github.com/go-redis/redis v6.15.9+incompatible h1:K0pv1D7EQUjfyoMql+r/jZqCLizCGK github.com/go-redis/redis v6.15.9+incompatible/go.mod h1:NAIEuMOZ/fxfXJIrKDQDz8wamY7mA7PouImQ2Jvg6kA= github.com/go-redis/redis/v7 v7.4.0 h1:7obg6wUoj05T0EpY0o8B59S9w5yeMWql7sw2kwNW1x4= github.com/go-redis/redis/v7 v7.4.0/go.mod h1:JDNMw23GTyLNC4GZu9njt15ctBQVn7xjRfnwdHj/Dcg= +github.com/go-redis/redis/v8 v8.4.2/go.mod h1:A1tbYoHSa1fXwN+//ljcCYYJeLmVrwL9hbQN45Jdy0M= github.com/go-redis/redis/v8 v8.11.4/go.mod h1:2Z2wHZXdQpCDXEGzqMockDpNyYvi2l4Pxt6RJr792+w= github.com/go-redis/redis/v8 v8.11.5 h1:AcZZR7igkdvfVmQTPnu9WE37LRrO/YrBH5zWyjDC0oI= github.com/go-redsync/redsync/v4 v4.8.1 h1:rq2RvdTI0obznMdxKUWGdmmulo7lS9yCzb8fgDKOlbM= @@ -1765,6 +1768,7 @@ github.com/onsi/ginkgo v1.12.0/go.mod h1:oUhWkIvk5aDxtKvDDuw8gItl8pKl42LzjC9KZE0 github.com/onsi/ginkgo v1.12.1/go.mod h1:zj2OWP4+oCPe1qIXoGWkgMRwljMUYCdkwsT2108oapk= github.com/onsi/ginkgo v1.13.0/go.mod h1:+REjRxOmWfHCjfv9TTWB1jD1Frx4XydAD3zm1lskyM0= github.com/onsi/ginkgo v1.14.0/go.mod h1:iSB4RoI2tjJc9BBv4NKIKWKya62Rps+oPG/Lv9klQyY= +github.com/onsi/ginkgo v1.14.2/go.mod h1:iSB4RoI2tjJc9BBv4NKIKWKya62Rps+oPG/Lv9klQyY= github.com/onsi/ginkgo v1.15.2/go.mod h1:Dd6YFfwBW84ETqqtL0CPyPXillHgY6XhQH3uuCCTr/o= github.com/onsi/ginkgo v1.16.4/go.mod h1:dX+/inL/fNMqNlz0e9LfyB9TswhZpCVdJM/Z6Vvnwo0= github.com/onsi/ginkgo/v2 v2.1.3/go.mod h1:vw5CSIxN1JObi/U8gcbwft7ZxR2dgaR70JSE3/PpL4c= @@ -1927,8 +1931,9 @@ github.com/qustavo/sqlhooks/v2 v2.1.0/go.mod h1:aMREyKo7fOKTwiLuWPsaHRXEmtqG4yRE github.com/rafaeljusto/redigomock/v3 v3.1.2 h1:B4Y0XJQiPjpwYmkH55aratKX1VfR+JRqzmDKyZbC99o= github.com/rafaeljusto/redigomock/v3 v3.1.2/go.mod h1:F9zPqz8rMriScZkPtUiLJoLruYcpGo/XXREpeyasREM= github.com/rainycape/unidecode v0.0.0-20150907023854-cb7f23ec59be/go.mod h1:MIDFMn7db1kT65GmV94GzpX9Qdi7N/pQlwb+AN8wh+Q= -github.com/redis/go-redis/v9 v9.0.2 h1:BA426Zqe/7r56kCcvxYLWe1mkaz71LKF77GwgFzSxfE= github.com/redis/go-redis/v9 v9.0.2/go.mod h1:/xDTe9EF1LM61hek62Poq2nzQSGj0xSrEtEHbBQevps= +github.com/redis/go-redis/v9 v9.0.5 h1:CuQcn5HIEeK7BgElubPP8CGtE0KakrnbBSTLjathl5o= +github.com/redis/go-redis/v9 v9.0.5/go.mod h1:WqMKv5vnQbRuZstUwxQI195wHy+t4PuXDOjzMvcuQHk= github.com/rhnvrm/simples3 v0.6.1/go.mod h1:Y+3vYm2V7Y4VijFoJHHTrja6OgPrJ2cBti8dPGkC3sA= github.com/ricochet2200/go-disk-usage/du v0.0.0-20210707232629-ac9918953285 h1:d54EL9l+XteliUfUCGsEwwuk65dmmxX85VXF+9T6+50= github.com/ricochet2200/go-disk-usage/du v0.0.0-20210707232629-ac9918953285/go.mod h1:fxIDly1xtudczrZeOOlfaUvd2OPb2qZAPuWdU2BsBTk= @@ -2139,8 +2144,8 @@ github.com/tchap/go-patricia v2.2.6+incompatible/go.mod h1:bmLyhP68RS6kStMGxByiQ github.com/temoto/robotstxt v1.1.2 h1:W2pOjSJ6SWvldyEuiFXNxz3xZ8aiWX5LbfDiOFd7Fxg= github.com/temoto/robotstxt v1.1.2/go.mod h1:+1AmkuG3IYkh1kv0d2qEB9Le88ehNO0zwOr3ujewlOo= github.com/testcontainers/testcontainers-go v0.19.0/go.mod h1:3YsSoxK0rGEUzbGD4gUVt1Nm3GJpCIq94GX+2LSf3d4= -github.com/throttled/throttled/v2 v2.9.0 h1:DOkCb1el7NYzRoPb1pyeHVghsUoonVWEjmo34vrcp/8= -github.com/throttled/throttled/v2 v2.9.0/go.mod h1:0JHxhGAidPyqbgD4HF8Y1sNFfG0ffVXK6C8EpkNdLEM= +github.com/throttled/throttled/v2 v2.12.0 h1:IezKE1uHlYC/0Al05oZV6Ar+uN/znw3cy9J8banxhEY= +github.com/throttled/throttled/v2 v2.12.0/go.mod h1:+EAvrG2hZAQTx8oMpBu8fq6Xmm+d1P2luKK7fIY1Esc= github.com/tidwall/gjson v1.14.0 h1:6aeJ0bzojgWLa82gDQHcx3S0Lr/O51I9bJ5nv6JFx5w= github.com/tidwall/gjson v1.14.0/go.mod h1:/wbyibRr2FHMks5tjHJ5F8dMZh3AcwJEMf5vlfC0lxk= github.com/tidwall/match v1.1.1 h1:+Ho715JplO36QYgwN9PGYNhgZvoUSc9X2c80KVTi+GA= @@ -2354,6 +2359,7 @@ go.opentelemetry.io/contrib/propagators/jaeger v1.14.0 h1:j6Xah53xRDrR+K1c4Y1uVH go.opentelemetry.io/contrib/propagators/jaeger v1.14.0/go.mod h1:viOfwr1OqHmCF6G3KvKnnmpSJUX/rLzXztU18FC9ymU= go.opentelemetry.io/contrib/propagators/ot v1.14.0 h1:jqxznjMkz/3l4NUsYq4OMbP+zs5twBbCZwSlSt82KXo= go.opentelemetry.io/contrib/propagators/ot v1.14.0/go.mod h1:0qeUHgw3lmmZupgaft3m2/K6ip+pzqwePuW/lvU7ia4= +go.opentelemetry.io/otel v0.14.0/go.mod h1:vH5xEuwy7Rts0GNtsCW3HYQoZDY+OmBJ6t1bFGGlxgw= go.opentelemetry.io/otel v0.20.0/go.mod h1:Y3ugLH2oa81t5QO+Lty+zXf8zC9L26ax4Nzoxm/dooo= go.opentelemetry.io/otel v1.3.0/go.mod h1:PWIKzi6JCp7sM0k9yZ43VX+T345uNbAkDKwHVjb2PTs= go.opentelemetry.io/otel v1.6.3/go.mod h1:7BgNga5fNlF/iZjG06hM3yofffp0ofKCDwSXx1GC4dI= @@ -2837,6 +2843,7 @@ golang.org/x/text v0.3.4/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ= golang.org/x/text v0.3.5/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ= golang.org/x/text v0.3.6/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ= golang.org/x/text v0.3.7/go.mod h1:u+2+/6zg+i71rQMx5EYifcz6MCKuco9NR6JIITiCfzQ= +golang.org/x/text v0.3.8/go.mod h1:E6s5w1FMmriuDzIBO73fBruAKo1PCIq6d2Q6DHfQ8WQ= golang.org/x/text v0.4.0/go.mod h1:mrYo+phRRbMaCq/xk9113O4dZlRixOauAjOtrjsXDZ8= golang.org/x/text v0.5.0/go.mod h1:mrYo+phRRbMaCq/xk9113O4dZlRixOauAjOtrjsXDZ8= golang.org/x/text v0.7.0/go.mod h1:mrYo+phRRbMaCq/xk9113O4dZlRixOauAjOtrjsXDZ8= @@ -3225,7 +3232,6 @@ gopkg.in/yaml.v2 v2.2.2/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI= gopkg.in/yaml.v2 v2.2.3/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI= gopkg.in/yaml.v2 v2.2.4/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI= gopkg.in/yaml.v2 v2.2.5/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI= -gopkg.in/yaml.v2 v2.2.7/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI= gopkg.in/yaml.v2 v2.2.8/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI= gopkg.in/yaml.v2 v2.3.0/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI= gopkg.in/yaml.v2 v2.4.0 h1:D8xgwECY7CYvx+Y2n4sBz93Jn9JRvxdiyyo8CTfuKaY= From 9cb75793b4d273c9424b562249aec76daa94dc6b Mon Sep 17 00:00:00 2001 From: Keegan Carruthers-Smith Date: Mon, 10 Jul 2023 18:26:06 +0200 Subject: [PATCH 2/3] backcompat spray and pray --- dev/backcompat/defs.bzl | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/dev/backcompat/defs.bzl b/dev/backcompat/defs.bzl index db17ca4f82db..73da65975139 100644 --- a/dev/backcompat/defs.bzl +++ b/dev/backcompat/defs.bzl @@ -56,6 +56,7 @@ fi find . -type f -name "*.bazel" -exec $_sed_binary -i 's|@com_github_sourcegraph_conc|@back_compat_com_github_sourcegraph_conc|g' {} + find . -type f -name "*.bazel" -exec $_sed_binary -i 's|@com_github_sourcegraph_scip|@back_compat_com_github_sourcegraph_scip|g' {} + find . -type f -name "*.bazel" -exec $_sed_binary -i 's|@com_github_sourcegraph_zoekt|@back_compat_com_github_sourcegraph_zoekt|g' {} + +find . -type f -name "*.bazel" -exec $_sed_binary -i 's|@com_github_throttled_throttled_v2|@back_compat_com_github_throttled_throttled_v2|g' {} + """ def back_compat_defs(): @@ -114,6 +115,15 @@ def back_compat_defs(): version = "v0.0.0-20230620185637-63241cb1b17a", ) + go_repository( + name = "back_compat_com_github_throttled_throttled_v2", + build_file_proto_mode = "disable_global", + importpath = "github.com/throttled/throttled/v2", + sum = "h1:DOkCb1el7NYzRoPb1pyeHVghsUoonVWEjmo34vrcp/8=", + version = "v2.9.0", + ) + + # Now that we have declared a replacement for the two problematic go packages that # @sourcegraph_back_compat depends on, we can define the repository itself. Because it # comes with its Bazel rules (logical, that's just the current repository but with a different From 7bf85c2e5e85c77c0a86889ee92b8e8cc3d156b2 Mon Sep 17 00:00:00 2001 From: Keegan Carruthers-Smith Date: Tue, 11 Jul 2023 16:28:03 +0200 Subject: [PATCH 3/3] use RateLimitCtx and remove unused RateLimiter --- cmd/frontend/graphqlbackend/rate_limit.go | 38 ++----------------- .../graphqlbackend/rate_limit_test.go | 5 ++- cmd/frontend/internal/httpapi/api_test.go | 2 +- cmd/frontend/internal/httpapi/graphql.go | 2 +- 4 files changed, 9 insertions(+), 38 deletions(-) diff --git a/cmd/frontend/graphqlbackend/rate_limit.go b/cmd/frontend/graphqlbackend/rate_limit.go index b3bd7bf95dc3..a84eb765c821 100644 --- a/cmd/frontend/graphqlbackend/rate_limit.go +++ b/cmd/frontend/graphqlbackend/rate_limit.go @@ -1,6 +1,7 @@ package graphqlbackend import ( + "context" "strconv" "sync/atomic" @@ -369,7 +370,7 @@ type LimiterArgs struct { } type Limiter interface { - RateLimit(key string, quantity int, args LimiterArgs) (bool, throttled.RateLimitResult, error) + RateLimit(ctx context.Context, key string, quantity int, args LimiterArgs) (bool, throttled.RateLimitResult, error) } type LimitWatcher interface { @@ -434,40 +435,9 @@ type BasicLimiter struct { // RateLimit limits unauthenticated requests to the GraphQL API with an equal // quantity of 1. -func (bl *BasicLimiter) RateLimit(_ string, _ int, args LimiterArgs) (bool, throttled.RateLimitResult, error) { +func (bl *BasicLimiter) RateLimit(ctx context.Context, _ string, _ int, args LimiterArgs) (bool, throttled.RateLimitResult, error) { if args.Anonymous && args.RequestName == "unknown" && args.RequestSource == trace.SourceOther && bl.GCRARateLimiterCtx != nil { - return bl.GCRARateLimiterCtx.RateLimit("basic", 1) + return bl.GCRARateLimiterCtx.RateLimitCtx(ctx, "basic", 1) } return false, throttled.RateLimitResult{}, nil } - -type RateLimiter struct { - enabled bool - ipLimiter *throttled.GCRARateLimiterCtx - userLimiter *throttled.GCRARateLimiterCtx - overrides map[string]limiter -} - -func (rl *RateLimiter) RateLimit(uid string, cost int, args LimiterArgs) (bool, throttled.RateLimitResult, error) { - if r, ok := rl.overrides[uid]; ok { - return r.RateLimit(uid, cost) - } - if args.IsIP { - return rl.ipLimiter.RateLimit(uid, cost) - } - return rl.userLimiter.RateLimit(uid, cost) -} - -type limiter interface { - RateLimit(string, int) (bool, throttled.RateLimitResult, error) -} - -// fixedLimiter is a rate limiter that always returns the same result -type fixedLimiter struct { - limited bool - result throttled.RateLimitResult -} - -func (f *fixedLimiter) RateLimit(string, int) (bool, throttled.RateLimitResult, error) { - return f.limited, f.result, nil -} diff --git a/cmd/frontend/graphqlbackend/rate_limit_test.go b/cmd/frontend/graphqlbackend/rate_limit_test.go index fc47d1c6b235..f4803a5b7e7a 100644 --- a/cmd/frontend/graphqlbackend/rate_limit_test.go +++ b/cmd/frontend/graphqlbackend/rate_limit_test.go @@ -1,6 +1,7 @@ package graphqlbackend import ( + "context" "fmt" "testing" @@ -458,7 +459,7 @@ func TestBasicLimiter(t *testing.T) { } // 1st call should not be limited. - limited, _, err := limiter.RateLimit("", 1, limiterArgs) + limited, _, err := limiter.RateLimit(context.Background(), "", 1, limiterArgs) if err != nil { t.Fatal(err) } @@ -467,7 +468,7 @@ func TestBasicLimiter(t *testing.T) { } // 2nd call should be limited. - limited, _, err = limiter.RateLimit("", 1, limiterArgs) + limited, _, err = limiter.RateLimit(context.Background(), "", 1, limiterArgs) if err != nil { t.Fatal(err) } diff --git a/cmd/frontend/internal/httpapi/api_test.go b/cmd/frontend/internal/httpapi/api_test.go index 148211dc01ca..2606d340351b 100644 --- a/cmd/frontend/internal/httpapi/api_test.go +++ b/cmd/frontend/internal/httpapi/api_test.go @@ -24,7 +24,7 @@ func init() { func newTest(t *testing.T) *httptestutil.Client { logger := logtest.Scoped(t) enterpriseServices := enterprise.DefaultServices() - rateLimitStore, _ := memstore.New(1024) + rateLimitStore, _ := memstore.NewCtx(1024) rateLimiter := graphqlbackend.NewBasicLimitWatcher(logger, rateLimitStore) db := database.NewMockDB() diff --git a/cmd/frontend/internal/httpapi/graphql.go b/cmd/frontend/internal/httpapi/graphql.go index f197cbc187d0..dbc7860a33f5 100644 --- a/cmd/frontend/internal/httpapi/graphql.go +++ b/cmd/frontend/internal/httpapi/graphql.go @@ -92,7 +92,7 @@ func serveGraphQL(logger sglog.Logger, schema *graphql.Schema, rlw graphqlbacken traceData.cost = cost if rl, enabled := rlw.Get(); enabled && cost != nil { - limited, result, err := rl.RateLimit(uid, cost.FieldCount, graphqlbackend.LimiterArgs{ + limited, result, err := rl.RateLimit(r.Context(), uid, cost.FieldCount, graphqlbackend.LimiterArgs{ IsIP: isIP, Anonymous: anonymous, RequestName: requestName,