storage.py

BitTorrentABC-Linux-V.2.4.3源码

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

from sha import sha
from cStringIO import StringIO
from threading import Lock
from time import time
from bisect import bisect_right
true = 1
false = 0

DEBUG = false

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

class Storage:
    def __init__(self, files, open, exists, getsize, statusfunc, doneflag, config, failed):
        # can raise IOError and ValueError
        self.doneflag = doneflag
        self.failed = failed
        self.ranges = []
        numfiles = 0
        total = 0l
        so_far = 0l
        self.handles = {}
        self.whandles = {}
        self.tops = {}
        self.lock = Lock()

        for file, length in files:
            #if doneflag.isSet():    # bail out if doneflag is set
            #    return
            if length != 0:
                self.ranges.append((total, total + length, file))
                numfiles += 1
                total += length
                if exists(file):
                    l = getsize(file)
                    if l > length:
                        h = open(file, 'rb+')
                        h.truncate(length)
                        h.close()
                        l = length
                else:
                    l = 0
                    open(file, 'wb+').close()
                    
                self.tops[file] = l
                so_far += l

        self.begins = [i[0] for i in self.ranges]
        self.total_length = total


        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
                    

    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 set_readonly(self):
        # may raise IOError or OSError
        for file in self.whandles.keys():
            self.lock.acquire()
            old = self.handles[file]
            old.flush()
            old.close()
            self.handles[file] = open(file, 'rb')
            self.lock.release()
        self.whandles = {}

    def get_total_length(self):
        return self.total_length


    def _get_file_handle(self, file, for_write):
        if self.handles.has_key(file):
            if for_write and not self.whandles.has_key(file):
                f = self.handles[file]
                del self.handles[file]
                f.close()
                try:
                    self.handles[file] = open(file, 'rb+')
                except IOError, e:
                    self.failed('unable to reopen '+file+': '+str(e))
                    return None
                self.whandles[file] = 1
                
            if self.handlebuffer is not None:
                pos = self.handlebuffer.index(file)
                if pos < len(self.handlebuffer)-1:
                    del self.handlebuffer[pos]
                    self.handlebuffer.append(file)
        else:
            try:
                if for_write:
                    self.handles[file] = open(file, 'rb+')
                    self.whandles[file] = 1
                else:
                    self.handles[file] = open(file, 'rb')
            except IOError, e:
                self.failed('unable to open '+file+': '+str(e))
                return None
            
            if self.handlebuffer is not None:
                self.handlebuffer.append(file)
                if len(self.handlebuffer) > self.max_files_open:
                    x = self.handlebuffer.pop(0)
                    if self.whandles.has_key(x):
                        del self.whandles[x]
                        self.handles[x].flush()
                    f = self.handles[x]
                    del self.handles[x]
                    f.close()

        return self.handles[file]

    def _intervals(self, pos, amount):
        r = []
        stop = pos + amount
        p = bisect_right(self.begins, pos) - 1
        while p < len(self.ranges) and self.ranges[p][0] < stop:
            begin, end, file = self.ranges[p]
            r.append((file, max(pos, begin) - begin, min(end, stop) - begin))
            p += 1
        return r

    def read(self, pos, amount):
        r = []
        for file, pos, end in self._intervals(pos, amount):
            self.lock.acquire()
            h = self._get_file_handle(file, false)
            h.seek(pos)
            r.append(h.read(end - pos))
            self.lock.release()
        return ''.join(r)

    def write(self, pos, s):
        # might raise an IOError
        total = 0
        for file, begin, end in self._intervals(pos, len(s)):
            self.lock.acquire()
            h = self._get_file_handle(file, true)
            if h is None:
                return
            h.seek(begin)
            h.write(s[total: total + end - begin])
            self.lock.release()
            total += end - begin

    def top_off(self):
        for begin, end, file in self.ranges:
            l = end - begin
            if l < self.tops.get(file, 0):
                self.lock.acquire()
                h = self._get_file_handle(file, true)
                if h is None:
                    return None
                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 h in self.handles.values():
            try:
                h.close()
            except:
                pass

def lrange(a, b, c):
    r = []
    while a < b:
        r.append(a)
        a += c
    return r

# everything below is for testing

from fakeopen import FakeOpen

def test_Storage_simple():
    f = FakeOpen()
    m = Storage([('a', 5)], f.open, f.exists, f.getsize, dummy_status)
    assert f.files.keys() == ['a']
    m.write(0, 'abc')
    assert m.read(0, 3) == 'abc'
    m.write(2, 'abc')
    assert m.read(2, 3) == 'abc'
    m.write(1, 'abc')
    assert m.read(0, 5) == 'aabcc'
    
def test_Storage_multiple():
    f = FakeOpen()
    m = Storage([('a', 5), ('2', 4), ('c', 3)], 
        f.open, f.exists, f.getsize, dummy_status)
    x = f.files.keys()
    x.sort()
    assert x == ['2', 'a', 'c']
    m.write(3, 'abc')
    assert m.read(3, 3) == 'abc'
    m.write(5, 'ab')
    assert m.read(4, 3) == 'bab'
    m.write(3, 'pqrstuvw')
    assert m.read(3, 8) == 'pqrstuvw'
    m.write(3, 'abcdef')
    assert m.read(3, 7) == 'abcdefv'

def test_Storage_zero():
    f = FakeOpen()
    Storage([('a', 0)], f.open, f.exists, f.getsize, dummy_status)
    assert f.files == {'a': []}

def test_resume_zero():
    f = FakeOpen({'a': ''})
    Storage([('a', 0)], f.open, f.exists, f.getsize, dummy_status)
    assert f.files == {'a': []}

def test_Storage_with_zero():
    f = FakeOpen()
    m = Storage([('a', 3), ('b', 0), ('c', 3)], 
        f.open, f.exists, f.getsize, dummy_status)
    m.write(2, 'abc')
    assert m.read(2, 3) == 'abc'
    x = f.files.keys()
    x.sort()
    assert x == ['a', 'b', 'c']
    assert len(f.files['a']) == 3
    assert len(f.files['b']) == 0

def test_Storage_resume():
    f = FakeOpen({'a': 'abc'})
    m = Storage([('a', 4)], 
        f.open, f.exists, f.getsize, dummy_status)
    assert f.files.keys() == ['a']
    assert m.read(0, 3) == 'abc'

def test_Storage_mixed_resume():
    f = FakeOpen({'b': 'abc'})
    m = Storage([('a', 3), ('b', 4)], 
        f.open, f.exists, f.getsize, dummy_status)
    x = f.files.keys()
    x.sort()
    assert x == ['a', 'b']
    assert m.read(3, 3) == 'abc'