async def test_token_bucket_refills_itself(): CAPACITY = 50 TOKENS_PER_SECOND = 1000 bucket = TokenBucket(TOKENS_PER_SECOND, CAPACITY) # consume all of the tokens for _ in range(CAPACITY): await bucket.take() # enough time for the bucket to fully refill start_at = time.perf_counter() time_to_refill = CAPACITY / TOKENS_PER_SECOND while time.perf_counter() - start_at < time_to_refill: await asyncio.sleep(time_to_refill) # This should take roughly zero time start_at = time.perf_counter() for _ in range(CAPACITY): await bucket.take() end_at = time.perf_counter() delta = end_at - start_at await assert_close_to_zero(delta, CAPACITY)
async def measure_zero(iterations): bucket = TokenBucket(1, iterations) start_at = time.perf_counter() for _ in range(iterations): await bucket.take() end_at = time.perf_counter() return end_at - start_at
def __init__(self, chain: BaseAsyncChain, db: BaseAsyncChainDB, peer_pool: ETHPeerPool, header_syncer: HeaderSyncerAPI, block_importer: BaseBlockImporter, token: CancelToken = None) -> None: super().__init__(chain, db, peer_pool, header_syncer, token) # track when block bodies are downloaded, so that blocks can be imported self._block_import_tracker = OrderedTaskPreparation( BlockImportPrereqs, id_extractor=attrgetter('hash'), # make sure that a block is not imported until the parent block is imported dependency_extractor=attrgetter('parent_hash'), ) self._block_importer = block_importer # Track if any headers have been received yet self._got_first_header = asyncio.Event() # Rate limit the block import logs self._import_log_limiter = TokenBucket( 0.33, # show about one log per 3 seconds 5, # burst up to 5 logs after a lag )
def __init__(self, chain: AsyncChainAPI, db: BaseAsyncChainDB, peer_pool: ETHPeerPool, header_syncer: HeaderSyncerAPI, block_importer: BaseBlockImporter, token: CancelToken = None) -> None: super().__init__(chain, db, peer_pool, header_syncer, token) # track when block bodies are downloaded, so that blocks can be imported self._block_import_tracker = OrderedTaskPreparation( BlockImportPrereqs, id_extractor=attrgetter('hash'), # make sure that a block is not imported until the parent block is imported dependency_extractor=attrgetter('parent_hash'), # Avoid problems by keeping twice as much data as the import queue size max_depth=BLOCK_IMPORT_QUEUE_SIZE * 2, ) self._block_importer = block_importer # Track if any headers have been received yet self._got_first_header = asyncio.Event() # Rate limit the block import logs self._import_log_limiter = TokenBucket( 0.33, # show about one log per 3 seconds 5, # burst up to 5 logs after a lag ) # the queue of blocks that are downloaded and ready to be imported self._import_queue: 'asyncio.Queue[BlockAPI]' = asyncio.Queue(BLOCK_IMPORT_QUEUE_SIZE) self._import_active = asyncio.Lock()
async def maybe_connect_more_peers(self) -> None: rate_limiter = TokenBucket( rate=1 / PEER_CONNECT_INTERVAL, capacity=MAX_SEQUENTIAL_PEER_CONNECT, ) while self.is_operational: if self.is_full: await self.sleep(PEER_CONNECT_INTERVAL) continue await self.wait(rate_limiter.take()) try: await asyncio.gather(*(self._add_peers_from_backend(backend) for backend in self.peer_backends)) except OperationCancelled: break except asyncio.CancelledError: # no need to log this exception, this is expected raise except Exception: self.logger.exception( "unexpected error during peer connection") # Continue trying to connect to peers, even if there was a # surprising failure during one of the attempts. continue
def test_token_bucket_take_nowait(): bucket = TokenBucket(1, 10) assert bucket.can_take(10) bucket.take_nowait(10) assert not bucket.can_take(1) with pytest.raises(NotEnoughTokens): bucket.take_nowait(1)
async def test_token_bucket_can_take(): bucket = TokenBucket(1, 10) assert bucket.can_take() is True # can take 1 assert bucket.can_take( bucket.get_num_tokens()) is True # can take full capacity await bucket.take(10) # empty the bucket assert bucket.can_take() is False
def __init__(self, peer: BasePeer, response_msg_type: Type[CommandAPI], token: CancelToken) -> None: super().__init__(token) self._peer = peer self.response_msg_type = response_msg_type self._lock = asyncio.Lock() # token bucket for limiting timeouts. # - Refills at 1-token every 5 minutes # - Max capacity of 3 tokens self.timeout_bucket = TokenBucket(TIMEOUT_BUCKET_RATE, TIMEOUT_BUCKET_CAPACITY)
async def test_token_bucket_get_num_tokens(): bucket = TokenBucket(1, 10) # starts at full capacity assert bucket.get_num_tokens() == 10 await bucket.take(5) assert 5 <= bucket.get_num_tokens() <= 5.1 await bucket.take(bucket.get_num_tokens()) assert 0 <= bucket.get_num_tokens() <= 0.1
async def test_token_bucket_initial_tokens(): CAPACITY = 10 bucket = TokenBucket(1000, CAPACITY) start_at = time.perf_counter() for _ in range(CAPACITY): await bucket.take() end_at = time.perf_counter() delta = end_at - start_at await assert_close_to_zero(delta, CAPACITY)
async def maybe_connect_more_peers(self) -> None: rate_limiter = TokenBucket( rate=1 / PEER_CONNECT_INTERVAL, capacity=MAX_SEQUENTIAL_PEER_CONNECT, ) # We set this to 0 so that upon startup (when our RoutingTable will have only a few # entries) we use the less restrictive filter function and get as many connection # candidates as possible. last_candidates_count = 0 while self.manager.is_running: if self.is_full: await asyncio.sleep(PEER_CONNECT_INTERVAL) continue await rate_limiter.take() if last_candidates_count >= self.available_slots: head = await self.get_chain_head() genesis_hash = await self.get_genesis_hash() fork_blocks = extract_fork_blocks(self.vm_configuration) should_skip = functools.partial( skip_candidate_if_on_list_or_fork_mismatch, genesis_hash, head.block_number, fork_blocks, ) else: self.logger.debug( "Didn't get enough candidates last time, falling back to skipping " "only peers that are blacklisted or already connected to") should_skip = skip_candidate_if_on_list # type: ignore try: candidate_counts = await asyncio.gather( *(self._add_peers_from_backend(backend, should_skip) for backend in self.peer_backends)) last_candidates_count = sum(candidate_counts) except OperationCancelled: # FIXME: We may no longer need this; need to confirm that none of the tasks we # create use BaseService. break except asyncio.CancelledError: # no need to log this exception, this is expected raise except Exception: self.logger.exception( "unexpected error during peer connection") # Continue trying to connect to peers, even if there was a # surprising failure during one of the attempts. continue
def __init__(self, peer: BasePeer, listening_for: Type[CommandAPI], cancel_token: CancelToken) -> None: self._peer = peer self._cancel_token = cancel_token self._response_command_type = listening_for # This TokenBucket allows for the occasional invalid response at a # maximum rate of 1-per-10-minutes and allowing up to two in quick # succession. We *allow* invalid responses because the ETH protocol # doesn't have strong correlation between request/response and certain # networking conditions can result in us interpreting a legitimate # message as an invalid response if messages arrive out of order or # late. self._invalid_response_bucket = TokenBucket(1 / 600, 2)
async def test_token_bucket_initial_tokens(): bucket = TokenBucket(1000, 10) start_at = time.perf_counter() for _ in range(10): await bucket.take() end_at = time.perf_counter() delta = end_at - start_at # since the bucket starts out full the loop # should take near zero time expected = await measure_zero(10) # drift is allowed to be up to 1000% since we're working with very small # numbers. assert_fuzzy_equal(delta, expected, allowed_drift=10)
async def maybe_connect_more_peers(self) -> None: rate_limiter = TokenBucket( rate=1 / PEER_CONNECT_INTERVAL, capacity=MAX_SEQUENTIAL_PEER_CONNECT, ) while self.is_operational: if self.is_full: await self.sleep(PEER_CONNECT_INTERVAL) continue await self.wait(rate_limiter.take()) await self.wait( asyncio.gather(*(self._add_peers_from_backend(backend) for backend in self.peer_backends)))
async def test_token_bucket_hits_limit(): bucket = TokenBucket(1000, 10) bucket.take_nowait(10) start_at = time.perf_counter() # first 10 tokens should be roughly instant # next 10 tokens should each take 1/1000th second each to generate. while True: if bucket.can_take(10): break else: await asyncio.sleep(0) end_at = time.perf_counter() # we use a zero-measure of 20 to account for the loop overhead. zero = await measure_zero(10) expected_delta = 10 / 1000 + zero delta = end_at - start_at # allow up to 10% difference in expected time assert_fuzzy_equal(delta, expected_delta, allowed_drift=0.1)
async def maybe_connect_more_peers(self) -> None: rate_limiter = TokenBucket( rate=1 / PEER_CONNECT_INTERVAL, capacity=MAX_SEQUENTIAL_PEER_CONNECT, ) # We set this to 0 so that upon startup (when our RoutingTable will have only a few # entries) we use the less restrictive filter function and get as many connection # candidates as possible. last_candidates_count = 0 while self.manager.is_running: if self.is_full: await asyncio.sleep(PEER_CONNECT_INTERVAL) continue await rate_limiter.take() if last_candidates_count >= self.available_slots: head = await self.get_chain_head() genesis_hash = await self.get_genesis_hash() fork_blocks = extract_fork_blocks(self.vm_configuration) should_skip = functools.partial( skip_candidate_if_on_list_or_fork_mismatch, genesis_hash, head.block_number, fork_blocks, ) else: self.logger.debug( "Didn't get enough candidates last time, falling back to skipping " "only peers that are blacklisted or already connected to") should_skip = skip_candidate_if_on_list # type: ignore candidate_counts = await asyncio.gather(*( self._add_peers_from_backend(backend, should_skip) for backend in self.peer_backends )) last_candidates_count = sum(candidate_counts)
async def test_token_bucket_hits_limit(): CAPACITY = 50 TOKENS_PER_SECOND = 1000 bucket = TokenBucket(TOKENS_PER_SECOND, CAPACITY) bucket.take_nowait(CAPACITY) start_at = time.perf_counter() # first CAPACITY tokens should be roughly instant # next CAPACITY tokens should each take 1/TOKENS_PER_SECOND second each to generate. while True: if bucket.can_take(CAPACITY): break else: await asyncio.sleep(0) end_at = time.perf_counter() # we use a zero-measure of CAPACITY loops to account for the loop overhead. zero = await measure_zero(CAPACITY) expected_delta = CAPACITY / TOKENS_PER_SECOND + zero delta = end_at - start_at # allow up to 10% difference in expected time assert_fuzzy_equal(delta, expected_delta, allowed_drift=0.1)
async def test_token_bucket_refills_itself(): bucket = TokenBucket(1000, 10) # consume all of the tokens for _ in range(10): await bucket.take() # enough time for the bucket to fully refill await asyncio.sleep(20 / 1000) start_at = time.perf_counter() for _ in range(10): await bucket.take() end_at = time.perf_counter() delta = end_at - start_at # since the capacity should have been fully refilled, second loop time # should take near zero time expected = await measure_zero(10) # drift is allowed to be up to 300% since we're working with very small # numbers, and the performance in CI varies widely. assert_fuzzy_equal(delta, expected, allowed_drift=3)