Fix excessive memory usage in local storage range downloads

CHANGES.rst

@@ -66,7 +66,7 @@ Compute
   (GITHUB-1893)
   [Sanjay Santhanam - @Sanjays2402]
 
-- [UpCloud] 
+- [UpCloud]
 
   Add new functions to complete the Upcloud driver and move it to the
   last API version 1.3
@@ -76,6 +76,11 @@ Compute
 Storage
 ~~~~~~~
 
+- [Local] Stream range downloads in bounded chunks without reading the entire
+  source file or buffering the complete range before writing it to disk.
+  (GITHUB-2191)
+  [Neha Mahesh - @mneha05]
+
 - [Azure Blobs] Fix SAS signatures for percent-encoded object paths by using
   the decoded path in the canonicalized resource while preserving the encoded
   path in the returned URL.

libcloud/storage/drivers/local.py

@@ -27,9 +27,9 @@
 
 from libcloud.utils.py3 import u, relpath
 from libcloud.common.base import Connection
-from libcloud.utils.files import read_in_chunks, exhaust_iterator
+from libcloud.utils.files import read_in_chunks
 from libcloud.common.types import LibcloudError
-from libcloud.storage.base import Object, Container, StorageDriver
+from libcloud.storage.base import CHUNK_SIZE, Object, Container, StorageDriver
 from libcloud.storage.types import (
     ObjectError,
     ObjectDoesNotExistError,
@@ -500,28 +500,36 @@ def download_object_range(
         )
 
         with open(file_path, "wb") as fp:
-            fp.write(exhaust_iterator(iterator))
+            for chunk in iterator:
+                fp.write(chunk)
 
         return True
 
     def download_object_range_as_stream(self, obj, start_bytes, end_bytes=None, chunk_size=None):
         self._validate_start_and_end_bytes(start_bytes=start_bytes, end_bytes=end_bytes)
+        chunk_size = chunk_size or CHUNK_SIZE
+        if chunk_size < 0:
+            raise ValueError("chunk_size must be greater than 0")
 
         path = self.get_object_cdn_url(obj)
         with open(path, "rb") as obj_file:
-            file_size = len(obj_file.read())
+            file_size = os.fstat(obj_file.fileno()).st_size
 
             if end_bytes and end_bytes > file_size:
                 raise ValueError("end_bytes is larger than file size")
 
             if end_bytes is None:
-                read_bytes = (file_size - start_bytes) + 1
+                read_bytes = file_size - start_bytes
             else:
                 read_bytes = end_bytes - start_bytes
 
             obj_file.seek(start_bytes)
-            data = obj_file.read(read_bytes)
-            yield data
+            while read_bytes > 0:
+                data = obj_file.read(min(chunk_size, read_bytes))
+                if not data:
+                    break
+                read_bytes -= len(data)
+                yield data
 
     def upload_object(
         self,

libcloud/test/storage/test_local.py

@@ -22,10 +22,11 @@
 import tempfile
 import unittest
 import multiprocessing
+from unittest import mock
 
 from libcloud.utils.files import exhaust_iterator
 from libcloud.common.types import LibcloudError
-from libcloud.storage.base import Object, Container
+from libcloud.storage.base import CHUNK_SIZE, Object, Container
 from libcloud.storage.types import (
     ContainerIsNotEmptyError,
     InvalidContainerNameError,
@@ -603,6 +604,115 @@ def test_download_object_range_invalid_values(self):
             end_bytes=5,
         )
 
+    def test_download_object_range_as_stream_respects_chunk_size(self):
+        content = bytes(range(256)) * 100
+        container = self.driver.create_container("range-chunks")
+        path = os.path.join(self.key, container.name, "object")
+        with open(path, "wb") as fp:
+            fp.write(content)
+        obj = container.get_object("object")
+
+        for chunk_size in (3, None, 0):
+            for end_bytes in (20005, None):
+                with self.subTest(chunk_size=chunk_size, end_bytes=end_bytes):
+                    chunks = list(
+                        self.driver.download_object_range_as_stream(
+                            obj, start_bytes=5, end_bytes=end_bytes, chunk_size=chunk_size
+                        )
+                    )
+                    self.assertEqual(b"".join(chunks), content[5:end_bytes])
+                    self.assertTrue(
+                        all(0 < len(chunk) <= (chunk_size or CHUNK_SIZE) for chunk in chunks)
+                    )
+
+    def test_download_object_range_as_stream_reads_only_requested_bytes(self):
+        content = b"0123456789" * 100
+        container = self.driver.create_container("range-reads")
+        path = os.path.join(self.key, container.name, "object")
+        with open(path, "wb") as fp:
+            fp.write(content)
+        obj = container.get_object("object")
+
+        with open(path, "rb") as fp:
+            with mock.patch.object(fp, "read", wraps=fp.read) as read:
+                with mock.patch(
+                    "libcloud.storage.drivers.local.open", return_value=fp, create=True
+                ):
+                    chunks = list(
+                        self.driver.download_object_range_as_stream(
+                            obj, start_bytes=5, end_bytes=12, chunk_size=3
+                        )
+                    )
+
+        self.assertEqual(b"".join(chunks), content[5:12])
+        self.assertEqual(read.call_args_list, [mock.call(3), mock.call(3), mock.call(1)])
+
+    def test_download_object_range_writes_chunks_before_reading_next(self):
+        container = self.driver.create_container("range-writes")
+        path = os.path.join(self.key, container.name, "object")
+        with open(path, "wb") as fp:
+            fp.write(b"0123456789")
+        obj = container.get_object("object")
+        destination = os.path.join(self.key, "download")
+
+        with open(destination, "wb") as fp:
+            with mock.patch.object(fp, "write", wraps=fp.write) as write:
+
+                def chunks():
+                    yield b"123"
+                    write.assert_called_once_with(b"123")
+                    yield b"45"
+
+                with mock.patch.object(
+                    self.driver, "download_object_range_as_stream", return_value=chunks()
+                ):
+                    with mock.patch(
+                        "libcloud.storage.drivers.local.open", return_value=fp, create=True
+                    ):
+                        self.assertTrue(
+                            self.driver.download_object_range(
+                                obj,
+                                destination,
+                                start_bytes=1,
+                                end_bytes=6,
+                                overwrite_existing=True,
+                            )
+                        )
+
+        with open(destination, "rb") as fp:
+            self.assertEqual(fp.read(), b"12345")
+
+    def test_download_object_range_as_stream_rejects_negative_chunk_size(self):
+        container = self.driver.create_container("range-negative-chunk")
+        path = os.path.join(self.key, container.name, "object")
+        with open(path, "wb") as fp:
+            fp.write(b"0123456789")
+        obj = container.get_object("object")
+
+        stream = self.driver.download_object_range_as_stream(obj, start_bytes=1, chunk_size=-1)
+        self.assertRaisesRegex(ValueError, "chunk_size must be greater than 0", next, stream)
+
+    def test_download_object_range_as_stream_stops_at_early_eof(self):
+        container = self.driver.create_container("range-early-eof")
+        path = os.path.join(self.key, container.name, "object")
+        with open(path, "wb") as fp:
+            fp.write(b"0123456789")
+        obj = container.get_object("object")
+
+        with open(path, "rb") as fp:
+            with mock.patch.object(fp, "read", side_effect=[b"56", b""]) as read:
+                with mock.patch(
+                    "libcloud.storage.drivers.local.open", return_value=fp, create=True
+                ):
+                    chunks = list(
+                        self.driver.download_object_range_as_stream(
+                            obj, start_bytes=5, end_bytes=9, chunk_size=3
+                        )
+                    )
+
+        self.assertEqual(chunks, [b"56"])
+        self.assertEqual(read.call_args_list, [mock.call(3), mock.call(2)])
+
     def test_download_object_range_as_stream_invalid_values(self):
         content = b"0123456789123456789"
         tmppath = self.make_tmp_file(content=content)