gsutil/gslib/boto_resumable_upload.py

# -*- coding: utf-8 -*-
# Copyright 2010 Google Inc. All Rights Reserved.
#
# Permission is hereby granted, free of charge, to any person obtaining a
# copy of this software and associated documentation files (the
# "Software"), to deal in the Software without restriction, including
# without limitation the rights to use, copy, modify, merge, publish, dis-
# tribute, sublicense, and/or sell copies of the Software, and to permit
# persons to whom the Software is furnished to do so, subject to the fol-
# lowing conditions:
#
# The above copyright notice and this permission notice shall be included
# in all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
# OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABIL-
# ITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT
# SHALL THE AUTHOR BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY,
# WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
# IN THE SOFTWARE.
"""Boto translation layer for resumable uploads.

See http://code.google.com/apis/storage/docs/developer-guide.html#resumable
for details.

Resumable uploads will retry interrupted uploads, resuming at the byte
count completed by the last upload attempt. If too many retries happen with
no progress (per configurable num_retries param), the upload will be
aborted in the current process.

Unlike the boto implementation of resumable upload handler, this class does
not directly interact with tracker files.

Originally Google wrote and contributed this code to the boto project,
then copied that code back into gsutil on the release of gsutil 4.0 which
supports both boto and non-boto codepaths for resumable uploads.  Any bug
fixes made to this file should also be integrated to resumable_upload_handler.py
in boto, where applicable.

TODO: gsutil-beta: Add a similar comment to the boto code.
"""

from __future__ import absolute_import

import errno
import httplib
import random
import re
import socket
import time
import urlparse
from boto import UserAgent
from boto.connection import AWSAuthConnection
from boto.exception import ResumableTransferDisposition
from boto.exception import ResumableUploadException
from gslib.exception import InvalidUrlError
from gslib.util import GetMaxRetryDelay
from gslib.util import GetNumRetries
from gslib.util import XML_PROGRESS_CALLBACKS


class BotoResumableUpload(object):
  """Upload helper class for resumable uploads via boto."""

  BUFFER_SIZE = 8192
  RETRYABLE_EXCEPTIONS = (httplib.HTTPException, IOError, socket.error,
                          socket.gaierror)

  # (start, end) response indicating service has nothing (upload protocol uses
  # inclusive numbering).
  SERVICE_HAS_NOTHING = (0, -1)

  def __init__(self, tracker_callback, logger,
               resume_url=None, num_retries=None):
    """Constructor. Instantiate once for each uploaded file.

    Args:
      tracker_callback: Callback function that takes a string argument.  Used
                        by caller to track this upload across upload
                        interruption.
      logger: logging.logger instance to use for debug messages.
      resume_url: If present, attempt to resume the upload at this URL.
      num_retries: Number of times to retry the upload making no progress.
                   This count resets every time we make progress, so the upload
                   can span many more than this number of retries.
    """
    if resume_url:
      self._SetUploadUrl(resume_url)
    else:
      self.upload_url = None
    self.num_retries = num_retries
    self.service_has_bytes = 0  # Byte count at last service check.
    # Save upload_start_point in instance state so caller can find how
    # much was transferred by this ResumableUploadHandler (across retries).
    self.upload_start_point = None
    self.tracker_callback = tracker_callback
    self.logger = logger

  def _SetUploadUrl(self, url):
    """Saves URL and resets upload state.

    Called when we start a new resumable upload or get a new tracker
    URL for the upload.

    Args:
      url: URL string for the upload.

    Raises InvalidUrlError if URL is syntactically invalid.
    """
    parse_result = urlparse.urlparse(url)
    if (parse_result.scheme.lower() not in ['http', 'https'] or
        not parse_result.netloc):
      raise InvalidUrlError('Invalid upload URL (%s)' % url)
    self.upload_url = url
    self.upload_url_host = parse_result.netloc
    self.upload_url_path = '%s?%s' % (
        parse_result.path, parse_result.query)
    self.service_has_bytes = 0

  def _BuildContentRangeHeader(self, range_spec='*', length_spec='*'):
    return 'bytes %s/%s' % (range_spec, length_spec)

  def _QueryServiceState(self, conn, file_length):
    """Queries service to find out state of given upload.

    Note that this method really just makes special case use of the
    fact that the upload service always returns the current start/end
    state whenever a PUT doesn't complete.

    Args:
      conn: HTTPConnection to use for the query.
      file_length: Total length of the file.

    Returns:
      HTTP response from sending request.

    Raises:
      ResumableUploadException if problem querying service.
    """
    # Send an empty PUT so that service replies with this resumable
    # transfer's state.
    put_headers = {}
    put_headers['Content-Range'] = (
        self._BuildContentRangeHeader('*', file_length))
    put_headers['Content-Length'] = '0'
    return AWSAuthConnection.make_request(
        conn, 'PUT', path=self.upload_url_path, auth_path=self.upload_url_path,
        headers=put_headers, host=self.upload_url_host)

  def _QueryServicePos(self, conn, file_length):
    """Queries service to find out what bytes it currently has.

    Args:
      conn: HTTPConnection to use for the query.
      file_length: Total length of the file.

    Returns:
      (service_start, service_end), where the values are inclusive.
      For example, (0, 2) would mean that the service has bytes 0, 1, *and* 2.

    Raises:
      ResumableUploadException if problem querying service.
    """
    resp = self._QueryServiceState(conn, file_length)
    if resp.status == 200:
      # To handle the boundary condition where the service has the complete
      # file, we return (service_start, file_length-1). That way the
      # calling code can always simply read up through service_end. (If we
      # didn't handle this boundary condition here, the caller would have
      # to check whether service_end == file_length and read one fewer byte
      # in that case.)
      return (0, file_length - 1)  # Completed upload.
    if resp.status != 308:
      # This means the service didn't have any state for the given
      # upload ID, which can happen (for example) if the caller saved
      # the upload URL to a file and then tried to restart the transfer
      # after that upload ID has gone stale. In that case we need to
      # start a new transfer (and the caller will then save the new
      # upload URL to the tracker file).
      raise ResumableUploadException(
          'Got non-308 response (%s) from service state query' %
          resp.status, ResumableTransferDisposition.START_OVER)
    got_valid_response = False
    range_spec = resp.getheader('range')
    if range_spec:
      # Parse 'bytes=<from>-<to>' range_spec.
      m = re.search(r'bytes=(\d+)-(\d+)', range_spec)
      if m:
        service_start = long(m.group(1))
        service_end = long(m.group(2))
        got_valid_response = True
    else:
      # No Range header, which means the service does not yet have
      # any bytes. Note that the Range header uses inclusive 'from'
      # and 'to' values. Since Range 0-0 would mean that the service
      # has byte 0, omitting the Range header is used to indicate that
      # the service doesn't have any bytes.
      return self.SERVICE_HAS_NOTHING
    if not got_valid_response:
      raise ResumableUploadException(
          'Couldn\'t parse upload service state query response (%s)' %
          str(resp.getheaders()), ResumableTransferDisposition.START_OVER)
    if conn.debug >= 1:
      self.logger.debug('Service has: Range: %d - %d.', service_start,
                        service_end)
    return (service_start, service_end)

  def _StartNewResumableUpload(self, key, headers=None):
    """Starts a new resumable upload.

    Args:
      key: Boto Key representing the object to upload.
      headers: Headers to use in the upload requests.

    Raises:
      ResumableUploadException if any errors occur.
    """
    conn = key.bucket.connection
    if conn.debug >= 1:
      self.logger.debug('Starting new resumable upload.')
    self.service_has_bytes = 0

    # Start a new resumable upload by sending a POST request with an
    # empty body and the "X-Goog-Resumable: start" header. Include any
    # caller-provided headers (e.g., Content-Type) EXCEPT Content-Length
    # (and raise an exception if they tried to pass one, since it's
    # a semantic error to specify it at this point, and if we were to
    # include one now it would cause the service to expect that many
    # bytes; the POST doesn't include the actual file bytes  We set
    # the Content-Length in the subsequent PUT, based on the uploaded
    # file size.
    post_headers = {}
    for k in headers:
      if k.lower() == 'content-length':
        raise ResumableUploadException(
            'Attempt to specify Content-Length header (disallowed)',
            ResumableTransferDisposition.ABORT)
      post_headers[k] = headers[k]
    post_headers[conn.provider.resumable_upload_header] = 'start'

    resp = conn.make_request(
        'POST', key.bucket.name, key.name, post_headers)
    # Get upload URL from response 'Location' header.
    body = resp.read()

    # Check for various status conditions.
    if resp.status in [429, 500, 503]:
      # Retry after a delay.
      raise ResumableUploadException(
          'Got status %d from attempt to start resumable upload. '
          'Will wait/retry' % resp.status,
          ResumableTransferDisposition.WAIT_BEFORE_RETRY)
    elif resp.status != 200 and resp.status != 201:
      raise ResumableUploadException(
          'Got status %d from attempt to start resumable upload. '
          'Aborting' % resp.status,
          ResumableTransferDisposition.ABORT)

    # Else we got 200 or 201 response code, indicating the resumable
    # upload was created.
    upload_url = resp.getheader('Location')
    if not upload_url:
      raise ResumableUploadException(
          'No resumable upload URL found in resumable initiation '
          'POST response (%s)' % body,
          ResumableTransferDisposition.WAIT_BEFORE_RETRY)
    self._SetUploadUrl(upload_url)
    self.tracker_callback(upload_url)

  def _UploadFileBytes(self, conn, http_conn, fp, file_length,
                       total_bytes_uploaded, cb, num_cb, headers):
    """Attempts to upload file bytes.

    Makes a single attempt using an existing resumable upload connection.

    Args:
      conn: HTTPConnection from the boto Key.
      http_conn: Separate HTTPConnection for the transfer.
      fp: File pointer containing bytes to upload.
      file_length: Total length of the file.
      total_bytes_uploaded: The total number of bytes uploaded.
      cb: Progress callback function that takes (progress, total_size).
      num_cb: Granularity of the callback (maximum number of times the
              callback will be called during the file transfer). If negative,
              perform callback with each buffer read.
      headers: Headers to be used in the upload requests.

    Returns:
      (etag, generation, metageneration) from service upon success.

    Raises:
      ResumableUploadException if any problems occur.
    """
    buf = fp.read(self.BUFFER_SIZE)
    if cb:
      # The cb_count represents the number of full buffers to send between
      # cb executions.
      if num_cb > 2:
        cb_count = file_length / self.BUFFER_SIZE / (num_cb-2)
      elif num_cb < 0:
        cb_count = -1
      else:
        cb_count = 0
      i = 0
      cb(total_bytes_uploaded, file_length)

    # Build resumable upload headers for the transfer. Don't send a
    # Content-Range header if the file is 0 bytes long, because the
    # resumable upload protocol uses an *inclusive* end-range (so, sending
    # 'bytes 0-0/1' would actually mean you're sending a 1-byte file).
    put_headers = headers.copy() if headers else {}
    if file_length:
      if total_bytes_uploaded == file_length:
        range_header = self._BuildContentRangeHeader(
            '*', file_length)
      else:
        range_header = self._BuildContentRangeHeader(
            '%d-%d' % (total_bytes_uploaded, file_length - 1),
            file_length)
      put_headers['Content-Range'] = range_header
    # Set Content-Length to the total bytes we'll send with this PUT.
    put_headers['Content-Length'] = str(file_length - total_bytes_uploaded)
    http_request = AWSAuthConnection.build_base_http_request(
        conn, 'PUT', path=self.upload_url_path, auth_path=None,
        headers=put_headers, host=self.upload_url_host)
    http_conn.putrequest('PUT', http_request.path)
    for k in put_headers:
      http_conn.putheader(k, put_headers[k])
    http_conn.endheaders()

    # Turn off debug on http connection so upload content isn't included
    # in debug stream.
    http_conn.set_debuglevel(0)
    while buf:
      http_conn.send(buf)
      total_bytes_uploaded += len(buf)
      if cb:
        i += 1
        if i == cb_count or cb_count == -1:
          cb(total_bytes_uploaded, file_length)
          i = 0
      buf = fp.read(self.BUFFER_SIZE)
    http_conn.set_debuglevel(conn.debug)
    if cb:
      cb(total_bytes_uploaded, file_length)
    if total_bytes_uploaded != file_length:
      # Abort (and delete the tracker file) so if the user retries
      # they'll start a new resumable upload rather than potentially
      # attempting to pick back up later where we left off.
      raise ResumableUploadException(
          'File changed during upload: EOF at %d bytes of %d byte file.' %
          (total_bytes_uploaded, file_length),
          ResumableTransferDisposition.ABORT)
    resp = http_conn.getresponse()
    # Restore http connection debug level.
    http_conn.set_debuglevel(conn.debug)

    if resp.status == 200:
      # Success.
      return (resp.getheader('etag'),
              resp.getheader('x-goog-generation'),
              resp.getheader('x-goog-metageneration'))
    # Retry timeout (408) and status 429, 500 and 503 errors after a delay.
    elif resp.status in [408, 429, 500, 503]:
      disposition = ResumableTransferDisposition.WAIT_BEFORE_RETRY
    else:
      # Catch all for any other error codes.
      disposition = ResumableTransferDisposition.ABORT
    raise ResumableUploadException('Got response code %d while attempting '
                                   'upload (%s)' %
                                   (resp.status, resp.reason), disposition)

  def _AttemptResumableUpload(self, key, fp, file_length, headers, cb,
                              num_cb):
    """Attempts a resumable upload.

    Args:
      key: Boto key representing object to upload.
      fp: File pointer containing upload bytes.
      file_length: Total length of the upload.
      headers: Headers to be used in upload requests.
      cb: Progress callback function that takes (progress, total_size).
      num_cb: Granularity of the callback (maximum number of times the
              callback will be called during the file transfer). If negative,
              perform callback with each buffer read.

    Returns:
      (etag, generation, metageneration) from service upon success.

    Raises:
      ResumableUploadException if any problems occur.
    """
    (service_start, service_end) = self.SERVICE_HAS_NOTHING
    conn = key.bucket.connection
    if self.upload_url:
      # Try to resume existing resumable upload.
      try:
        (service_start, service_end) = (
            self._QueryServicePos(conn, file_length))
        self.service_has_bytes = service_start
        if conn.debug >= 1:
          self.logger.debug('Resuming transfer.')
      except ResumableUploadException, e:
        if conn.debug >= 1:
          self.logger.debug('Unable to resume transfer (%s).', e.message)
        self._StartNewResumableUpload(key, headers)
    else:
      self._StartNewResumableUpload(key, headers)

    # upload_start_point allows the code that instantiated the
    # ResumableUploadHandler to find out the point from which it started
    # uploading (e.g., so it can correctly compute throughput).
    if self.upload_start_point is None:
      self.upload_start_point = service_end

    total_bytes_uploaded = service_end + 1

    # Start reading from the file based upon the number of bytes that the
    # server has so far.
    if total_bytes_uploaded < file_length:
      fp.seek(total_bytes_uploaded)

    conn = key.bucket.connection

    # Get a new HTTP connection (vs conn.get_http_connection(), which reuses
    # pool connections) because httplib requires a new HTTP connection per
    # transaction. (Without this, calling http_conn.getresponse() would get
    # "ResponseNotReady".)
    http_conn = conn.new_http_connection(self.upload_url_host, conn.port,
                                         conn.is_secure)
    http_conn.set_debuglevel(conn.debug)

    # Make sure to close http_conn at end so if a local file read
    # failure occurs partway through service will terminate current upload
    # and can report that progress on next attempt.
    try:
      return self._UploadFileBytes(conn, http_conn, fp, file_length,
                                   total_bytes_uploaded, cb, num_cb,
                                   headers)
    except (ResumableUploadException, socket.error):
      resp = self._QueryServiceState(conn, file_length)
      if resp.status == 400:
        raise ResumableUploadException(
            'Got 400 response from service state query after failed resumable '
            'upload attempt. This can happen for various reasons, including '
            'specifying an invalid request (e.g., an invalid canned ACL) or '
            'if the file size changed between upload attempts',
            ResumableTransferDisposition.ABORT)
      else:
        raise
    finally:
      http_conn.close()

  def HandleResumableUploadException(self, e, debug):
    if e.disposition == ResumableTransferDisposition.ABORT_CUR_PROCESS:
      if debug >= 1:
        self.logger.debug('Caught non-retryable ResumableUploadException (%s); '
                          'aborting but retaining tracker file', e.message)
      raise
    elif e.disposition == ResumableTransferDisposition.ABORT:
      if debug >= 1:
        self.logger.debug('Caught non-retryable ResumableUploadException (%s); '
                          'aborting and removing tracker file', e.message)
      raise
    elif e.disposition == ResumableTransferDisposition.START_OVER:
      raise
    else:
      if debug >= 1:
        self.logger.debug(
            'Caught ResumableUploadException (%s) - will retry', e.message)

  def TrackProgressLessIterations(self, service_had_bytes_before_attempt,
                                  debug=0):
    """Tracks the number of iterations without progress.

    Performs randomized exponential backoff.

    Args:
      service_had_bytes_before_attempt: Number of bytes the service had prior
                                       to this upload attempt.
      debug: debug level 0..3
    """
    # At this point we had a re-tryable failure; see if made progress.
    if self.service_has_bytes > service_had_bytes_before_attempt:
      self.progress_less_iterations = 0   # If progress, reset counter.
    else:
      self.progress_less_iterations += 1

    if self.progress_less_iterations > self.num_retries:
      # Don't retry any longer in the current process.
      raise ResumableUploadException(
          'Too many resumable upload attempts failed without '
          'progress. You might try this upload again later',
          ResumableTransferDisposition.ABORT_CUR_PROCESS)

    # Use binary exponential backoff to desynchronize client requests.
    sleep_time_secs = min(random.random() * (2**self.progress_less_iterations),
                          GetMaxRetryDelay())
    if debug >= 1:
      self.logger.debug('Got retryable failure (%d progress-less in a row).\n'
                        'Sleeping %3.1f seconds before re-trying',
                        self.progress_less_iterations, sleep_time_secs)
    time.sleep(sleep_time_secs)

  def SendFile(self, key, fp, size, headers, canned_acl=None, cb=None,
               num_cb=XML_PROGRESS_CALLBACKS):
    """Upload a file to a key into a bucket on GS, resumable upload protocol.

    Args:
      key: `boto.s3.key.Key` or subclass representing the upload destination.
      fp: File pointer to upload
      size: Size of the file to upload.
      headers: The headers to pass along with the PUT request
      canned_acl: Optional canned ACL to apply to object.
      cb: Callback function that will be called to report progress on
          the upload.  The callback should accept two integer parameters, the
          first representing the number of bytes that have been successfully
          transmitted to GS, and the second representing the total number of
          bytes that need to be transmitted.
      num_cb: (optional) If a callback is specified with the cb parameter, this
              parameter determines the granularity of the callback by defining
              the maximum number of times the callback will be called during the
              file transfer. Providing a negative integer will cause your
              callback to be called with each buffer read.

    Raises:
      ResumableUploadException if a problem occurs during the transfer.
    """

    if not headers:
      headers = {}
    # If Content-Type header is present and set to None, remove it.
    # This is gsutil's way of asking boto to refrain from auto-generating
    # that header.
    content_type = 'Content-Type'
    if content_type in headers and headers[content_type] is None:
      del headers[content_type]

    if canned_acl:
      headers[key.provider.acl_header] = canned_acl

    headers['User-Agent'] = UserAgent

    file_length = size
    debug = key.bucket.connection.debug

    # Use num-retries from constructor if one was provided; else check
    # for a value specified in the boto config file; else default to 5.
    if self.num_retries is None:
      self.num_retries = GetNumRetries()
    self.progress_less_iterations = 0

    while True:  # Retry as long as we're making progress.
      service_had_bytes_before_attempt = self.service_has_bytes
      try:
        # Save generation and metageneration in class state so caller
        # can find these values, for use in preconditions of future
        # operations on the uploaded object.
        (_, self.generation, self.metageneration) = (
            self._AttemptResumableUpload(key, fp, file_length,
                                         headers, cb, num_cb))

        key.generation = self.generation
        if debug >= 1:
          self.logger.debug('Resumable upload complete.')
        return
      except self.RETRYABLE_EXCEPTIONS, e:
        if debug >= 1:
          self.logger.debug('Caught exception (%s)', e.__repr__())
        if isinstance(e, IOError) and e.errno == errno.EPIPE:
          # Broken pipe error causes httplib to immediately
          # close the socket (http://bugs.python.org/issue5542),
          # so we need to close the connection before we resume
          # the upload (which will cause a new connection to be
          # opened the next time an HTTP request is sent).
          key.bucket.connection.connection.close()
      except ResumableUploadException, e:
        self.HandleResumableUploadException(e, debug)

      self.TrackProgressLessIterations(service_had_bytes_before_attempt,
                                       debug=debug)