Merge pull request #341 from CloudVE/atomic-ranged-download

Never assemble a download at its destination path

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

cloudbridge/base/resources.py

@@ -9,6 +9,7 @@ import os
 import queue
 import re
 import shutil
+import threading
 import time
 import uuid
 from concurrent.futures import FIRST_COMPLETED
@@ -886,14 +887,43 @@ class BaseBucketObject(BaseCloudResource, BucketObject):
 
     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 @@ class BaseBucketObject(BaseCloudResource, BucketObject):
         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 @@ class BaseBucketObject(BaseCloudResource, BucketObject):
 
         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 @@ class BaseBucketObject(BaseCloudResource, BucketObject):
                 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)

cloudbridge/interfaces/resources.py

@@ -2388,8 +2388,12 @@ class BucketObject(CloudResource):
         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

cloudbridge/providers/aws/resources.py

@@ -944,8 +944,8 @@ class AWSBucketObject(BaseBucketObject):
         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.

cloudbridge/providers/azure/resources.py

@@ -302,8 +302,8 @@ class AzureBucketObject(BaseBucketObject):
             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.

tests/test_download_driver.py

@@ -9,6 +9,7 @@ driver is exercised here directly against in-memory fakes so it has coverage
 in CI without cloud credentials.
 """
 import os
+import shutil
 import tempfile
 import threading
 import unittest
@@ -32,12 +33,15 @@ class _Recorder:
         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 @@ class DownloadDriverTestCase(unittest.TestCase):
             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)