import logging import asyncio from asyncio import CancelledError from time import time from typing import List from torba.rpc import RPCSession as BaseClientSession, Connector, RPCError from torba import __version__ from torba.stream import StreamController log = logging.getLogger(__name__) class ClientSession(BaseClientSession): def __init__(self, *args, network, server, **kwargs): self.network = network self.server = server super().__init__(*args, **kwargs) self._on_disconnect_controller = StreamController() self.on_disconnected = self._on_disconnect_controller.stream self.bw_limit = self.framer.max_size = self.max_errors = 1 << 32 self.max_seconds_idle = 60 self.ping_task = None async def send_request(self, method, args=()): try: return await super().send_request(method, args) except RPCError as e: log.warning("Wallet server returned an error. Code: %s Message: %s", *e.args) raise e async def ping_forever(self): # TODO: change to 'ping' on newer protocol (above 1.2) while not self.is_closing(): if (time() - self.last_send) > self.max_seconds_idle: await self.send_request('server.banner') await asyncio.sleep(self.max_seconds_idle//3) async def create_connection(self, timeout=6): connector = Connector(lambda: self, *self.server) await asyncio.wait_for(connector.create_connection(), timeout=timeout) self.ping_task = asyncio.create_task(self.ping_forever()) async def handle_request(self, request): controller = self.network.subscription_controllers[request.method] controller.add(request.args) def connection_lost(self, exc): super().connection_lost(exc) self._on_disconnect_controller.add(True) if self.ping_task: self.ping_task.cancel() class BaseNetwork: def __init__(self, ledger): self.config = ledger.config self.client: ClientSession = None self.session_pool: SessionPool = None self.running = False self._on_connected_controller = StreamController() self.on_connected = self._on_connected_controller.stream self._on_header_controller = StreamController() self.on_header = self._on_header_controller.stream self._on_status_controller = StreamController() self.on_status = self._on_status_controller.stream self.subscription_controllers = { 'blockchain.headers.subscribe': self._on_header_controller, 'blockchain.address.subscribe': self._on_status_controller, } async def start(self): self.running = True connect_timeout = self.config.get('connect_timeout', 6) self.session_pool = SessionPool(network=self, timeout=connect_timeout) self.session_pool.start(self.config['default_servers']) while True: try: self.client = await self.pick_fastest_session() if self.is_connected: await self.ensure_server_version() log.info("Successfully connected to SPV wallet server: %s:%d", *self.client.server) self._on_connected_controller.add(True) await self.client.on_disconnected.first except CancelledError: self.running = False except asyncio.TimeoutError: log.warning("Timed out while trying to find a server!") except Exception: # pylint: disable=broad-except log.exception("Exception while trying to find a server!") if not self.running: return elif self.client: await self.client.close() self.client.connection.cancel_pending_requests() async def stop(self): self.running = False if self.session_pool: self.session_pool.stop() if self.is_connected: disconnected = self.client.on_disconnected.first await self.client.close() await disconnected @property def is_connected(self): return self.client is not None and not self.client.is_closing() def rpc(self, list_or_method, args): if self.is_connected: return self.client.send_request(list_or_method, args) else: raise ConnectionError("Attempting to send rpc request when connection is not available.") async def pick_fastest_session(self): sessions = await self.session_pool.get_online_sessions() done, pending = await asyncio.wait([ self.probe_session(session) for session in sessions if not session.is_closing() ], return_when='FIRST_COMPLETED') for task in pending: task.cancel() for session in done: return await session async def probe_session(self, session: ClientSession): await session.send_request('server.banner') return session def ensure_server_version(self, required='1.2'): return self.rpc('server.version', [__version__, required]) def broadcast(self, raw_transaction): return self.rpc('blockchain.transaction.broadcast', [raw_transaction]) def get_history(self, address): return self.rpc('blockchain.address.get_history', [address]) def get_transaction(self, tx_hash): return self.rpc('blockchain.transaction.get', [tx_hash]) def get_transaction_height(self, tx_hash): return self.rpc('blockchain.transaction.get_height', [tx_hash]) def get_merkle(self, tx_hash, height): return self.rpc('blockchain.transaction.get_merkle', [tx_hash, height]) def get_headers(self, height, count=10000): return self.rpc('blockchain.block.headers', [height, count]) def subscribe_headers(self): return self.rpc('blockchain.headers.subscribe', [True]) def subscribe_address(self, address): return self.rpc('blockchain.address.subscribe', [address]) class SessionPool: def __init__(self, network: BaseNetwork, timeout: float): self.network = network self.sessions: List[ClientSession] = [] self._dead_servers: List[ClientSession] = [] self.maintain_connections_task = None self.timeout = timeout # triggered when the master server is out, to speed up reconnect self._lost_master = asyncio.Event() @property def online(self): for session in self.sessions: if not session.is_closing(): return True return False def start(self, default_servers): self.sessions = [ ClientSession(network=self.network, server=server) for server in default_servers ] self.maintain_connections_task = asyncio.create_task(self.ensure_connections()) def stop(self): if self.maintain_connections_task: self.maintain_connections_task.cancel() for session in self.sessions: if not session.is_closing(): session.abort() self.sessions, self._dead_servers, self.maintain_connections_task = [], [], None async def ensure_connections(self): while True: await asyncio.gather(*[ self.ensure_connection(session) for session in self.sessions ], return_exceptions=True) await asyncio.wait([asyncio.sleep(3), self._lost_master.wait()], return_when='FIRST_COMPLETED') self._lost_master.clear() if not self.sessions: self.sessions.extend(self._dead_servers) self._dead_servers = [] async def ensure_connection(self, session): if not session.is_closing(): return try: return await session.create_connection(self.timeout) except asyncio.TimeoutError: log.warning("Timeout connecting to %s:%d", *session.server) except asyncio.CancelledError: # pylint: disable=try-except-raise raise except Exception as err: # pylint: disable=broad-except if 'Connect call failed' in str(err): log.warning("Could not connect to %s:%d", *session.server) else: log.exception("Connecting to %s:%d raised an exception:", *session.server) self._dead_servers.append(session) self.sessions.remove(session) async def get_online_sessions(self): self._lost_master.set() while not self.online: await asyncio.sleep(0.1) return self.sessions