gsutil/gslib/daisy_chain_wrapper.py

# -*- coding: utf-8 -*-
# Copyright 2014 Google Inc. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Wrapper for use in daisy-chained copies."""

from collections import deque
import os
import threading
import time

from gslib.cloud_api import BadRequestException
from gslib.cloud_api import CloudApi
from gslib.util import CreateLock
from gslib.util import TRANSFER_BUFFER_SIZE


# This controls the amount of bytes downloaded per download request.
# We do not buffer this many bytes in memory at a time - that is controlled by
# DaisyChainWrapper.max_buffer_size. This is the upper bound of bytes that may
# be unnecessarily downloaded if there is a break in the resumable upload.
_DEFAULT_DOWNLOAD_CHUNK_SIZE = 1024*1024*100


class BufferWrapper(object):
  """Wraps the download file pointer to use our in-memory buffer."""

  def __init__(self, daisy_chain_wrapper):
    """Provides a buffered write interface for a file download.

    Args:
      daisy_chain_wrapper: DaisyChainWrapper instance to use for buffer and
                           locking.
    """
    self.daisy_chain_wrapper = daisy_chain_wrapper

  def write(self, data):  # pylint: disable=invalid-name
    """Waits for space in the buffer, then writes data to the buffer."""
    while True:
      with self.daisy_chain_wrapper.lock:
        if (self.daisy_chain_wrapper.bytes_buffered <
            self.daisy_chain_wrapper.max_buffer_size):
          break
      # Buffer was full, yield thread priority so the upload can pull from it.
      time.sleep(0)
    data_len = len(data)
    if data_len:
      with self.daisy_chain_wrapper.lock:
        self.daisy_chain_wrapper.buffer.append(data)
        self.daisy_chain_wrapper.bytes_buffered += data_len


class DaisyChainWrapper(object):
  """Wrapper class for daisy-chaining a cloud download to an upload.

  This class instantiates a BufferWrapper object to buffer the download into
  memory, consuming a maximum of max_buffer_size. It implements intelligent
  behavior around read and seek that allow for all of the operations necessary
  to copy a file.

  This class is coupled with the XML and JSON implementations in that it
  expects that small buffers (maximum of TRANSFER_BUFFER_SIZE) in size will be
  used.
  """

  def __init__(self, src_url, src_obj_size, gsutil_api, progress_callback=None,
               download_chunk_size=_DEFAULT_DOWNLOAD_CHUNK_SIZE):
    """Initializes the daisy chain wrapper.

    Args:
      src_url: Source CloudUrl to copy from.
      src_obj_size: Size of source object.
      gsutil_api: gsutil Cloud API to use for the copy.
      progress_callback: Optional callback function for progress notifications
          for the download thread. Receives calls with arguments
          (bytes_transferred, total_size).
      download_chunk_size: Integer number of bytes to download per
          GetObjectMedia request. This is the upper bound of bytes that may be
          unnecessarily downloaded if there is a break in the resumable upload.

    """
    # Current read position for the upload file pointer.
    self.position = 0
    self.buffer = deque()

    self.bytes_buffered = 0
    # Maximum amount of bytes in memory at a time.
    self.max_buffer_size = 1024 * 1024  # 1 MiB

    self._download_chunk_size = download_chunk_size

    # We save one buffer's worth of data as a special case for boto,
    # which seeks back one buffer and rereads to compute hashes. This is
    # unnecessary because we can just compare cloud hash digests at the end,
    # but it allows this to work without modfiying boto.
    self.last_position = 0
    self.last_data = None

    # Protects buffer, position, bytes_buffered, last_position, and last_data.
    self.lock = CreateLock()

    # Protects download_exception.
    self.download_exception_lock = CreateLock()

    self.src_obj_size = src_obj_size
    self.src_url = src_url

    # This is safe to use the upload and download thread because the download
    # thread calls only GetObjectMedia, which creates a new HTTP connection
    # independent of gsutil_api. Thus, it will not share an HTTP connection
    # with the upload.
    self.gsutil_api = gsutil_api

    # If self.download_thread dies due to an exception, it is saved here so
    # that it can also be raised in the upload thread.
    self.download_exception = None
    self.download_thread = None
    self.progress_callback = progress_callback
    self.stop_download = threading.Event()
    self.StartDownloadThread(progress_callback=self.progress_callback)

  def StartDownloadThread(self, start_byte=0, progress_callback=None):
    """Starts the download thread for the source object (from start_byte)."""

    def PerformDownload(start_byte, progress_callback):
      """Downloads the source object in chunks.

      This function checks the stop_download event and exits early if it is set.
      It should be set when there is an error during the daisy-chain upload,
      then this function can be called again with the upload's current position
      as start_byte.

      Args:
        start_byte: Byte from which to begin the download.
        progress_callback: Optional callback function for progress
            notifications. Receives calls with arguments
            (bytes_transferred, total_size).
      """
      # TODO: Support resumable downloads. This would require the BufferWrapper
      # object to support seek() and tell() which requires coordination with
      # the upload.
      try:
        while start_byte + self._download_chunk_size < self.src_obj_size:
          self.gsutil_api.GetObjectMedia(
              self.src_url.bucket_name, self.src_url.object_name,
              BufferWrapper(self), start_byte=start_byte,
              end_byte=start_byte + self._download_chunk_size - 1,
              generation=self.src_url.generation, object_size=self.src_obj_size,
              download_strategy=CloudApi.DownloadStrategy.ONE_SHOT,
              provider=self.src_url.scheme, progress_callback=progress_callback)
          if self.stop_download.is_set():
            # Download thread needs to be restarted, so exit.
            self.stop_download.clear()
            return
          start_byte += self._download_chunk_size
        self.gsutil_api.GetObjectMedia(
            self.src_url.bucket_name, self.src_url.object_name,
            BufferWrapper(self), start_byte=start_byte,
            generation=self.src_url.generation, object_size=self.src_obj_size,
            download_strategy=CloudApi.DownloadStrategy.ONE_SHOT,
            provider=self.src_url.scheme, progress_callback=progress_callback)
      # We catch all exceptions here because we want to store them.
      except Exception, e:  # pylint: disable=broad-except
        # Save the exception so that it can be seen in the upload thread.
        with self.download_exception_lock:
          self.download_exception = e
          raise

    # TODO: If we do gzip encoding transforms mid-transfer, this will fail.
    self.download_thread = threading.Thread(
        target=PerformDownload,
        args=(start_byte, progress_callback))
    self.download_thread.start()

  def read(self, amt=None):  # pylint: disable=invalid-name
    """Exposes a stream from the in-memory buffer to the upload."""
    if self.position == self.src_obj_size or amt == 0:
      # If there is no data left or 0 bytes were requested, return an empty
      # string so callers can call still call len() and read(0).
      return ''
    if amt is None or amt > TRANSFER_BUFFER_SIZE:
      raise BadRequestException(
          'Invalid HTTP read size %s during daisy chain operation, '
          'expected <= %s.' % (amt, TRANSFER_BUFFER_SIZE))

    while True:
      with self.lock:
        if self.buffer:
          break
        with self.download_exception_lock:
          if self.download_exception:
            # Download thread died, so we will never recover. Raise the
            # exception that killed it.
            raise self.download_exception  # pylint: disable=raising-bad-type
      # Buffer was empty, yield thread priority so the download thread can fill.
      time.sleep(0)
    with self.lock:
      # TODO: Need to handle the caller requesting less than a
      # transfer_buffer_size worth of data.
      data = self.buffer.popleft()
      self.last_position = self.position
      self.last_data = data
      data_len = len(data)
      self.position += data_len
      self.bytes_buffered -= data_len
    if data_len > amt:
      raise BadRequestException(
          'Invalid read during daisy chain operation, got data of size '
          '%s, expected size %s.' % (data_len, amt))
    return data

  def tell(self):  # pylint: disable=invalid-name
    with self.lock:
      return self.position

  def seek(self, offset, whence=os.SEEK_SET):  # pylint: disable=invalid-name
    restart_download = False
    if whence == os.SEEK_END:
      if offset:
        raise IOError(
            'Invalid seek during daisy chain operation. Non-zero offset %s '
            'from os.SEEK_END is not supported' % offset)
      with self.lock:
        self.last_position = self.position
        self.last_data = None
        # Safe because we check position against src_obj_size in read.
        self.position = self.src_obj_size
    elif whence == os.SEEK_SET:
      with self.lock:
        if offset == self.position:
          pass
        elif offset == self.last_position:
          self.position = self.last_position
          if self.last_data:
            # If we seek to end and then back, we won't have last_data; we'll
            # get it on the next call to read.
            self.buffer.appendleft(self.last_data)
            self.bytes_buffered += len(self.last_data)
        else:
          # Once a download is complete, boto seeks to 0 and re-reads to
          # compute the hash if an md5 isn't already present (for example a GCS
          # composite object), so we have to re-download the whole object.
          # Also, when daisy-chaining to a resumable upload, on error the
          # service may have received any number of the bytes; the download
          # needs to be restarted from that point.
          restart_download = True

      if restart_download:
        self.stop_download.set()

        # Consume any remaining bytes in the download thread so that
        # the thread can exit, then restart the thread at the desired position.
        while self.download_thread.is_alive():
          with self.lock:
            while self.bytes_buffered:
              self.bytes_buffered -= len(self.buffer.popleft())
          time.sleep(0)

        with self.lock:
          self.position = offset
          self.buffer = deque()
          self.bytes_buffered = 0
          self.last_position = 0
          self.last_data = None
        self.StartDownloadThread(start_byte=offset,
                                 progress_callback=self.progress_callback)
    else:
      raise IOError('Daisy-chain download wrapper does not support '
                    'seek mode %s' % whence)

  def seekable(self):  # pylint: disable=invalid-name
    return True