feat: add instance id for scale

Signed-off-by: binaryYuki <60097976+binaryYuki@users.noreply.github.com>

_db.py

@@ -173,8 +173,12 @@ async def init_db():
             logging.info("重建数据库表, 原因: %s", str(e))
             # # 删除所有表
             # await conn.run_sync(Base.metadata.drop_all)
-            # # 创建所有表
-            await conn.run_sync(Base.metadata.create_all)
+            if os.environ.get("REBUILD_DB", False):
+                await conn.run_sync(Base.metadata.drop_all)
+                await conn.run_sync(Base.metadata.create_all)
+            else:
+                # 创建所有表
+                await conn.run_sync(Base.metadata.create_all)
         except Exception as e:
             raise RuntimeError(f"Database initialization failed: {str(e)}")
 

_redis.py

@@ -51,7 +51,6 @@ async def get_keys_by_pattern(pattern: str) -> list:
             return []
 
 
-
 # Set a key-value pair in Redis
 async def set_key(key: str, value: str, ex: Optional[int] = None) -> bool:
     """
@@ -107,3 +106,38 @@ async def key_exists(key: str) -> bool:
     except redis.RedisError as e:
         print(f"Error checking key in Redis: {e}")
         return False
+
+
+async def sadd_key(key: str, value: str) -> bool:
+    """
+    Add a value to a Redis set.
+    """
+    try:
+        await redis_client.sadd(key, value)
+        return True
+    except redis.RedisError as e:
+        print(f"Error adding value to Redis set: {e}")
+        return False
+
+
+async def srem_key(key: str, value: str) -> bool:
+    """
+    Remove a value from a Redis set.
+    """
+    try:
+        await redis_client.srem(key, value)
+        return True
+    except redis.RedisError as e:
+        print(f"Error removing value from Redis set: {e}")
+        return False
+
+
+async def scard_key(key: str) -> int:
+    """
+    Get the number of elements in a Redis set.
+    """
+    try:
+        return await redis_client.scard(key)
+    except redis.RedisError as e:
+        print(f"Error getting set cardinality from Redis: {e}")
+        return 0

app.py

@@ -4,6 +4,7 @@
 import os
 import random
 import subprocess
+import uuid
 from contextlib import asynccontextmanager
 
 import httpx
@@ -33,6 +34,8 @@
 logging.basicConfig(level=logging.getLevelName(loglevel))
 logger = logging.getLogger(__name__)
 
+INSTANCE_ID = str(uuid.uuid4())  # 为每个实例生成唯一ID
+
 
 @repeat_every(seconds=60 * 60, wait_first=True)
 async def testPushServer():
@@ -42,6 +45,18 @@ async def testPushServer():
             await redis_set_key("server_status", "running")
 
 
+async def initInstance():
+    """
+    初始化实例
+    """
+    # 初始化实例 从 redis 里面拿到所有的实例 例如 instance:1 instance:2 所以先获取所有的实例 再+1
+    instance_list = await redis_client.get("instance_list")
+    instance_list = list(instance_list) if instance_list else []
+    instance_list.append(INSTANCE_ID)
+    await redis_client.set("instance_list", str(instance_list))
+
+
+
 @asynccontextmanager
 async def lifespan(_: FastAPI):
     """
@@ -63,6 +78,7 @@ async def lifespan(_: FastAPI):
         await init_db()
         logger.info("MySQL connection established")
     await init_crypto()
+    await initInstance()
     await testPushServer()
     await pushTaskExecQueue()
     yield
@@ -93,6 +109,7 @@ async def index():
         "author": "binaryYuki <noreply.tzpro.xyz>",
         "arch": subprocess.run(['uname', '-m'], stdout=subprocess.PIPE).stdout.decode().strip(),
         "commit": os.getenv("COMMIT_ID", ""),
+        "instance_id": INSTANCE_ID,
     }
 
     # 将字典转换为 JSON 字符串并格式化
@@ -135,6 +152,7 @@ async def healthz():
                                      "redis_hint": redisHint if not redisStatus else "",
                                      "mysql_hint": mysqlHint if not mysqlStatus else ""})
 
+
 secret_key = os.environ.get("SESSION_SECRET")
 if not secret_key:
     secret_key = binascii.hexlify(random.randbytes(16)).decode('utf-8')