diff --git a/CHANGELOG.rst b/CHANGELOG.rst index ec191253..adb350e7 100644 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -19,6 +19,18 @@ a full 30. Measured against Route53, INSYNC was reached inside the first poll interval every time, making the granularity the entire cost. The waiter now polls every 5 seconds while keeping the same ~30 minute ceiling. +* **Downloads no longer assemble the object at the destination path.** + ``BucketObject.download_to_file`` builds the file out of the way and moves it + into place once complete, so the destination only ever holds a whole object. + Previously the generic ranged driver (used by GCP and OpenStack Swift) + created the destination up front and reopened it for every range, so anything + that replaced that path mid-transfer - notably a second download of the same + object to the same path, as a download cache does - could truncate the + in-progress file or make the next range fail with ``FileNotFoundError``. A + failed transfer no longer deletes an existing file at the destination either, + and the Azure downloader (which wrote in place) gains the same guarantee. + Ranges are now also written through a single file handle rather than + reopening the path per range. ## Build and CI * The AWS cloud integration job now requests a 3 hour OIDC session instead of diff --git a/cloudbridge/base/resources.py b/cloudbridge/base/resources.py index 911aebae..86b95456 100644 --- a/cloudbridge/base/resources.py +++ b/cloudbridge/base/resources.py @@ -9,6 +9,7 @@ import queue import re import shutil +import threading import time import uuid from concurrent.futures import FIRST_COMPLETED @@ -886,14 +887,43 @@ def save_content(self, target_stream: IO[bytes]) -> None: def download_to_file(self, path: str, config: TransferConfig | None = None) -> None: + # Assemble the object in a private file alongside the destination and + # rename it into place once complete, so ``path`` only ever holds a + # whole object. Callers commonly download every copy of an object to + # one well-known path (a cache entry, say), so writing in place would + # let concurrent downloads truncate each other's file - or rename it + # away mid-transfer - and would destroy a previously downloaded copy + # when a transfer fails. + part_path = f"{path}.{uuid.uuid4().hex}.cbpart" + try: + self._download_to_path(part_path, config) + os.replace(part_path, path) + except BaseException: + try: + os.remove(part_path) + except OSError: + pass + raise + + def _download_to_path(self, path: str, + config: TransferConfig | None = None) -> None: + """ + Write this object's content to ``path``, which the caller owns. + + Providers with an efficient, thread-safe native downloader (e.g. AWS + via boto3's ``download_file``, Azure via ``download_blob``) override + this to use it; the default implementation streams small objects and + fetches larger ones as parallel ranged reads. + """ size = self.size if size <= self._multipart_threshold(config): with open(path, 'wb') as f: self.save_content(f) return - self._download_ranged(path, size, config) + with open(path, 'w+b') as f: + self._download_ranged(f, size, config) - def _download_ranged(self, path: str, size: int, + def _download_ranged(self, target: IO[bytes], size: int, config: TransferConfig | None = None) -> None: """ Fetch the object as ranged reads across a bounded thread pool, @@ -902,40 +932,26 @@ def _download_ranged(self, path: str, size: int, To stay safe even on providers whose SDK client/connection is not thread-safe, each worker reads through its own cloned provider (see :meth:`.CloudProvider.clone`), so no provider state is shared between - threads. Memory is bounded to ~concurrency * part_size. On any - failure the partial file is removed and the error re-raised. - - Providers with an efficient, thread-safe native downloader (e.g. AWS - via boto3's ``download_file``, Azure via ``download_blob``) override - ``download_to_file`` to use it directly. + threads. Memory is bounded to ~concurrency * part_size. """ part_size = self._multipart_part_size(config) if part_size < 1: raise InvalidValueException('part_size', part_size) concurrency = max(1, self._multipart_max_concurrency(config)) + target.truncate(size) ranges = [(offset, min(part_size, size - offset)) for offset in range(0, size, part_size)] - try: - with open(path, 'wb') as f: - f.truncate(size) - if concurrency == 1: - bucket_objects = self._bucket_objects - with open(path, 'r+b') as f: - for offset, length in ranges: - f.seek(offset) - f.write(bucket_objects.download_range( - self.bucket, self.name, offset, length)) - else: - self._download_ranges_concurrently(path, ranges, concurrency) - except Exception: - try: - os.remove(path) - except OSError: - pass - raise + if concurrency == 1: + bucket_objects = self._bucket_objects + for offset, length in ranges: + target.seek(offset) + target.write(bucket_objects.download_range( + self.bucket, self.name, offset, length)) + else: + self._download_ranges_concurrently(target, ranges, concurrency) def _download_ranges_concurrently( - self, path: str, ranges: list[tuple[int, int]], + self, target: IO[bytes], ranges: list[tuple[int, int]], concurrency: int) -> None: # A pool of cloned bucket-object services, one per worker, so each # thread touches an isolated provider/connection. @@ -947,6 +963,7 @@ def _download_ranges_concurrently( bucket = self.bucket name = self.name + write_lock = threading.Lock() def fetch_one(offset: int, length: int) -> None: service = clones.get() @@ -954,13 +971,14 @@ def fetch_one(offset: int, length: int) -> None: data = service.download_range(bucket, name, offset, length) finally: clones.put(service) - # Each worker writes through its own handle at its own offset; - # ranges never overlap, so no locking is needed. Data is released - # as soon as it is written, bounding memory to - # ~concurrency * part_size. - with open(path, 'r+b') as f: - f.seek(offset) - f.write(data) + # Ranges are fetched in parallel but written through the one + # handle the caller opened, so a range can never be written to a + # file that has since been replaced. Serializing the writes costs + # little next to the fetches, and the data is released as soon as + # it is written, bounding memory to ~concurrency * part_size. + with write_lock: + target.seek(offset) + target.write(data) with ThreadPoolExecutor(max_workers=concurrency) as executor: futures = [executor.submit(fetch_one, offset, length) diff --git a/cloudbridge/interfaces/resources.py b/cloudbridge/interfaces/resources.py index 09d12fe9..3732ad33 100644 --- a/cloudbridge/interfaces/resources.py +++ b/cloudbridge/interfaces/resources.py @@ -2388,8 +2388,12 @@ def download_to_file(self, path: str, remain single-stream alternatives for arbitrary target streams. :type path: ``str`` - :param path: Local path to write the object's content to. An existing - file is overwritten; on failure no partial file is left behind. + :param path: Local path to write the object's content to. The object + is assembled out of the way and moved into place once complete, + so ``path`` never holds a partial object: an existing file is + replaced atomically, and a failed transfer leaves it untouched. + Concurrent downloads to one path are therefore safe, with the + last to complete winning. :type config: :class:`.TransferConfig` :param config: Optional per-call transfer tuning (threshold, part diff --git a/cloudbridge/providers/aws/resources.py b/cloudbridge/providers/aws/resources.py index 8299c9f7..de72cd1b 100644 --- a/cloudbridge/providers/aws/resources.py +++ b/cloudbridge/providers/aws/resources.py @@ -944,8 +944,8 @@ def upload_from_file(self, path: str, self._obj.upload_file(path, Config=transfer_config) return self - def download_to_file(self, path: str, - config: TransferConfig | None = None) -> None: + def _download_to_path(self, path: str, + config: TransferConfig | None = None) -> None: # boto3's TransferManager downloads large objects as parallel ranged # GETs with a thread-safe client, so the transparent ranged path # delegates to it rather than CloudBridge's generic clone-pool driver. diff --git a/cloudbridge/providers/azure/resources.py b/cloudbridge/providers/azure/resources.py index 251d76b6..c9035a64 100644 --- a/cloudbridge/providers/azure/resources.py +++ b/cloudbridge/providers/azure/resources.py @@ -302,8 +302,8 @@ def _upload_multipart(self, stream: IO[bytes], max_concurrency=self._multipart_max_concurrency(config)) return self - def download_to_file(self, path: str, - config: TransferConfig | None = None) -> None: + def _download_to_path(self, path: str, + config: TransferConfig | None = None) -> None: # azure-storage-blob's downloader fetches block ranges concurrently # with a thread-safe client, so delegate to it rather than # CloudBridge's generic clone-pool driver. diff --git a/tests/test_download_driver.py b/tests/test_download_driver.py index 773185c9..a8f74c18 100644 --- a/tests/test_download_driver.py +++ b/tests/test_download_driver.py @@ -9,6 +9,7 @@ in CI without cloud credentials. """ import os +import shutil import tempfile import threading import unittest @@ -32,12 +33,15 @@ def __init__(self, content): self.active = 0 self.max_active = 0 self.fail_on_offset = None # offset that should raise + self.on_serve = None # hook called as each range is served def serve_range(self, service, offset, length): with self._lock: self.active += 1 self.max_active = max(self.max_active, self.active) try: + if self.on_serve: + self.on_serve() if self.fail_on_offset == offset: raise RuntimeError("boom at offset %d" % offset) # Hold briefly so concurrent fetches genuinely overlap. @@ -222,6 +226,72 @@ def test_removes_partial_file_and_raises_on_range_failure(self): if os.path.exists(path): os.remove(path) + def test_destination_only_appears_once_complete(self): + content = bytes(range(256)) + recorder = _Recorder(content) + driver = self._driver( + recorder, threshold=1, part_size=16, concurrency=3) + fd, path = tempfile.mkstemp() + os.close(fd) + os.remove(path) + seen_early = [] + recorder.on_serve = lambda: seen_early.append(os.path.exists(path)) + try: + driver.download_to_file(path) + with open(path, 'rb') as f: + self.assertEqual(f.read(), content) + finally: + if os.path.exists(path): + os.remove(path) + # A partially written object is never visible at the destination. + self.assertTrue(seen_early) + self.assertNotIn(True, seen_early) + + def test_survives_concurrent_downloader_taking_the_destination(self): + # Galaxy gives every download of a dataset the same cache .tmp path, + # so a second download of the same dataset can rename the destination + # away while this one is still fetching ranges. + content = bytes(range(256)) + recorder = _Recorder(content) + driver = self._driver( + recorder, threshold=1, part_size=16, concurrency=3) + directory = tempfile.mkdtemp() + path = os.path.join(directory, 'dataset.dat') + taken = os.path.join(directory, 'taken.dat') + + def steal_destination(): + if os.path.exists(path): + os.replace(path, taken) + + recorder.on_serve = steal_destination + try: + driver.download_to_file(path) + with open(path, 'rb') as f: + self.assertEqual(f.read(), content) + finally: + shutil.rmtree(directory) + + def test_failed_download_leaves_an_existing_destination_intact(self): + content = bytes(range(64)) + recorder = _Recorder(content) + recorder.fail_on_offset = 16 + driver = self._driver( + recorder, threshold=1, part_size=16, concurrency=2) + directory = tempfile.mkdtemp() + path = os.path.join(directory, 'dataset.dat') + with open(path, 'wb') as f: + f.write(b'previously cached') + try: + with self.assertRaises(Exception): + driver.download_to_file(path) + # The cached copy survives a failed refetch, and no scratch file + # is left behind next to it. + with open(path, 'rb') as f: + self.assertEqual(f.read(), b'previously cached') + self.assertEqual(os.listdir(directory), ['dataset.dat']) + finally: + shutil.rmtree(directory) + def test_part_size_must_be_positive(self): content = bytes(range(16)) recorder = _Recorder(content)