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)