Add back concurrency limits

Author: ponyisi

servicex/minio_adapter.py

@@ -28,22 +28,25 @@
 import os.path
 from hashlib import sha1
 from pathlib import Path
-from typing import List
+from typing import List, Optional
 
 from tenacity import retry, stop_after_attempt, wait_random_exponential
 
 import aioboto3
 from boto3.s3.transfer import TransferConfig
+import asyncio
 
 from servicex.models import ResultFile, TransformStatus
 
-_transferconfig = TransferConfig(max_concurrency=10)
+_transferconfig = TransferConfig(max_concurrency=5)
+_sem = asyncio.Semaphore(10)
+_bucket_list_sem = asyncio.Semaphore(5)
 
 
 def init_s3_config(concurrency: int = 10):
     "Update the number of concurrent connections"
-    global _transferconfig
-    _transferconfig = TransferConfig(max_concurrency=concurrency)
+    global _sem
+    _sem = asyncio.Semaphore(concurrency)
 
 
 def _sanitize_filename(fname: str):
@@ -85,24 +88,31 @@ def for_transform(cls, transform: TransformStatus):
         stop=stop_after_attempt(3), wait=wait_random_exponential(max=60), reraise=True
     )
     async def list_bucket(self) -> List[ResultFile]:
-        async with self.minio.resource("s3", endpoint_url=self.endpoint_host) as s3:
-            bucket = await s3.Bucket(self.bucket)
-            objects = bucket.objects.all()
-            return [
-                ResultFile(
-                    filename=obj.key,
-                    size=await obj.size,
-                    extension=obj.key.split(".")[-1],
-                )
-                async for obj in objects
-                if not obj.key.endswith("/")
-            ]
+        async with _bucket_list_sem:
+            async with self.minio.client("s3", endpoint_url=self.endpoint_host) as s3:
+                paginator = s3.get_paginator("list_objects_v2")
+                pagination = paginator.paginate(Bucket=self.bucket)
+                listing = await pagination.build_full_result()
+                rv = [
+                    ResultFile(
+                        filename=_["Key"],
+                        size=_["Size"],
+                        extension=_["Key"].split(".")[-1],
+                    )
+                    for _ in listing["Contents"]
+                    if not _["Key"].endswith("/")
+                ]
+                return rv
 
     @retry(
         stop=stop_after_attempt(3), wait=wait_random_exponential(max=60), reraise=True
     )
     async def download_file(
-        self, object_name: str, local_dir: str, shorten_filename: bool = False
+        self,
+        object_name: str,
+        local_dir: str,
+        shorten_filename: bool = False,
+        expected_size: Optional[int] = None,
     ) -> Path:
         os.makedirs(local_dir, exist_ok=True)
         path = Path(
@@ -114,16 +124,26 @@ async def download_file(
             )
         )
 
-        async with self.minio.resource("s3", endpoint_url=self.endpoint_host) as s3:
-            obj = await s3.Object(self.bucket, object_name)
-            remotesize = await obj.content_length
+        async with self.minio.client("s3", endpoint_url=self.endpoint_host) as s3:
+            if expected_size is not None:
+                remotesize = expected_size
+            else:
+                async with _sem:
+                    info = await s3.head_object(Bucket=self.bucket, Key=object_name)
+                    remotesize = info["ContentLength"]
             if path.exists():
                 # if file size is the same, let's not download anything
                 # maybe move to a better verification mechanism with e-tags in the future
                 localsize = path.stat().st_size
                 if localsize == remotesize:
                     return path.resolve()
-            await obj.download_file(path.as_posix(), Config=_transferconfig)
+            async with _sem:
+                await s3.download_file(
+                    Bucket=self.bucket,
+                    Key=object_name,
+                    Filename=path.as_posix(),
+                    Config=_transferconfig,
+                )
             localsize = path.stat().st_size
             if localsize != remotesize:
                 raise RuntimeError(f"Download of {object_name} failed")

servicex/query_core.py

@@ -533,9 +533,13 @@ async def download_file(
             progress: Progress,
             download_progress: TaskID,
             shorten_filename: bool = False,
+            expected_size: Optional[int] = None,
         ):
             downloaded_filename = await minio.download_file(
-                filename, self.download_path, shorten_filename=shorten_filename
+                filename,
+                self.download_path,
+                shorten_filename=shorten_filename,
+                expected_size=expected_size,
             )
             result_uris.append(downloaded_filename.as_posix())
             progress.advance(task_id=download_progress, task_type="Download")
@@ -580,6 +584,7 @@ async def get_signed_url(
                                             progress,
                                             download_progress,
                                             shorten_filename=self.configuration.shortened_downloaded_filename,  # NOQA: E501
+                                            expected_size=file.size,
                                         )
                                     )
                                 )  # NOQA 501

tests/test_minio_adapter.py

@@ -112,6 +112,19 @@ async def test_download_file(minio_adapter, populate_bucket):
     result.unlink()  # it should exist, from above ...
 
 
+@pytest.mark.parametrize("populate_bucket", ["test.txt"], indirect=True)
+@pytest.mark.asyncio
+async def test_download_file_with_expected_size(minio_adapter, populate_bucket):
+    info = await minio_adapter.list_bucket()
+    result = await minio_adapter.download_file(
+        "test.txt", local_dir="/tmp/foo", expected_size=info[0].size
+    )
+    assert str(result).endswith("test.txt")
+    assert result.exists()
+    assert result.read_bytes() == (b"\x01" * 10)
+    result.unlink()  # it should exist, from above ...
+
+
 @pytest.mark.parametrize("populate_bucket", ["t::est.txt"], indirect=True)
 @pytest.mark.asyncio
 async def test_download_bad_filename(minio_adapter, populate_bucket):

tests/test_servicex_dataset.py

@@ -216,7 +216,7 @@ async def test_submit(mocker):
     mock_minio = AsyncMock()
     mock_minio.list_bucket = AsyncMock(side_effect=[[file1], [file1, file2]])
     mock_minio.download_file = AsyncMock(
-        side_effect=lambda a, _, shorten_filename: PurePath(a)
+        side_effect=lambda a, _, shorten_filename, expected_size: PurePath(a)
     )
 
     mock_cache = mocker.MagicMock(QueryCache)
@@ -260,7 +260,7 @@ async def test_submit_partial_success(mocker):
     mock_minio = AsyncMock()
     mock_minio.list_bucket = AsyncMock(side_effect=[[file1], [file1]])
     mock_minio.download_file = AsyncMock(
-        side_effect=lambda a, _, shorten_filename: PurePath(a)
+        side_effect=lambda a, _, shorten_filename, expected_size: PurePath(a)
     )
 
     mock_cache = mocker.MagicMock(QueryCache)
@@ -303,7 +303,7 @@ async def test_use_of_cache(mocker):
     mock_minio = AsyncMock()
     mock_minio.list_bucket = AsyncMock(return_value=[file1, file2])
     mock_minio.download_file = AsyncMock(
-        side_effect=lambda a, _, shorten_filename: PurePath(a)
+        side_effect=lambda a, _, shorten_filename, expected_size: PurePath(a)
     )
     mock_minio.get_signed_url = AsyncMock(side_effect=["http://file1", "http://file2"])