import asyncio import logging import random import time import typing import requests from pytoniq_core import BlockIdExt, Block, Address, Account, ShardAccount, SimpleAccount, ShardDescr, Transaction, Cell from pytoniq_core.tlb.block import BinTree from .client import LiteClient, LiteClientError, LiteServerError class BalancerError(LiteClientError): pass class LiteBalancer: def __init__(self, peers: typing.List[LiteClient], timeout: int = 10): self._peers = peers self._alive_peers: typing.Set[int] = set() self._archival_peers = set() self._checker: asyncio.Task = None self._mc_blocks = {} # {index: masterchain_seqno} self._av_resp_time = {} # {index: average_response_time} self._total_req_num = {} # {index: successful_requests_num} self._current_req_num = {} # {index: current_waiting_requests_num} self._logger = logging.getLogger(self.__class__.__name__) self.inited = False self.max_req_per_peer = 100 self.max_retries = 1 self.timeout = timeout @property def peers_num(self): return len(self._peers) @property def alive_peers_num(self): return len(self._alive_peers) @property def archival_peers_num(self): return len(self._archival_peers) @property def last_mc_block(self): seqno = self._find_consensus_block() for p in self._peers: if p.last_mc_block is not None and p.last_mc_block.seqno == seqno: return p.last_mc_block return None def set_max_retries(self, retries_num: int) -> None: self.max_retries = retries_num async def start_up(self): have_blockstore = False tasks = [] for client in self._peers: if client.trust_level >= 1: tasks.append(self._connect_to_peer(client)) elif have_blockstore: tasks.append(self._connect_to_peer(client)) else: # so we can verify blocks proof link only once connected = await self._connect_to_peer(client) if connected: have_blockstore = True async def f(): return connected tasks.append(f()) result = await asyncio.gather(*tasks) for i, client in enumerate(self._peers): if result[i]: self._alive_peers.add(i) await self._find_archives() self._checker = asyncio.create_task(self._check_peers()) self._delete_unsync_peers() self.inited = True async def _find_archives(self): tasks = [] inds = [] for i in self._alive_peers: tasks.append(self.check_archive(self._peers[i])) inds.append(i) result = await asyncio.gather(*tasks, return_exceptions=True) for i, r in enumerate(result): if isinstance(r, bool): if r: self._archival_peers.add(inds[i]) else: self._archival_peers.discard(inds[i]) # almost impossible case when peer becomes unarchival elif isinstance(r, Exception): self._logger.info(f'Failed to check peer {inds[i]} on archival: {r}') @staticmethod async def check_archive(peer: LiteClient): try: blk, _ = await peer.lookup_block(wc=-1, shard=-2**63, seqno=random.randint(1, 1024)) # as ton-http-api does, but maybe need to ask the first block. todo return True except LiteClientError: return False async def _connect_to_peer(self, client: LiteClient): self._check_errors(client) if client.listener is not None and not client.listener.done(): client.listener.cancel() while not client.listener.done(): await asyncio.sleep(0) try: if client.trust_level >= 1: await asyncio.wait_for(client.connect(), 3) else: await client.connect() return True except asyncio.TimeoutError: return False except Exception as e: self._logger.debug(f'Failed to connect to the peer {client.server.get_key_id().hex()}: {e}') return False finally: self._check_errors(client) async def _ping_peer(self, peer: LiteClient): try: await asyncio.wait_for(peer.get_masterchain_info(), 3) return True except asyncio.TimeoutError: return False except Exception as e: self._logger.debug(f'Failed to ping peer {peer.server.get_key_id().hex()}: {e}') def _check_errors(self, client: LiteClient): for task in [client.updater, client.pinger, client.listener]: if task is not None and task.done(): if not task.cancelled(): self._logger.debug(f'client task {task} failed with exception: {task.exception()}') return True return False async def _check_peers(self): while True: await asyncio.sleep(3) for i, client in enumerate(self._peers): self._delete_unsync_peers() client: LiteClient if client.inited: if self._check_errors(client): self._alive_peers.discard(i) await client.close() if await self._connect_to_peer(client): self._alive_peers.add(i) else: self._alive_peers.discard(i) continue ping_res = await self._ping_peer(client) if ping_res: self._alive_peers.add(i) else: self._alive_peers.discard(i) else: if await self._connect_to_peer(client): self._alive_peers.add(i) else: self._alive_peers.discard(i) async def connect(self): raise BalancerError(f'Use start_up()') def _build_priority_list(self, only_archive: bool = False): sorted_peers = sorted( list(self._alive_peers) if not only_archive else list(self._archival_peers), key=lambda e: (self._mc_blocks.get(e, 0), -self._av_resp_time.get(e, self.timeout * 1000)), reverse=True ) # first peers are with biggest masterchain seqno and lowest avg time response return sorted_peers def _choose_peer(self, only_archive: bool = False): peers = self._build_priority_list(only_archive) min_req = float('inf') for p in peers: peer_req = self._current_req_num.get(p, 0) if peer_req <= self.max_req_per_peer: return p if peer_req < min_req: min_req = peer_req for p in peers: peer_req = self._current_req_num.get(p, 0) if peer_req <= min_req: return p return peers[0] # should never happen @staticmethod def _calc_new_average(old_average: int, n: int, new_value: int): return (old_average * (n - 1) + new_value) / n def _update_average_request_time(self, ls_index: int, req_time: int): old = self._av_resp_time.get(ls_index, 0) req_num = self._total_req_num.get(ls_index, 0) req_num += 1 self._av_resp_time[ls_index] = self._calc_new_average(old, req_num, req_time) self._total_req_num[ls_index] = req_num def _update_mc_seqno(self, ls_index: int): client = self._peers[ls_index] blk = client.last_mc_block if blk and self._mc_blocks.get(ls_index, 0) < blk.seqno: self._mc_blocks[ls_index] = blk.seqno def _update_mc_seqnos(self): for i in range(len(self._peers)): self._update_mc_seqno(i) def _find_consensus_block(self): self._update_mc_seqnos() seqnos = sorted(self._mc_blocks.values(), reverse=True) if not seqnos: return 0 return seqnos[len(seqnos) * 2 // 3] # block that knows at least 2/3 liteservers def _delete_unsync_peers(self): cons_block = self._find_consensus_block() for i in list(self._alive_peers): if self._mc_blocks.get(i, 0) < cons_block: self._alive_peers.discard(i) async def execute_method(self, method_name_: str, *args, **kwargs) -> typing.Union[dict, typing.Any]: only_archive = kwargs.pop('only_archive', False) choose_random = kwargs.pop('choose_random', False) retry = False i = 0 while i < self.max_retries or retry: retry = False if not len(self._alive_peers): raise BalancerError(f'have no alive peers') if only_archive and choose_random: raise BalancerError('Currently you cant execute method for both random and archive peer') if only_archive and not len(self._archival_peers): await self._find_archives() # give one more chance to find if not len(self._archival_peers): raise BalancerError(f'have no alive archive peers') self._update_mc_seqnos() if choose_random: ind = random.choice(list(self._alive_peers)) else: ind = self._choose_peer(only_archive) peer: LiteClient = self._peers[ind] peer_meth = getattr(peer, method_name_, None) self._current_req_num[ind] = self._current_req_num.get(ind, 0) + 1 s = time.time_ns() if not peer_meth: raise BalancerError('Unknown method for peer') try: resp = await peer_meth(*args, **kwargs) self._update_average_request_time(ind, (time.time_ns() - s) // 10**6) # provide milliseconds return resp except asyncio.TimeoutError: self._update_average_request_time(ind, self.timeout * 10**6) # provide milliseconds self._alive_peers.discard(ind) continue except LiteServerError as e: if e.message == 'timeout': self._update_average_request_time(ind, self.timeout * 10 ** 6) self._alive_peers.discard(ind) continue raise e except ConnectionError: # if socket is dead we just try another peer and somewhere in future will reconnect to this one self._alive_peers.discard(ind) retry = True continue finally: self._current_req_num[ind] -= 1 raise asyncio.TimeoutError() @staticmethod def _get_args(locals_: dict): a = locals_.copy() for k in list(a.keys()): if k.startswith('_'): a.pop(k) a.pop('self') kwargs = a.pop('kwargs', {}) a |= kwargs return a """CODE BELOW IS AUTOGENERATED. DO NOT EDIT MANUALLY""" async def get_masterchain_info(self, **kwargs): return await self.execute_method('get_masterchain_info', **self._get_args(locals())) async def raw_wait_masterchain_seqno(self, seqno: int, timeout_ms: int, suffix: bytes = b'', **kwargs): return await self.execute_method('raw_wait_masterchain_seqno', **self._get_args(locals())) async def wait_masterchain_seqno(self, seqno: int, timeout_ms: int, schema_name: str, data: dict = None, **kwargs): return await self.execute_method('wait_masterchain_seqno', **self._get_args(locals())) async def get_masterchain_info_ext(self, **kwargs): return await self.execute_method('get_masterchain_info_ext', **self._get_args(locals())) async def get_time(self, **kwargs): return await self.execute_method('get_time', **self._get_args(locals())) async def get_version(self, **kwargs): return await self.execute_method('get_version', **self._get_args(locals())) async def get_state(self, wc: int, shard: typing.Optional[int], seqno: int, root_hash: typing.Union[str, bytes], file_hash: typing.Union[str, bytes] , **kwargs) -> dict: return await self.execute_method('get_state', **self._get_args(locals())) async def raw_get_block_header(self, block: BlockIdExt, **kwargs) -> Block: return await self.execute_method('raw_get_block_header', **self._get_args(locals())) async def get_block_header(self, wc: int, shard: typing.Optional[int], seqno: int, root_hash: typing.Union[str, bytes], file_hash: typing.Union[str, bytes] , **kwargs) -> Block: return await self.execute_method('get_block_header', **self._get_args(locals())) async def lookup_block(self, wc: int, shard: int, seqno: int = -1, lt: typing.Optional[int] = None, utime: typing.Optional[int] = None, **kwargs) -> typing.Tuple[BlockIdExt, Block]: return await self.execute_method('lookup_block', **self._get_args(locals())) async def raw_get_block(self, block: BlockIdExt, **kwargs) -> Block: return await self.execute_method('raw_get_block', **self._get_args(locals())) async def get_block(self, wc: int, shard: typing.Optional[int], seqno: int, root_hash: typing.Union[str, bytes], file_hash: typing.Union[str, bytes], **kwargs) -> Block: return await self.execute_method('get_block', **self._get_args(locals())) async def raw_get_account_state(self, address: typing.Union[str, Address], block: typing.Optional[BlockIdExt] = None , **kwargs) -> typing.Tuple[typing.Optional[Account], typing.Optional[ShardAccount]]: return await self.execute_method('raw_get_account_state', **self._get_args(locals())) async def get_account_state(self, address: typing.Union[str, Address], **kwargs) -> SimpleAccount: return await self.execute_method('get_account_state', **self._get_args(locals())) async def run_get_method(self, address: typing.Union[Address, str], method: typing.Union[int, str], stack: list, block: BlockIdExt = None , **kwargs) -> list: return await self.execute_method('run_get_method', **self._get_args(locals())) async def run_get_method_remote(self, address: typing.Union[Address, str], method: typing.Union[int, str], stack: list, block: BlockIdExt = None , **kwargs) -> list: return await self.execute_method('run_get_method_remote', **self._get_args(locals())) async def run_get_method_local(self, address: typing.Union[Address, str], method: typing.Union[int, str], stack: list, block: BlockIdExt = None, gas_limit: int = 300000, **kwargs) -> list: return await self.execute_method('run_get_method_local', **self._get_args(locals())) async def raw_get_shard_info(self, block: typing.Optional[BlockIdExt] = None, wc: int = 0, shard: int = -9223372036854775808, exact: bool = True , **kwargs) -> ShardDescr: return await self.execute_method('raw_get_shard_info', **self._get_args(locals())) async def raw_get_all_shards_info(self, block: typing.Optional[BlockIdExt] = None, **kwargs) -> typing.Dict[int, BinTree]: return await self.execute_method('raw_get_all_shards_info', **self._get_args(locals())) async def get_all_shards_info(self, block: typing.Optional[BlockIdExt] = None, **kwargs) -> typing.List[BlockIdExt]: return await self.execute_method('get_all_shards_info', **self._get_args(locals())) async def get_one_transaction(self, address: typing.Union[Address, str], lt: int, block: BlockIdExt , **kwargs) -> typing.Optional[Transaction]: return await self.execute_method('get_one_transaction', **self._get_args(locals())) async def raw_get_transactions(self, address: typing.Union[Address, str], count: int, from_lt: int = None, from_hash: typing.Optional[bytes] = None , **kwargs) -> typing.Tuple[typing.List[Transaction], typing.List[BlockIdExt]]: return await self.execute_method('raw_get_transactions', **self._get_args(locals())) async def get_transactions(self, address: typing.Union[Address, str], count: int, from_lt: int = None, from_hash: typing.Optional[bytes] = None, to_lt: int = 0 , **kwargs) -> typing.List[Transaction]: return await self.execute_method('get_transactions', **self._get_args(locals())) async def raw_get_block_transactions(self, block: BlockIdExt, count: int = 1024, **kwargs) -> typing.List[dict]: return await self.execute_method('raw_get_block_transactions', **self._get_args(locals())) async def raw_get_block_transactions_ext(self, block: BlockIdExt, count: int = 1024, **kwargs) -> typing.List[Transaction]: return await self.execute_method('raw_get_block_transactions_ext', **self._get_args(locals())) async def raw_get_mc_block_proof(self, known_block: BlockIdExt, target_block: typing.Optional[BlockIdExt] = None, return_best_key_block=False , **kwargs) -> typing.Tuple[ bool, BlockIdExt, typing.Optional[BlockIdExt], typing.Optional[int] ]: return await self.execute_method('raw_get_mc_block_proof', **self._get_args(locals())) async def get_mc_block_proof(self, known_block: BlockIdExt, target_block: BlockIdExt, return_best_key_block=False , **kwargs) -> typing.Tuple[typing.Optional[BlockIdExt], int]: return await self.execute_method('get_mc_block_proof', **self._get_args(locals())) async def prove_block(self, target_block: BlockIdExt, **kwargs) -> None: return await self.execute_method('prove_block', **self._get_args(locals())) async def get_config_all(self, blk: typing.Optional[BlockIdExt] = None, **kwargs) -> dict: return await self.execute_method('get_config_all', **self._get_args(locals())) async def get_config_params(self, params: typing.List[int], blk: typing.Optional[BlockIdExt] = None, **kwargs) -> dict: return await self.execute_method('get_config_params', **self._get_args(locals())) async def get_libraries(self, library_list: typing.List[typing.Union[bytes, str]], **kwargs) -> typing.Dict[str, typing.Optional[Cell]]: return await self.execute_method('get_libraries', **self._get_args(locals())) async def get_out_msg_queue_sizes(self, wc: int = None, shard: int = None, **kwargs): return await self.execute_method('get_out_msg_queue_sizes', **self._get_args(locals())) async def nonfinal_get_validator_groups(self, wc: int = None, shard: int = None, **kwargs): return await self.execute_method('nonfinal_get_validator_groups', **self._get_args(locals())) async def nonfinal_raw_get_candidate(self, candidate_id: dict, **kwargs): return await self.execute_method('nonfinal_raw_get_candidate', **self._get_args(locals())) async def nonfinal_get_candidate(self, candidate_id: dict, **kwargs): return await self.execute_method('nonfinal_get_candidate', **self._get_args(locals())) async def get_shard_block_proof(self, blk: BlockIdExt, prove_mc: bool = False, **kwargs): return await self.execute_method('get_shard_block_proof', **self._get_args(locals())) """CODE ABOVE IS AUTOGENERATED. DO NOT EDIT MANUALLY""" async def raw_send_message(self, message: bytes, **kwargs): _tasks = [] _choose_random = kwargs.pop('choose_random', True) _k = 4 if len(self._alive_peers) < 12 else len(self._alive_peers) // 3 for _ in range(_k): # distribute external to approximately 1/3 alive peers, but not less than 4 _tasks.append(self.execute_method('raw_send_message', choose_random=_choose_random, **self._get_args(locals()))) result = await asyncio.gather(*_tasks, return_exceptions=True) success = False exc = None for r in result: if isinstance(r, Exception): exc = r elif r == 1: success = True break if success: return 1 raise exc # raise last exception async def close_all(self): for peer in self._peers: self._check_errors(peer) if peer.inited: await peer.close() self._checker.cancel() while not self._checker.done(): await asyncio.sleep(0) self.inited = False async def close(self): raise BalancerError('Use .close_all()') @classmethod def from_config(cls, config: dict, trust_level: int = 2, timeout: int = 10): clients = [] for i in range(len(config['liteservers'])): clients.append(LiteClient.from_config(config, i, trust_level, timeout)) return cls(clients) @classmethod def from_mainnet_config(cls, trust_level: int = 0, timeout: int = 10): config = requests.get('https://ton.org/global-config.json').json() return cls.from_config(config, trust_level, timeout) @classmethod def from_testnet_config(cls, trust_level: int = 0, timeout: int = 10): config = requests.get('https://ton.org/testnet-global.config.json').json() return cls.from_config(config, trust_level, timeout) async def __aenter__(self): await self.start_up() return self async def __aexit__(self, exc_type, exc_val, exc_tb): await self.close_all() if exc_type: return False return True