You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

712 lines
25 KiB

4 years ago
  1. # Copyright 2016 Amazon.com, Inc. or its affiliates. All Rights Reserved.
  2. #
  3. # Licensed under the Apache License, Version 2.0 (the "License"). You
  4. # may not use this file except in compliance with the License. A copy of
  5. # the License is located at
  6. #
  7. # http://aws.amazon.com/apache2.0/
  8. #
  9. # or in the "license" file accompanying this file. This file is
  10. # distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
  11. # ANY KIND, either express or implied. See the License for the specific
  12. # language governing permissions and limitations under the License.
  13. import random
  14. import time
  15. import functools
  16. import math
  17. import os
  18. import stat
  19. import string
  20. import logging
  21. import threading
  22. import io
  23. from collections import defaultdict
  24. from s3transfer.compat import rename_file
  25. from s3transfer.compat import seekable
  26. MAX_PARTS = 10000
  27. # The maximum file size you can upload via S3 per request.
  28. # See: http://docs.aws.amazon.com/AmazonS3/latest/dev/UploadingObjects.html
  29. # and: http://docs.aws.amazon.com/AmazonS3/latest/dev/qfacts.html
  30. MAX_SINGLE_UPLOAD_SIZE = 5 * (1024 ** 3)
  31. MIN_UPLOAD_CHUNKSIZE = 5 * (1024 ** 2)
  32. logger = logging.getLogger(__name__)
  33. def random_file_extension(num_digits=8):
  34. return ''.join(random.choice(string.hexdigits) for _ in range(num_digits))
  35. def signal_not_transferring(request, operation_name, **kwargs):
  36. if operation_name in ['PutObject', 'UploadPart'] and \
  37. hasattr(request.body, 'signal_not_transferring'):
  38. request.body.signal_not_transferring()
  39. def signal_transferring(request, operation_name, **kwargs):
  40. if operation_name in ['PutObject', 'UploadPart'] and \
  41. hasattr(request.body, 'signal_transferring'):
  42. request.body.signal_transferring()
  43. def calculate_range_parameter(part_size, part_index, num_parts,
  44. total_size=None):
  45. """Calculate the range parameter for multipart downloads/copies
  46. :type part_size: int
  47. :param part_size: The size of the part
  48. :type part_index: int
  49. :param part_index: The index for which this parts starts. This index starts
  50. at zero
  51. :type num_parts: int
  52. :param num_parts: The total number of parts in the transfer
  53. :returns: The value to use for Range parameter on downloads or
  54. the CopySourceRange parameter for copies
  55. """
  56. # Used to calculate the Range parameter
  57. start_range = part_index * part_size
  58. if part_index == num_parts - 1:
  59. end_range = ''
  60. if total_size is not None:
  61. end_range = str(total_size - 1)
  62. else:
  63. end_range = start_range + part_size - 1
  64. range_param = 'bytes=%s-%s' % (start_range, end_range)
  65. return range_param
  66. def get_callbacks(transfer_future, callback_type):
  67. """Retrieves callbacks from a subscriber
  68. :type transfer_future: s3transfer.futures.TransferFuture
  69. :param transfer_future: The transfer future the subscriber is associated
  70. to.
  71. :type callback_type: str
  72. :param callback_type: The type of callback to retrieve from the subscriber.
  73. Valid types include:
  74. * 'queued'
  75. * 'progress'
  76. * 'done'
  77. :returns: A list of callbacks for the type specified. All callbacks are
  78. preinjected with the transfer future.
  79. """
  80. callbacks = []
  81. for subscriber in transfer_future.meta.call_args.subscribers:
  82. callback_name = 'on_' + callback_type
  83. if hasattr(subscriber, callback_name):
  84. callbacks.append(
  85. functools.partial(
  86. getattr(subscriber, callback_name),
  87. future=transfer_future
  88. )
  89. )
  90. return callbacks
  91. def invoke_progress_callbacks(callbacks, bytes_transferred):
  92. """Calls all progress callbacks
  93. :param callbacks: A list of progress callbacks to invoke
  94. :param bytes_transferred: The number of bytes transferred. This is passed
  95. to the callbacks. If no bytes were transferred the callbacks will not
  96. be invoked because no progress was achieved. It is also possible
  97. to receive a negative amount which comes from retrying a transfer
  98. request.
  99. """
  100. # Only invoke the callbacks if bytes were actually transferred.
  101. if bytes_transferred:
  102. for callback in callbacks:
  103. callback(bytes_transferred=bytes_transferred)
  104. def get_filtered_dict(original_dict, whitelisted_keys):
  105. """Gets a dictionary filtered by whitelisted keys
  106. :param original_dict: The original dictionary of arguments to source keys
  107. and values.
  108. :param whitelisted_key: A list of keys to include in the filtered
  109. dictionary.
  110. :returns: A dictionary containing key/values from the original dictionary
  111. whose key was included in the whitelist
  112. """
  113. filtered_dict = {}
  114. for key, value in original_dict.items():
  115. if key in whitelisted_keys:
  116. filtered_dict[key] = value
  117. return filtered_dict
  118. class CallArgs(object):
  119. def __init__(self, **kwargs):
  120. """A class that records call arguments
  121. The call arguments must be passed as keyword arguments. It will set
  122. each keyword argument as an attribute of the object along with its
  123. associated value.
  124. """
  125. for arg, value in kwargs.items():
  126. setattr(self, arg, value)
  127. class FunctionContainer(object):
  128. """An object that contains a function and any args or kwargs to call it
  129. When called the provided function will be called with provided args
  130. and kwargs.
  131. """
  132. def __init__(self, func, *args, **kwargs):
  133. self._func = func
  134. self._args = args
  135. self._kwargs = kwargs
  136. def __repr__(self):
  137. return 'Function: %s with args %s and kwargs %s' % (
  138. self._func, self._args, self._kwargs)
  139. def __call__(self):
  140. return self._func(*self._args, **self._kwargs)
  141. class CountCallbackInvoker(object):
  142. """An abstraction to invoke a callback when a shared count reaches zero
  143. :param callback: Callback invoke when finalized count reaches zero
  144. """
  145. def __init__(self, callback):
  146. self._lock = threading.Lock()
  147. self._callback = callback
  148. self._count = 0
  149. self._is_finalized = False
  150. @property
  151. def current_count(self):
  152. with self._lock:
  153. return self._count
  154. def increment(self):
  155. """Increment the count by one"""
  156. with self._lock:
  157. if self._is_finalized:
  158. raise RuntimeError(
  159. 'Counter has been finalized it can no longer be '
  160. 'incremented.'
  161. )
  162. self._count += 1
  163. def decrement(self):
  164. """Decrement the count by one"""
  165. with self._lock:
  166. if self._count == 0:
  167. raise RuntimeError(
  168. 'Counter is at zero. It cannot dip below zero')
  169. self._count -= 1
  170. if self._is_finalized and self._count == 0:
  171. self._callback()
  172. def finalize(self):
  173. """Finalize the counter
  174. Once finalized, the counter never be incremented and the callback
  175. can be invoked once the count reaches zero
  176. """
  177. with self._lock:
  178. self._is_finalized = True
  179. if self._count == 0:
  180. self._callback()
  181. class OSUtils(object):
  182. def get_file_size(self, filename):
  183. return os.path.getsize(filename)
  184. def open_file_chunk_reader(self, filename, start_byte, size, callbacks):
  185. return ReadFileChunk.from_filename(filename, start_byte,
  186. size, callbacks,
  187. enable_callbacks=False)
  188. def open_file_chunk_reader_from_fileobj(self, fileobj, chunk_size,
  189. full_file_size, callbacks,
  190. close_callbacks=None):
  191. return ReadFileChunk(
  192. fileobj, chunk_size, full_file_size,
  193. callbacks=callbacks, enable_callbacks=False,
  194. close_callbacks=close_callbacks)
  195. def open(self, filename, mode):
  196. return open(filename, mode)
  197. def remove_file(self, filename):
  198. """Remove a file, noop if file does not exist."""
  199. # Unlike os.remove, if the file does not exist,
  200. # then this method does nothing.
  201. try:
  202. os.remove(filename)
  203. except OSError:
  204. pass
  205. def rename_file(self, current_filename, new_filename):
  206. rename_file(current_filename, new_filename)
  207. def is_special_file(cls, filename):
  208. """Checks to see if a file is a special UNIX file.
  209. It checks if the file is a character special device, block special
  210. device, FIFO, or socket.
  211. :param filename: Name of the file
  212. :returns: True if the file is a special file. False, if is not.
  213. """
  214. # If it does not exist, it must be a new file so it cannot be
  215. # a special file.
  216. if not os.path.exists(filename):
  217. return False
  218. mode = os.stat(filename).st_mode
  219. # Character special device.
  220. if stat.S_ISCHR(mode):
  221. return True
  222. # Block special device
  223. if stat.S_ISBLK(mode):
  224. return True
  225. # Named pipe / FIFO
  226. if stat.S_ISFIFO(mode):
  227. return True
  228. # Socket.
  229. if stat.S_ISSOCK(mode):
  230. return True
  231. return False
  232. class DeferredOpenFile(object):
  233. def __init__(self, filename, start_byte=0, mode='rb', open_function=open):
  234. """A class that defers the opening of a file till needed
  235. This is useful for deferring opening of a file till it is needed
  236. in a separate thread, as there is a limit of how many open files
  237. there can be in a single thread for most operating systems. The
  238. file gets opened in the following methods: ``read()``, ``seek()``,
  239. and ``__enter__()``
  240. :type filename: str
  241. :param filename: The name of the file to open
  242. :type start_byte: int
  243. :param start_byte: The byte to seek to when the file is opened.
  244. :type mode: str
  245. :param mode: The mode to use to open the file
  246. :type open_function: function
  247. :param open_function: The function to use to open the file
  248. """
  249. self._filename = filename
  250. self._fileobj = None
  251. self._start_byte = start_byte
  252. self._mode = mode
  253. self._open_function = open_function
  254. def _open_if_needed(self):
  255. if self._fileobj is None:
  256. self._fileobj = self._open_function(self._filename, self._mode)
  257. if self._start_byte != 0:
  258. self._fileobj.seek(self._start_byte)
  259. @property
  260. def name(self):
  261. return self._filename
  262. def read(self, amount=None):
  263. self._open_if_needed()
  264. return self._fileobj.read(amount)
  265. def write(self, data):
  266. self._open_if_needed()
  267. self._fileobj.write(data)
  268. def seek(self, where):
  269. self._open_if_needed()
  270. self._fileobj.seek(where)
  271. def tell(self):
  272. if self._fileobj is None:
  273. return self._start_byte
  274. return self._fileobj.tell()
  275. def close(self):
  276. if self._fileobj:
  277. self._fileobj.close()
  278. def __enter__(self):
  279. self._open_if_needed()
  280. return self
  281. def __exit__(self, *args, **kwargs):
  282. self.close()
  283. class ReadFileChunk(object):
  284. def __init__(self, fileobj, chunk_size, full_file_size,
  285. callbacks=None, enable_callbacks=True, close_callbacks=None):
  286. """
  287. Given a file object shown below::
  288. |___________________________________________________|
  289. 0 | | full_file_size
  290. |----chunk_size---|
  291. f.tell()
  292. :type fileobj: file
  293. :param fileobj: File like object
  294. :type chunk_size: int
  295. :param chunk_size: The max chunk size to read. Trying to read
  296. pass the end of the chunk size will behave like you've
  297. reached the end of the file.
  298. :type full_file_size: int
  299. :param full_file_size: The entire content length associated
  300. with ``fileobj``.
  301. :type callbacks: A list of function(amount_read)
  302. :param callbacks: Called whenever data is read from this object in the
  303. order provided.
  304. :type enable_callbacks: boolean
  305. :param enable_callbacks: True if to run callbacks. Otherwise, do not
  306. run callbacks
  307. :type close_callbacks: A list of function()
  308. :param close_callbacks: Called when close is called. The function
  309. should take no arguments.
  310. """
  311. self._fileobj = fileobj
  312. self._start_byte = self._fileobj.tell()
  313. self._size = self._calculate_file_size(
  314. self._fileobj, requested_size=chunk_size,
  315. start_byte=self._start_byte, actual_file_size=full_file_size)
  316. self._amount_read = 0
  317. self._callbacks = callbacks
  318. if callbacks is None:
  319. self._callbacks = []
  320. self._callbacks_enabled = enable_callbacks
  321. self._close_callbacks = close_callbacks
  322. if close_callbacks is None:
  323. self._close_callbacks = close_callbacks
  324. @classmethod
  325. def from_filename(cls, filename, start_byte, chunk_size, callbacks=None,
  326. enable_callbacks=True):
  327. """Convenience factory function to create from a filename.
  328. :type start_byte: int
  329. :param start_byte: The first byte from which to start reading.
  330. :type chunk_size: int
  331. :param chunk_size: The max chunk size to read. Trying to read
  332. pass the end of the chunk size will behave like you've
  333. reached the end of the file.
  334. :type full_file_size: int
  335. :param full_file_size: The entire content length associated
  336. with ``fileobj``.
  337. :type callbacks: function(amount_read)
  338. :param callbacks: Called whenever data is read from this object.
  339. :type enable_callbacks: bool
  340. :param enable_callbacks: Indicate whether to invoke callback
  341. during read() calls.
  342. :rtype: ``ReadFileChunk``
  343. :return: A new instance of ``ReadFileChunk``
  344. """
  345. f = open(filename, 'rb')
  346. f.seek(start_byte)
  347. file_size = os.fstat(f.fileno()).st_size
  348. return cls(f, chunk_size, file_size, callbacks, enable_callbacks)
  349. def _calculate_file_size(self, fileobj, requested_size, start_byte,
  350. actual_file_size):
  351. max_chunk_size = actual_file_size - start_byte
  352. return min(max_chunk_size, requested_size)
  353. def read(self, amount=None):
  354. if amount is None:
  355. amount_to_read = self._size - self._amount_read
  356. else:
  357. amount_to_read = min(self._size - self._amount_read, amount)
  358. data = self._fileobj.read(amount_to_read)
  359. self._amount_read += len(data)
  360. if self._callbacks is not None and self._callbacks_enabled:
  361. invoke_progress_callbacks(self._callbacks, len(data))
  362. return data
  363. def signal_transferring(self):
  364. self.enable_callback()
  365. if hasattr(self._fileobj, 'signal_transferring'):
  366. self._fileobj.signal_transferring()
  367. def signal_not_transferring(self):
  368. self.disable_callback()
  369. if hasattr(self._fileobj, 'signal_not_transferring'):
  370. self._fileobj.signal_not_transferring()
  371. def enable_callback(self):
  372. self._callbacks_enabled = True
  373. def disable_callback(self):
  374. self._callbacks_enabled = False
  375. def seek(self, where):
  376. self._fileobj.seek(self._start_byte + where)
  377. if self._callbacks is not None and self._callbacks_enabled:
  378. # To also rewind the callback() for an accurate progress report
  379. invoke_progress_callbacks(
  380. self._callbacks, bytes_transferred=where - self._amount_read)
  381. self._amount_read = where
  382. def close(self):
  383. if self._close_callbacks is not None and self._callbacks_enabled:
  384. for callback in self._close_callbacks:
  385. callback()
  386. self._fileobj.close()
  387. def tell(self):
  388. return self._amount_read
  389. def __len__(self):
  390. # __len__ is defined because requests will try to determine the length
  391. # of the stream to set a content length. In the normal case
  392. # of the file it will just stat the file, but we need to change that
  393. # behavior. By providing a __len__, requests will use that instead
  394. # of stat'ing the file.
  395. return self._size
  396. def __enter__(self):
  397. return self
  398. def __exit__(self, *args, **kwargs):
  399. self.close()
  400. def __iter__(self):
  401. # This is a workaround for http://bugs.python.org/issue17575
  402. # Basically httplib will try to iterate over the contents, even
  403. # if its a file like object. This wasn't noticed because we've
  404. # already exhausted the stream so iterating over the file immediately
  405. # stops, which is what we're simulating here.
  406. return iter([])
  407. class StreamReaderProgress(object):
  408. """Wrapper for a read only stream that adds progress callbacks."""
  409. def __init__(self, stream, callbacks=None):
  410. self._stream = stream
  411. self._callbacks = callbacks
  412. if callbacks is None:
  413. self._callbacks = []
  414. def read(self, *args, **kwargs):
  415. value = self._stream.read(*args, **kwargs)
  416. invoke_progress_callbacks(self._callbacks, len(value))
  417. return value
  418. class NoResourcesAvailable(Exception):
  419. pass
  420. class TaskSemaphore(object):
  421. def __init__(self, count):
  422. """A semaphore for the purpose of limiting the number of tasks
  423. :param count: The size of semaphore
  424. """
  425. self._semaphore = threading.Semaphore(count)
  426. def acquire(self, tag, blocking=True):
  427. """Acquire the semaphore
  428. :param tag: A tag identifying what is acquiring the semaphore. Note
  429. that this is not really needed to directly use this class but is
  430. needed for API compatibility with the SlidingWindowSemaphore
  431. implementation.
  432. :param block: If True, block until it can be acquired. If False,
  433. do not block and raise an exception if cannot be aquired.
  434. :returns: A token (can be None) to use when releasing the semaphore
  435. """
  436. logger.debug("Acquiring %s", tag)
  437. if not self._semaphore.acquire(blocking):
  438. raise NoResourcesAvailable("Cannot acquire tag '%s'" % tag)
  439. def release(self, tag, acquire_token):
  440. """Release the semaphore
  441. :param tag: A tag identifying what is releasing the semaphore
  442. :param acquire_token: The token returned from when the semaphore was
  443. acquired. Note that this is not really needed to directly use this
  444. class but is needed for API compatibility with the
  445. SlidingWindowSemaphore implementation.
  446. """
  447. logger.debug("Releasing acquire %s/%s" % (tag, acquire_token))
  448. self._semaphore.release()
  449. class SlidingWindowSemaphore(TaskSemaphore):
  450. """A semaphore used to coordinate sequential resource access.
  451. This class is similar to the stdlib BoundedSemaphore:
  452. * It's initialized with a count.
  453. * Each call to ``acquire()`` decrements the counter.
  454. * If the count is at zero, then ``acquire()`` will either block until the
  455. count increases, or if ``blocking=False``, then it will raise
  456. a NoResourcesAvailable exception indicating that it failed to acquire the
  457. semaphore.
  458. The main difference is that this semaphore is used to limit
  459. access to a resource that requires sequential access. For example,
  460. if I want to access resource R that has 20 subresources R_0 - R_19,
  461. this semaphore can also enforce that you only have a max range of
  462. 10 at any given point in time. You must also specify a tag name
  463. when you acquire the semaphore. The sliding window semantics apply
  464. on a per tag basis. The internal count will only be incremented
  465. when the minimum sequence number for a tag is released.
  466. """
  467. def __init__(self, count):
  468. self._count = count
  469. # Dict[tag, next_sequence_number].
  470. self._tag_sequences = defaultdict(int)
  471. self._lowest_sequence = {}
  472. self._lock = threading.Lock()
  473. self._condition = threading.Condition(self._lock)
  474. # Dict[tag, List[sequence_number]]
  475. self._pending_release = {}
  476. def current_count(self):
  477. with self._lock:
  478. return self._count
  479. def acquire(self, tag, blocking=True):
  480. logger.debug("Acquiring %s", tag)
  481. self._condition.acquire()
  482. try:
  483. if self._count == 0:
  484. if not blocking:
  485. raise NoResourcesAvailable("Cannot acquire tag '%s'" % tag)
  486. else:
  487. while self._count == 0:
  488. self._condition.wait()
  489. # self._count is no longer zero.
  490. # First, check if this is the first time we're seeing this tag.
  491. sequence_number = self._tag_sequences[tag]
  492. if sequence_number == 0:
  493. # First time seeing the tag, so record we're at 0.
  494. self._lowest_sequence[tag] = sequence_number
  495. self._tag_sequences[tag] += 1
  496. self._count -= 1
  497. return sequence_number
  498. finally:
  499. self._condition.release()
  500. def release(self, tag, acquire_token):
  501. sequence_number = acquire_token
  502. logger.debug("Releasing acquire %s/%s", tag, sequence_number)
  503. self._condition.acquire()
  504. try:
  505. if tag not in self._tag_sequences:
  506. raise ValueError("Attempted to release unknown tag: %s" % tag)
  507. max_sequence = self._tag_sequences[tag]
  508. if self._lowest_sequence[tag] == sequence_number:
  509. # We can immediately process this request and free up
  510. # resources.
  511. self._lowest_sequence[tag] += 1
  512. self._count += 1
  513. self._condition.notify()
  514. queued = self._pending_release.get(tag, [])
  515. while queued:
  516. if self._lowest_sequence[tag] == queued[-1]:
  517. queued.pop()
  518. self._lowest_sequence[tag] += 1
  519. self._count += 1
  520. else:
  521. break
  522. elif self._lowest_sequence[tag] < sequence_number < max_sequence:
  523. # We can't do anything right now because we're still waiting
  524. # for the min sequence for the tag to be released. We have
  525. # to queue this for pending release.
  526. self._pending_release.setdefault(
  527. tag, []).append(sequence_number)
  528. self._pending_release[tag].sort(reverse=True)
  529. else:
  530. raise ValueError(
  531. "Attempted to release unknown sequence number "
  532. "%s for tag: %s" % (sequence_number, tag))
  533. finally:
  534. self._condition.release()
  535. class ChunksizeAdjuster(object):
  536. def __init__(self, max_size=MAX_SINGLE_UPLOAD_SIZE,
  537. min_size=MIN_UPLOAD_CHUNKSIZE, max_parts=MAX_PARTS):
  538. self.max_size = max_size
  539. self.min_size = min_size
  540. self.max_parts = max_parts
  541. def adjust_chunksize(self, current_chunksize, file_size=None):
  542. """Get a chunksize close to current that fits within all S3 limits.
  543. :type current_chunksize: int
  544. :param current_chunksize: The currently configured chunksize.
  545. :type file_size: int or None
  546. :param file_size: The size of the file to upload. This might be None
  547. if the object being transferred has an unknown size.
  548. :returns: A valid chunksize that fits within configured limits.
  549. """
  550. chunksize = current_chunksize
  551. if file_size is not None:
  552. chunksize = self._adjust_for_max_parts(chunksize, file_size)
  553. return self._adjust_for_chunksize_limits(chunksize)
  554. def _adjust_for_chunksize_limits(self, current_chunksize):
  555. if current_chunksize > self.max_size:
  556. logger.debug(
  557. "Chunksize greater than maximum chunksize. "
  558. "Setting to %s from %s." % (self.max_size, current_chunksize))
  559. return self.max_size
  560. elif current_chunksize < self.min_size:
  561. logger.debug(
  562. "Chunksize less than minimum chunksize. "
  563. "Setting to %s from %s." % (self.min_size, current_chunksize))
  564. return self.min_size
  565. else:
  566. return current_chunksize
  567. def _adjust_for_max_parts(self, current_chunksize, file_size):
  568. chunksize = current_chunksize
  569. num_parts = int(math.ceil(file_size / float(chunksize)))
  570. while num_parts > self.max_parts:
  571. chunksize *= 2
  572. num_parts = int(math.ceil(file_size / float(chunksize)))
  573. if chunksize != current_chunksize:
  574. logger.debug(
  575. "Chunksize would result in the number of parts exceeding the "
  576. "maximum. Setting to %s from %s." %
  577. (chunksize, current_chunksize))
  578. return chunksize