func setUpRateLimiting(
in gcs.Bucket,
opRateLimitHz float64,
egressBandwidthLimit float64) (out gcs.Bucket, err error) {
// If no rate limiting has been requested, just return the bucket.
if !(opRateLimitHz > 0 || egressBandwidthLimit > 0) {
out = in
return
}
// Treat a disabled limit as a very large one.
if !(opRateLimitHz > 0) {
opRateLimitHz = 1e15
}
if !(egressBandwidthLimit > 0) {
egressBandwidthLimit = 1e15
}
// Choose token bucket capacities.
const window = 30 * time.Second
opCapacity, err := ratelimit.ChooseTokenBucketCapacity(
opRateLimitHz,
window)
if err != nil {
err = fmt.Errorf("Choosing operation token bucket capacity: %v", err)
return
}
egressCapacity, err := ratelimit.ChooseTokenBucketCapacity(
egressBandwidthLimit,
window)
if err != nil {
err = fmt.Errorf("Choosing egress bandwidth token bucket capacity: %v", err)
return
}
// Create the throttles.
opThrottle := ratelimit.NewThrottle(opRateLimitHz, opCapacity)
egressThrottle := ratelimit.NewThrottle(egressBandwidthLimit, egressCapacity)
// And the bucket.
out = ratelimit.NewThrottledBucket(
opThrottle,
egressThrottle,
in)
return
}
func (t *ThrottleTest) IntegrationTest() {
runtime.GOMAXPROCS(runtime.NumCPU())
const perCaseDuration = 1 * time.Second
// Set up several test cases where we have N goroutines simulating arrival of
// packets at a given rate, asking a token bucket when to admit them.
testCases := []struct {
numActors int
arrivalRateHz float64
limitRateHz float64
}{
// Single actor
{1, 150, 200},
{1, 200, 200},
{1, 250, 200},
// Multiple actors
{4, 150, 200},
{4, 200, 200},
{4, 250, 200},
}
// Run each test case.
for i, tc := range testCases {
// Create a throttle.
capacity, err := ratelimit.ChooseTokenBucketCapacity(
tc.limitRateHz,
perCaseDuration)
AssertEq(nil, err)
throttle := ratelimit.NewThrottle(tc.limitRateHz, capacity)
// Start workers.
var wg sync.WaitGroup
var totalProcessed uint64
ctx, _ := context.WithDeadline(
context.Background(),
time.Now().Add(perCaseDuration))
for i := 0; i < tc.numActors; i++ {
wg.Add(1)
go func() {
defer wg.Done()
processed := processArrivals(
ctx,
throttle,
tc.arrivalRateHz/float64(tc.numActors),
perCaseDuration)
atomic.AddUint64(&totalProcessed, processed)
}()
}
// Wait for them all to finish.
wg.Wait()
// We should have processed about the correct number of arrivals.
smallerRateHz := tc.arrivalRateHz
if smallerRateHz > tc.limitRateHz {
smallerRateHz = tc.limitRateHz
}
expected := smallerRateHz * (float64(perCaseDuration) / float64(time.Second))
ExpectThat(
totalProcessed,
AllOf(
GreaterThan(expected*0.90),
LessThan(expected*1.10)),
"Test case %d. expected: %f",
i,
expected)
}
}