Ejemplo n.º 1
0
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, targeting only a few percent error in each
	// window of the given size.
	const window = 8 * time.Hour

	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
}
Ejemplo n.º 2
0
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)
	}
}