storagewrapper.py

一个多点下载且源码公开的P2P软件

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

from BitTornado.bitfield import Bitfield
from sha import sha
from BitTornado.clock import clock
from traceback import print_exc
from random import randrange
try:
    True
except:
    True = 1
    False = 0
try:
    from bisect import insort
except:
    def insort(l, item):
        l.append(item)
        l.sort()

DEBUG = False

STATS_INTERVAL = 0.2

def dummy_status(fractionDone = None, activity = None):
    pass

class Olist:
    def __init__(self, l = []):
        self.d = {}
        for i in l:
            self.d[i] = 1
    def __len__(self):
        return len(self.d)
    def includes(self, i):
        return self.d.has_key(i)
    def add(self, i):
        self.d[i] = 1
    def extend(self, l):
        for i in l:
            self.d[i] = 1
    def pop(self, n=0):
        # assert self.d
        k = self.d.keys()
        if n == 0:
            i = min(k)
        elif n == -1:
            i = max(k)
        else:
            k.sort()
            i = k[n]
        del self.d[i]
        return i
    def remove(self, i):
        if self.d.has_key(i):
            del self.d[i]

class fakeflag:
    def __init__(self, state=False):
        self.state = state
    def wait(self):
        pass
    def isSet(self):
        return self.state


class StorageWrapper:
    def __init__(self, storage, request_size, hashes, 
            piece_size, finished, failed, 
            statusfunc = dummy_status, flag = fakeflag(), check_hashes = True,
            data_flunked = lambda x: None, backfunc = None,
            config = {}, unpauseflag = fakeflag(True) ):
        self.storage = storage
        self.request_size = long(request_size)
        self.hashes = hashes
        self.piece_size = long(piece_size)
        self.piece_length = long(piece_size)
        self.finished = finished
        self.failed = failed
        self.statusfunc = statusfunc
        self.flag = flag
        self.check_hashes = check_hashes
        self.data_flunked = data_flunked
        self.backfunc = backfunc
        self.config = config
        self.unpauseflag = unpauseflag
        
        self.alloc_type = config.get('alloc_type','normal')
        self.double_check = config.get('double_check', 0)
        self.triple_check = config.get('triple_check', 0)
        if self.triple_check:
            self.double_check = True
        self.bgalloc_enabled = False
        self.bgalloc_active = False
        self.total_length = storage.get_total_length()
        self.amount_left = self.total_length
        if self.total_length <= self.piece_size * (len(hashes) - 1):
            raise ValueError, 'bad data in responsefile - total too small'
        if self.total_length > self.piece_size * len(hashes):
            raise ValueError, 'bad data in responsefile - total too big'
        self.numactive = [0] * len(hashes)
        self.inactive_requests = [1] * len(hashes)
        self.amount_inactive = self.total_length
        self.amount_obtained = 0
        self.amount_desired = self.total_length
        self.have = Bitfield(len(hashes))
        self.have_cloaked_data = None
        self.blocked = [False] * len(hashes)
        self.blocked_holes = []
        self.blocked_movein = Olist()
        self.blocked_moveout = Olist()
        self.waschecked = [False] * len(hashes)
        self.places = {}
        self.holes = []
        self.stat_active = {}
        self.stat_new = {}
        self.dirty = {}
        self.stat_numflunked = 0
        self.stat_numdownloaded = 0
        self.stat_numfound = 0
        self.download_history = {}
        self.failed_pieces = {}
        self.out_of_place = 0
        self.write_buf_max = config['write_buffer_size']*1048576L
        self.write_buf_size = 0L
        self.write_buf = {}   # structure:  piece: [(start, data), ...]
        self.write_buf_list = []

        self.initialize_tasks = [
            ['checking existing data', 0, self.init_hashcheck, self.hashcheckfunc],
            ['moving data', 1, self.init_movedata, self.movedatafunc],
            ['allocating disk space', 1, self.init_alloc, self.allocfunc] ]

        self.backfunc(self._bgalloc,0.1)
        self.backfunc(self._bgsync,max(self.config['auto_flush']*60,60))

    def _bgsync(self):
        if self.config['auto_flush']:
            self.sync()
        self.backfunc(self._bgsync,max(self.config['auto_flush']*60,60))


    def old_style_init(self):
        while self.initialize_tasks:
            msg, done, init, next = self.initialize_tasks.pop(0)
            if init():
                self.statusfunc(activity = msg, fractionDone = done)
                t = clock() + STATS_INTERVAL
                x = 0
                while x is not None:
                    if t < clock():
                        t = clock() + STATS_INTERVAL
                        self.statusfunc(fractionDone = x)
                    self.unpauseflag.wait()
                    if self.flag.isSet():
                        return False
                    x = next()

        self.statusfunc(fractionDone = 0)
        return True


    def initialize(self, donefunc, statusfunc = None):
        self.initialize_done = donefunc
        if statusfunc is None:
            statusfunc = self.statusfunc
        self.initialize_status = statusfunc
        self.initialize_next = None
            
        self.backfunc(self._initialize)

    def _initialize(self):
        if not self.unpauseflag.isSet():
            self.backfunc(self._initialize, 1)
            return
        
        if self.initialize_next:
            x = self.initialize_next()
            if x is None:
                self.initialize_next = None
            else:
                self.initialize_status(fractionDone = x)
        else:
            if not self.initialize_tasks:
                self.initialize_done()
                return
            msg, done, init, next = self.initialize_tasks.pop(0)
            if init():
                self.initialize_status(activity = msg, fractionDone = done)
                self.initialize_next = next

        self.backfunc(self._initialize)


    def init_hashcheck(self):
        if self.flag.isSet():
            return False
        self.check_list = []
        if len(self.hashes) == 0 or self.amount_left == 0:
            self.check_total = 0
            self.finished()
            return False

        self.check_targets = {}
        got = {}
        for p,v in self.places.items():
            assert not got.has_key(v)
            got[v] = 1
        for i in xrange(len(self.hashes)):
            if self.places.has_key(i):  # restored from pickled
                self.check_targets[self.hashes[i]] = []
                if self.places[i] == i:
                    continue
                else:
                    assert not got.has_key(i)
                    self.out_of_place += 1
            if got.has_key(i):
                continue
            if self._waspre(i):
                if self.blocked[i]:
                    self.places[i] = i
                else:
                    self.check_list.append(i)
                continue
            if not self.check_hashes:
                self.failed('told file complete on start-up, but data is missing')
                return False
            self.holes.append(i)
            if self.blocked[i] or self.check_targets.has_key(self.hashes[i]):
                self.check_targets[self.hashes[i]] = [] # in case of a hash collision, discard
            else:
                self.check_targets[self.hashes[i]] = [i]
        self.check_total = len(self.check_list)
        self.check_numchecked = 0.0
        self.lastlen = self._piecelen(len(self.hashes) - 1)
        self.numchecked = 0.0
        return self.check_total > 0

    def _markgot(self, piece, pos):
        if DEBUG:
            print str(piece)+' at '+str(pos)
        self.places[piece] = pos
        self.have[piece] = True
        len = self._piecelen(piece)
        self.amount_obtained += len
        self.amount_left -= len
        self.amount_inactive -= len
        self.inactive_requests[piece] = None
        self.waschecked[piece] = self.check_hashes
        self.stat_numfound += 1

    def hashcheckfunc(self):
        if self.flag.isSet():
            return None
        if not self.check_list:
            return None
        
        i = self.check_list.pop(0)
        if not self.check_hashes:
            self._markgot(i, i)
        else:
            d1 = self.read_raw(i,0,self.lastlen)
            if d1 is None:
                return None
            sh = sha(d1[:])
            d1.release()
            sp = sh.digest()
            d2 = self.read_raw(i,self.lastlen,self._piecelen(i)-self.lastlen)
            if d2 is None:
                return None
            sh.update(d2[:])
            d2.release()
            s = sh.digest()
            if s == self.hashes[i]:
                self._markgot(i, i)
            elif ( self.check_targets.get(s)
                   and self._piecelen(i) == self._piecelen(self.check_targets[s][-1]) ):
                self._markgot(self.check_targets[s].pop(), i)
                self.out_of_place += 1
            elif ( not self.have[-1] and sp == self.hashes[-1]
                   and (i == len(self.hashes) - 1
                        or not self._waspre(len(self.hashes) - 1)) ):
                self._markgot(len(self.hashes) - 1, i)
                self.out_of_place += 1
            else:
                self.places[i] = i
        self.numchecked += 1
        if self.amount_left == 0:
            self.finished()
        return (self.numchecked / self.check_total)


    def init_movedata(self):
        if self.flag.isSet():
            return False
        if self.alloc_type != 'sparse':
            return False
        self.storage.top_off()  # sets file lengths to their final size
        self.movelist = []
        if self.out_of_place == 0:
            for i in self.holes:
                self.places[i] = i
            self.holes = []
            return False
        self.tomove = float(self.out_of_place)
        for i in xrange(len(self.hashes)):
            if not self.places.has_key(i):
                self.places[i] = i
            elif self.places[i] != i:
                self.movelist.append(i)
        self.holes = []
        return True

    def movedatafunc(self):
        if self.flag.isSet():
            return None
        if not self.movelist:
            return None
        i = self.movelist.pop(0)
        old = self.read_raw(self.places[i], 0, self._piecelen(i))
        if old is None:
            return None
        if not self.write_raw(i, 0, old):
            return None
        if self.double_check and self.have[i]:
            if self.triple_check:
                old.release()
                old = self.read_raw( i, 0, self._piecelen(i),
                                            flush_first = True )
                if old is None:
                    return None
            if sha(old[:]).digest() != self.hashes[i]:
                self.failed('download corrupted; please restart and resume')
                return None
        old.release()

        self.places[i] = i
        self.tomove -= 1
        return (self.tomove / self.out_of_place)

        
    def init_alloc(self):
        if self.flag.isSet():
            return False
        if not self.holes:
            return False
        self.numholes = float(len(self.holes))
        self.alloc_buf = chr(0xFF) * self.piece_size
        if self.alloc_type == 'pre-allocate':