chore: 优化s3获取文件列表慢的问题，和支持更多的进度条显示

Author: tianxiu2b2t

core/__init__.py

@@ -1,7 +1,6 @@
 import asyncio
 import base64
 import hashlib
-import hmac
 import io
 import time
 from typing import Optional
@@ -10,11 +9,11 @@
 import fastapi
 
 from core import units
-from core.abc import ResponseFile, ResponseFileLocal, ResponseFileMemory, ResponseFileRemote
+from core.abc import ResponseFileLocal, ResponseFileMemory, ResponseFileRemote
 
 from .locale import load_languages
 from .cluster import ClusterManager
-from .config import API_VERSION, ROOT_PATH, VERSION, cfg
+from .config import API_VERSION, VERSION, cfg
 from .logger import logger
 from .utils import runtime
 from . import web

core/cluster.py

@@ -402,12 +402,12 @@ def __init__(
         self._missing_files = missing_files
         self._clusters = clusters
         self._storages = storages
-        self._pbar = tqdm.tqdm(
-            total=sum(f.size for f in missing_files),
-            desc="Download",
-            unit="b",
-            unit_divisor=1024,
+        self._pbar = utils.MultiTQDM(
+            total=sum(missing_file.size for missing_file in missing_files),
+            description="Download",
+            unit="B",
             unit_scale=True,
+            unit_divisor=1024,
         )
         self._failed = 0
         self._success = 0
@@ -430,13 +430,16 @@ async def download(self):
         configuration = await self.get_configurations()
         logger.tinfo("download.configuration", source=configuration.source, concurrency=configuration.concurrency)
         async with anyio.create_task_group() as task_group:
-            missfiles = deque(self._missing_files)
+            queue = utils.Queue()
+            for file in self._missing_files:
+                queue.put_item(file)
             for _ in range(configuration.concurrency):
-                task_group.start_soon(self._download_files, missfiles)
+                task_group.start_soon(self._download_files, queue, _)
 
     async def _download_files(
         self,
-        files: deque[BMCLAPIFile]
+        files: utils.Queue[BMCLAPIFile],
+        worker_id: int
     ):
         async with aiohttp.ClientSession(
             base_url=cfg.base_url,
@@ -445,14 +448,21 @@ async def _download_files(
                 "User-Agent": USER_AGENT,
             }
         ) as session:
-            while len(files) != 0:
-                file = files.popleft()
-                await self._download_file(file, session)
+            with self._pbar.sub(0, f"Worker {worker_id}", unit="B", unit_scale=True, unit_divisor=1024) as pbar:
+                while len(files) != 0:
+                    file = await files.get_item()
+                    pbar._tqdm.total = file.size
+                    pbar._tqdm.n = 0
+                    pbar._tqdm.set_description_str(file.path)
+                    pbar._tqdm.refresh()
+                    pbar._tqdm.update(0)
+                    await self._download_file(file, session, pbar)
 
     async def _download_file(
         self,
         file: BMCLAPIFile,
-        session: aiohttp.ClientSession
+        session: aiohttp.ClientSession,
+        pbar: utils.SubTQDM
     ):
         last_error = None
         for _ in range(10):
@@ -461,6 +471,7 @@ async def _download_file(
             with tempfile.NamedTemporaryFile(
                 dir=self._cache_dir,
             ) as tmp_file:
+                print(pbar.position)
                 try:
                     async with session.get(
                         file.path
@@ -471,13 +482,15 @@ async def _download_file(
                             inc = len(data)
                             size += inc
                             self._pbar.update(inc)
+                            pbar.update(inc)
                         if hash.hexdigest() != file.hash or size != file.size:
                             await anyio.sleep(50)
                             raise Exception(f"hash mismatch, got {hash.hexdigest()} expected {file.hash}")
 
                 except Exception as e:
                     last_error = e
                     self._pbar.update(-size)
+                    pbar.update(-size)
                     self.update_failed()
                     continue
                 self.update_success()
@@ -585,11 +598,15 @@ async def sync(self):
         logger.tinfo("cluster.sync.files", count=len(files), size=units.format_bytes(total_size), last_modified=units.format_datetime_from_timestamp(last_modified))
 
         check_storages = [CheckStorage(storage) for storage in self.storages.storages]
-        missing_files: set[BMCLAPIFile] = set().union(
-            *await utils.gather(*[
-                check_storage.get_missing_files(files) for check_storage in check_storages
-            ])
-        )
+        with utils.MultiTQDM(
+            len(check_storages),
+            description="Listing files"
+        ) as pbar:
+            missing_files: set[BMCLAPIFile] = set().union(
+                *await utils.gather(*[
+                    check_storage.get_missing_files(files, pbar) for check_storage in check_storages
+                ])
+            )
         if len(missing_files) > 0:
             logger.tinfo("cluster.sync.missing_files", count=len(missing_files), size=units.format_bytes(sum([f.size for f in missing_files])))
             download_manager = DownloadManager(missing_files, self.clusters, check_storages)

core/storage/__init__.py

@@ -109,9 +109,9 @@ def __init__(
         self.files: dict[str, FileInfo] = {}
         self.missing_files: set[BMCLAPIFile] = set()
 
-    async def get_missing_files(self, bmclapi_files: set[BMCLAPIFile]):
+    async def get_missing_files(self, bmclapi_files: set[BMCLAPIFile], muitlpbar: utils.MultiTQDM):
 
-        for file in await self.storage.list_download_files():
+        for file in await self.storage.list_download_files(muitlpbar):
             self.files[file.name] = file
 
         # start check

core/storage/abc.py

@@ -68,16 +68,22 @@ async def list_files(
 
     async def list_download_files(
         self,
+        muitlpbar: utils.MultiTQDM
     ):
         res: list[FileInfo] = []
-        async def works(root_ids: list[int]):
-            for root_id in root_ids:
-                files = await self.list_files(f"download/{root_id:02x}")
-                res.extend(files)
-        async with anyio.create_task_group() as task_group:
-            work = utils.split_workload(list(RANGE), 10)
-            for w in work:
-                task_group.start_soon(works, w)
+        with muitlpbar.sub(
+            256,
+            description=f"Listing files in {self.name}({self.type})"
+        ) as pbar:
+            async def works(root_ids: list[int]):
+                for root_id in root_ids:
+                    files = await self.list_files(f"download/{root_id:02x}")
+                    res.extend(files)
+                    pbar.update(1)
+            async with anyio.create_task_group() as task_group:
+                work = utils.split_workload(list(RANGE), 10)
+                for w in work:
+                    task_group.start_soon(works, w)
         return res
 
     @abc.abstractmethod

core/storage/alist.py

@@ -24,6 +24,7 @@ def __repr__(self):
         return f"<AlistResponse code={self.code} data={self.data} message={self.message}>"
 
 class AlistStorage(abc.Storage):
+    type = "alist"
     def __init__(
         self,
         name: str,

core/storage/s3.py

@@ -1,5 +1,7 @@
+import inspect
 import tempfile
 import time
+from typing import Any
 import aioboto3.session
 import anyio.abc
 import anyio.to_thread
@@ -37,6 +39,7 @@ def __repr__(
         return f"S3Response(metadata={self.metadata}, data={self.raw_data})"
 
 class S3Storage(abc.Storage):
+    type = "s3"
     def __init__(
         self,
         name: str,
@@ -53,12 +56,21 @@ def __init__(
         self.bucket = bucket
         self.access_key = access_key
         self.secret_key = secret_key
+        self.region = kwargs.get("region")
         self.custom_s3_host = kwargs.get("custom_s3_host", "")
         self.public_endpoint = kwargs.get("public_endpoint", "")
         self.session = aioboto3.Session()
         self.list_lock = anyio.Lock()
         self.cache_list_bucket: dict[str, abc.FileInfo] = {}
         self.last_cache: float = 0
+        self._config = {
+            "endpoint_url": self.endpoint,
+            "aws_access_key_id": self.access_key,
+            "aws_secret_access_key": self.secret_key,
+        }
+        if self.region:
+            self._config["region_name"] = self.region
+            
 
     async def setup(
         self,
@@ -71,37 +83,52 @@ async def setup(
     async def list_bucket(
         self,
     ):
-        async with self.list_lock:
-            if time.perf_counter() - self.last_cache < 60:
-                return
-            async with self.session.resource(
-                "s3",
-                endpoint_url=self.endpoint,
-                aws_access_key_id=self.access_key,
-                aws_secret_access_key=self.secret_key,
-            ) as resource:
-                bucket = await resource.Bucket(self.bucket)
-                self.cache_list_bucket = {}
-                async for obj in bucket.objects.all():
-                    cp = abc.CPath("/" + obj.key)
-                    self.cache_list_bucket[str(cp)] = abc.FileInfo(
-                        path=str(cp),
-                        name=cp.name,
-                        size=await obj.size,
-                    )
-                self.last_cache = time.perf_counter()
+        ...
 
     async def list_files(
         self,
         path: str
     ) -> list[abc.FileInfo]:
-        await self.list_bucket()
-        # find by keys
         p = str(self.path / path)
         res = []
-        for key in self.cache_list_bucket.keys():
-            if str(abc.CPath(key).parents[-1]) == p:
-                res.append(self.cache_list_bucket[key])
+        async with self.session.client(
+            "s3",
+            endpoint_url=self.endpoint,
+            aws_access_key_id=self.access_key,
+            aws_secret_access_key=self.secret_key,
+            region_name=self.region
+        ) as client: # type: ignore
+            continuation_token = None
+            while True:
+                kwargs = {
+                    "Bucket": self.bucket,
+                    "Prefix": p[1:],
+                    #"Delimiter": "/",  # 使用分隔符来模拟文件夹结构
+                    #"MaxKeys": 1000
+                }
+                if continuation_token:
+                    kwargs["ContinuationToken"] = continuation_token
+
+                response = await client.list_objects_v2(**kwargs)
+                contents = response.get("Contents", [])
+                for content in contents:
+                    file_path = f"/{content['Key']}"
+                    if "/" in file_path:
+                        file_name = file_path.rsplit("/", 1)[1]
+                    else:
+                        file_name = file_path[1:]
+                    res.append(abc.FileInfo(
+                        name=file_name,
+                        size=content["Size"],
+                        path=f'/{content["Key"]}',
+                    ))
+
+                #res.extend(response.get("Contents", []))  # 添加文件
+                #res.extend(response.get("CommonPrefixes", []))  # 添加子目录
+
+                if "NextContinuationToken" not in response:
+                    break
+                continuation_token = response["NextContinuationToken"]
         return res
 
 
@@ -115,6 +142,7 @@ async def upload(
             endpoint_url=self.endpoint,
             aws_access_key_id=self.access_key,
             aws_secret_access_key=self.secret_key,
+            region_name=self.region
         ) as resource:
             bucket = await resource.Bucket(self.bucket)
             obj = await bucket.Object(str(self.path / path))
@@ -152,6 +180,7 @@ async def get_response_file(self, hash: str) -> abc.ResponseFile:
                 endpoint_url=self.endpoint,
                 aws_access_key_id=self.access_key,
                 aws_secret_access_key=self.secret_key,
+                region_name=self.region
             ) as client: # type: ignore
                 url = await client.generate_presigned_url(
                     ClientMethod="get_object",
@@ -182,6 +211,7 @@ async def get_response_file(self, hash: str) -> abc.ResponseFile:
             endpoint_url=self.endpoint,
             aws_access_key_id=self.access_key,
             aws_secret_access_key=self.secret_key,
+            region_name=self.region
         ) as resource:
             bucket = await resource.Bucket(self.bucket)
             obj = await bucket.Object(cpath)

core/storage/webdav.py

@@ -13,6 +13,7 @@
 import aiowebdav
 
 class WebDavStorage(abc.Storage):
+    type = "webdav"
     def __init__(
         self,
         name: str,