初始化提交

app/controllers/manager/base_manager.py

@@ -0,0 +1,87 @@
+import threading
+from typing import Any, Callable, Dict
+
+from loguru import logger
+
+
+class TaskQueueFullError(ValueError):
+    pass
+
+
+class TaskManager:
+    def __init__(self, max_concurrent_tasks: int, max_queued_tasks: int = 100):
+        self.max_concurrent_tasks = max_concurrent_tasks
+        self.max_queued_tasks = max_queued_tasks
+        self.current_tasks = 0
+        self.lock = threading.Lock()
+        self.queue = self.create_queue()
+
+    def create_queue(self):
+        raise NotImplementedError()
+
+    def add_task(self, func: Callable, *args: Any, **kwargs: Any):
+        with self.lock:
+            if self.current_tasks < self.max_concurrent_tasks:
+                logger.info(
+                    f"add task: {func.__name__}, current_tasks: {self.current_tasks}"
+                )
+                self.execute_task(func, *args, **kwargs)
+            else:
+                queue_size = self.queue_size()
+                # 并发数已满时才进入排队。队列必须有上限，否则匿名接口可以持续
+                # 堆积任务对象和请求参数，最终造成内存耗尽或第三方 API 成本失控。
+                if queue_size >= self.max_queued_tasks:
+                    logger.warning(
+                        f"reject task: {func.__name__}, queue_size: {queue_size}, "
+                        f"max_queued_tasks: {self.max_queued_tasks}"
+                    )
+                    raise TaskQueueFullError("task queue is full, please try again later")
+
+                logger.info(
+                    f"enqueue task: {func.__name__}, current_tasks: {self.current_tasks}, "
+                    f"queue_size: {queue_size}"
+                )
+                self.enqueue({"func": func, "args": args, "kwargs": kwargs})
+
+    def execute_task(self, func: Callable, *args: Any, **kwargs: Any):
+        thread = threading.Thread(
+            target=self.run_task, args=(func, *args), kwargs=kwargs
+        )
+        thread.start()
+
+    def run_task(self, func: Callable, *args: Any, **kwargs: Any):
+        try:
+            with self.lock:
+                self.current_tasks += 1
+            func(*args, **kwargs)  # call the function here, passing *args and **kwargs.
+        finally:
+            self.task_done()
+
+    def check_queue(self):
+        with self.lock:
+            if (
+                self.current_tasks < self.max_concurrent_tasks
+                and not self.is_queue_empty()
+            ):
+                task_info = self.dequeue()
+                func = task_info["func"]
+                args = task_info.get("args", ())
+                kwargs = task_info.get("kwargs", {})
+                self.execute_task(func, *args, **kwargs)
+
+    def task_done(self):
+        with self.lock:
+            self.current_tasks -= 1
+        self.check_queue()
+
+    def enqueue(self, task: Dict):
+        raise NotImplementedError()
+
+    def dequeue(self):
+        raise NotImplementedError()
+
+    def is_queue_empty(self):
+        raise NotImplementedError()
+
+    def queue_size(self):
+        raise NotImplementedError()

app/controllers/manager/memory_manager.py

@@ -0,0 +1,21 @@
+from queue import Queue
+from typing import Dict
+
+from app.controllers.manager.base_manager import TaskManager
+
+
+class InMemoryTaskManager(TaskManager):
+    def create_queue(self):
+        return Queue(maxsize=self.max_queued_tasks)
+
+    def enqueue(self, task: Dict):
+        self.queue.put(task)
+
+    def dequeue(self):
+        return self.queue.get()
+
+    def is_queue_empty(self):
+        return self.queue.empty()
+
+    def queue_size(self):
+        return self.queue.qsize()

app/controllers/manager/redis_manager.py

@@ -0,0 +1,64 @@
+import json
+from typing import Dict
+
+import redis
+
+from app.controllers.manager.base_manager import TaskManager
+from app.models.schema import VideoParams
+from app.services import task as tm
+
+FUNC_MAP = {
+    "start": tm.start,
+    # 'start_test': tm.start_test
+}
+
+
+class RedisTaskManager(TaskManager):
+    def __init__(
+        self,
+        max_concurrent_tasks: int,
+        redis_url: str,
+        max_queued_tasks: int = 100,
+    ):
+        self.redis_client = redis.Redis.from_url(redis_url)
+        super().__init__(max_concurrent_tasks, max_queued_tasks=max_queued_tasks)
+
+    def create_queue(self):
+        return "task_queue"
+
+    def enqueue(self, task: Dict):
+        task_with_serializable_params = task.copy()
+
+        if "params" in task["kwargs"] and isinstance(
+            task["kwargs"]["params"], VideoParams
+        ):
+            task_with_serializable_params["kwargs"]["params"] = task["kwargs"][
+                "params"
+            ].dict()
+
+        # 将函数对象转换为其名称
+        task_with_serializable_params["func"] = task["func"].__name__
+        self.redis_client.rpush(self.queue, json.dumps(task_with_serializable_params))
+
+    def dequeue(self):
+        task_json = self.redis_client.lpop(self.queue)
+        if task_json:
+            task_info = json.loads(task_json)
+            # 将函数名称转换回函数对象
+            task_info["func"] = FUNC_MAP[task_info["func"]]
+
+            if "params" in task_info["kwargs"] and isinstance(
+                task_info["kwargs"]["params"], dict
+            ):
+                task_info["kwargs"]["params"] = VideoParams(
+                    **task_info["kwargs"]["params"]
+                )
+
+            return task_info
+        return None
+
+    def is_queue_empty(self):
+        return self.redis_client.llen(self.queue) == 0
+
+    def queue_size(self):
+        return self.redis_client.llen(self.queue)