redis 连接池

yyetsweb/databases/__init__.py

@@ -9,7 +9,9 @@ import os
 import pathlib
 import sys
 
+import fakeredis
 import pymongo
+import redis
 
 DOUBAN_SEARCH = "https://www.douban.com/search?cat=1002&q={}"
 DOUBAN_DETAIL = "https://movie.douban.com/subject/{}/"
@@ -24,7 +26,7 @@ logging.info(
     (BD2020, XL720, NewzmzOnline, ZhuixinfanOnline, ZimuxiaOnline),
 )
 
-client = pymongo.MongoClient(
+mongo_client = pymongo.MongoClient(
     host=os.getenv("MONGO", "localhost"),
     connect=True,
     connectTimeoutMS=5000,
@@ -33,4 +35,15 @@ client = pymongo.MongoClient(
     minPoolSize=50,
     maxIdleTimeMS=600000,
 )
-db = client["zimuzu"]
+db = mongo_client["zimuzu"]
+
+try:
+    redis_client = redis.StrictRedis(
+        host=os.getenv("REDIS", "localhost"),
+        decode_responses=True,
+        max_connections=100,
+    )
+    redis_client.ping()
+except redis.exceptions.ConnectionError:
+    logging.warning("%s Using fakeredis now... %s", "#" * 10, "#" * 10)
+    redis_client = fakeredis.FakeStrictRedis()

yyetsweb/databases/base.py

@@ -5,18 +5,15 @@ import logging
 import os
 import time
 
-import fakeredis
 import meilisearch
-import redis
 
-from databases import db
-
-faker_redis = fakeredis.FakeStrictRedis()
+from databases import db, redis_client
 
 
 class Mongo:
     def __init__(self):
         self.db = db
+        super().__init__()
 
     def is_admin(self, username: str) -> bool:
         data = self.db["users"].find_one({"username": username, "group": {"$in": ["admin"]}})
@@ -34,17 +31,9 @@ class Mongo:
 
 class Redis:
     def __init__(self):
-        try:
-            self.r = redis.StrictRedis(host=os.getenv("REDIS", "localhost"), decode_responses=True)
-            self.r.ping()
-        except redis.exceptions.ConnectionError:
-            logging.warning("%s Using fakeredis now... %s", "#" * 10, "#" * 10)
-            self.r = faker_redis
+        self.r = redis_client
         super().__init__()
 
-    def __del__(self):
-        self.r.close()
-
     @classmethod
     def cache(cls, timeout: int):
         def func(fun):

yyetsweb/databases/other.py

@@ -137,7 +137,7 @@ class SpamProcess(Mongo):
         logging.info("Telegram response: %s", resp)
 
 
-class Other(Mongo):
+class Other(Mongo, Redis):
     def reset_top(self):
         # before resetting, save top data to history
         json_data = requests.get("http://127.0.0.1:8888/api/top").json()
@@ -159,12 +159,10 @@ class Other(Mongo):
 
     def import_ban_user(self):
         usernames = self.db["users"].find({"status.disable": True}, projection={"username": True})
-        r = Redis().r
-        r.delete("user_blacklist")
+        self.r.delete("user_blacklist")
         logging.info("Importing ban users to redis...%s", usernames)
         for username in [u["username"] for u in usernames]:
-            r.hset("user_blacklist", username, 100)
-        r.close()
+            self.r.hset("user_blacklist", username, 100)
 
     def fill_user_hash(self):
         users = self.db["users"].find({"hash": {"$exists": False}}, projection={"username": True})

yyetsweb/databases/resources.py

@@ -240,12 +240,10 @@ class Top(Mongo):
         return all_data
 
 
-class ResourceLatest(Mongo):
-    @staticmethod
-    def get_latest_resource() -> dict:
-        redis = Redis().r
+class ResourceLatest(Mongo, Redis):
+    def get_latest_resource(self) -> dict:
         key = "latest-resource"
-        latest = redis.get(key)
+        latest = self.r.get(key)
         if latest:
             logging.info("Cache hit for latest resource")
             latest = json.loads(latest)
@@ -253,7 +251,7 @@ class ResourceLatest(Mongo):
         else:
             logging.warning("Cache miss for latest resource")
             latest = ResourceLatest().query_db()
-            redis.set(key, json.dumps(latest, ensure_ascii=False))
+            self.r.set(key, json.dumps(latest, ensure_ascii=False))
         return latest
 
     def query_db(self) -> dict:
@@ -286,10 +284,9 @@ class ResourceLatest(Mongo):
         return dict(data=ok)
 
     def refresh_latest_resource(self):
-        redis = Redis().r
         logging.info("Getting new resources...")
         latest = self.query_db()
-        redis.set("latest-resource", json.dumps(latest, ensure_ascii=False))
+        self.r.set("latest-resource", json.dumps(latest, ensure_ascii=False))
         logging.info("latest-resource data refreshed.")
 
 

yyetsweb/databases/user.py

@@ -129,7 +129,6 @@ class User(Mongo, Redis):
         )
 
     def update_user_info(self, username: str, data: dict) -> dict:
-        redis = Redis().r
         valid_fields = ["email"]
         valid_data = {}
         for field in valid_fields:
@@ -147,7 +146,7 @@ class User(Mongo, Redis):
             # rate limit
             user_email = valid_data.get("email")
             timeout_key = f"timeout-{username}"
-            if redis.get(timeout_key):
+            if self.r.get(timeout_key):
                 return {
                     "status_code": HTTPStatus.TOO_MANY_REQUESTS,
                     "status": False,
@@ -162,9 +161,9 @@ class User(Mongo, Redis):
             context = {"username": username, "text": text}
             send_mail(user_email, subject, context)
             # 发送成功才设置缓存
-            redis.set(timeout_key, username, ex=1800)
-            redis.hset(user_email, mapping={"code": verify_code, "wrong": 0})
-            redis.expire(user_email, 24 * 3600)
+            self.r.set(timeout_key, username, ex=1800)
+            self.r.hset(user_email, mapping={"code": verify_code, "wrong": 0})
+            self.r.expire(user_email, 24 * 3600)
 
         self.db["users"].update_one({"username": username}, {"$set": valid_data})
         return {
@@ -199,11 +198,10 @@ class UserAvatar(User, Mongo):
             return {"image": None, "content_type": None}
 
 
-class UserEmail(Mongo):
+class UserEmail(Mongo, Redis):
     def verify_email(self, username, code):
-        r = Redis().r
         email = self.db["users"].find_one({"username": username})["email"]["address"]
-        verify_data = r.hgetall(email)
+        verify_data = self.r.hgetall(email)
         wrong_count = int(verify_data["wrong"])
         MAX = 10
         if wrong_count >= MAX:
@@ -219,8 +217,8 @@ class UserEmail(Mongo):
         correct_code = verify_data["code"]
 
         if correct_code == code:
-            r.expire(email, 0)
-            r.expire(f"timeout-{email}", 0)
+            self.r.expire(email, 0)
+            self.r.expire(f"timeout-{email}", 0)
             self.db["users"].update_one({"username": username}, {"$set": {"email.verified": True}})
             return {
                 "status": True,
@@ -228,7 +226,7 @@ class UserEmail(Mongo):
                 "message": "邮箱已经验证成功",
             }
         else:
-            r.hset(email, "wrong", wrong_count + 1)
+            self.r.hset(email, "wrong", wrong_count + 1)
             return {
                 "status": False,
                 "status_code": HTTPStatus.FORBIDDEN,

yyetsweb/handlers/base.py

@@ -12,7 +12,7 @@ from pathlib import Path
 from tornado import gen, web
 from tornado.concurrent import run_on_executor
 
-from databases.base import Redis
+from databases.base import redis_client
 from handlers import cf
 
 index = pathlib.Path(__file__).parent.parent.joinpath("templates", "index.html").as_posix()
@@ -34,7 +34,7 @@ class BaseHandler(web.RequestHandler):
         module = importlib.import_module(f"databases.{self.filename}")
 
         self.instance = getattr(module, class_name, lambda: 1)()
-        self.r = Redis().r
+        self.r = redis_client
 
     def add_tauri(self):
         origin = self.request.headers.get("origin", "")