'commit'

DouYin/dy-downloader/control/__init__.py

@@ -0,0 +1,5 @@
+from .rate_limiter import RateLimiter
+from .retry_handler import RetryHandler
+from .queue_manager import QueueManager
+
+__all__ = ['RateLimiter', 'RetryHandler', 'QueueManager']

DouYin/dy-downloader/control/queue_manager.py

@@ -0,0 +1,37 @@
+import asyncio
+from typing import List, Callable, Any, TypeVar
+from utils.logger import setup_logger
+
+logger = setup_logger('QueueManager')
+
+T = TypeVar('T')
+
+
+class QueueManager:
+    def __init__(self, max_workers: int = 5):
+        self.max_workers = max_workers
+        self.semaphore = asyncio.Semaphore(max_workers)
+
+    async def process_tasks(self, tasks: List[Callable], *args, **kwargs) -> List[Any]:
+        async def _task_wrapper(task):
+            async with self.semaphore:
+                try:
+                    return await task(*args, **kwargs)
+                except Exception as e:
+                    logger.error(f"Task failed: {e}")
+                    return None
+
+        results = await asyncio.gather(*[_task_wrapper(task) for task in tasks], return_exceptions=True)
+        return results
+
+    async def download_batch(self, download_func: Callable, items: List[Any]) -> List[Any]:
+        async def _download_wrapper(item):
+            async with self.semaphore:
+                try:
+                    return await download_func(item)
+                except Exception as e:
+                    logger.error(f"Download failed for item: {e}")
+                    return {'status': 'error', 'error': str(e), 'item': item}
+
+        results = await asyncio.gather(*[_download_wrapper(item) for item in items], return_exceptions=False)
+        return results

DouYin/dy-downloader/control/rate_limiter.py

@@ -0,0 +1,21 @@
+import asyncio
+import time
+
+
+class RateLimiter:
+    def __init__(self, max_per_second: float = 2):
+        self.max_per_second = max_per_second
+        self.min_interval = 1.0 / max_per_second
+        self.last_request = 0
+        self._lock = asyncio.Lock()
+
+    async def acquire(self):
+        async with self._lock:
+            current = time.time()
+            time_since_last = current - self.last_request
+
+            if time_since_last < self.min_interval:
+                wait_time = self.min_interval - time_since_last
+                await asyncio.sleep(wait_time)
+
+            self.last_request = time.time()

DouYin/dy-downloader/control/retry_handler.py

@@ -0,0 +1,29 @@
+import asyncio
+from typing import Callable, Any, TypeVar
+from utils.logger import setup_logger
+
+logger = setup_logger('RetryHandler')
+
+T = TypeVar('T')
+
+
+class RetryHandler:
+    def __init__(self, max_retries: int = 3):
+        self.max_retries = max_retries
+        self.retry_delays = [1, 2, 5]
+
+    async def execute_with_retry(self, func: Callable[..., T], *args, **kwargs) -> T:
+        last_error = None
+
+        for attempt in range(self.max_retries):
+            try:
+                return await func(*args, **kwargs)
+            except Exception as e:
+                last_error = e
+                if attempt < self.max_retries - 1:
+                    delay = self.retry_delays[min(attempt, len(self.retry_delays) - 1)]
+                    logger.warning(f"Attempt {attempt + 1} failed: {e}, retrying in {delay}s...")
+                    await asyncio.sleep(delay)
+
+        logger.error(f"All {self.max_retries} attempts failed: {last_error}")
+        raise last_error