-
Notifications
You must be signed in to change notification settings - Fork 426
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Removing rate limit dependency (#1253)
* local changes * local changes * local changes * local changes * removed token bucket and added unit tests * fixing lint * small fix- renaming * small fix- renaming * fix lint * adding licence * testing * back to changes * fixing comments * removing throttle_test * adding throttle test * fixing comment * fixing lint * fixed comment * lint tests * lint tests * fixing comment * empty commit * empty commit
- Loading branch information
Showing
10 changed files
with
1,058 additions
and
10 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,83 @@ | ||
| // Copyright 2023 Google Inc. All Rights Reserved. | ||
| // | ||
| // Licensed under the Apache License, Version 2.0 (the "License"); | ||
| // you may not use this file except in compliance with the License. | ||
| // You may obtain a copy of the License at | ||
| // | ||
| // http://www.apache.org/licenses/LICENSE-2.0 | ||
| // | ||
| // Unless required by applicable law or agreed to in writing, software | ||
| // distributed under the License is distributed on an "AS IS" BASIS, | ||
| // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||
| // See the License for the specific language governing permissions and | ||
| // limitations under the License. | ||
|
|
||
| package ratelimit | ||
|
|
||
| import ( | ||
| "fmt" | ||
| "math" | ||
| "time" | ||
| ) | ||
|
|
||
| // Choose a limiter capacity that ensures that the action gated by the | ||
| // limiter will be limited to within a few percent of `rateHz * window` | ||
| // for any window of the given size. | ||
| // | ||
| // This is not be possible for all rates and windows. In that case, an error | ||
| // will be returned. | ||
| func ChooseLimiterCapacity( | ||
| rateHz float64, | ||
| window time.Duration) (capacity uint64, err error) { | ||
| // Check that the input is reasonable. | ||
| if rateHz <= 0 || math.IsInf(rateHz, 0) { | ||
| err = fmt.Errorf("Illegal rate: %f", rateHz) | ||
| return | ||
| } | ||
|
|
||
| if window <= 0 { | ||
| err = fmt.Errorf("Illegal window: %v", window) | ||
| return | ||
| } | ||
|
|
||
| // We cannot help but allow the rate to exceed the configured maximum by some | ||
| // factor in an arbitrary window, no matter how small we scale the max | ||
| // accumulated credit -- the bucket may be full at the start of the window, | ||
| // be immediately exhausted, then be repeatedly exhausted just before filling | ||
| // throughout the window. | ||
| // | ||
| // For example: let the window W = 10 seconds, and the bandwidth B = 20 MiB/s. | ||
| // Set the max accumulated credit C = W*B/2 = 100 MiB. Then this | ||
| // sequence of events is allowed: | ||
| // | ||
| // * T=0: Allow through 100 MiB. | ||
| // * T=4.999999: Allow through nearly 100 MiB. | ||
| // * T=9.999999: Allow through nearly 100 MiB. | ||
| // | ||
| // Above we allow through nearly 300 MiB, exceeding the allowed bytes for the | ||
| // window by nearly 50%. Note however that this trend cannot continue into | ||
| // the next window, so this must be a transient spike. | ||
| // | ||
| // In general if we set C <= W*B/N, then we're off by no more than a factor | ||
| // of (N+1)/N within any window of size W. | ||
| // | ||
| // Choose a reasonable N. | ||
| const N = 50 // At most 2% error | ||
|
|
||
| w := float64(window) / float64(time.Second) | ||
| capacityFloat := math.Floor(w * rateHz / N) | ||
| if !(capacityFloat >= 1 && capacityFloat < float64(math.MaxUint64)) { | ||
| err = fmt.Errorf( | ||
| "Can't use a token bucket to limit to %f Hz over a window of %v "+ | ||
| "(result is a capacity of %f)", | ||
| rateHz, | ||
| window, | ||
| capacityFloat) | ||
|
|
||
| return | ||
| } | ||
|
|
||
| capacity = uint64(capacityFloat) | ||
|
|
||
| return | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,96 @@ | ||
| // Copyright 2023 Google Inc. All Rights Reserved. | ||
| // | ||
| // Licensed under the Apache License, Version 2.0 (the "License"); | ||
| // you may not use this file except in compliance with the License. | ||
| // You may obtain a copy of the License at | ||
| // | ||
| // http://www.apache.org/licenses/LICENSE-2.0 | ||
| // | ||
| // Unless required by applicable law or agreed to in writing, software | ||
| // distributed under the License is distributed on an "AS IS" BASIS, | ||
| // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||
| // See the License for the specific language governing permissions and | ||
| // limitations under the License. | ||
|
|
||
| package ratelimit | ||
|
|
||
| import ( | ||
| "fmt" | ||
| "testing" | ||
| "time" | ||
|
|
||
| . "github.com/jacobsa/ogletest" | ||
| ) | ||
|
|
||
| func TestLimiterCapacity(t *testing.T) { RunTests(t) } | ||
|
|
||
| //////////////////////////////////////////////////////////////////////// | ||
| // Boilerplate | ||
| //////////////////////////////////////////////////////////////////////// | ||
|
|
||
| type LimiterCapacityTest struct { | ||
| } | ||
|
|
||
| func init() { RegisterTestSuite(&LimiterCapacityTest{}) } | ||
|
|
||
| func rateLessThanOrEqualToZero(rate float64) { | ||
| _, err := ChooseLimiterCapacity(rate, 30) | ||
|
|
||
| expectedError := fmt.Errorf("Illegal rate: %f", rate) | ||
|
|
||
| AssertEq(expectedError.Error(), err.Error()) | ||
| } | ||
|
|
||
| func (t *LimiterCapacityTest) TestRateLessThanZero() { | ||
| var negativeRateHz float64 = -1 | ||
|
|
||
| rateLessThanOrEqualToZero(negativeRateHz) | ||
| } | ||
|
|
||
| func (t *LimiterCapacityTest) TestRateEqualToZero() { | ||
| var zeroRateHz float64 = 0 | ||
|
|
||
| rateLessThanOrEqualToZero(zeroRateHz) | ||
| } | ||
|
|
||
| func windowLessThanOrEqualToZero(window time.Duration) { | ||
| _, err := ChooseLimiterCapacity(1, window) | ||
|
|
||
| expectedError := fmt.Errorf("Illegal window: %v", window) | ||
|
|
||
| AssertEq(expectedError.Error(), err.Error()) | ||
| } | ||
|
|
||
| func (t *LimiterCapacityTest) TestWindowLessThanZero() { | ||
| var negativeWindow time.Duration = -1 | ||
|
|
||
| windowLessThanOrEqualToZero(negativeWindow) | ||
| } | ||
|
|
||
| func (t *LimiterCapacityTest) TestWindowEqualToZero() { | ||
| var zeroWindow time.Duration = 0 | ||
|
|
||
| windowLessThanOrEqualToZero(zeroWindow) | ||
| } | ||
|
|
||
| func (t *LimiterCapacityTest) TestCapacityEqualToZero() { | ||
| var rate = 0.5 | ||
| var window time.Duration = 1 | ||
|
|
||
| capacity, err := ChooseLimiterCapacity(rate, window) | ||
|
|
||
| expectedError := fmt.Errorf( | ||
| "Can't use a token bucket to limit to %f Hz over a window of %v (result is a capacity of %f)", rate, window, float64(capacity)) | ||
| AssertEq(expectedError.Error(), err.Error()) | ||
| } | ||
|
|
||
| func (t *LimiterCapacityTest) TestExpectedCapacity() { | ||
| var rate float64 = 20 | ||
| var window = 10 * time.Second | ||
|
|
||
| capacity, err := ChooseLimiterCapacity(rate, window) | ||
| // capacity = floor((20.0 * 10)/50) = floor(4.0) = 4 | ||
|
|
||
| ExpectEq(nil, err) | ||
| ExpectEq(4, capacity) | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,58 @@ | ||
| // Copyright 2023 Google Inc. All Rights Reserved. | ||
| // | ||
| // Licensed under the Apache License, Version 2.0 (the "License"); | ||
| // you may not use this file except in compliance with the License. | ||
| // You may obtain a copy of the License at | ||
| // | ||
| // http://www.apache.org/licenses/LICENSE-2.0 | ||
| // | ||
| // Unless required by applicable law or agreed to in writing, software | ||
| // distributed under the License is distributed on an "AS IS" BASIS, | ||
| // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||
| // See the License for the specific language governing permissions and | ||
| // limitations under the License. | ||
|
|
||
| package ratelimit | ||
|
|
||
| import ( | ||
| "golang.org/x/net/context" | ||
| "golang.org/x/time/rate" | ||
| ) | ||
|
|
||
| // A simple interface for limiting the rate of some event. Unlike TokenBucket, | ||
| // does not allow the user control over what time means. | ||
| // | ||
| // Safe for concurrent access. | ||
| type Throttle interface { | ||
| // Return the maximum number of tokens that can be requested in a call to | ||
| // Wait. | ||
| Capacity() (c uint64) | ||
|
|
||
| // Acquire the given number of tokens from the underlying token bucket, then | ||
| // sleep until when it says to wake. If the context is cancelled before then, | ||
| // return early with an error. | ||
| // | ||
| // REQUIRES: tokens <= capacity | ||
| Wait(ctx context.Context, tokens uint64) (err error) | ||
| } | ||
|
|
||
| type limiter struct { | ||
| *rate.Limiter | ||
| } | ||
|
|
||
| func NewThrottle( | ||
| rateHz float64, | ||
| capacity uint64) (t Throttle) { | ||
| t = &limiter{rate.NewLimiter(rate.Limit(rateHz), int(capacity))} | ||
| return | ||
| } | ||
|
|
||
| func (l *limiter) Capacity() (c uint64) { | ||
| return uint64(l.Burst()) | ||
| } | ||
|
|
||
| func (l *limiter) Wait( | ||
| ctx context.Context, | ||
| tokens uint64) (err error) { | ||
| return l.WaitN(ctx, int(tokens)) | ||
| } |
Oops, something went wrong.