4823: Fix tests broken by prior refactoring. Renamed 'api.py' to 'apisetup.py'
[arvados.git] / sdk / python / arvados / arvfile.py
index 456602fb9055017cbc3af862e75b473a85c6665c..d9d9cd287a3b2a81ae85ec283bc06c840f8f81b2 100644 (file)
@@ -2,7 +2,7 @@ import functools
 import os
 import zlib
 import bz2
-from .ranges import *
+from ._ranges import locators_and_ranges, replace_range, Range
 from arvados.retry import retry_method
 import config
 import hashlib
@@ -11,7 +11,9 @@ import threading
 import Queue
 import copy
 import errno
-from .errors import KeepWriteError
+from .errors import KeepWriteError, AssertionError
+from .keep import KeepLocator
+from _normalize_stream import normalize_stream
 
 def split(path):
     """Separate the stream name and file name in a /-separated stream path and
@@ -26,7 +28,7 @@ def split(path):
         stream_name, file_name = '.', path
     return stream_name, file_name
 
-class ArvadosFileBase(object):
+class _FileLikeObjectBase(object):
     def __init__(self, name, mode):
         self.name = name
         self.mode = mode
@@ -55,7 +57,7 @@ class ArvadosFileBase(object):
         self.closed = True
 
 
-class ArvadosFileReaderBase(ArvadosFileBase):
+class ArvadosFileReaderBase(_FileLikeObjectBase):
     class _NameAttribute(str):
         # The Python file API provides a plain .name attribute.
         # Older SDK provided a name() method.
@@ -79,7 +81,7 @@ class ArvadosFileReaderBase(ArvadosFileBase):
     def decompressed_name(self):
         return re.sub('\.(bz2|gz)$', '', self.name)
 
-    @ArvadosFileBase._before_close
+    @_FileLikeObjectBase._before_close
     def seek(self, pos, whence=os.SEEK_CUR):
         if whence == os.SEEK_CUR:
             pos += self._filepos
@@ -90,7 +92,7 @@ class ArvadosFileReaderBase(ArvadosFileBase):
     def tell(self):
         return self._filepos
 
-    @ArvadosFileBase._before_close
+    @_FileLikeObjectBase._before_close
     @retry_method
     def readall(self, size=2**20, num_retries=None):
         while True:
@@ -99,7 +101,7 @@ class ArvadosFileReaderBase(ArvadosFileBase):
                 break
             yield data
 
-    @ArvadosFileBase._before_close
+    @_FileLikeObjectBase._before_close
     @retry_method
     def readline(self, size=float('inf'), num_retries=None):
         cache_pos, cache_data = self._readline_cache
@@ -123,7 +125,7 @@ class ArvadosFileReaderBase(ArvadosFileBase):
         self._readline_cache = (self.tell(), data[nextline_index:])
         return data[:nextline_index]
 
-    @ArvadosFileBase._before_close
+    @_FileLikeObjectBase._before_close
     @retry_method
     def decompress(self, decompress, size, num_retries=None):
         for segment in self.readall(size, num_retries):
@@ -131,7 +133,7 @@ class ArvadosFileReaderBase(ArvadosFileBase):
             if data:
                 yield data
 
-    @ArvadosFileBase._before_close
+    @_FileLikeObjectBase._before_close
     @retry_method
     def readall_decompressed(self, size=2**20, num_retries=None):
         self.seek(0)
@@ -146,7 +148,7 @@ class ArvadosFileReaderBase(ArvadosFileBase):
         else:
             return self.readall(size, num_retries=num_retries)
 
-    @ArvadosFileBase._before_close
+    @_FileLikeObjectBase._before_close
     @retry_method
     def readlines(self, sizehint=float('inf'), num_retries=None):
         data = []
@@ -181,7 +183,7 @@ class StreamFileReader(ArvadosFileReaderBase):
         n = self.segments[-1]
         return n.range_start + n.range_size
 
-    @ArvadosFileBase._before_close
+    @_FileLikeObjectBase._before_close
     @retry_method
     def read(self, size, num_retries=None):
         """Read up to 'size' bytes from the stream, starting at the current file position"""
@@ -199,7 +201,7 @@ class StreamFileReader(ArvadosFileReaderBase):
         self._filepos += len(data)
         return data
 
-    @ArvadosFileBase._before_close
+    @_FileLikeObjectBase._before_close
     @retry_method
     def readfrom(self, start, size, num_retries=None):
         """Read up to 'size' bytes from the stream, starting at 'start'"""
@@ -213,14 +215,20 @@ class StreamFileReader(ArvadosFileReaderBase):
         return ''.join(data)
 
     def as_manifest(self):
-        from stream import normalize_stream
         segs = []
         for r in self.segments:
             segs.extend(self._stream.locators_and_ranges(r.locator, r.range_size))
         return " ".join(normalize_stream(".", {self.name: segs})) + "\n"
 
 
-class BufferBlock(object):
+def synchronized(orig_func):
+    @functools.wraps(orig_func)
+    def synchronized_wrapper(self, *args, **kwargs):
+        with self.lock:
+            return orig_func(self, *args, **kwargs)
+    return synchronized_wrapper
+
+class _BufferBlock(object):
     """A BufferBlock is a stand-in for a Keep block that is in the process of being
     written.
 
@@ -261,10 +269,12 @@ class BufferBlock(object):
         self.buffer_block = bytearray(starting_capacity)
         self.buffer_view = memoryview(self.buffer_block)
         self.write_pointer = 0
-        self.state = BufferBlock.WRITABLE
+        self._state = _BufferBlock.WRITABLE
         self._locator = None
         self.owner = owner
+        self.lock = threading.Lock()
 
+    @synchronized
     def append(self, data):
         """Append some data to the buffer.
 
@@ -272,7 +282,7 @@ class BufferBlock(object):
         buffer, doubling capacity as needed to accomdate all the data.
 
         """
-        if self.state == BufferBlock.WRITABLE:
+        if self._state == _BufferBlock.WRITABLE:
             while (self.write_pointer+len(data)) > len(self.buffer_block):
                 new_buffer_block = bytearray(len(self.buffer_block) * 2)
                 new_buffer_block[0:self.write_pointer] = self.buffer_block[0:self.write_pointer]
@@ -284,10 +294,27 @@ class BufferBlock(object):
         else:
             raise AssertionError("Buffer block is not writable")
 
+    @synchronized
+    def set_state(self, nextstate, loc=None):
+        if ((self._state == _BufferBlock.WRITABLE and nextstate == _BufferBlock.PENDING) or
+            (self._state == _BufferBlock.PENDING and nextstate == _BufferBlock.COMMITTED)):
+            self._state = nextstate
+            if self._state == _BufferBlock.COMMITTED:
+                self._locator = loc
+                self.buffer_view = None
+                self.buffer_block = None
+        else:
+            raise AssertionError("Invalid state change from %s to %s" % (self.state, state))
+
+    @synchronized
+    def state(self):
+        return self._state
+
     def size(self):
         """The amount of data written to the buffer."""
         return self.write_pointer
 
+    @synchronized
     def locator(self):
         """The Keep locator for this buffer's contents."""
         if self._locator is None:
@@ -295,13 +322,6 @@ class BufferBlock(object):
         return self._locator
 
 
-def synchronized(orig_func):
-    @functools.wraps(orig_func)
-    def synchronized_wrapper(self, *args, **kwargs):
-        with self.lock:
-            return orig_func(self, *args, **kwargs)
-    return synchronized_wrapper
-
 class NoopLock(object):
     def __enter__(self):
         return self
@@ -328,7 +348,7 @@ def must_be_writable(orig_func):
     return must_be_writable_wrapper
 
 
-class BlockManager(object):
+class _BlockManager(object):
     """BlockManager handles buffer blocks, background block uploads, and background
     block prefetch for a Collection of ArvadosFiles.
 
@@ -363,26 +383,29 @@ class BlockManager(object):
         """
         if blockid is None:
             blockid = "bufferblock%i" % len(self._bufferblocks)
-        bufferblock = BufferBlock(blockid, starting_capacity=starting_capacity, owner=owner)
+        bufferblock = _BufferBlock(blockid, starting_capacity=starting_capacity, owner=owner)
         self._bufferblocks[bufferblock.blockid] = bufferblock
         return bufferblock
 
     @synchronized
-    def dup_block(self, blockid, owner):
+    def dup_block(self, block, owner):
         """Create a new bufferblock in WRITABLE state, initialized with the content of
         an existing bufferblock.
 
-        :blockid:
-          the block to copy.  May be an existing buffer block id.
+        :block:
+          the buffer block to copy.
 
         :owner:
           ArvadosFile that owns the new block
 
         """
         new_blockid = "bufferblock%i" % len(self._bufferblocks)
-        block = self._bufferblocks[blockid]
-        bufferblock = BufferBlock(new_blockid, len(block), owner)
-        bufferblock.append(block)
+        with block.lock:
+            if block._state == _BufferBlock.COMMITTED:
+                raise AssertionError("Can only duplicate a writable or pending buffer block")
+
+            bufferblock = _BufferBlock(new_blockid, block.size(), owner)
+            bufferblock.append(block.buffer_view[0:block.size()])
         self._bufferblocks[bufferblock.blockid] = bufferblock
         return bufferblock
 
@@ -427,10 +450,9 @@ class BlockManager(object):
                     bufferblock = self._put_queue.get()
                     if bufferblock is None:
                         return
-                    bufferblock._locator = self._keep.put(bufferblock.buffer_view[0:bufferblock.write_pointer].tobytes())
-                    bufferblock.state = BufferBlock.COMMITTED
-                    bufferblock.buffer_view = None
-                    bufferblock.buffer_block = None
+                    loc = self._keep.put(bufferblock.buffer_view[0:bufferblock.write_pointer].tobytes())
+                    bufferblock.set_state(_BufferBlock.COMMITTED, loc)
+
                 except Exception as e:
                     print e
                     self._put_errors.put((bufferblock.locator(), e))
@@ -462,10 +484,14 @@ class BlockManager(object):
                     thread.start()
 
         # Mark the block as PENDING so to disallow any more appends.
-        block.state = BufferBlock.PENDING
+        block.set_state(_BufferBlock.PENDING)
         self._put_queue.put(block)
 
-    def get_block(self, locator, num_retries, cache_only=False):
+    @synchronized
+    def get_bufferblock(self, locator):
+        return self._bufferblocks.get(locator)
+
+    def get_block_contents(self, locator, num_retries, cache_only=False):
         """Fetch a block.
 
         First checks to see if the locator is a BufferBlock and return that, if
@@ -475,7 +501,7 @@ class BlockManager(object):
         with self.lock:
             if locator in self._bufferblocks:
                 bufferblock = self._bufferblocks[locator]
-                if bufferblock.state != BufferBlock.COMMITTED:
+                if bufferblock.state() != _BufferBlock.COMMITTED:
                     return bufferblock.buffer_view[0:bufferblock.write_pointer].tobytes()
                 else:
                     locator = bufferblock._locator
@@ -496,12 +522,13 @@ class BlockManager(object):
             items = self._bufferblocks.items()
 
         for k,v in items:
-            if v.state == BufferBlock.WRITABLE:
+            if v.state() == _BufferBlock.WRITABLE:
                 self.commit_bufferblock(v)
 
         with self.lock:
             if self._put_queue is not None:
                 self._put_queue.join()
+
                 if not self._put_errors.empty():
                     err = []
                     try:
@@ -587,24 +614,29 @@ class ArvadosFile(object):
     def clone(self, new_parent):
         """Make a copy of this file."""
         cp = ArvadosFile(new_parent)
-
-        map_loc = {}
-        for r in self._segments:
-            new_loc = r.locator
-            if self.parent._my_block_manager().is_bufferblock(r.locator):
-                if r.locator not in map_loc:
-                    map_loc[r.locator] = self.parent._my_block_manager().dup_block(r.locator, cp).blockid
-                new_loc = map_loc[r.locator]
-
-            cp._segments.append(Range(new_loc, r.range_start, r.range_size, r.segment_offset))
-
+        cp.replace_contents(self)
         return cp
 
     @must_be_writable
     @synchronized
     def replace_contents(self, other):
         """Replace segments of this file with segments from another `ArvadosFile` object."""
-        self._segments = other.segments()
+
+        map_loc = {}
+        self._segments = []
+        for other_segment in other.segments():
+            new_loc = other_segment.locator
+            if other.parent._my_block_manager().is_bufferblock(other_segment.locator):
+                if other_segment.locator not in map_loc:
+                    bufferblock = other.parent._my_block_manager().get_bufferblock(other_segment.locator)
+                    if bufferblock.state() != _BufferBlock.WRITABLE:
+                        map_loc[other_segment.locator] = bufferblock.locator()
+                    else:
+                        map_loc[other_segment.locator] = self.parent._my_block_manager().dup_block(bufferblock, self).blockid
+                new_loc = map_loc[other_segment.locator]
+
+            self._segments.append(Range(new_loc, other_segment.range_start, other_segment.range_size, other_segment.segment_offset))
+
         self._modified = True
 
     def __eq__(self, other):
@@ -613,9 +645,29 @@ class ArvadosFile(object):
         if not isinstance(other, ArvadosFile):
             return False
 
-        s = other.segments()
+        othersegs = other.segments()
         with self.lock:
-            return self._segments == s
+            if len(self._segments) != len(othersegs):
+                return False
+            for i in xrange(0, len(othersegs)):
+                seg1 = self._segments[i]
+                seg2 = othersegs[i]
+                loc1 = seg1.locator
+                loc2 = seg2.locator
+
+                if self.parent._my_block_manager().is_bufferblock(loc1):
+                    loc1 = self.parent._my_block_manager().get_bufferblock(loc1).locator()
+
+                if other.parent._my_block_manager().is_bufferblock(loc2):
+                    loc2 = other.parent._my_block_manager().get_bufferblock(loc2).locator()
+
+                if (KeepLocator(loc1).stripped() != KeepLocator(loc2).stripped() or
+                    seg1.range_start != seg2.range_start or
+                    seg1.range_size != seg2.range_size or
+                    seg1.segment_offset != seg2.segment_offset):
+                    return False
+
+        return True
 
     def __ne__(self, other):
         return not self.__eq__(other)
@@ -674,7 +726,7 @@ class ArvadosFile(object):
 
         data = []
         for lr in readsegs:
-            block = self.parent._my_block_manager().get_block(lr.locator, num_retries=num_retries, cache_only=bool(data))
+            block = self.parent._my_block_manager().get_block_contents(lr.locator, num_retries=num_retries, cache_only=bool(data))
             if block:
                 data.append(block[lr.segment_offset:lr.segment_offset+lr.segment_size])
             else:
@@ -727,7 +779,7 @@ class ArvadosFile(object):
 
         self._modified = True
 
-        if self._current_bblock is None or self._current_bblock.state != BufferBlock.WRITABLE:
+        if self._current_bblock is None or self._current_bblock.state() != _BufferBlock.WRITABLE:
             self._current_bblock = self.parent._my_block_manager().alloc_bufferblock(owner=self)
 
         if (self._current_bblock.size() + len(data)) > config.KEEP_BLOCK_SIZE:
@@ -766,6 +818,26 @@ class ArvadosFile(object):
         else:
             return 0
 
+
+    @synchronized
+    def manifest_text(self, stream_name=".", portable_locators=False, normalize=False):
+        buf = ""
+        item = self
+        filestream = []
+        for segment in item.segments:
+            loc = segment.locator
+            if loc.startswith("bufferblock"):
+                loc = item._bufferblocks[loc].calculate_locator()
+            if portable_locators:
+                loc = KeepLocator(loc).stripped()
+            filestream.append(LocatorAndRange(loc, locator_block_size(loc),
+                                 segment.segment_offset, segment.range_size))
+        stream[stream_name] = filestream
+        buf += ' '.join(normalize_stream(stream_name, stream))
+        buf += "\n"
+        return buf
+
+
 class ArvadosFileReader(ArvadosFileReaderBase):
     """Wraps ArvadosFile in a file-like object supporting reading only.
 
@@ -781,7 +853,10 @@ class ArvadosFileReader(ArvadosFileReaderBase):
     def size(self):
         return self.arvadosfile.size()
 
-    @ArvadosFileBase._before_close
+    def stream_name(self):
+        return self.arvadosfile.parent.stream_name()
+
+    @_FileLikeObjectBase._before_close
     @retry_method
     def read(self, size, num_retries=None):
         """Read up to `size` bytes from the stream, starting at the current file position."""
@@ -789,7 +864,7 @@ class ArvadosFileReader(ArvadosFileReaderBase):
         self._filepos += len(data)
         return data
 
-    @ArvadosFileBase._before_close
+    @_FileLikeObjectBase._before_close
     @retry_method
     def readfrom(self, offset, size, num_retries=None):
         """Read up to `size` bytes from the stream, starting at the current file position."""
@@ -810,7 +885,7 @@ class ArvadosFileWriter(ArvadosFileReader):
     def __init__(self, arvadosfile, name, mode, num_retries=None):
         super(ArvadosFileWriter, self).__init__(arvadosfile, name, mode, num_retries=num_retries)
 
-    @ArvadosFileBase._before_close
+    @_FileLikeObjectBase._before_close
     @retry_method
     def write(self, data, num_retries=None):
         if self.mode[0] == "a":
@@ -819,7 +894,7 @@ class ArvadosFileWriter(ArvadosFileReader):
             self.arvadosfile.writeto(self._filepos, data, num_retries)
             self._filepos += len(data)
 
-    @ArvadosFileBase._before_close
+    @_FileLikeObjectBase._before_close
     @retry_method
     def writelines(self, seq, num_retries=None):
         for s in seq: