storage.py

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

            while pos < end:
                length = min(end-pos, MAXREADSIZE)
                data = h.read(length)
                if len(data) != length:
                    raise IOError('error reading data from '+file)
                r.append(data)
                pos += length
            self.lock.release()
        return r

    def write(self, pos, s):
        # might raise an IOError
        total = 0
        for file, begin, end in self._intervals(pos, len(s)):
            if DEBUG:
                print 'writing '+file+' from '+str(pos)+' to '+str(end)
            self.lock.acquire()
            h = self._get_file_handle(file, True)
            h.seek(begin)
            h.write(s[total: total + end - begin])
            self.lock.release()
            total += end - begin

    def top_off(self):
        for begin, end, offset, file in self.ranges:
            l = offset + end - begin
            if l > self.tops.get(file, 0):
                self.lock.acquire()
                h = self._get_file_handle(file, True)
                h.seek(l-1)
                h.write(chr(0xFF))
                self.lock.release()

    def flush(self):
        # may raise IOError or OSError
        for file in self.whandles.keys():
            self.lock.acquire()
            self.handles[file].flush()
            self.lock.release()

    def close(self):
        for file, f in self.handles.items():
            try:
                self.unlock_file(file, f)
            except:
                pass
            try:
                f.close()
            except:
                pass
        self.handles = {}
        self.whandles = {}
        self.handlebuffer = None


    def _get_disabled_ranges(self, f):
        if not self.file_ranges[f]:
            return ((),(),())
        r = self.disabled_ranges[f]
        if r:
            return r
        start, end, offset, file = self.file_ranges[f]
        if DEBUG:
            print 'calculating disabled range for '+self.files[f][0]
            print 'bytes: '+str(start)+'-'+str(end)
            print 'file spans pieces '+str(int(start/self.piece_length))+'-'+str(int((end-1)/self.piece_length)+1)
        pieces = range( int(start/self.piece_length),
                        int((end-1)/self.piece_length)+1 )
        offset = 0
        disabled_files = []
        if len(pieces) == 1:
            if ( start % self.piece_length == 0
                 and end % self.piece_length == 0 ):   # happens to be a single,
                                                       # perfect piece
                working_range = [(start, end, offset, file)]
                update_pieces = []
            else:
                midfile = os.path.join(self.bufferdir,str(f))
                working_range = [(start, end, 0, midfile)]
                disabled_files.append((midfile, start, end))
                length = end - start
                self.sizes[midfile] = length
                piece = pieces[0]
                update_pieces = [(piece, start-(piece*self.piece_length), length)]
        else:
            update_pieces = []
            if start % self.piece_length != 0:  # doesn't begin on an even piece boundary
                end_b = pieces[1]*self.piece_length
                startfile = os.path.join(self.bufferdir,str(f)+'b')
                working_range_b = [ ( start, end_b, 0, startfile ) ]
                disabled_files.append((startfile, start, end_b))
                length = end_b - start
                self.sizes[startfile] = length
                offset = length
                piece = pieces.pop(0)
                update_pieces.append((piece, start-(piece*self.piece_length), length))
            else:
                working_range_b = []
            if f  != len(self.files)-1 and end % self.piece_length != 0:
                                                # doesn't end on an even piece boundary
                start_e = pieces[-1] * self.piece_length
                endfile = os.path.join(self.bufferdir,str(f)+'e')
                working_range_e = [ ( start_e, end, 0, endfile ) ]
                disabled_files.append((endfile, start_e, end))
                length = end - start_e
                self.sizes[endfile] = length
                piece = pieces.pop(-1)
                update_pieces.append((piece, 0, length))
            else:
                working_range_e = []
            if pieces:
                working_range_m = [ ( pieces[0]*self.piece_length,
                                      (pieces[-1]+1)*self.piece_length,
                                      offset, file ) ]
            else:
                working_range_m = []
            working_range = working_range_b + working_range_m + working_range_e

        if DEBUG:            
            print str(working_range)
            print str(update_pieces)
        r = (tuple(working_range), tuple(update_pieces), tuple(disabled_files))
        self.disabled_ranges[f] = r
        return r
        

    def set_bufferdir(self, dir):
        self.bufferdir = dir

    def enable_file(self, f):
        if not self.disabled[f]:
            return
        self.disabled[f] = False
        r = self.file_ranges[f]
        if not r:
            return
        file = r[3]
        if not exists(file):
            h = open(file, 'wb+')
            h.flush()
            h.close()
        if not self.tops.has_key(file):
            self.tops[file] = getsize(file)
        if not self.mtimes.has_key(file):
            self.mtimes[file] = getmtime(file)
        self.working_ranges[f] = [r]

    def disable_file(self, f):
        if self.disabled[f]:
            return
        self.disabled[f] = True
        r = self._get_disabled_ranges(f)
        if not r:
            return
        for file, begin, end in r[2]:
            if not os.path.isdir(self.bufferdir):
                os.makedirs(self.bufferdir)
            if not exists(file):
                h = open(file, 'wb+')
                h.flush()
                h.close()
            if not self.tops.has_key(file):
                self.tops[file] = getsize(file)
            if not self.mtimes.has_key(file):
                self.mtimes[file] = getmtime(file)
        self.working_ranges[f] = r[0]

    reset_file_status = _reset_ranges


    def get_piece_update_list(self, f):
        return self._get_disabled_ranges(f)[1]


    def delete_file(self, f):
        try:
            os.remove(self.files[f][0])
        except:
            pass


    '''
    Pickled data format:

    d['files'] = [ file #, size, mtime {, file #, size, mtime...} ]
                    file # in torrent, and the size and last modification
                    time for those files.  Missing files are either empty
                    or disabled.
    d['partial files'] = [ name, size, mtime... ]
                    Names, sizes and last modification times of files containing
                    partial piece data.  Filenames go by the following convention:
                    {file #, 0-based}{nothing, "b" or "e"}
                    eg: "0e" "3" "4b" "4e"
                    Where "b" specifies the partial data for the first piece in
                    the file, "e" the last piece, and no letter signifying that
                    the file is disabled but is smaller than one piece, and that
                    all the data is cached inside so adjacent files may be
                    verified.
    '''
    def pickle(self):
        files = []
        pfiles = []
        for i in xrange(len(self.files)):
            if not self.files[i][1]:    # length == 0
                continue
            if self.disabled[i]:
                for file, start, end in self._get_disabled_ranges(i)[2]:
                    pfiles.extend([basename(file),getsize(file),int(getmtime(file))])
                continue
            file = self.files[i][0]
            files.extend([i,getsize(file),int(getmtime(file))])
        return {'files': files, 'partial files': pfiles}


    def unpickle(self, data):
        # assume all previously-disabled files have already been disabled
        try:
            files = {}
            pfiles = {}
            l = data['files']
            assert len(l) % 3 == 0
            l = [l[x:x+3] for x in xrange(0,len(l),3)]
            for f, size, mtime in l:
                files[f] = (size, mtime)
            l = data.get('partial files',[])
            assert len(l) % 3 == 0
            l = [l[x:x+3] for x in xrange(0,len(l),3)]
            for file, size, mtime in l:
                pfiles[file] = (size, mtime)

            valid_pieces = {}
            for i in xrange(len(self.files)):
                if self.disabled[i]:
                    continue
                r = self.file_ranges[i]
                if not r:
                    continue
                start, end, offset, file =r
                if DEBUG:
                    print 'adding '+file
                for p in xrange( int(start/self.piece_length),
                                 int((end-1)/self.piece_length)+1 ):
                    valid_pieces[p] = 1

            if DEBUG:
                print valid_pieces.keys()
            
            def test(old, size, mtime):
                oldsize, oldmtime = old
                if size != oldsize:
                    return False
                if mtime > oldmtime+1:
                    return False
                if mtime < oldmtime-1:
                    return False
                return True

            for i in xrange(len(self.files)):
                if self.disabled[i]:
                    for file, start, end in self._get_disabled_ranges(i)[2]:
                        f1 = basename(file)
                        if ( not pfiles.has_key(f1)
                             or not test(pfiles[f1],getsize(file),getmtime(file)) ):
                            if DEBUG:
                                print 'removing '+file
                            for p in xrange( int(start/self.piece_length),
                                             int((end-1)/self.piece_length)+1 ):
                                if valid_pieces.has_key(p):
                                    del valid_pieces[p]
                    continue
                file, size = self.files[i]
                if not size:
                    continue
                if ( not files.has_key(i)
                     or not test(files[i],getsize(file),getmtime(file)) ):
                    start, end, offset, file = self.file_ranges[i]
                    if DEBUG:
                        print 'removing '+file
                    for p in xrange( int(start/self.piece_length),
                                     int((end-1)/self.piece_length)+1 ):
                        if valid_pieces.has_key(p):
                            del valid_pieces[p]
        except:
            if DEBUG:
                print_exc()
            return []

        if DEBUG:
            print valid_pieces.keys()                        
        return valid_pieces.keys()