import random import logging from twisted.internet import defer, reactor from lbrynet import conf from lbrynet.core.client.ClientProtocol import ClientProtocolFactory from lbrynet.core.Error import InsufficientFundsError from lbrynet.core import utils log = logging.getLogger(__name__) class PeerConnectionHandler: def __init__(self, request_creators, factory): self.request_creators = request_creators self.factory = factory self.connection = None class ConnectionManager: #implements(interfaces.IConnectionManager) MANAGE_CALL_INTERVAL_SEC = 5 TCP_CONNECT_TIMEOUT = 15 def __init__(self, downloader, rate_limiter, primary_request_creators, secondary_request_creators): self.seek_head_blob_first = conf.settings['seek_head_blob_first'] self.max_connections_per_stream = conf.settings['max_connections_per_stream'] self.downloader = downloader self.rate_limiter = rate_limiter self._primary_request_creators = primary_request_creators self._secondary_request_creators = secondary_request_creators self._peer_connections = {} # {Peer: PeerConnectionHandler} self._connections_closing = {} # {Peer: deferred (fired when the connection is closed)} self._next_manage_call = None # a deferred that gets fired when a _manage call is set self._manage_deferred = None self.stopped = True log.debug("%s initialized", self._get_log_name()) # this identifies what the connection manager is for, # used for logging purposes only def _get_log_name(self): out = 'Connection Manager Unknown' if hasattr(self.downloader, 'stream_name'): out = 'Connection Manager '+self.downloader.stream_name elif hasattr(self.downloader, 'blob_hash'): out = 'Connection Manager '+self.downloader.blob_hash return out def _start(self): self.stopped = False if self._next_manage_call is not None and self._next_manage_call.active() is True: self._next_manage_call.cancel() def start(self): log.debug("%s starting", self._get_log_name()) self._start() self._next_manage_call = utils.call_later(0, self.manage) return defer.succeed(True) @defer.inlineCallbacks def stop(self): log.debug("%s stopping", self._get_log_name()) self.stopped = True # wait for the current manage call to finish if self._manage_deferred: yield self._manage_deferred # in case we stopped between manage calls, cancel the next one if self._next_manage_call and self._next_manage_call.active(): self._next_manage_call.cancel() self._next_manage_call = None yield self._close_peers() def num_peer_connections(self): return len(self._peer_connections) def _close_peers(self): def disconnect_peer(p): d = defer.Deferred() self._connections_closing[p] = d self._peer_connections[p].connection.disconnect() if p in self._peer_connections: del self._peer_connections[p] return d def close_connection(p): log.debug("%s Abruptly closing a connection to %s due to downloading being paused", self._get_log_name(), p) if self._peer_connections[p].factory.p is not None: d = self._peer_connections[p].factory.p.cancel_requests() else: d = defer.succeed(True) d.addBoth(lambda _: disconnect_peer(p)) return d closing_deferreds = [close_connection(peer) for peer in self._peer_connections.keys()] return defer.DeferredList(closing_deferreds) @defer.inlineCallbacks def get_next_request(self, peer, protocol): log.debug("%s Trying to get the next request for peer %s", self._get_log_name(), peer) if not peer in self._peer_connections or self.stopped is True: log.debug("%s The peer %s has already been told to shut down.", self._get_log_name(), peer) defer.returnValue(False) requests = yield self._send_primary_requests(peer, protocol) have_request = any(r[1] for r in requests if r[0] is True) if have_request: yield self._send_secondary_requests(peer, protocol) defer.returnValue(have_request) def _send_primary_requests(self, peer, protocol): def handle_error(err): err.trap(InsufficientFundsError) self.downloader.insufficient_funds(err) return False def check_if_request_sent(request_sent, request_creator): if peer not in self._peer_connections: # This can happen if the connection is told to close return False if request_sent is False: if request_creator in self._peer_connections[peer].request_creators: self._peer_connections[peer].request_creators.remove(request_creator) else: if not request_creator in self._peer_connections[peer].request_creators: self._peer_connections[peer].request_creators.append(request_creator) return request_sent ds = [] for p_r_c in self._primary_request_creators: d = p_r_c.send_next_request(peer, protocol) d.addErrback(handle_error) d.addCallback(check_if_request_sent, p_r_c) ds.append(d) return defer.DeferredList(ds, fireOnOneErrback=True) def _send_secondary_requests(self, peer, protocol): ds = [ s_r_c.send_next_request(peer, protocol) for s_r_c in self._secondary_request_creators ] return defer.DeferredList(ds) @defer.inlineCallbacks def manage(self, schedule_next_call=True): self._manage_deferred = defer.Deferred() if len(self._peer_connections) < self.max_connections_per_stream: log.debug("%s have %d connections, looking for %d", self._get_log_name(), len(self._peer_connections), self.max_connections_per_stream) peers = yield self._get_new_peers() for peer in peers: self._connect_to_peer(peer) self._manage_deferred.callback(None) self._manage_deferred = None if not self.stopped and schedule_next_call: self._next_manage_call = utils.call_later(self.MANAGE_CALL_INTERVAL_SEC, self.manage) def return_shuffled_peers_not_connected_to(self, peers, new_conns_needed): out = [peer for peer in peers if peer not in self._peer_connections] random.shuffle(out) return out[0:new_conns_needed] @defer.inlineCallbacks def _get_new_peers(self): new_conns_needed = self.max_connections_per_stream - len(self._peer_connections) if new_conns_needed < 1: defer.returnValue([]) # we always get the peer from the first request creator # must be a type BlobRequester... request_creator = self._primary_request_creators[0] log.debug("%s Trying to get a new peer to connect to", self._get_log_name()) # find peers for the head blob if configured to do so if self.seek_head_blob_first: try: peers = yield request_creator.get_new_peers_for_head_blob() peers = self.return_shuffled_peers_not_connected_to(peers, new_conns_needed) except KeyError: log.warning("%s does not have a head blob", self._get_log_name()) peers = [] else: peers = [] # we didn't find any new peers on the head blob, # we have to look for the first unavailable blob if not peers: peers = yield request_creator.get_new_peers_for_next_unavailable() peers = self.return_shuffled_peers_not_connected_to(peers, new_conns_needed) log.debug("%s Got a list of peers to choose from: %s", self._get_log_name(), peers) log.debug("%s Current connections: %s", self._get_log_name(), self._peer_connections.keys()) log.debug("%s List of connection states: %s", self._get_log_name(), [p_c_h.connection.state for p_c_h in self._peer_connections.values()]) defer.returnValue(peers) def _connect_to_peer(self, peer): if self.stopped: return log.debug("%s Trying to connect to %s", self._get_log_name(), peer) factory = ClientProtocolFactory(peer, self.rate_limiter, self) factory.connection_was_made_deferred.addCallback( lambda c_was_made: self._peer_disconnected(c_was_made, peer)) self._peer_connections[peer] = PeerConnectionHandler(self._primary_request_creators[:], factory) connection = reactor.connectTCP(peer.host, peer.port, factory, timeout=self.TCP_CONNECT_TIMEOUT) self._peer_connections[peer].connection = connection def _peer_disconnected(self, connection_was_made, peer): log.debug("%s protocol disconnected for %s", self._get_log_name(), peer) if peer in self._peer_connections: del self._peer_connections[peer] if peer in self._connections_closing: d = self._connections_closing[peer] del self._connections_closing[peer] d.callback(True) return connection_was_made