downloader.py

BitTorrentABC-Linux-V.2.4.3源码

# Written by Bram Cohen
# see LICENSE.txt for license information

from CurrentRateMeasure import Measure
from random import shuffle
from time import time
from math import sqrt
true = 1
false = 0

EXPIRE_TIME = 60 * 60

class SingleDownload:
    def __init__(self, downloader, connection):
        self.downloader = downloader
        self.unhave = downloader.numpieces
        self.connection = connection
        self.choked = true
        self.interested = false
        self.active_requests = []
        self.measure = Measure(downloader.max_rate_period)
        self.peermeasure = Measure(downloader.max_rate_period)
        self.have = [false] * downloader.numpieces
        self.last = 0
        self.example_interest = None

    def _backlog(self):
        return min(2+int(self.measure.get_rate()/8000), self.downloader.backlog)

    def disconnected(self):
        self.downloader.lost_peer(self)
        for i in xrange(len(self.have)):
            if self.have[i]:
                self.downloader.picker.lost_have(i)
        if self.unhave == 0 and self.downloader.storage.is_endgame():
            self.downloader.add_disconnected_seed(self.connection.get_id())
        self._letgo()

    def _letgo(self):
        if not self.active_requests:
            return
        if self.downloader.storage.is_endgame():
            self.active_requests = []
            return
        lost = []
        for index, begin, length in self.active_requests:
            self.downloader.storage.request_lost(index, begin, length)
            if index not in lost:
                lost.append(index)
        self.active_requests = []
        ds = [d for d in self.downloader.downloads if not d.choked]
        shuffle(ds)
        for d in ds:
            d._request_more(lost)
        for d in self.downloader.downloads:
            if d.choked and not d.interested:
                for l in lost:
                    if d.have[l] and self.downloader.storage.do_I_have_requests(l):
                        d.interested = true
                        d.connection.send_interested()
                        break

    def got_choke(self):
        if not self.choked:
            self.choked = true
            self._letgo()

    def got_unchoke(self):
        if self.choked:
            self.choked = false
            if self.interested:
                self._request_more()

    def is_choked(self):
        return self.choked

    def is_interested(self):
        return self.interested

    def got_piece(self, index, begin, piece):
        try:
            self.active_requests.remove((index, begin, len(piece)))
        except ValueError:
            return false
        if self.downloader.storage.is_endgame():
            self.downloader.all_requests.remove((index, begin, len(piece)))
        self.last = time()
        self.measure.update_rate(len(piece))
        self.downloader.measurefunc(len(piece))
        self.downloader.downmeasure.update_rate(len(piece))
        if not self.downloader.storage.piece_came_in(index, begin, piece):
            self.downloader.piece_flunked(index)
            return false
        if self.downloader.storage.do_I_have(index):
            self.downloader.picker.complete(index)
        if self.downloader.storage.is_endgame():
            for d in self.downloader.downloads:
                if d is not self:
                    try:
                        d.active_requests.remove((index, begin, len(piece)))
                    except ValueError:
                        continue
                    d.connection.send_cancel(index, begin, len(piece))
                    d.fix_download_endgame()
        self._request_more()
        if self.downloader.picker.am_I_complete():
            for d in [i for i in self.downloader.downloads if i.unhave == 0]:
                d.connection.send_have(index)   # be nice, tell the other seed you completed
                self.downloader.add_disconnected_seed(d.connection.get_id())
                d.connection.close()
        return self.downloader.storage.do_I_have(index)

    def _want(self, index):
        return self.have[index] and self.downloader.storage.do_I_have_requests(index)

    def _request_more(self, indices = None):
        assert not self.choked
        if len(self.active_requests) >= self._backlog():
            return
        if self.downloader.storage.is_endgame():
            self.fix_download_endgame()
            return
        lost_interests = []
        while len(self.active_requests) < self._backlog():
            if indices is None:
                interest = self.downloader.picker.next(self._want)
            else:
                interest = None
                for i in indices:
                    if self.have[i] and self.downloader.storage.do_I_have_requests(i):
                        interest = i
                        break
            if interest is None:
                break
            if not self.interested:
                self.interested = true
                self.connection.send_interested()
            self.example_interest = interest
            begin, length = self.downloader.storage.new_request(interest)
            self.downloader.picker.requested(interest)
            self.active_requests.append((interest, begin, length))
            self.connection.send_request(interest, begin, length)
            if not self.downloader.storage.do_I_have_requests(interest):
                lost_interests.append(interest)
        if not self.active_requests and self.interested:
            self.interested = false
            self.connection.send_not_interested()
        if lost_interests:
            for d in self.downloader.downloads:
                if d.active_requests or not d.interested:
                    continue
                if d.example_interest is not None and self.downloader.storage.do_I_have_requests(d.example_interest):
                    continue
                for lost in lost_interests:
                    if d.have[lost]:
                        break
                else:
                    continue
                interest = self.downloader.picker.next(d._want)
                if interest is None:
                    d.interested = false
                    d.connection.send_not_interested()
                else:
                    d.example_interest = interest
        if self.downloader.storage.is_endgame():
            self.downloader.all_requests = []
            for d in self.downloader.downloads:
                self.downloader.all_requests.extend(d.active_requests)
            for d in self.downloader.downloads:
                d.fix_download_endgame()

    def fix_download_endgame(self):
        want = [a for a in self.downloader.all_requests if self.have[a[0]] and a not in self.active_requests]
        if self.interested and not self.active_requests and not want:
            self.interested = false
            self.connection.send_not_interested()
            return
        if not self.interested and want:
            self.interested = true
            self.connection.send_interested()
        if self.choked or len(self.active_requests) >= self._backlog():
            return
        shuffle(want)
        del want[self._backlog() - len(self.active_requests):]
        self.active_requests.extend(want)
        for piece, begin, length in want:
            self.connection.send_request(piece, begin, length)

    def got_have(self, index):
        if index == self.downloader.numpieces-1:
            self.downloader.totalmeasure.update_rate(self.downloader.storage.total_length-(self.downloader.numpieces-1)*self.downloader.storage.piece_length)
            self.peermeasure.update_rate(self.downloader.storage.total_length-(self.downloader.numpieces-1)*self.downloader.storage.piece_length)
        else:
            self.downloader.totalmeasure.update_rate(self.downloader.storage.piece_length)
            self.peermeasure.update_rate(self.downloader.storage.piece_length)
        if self.have[index]:
            return
        self.have[index] = true
        self.unhave -= 1
        self.downloader.picker.got_have(index)
        if self.downloader.picker.am_I_complete() and self.unhave == 0:
            self.downloader.add_disconnected_seed(self.connection.get_id())
            self.connection.close()
            return
        if self.downloader.storage.do_I_have_requests(index):
            if not self.choked:
                self._request_more([index])
            else:
                if not self.interested:
                    self.interested = true
                    self.connection.send_interested()

    def got_have_bitfield(self, have):
        self.have = have
        for i in xrange(len(have)):
            if have[i]:
                self.unhave -= 1
                self.downloader.picker.got_have(i)
        if self.downloader.picker.am_I_complete() and self.unhave == 0:
            self.connection.send_bitfield(have)     # be nice, show you're a seed too
            self.connection.close()
            self.downloader.add_disconnected_seed(self.connection.get_id())
            return
        if self.downloader.storage.is_endgame():
            for piece, begin, length in self.downloader.all_requests:
                if self.have[piece]:
                    self.interested = true
                    self.connection.send_interested()
                    return
        for i in xrange(len(have)):
            if have[i] and self.downloader.storage.do_I_have_requests(i):
                self.interested = true
                self.connection.send_interested()
                return

    def get_rate(self):
        return self.measure.get_rate()

    def is_snubbed(self):
        return time() - self.last > self.downloader.snub_time


class Downloader:
    def __init__(self, storage, picker, backlog, max_rate_period, numpieces, 
            downmeasure, snub_time, measurefunc = lambda x: None):
        self.storage = storage
        self.picker = picker
        self.backlog = backlog
        self.max_rate_period = max_rate_period
        self.downmeasure = downmeasure
        self.totalmeasure = Measure(max_rate_period*storage.piece_length/storage.request_size)
        self.numpieces = numpieces
        self.snub_time = snub_time
        self.measurefunc = measurefunc
        self.disconnectedseeds = {}
        self.downloads = []
        self.maxlistlen = long(sqrt(numpieces))

    def make_download(self, connection):
        self.downloads.append(SingleDownload(self, connection))
        return self.downloads[-1]

    def piece_flunked(self, index):
        if self.storage.is_endgame():
            if not self.downloads:  # must've been called externally
                self.storage.reset_endgame()
                return
            while self.storage.do_I_have_requests(index):
                nb, nl = self.storage.new_request(index)
                self.all_requests.append((index, nb, nl))
            for d in self.downloads:
                d.fix_download_endgame()
            return
        self.picker.bump(index)
        ds = [d for d in self.downloads if not d.choked]
        shuffle(ds)
        for d in ds:
            d._request_more([index])

    def has_downloaders(self):
        return len(self.downloads)

    def lost_peer(self, download):
        self.downloads.remove(download)
        if self.storage.is_endgame() and not self.downloads and self.all_requests: # all peers gone
            for index, begin, length in self.all_requests:
                self.storage.request_lost(index, begin, length)