| OLD | NEW |
| (Empty) |
| 1 # Copyright 2014 Google Inc. All Rights Reserved. | |
| 2 # | |
| 3 # Licensed under the Apache License, Version 2.0 (the "License"); | |
| 4 # you may not use this file except in compliance with the License. | |
| 5 # You may obtain a copy of the License at | |
| 6 # | |
| 7 # http://www.apache.org/licenses/LICENSE-2.0 | |
| 8 # | |
| 9 # Unless required by applicable law or agreed to in writing, software | |
| 10 # distributed under the License is distributed on an "AS IS" BASIS, | |
| 11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | |
| 12 # See the License for the specific language governing permissions and | |
| 13 # limitations under the License. | |
| 14 """Helper class for streaming resumable uploads.""" | |
| 15 | |
| 16 import collections | |
| 17 import os | |
| 18 | |
| 19 from gslib.exception import CommandException | |
| 20 from gslib.util import GetJsonResumableChunkSize | |
| 21 | |
| 22 | |
| 23 class ResumableStreamingJsonUploadWrapper(object): | |
| 24 """Wraps an input stream in a buffer for resumable uploads. | |
| 25 | |
| 26 This class takes a non-seekable input stream, buffers it, and exposes it | |
| 27 as a stream with limited seek capabilities such that it can be used in a | |
| 28 resumable JSON API upload. | |
| 29 | |
| 30 max_buffer_size bytes of buffering is supported. | |
| 31 """ | |
| 32 | |
| 33 def __init__(self, stream, max_buffer_size, test_small_buffer=False): | |
| 34 """Initializes the wrapper. | |
| 35 | |
| 36 Args: | |
| 37 stream: Input stream. | |
| 38 max_buffer_size: Maximum size of internal buffer; should be >= the chunk | |
| 39 size of the resumable upload API to ensure that at least one full | |
| 40 chunk write can be replayed in the event of a server error. | |
| 41 test_small_buffer: Skip check for buffer size vs. chunk size, for testing. | |
| 42 """ | |
| 43 self._orig_fp = stream | |
| 44 | |
| 45 if not test_small_buffer and max_buffer_size < GetJsonResumableChunkSize(): | |
| 46 raise CommandException('Resumable streaming upload created with buffer ' | |
| 47 'size %s, JSON resumable upload chunk size %s. ' | |
| 48 'Buffer size must be >= JSON resumable upload ' | |
| 49 'chunk size to ensure that uploads can be ' | |
| 50 'resumed.' % (max_buffer_size, | |
| 51 GetJsonResumableChunkSize())) | |
| 52 | |
| 53 self._max_buffer_size = max_buffer_size | |
| 54 self._buffer = collections.deque() | |
| 55 self._buffer_start = 0 | |
| 56 self._buffer_end = 0 | |
| 57 self._position = 0 | |
| 58 | |
| 59 def read(self, size=-1): # pylint: disable=invalid-name | |
| 60 """"Reads from the wrapped stream. | |
| 61 | |
| 62 Args: | |
| 63 size: The amount of bytes to read. If omitted or negative, the entire | |
| 64 contents of the stream will be read and returned. | |
| 65 | |
| 66 Returns: | |
| 67 Bytes from the wrapped stream. | |
| 68 """ | |
| 69 read_all_bytes = size is None or size < 0 | |
| 70 if read_all_bytes: | |
| 71 bytes_remaining = self._max_buffer_size | |
| 72 else: | |
| 73 bytes_remaining = size | |
| 74 data = b'' | |
| 75 buffered_data = [] | |
| 76 if self._position < self._buffer_end: | |
| 77 # There was a backwards seek, so read from the buffer first. | |
| 78 | |
| 79 # TODO: Performance test to validate if it is worth re-aligning | |
| 80 # the buffers in this case. Also, seeking through the buffer for | |
| 81 # each read on a long catch-up is probably not performant, but we'd | |
| 82 # need a more complex data structure than a deque to get around this. | |
| 83 pos_in_buffer = self._buffer_start | |
| 84 buffer_index = 0 | |
| 85 # First, find the start position in the buffer. | |
| 86 while pos_in_buffer + len(self._buffer[buffer_index]) < self._position: | |
| 87 # When this loop exits, buffer_index will refer to a buffer that | |
| 88 # has at least some overlap with self._position, and | |
| 89 # pos_in_buffer will be >= self._position | |
| 90 pos_in_buffer += len(self._buffer[buffer_index]) | |
| 91 buffer_index += 1 | |
| 92 | |
| 93 # Read until we've read enough or we're out of buffer. | |
| 94 while pos_in_buffer < self._buffer_end and bytes_remaining > 0: | |
| 95 buffer_len = len(self._buffer[buffer_index]) | |
| 96 # This describes how far into the current buffer self._position is. | |
| 97 offset_from_position = self._position - pos_in_buffer | |
| 98 bytes_available_this_buffer = buffer_len - offset_from_position | |
| 99 read_size = min(bytes_available_this_buffer, bytes_remaining) | |
| 100 buffered_data.append( | |
| 101 self._buffer[buffer_index] | |
| 102 [offset_from_position:offset_from_position + read_size]) | |
| 103 bytes_remaining -= read_size | |
| 104 pos_in_buffer += buffer_len | |
| 105 buffer_index += 1 | |
| 106 self._position += read_size | |
| 107 | |
| 108 # At this point we're guaranteed that if there are any bytes left to read, | |
| 109 # then self._position == self._buffer_end, and we can read from the | |
| 110 # wrapped stream if needed. | |
| 111 if read_all_bytes: | |
| 112 # TODO: The user is requesting reading until the end of an | |
| 113 # arbitrary length stream, which is bad we'll need to return data | |
| 114 # with no size limits; if the stream is sufficiently long, we could run | |
| 115 # out of memory. We could break this down into smaller reads and | |
| 116 # buffer it as we go, but we're still left returning the data all at | |
| 117 # once to the caller. We could raise, but for now trust the caller to | |
| 118 # be sane and have enough memory to hold the remaining stream contents. | |
| 119 new_data = self._orig_fp.read(size) | |
| 120 data_len = len(new_data) | |
| 121 if not buffered_data: | |
| 122 data = new_data | |
| 123 else: | |
| 124 buffered_data.append(new_data) | |
| 125 data = b''.join(buffered_data) | |
| 126 self._position += data_len | |
| 127 elif bytes_remaining: | |
| 128 new_data = self._orig_fp.read(bytes_remaining) | |
| 129 if not buffered_data: | |
| 130 data = new_data | |
| 131 else: | |
| 132 buffered_data.append(new_data) | |
| 133 data = b''.join(buffered_data) | |
| 134 data_len = len(new_data) | |
| 135 if data_len: | |
| 136 self._position += data_len | |
| 137 self._buffer.append(new_data) | |
| 138 self._buffer_end += data_len | |
| 139 oldest_data = None | |
| 140 while self._buffer_end - self._buffer_start > self._max_buffer_size: | |
| 141 oldest_data = self._buffer.popleft() | |
| 142 self._buffer_start += len(oldest_data) | |
| 143 if oldest_data: | |
| 144 refill_amount = self._max_buffer_size - (self._buffer_end - | |
| 145 self._buffer_start) | |
| 146 if refill_amount: | |
| 147 self._buffer.appendleft(oldest_data[-refill_amount:]) | |
| 148 self._buffer_start -= refill_amount | |
| 149 else: | |
| 150 data = b''.join(buffered_data) if buffered_data else b'' | |
| 151 | |
| 152 return data | |
| 153 | |
| 154 def tell(self): # pylint: disable=invalid-name | |
| 155 """Returns the current stream position.""" | |
| 156 return self._position | |
| 157 | |
| 158 def seekable(self): # pylint: disable=invalid-name | |
| 159 """Returns true since limited seek support exists.""" | |
| 160 return True | |
| 161 | |
| 162 def seek(self, offset, whence=os.SEEK_SET): # pylint: disable=invalid-name | |
| 163 """Seeks on the buffered stream. | |
| 164 | |
| 165 Args: | |
| 166 offset: The offset to seek to; must be within the buffer bounds. | |
| 167 whence: Must be os.SEEK_SET. | |
| 168 | |
| 169 Raises: | |
| 170 CommandException if an unsupported seek mode or position is used. | |
| 171 """ | |
| 172 if whence == os.SEEK_SET: | |
| 173 if offset < self._buffer_start or offset > self._buffer_end: | |
| 174 raise CommandException('Unable to resume upload because of limited ' | |
| 175 'buffering available for streaming uploads. ' | |
| 176 'Offset %s was requested, but only data from ' | |
| 177 '%s to %s is buffered.' % | |
| 178 (offset, self._buffer_start, self._buffer_end)) | |
| 179 # Move to a position within the buffer. | |
| 180 self._position = offset | |
| 181 elif whence == os.SEEK_END: | |
| 182 if offset > self._max_buffer_size: | |
| 183 raise CommandException('Invalid SEEK_END offset %s on streaming ' | |
| 184 'upload. Only %s can be buffered.' % | |
| 185 (offset, self._max_buffer_size)) | |
| 186 # Read to the end and rely on buffering to handle the offset. | |
| 187 while self.read(self._max_buffer_size): | |
| 188 pass | |
| 189 # Now we're at the end. | |
| 190 self._position -= offset | |
| 191 else: | |
| 192 raise CommandException('Invalid seek mode on streaming upload. ' | |
| 193 '(mode %s, offset %s)' % (whence, offset)) | |
| 194 | |
| 195 def close(self): # pylint: disable=invalid-name | |
| 196 return self._orig_fp.close() | |
| OLD | NEW |