Просмотр исходного кода

Never assemble a download at its destination path

download_to_file built the object at the destination: the generic ranged
driver (GCP and OpenStack Swift) created the file up front and reopened
it by path for every range, and the Azure downloader wrote in place.
Callers commonly download every copy of an object to one well-known path
- a download cache keyed by object, say - so a second download of the
same object could truncate the first's file, and renaming that path into
place mid-transfer left the other download reopening a path that no
longer existed, failing with FileNotFoundError. A failed transfer also
deleted whatever was already at the destination.

Assemble into a private sibling file and rename it into place once
complete, so the destination only ever holds a whole object and a failed
transfer leaves an existing one untouched. Providers now fill a caller-
owned path via _download_to_path and inherit that guarantee. Ranges are
written through the single handle the driver opened, so a range can
never land in a file that has since been replaced.
Nuwan Goonasekera 2 недель назад
Родитель
Сommit
836da4e3fa

+ 12 - 0
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

+ 52 - 34
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)

+ 6 - 2
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

+ 2 - 2
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.

+ 2 - 2
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.

+ 70 - 0
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)