From 6b728abbeb5e652b96d677d59ef352045d916282 Mon Sep 17 00:00:00 2001
From: jerryliang <jerryliang@jerryliang.top>
Date: Fri, 16 Jan 2026 17:33:22 +0800
Subject: [PATCH] =?UTF-8?q?perf:=20=E5=A4=8D=E7=94=A8ClientSession?=
 =?UTF-8?q?=E6=8F=90=E5=8D=87=E6=96=87=E4=BB=B6=E4=B8=8B=E8=BD=BD=E6=80=A7?=
 =?UTF-8?q?=E8=83=BD?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 core/storage.py | 41 ++++++++++++++++++++++++++---------------
 1 file changed, 26 insertions(+), 15 deletions(-)

diff --git a/core/storage.py b/core/storage.py
index 17d354b..f04d1d8 100644
--- a/core/storage.py
+++ b/core/storage.py
@@ -337,8 +337,11 @@ class S3FileStorage(FileStorageInterface):
                     ExpiresIn=3600,
                 )
             
+            # 创建ClientSession并传递给生成器复用
+            session = aiohttp.ClientSession()
+            
             async def stream_generator():
-                async with aiohttp.ClientSession() as session:
+                try:
                     async with session.get(link) as resp:
                         if resp.status != 200:
                             raise HTTPException(
@@ -352,6 +355,8 @@ class S3FileStorage(FileStorageInterface):
                             if not chunk:
                                 break
                             yield chunk
+                finally:
+                    await session.close()
             
             from fastapi.responses import StreamingResponse
             encoded_filename = quote(filename, safe='')
@@ -647,18 +652,20 @@ class OneDriveFileStorage(FileStorageInterface):
             
             content_length = file_code.size  # 默认使用数据库中的大小
             
+            # 创建ClientSession并复用
+            session = aiohttp.ClientSession()
+            
             # 尝试发送HEAD请求获取Content-Length
             try:
-                async with aiohttp.ClientSession() as session:
-                    async with session.head(link) as resp:
-                        if resp.status == 200 and 'Content-Length' in resp.headers:
-                            content_length = int(resp.headers['Content-Length'])
+                async with session.head(link) as resp:
+                    if resp.status == 200 and 'Content-Length' in resp.headers:
+                        content_length = int(resp.headers['Content-Length'])
             except Exception:
                 # 如果HEAD请求失败，继续使用默认大小
                 pass
             
             async def stream_generator():
-                async with aiohttp.ClientSession() as session:
+                try:
                     async with session.get(link) as resp:
                         if resp.status != 200:
                             raise HTTPException(
@@ -671,6 +678,8 @@ class OneDriveFileStorage(FileStorageInterface):
                             if not chunk:
                                 break
                             yield chunk
+                finally:
+                    await session.close()
             
             encoded_filename = quote(filename, safe='')
             headers = {
@@ -1080,22 +1089,22 @@ class WebDAVFileStorage(FileStorageInterface):
             url = self._build_url(await file_code.get_file_path())
             content_length = file_code.size  # 默认使用数据库中的大小
             
+            # 创建ClientSession并复用（包含认证头）
+            session = aiohttp.ClientSession(headers={
+                "Authorization": f"Basic {base64.b64encode(f'{settings.webdav_username}:{settings.webdav_password}'.encode()).decode()}"
+            })
+            
             # 尝试发送HEAD请求获取Content-Length
             try:
-                async with aiohttp.ClientSession(headers={
-                    "Authorization": f"Basic {base64.b64encode(f'{settings.webdav_username}:{settings.webdav_password}'.encode()).decode()}"
-                }) as session:
-                    async with session.head(url) as resp:
-                        if resp.status == 200 and 'Content-Length' in resp.headers:
-                            content_length = int(resp.headers['Content-Length'])
+                async with session.head(url) as resp:
+                    if resp.status == 200 and 'Content-Length' in resp.headers:
+                        content_length = int(resp.headers['Content-Length'])
             except Exception:
                 # 如果HEAD请求失败，继续使用默认大小
                 pass
             
             async def stream_generator():
-                async with aiohttp.ClientSession(headers={
-                    "Authorization": f"Basic {base64.b64encode(f'{settings.webdav_username}:{settings.webdav_password}'.encode()).decode()}"
-                }) as session:
+                try:
                     async with session.get(url) as resp:
                         if resp.status != 200:
                             raise HTTPException(
@@ -1108,6 +1117,8 @@ class WebDAVFileStorage(FileStorageInterface):
                             if not chunk:
                                 break
                             yield chunk
+                finally:
+                    await session.close()
             
             encoded_filename = quote(filename, safe='')
             headers = {