sdk/python/arvados/arvfile.py

   1 # Copyright (C) The Arvados Authors. All rights reserved.
   2 #
   3 # SPDX-License-Identifier: Apache-2.0
   4
   5 from __future__ import absolute_import
   6 from __future__ import division
   7 from future import standard_library
   8 from future.utils import listitems, listvalues
   9 standard_library.install_aliases()
  10 from builtins import range
  11 from builtins import object
  12 import bz2
  13 import collections
  14 import copy
  15 import errno
  16 import functools
  17 import hashlib
  18 import logging
  19 import os
  20 import queue
  21 import re
  22 import sys
  23 import threading
  24 import uuid
  25 import zlib
  26
  27 from . import config
  28 from .errors import KeepWriteError, AssertionError, ArgumentError
  29 from .keep import KeepLocator
  30 from ._normalize_stream import normalize_stream
  31 from ._ranges import locators_and_ranges, replace_range, Range, LocatorAndRange
  32 from .retry import retry_method
  33
  34 MOD = "mod"
  35 WRITE = "write"
  36
  37 _logger = logging.getLogger('arvados.arvfile')
  38
  39 def split(path):
  40     """split(path) -> streamname, filename
  41
  42     Separate the stream name and file name in a /-separated stream path and
  43     return a tuple (stream_name, file_name).  If no stream name is available,
  44     assume '.'.
  45
  46     """
  47     try:
  48         stream_name, file_name = path.rsplit('/', 1)
  49     except ValueError:  # No / in string
  50         stream_name, file_name = '.', path
  51     return stream_name, file_name
  52
  53
  54 class UnownedBlockError(Exception):
  55     """Raised when there's an writable block without an owner on the BlockManager."""
  56     pass
  57
  58
  59 class _FileLikeObjectBase(object):
  60     def __init__(self, name, mode):
  61         self.name = name
  62         self.mode = mode
  63         self.closed = False
  64
  65     @staticmethod
  66     def _before_close(orig_func):
  67         @functools.wraps(orig_func)
  68         def before_close_wrapper(self, *args, **kwargs):
  69             if self.closed:
  70                 raise ValueError("I/O operation on closed stream file")
  71             return orig_func(self, *args, **kwargs)
  72         return before_close_wrapper
  73
  74     def __enter__(self):
  75         return self
  76
  77     def __exit__(self, exc_type, exc_value, traceback):
  78         try:
  79             self.close()
  80         except Exception:
  81             if exc_type is None:
  82                 raise
  83
  84     def close(self):
  85         self.closed = True
  86
  87
  88 class ArvadosFileReaderBase(_FileLikeObjectBase):
  89     def __init__(self, name, mode, num_retries=None):
  90         super(ArvadosFileReaderBase, self).__init__(name, mode)
  91         self._filepos = 0
  92         self.num_retries = num_retries
  93         self._readline_cache = (None, None)
  94
  95     def __iter__(self):
  96         while True:
  97             data = self.readline()
  98             if not data:
  99                 break
 100             yield data
 101
 102     def decompressed_name(self):
 103         return re.sub('\.(bz2|gz)$', '', self.name)
 104
 105     @_FileLikeObjectBase._before_close
 106     def seek(self, pos, whence=os.SEEK_SET):
 107         if whence == os.SEEK_CUR:
 108             pos += self._filepos
 109         elif whence == os.SEEK_END:
 110             pos += self.size()
 111         if pos < 0:
 112             raise IOError(errno.EINVAL, "Tried to seek to negative file offset.")
 113         self._filepos = pos
 114         return self._filepos
 115
 116     def tell(self):
 117         return self._filepos
 118
 119     def readable(self):
 120         return True
 121
 122     def writable(self):
 123         return False
 124
 125     def seekable(self):
 126         return True
 127
 128     @_FileLikeObjectBase._before_close
 129     @retry_method
 130     def readall(self, size=2**20, num_retries=None):
 131         while True:
 132             data = self.read(size, num_retries=num_retries)
 133             if len(data) == 0:
 134                 break
 135             yield data
 136
 137     @_FileLikeObjectBase._before_close
 138     @retry_method
 139     def readline(self, size=float('inf'), num_retries=None):
 140         cache_pos, cache_data = self._readline_cache
 141         if self.tell() == cache_pos:
 142             data = [cache_data]
 143             self._filepos += len(cache_data)
 144         else:
 145             data = [b'']
 146         data_size = len(data[-1])
 147         while (data_size < size) and (b'\n' not in data[-1]):
 148             next_read = self.read(2 ** 20, num_retries=num_retries)
 149             if not next_read:
 150                 break
 151             data.append(next_read)
 152             data_size += len(next_read)
 153         data = b''.join(data)
 154         try:
 155             nextline_index = data.index(b'\n') + 1
 156         except ValueError:
 157             nextline_index = len(data)
 158         nextline_index = min(nextline_index, size)
 159         self._filepos -= len(data) - nextline_index
 160         self._readline_cache = (self.tell(), data[nextline_index:])
 161         return data[:nextline_index].decode()
 162
 163     @_FileLikeObjectBase._before_close
 164     @retry_method
 165     def decompress(self, decompress, size, num_retries=None):
 166         for segment in self.readall(size, num_retries=num_retries):
 167             data = decompress(segment)
 168             if data:
 169                 yield data
 170
 171     @_FileLikeObjectBase._before_close
 172     @retry_method
 173     def readall_decompressed(self, size=2**20, num_retries=None):
 174         self.seek(0)
 175         if self.name.endswith('.bz2'):
 176             dc = bz2.BZ2Decompressor()
 177             return self.decompress(dc.decompress, size,
 178                                    num_retries=num_retries)
 179         elif self.name.endswith('.gz'):
 180             dc = zlib.decompressobj(16+zlib.MAX_WBITS)
 181             return self.decompress(lambda segment: dc.decompress(dc.unconsumed_tail + segment),
 182                                    size, num_retries=num_retries)
 183         else:
 184             return self.readall(size, num_retries=num_retries)
 185
 186     @_FileLikeObjectBase._before_close
 187     @retry_method
 188     def readlines(self, sizehint=float('inf'), num_retries=None):
 189         data = []
 190         data_size = 0
 191         for s in self.readall(num_retries=num_retries):
 192             data.append(s)
 193             data_size += len(s)
 194             if data_size >= sizehint:
 195                 break
 196         return b''.join(data).decode().splitlines(True)
 197
 198     def size(self):
 199         raise IOError(errno.ENOSYS, "Not implemented")
 200
 201     def read(self, size, num_retries=None):
 202         raise IOError(errno.ENOSYS, "Not implemented")
 203
 204     def readfrom(self, start, size, num_retries=None):
 205         raise IOError(errno.ENOSYS, "Not implemented")
 206
 207
 208 class StreamFileReader(ArvadosFileReaderBase):
 209     class _NameAttribute(str):
 210         # The Python file API provides a plain .name attribute.
 211         # Older SDK provided a name() method.
 212         # This class provides both, for maximum compatibility.
 213         def __call__(self):
 214             return self
 215
 216     def __init__(self, stream, segments, name):
 217         super(StreamFileReader, self).__init__(self._NameAttribute(name), 'rb', num_retries=stream.num_retries)
 218         self._stream = stream
 219         self.segments = segments
 220
 221     def stream_name(self):
 222         return self._stream.name()
 223
 224     def size(self):
 225         n = self.segments[-1]
 226         return n.range_start + n.range_size
 227
 228     @_FileLikeObjectBase._before_close
 229     @retry_method
 230     def read(self, size, num_retries=None):
 231         """Read up to 'size' bytes from the stream, starting at the current file position"""
 232         if size == 0:
 233             return b''
 234
 235         data = b''
 236         available_chunks = locators_and_ranges(self.segments, self._filepos, size)
 237         if available_chunks:
 238             lr = available_chunks[0]
 239             data = self._stream.readfrom(lr.locator+lr.segment_offset,
 240                                          lr.segment_size,
 241                                          num_retries=num_retries)
 242
 243         self._filepos += len(data)
 244         return data
 245
 246     @_FileLikeObjectBase._before_close
 247     @retry_method
 248     def readfrom(self, start, size, num_retries=None):
 249         """Read up to 'size' bytes from the stream, starting at 'start'"""
 250         if size == 0:
 251             return b''
 252
 253         data = []
 254         for lr in locators_and_ranges(self.segments, start, size):
 255             data.append(self._stream.readfrom(lr.locator+lr.segment_offset, lr.segment_size,
 256                                               num_retries=num_retries))
 257         return b''.join(data)
 258
 259     def as_manifest(self):
 260         segs = []
 261         for r in self.segments:
 262             segs.extend(self._stream.locators_and_ranges(r.locator, r.range_size))
 263         return " ".join(normalize_stream(".", {self.name: segs})) + "\n"
 264
 265
 266 def synchronized(orig_func):
 267     @functools.wraps(orig_func)
 268     def synchronized_wrapper(self, *args, **kwargs):
 269         with self.lock:
 270             return orig_func(self, *args, **kwargs)
 271     return synchronized_wrapper
 272
 273
 274 class StateChangeError(Exception):
 275     def __init__(self, message, state, nextstate):
 276         super(StateChangeError, self).__init__(message)
 277         self.state = state
 278         self.nextstate = nextstate
 279
 280 class _BufferBlock(object):
 281     """A stand-in for a Keep block that is in the process of being written.
 282
 283     Writers can append to it, get the size, and compute the Keep locator.
 284     There are three valid states:
 285
 286     WRITABLE
 287       Can append to block.
 288
 289     PENDING
 290       Block is in the process of being uploaded to Keep, append is an error.
 291
 292     COMMITTED
 293       The block has been written to Keep, its internal buffer has been
 294       released, fetching the block will fetch it via keep client (since we
 295       discarded the internal copy), and identifiers referring to the BufferBlock
 296       can be replaced with the block locator.
 297
 298     """
 299
 300     WRITABLE = 0
 301     PENDING = 1
 302     COMMITTED = 2
 303     ERROR = 3
 304     DELETED = 4
 305
 306     def __init__(self, blockid, starting_capacity, owner):
 307         """
 308         :blockid:
 309           the identifier for this block
 310
 311         :starting_capacity:
 312           the initial buffer capacity
 313
 314         :owner:
 315           ArvadosFile that owns this block
 316
 317         """
 318         self.blockid = blockid
 319         self.buffer_block = bytearray(starting_capacity)
 320         self.buffer_view = memoryview(self.buffer_block)
 321         self.write_pointer = 0
 322         self._state = _BufferBlock.WRITABLE
 323         self._locator = None
 324         self.owner = owner
 325         self.lock = threading.Lock()
 326         self.wait_for_commit = threading.Event()
 327         self.error = None
 328
 329     @synchronized
 330     def append(self, data):
 331         """Append some data to the buffer.
 332
 333         Only valid if the block is in WRITABLE state.  Implements an expanding
 334         buffer, doubling capacity as needed to accomdate all the data.
 335
 336         """
 337         if self._state == _BufferBlock.WRITABLE:
 338             if not isinstance(data, bytes) and not isinstance(data, memoryview):
 339                 data = data.encode()
 340             while (self.write_pointer+len(data)) > len(self.buffer_block):
 341                 new_buffer_block = bytearray(len(self.buffer_block) * 2)
 342                 new_buffer_block[0:self.write_pointer] = self.buffer_block[0:self.write_pointer]
 343                 self.buffer_block = new_buffer_block
 344                 self.buffer_view = memoryview(self.buffer_block)
 345             self.buffer_view[self.write_pointer:self.write_pointer+len(data)] = data
 346             self.write_pointer += len(data)
 347             self._locator = None
 348         else:
 349             raise AssertionError("Buffer block is not writable")
 350
 351     STATE_TRANSITIONS = frozenset([
 352             (WRITABLE, PENDING),
 353             (PENDING, COMMITTED),
 354             (PENDING, ERROR),
 355             (ERROR, PENDING)])
 356
 357     @synchronized
 358     def set_state(self, nextstate, val=None):
 359         if (self._state, nextstate) not in self.STATE_TRANSITIONS:
 360             raise StateChangeError("Invalid state change from %s to %s" % (self._state, nextstate), self._state, nextstate)
 361         self._state = nextstate
 362
 363         if self._state == _BufferBlock.PENDING:
 364             self.wait_for_commit.clear()
 365
 366         if self._state == _BufferBlock.COMMITTED:
 367             self._locator = val
 368             self.buffer_view = None
 369             self.buffer_block = None
 370             self.wait_for_commit.set()
 371
 372         if self._state == _BufferBlock.ERROR:
 373             self.error = val
 374             self.wait_for_commit.set()
 375
 376     @synchronized
 377     def state(self):
 378         return self._state
 379
 380     def size(self):
 381         """The amount of data written to the buffer."""
 382         return self.write_pointer
 383
 384     @synchronized
 385     def locator(self):
 386         """The Keep locator for this buffer's contents."""
 387         if self._locator is None:
 388             self._locator = "%s+%i" % (hashlib.md5(self.buffer_view[0:self.write_pointer]).hexdigest(), self.size())
 389         return self._locator
 390
 391     @synchronized
 392     def clone(self, new_blockid, owner):
 393         if self._state == _BufferBlock.COMMITTED:
 394             raise AssertionError("Cannot duplicate committed buffer block")
 395         bufferblock = _BufferBlock(new_blockid, self.size(), owner)
 396         bufferblock.append(self.buffer_view[0:self.size()])
 397         return bufferblock
 398
 399     @synchronized
 400     def clear(self):
 401         self._state = _BufferBlock.DELETED
 402         self.owner = None
 403         self.buffer_block = None
 404         self.buffer_view = None
 405
 406     @synchronized
 407     def repack_writes(self):
 408         """Optimize buffer block by repacking segments in file sequence.
 409
 410         When the client makes random writes, they appear in the buffer block in
 411         the sequence they were written rather than the sequence they appear in
 412         the file.  This makes for inefficient, fragmented manifests.  Attempt
 413         to optimize by repacking writes in file sequence.
 414
 415         """
 416         if self._state != _BufferBlock.WRITABLE:
 417             raise AssertionError("Cannot repack non-writable block")
 418
 419         segs = self.owner.segments()
 420
 421         # Collect the segments that reference the buffer block.
 422         bufferblock_segs = [s for s in segs if s.locator == self.blockid]
 423
 424         # Collect total data referenced by segments (could be smaller than
 425         # bufferblock size if a portion of the file was written and
 426         # then overwritten).
 427         write_total = sum([s.range_size for s in bufferblock_segs])
 428
 429         if write_total < self.size() or len(bufferblock_segs) > 1:
 430             # If there's more than one segment referencing this block, it is
 431             # due to out-of-order writes and will produce a fragmented
 432             # manifest, so try to optimize by re-packing into a new buffer.
 433             contents = self.buffer_view[0:self.write_pointer].tobytes()
 434             new_bb = _BufferBlock(None, write_total, None)
 435             for t in bufferblock_segs:
 436                 new_bb.append(contents[t.segment_offset:t.segment_offset+t.range_size])
 437                 t.segment_offset = new_bb.size() - t.range_size
 438
 439             self.buffer_block = new_bb.buffer_block
 440             self.buffer_view = new_bb.buffer_view
 441             self.write_pointer = new_bb.write_pointer
 442             self._locator = None
 443             new_bb.clear()
 444             self.owner.set_segments(segs)
 445
 446     def __repr__(self):
 447         return "<BufferBlock %s>" % (self.blockid)
 448
 449
 450 class NoopLock(object):
 451     def __enter__(self):
 452         return self
 453
 454     def __exit__(self, exc_type, exc_value, traceback):
 455         pass
 456
 457     def acquire(self, blocking=False):
 458         pass
 459
 460     def release(self):
 461         pass
 462
 463
 464 def must_be_writable(orig_func):
 465     @functools.wraps(orig_func)
 466     def must_be_writable_wrapper(self, *args, **kwargs):
 467         if not self.writable():
 468             raise IOError(errno.EROFS, "Collection is read-only.")
 469         return orig_func(self, *args, **kwargs)
 470     return must_be_writable_wrapper
 471
 472
 473 class _BlockManager(object):
 474     """BlockManager handles buffer blocks.
 475
 476     Also handles background block uploads, and background block prefetch for a
 477     Collection of ArvadosFiles.
 478
 479     """
 480
 481     DEFAULT_PUT_THREADS = 2
 482     DEFAULT_GET_THREADS = 2
 483
 484     def __init__(self, keep, copies=None, put_threads=None, num_retries=None, storage_classes=[]):
 485         """keep: KeepClient object to use"""
 486         self._keep = keep
 487         self._bufferblocks = collections.OrderedDict()
 488         self._put_queue = None
 489         self._put_threads = None
 490         self._prefetch_queue = None
 491         self._prefetch_threads = None
 492         self.lock = threading.Lock()
 493         self.prefetch_enabled = True
 494         if put_threads:
 495             self.num_put_threads = put_threads
 496         else:
 497             self.num_put_threads = _BlockManager.DEFAULT_PUT_THREADS
 498         self.num_get_threads = _BlockManager.DEFAULT_GET_THREADS
 499         self.copies = copies
 500         self.storage_classes = storage_classes
 501         self._pending_write_size = 0
 502         self.threads_lock = threading.Lock()
 503         self.padding_block = None
 504         self.num_retries = num_retries
 505
 506     @synchronized
 507     def alloc_bufferblock(self, blockid=None, starting_capacity=2**14, owner=None):
 508         """Allocate a new, empty bufferblock in WRITABLE state and return it.
 509
 510         :blockid:
 511           optional block identifier, otherwise one will be automatically assigned
 512
 513         :starting_capacity:
 514           optional capacity, otherwise will use default capacity
 515
 516         :owner:
 517           ArvadosFile that owns this block
 518
 519         """
 520         return self._alloc_bufferblock(blockid, starting_capacity, owner)
 521
 522     def _alloc_bufferblock(self, blockid=None, starting_capacity=2**14, owner=None):
 523         if blockid is None:
 524             blockid = str(uuid.uuid4())
 525         bufferblock = _BufferBlock(blockid, starting_capacity=starting_capacity, owner=owner)
 526         self._bufferblocks[bufferblock.blockid] = bufferblock
 527         return bufferblock
 528
 529     @synchronized
 530     def dup_block(self, block, owner):
 531         """Create a new bufferblock initialized with the content of an existing bufferblock.
 532
 533         :block:
 534           the buffer block to copy.
 535
 536         :owner:
 537           ArvadosFile that owns the new block
 538
 539         """
 540         new_blockid = str(uuid.uuid4())
 541         bufferblock = block.clone(new_blockid, owner)
 542         self._bufferblocks[bufferblock.blockid] = bufferblock
 543         return bufferblock
 544
 545     @synchronized
 546     def is_bufferblock(self, locator):
 547         return locator in self._bufferblocks
 548
 549     def _commit_bufferblock_worker(self):
 550         """Background uploader thread."""
 551
 552         while True:
 553             try:
 554                 bufferblock = self._put_queue.get()
 555                 if bufferblock is None:
 556                     return
 557
 558                 if self.copies is None:
 559                     loc = self._keep.put(bufferblock.buffer_view[0:bufferblock.write_pointer].tobytes(), num_retries=self.num_retries, classes=self.storage_classes)
 560                 else:
 561                     loc = self._keep.put(bufferblock.buffer_view[0:bufferblock.write_pointer].tobytes(), num_retries=self.num_retries, copies=self.copies, classes=self.storage_classes)
 562                 bufferblock.set_state(_BufferBlock.COMMITTED, loc)
 563             except Exception as e:
 564                 bufferblock.set_state(_BufferBlock.ERROR, e)
 565             finally:
 566                 if self._put_queue is not None:
 567                     self._put_queue.task_done()
 568
 569     def start_put_threads(self):
 570         with self.threads_lock:
 571             if self._put_threads is None:
 572                 # Start uploader threads.
 573
 574                 # If we don't limit the Queue size, the upload queue can quickly
 575                 # grow to take up gigabytes of RAM if the writing process is
 576                 # generating data more quickly than it can be send to the Keep
 577                 # servers.
 578                 #
 579                 # With two upload threads and a queue size of 2, this means up to 4
 580                 # blocks pending.  If they are full 64 MiB blocks, that means up to
 581                 # 256 MiB of internal buffering, which is the same size as the
 582                 # default download block cache in KeepClient.
 583                 self._put_queue = queue.Queue(maxsize=2)
 584
 585                 self._put_threads = []
 586                 for i in range(0, self.num_put_threads):
 587                     thread = threading.Thread(target=self._commit_bufferblock_worker)
 588                     self._put_threads.append(thread)
 589                     thread.daemon = True
 590                     thread.start()
 591
 592     def _block_prefetch_worker(self):
 593         """The background downloader thread."""
 594         while True:
 595             try:
 596                 b = self._prefetch_queue.get()
 597                 if b is None:
 598                     return
 599                 self._keep.get(b)
 600             except Exception:
 601                 _logger.exception("Exception doing block prefetch")
 602
 603     @synchronized
 604     def start_get_threads(self):
 605         if self._prefetch_threads is None:
 606             self._prefetch_queue = queue.Queue()
 607             self._prefetch_threads = []
 608             for i in range(0, self.num_get_threads):
 609                 thread = threading.Thread(target=self._block_prefetch_worker)
 610                 self._prefetch_threads.append(thread)
 611                 thread.daemon = True
 612                 thread.start()
 613
 614
 615     @synchronized
 616     def stop_threads(self):
 617         """Shut down and wait for background upload and download threads to finish."""
 618
 619         if self._put_threads is not None:
 620             for t in self._put_threads:
 621                 self._put_queue.put(None)
 622             for t in self._put_threads:
 623                 t.join()
 624         self._put_threads = None
 625         self._put_queue = None
 626
 627         if self._prefetch_threads is not None:
 628             for t in self._prefetch_threads:
 629                 self._prefetch_queue.put(None)
 630             for t in self._prefetch_threads:
 631                 t.join()
 632         self._prefetch_threads = None
 633         self._prefetch_queue = None
 634
 635     def __enter__(self):
 636         return self
 637
 638     def __exit__(self, exc_type, exc_value, traceback):
 639         self.stop_threads()
 640
 641     @synchronized
 642     def repack_small_blocks(self, force=False, sync=False, closed_file_size=0):
 643         """Packs small blocks together before uploading"""
 644
 645         self._pending_write_size += closed_file_size
 646
 647         # Check if there are enough small blocks for filling up one in full
 648         if not (force or (self._pending_write_size >= config.KEEP_BLOCK_SIZE)):
 649             return
 650
 651         # Search blocks ready for getting packed together before being
 652         # committed to Keep.
 653         # A WRITABLE block always has an owner.
 654         # A WRITABLE block with its owner.closed() implies that its
 655         # size is <= KEEP_BLOCK_SIZE/2.
 656         try:
 657             small_blocks = [b for b in listvalues(self._bufferblocks)
 658                             if b.state() == _BufferBlock.WRITABLE and b.owner.closed()]
 659         except AttributeError:
 660             # Writable blocks without owner shouldn't exist.
 661             raise UnownedBlockError()
 662
 663         if len(small_blocks) <= 1:
 664             # Not enough small blocks for repacking
 665             return
 666
 667         for bb in small_blocks:
 668             bb.repack_writes()
 669
 670         # Update the pending write size count with its true value, just in case
 671         # some small file was opened, written and closed several times.
 672         self._pending_write_size = sum([b.size() for b in small_blocks])
 673
 674         if self._pending_write_size < config.KEEP_BLOCK_SIZE and not force:
 675             return
 676
 677         new_bb = self._alloc_bufferblock()
 678         new_bb.owner = []
 679         files = []
 680         while len(small_blocks) > 0 and (new_bb.write_pointer + small_blocks[0].size()) <= config.KEEP_BLOCK_SIZE:
 681             bb = small_blocks.pop(0)
 682             new_bb.owner.append(bb.owner)
 683             self._pending_write_size -= bb.size()
 684             new_bb.append(bb.buffer_view[0:bb.write_pointer].tobytes())
 685             files.append((bb, new_bb.write_pointer - bb.size()))
 686
 687         self.commit_bufferblock(new_bb, sync=sync)
 688
 689         for bb, new_bb_segment_offset in files:
 690             newsegs = bb.owner.segments()
 691             for s in newsegs:
 692                 if s.locator == bb.blockid:
 693                     s.locator = new_bb.blockid
 694                     s.segment_offset = new_bb_segment_offset+s.segment_offset
 695             bb.owner.set_segments(newsegs)
 696             self._delete_bufferblock(bb.blockid)
 697
 698     def commit_bufferblock(self, block, sync):
 699         """Initiate a background upload of a bufferblock.
 700
 701         :block:
 702           The block object to upload
 703
 704         :sync:
 705           If `sync` is True, upload the block synchronously.
 706           If `sync` is False, upload the block asynchronously.  This will
 707           return immediately unless the upload queue is at capacity, in
 708           which case it will wait on an upload queue slot.
 709
 710         """
 711         try:
 712             # Mark the block as PENDING so to disallow any more appends.
 713             block.set_state(_BufferBlock.PENDING)
 714         except StateChangeError as e:
 715             if e.state == _BufferBlock.PENDING:
 716                 if sync:
 717                     block.wait_for_commit.wait()
 718                 else:
 719                     return
 720             if block.state() == _BufferBlock.COMMITTED:
 721                 return
 722             elif block.state() == _BufferBlock.ERROR:
 723                 raise block.error
 724             else:
 725                 raise
 726
 727         if sync:
 728             try:
 729                 if self.copies is None:
 730                     loc = self._keep.put(block.buffer_view[0:block.write_pointer].tobytes(), num_retries=self.num_retries, classes=self.storage_classes)
 731                 else:
 732                     loc = self._keep.put(block.buffer_view[0:block.write_pointer].tobytes(), num_retries=self.num_retries, copies=self.copies, classes=self.storage_classes)
 733                 block.set_state(_BufferBlock.COMMITTED, loc)
 734             except Exception as e:
 735                 block.set_state(_BufferBlock.ERROR, e)
 736                 raise
 737         else:
 738             self.start_put_threads()
 739             self._put_queue.put(block)
 740
 741     @synchronized
 742     def get_bufferblock(self, locator):
 743         return self._bufferblocks.get(locator)
 744
 745     @synchronized
 746     def get_padding_block(self):
 747         """Get a bufferblock 64 MB in size consisting of all zeros, used as padding
 748         when using truncate() to extend the size of a file.
 749
 750         For reference (and possible future optimization), the md5sum of the
 751         padding block is: 7f614da9329cd3aebf59b91aadc30bf0+67108864
 752
 753         """
 754
 755         if self.padding_block is None:
 756             self.padding_block = self._alloc_bufferblock(starting_capacity=config.KEEP_BLOCK_SIZE)
 757             self.padding_block.write_pointer = config.KEEP_BLOCK_SIZE
 758             self.commit_bufferblock(self.padding_block, False)
 759         return self.padding_block
 760
 761     @synchronized
 762     def delete_bufferblock(self, locator):
 763         self._delete_bufferblock(locator)
 764
 765     def _delete_bufferblock(self, locator):
 766         bb = self._bufferblocks[locator]
 767         bb.clear()
 768         del self._bufferblocks[locator]
 769
 770     def get_block_contents(self, locator, num_retries, cache_only=False):
 771         """Fetch a block.
 772
 773         First checks to see if the locator is a BufferBlock and return that, if
 774         not, passes the request through to KeepClient.get().
 775
 776         """
 777         with self.lock:
 778             if locator in self._bufferblocks:
 779                 bufferblock = self._bufferblocks[locator]
 780                 if bufferblock.state() != _BufferBlock.COMMITTED:
 781                     return bufferblock.buffer_view[0:bufferblock.write_pointer].tobytes()
 782                 else:
 783                     locator = bufferblock._locator
 784         if cache_only:
 785             return self._keep.get_from_cache(locator)
 786         else:
 787             return self._keep.get(locator, num_retries=num_retries)
 788
 789     def commit_all(self):
 790         """Commit all outstanding buffer blocks.
 791
 792         This is a synchronous call, and will not return until all buffer blocks
 793         are uploaded.  Raises KeepWriteError() if any blocks failed to upload.
 794
 795         """
 796         self.repack_small_blocks(force=True, sync=True)
 797
 798         with self.lock:
 799             items = listitems(self._bufferblocks)
 800
 801         for k,v in items:
 802             if v.state() != _BufferBlock.COMMITTED and v.owner:
 803                 # Ignore blocks with a list of owners, as if they're not in COMMITTED
 804                 # state, they're already being committed asynchronously.
 805                 if isinstance(v.owner, ArvadosFile):
 806                     v.owner.flush(sync=False)
 807
 808         with self.lock:
 809             if self._put_queue is not None:
 810                 self._put_queue.join()
 811
 812                 err = []
 813                 for k,v in items:
 814                     if v.state() == _BufferBlock.ERROR:
 815                         err.append((v.locator(), v.error))
 816                 if err:
 817                     raise KeepWriteError("Error writing some blocks", err, label="block")
 818
 819         for k,v in items:
 820             # flush again with sync=True to remove committed bufferblocks from
 821             # the segments.
 822             if v.owner:
 823                 if isinstance(v.owner, ArvadosFile):
 824                     v.owner.flush(sync=True)
 825                 elif isinstance(v.owner, list) and len(v.owner) > 0:
 826                     # This bufferblock is referenced by many files as a result
 827                     # of repacking small blocks, so don't delete it when flushing
 828                     # its owners, just do it after flushing them all.
 829                     for owner in v.owner:
 830                         owner.flush(sync=True)
 831                     self.delete_bufferblock(k)
 832
 833     def block_prefetch(self, locator):
 834         """Initiate a background download of a block.
 835
 836         This assumes that the underlying KeepClient implements a block cache,
 837         so repeated requests for the same block will not result in repeated
 838         downloads (unless the block is evicted from the cache.)  This method
 839         does not block.
 840
 841         """
 842
 843         if not self.prefetch_enabled:
 844             return
 845
 846         if self._keep.get_from_cache(locator) is not None:
 847             return
 848
 849         with self.lock:
 850             if locator in self._bufferblocks:
 851                 return
 852
 853         self.start_get_threads()
 854         self._prefetch_queue.put(locator)
 855
 856
 857 class ArvadosFile(object):
 858     """Represent a file in a Collection.
 859
 860     ArvadosFile manages the underlying representation of a file in Keep as a
 861     sequence of segments spanning a set of blocks, and implements random
 862     read/write access.
 863
 864     This object may be accessed from multiple threads.
 865
 866     """
 867
 868     __slots__ = ('parent', 'name', '_writers', '_committed',
 869                  '_segments', 'lock', '_current_bblock', 'fuse_entry')
 870
 871     def __init__(self, parent, name, stream=[], segments=[]):
 872         """
 873         ArvadosFile constructor.
 874
 875         :stream:
 876           a list of Range objects representing a block stream
 877
 878         :segments:
 879           a list of Range objects representing segments
 880         """
 881         self.parent = parent
 882         self.name = name
 883         self._writers = set()
 884         self._committed = False
 885         self._segments = []
 886         self.lock = parent.root_collection().lock
 887         for s in segments:
 888             self._add_segment(stream, s.locator, s.range_size)
 889         self._current_bblock = None
 890
 891     def writable(self):
 892         return self.parent.writable()
 893
 894     @synchronized
 895     def permission_expired(self, as_of_dt=None):
 896         """Returns True if any of the segment's locators is expired"""
 897         for r in self._segments:
 898             if KeepLocator(r.locator).permission_expired(as_of_dt):
 899                 return True
 900         return False
 901
 902     @synchronized
 903     def has_remote_blocks(self):
 904         """Returns True if any of the segment's locators has a +R signature"""
 905
 906         for s in self._segments:
 907             if '+R' in s.locator:
 908                 return True
 909         return False
 910
 911     @synchronized
 912     def _copy_remote_blocks(self, remote_blocks={}):
 913         """Ask Keep to copy remote blocks and point to their local copies.
 914
 915         This is called from the parent Collection.
 916
 917         :remote_blocks:
 918             Shared cache of remote to local block mappings. This is used to avoid
 919             doing extra work when blocks are shared by more than one file in
 920             different subdirectories.
 921         """
 922
 923         for s in self._segments:
 924             if '+R' in s.locator:
 925                 try:
 926                     loc = remote_blocks[s.locator]
 927                 except KeyError:
 928                     loc = self.parent._my_keep().refresh_signature(s.locator)
 929                     remote_blocks[s.locator] = loc
 930                 s.locator = loc
 931                 self.parent.set_committed(False)
 932         return remote_blocks
 933
 934     @synchronized
 935     def segments(self):
 936         return copy.copy(self._segments)
 937
 938     @synchronized
 939     def clone(self, new_parent, new_name):
 940         """Make a copy of this file."""
 941         cp = ArvadosFile(new_parent, new_name)
 942         cp.replace_contents(self)
 943         return cp
 944
 945     @must_be_writable
 946     @synchronized
 947     def replace_contents(self, other):
 948         """Replace segments of this file with segments from another `ArvadosFile` object."""
 949
 950         map_loc = {}
 951         self._segments = []
 952         for other_segment in other.segments():
 953             new_loc = other_segment.locator
 954             if other.parent._my_block_manager().is_bufferblock(other_segment.locator):
 955                 if other_segment.locator not in map_loc:
 956                     bufferblock = other.parent._my_block_manager().get_bufferblock(other_segment.locator)
 957                     if bufferblock.state() != _BufferBlock.WRITABLE:
 958                         map_loc[other_segment.locator] = bufferblock.locator()
 959                     else:
 960                         map_loc[other_segment.locator] = self.parent._my_block_manager().dup_block(bufferblock, self).blockid
 961                 new_loc = map_loc[other_segment.locator]
 962
 963             self._segments.append(Range(new_loc, other_segment.range_start, other_segment.range_size, other_segment.segment_offset))
 964
 965         self.set_committed(False)
 966
 967     def __eq__(self, other):
 968         if other is self:
 969             return True
 970         if not isinstance(other, ArvadosFile):
 971             return False
 972
 973         othersegs = other.segments()
 974         with self.lock:
 975             if len(self._segments) != len(othersegs):
 976                 return False
 977             for i in range(0, len(othersegs)):
 978                 seg1 = self._segments[i]
 979                 seg2 = othersegs[i]
 980                 loc1 = seg1.locator
 981                 loc2 = seg2.locator
 982
 983                 if self.parent._my_block_manager().is_bufferblock(loc1):
 984                     loc1 = self.parent._my_block_manager().get_bufferblock(loc1).locator()
 985
 986                 if other.parent._my_block_manager().is_bufferblock(loc2):
 987                     loc2 = other.parent._my_block_manager().get_bufferblock(loc2).locator()
 988
 989                 if (KeepLocator(loc1).stripped() != KeepLocator(loc2).stripped() or
 990                     seg1.range_start != seg2.range_start or
 991                     seg1.range_size != seg2.range_size or
 992                     seg1.segment_offset != seg2.segment_offset):
 993                     return False
 994
 995         return True
 996
 997     def __ne__(self, other):
 998         return not self.__eq__(other)
 999
1000     @synchronized
1001     def set_segments(self, segs):
1002         self._segments = segs
1003
1004     @synchronized
1005     def set_committed(self, value=True):
1006         """Set committed flag.
1007
1008         If value is True, set committed to be True.
1009
1010         If value is False, set committed to be False for this and all parents.
1011         """
1012         if value == self._committed:
1013             return
1014         self._committed = value
1015         if self._committed is False and self.parent is not None:
1016             self.parent.set_committed(False)
1017
1018     @synchronized
1019     def committed(self):
1020         """Get whether this is committed or not."""
1021         return self._committed
1022
1023     @synchronized
1024     def add_writer(self, writer):
1025         """Add an ArvadosFileWriter reference to the list of writers"""
1026         if isinstance(writer, ArvadosFileWriter):
1027             self._writers.add(writer)
1028
1029     @synchronized
1030     def remove_writer(self, writer, flush):
1031         """
1032         Called from ArvadosFileWriter.close(). Remove a writer reference from the list
1033         and do some block maintenance tasks.
1034         """
1035         self._writers.remove(writer)
1036
1037         if flush or self.size() > config.KEEP_BLOCK_SIZE // 2:
1038             # File writer closed, not small enough for repacking
1039             self.flush()
1040         elif self.closed():
1041             # All writers closed and size is adequate for repacking
1042             self.parent._my_block_manager().repack_small_blocks(closed_file_size=self.size())
1043
1044     def closed(self):
1045         """
1046         Get whether this is closed or not. When the writers list is empty, the file
1047         is supposed to be closed.
1048         """
1049         return len(self._writers) == 0
1050
1051     @must_be_writable
1052     @synchronized
1053     def truncate(self, size):
1054         """Shrink or expand the size of the file.
1055
1056         If `size` is less than the size of the file, the file contents after
1057         `size` will be discarded.  If `size` is greater than the current size
1058         of the file, it will be filled with zero bytes.
1059
1060         """
1061         if size < self.size():
1062             new_segs = []
1063             for r in self._segments:
1064                 range_end = r.range_start+r.range_size
1065                 if r.range_start >= size:
1066                     # segment is past the trucate size, all done
1067                     break
1068                 elif size < range_end:
1069                     nr = Range(r.locator, r.range_start, size - r.range_start, 0)
1070                     nr.segment_offset = r.segment_offset
1071                     new_segs.append(nr)
1072                     break
1073                 else:
1074                     new_segs.append(r)
1075
1076             self._segments = new_segs
1077             self.set_committed(False)
1078         elif size > self.size():
1079             padding = self.parent._my_block_manager().get_padding_block()
1080             diff = size - self.size()
1081             while diff > config.KEEP_BLOCK_SIZE:
1082                 self._segments.append(Range(padding.blockid, self.size(), config.KEEP_BLOCK_SIZE, 0))
1083                 diff -= config.KEEP_BLOCK_SIZE
1084             if diff > 0:
1085                 self._segments.append(Range(padding.blockid, self.size(), diff, 0))
1086             self.set_committed(False)
1087         else:
1088             # size == self.size()
1089             pass
1090
1091     def readfrom(self, offset, size, num_retries, exact=False):
1092         """Read up to `size` bytes from the file starting at `offset`.
1093
1094         :exact:
1095          If False (default), return less data than requested if the read
1096          crosses a block boundary and the next block isn't cached.  If True,
1097          only return less data than requested when hitting EOF.
1098         """
1099
1100         with self.lock:
1101             if size == 0 or offset >= self.size():
1102                 return b''
1103             readsegs = locators_and_ranges(self._segments, offset, size)
1104             prefetch = locators_and_ranges(self._segments, offset + size, config.KEEP_BLOCK_SIZE, limit=32)
1105
1106         locs = set()
1107         data = []
1108         for lr in readsegs:
1109             block = self.parent._my_block_manager().get_block_contents(lr.locator, num_retries=num_retries, cache_only=(bool(data) and not exact))
1110             if block:
1111                 blockview = memoryview(block)
1112                 data.append(blockview[lr.segment_offset:lr.segment_offset+lr.segment_size].tobytes())
1113                 locs.add(lr.locator)
1114             else:
1115                 break
1116
1117         for lr in prefetch:
1118             if lr.locator not in locs:
1119                 self.parent._my_block_manager().block_prefetch(lr.locator)
1120                 locs.add(lr.locator)
1121
1122         return b''.join(data)
1123
1124     @must_be_writable
1125     @synchronized
1126     def writeto(self, offset, data, num_retries):
1127         """Write `data` to the file starting at `offset`.
1128
1129         This will update existing bytes and/or extend the size of the file as
1130         necessary.
1131
1132         """
1133         if not isinstance(data, bytes) and not isinstance(data, memoryview):
1134             data = data.encode()
1135         if len(data) == 0:
1136             return
1137
1138         if offset > self.size():
1139             self.truncate(offset)
1140
1141         if len(data) > config.KEEP_BLOCK_SIZE:
1142             # Chunk it up into smaller writes
1143             n = 0
1144             dataview = memoryview(data)
1145             while n < len(data):
1146                 self.writeto(offset+n, dataview[n:n + config.KEEP_BLOCK_SIZE].tobytes(), num_retries)
1147                 n += config.KEEP_BLOCK_SIZE
1148             return
1149
1150         self.set_committed(False)
1151
1152         if self._current_bblock is None or self._current_bblock.state() != _BufferBlock.WRITABLE:
1153             self._current_bblock = self.parent._my_block_manager().alloc_bufferblock(owner=self)
1154
1155         if (self._current_bblock.size() + len(data)) > config.KEEP_BLOCK_SIZE:
1156             self._current_bblock.repack_writes()
1157             if (self._current_bblock.size() + len(data)) > config.KEEP_BLOCK_SIZE:
1158                 self.parent._my_block_manager().commit_bufferblock(self._current_bblock, sync=False)
1159                 self._current_bblock = self.parent._my_block_manager().alloc_bufferblock(owner=self)
1160
1161         self._current_bblock.append(data)
1162
1163         replace_range(self._segments, offset, len(data), self._current_bblock.blockid, self._current_bblock.write_pointer - len(data))
1164
1165         self.parent.notify(WRITE, self.parent, self.name, (self, self))
1166
1167         return len(data)
1168
1169     @synchronized
1170     def flush(self, sync=True, num_retries=0):
1171         """Flush the current bufferblock to Keep.
1172
1173         :sync:
1174           If True, commit block synchronously, wait until buffer block has been written.
1175           If False, commit block asynchronously, return immediately after putting block into
1176           the keep put queue.
1177         """
1178         if self.committed():
1179             return
1180
1181         if self._current_bblock and self._current_bblock.state() != _BufferBlock.COMMITTED:
1182             if self._current_bblock.state() == _BufferBlock.WRITABLE:
1183                 self._current_bblock.repack_writes()
1184             if self._current_bblock.state() != _BufferBlock.DELETED:
1185                 self.parent._my_block_manager().commit_bufferblock(self._current_bblock, sync=sync)
1186
1187         if sync:
1188             to_delete = set()
1189             for s in self._segments:
1190                 bb = self.parent._my_block_manager().get_bufferblock(s.locator)
1191                 if bb:
1192                     if bb.state() != _BufferBlock.COMMITTED:
1193                         self.parent._my_block_manager().commit_bufferblock(bb, sync=True)
1194                     to_delete.add(s.locator)
1195                     s.locator = bb.locator()
1196             for s in to_delete:
1197                 # Don't delete the bufferblock if it's owned by many files. It'll be
1198                 # deleted after all of its owners are flush()ed.
1199                 if self.parent._my_block_manager().get_bufferblock(s).owner is self:
1200                     self.parent._my_block_manager().delete_bufferblock(s)
1201
1202         self.parent.notify(MOD, self.parent, self.name, (self, self))
1203
1204     @must_be_writable
1205     @synchronized
1206     def add_segment(self, blocks, pos, size):
1207         """Add a segment to the end of the file.
1208
1209         `pos` and `offset` reference a section of the stream described by
1210         `blocks` (a list of Range objects)
1211
1212         """
1213         self._add_segment(blocks, pos, size)
1214
1215     def _add_segment(self, blocks, pos, size):
1216         """Internal implementation of add_segment."""
1217         self.set_committed(False)
1218         for lr in locators_and_ranges(blocks, pos, size):
1219             last = self._segments[-1] if self._segments else Range(0, 0, 0, 0)
1220             r = Range(lr.locator, last.range_start+last.range_size, lr.segment_size, lr.segment_offset)
1221             self._segments.append(r)
1222
1223     @synchronized
1224     def size(self):
1225         """Get the file size."""
1226         if self._segments:
1227             n = self._segments[-1]
1228             return n.range_start + n.range_size
1229         else:
1230             return 0
1231
1232     @synchronized
1233     def manifest_text(self, stream_name=".", portable_locators=False,
1234                       normalize=False, only_committed=False):
1235         buf = ""
1236         filestream = []
1237         for segment in self._segments:
1238             loc = segment.locator
1239             if self.parent._my_block_manager().is_bufferblock(loc):
1240                 if only_committed:
1241                     continue
1242                 loc = self.parent._my_block_manager().get_bufferblock(loc).locator()
1243             if portable_locators:
1244                 loc = KeepLocator(loc).stripped()
1245             filestream.append(LocatorAndRange(loc, KeepLocator(loc).size,
1246                                  segment.segment_offset, segment.range_size))
1247         buf += ' '.join(normalize_stream(stream_name, {self.name: filestream}))
1248         buf += "\n"
1249         return buf
1250
1251     @must_be_writable
1252     @synchronized
1253     def _reparent(self, newparent, newname):
1254         self.set_committed(False)
1255         self.flush(sync=True)
1256         self.parent.remove(self.name)
1257         self.parent = newparent
1258         self.name = newname
1259         self.lock = self.parent.root_collection().lock
1260
1261
1262 class ArvadosFileReader(ArvadosFileReaderBase):
1263     """Wraps ArvadosFile in a file-like object supporting reading only.
1264
1265     Be aware that this class is NOT thread safe as there is no locking around
1266     updating file pointer.
1267
1268     """
1269
1270     def __init__(self, arvadosfile, mode="r", num_retries=None):
1271         super(ArvadosFileReader, self).__init__(arvadosfile.name, mode=mode, num_retries=num_retries)
1272         self.arvadosfile = arvadosfile
1273
1274     def size(self):
1275         return self.arvadosfile.size()
1276
1277     def stream_name(self):
1278         return self.arvadosfile.parent.stream_name()
1279
1280     def readinto(self, b):
1281         data = self.read(len(b))
1282         b[:len(data)] = data
1283         return len(data)
1284
1285     @_FileLikeObjectBase._before_close
1286     @retry_method
1287     def read(self, size=None, num_retries=None):
1288         """Read up to `size` bytes from the file and return the result.
1289
1290         Starts at the current file position.  If `size` is None, read the
1291         entire remainder of the file.
1292         """
1293         if size is None:
1294             data = []
1295             rd = self.arvadosfile.readfrom(self._filepos, config.KEEP_BLOCK_SIZE, num_retries)
1296             while rd:
1297                 data.append(rd)
1298                 self._filepos += len(rd)
1299                 rd = self.arvadosfile.readfrom(self._filepos, config.KEEP_BLOCK_SIZE, num_retries)
1300             return b''.join(data)
1301         else:
1302             data = self.arvadosfile.readfrom(self._filepos, size, num_retries, exact=True)
1303             self._filepos += len(data)
1304             return data
1305
1306     @_FileLikeObjectBase._before_close
1307     @retry_method
1308     def readfrom(self, offset, size, num_retries=None):
1309         """Read up to `size` bytes from the stream, starting at the specified file offset.
1310
1311         This method does not change the file position.
1312         """
1313         return self.arvadosfile.readfrom(offset, size, num_retries)
1314
1315     def flush(self):
1316         pass
1317
1318
1319 class ArvadosFileWriter(ArvadosFileReader):
1320     """Wraps ArvadosFile in a file-like object supporting both reading and writing.
1321
1322     Be aware that this class is NOT thread safe as there is no locking around
1323     updating file pointer.
1324
1325     """
1326
1327     def __init__(self, arvadosfile, mode, num_retries=None):
1328         super(ArvadosFileWriter, self).__init__(arvadosfile, mode=mode, num_retries=num_retries)
1329         self.arvadosfile.add_writer(self)
1330
1331     def writable(self):
1332         return True
1333
1334     @_FileLikeObjectBase._before_close
1335     @retry_method
1336     def write(self, data, num_retries=None):
1337         if self.mode[0] == "a":
1338             self._filepos = self.size()
1339         self.arvadosfile.writeto(self._filepos, data, num_retries)
1340         self._filepos += len(data)
1341         return len(data)
1342
1343     @_FileLikeObjectBase._before_close
1344     @retry_method
1345     def writelines(self, seq, num_retries=None):
1346         for s in seq:
1347             self.write(s, num_retries=num_retries)
1348
1349     @_FileLikeObjectBase._before_close
1350     def truncate(self, size=None):
1351         if size is None:
1352             size = self._filepos
1353         self.arvadosfile.truncate(size)
1354
1355     @_FileLikeObjectBase._before_close
1356     def flush(self):
1357         self.arvadosfile.flush()
1358
1359     def close(self, flush=True):
1360         if not self.closed:
1361             self.arvadosfile.remove_writer(self, flush)
1362             super(ArvadosFileWriter, self).close()
1363
1364
1365 class WrappableFile(object):
1366     """An interface to an Arvados file that's compatible with io wrappers.
1367
1368     """
1369     def __init__(self, f):
1370         self.f = f
1371         self.closed = False
1372     def close(self):
1373         self.closed = True
1374         return self.f.close()
1375     def flush(self):
1376         return self.f.flush()
1377     def read(self, *args, **kwargs):
1378         return self.f.read(*args, **kwargs)
1379     def readable(self):
1380         return self.f.readable()
1381     def readinto(self, *args, **kwargs):
1382         return self.f.readinto(*args, **kwargs)
1383     def seek(self, *args, **kwargs):
1384         return self.f.seek(*args, **kwargs)
1385     def seekable(self):
1386         return self.f.seekable()
1387     def tell(self):
1388         return self.f.tell()
1389     def writable(self):
1390         return self.f.writable()
1391     def write(self, *args, **kwargs):
1392         return self.f.write(*args, **kwargs)