def test_can_change_max_rate_while_blocking(self): # This isn't a stress test, we just want to verify we can change # the rate at which we acquire a token. min_rate = 0.1 max_rate = 1 token_bucket = bucket.TokenBucket( min_rate=min_rate, max_rate=max_rate, clock=self.create_clock(), ) # First we'll set the max_rate to 0.1 (min_rate). This means that # it will take 10 seconds to accumulate a single token. We'll start # a thread and have it acquire() a token. # Then in the main thread we'll change the max_rate to something # really quick (e.g 100). We should immediately get a token back. # This is going to be timing sensitive, but we can verify that # as long as it doesn't take 10 seconds to get a token, we were # able to update the rate as needed. thread = threading.Thread(target=token_bucket.acquire) token_bucket.max_rate = min_rate start_time = time.time() thread.start() # This shouldn't block the main thread. token_bucket.max_rate = 100 thread.join() end_time = time.time() self.assertLessEqual(end_time - start_time, 1.0 / min_rate)
def register_retry_handler(client): clock = bucket.Clock() rate_adjustor = throttling.CubicCalculator(starting_max_rate=0, start_time=clock.current_time()) token_bucket = bucket.TokenBucket(max_rate=1, clock=clock) rate_clocker = RateClocker(clock) throttling_detector = standard.ThrottlingErrorDetector( retry_event_adapter=standard.RetryEventAdapter(), ) limiter = ClientRateLimiter( rate_adjustor=rate_adjustor, rate_clocker=rate_clocker, token_bucket=token_bucket, throttling_detector=throttling_detector, clock=clock, ) client.meta.events.register( 'before-send', limiter.on_sending_request, ) client.meta.events.register( 'needs-retry', limiter.on_receiving_response, ) return limiter
def create_token_bucket(self, max_rate=10, min_rate=0.1): return bucket.TokenBucket(max_rate=max_rate, clock=self.clock, min_rate=min_rate)