storage.py

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

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

from BitTornado.piecebuffer import BufferPool
from threading import Lock
from time import time, strftime, localtime
import os
from os.path import exists, getsize, getmtime, basename
from traceback import print_exc
try:
    from os import fsync
except ImportError:
    fsync = lambda x: None
from bisect import bisect
    
try:
    True
except:
    True = 1
    False = 0

DEBUG = False

MAXREADSIZE = 32768
MAXLOCKSIZE = 1000000000L
MAXLOCKRANGE = 3999999999L   # only lock first 4 gig of file

_pool = BufferPool()
PieceBuffer = _pool.new

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

class Storage:
    def __init__(self, files, piece_length, doneflag, config,
                 disabled_files = None):
        # can raise IOError and ValueError
        self.files = files
        self.piece_length = piece_length
        self.doneflag = doneflag
        self.disabled = [False] * len(files)
        self.file_ranges = []
        self.disabled_ranges = []
        self.working_ranges = []
        numfiles = 0
        total = 0l
        so_far = 0l
        self.handles = {}
        self.whandles = {}
        self.tops = {}
        self.sizes = {}
        self.mtimes = {}
        if config.get('lock_files', True):
            self.lock_file, self.unlock_file = self._lock_file, self._unlock_file
        else:
            self.lock_file, self.unlock_file = lambda x1,x2: None, lambda x1,x2: None
        self.lock_while_reading = config.get('lock_while_reading', False)
        self.lock = Lock()

        if not disabled_files:
            disabled_files = [False] * len(files)

        for i in xrange(len(files)):
            file, length = files[i]
            if doneflag.isSet():    # bail out if doneflag is set
                return
            self.disabled_ranges.append(None)
            if length == 0:
                self.file_ranges.append(None)
                self.working_ranges.append([])
            else:
                range = (total, total + length, 0, file)
                self.file_ranges.append(range)
                self.working_ranges.append([range])
                numfiles += 1
                total += length
                if disabled_files[i]:
                    l = 0
                else:
                    if exists(file):
                        l = getsize(file)
                        if l > length:
                            h = open(file, 'rb+')
                            h.truncate(length)
                            h.flush()
                            h.close()
                            l = length
                    else:
                        l = 0
                        h = open(file, 'wb+')
                        h.flush()
                        h.close()
                    self.mtimes[file] = getmtime(file)
                self.tops[file] = l
                self.sizes[file] = length
                so_far += l

        self.total_length = total
        self._reset_ranges()

        self.max_files_open = config['max_files_open']
        if self.max_files_open > 0 and numfiles > self.max_files_open:
            self.handlebuffer = []
        else:
            self.handlebuffer = None


    if os.name == 'nt':
        def _lock_file(self, name, f):
            import msvcrt
            for p in range(0, min(self.sizes[name],MAXLOCKRANGE), MAXLOCKSIZE):
                f.seek(p)
                msvcrt.locking(f.fileno(), msvcrt.LK_LOCK,
                               min(MAXLOCKSIZE,self.sizes[name]-p))

        def _unlock_file(self, name, f):
            import msvcrt
            for p in range(0, min(self.sizes[name],MAXLOCKRANGE), MAXLOCKSIZE):
                f.seek(p)
                msvcrt.locking(f.fileno(), msvcrt.LK_UNLCK,
                               min(MAXLOCKSIZE,self.sizes[name]-p))

    elif os.name == 'posix':
        def _lock_file(self, name, f):
            import fcntl
            fcntl.flock(f.fileno(), fcntl.LOCK_EX)

        def _unlock_file(self, name, f):
            import fcntl
            fcntl.flock(f.fileno(), fcntl.LOCK_UN)

    else:
        def _lock_file(self, name, f):
            pass
        def _unlock_file(self, name, f):
            pass


    def was_preallocated(self, pos, length):
        for file, begin, end in self._intervals(pos, length):
            if self.tops.get(file, 0) < end:
                return False
        return True


    def _sync(self, file):
        self._close(file)
        if self.handlebuffer:
            self.handlebuffer.remove(file)

    def sync(self):
        # may raise IOError or OSError
        for file in self.whandles.keys():
            self._sync(file)


    def set_readonly(self, f=None):
        if f is None:
            self.sync()
            return
        file = self.files[f][0]
        if self.whandles.has_key(file):
            self._sync(file)
            

    def get_total_length(self):
        return self.total_length


    def _open(self, file, mode):
        if self.mtimes.has_key(file):
            try:
              if self.handlebuffer is not None:
                assert getsize(file) == self.tops[file]
                newmtime = getmtime(file)
                oldmtime = self.mtimes[file]
                assert newmtime <= oldmtime+1
                assert newmtime >= oldmtime-1
            except:
                if DEBUG:
                    print ( file+' modified: '
                            +strftime('(%x %X)',localtime(self.mtimes[file]))
                            +strftime(' != (%x %X) ?',localtime(getmtime(file))) )
                raise IOError('modified during download')
        try:
            return open(file, mode)
        except:
            if DEBUG:
                print_exc()
            raise


    def _close(self, file):
        f = self.handles[file]
        del self.handles[file]
        if self.whandles.has_key(file):
            del self.whandles[file]
            f.flush()
            self.unlock_file(file, f)
            f.close()
            self.tops[file] = getsize(file)
            self.mtimes[file] = getmtime(file)
        else:
            if self.lock_while_reading:
                self.unlock_file(file, f)
            f.close()


    def _close_file(self, file):
        if not self.handles.has_key(file):
            return
        self._close(file)
        if self.handlebuffer:
            self.handlebuffer.remove(file)
        

    def _get_file_handle(self, file, for_write):
        if self.handles.has_key(file):
            if for_write and not self.whandles.has_key(file):
                self._close(file)
                try:
                    f = self._open(file, 'rb+')
                    self.handles[file] = f
                    self.whandles[file] = 1
                    self.lock_file(file, f)
                except (IOError, OSError), e:
                    if DEBUG:
                        print_exc()
                    raise IOError('unable to reopen '+file+': '+str(e))

            if self.handlebuffer:
                if self.handlebuffer[-1] != file:
                    self.handlebuffer.remove(file)
                    self.handlebuffer.append(file)
            elif self.handlebuffer is not None:
                self.handlebuffer.append(file)
        else:
            try:
                if for_write:
                    f = self._open(file, 'rb+')
                    self.handles[file] = f
                    self.whandles[file] = 1
                    self.lock_file(file, f)
                else:
                    f = self._open(file, 'rb')
                    self.handles[file] = f
                    if self.lock_while_reading:
                        self.lock_file(file, f)
            except (IOError, OSError), e:
                if DEBUG:
                    print_exc()
                raise IOError('unable to open '+file+': '+str(e))
            
            if self.handlebuffer is not None:
                self.handlebuffer.append(file)
                if len(self.handlebuffer) > self.max_files_open:
                    self._close(self.handlebuffer.pop(0))

        return self.handles[file]


    def _reset_ranges(self):
        self.ranges = []
        for l in self.working_ranges:
            self.ranges.extend(l)
            self.begins = [i[0] for i in self.ranges]

    def _intervals(self, pos, amount):
        r = []
        stop = pos + amount
        p = bisect(self.begins, pos) - 1
        while p < len(self.ranges):
            begin, end, offset, file = self.ranges[p]
            if begin >= stop:
                break
            r.append(( file,
                       offset + max(pos, begin) - begin,
                       offset + min(end, stop) - begin   ))
            p += 1
        return r


    def read(self, pos, amount, flush_first = False):
        r = PieceBuffer()
        for file, pos, end in self._intervals(pos, amount):
            if DEBUG:
                print 'reading '+file+' from '+str(pos)+' to '+str(end)
            self.lock.acquire()
            h = self._get_file_handle(file, False)
            if flush_first and self.whandles.has_key(file):
                h.flush()
                fsync(h)
            h.seek(pos)