fix: 避免异常崩溃产生多个线程池

Author: z0z0r4

mcim_sync/tasks/__init__.py

@@ -1,11 +1,19 @@
 from concurrent.futures import ThreadPoolExecutor
+from contextlib import contextmanager
 
-
+@contextmanager
 def create_tasks_pool(sync_function, data, max_workers, thread_name_prefix):
+    # 创建线程池
     thread_pool = ThreadPoolExecutor(
-        max_workers=max_workers, thread_name_prefix=thread_name_prefix
+        max_workers=max_workers,
+        thread_name_prefix=thread_name_prefix
     )
-    futures = [
-        thread_pool.submit(sync_function, item) for item in data
-    ]
-    return thread_pool, futures
+    
+    try:
+        # 提交所有任务
+        futures = [
+            thread_pool.submit(sync_function, item) for item in data
+        ]
+        yield futures
+    finally:
+        thread_pool.shutdown(wait=True)

mcim_sync/tasks/curseforge.py

@@ -45,27 +45,23 @@ def refresh_curseforge_with_modify_date() -> bool:
     log.info(f"Curseforge expired data fetched: {len(curseforge_expired_modids)}")
     log.info("Start syncing CurseForge expired data...")
 
-    curseforge_pool, curseforge_futures = create_tasks_pool(
-        sync_mod,  # 需要 ProjectDetail 返回值
+    with create_tasks_pool(
+        sync_mod,
         curseforge_expired_modids,
         MAX_WORKERS,
         "refresh_curseforge",
-    )
-    projects_detail_info: List[ProjectDetail] = []
-    for future in as_completed(curseforge_futures):
-        result = future.result()
-        if result:
-            projects_detail_info.append(result)
-    else:
-        curseforge_pool.shutdown()
+    ) as curseforge_futures:
+        projects_detail_info: List[ProjectDetail] = []
+        for future in as_completed(curseforge_futures):
+            result = future.result()
+            if result:
+                projects_detail_info.append(result)
 
     success_modids = [project.id for project in projects_detail_info if project]
-    
+
     failed_count = len(curseforge_expired_modids) - len(success_modids)
     failed_modids = [
-        modid
-        for modid in curseforge_expired_modids
-        if modid not in success_modids
+        modid for modid in curseforge_expired_modids if modid not in success_modids
     ]
 
     log.info(
@@ -112,18 +108,15 @@ def sync_curseforge_queue() -> bool:
     log.info(f"New modids: {new_modids}, count: {len(new_modids)}")
 
     if new_modids:
-        # pool, futures = create_tasks_pool(sync_mod, modids, MAX_WORKERS, "curseforge")
-        pool, futures = create_tasks_pool(
+        with create_tasks_pool(
             sync_mod, new_modids, MAX_WORKERS, "sync_curseforge_queue"
-        )  # https://github.com/mcmod-info-mirror/mcim-sync/issues/2
-
-        projects_detail_info = []
-        for future in as_completed(futures):
-            result = future.result()
-            if result:
-                projects_detail_info.append(result)
+        ) as futures:
+            projects_detail_info = []
+            for future in as_completed(futures):
+                result = future.result()
+                if result:
+                    projects_detail_info.append(result)
 
-        pool.shutdown()
         log.info(f"CurseForge queue sync finished, total: {len(modids)}")
 
         # clear queue
@@ -199,17 +192,15 @@ def sync_curseforge_by_search(class_ids: Optional[List[int]] = None) -> bool:
 
     log.info(f"CurseForge new modids fetched: {len(new_modids)}")
     if new_modids:
-        pool, futures = create_tasks_pool(
+        with create_tasks_pool(
             sync_mod, new_modids, MAX_WORKERS, "sync_curseforge_by_search"
-        )
-
-        projects_detail_info = []
-        for future in as_completed(futures):
-            result = future.result()
-            if result:
-                projects_detail_info.append(result)
+        ) as futures:
+            projects_detail_info = []
+            for future in as_completed(futures):
+                result = future.result()
+                if result:
+                    projects_detail_info.append(result)
 
-        pool.shutdown()
         log.info(f"CurseForge search sync finished, total: {len(new_modids)}")
 
         if config.telegram_bot:
@@ -231,28 +222,23 @@ def sync_curseforge_full():
     curseforge_data = fetch_all_curseforge_data()
     log.info(f"Curseforge data totally fetched: {len(curseforge_data)}")
 
-    curseforge_pool, curseforge_futures = create_tasks_pool(
+    with create_tasks_pool(
         sync_mod, curseforge_data, MAX_WORKERS, "curseforge_refresh_full"
-    )
-
-    log.info(
-        f"All {len(curseforge_futures)} tasks submitted, waiting for completion..."
-    )
+    ) as curseforge_futures:
+        log.info(
+            f"All {len(curseforge_futures)} tasks submitted, waiting for completion..."
+        )
 
-    projects_detail_info: List[ProjectDetail] = []
-    for future in as_completed(curseforge_futures):
-        result = future.result()
-        if result:
-            projects_detail_info.append(result)
-    else:
-        curseforge_pool.shutdown()
+        projects_detail_info: List[ProjectDetail] = []
+        for future in as_completed(curseforge_futures):
+            result = future.result()
+            if result:
+                projects_detail_info.append(result)
 
     success_modids = [project.id for project in projects_detail_info if project]
 
     failed_count = len(curseforge_data) - len(success_modids)
-    failed_modids = [
-        modid for modid in curseforge_data if modid not in success_modids
-    ]
+    failed_modids = [modid for modid in curseforge_data if modid not in success_modids]
 
     log.info(
         f"CurseForge full sync finished, total: {len(curseforge_data)}, "
@@ -268,5 +254,5 @@ def sync_curseforge_full():
         )
         notification.send_to_telegram()
         log.info("CurseForge refresh message sent to telegram.")
-    
-    return True
+
+    return True

mcim_sync/tasks/modrinth.py

@@ -23,7 +23,10 @@
     check_newest_search_result,
 )
 from mcim_sync.cleaner.modrinth import remove_projects
-from mcim_sync.fetcher.modrinth import fetch_expired_and_removed_modrinth_data, fetch_all_modrinth_data
+from mcim_sync.fetcher.modrinth import (
+    fetch_expired_and_removed_modrinth_data,
+    fetch_all_modrinth_data,
+)
 from mcim_sync.queues.modrinth import clear_modrinth_all_queues
 from mcim_sync.tasks import create_tasks_pool
 
@@ -52,20 +55,20 @@ def refresh_modrinth_with_modify_date() -> bool:
 
     # 刷新过期的 modrinth 数据
     log.info("Start syncing Modrinth expired data...")
-    modrinth_pool, modrinth_futures = create_tasks_pool(
+    with create_tasks_pool(
         sync_project,  # 需要 ProjectDetail 返回值
         modrinth_expired_data,
         MAX_WORKERS,
         "refresh_modrinth",
-    )
-
-    projects_detail_info = []
-    for future in as_completed(modrinth_futures):
-        result = future.result()
-        if result:
-            projects_detail_info.append(result)
-    else:
-        modrinth_pool.shutdown()
+    ) as modrinth_futures:
+        log.info(
+            f"All {len(modrinth_futures)} tasks submitted, waiting for completion..."
+        )
+        projects_detail_info = []
+        for future in as_completed(modrinth_futures):
+            result = future.result()
+            if result:
+                projects_detail_info.append(result)
 
     if config.telegram_bot:
         notification = RefreshNotification(
@@ -85,13 +88,19 @@ def fetch_modrinth_not_found_ids_from_queue():
     project_ids = []
     avaliable_project_ids = check_modrinth_project_ids_available()
     project_ids.extend(avaliable_project_ids)
-    log.info(f"Modrinth project ids queue available project_ids: {len(avaliable_project_ids)}")
+    log.info(
+        f"Modrinth project ids queue available project_ids: {len(avaliable_project_ids)}"
+    )
     avaliable_project_ids = check_modrinth_version_ids_available()
     project_ids.extend(avaliable_project_ids)
-    log.info(f"Modrinth version ids queue available project_ids: {len(avaliable_project_ids)}")
+    log.info(
+        f"Modrinth version ids queue available project_ids: {len(avaliable_project_ids)}"
+    )
     avaliable_project_ids = check_modrinth_hashes_available()
     project_ids.extend(avaliable_project_ids)
-    log.info(f"Modrinth hashes queue available project_ids: {len(avaliable_project_ids)}")
+    log.info(
+        f"Modrinth hashes queue available project_ids: {len(avaliable_project_ids)}"
+    )
     return project_ids
 
 
@@ -107,22 +116,22 @@ def sync_modrinth_queue() -> bool:
     log.info(f"New project ids: {new_project_ids}, count: {len(new_project_ids)}")
 
     if new_project_ids:
-        pool, futures = create_tasks_pool(
-            # sync_project, project_ids, MAX_WORKERS, "modrinth"
-            sync_project,
-            new_project_ids,
-            MAX_WORKERS,
-            "sync_modrinth_by_queue",  # https://github.com/mcmod-info-mirror/mcim-sync/issues/2
-        )
-
-        projects_detail_info = []
+        with (
+            create_tasks_pool(
+                # sync_project, project_ids, MAX_WORKERS, "modrinth"
+                sync_project,
+                new_project_ids,
+                MAX_WORKERS,
+                "sync_modrinth_by_queue",  # https://github.com/mcmod-info-mirror/mcim-sync/issues/2
+            ) as futures
+        ):
+            projects_detail_info = []
+
+            for future in as_completed(futures):
+                result = future.result()
+                if result:
+                    projects_detail_info.append(result)
 
-        for future in as_completed(futures):
-            result = future.result()
-            if result:
-                projects_detail_info.append(result)
-
-        pool.shutdown()
         log.info(f"Modrinth queue sync finished, total: {len(project_ids)}")
 
         # clear queue
@@ -150,20 +159,18 @@ def sync_modrinth_by_search():
     new_project_ids = check_newest_search_result()
     log.info(f"Modrinth project ids fetched: {len(new_project_ids)}")
     if new_project_ids:
-        pool, futures = create_tasks_pool(
+        with create_tasks_pool(
             sync_project,
             new_project_ids,
             MAX_WORKERS,
             "sync_modrinth_by_search",
-        )
+        ) as futures:
+            projects_detail_info = []
+            for future in as_completed(futures):
+                result = future.result()
+                if result:
+                    projects_detail_info.append(result)
 
-        projects_detail_info = []
-        for future in as_completed(futures):
-            result = future.result()
-            if result:
-                projects_detail_info.append(result)
-
-        pool.shutdown()
         log.info(
             f"Modrinth sync new project by search finished, total: {len(new_project_ids)}"
         )
@@ -200,6 +207,7 @@ def refresh_modrinth_tags():
         log.info("All Message sent to telegram.")
     return True
 
+
 def refresh_modrinth_full():
     """
     刷新 modrinth 所有数据
@@ -209,26 +217,25 @@ def refresh_modrinth_full():
     modrinth_data = fetch_all_modrinth_data()
     log.info(f"Modrinth data totally fetched: {len(modrinth_data)}")
 
-    modrinth_pool, modrinth_futures = create_tasks_pool(
+    with create_tasks_pool(
         sync_project, modrinth_data, MAX_WORKERS, "modrinth_refresh_full"
-    )
-
-    log.info(
-        f"All {len(modrinth_futures)} tasks submitted, waiting for completion..."
-    )
+    ) as modrinth_futures:
+        log.info(
+            f"All {len(modrinth_futures)} tasks submitted, waiting for completion..."
+        )
 
-    projects_detail_info = []
-    for future in as_completed(modrinth_futures):
-        result = future.result()
-        if result:
-            projects_detail_info.append(result)
-    else:
-        modrinth_pool.shutdown()
+        projects_detail_info = []
+        for future in as_completed(modrinth_futures):
+            result = future.result()
+            if result:
+                projects_detail_info.append(result)
 
     success_project_ids = [project.id for project in projects_detail_info if project]
-    
+
     failed_count = len(modrinth_data) - len(success_project_ids)
-    failed_project_ids = [project.id for project in modrinth_data if project.id not in success_project_ids]
+    failed_project_ids = [
+        project.id for project in modrinth_data if project.id not in success_project_ids
+    ]
 
     log.info(
         f"Modrinth full sync finished, total: {len(modrinth_data)}, "
@@ -244,5 +251,5 @@ def refresh_modrinth_full():
         )
         notice.send_to_telegram()
         log.info("Modrinth full refresh message sent to telegram.")
-    
-    return True
+
+    return True