Example #1
0
class BtConnection(Connection):
    def __init__(
        self,
        torrentIdent,
        globalStatus,
        connStatsCache,
        connStatus,
        remotePeerId,
        scheduler,
        conn,
        direction,
        remotePeerAddr,
        inMeasureParent,
        outMeasureParent,
        outLimiter,
        inLimiter,
    ):

        log = Logger("BtConnection", "%-6s - %-6s - ", torrentIdent, conn.fileno())
        inRate, outRate = connStatsCache.get(torrentIdent, remotePeerAddr, remotePeerId)

        Connection.__init__(
            self,
            connStatus,
            scheduler,
            conn,
            direction,
            remotePeerAddr,
            inRate,
            outRate,
            inMeasureParent,
            outMeasureParent,
            outLimiter,
            inLimiter,
            Messages.getMessageLength,
            Messages.decodeMessage,
            4,
            140000,
            Messages.generateKeepAlive,
            log,
        )

        # ident
        self.torrentIdent = torrentIdent

        # peer
        self.remotePeerId = remotePeerId
        self.remoteClient = peerIdToClient(remotePeerId)

        # conn stats cache
        self.connStatsCache = connStatsCache

        # piece status
        self.status = Status(globalStatus.getPieceAmount(), globalStatus)

        # choke and interest state
        self.localInterest = False
        self.remoteInterest = False
        self.localChoke = True
        self.remoteChoke = True

        # requests
        self.outRequestsInFlight = 0
        self.outRequestQueue = []
        self.outRequestHandles = {}
        self.maxInRequests = self._calculateMaxAmountOfInRequests()
        self.inRequestQueue = []
        self.inRequestInfo = {}

        # events
        self.requestTimeoutEvent = None

    ##internal functions - socket

    def _gotMessage(self, msg):
        if msg[0] == 7:
            # a piece part
            if self._hasThisInRequest(msg[1][0], msg[1][1], len(msg[1][2])):
                # a valid one even
                assert self.requestTimeoutEvent is not None, "got data for a request but no request timeout exists?!"
                if self._amountOfInRequests() == 1:
                    self.sched.removeEvent(self.requestTimeoutEvent)
                    self.requestTimeoutEvent = None
                else:
                    self.sched.rescheduleEvent(self.requestTimeoutEvent, timedelta=120)
                self.inRate.updatePayloadCounter(len(msg[1][2]))

    def _close(self):
        Connection._close(self)

        # store inRate and outRate
        self.connStatsCache.store(self.torrentIdent, self.remotePeerAddr, self.remotePeerId, self.inRate, self.outRate)

        # clear piece status
        self.status.clear()

        # set interest and choke states to defaults
        self.localInterest = False
        self.remoteInterest = False
        self.localChoke = True
        self.remoteChoke = True

        # remove requests
        self._delAllOutRequests()
        self._delAllInRequests()

    ##internal functions - inrequests

    def _addInRequest(self, pieceIndex, offset, length, callback=None, callbackArgs=[], callbackKw={}):
        assert self.remoteChoke == False, "requesting but choked?!"

        # add timeout
        if self.requestTimeoutEvent is None:
            self.requestTimeoutEvent = self.sched.scheduleEvent(
                self.timeout, timedelta=120, funcArgs=["request timed out"]
            )

        # add request
        messageId = self._queueSend(Messages.generateRequest(pieceIndex, offset, length))
        inRequest = (pieceIndex, offset, length)
        assert not inRequest in self.inRequestInfo, "queueing an already queued request?!"
        self.inRequestQueue.append(inRequest)
        self.inRequestInfo[inRequest] = {
            "messageId": messageId,
            "func": callback,
            "funcArgs": callbackArgs,
            "funcKw": callbackKw,
        }

    def _getInRequestsOfPiece(self, pieceIndex):
        requests = set([inRequest[1] for inRequest in self.inRequestQueue if inRequest[0] == pieceIndex])
        return requests

    def _finishedInRequest(self, pieceIndex, offset, length):
        # try to find the request and delete it if found
        inRequest = (pieceIndex, offset, length)
        self.inRequestQueue.remove(inRequest)
        del self.inRequestInfo[inRequest]

    def _cancelInRequest(self, pieceIndex, offset, length):
        # try to find the request, send cancel and then delete it
        inRequest = (pieceIndex, offset, length)
        self.inRequestQueue.remove(inRequest)
        requestInfo = self.inRequestInfo.pop(inRequest)
        if not self._abortSend(requestInfo["messageId"]):
            # the request was already send
            self._queueSend(Messages.generateCancel(pieceIndex, offset, length))

    def _delAllInRequests(self):
        for requestInfo in self.inRequestInfo.itervalues():
            # call callback of failed request
            self._abortSend(requestInfo["messageId"])
            if requestInfo["func"] is not None:
                apply(requestInfo["func"], requestInfo["funcArgs"], requestInfo["funcKw"])
        self.inRequestQueue = []
        self.inRequestInfo = {}

    def _hasThisInRequest(self, pieceIndex, offset, length):
        return (pieceIndex, offset, length) in self.inRequestInfo

    def _amountOfInRequests(self):
        return len(self.inRequestQueue)

    def _calculateMaxAmountOfInRequests(self):
        if self.remoteClient.startswith("Azureus"):
            limit = 16

        elif self.remoteClient.startswith("I2PRufus"):
            limit = 16

        elif self.remoteClient.startswith("I2PSnark"):
            limit = 32

        elif self.remoteClient.startswith("PyBit") and " " in self.remoteClient:
            version = tuple((int(digit) for digit in self.remoteClient.split(" ")[1].split(".")))
            if version < (0, 0, 9):
                limit = 32
            else:
                limit = 64

        elif self.remoteClient.startswith("Robert"):
            limit = 31

        else:
            # I2P-Bt and unknown
            limit = 16

        return limit

    def _getMaxAmountOfInRequests(self):
        return self.maxInRequests

    ##internal functions - outrequests

    def _sendOutRequest(self):
        # queue one outrequest in the outbuffer
        outRequest = self.outRequestQueue.pop(0)
        try:
            # try to get data
            data = self.outRequestHandles.pop(outRequest)()
        except StorageException:
            # failed to get data
            self.log.error("Failed to get data for outrequest:\n%s", logTraceback())
            data = None
            self._fail("could not get data for outrequest")

        if data is not None:
            # got data
            if not len(data) == outRequest[2]:
                self.log.error("Did not get enough data for outrequest: expected %i, got %i!", outRequest[2], len(data))
                self._fail("could not get data for outrequest")
            else:
                message = Messages.generatePiece(outRequest[0], outRequest[1], data)
                self.outRequestsInFlight += 1
                self._queueSend(message, self._outRequestGotSend, [outRequest[2]])

    def _outRequestGotSend(self, dataSize):
        self.outRate.updatePayloadCounter(dataSize)
        self.outRequestsInFlight -= 1
        assert self.outRequestsInFlight == 0, "multiple out requests in flight?!"
        assert len(self.outRequestQueue) == len(
            self.outRequestHandles
        ), "out of sync: queue length %i but %i handles!" % (len(self.outRequestQueue), len(self.outRequestHandles))
        if len(self.outRequestQueue) > 0 and self.outRequestsInFlight == 0:
            self._sendOutRequest()

    def _addOutRequest(self, pieceIndex, offset, length, dataHandle):
        self.outRequestQueue.append((pieceIndex, offset, length))
        self.outRequestHandles[(pieceIndex, offset, length)] = dataHandle
        if self.outRequestsInFlight == 0:
            # no outrequest is currently being send, send one directly
            self._sendOutRequest()

    def _hasThisOutRequest(self, pieceIndex, offset, length):
        return (pieceIndex, offset, length) in self.outRequestHandles

    def _getAmountOfOutRequests(self):
        return len(self.outRequestQueue)

    def _delOutRequest(self, pieceIndex, offset, length):
        # try to find the request and delete it if found
        outRequest = (pieceIndex, offset, length)
        if outRequest in self.outRequestHandles:
            self.outRequestQueue.remove(outRequest)
            del self.outRequestHandles[outRequest]

    def _delAllOutRequests(self):
        self.outRequestQueue = []
        self.outRequestHandles.clear()

    ##internal functions - choking and interest

    def _setLocalInterest(self, value):
        if value == False and self.localInterest == True:
            # changing to negative, we were interested before
            assert self._amountOfInRequests() == 0, "Local Requests are running and we are not interested?"
            # self._cancelAllInRequests()
            self._queueSend(Messages.generateNotInterested())
            self.localInterest = value
        elif value == True and self.localInterest == False:
            self._queueSend(Messages.generateInterested())
            self.localInterest = value

    def _setRemoteInterest(self, value):
        if value == False and self.remoteInterest == True:
            self.setLocalChoke(True)
            self.remoteInterest = value

        elif value == True and self.remoteInterest == False:
            self.remoteInterest = value

    def _setLocalChoke(self, value):
        if value == True and self.localChoke == False:
            # choking
            self._delAllOutRequests()
            self._queueSend(Messages.generateChoke())
            self.localChoke = value

        elif value == False and self.localChoke == True:
            self._queueSend(Messages.generateUnChoke())
            self.localChoke = value

    def _setRemoteChoke(self, value):
        if value == True and self.remoteChoke == False:
            # choked us, delete all incomming requests
            self._delAllInRequests()
            self.remoteChoke = value

        elif value == False and self.remoteChoke == True:
            self.remoteChoke = value

    ##internal functions - other

    def _getScore(self):
        ratio = self._getPayloadRatio()
        score = ratio + (ratio * self.inRate.getAveragePayloadRate())
        return score

    ##external functions - choking and interested

    def localChoked(self):
        self.lock.acquire()
        value = self.localChoke
        self.lock.release()
        return value

    def remoteChoked(self):
        self.lock.acquire()
        value = self.remoteChoke
        self.lock.release()
        return value

    def localInterested(self):
        self.lock.acquire()
        value = self.localInterest
        self.lock.release()
        return value

    def remoteInterested(self):
        self.lock.acquire()
        value = self.remoteInterest
        self.lock.release()
        return value

    def setLocalInterest(self, value):
        self.lock.acquire()
        if not self.closed:
            self._setLocalInterest(value)
        self.lock.release()

    def setRemoteInterest(self, value):
        self.lock.acquire()
        if not self.closed:
            self._setRemoteInterest(value)
        self.lock.release()

    def setLocalChoke(self, value):
        self.lock.acquire()
        if not self.closed:
            self._setLocalChoke(value)
        self.lock.release()

    def setRemoteChoke(self, value):
        self.lock.acquire()
        if not self.closed:
            self._setRemoteChoke(value)
        self.lock.release()

    ##external functions - inrequests

    def addInRequest(self, pieceIndex, offset, length, failFunc=None, failFuncArgs=[], failFuncKw={}):
        self.lock.acquire()
        assert not self.remoteChoke, "uhm, we are not allowed to make requests?!"
        if not self.closed:
            self._addInRequest(pieceIndex, offset, length, failFunc, failFuncArgs, failFuncKw)
        self.lock.release()

    def getInRequestsOfPiece(self, pieceIndex):
        self.lock.acquire()
        requests = self._getInRequestsOfPiece(pieceIndex)
        self.lock.release()
        return requests

    def finishedInRequest(self, pieceIndex, offset, length):
        self.lock.acquire()
        if not self.closed:
            self._finishedInRequest(pieceIndex, offset, length)
        self.lock.release()

    def cancelInRequest(self, pieceIndex, offset, length):
        self.lock.acquire()
        if not self.closed:
            self._cancelInRequest(pieceIndex, offset, length)
        self.lock.release()

    def hasThisInRequest(self, pieceIndex, offset, length):
        self.lock.acquire()
        value = self._hasThisInRequest(pieceIndex, offset, length)
        self.lock.release()
        return value

    def getAmountOfInRequests(self):
        self.lock.acquire()
        value = self._amountOfInRequests()
        self.lock.release()
        return value

    def getMaxAmountOfInRequests(self):
        self.lock.acquire()
        value = self._getMaxAmountOfInRequests()
        self.lock.release()
        return value

    ##internal functions - outrequests

    def addOutRequest(self, pieceIndex, offset, length, dataHandle):
        self.lock.acquire()
        if not self.closed:
            self._addOutRequest(pieceIndex, offset, length, dataHandle)
        self.lock.release()

    def hasThisOutRequest(self, pieceIndex, offset, length):
        self.lock.acquire()
        value = self._hasThisOutRequest(pieceIndex, offset, length)
        self.lock.release()
        return value

    def delOutRequest(self, pieceIndex, offset, length):
        self.lock.acquire()
        self._delOutRequest(pieceIndex, offset, length)
        self.lock.release()

    def getAmountOfOutRequests(self):
        self.lock.acquire()
        value = self._getAmountOfOutRequests()
        self.lock.release()
        return value

    ##external functions - get info

    def getStatus(self):
        self.lock.acquire()
        obj = self.status
        self.lock.release()
        return obj

    def getScore(self):
        self.lock.acquire()
        score = self._getScore()
        self.lock.release()
        return score

    def getRemotePeerId(self):
        self.lock.acquire()
        value = self.remotePeerId
        self.lock.release()
        return value

    def getTorrentIdent(self):
        self.lock.acquire()
        value = self.torrentIdent
        self.lock.release()
        return value

    ##external functions - stats

    def getStats(self):
        self.lock.acquire()
        stats = {}
        stats["id"] = self.connIdent
        stats["addr"] = self.conn.getpeername()
        stats["direction"] = self.direction
        stats["connectedInterval"] = time() - self.connectTime
        stats["totalConnectedInterval"] = self.inRate.getTotalRunTime()
        stats["peerProgress"] = self.status.getPercent()
        stats["peerClient"] = self.remoteClient
        stats["inRawBytes"] = self.inRate.getTotalTransferedBytes()
        stats["outRawBytes"] = self.outRate.getTotalTransferedBytes()
        stats["inPayloadBytes"] = self.inRate.getTotalTransferedPayloadBytes()
        stats["outPayloadBytes"] = self.outRate.getTotalTransferedPayloadBytes()
        stats["inRawSpeed"] = self.inRate.getCurrentRate()
        stats["outRawSpeed"] = self.outRate.getCurrentRate()
        stats["localInterest"] = self.localInterest
        stats["remoteInterest"] = self.remoteInterest
        stats["localChoke"] = self.localChoke
        stats["remoteChoke"] = self.remoteChoke
        stats["localRequestCount"] = len(self.inRequestQueue)
        stats["remoteRequestCount"] = self.outRequestsInFlight + len(self.outRequestQueue)
        stats["avgInRawSpeed"] = self.inRate.getAverageRate() * 1024
        stats["avgOutRawSpeed"] = self.outRate.getAverageRate() * 1024
        stats["avgInPayloadSpeed"] = self.inRate.getAveragePayloadRate() * 1024
        stats["avgOutPayloadSpeed"] = self.outRate.getAveragePayloadRate() * 1024
        stats["score"] = self._getScore()
        stats["payloadRatio"] = self._getPayloadRatio()
        stats["protocolOverhead"] = (
            100.0 * (stats["inRawBytes"] + stats["outRawBytes"] - stats["inPayloadBytes"] - stats["outPayloadBytes"])
        ) / max(stats["inPayloadBytes"] + stats["outPayloadBytes"], 1.0)
        self.lock.release()
        return stats