| 1 | import ast |
| 2 | import copy |
| 3 | import threading |
| 4 | from abc import ABC, abstractmethod |
| 5 | |
| 6 | from app.config import config |
| 7 | from app.models import const |
| 8 | |
| 9 | |
| 10 | # Base class for state management |
| 11 | class BaseState(ABC): |
| 12 | @abstractmethod |
| 13 | def update_task(self, task_id: str, state: int, progress: int = 0, **kwargs): |
| 14 | pass |
| 15 | |
| 16 | @abstractmethod |
| 17 | def get_task(self, task_id: str): |
| 18 | pass |
| 19 | |
| 20 | @abstractmethod |
| 21 | def get_all_tasks(self, page: int, page_size: int): |
| 22 | pass |
| 23 | |
| 24 | |
| 25 | # Memory state management |
| 26 | class MemoryState(BaseState): |
| 27 | def __init__(self): |
| 28 | self._tasks = {} |
| 29 | self._lock = threading.RLock() |
| 30 | |
| 31 | def get_all_tasks(self, page: int, page_size: int): |
| 32 | start = (page - 1) * page_size |
| 33 | end = start + page_size |
| 34 | with self._lock: |
| 35 | tasks = [copy.deepcopy(task) for task in self._tasks.values()] |
| 36 | total = len(tasks) |
| 37 | return tasks[start:end], total |
| 38 | |
| 39 | def update_task( |
| 40 | self, |
| 41 | task_id: str, |
| 42 | state: int = const.TASK_STATE_PROCESSING, |
| 43 | progress: int = 0, |
| 44 | **kwargs, |
| 45 | ): |
| 46 | progress = int(progress) |
| 47 | if progress > 100: |
| 48 | progress = 100 |
| 49 | |
| 50 | with self._lock: |
| 51 | self._tasks[task_id] = { |
| 52 | "task_id": task_id, |
| 53 | "state": state, |
| 54 | "progress": progress, |
| 55 | **kwargs, |
| 56 | } |
| 57 | |
| 58 | def get_task(self, task_id: str): |
| 59 | with self._lock: |
| 60 | task = self._tasks.get(task_id, None) |
| 61 | return copy.deepcopy(task) if task is not None else None |
| 62 | |
| 63 | def delete_task(self, task_id: str): |
| 64 | with self._lock: |
| 65 | self._tasks.pop(task_id, None) |
| 66 | |
| 67 | |
| 68 | # Redis state management |
| 69 | class RedisState(BaseState): |
| 70 | """ |
| 71 | Redis-backed task state. |
| 72 | |
| 73 | Trust boundary: Redis is expected to be private to this application. Task |
| 74 | values are written by MoneyPrinterTurbo and converted back from strings for |
| 75 | compatibility with existing state records. Do not expose this Redis database |
| 76 | to untrusted writers without replacing deserialization with a stricter |
| 77 | schema-based format. |
| 78 | """ |
| 79 | |
| 80 | def __init__(self, host="localhost", port=6379, db=0, password=None): |
| 81 | import redis |
| 82 | |
| 83 | self._redis = redis.StrictRedis(host=host, port=port, db=db, password=password) |
| 84 | |
| 85 | def get_all_tasks(self, page: int, page_size: int): |
| 86 | start = (page - 1) * page_size |
| 87 | end = start + page_size |
| 88 | tasks = [] |
| 89 | cursor = 0 |
| 90 | total = 0 |
| 91 | while True: |
| 92 | cursor, keys = self._redis.scan(cursor, count=page_size) |
| 93 | batch_start = total |
| 94 | batch_size = len(keys) |
| 95 | total += batch_size |
| 96 | |
| 97 | # Redis SCAN 是分批返回 key。分页切片必须基于“当前批次起始索引” |
| 98 | # 计算，而不能用累积后的 total 反推，否则第一页会切到空数组， |
| 99 | # 第二页也可能只返回部分数据。 |
| 100 | if batch_start < end and total > start: |
| 101 | slice_start = max(0, start - batch_start) |
| 102 | slice_end = min(batch_size, end - batch_start) |
| 103 | for key in keys[slice_start:slice_end]: |
| 104 | task_data = self._redis.hgetall(key) |
| 105 | task = { |
| 106 | k.decode("utf-8"): self._convert_to_original_type(v) |
| 107 | for k, v in task_data.items() |
| 108 | } |
| 109 | tasks.append(task) |
| 110 | |
| 111 | # 即使当前页已经取满，也要继续 SCAN 到 cursor=0， |
| 112 | # 因为调用方需要准确 total 来渲染分页信息。 |
| 113 | if cursor == 0: |
| 114 | break |
| 115 | return tasks, total |
| 116 | |
| 117 | def update_task( |
| 118 | self, |
| 119 | task_id: str, |
| 120 | state: int = const.TASK_STATE_PROCESSING, |
| 121 | progress: int = 0, |
| 122 | **kwargs, |
| 123 | ): |
| 124 | progress = int(progress) |
| 125 | if progress > 100: |
| 126 | progress = 100 |
| 127 | |
| 128 | fields = { |
| 129 | "task_id": task_id, |
| 130 | "state": state, |
| 131 | "progress": progress, |
| 132 | **kwargs, |
| 133 | } |
| 134 | |
| 135 | for field, value in fields.items(): |
| 136 | self._redis.hset(task_id, field, str(value)) |
| 137 | |
| 138 | def get_task(self, task_id: str): |
| 139 | task_data = self._redis.hgetall(task_id) |
| 140 | if not task_data: |
| 141 | return None |
| 142 | |
| 143 | task = { |
| 144 | key.decode("utf-8"): self._convert_to_original_type(value) |
| 145 | for key, value in task_data.items() |
| 146 | } |
| 147 | return task |
| 148 | |
| 149 | def delete_task(self, task_id: str): |
| 150 | self._redis.delete(task_id) |
| 151 | |
| 152 | @staticmethod |
| 153 | def _convert_to_original_type(value): |
| 154 | """ |
| 155 | Convert values written by this application back to common Python types. |
| 156 | |
| 157 | This compatibility parser assumes Redis is inside the application's |
| 158 | trust boundary. If Redis can be written by untrusted clients, task state |
| 159 | should move to a strict JSON/schema parser instead of open-ended literal |
| 160 | conversion. |
| 161 | """ |
| 162 | value_str = value.decode("utf-8") |
| 163 | |
| 164 | try: |
| 165 | # try to convert byte string array to list |
| 166 | return ast.literal_eval(value_str) |
| 167 | except (ValueError, SyntaxError): |
| 168 | pass |
| 169 | |
| 170 | if value_str.isdigit(): |
| 171 | return int(value_str) |
| 172 | # Add more conversions here if needed |
| 173 | return value_str |
| 174 | |
| 175 | |
| 176 | # Global state |
| 177 | _enable_redis = config.app.get("enable_redis", False) |
| 178 | _redis_host = config.app.get("redis_host", "localhost") |
| 179 | _redis_port = config.app.get("redis_port", 6379) |
| 180 | _redis_db = config.app.get("redis_db", 0) |
| 181 | _redis_password = config.app.get("redis_password", None) |
| 182 | |
| 183 | state = ( |
| 184 | RedisState( |
| 185 | host=_redis_host, port=_redis_port, db=_redis_db, password=_redis_password |
| 186 | ) |
| 187 | if _enable_redis |
| 188 | else MemoryState() |
| 189 | ) |
| 190 |