tools/telemetry/third_party/gsutilz/gslib/daisy_chain_wrapper.py

   1 # -*- coding: utf-8 -*-
   2 # Copyright 2014 Google Inc. All Rights Reserved.
   3 #
   4 # Licensed under the Apache License, Version 2.0 (the "License");
   5 # you may not use this file except in compliance with the License.
   6 # You may obtain a copy of the License at
   7 #
   8 #     http://www.apache.org/licenses/LICENSE-2.0
   9 #
  10 # Unless required by applicable law or agreed to in writing, software
  11 # distributed under the License is distributed on an "AS IS" BASIS,
  12 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  13 # See the License for the specific language governing permissions and
  14 # limitations under the License.
  15 """Wrapper for use in daisy-chained copies."""
  16
  17 from collections import deque
  18 import os
  19 import threading
  20 import time
  21
  22 from gslib.cloud_api import BadRequestException
  23 from gslib.cloud_api import CloudApi
  24 from gslib.util import CreateLock
  25 from gslib.util import TRANSFER_BUFFER_SIZE
  26
  27
  28 # This controls the amount of bytes downloaded per download request.
  29 # We do not buffer this many bytes in memory at a time - that is controlled by
  30 # DaisyChainWrapper.max_buffer_size. This is the upper bound of bytes that may
  31 # be unnecessarily downloaded if there is a break in the resumable upload.
  32 _DEFAULT_DOWNLOAD_CHUNK_SIZE = 1024*1024*100
  33
  34
  35 class BufferWrapper(object):
  36   """Wraps the download file pointer to use our in-memory buffer."""
  37
  38   def __init__(self, daisy_chain_wrapper):
  39     """Provides a buffered write interface for a file download.
  40
  41     Args:
  42       daisy_chain_wrapper: DaisyChainWrapper instance to use for buffer and
  43                            locking.
  44     """
  45     self.daisy_chain_wrapper = daisy_chain_wrapper
  46
  47   def write(self, data):  # pylint: disable=invalid-name
  48     """Waits for space in the buffer, then writes data to the buffer."""
  49     while True:
  50       with self.daisy_chain_wrapper.lock:
  51         if (self.daisy_chain_wrapper.bytes_buffered <
  52             self.daisy_chain_wrapper.max_buffer_size):
  53           break
  54       # Buffer was full, yield thread priority so the upload can pull from it.
  55       time.sleep(0)
  56     data_len = len(data)
  57     if data_len:
  58       with self.daisy_chain_wrapper.lock:
  59         self.daisy_chain_wrapper.buffer.append(data)
  60         self.daisy_chain_wrapper.bytes_buffered += data_len
  61
  62
  63 class DaisyChainWrapper(object):
  64   """Wrapper class for daisy-chaining a cloud download to an upload.
  65
  66   This class instantiates a BufferWrapper object to buffer the download into
  67   memory, consuming a maximum of max_buffer_size. It implements intelligent
  68   behavior around read and seek that allow for all of the operations necessary
  69   to copy a file.
  70
  71   This class is coupled with the XML and JSON implementations in that it
  72   expects that small buffers (maximum of TRANSFER_BUFFER_SIZE) in size will be
  73   used.
  74   """
  75
  76   def __init__(self, src_url, src_obj_size, gsutil_api, progress_callback=None,
  77                download_chunk_size=_DEFAULT_DOWNLOAD_CHUNK_SIZE):
  78     """Initializes the daisy chain wrapper.
  79
  80     Args:
  81       src_url: Source CloudUrl to copy from.
  82       src_obj_size: Size of source object.
  83       gsutil_api: gsutil Cloud API to use for the copy.
  84       progress_callback: Optional callback function for progress notifications
  85           for the download thread. Receives calls with arguments
  86           (bytes_transferred, total_size).
  87       download_chunk_size: Integer number of bytes to download per
  88           GetObjectMedia request. This is the upper bound of bytes that may be
  89           unnecessarily downloaded if there is a break in the resumable upload.
  90
  91     """
  92     # Current read position for the upload file pointer.
  93     self.position = 0
  94     self.buffer = deque()
  95
  96     self.bytes_buffered = 0
  97     # Maximum amount of bytes in memory at a time.
  98     self.max_buffer_size = 1024 * 1024  # 1 MiB
  99
 100     self._download_chunk_size = download_chunk_size
 101
 102     # We save one buffer's worth of data as a special case for boto,
 103     # which seeks back one buffer and rereads to compute hashes. This is
 104     # unnecessary because we can just compare cloud hash digests at the end,
 105     # but it allows this to work without modfiying boto.
 106     self.last_position = 0
 107     self.last_data = None
 108
 109     # Protects buffer, position, bytes_buffered, last_position, and last_data.
 110     self.lock = CreateLock()
 111
 112     # Protects download_exception.
 113     self.download_exception_lock = CreateLock()
 114
 115     self.src_obj_size = src_obj_size
 116     self.src_url = src_url
 117
 118     # This is safe to use the upload and download thread because the download
 119     # thread calls only GetObjectMedia, which creates a new HTTP connection
 120     # independent of gsutil_api. Thus, it will not share an HTTP connection
 121     # with the upload.
 122     self.gsutil_api = gsutil_api
 123
 124     # If self.download_thread dies due to an exception, it is saved here so
 125     # that it can also be raised in the upload thread.
 126     self.download_exception = None
 127     self.download_thread = None
 128     self.progress_callback = progress_callback
 129     self.stop_download = threading.Event()
 130     self.StartDownloadThread(progress_callback=self.progress_callback)
 131
 132   def StartDownloadThread(self, start_byte=0, progress_callback=None):
 133     """Starts the download thread for the source object (from start_byte)."""
 134
 135     def PerformDownload(start_byte, progress_callback):
 136       """Downloads the source object in chunks.
 137
 138       This function checks the stop_download event and exits early if it is set.
 139       It should be set when there is an error during the daisy-chain upload,
 140       then this function can be called again with the upload's current position
 141       as start_byte.
 142
 143       Args:
 144         start_byte: Byte from which to begin the download.
 145         progress_callback: Optional callback function for progress
 146             notifications. Receives calls with arguments
 147             (bytes_transferred, total_size).
 148       """
 149       # TODO: Support resumable downloads. This would require the BufferWrapper
 150       # object to support seek() and tell() which requires coordination with
 151       # the upload.
 152       try:
 153         while start_byte + self._download_chunk_size < self.src_obj_size:
 154           self.gsutil_api.GetObjectMedia(
 155               self.src_url.bucket_name, self.src_url.object_name,
 156               BufferWrapper(self), start_byte=start_byte,
 157               end_byte=start_byte + self._download_chunk_size - 1,
 158               generation=self.src_url.generation, object_size=self.src_obj_size,
 159               download_strategy=CloudApi.DownloadStrategy.ONE_SHOT,
 160               provider=self.src_url.scheme, progress_callback=progress_callback)
 161           if self.stop_download.is_set():
 162             # Download thread needs to be restarted, so exit.
 163             self.stop_download.clear()
 164             return
 165           start_byte += self._download_chunk_size
 166         self.gsutil_api.GetObjectMedia(
 167             self.src_url.bucket_name, self.src_url.object_name,
 168             BufferWrapper(self), start_byte=start_byte,
 169             generation=self.src_url.generation, object_size=self.src_obj_size,
 170             download_strategy=CloudApi.DownloadStrategy.ONE_SHOT,
 171             provider=self.src_url.scheme, progress_callback=progress_callback)
 172       # We catch all exceptions here because we want to store them.
 173       except Exception, e:  # pylint: disable=broad-except
 174         # Save the exception so that it can be seen in the upload thread.
 175         with self.download_exception_lock:
 176           self.download_exception = e
 177           raise
 178
 179     # TODO: If we do gzip encoding transforms mid-transfer, this will fail.
 180     self.download_thread = threading.Thread(
 181         target=PerformDownload,
 182         args=(start_byte, progress_callback))
 183     self.download_thread.start()
 184
 185   def read(self, amt=None):  # pylint: disable=invalid-name
 186     """Exposes a stream from the in-memory buffer to the upload."""
 187     if self.position == self.src_obj_size or amt == 0:
 188       # If there is no data left or 0 bytes were requested, return an empty
 189       # string so callers can call still call len() and read(0).
 190       return ''
 191     if amt is None or amt > TRANSFER_BUFFER_SIZE:
 192       raise BadRequestException(
 193           'Invalid HTTP read size %s during daisy chain operation, '
 194           'expected <= %s.' % (amt, TRANSFER_BUFFER_SIZE))
 195
 196     while True:
 197       with self.lock:
 198         if self.buffer:
 199           break
 200         with self.download_exception_lock:
 201           if self.download_exception:
 202             # Download thread died, so we will never recover. Raise the
 203             # exception that killed it.
 204             raise self.download_exception  # pylint: disable=raising-bad-type
 205       # Buffer was empty, yield thread priority so the download thread can fill.
 206       time.sleep(0)
 207     with self.lock:
 208       # TODO: Need to handle the caller requesting less than a
 209       # transfer_buffer_size worth of data.
 210       data = self.buffer.popleft()
 211       self.last_position = self.position
 212       self.last_data = data
 213       data_len = len(data)
 214       self.position += data_len
 215       self.bytes_buffered -= data_len
 216     if data_len > amt:
 217       raise BadRequestException(
 218           'Invalid read during daisy chain operation, got data of size '
 219           '%s, expected size %s.' % (data_len, amt))
 220     return data
 221
 222   def tell(self):  # pylint: disable=invalid-name
 223     with self.lock:
 224       return self.position
 225
 226   def seek(self, offset, whence=os.SEEK_SET):  # pylint: disable=invalid-name
 227     restart_download = False
 228     if whence == os.SEEK_END:
 229       if offset:
 230         raise IOError(
 231             'Invalid seek during daisy chain operation. Non-zero offset %s '
 232             'from os.SEEK_END is not supported' % offset)
 233       with self.lock:
 234         self.last_position = self.position
 235         self.last_data = None
 236         # Safe because we check position against src_obj_size in read.
 237         self.position = self.src_obj_size
 238     elif whence == os.SEEK_SET:
 239       with self.lock:
 240         if offset == self.position:
 241           pass
 242         elif offset == self.last_position:
 243           self.position = self.last_position
 244           if self.last_data:
 245             # If we seek to end and then back, we won't have last_data; we'll
 246             # get it on the next call to read.
 247             self.buffer.appendleft(self.last_data)
 248             self.bytes_buffered += len(self.last_data)
 249         else:
 250           # Once a download is complete, boto seeks to 0 and re-reads to
 251           # compute the hash if an md5 isn't already present (for example a GCS
 252           # composite object), so we have to re-download the whole object.
 253           # Also, when daisy-chaining to a resumable upload, on error the
 254           # service may have received any number of the bytes; the download
 255           # needs to be restarted from that point.
 256           restart_download = True
 257
 258       if restart_download:
 259         self.stop_download.set()
 260
 261         # Consume any remaining bytes in the download thread so that
 262         # the thread can exit, then restart the thread at the desired position.
 263         while self.download_thread.is_alive():
 264           with self.lock:
 265             while self.bytes_buffered:
 266               self.bytes_buffered -= len(self.buffer.popleft())
 267           time.sleep(0)
 268
 269         with self.lock:
 270           self.position = offset
 271           self.buffer = deque()
 272           self.bytes_buffered = 0
 273           self.last_position = 0
 274           self.last_data = None
 275         self.StartDownloadThread(start_byte=offset,
 276                                  progress_callback=self.progress_callback)
 277     else:
 278       raise IOError('Daisy-chain download wrapper does not support '
 279                     'seek mode %s' % whence)
 280
 281   def seekable(self):  # pylint: disable=invalid-name
 282     return True