perf: 不再缓存大量不常用图片

MistEO · Aug 12, 2023 · fbf733e · fbf733e
1 parent 1676317
commit fbf733e
Show file tree

Hide file tree

Showing 4 changed files with 47 additions and 31 deletions.
diff --git a/src/common/config/__init__.py b/src/common/config/__init__.py
@@ -15,7 +15,7 @@ class Config(ABC):
     @classmethod
     def _get_config_mongo(cls) -> Collection:
         if cls._config_mongo is None:
-            mongo_client = pymongo.MongoClient('127.0.0.1', 27017, w=0)
+            mongo_client = pymongo.MongoClient('127.0.0.1', 27017)
             mongo_db = mongo_client['PallasBot']
             cls._config_mongo = mongo_db[cls._table]
             cls._config_mongo.create_index(name='{}_index'.format(cls._key),

diff --git a/src/common/utils/media_cache/__init__.py b/src/common/utils/media_cache/__init__.py
@@ -2,10 +2,10 @@
 import httpx
 import base64
 import re
-from datetime import datetime
+from datetime import datetime, timedelta
 from typing import Optional
 
-mongo_client = pymongo.MongoClient('127.0.0.1', 27017, w=0)
+mongo_client = pymongo.MongoClient('127.0.0.1', 27017)
 mongo_db = mongo_client['PallasBot']
 
 image_cache = mongo_db['image_cache']
@@ -16,43 +16,59 @@
 async def insert_image(image_seg):
     cq_code = re.sub(r"\.image,.+?\]", ".image]", str(image_seg))
 
+    db_filter = {'cq_code': cq_code}
+
     idate = int(str(datetime.now().date()).replace('-', ''))
-    cache = image_cache.find_one({'cq_code': cq_code})
+    db_update = {
+        '$inc': {'ref_times': 1},
+        '$set': {'date': idate},
+    }
+
+    cache = image_cache.find_one(db_filter)
+
+    ref_times = 0
     if cache:
-        image_cache.update_one({'cq_code': cq_code}, {
-            '$inc': {'ref_times': 1},
-            '$set': {'date': idate},
-        })
-        return
-
-    url = image_seg.data["url"]
-    async with httpx.AsyncClient() as client:
-        rsp = await client.get(url)
-
-    if rsp.status_code != 200:
-        return
-
-    base64_data = base64.b64encode(rsp.content)
-    base64_data = base64_data.decode()
-    image_cache.update_one({'cq_code': cq_code},
-                           {'$set': {
-                               'cq_code': cq_code,
-                               'base64_data': base64_data,
-                               'ref_times': 1,
-                               'date': idate
-                           }},
-                           upsert=True)
+        if "ref_times" in cache:
+            ref_times = cache["ref_times"]
+        else:
+            ref_times = 1
+
+    # 不是经常收到的图不缓存，不然会占用大量空间
+    if ref_times > 2 and 'base64_data' not in cache:
+        url = image_seg.data["url"]
+        async with httpx.AsyncClient() as client:
+            rsp = await client.get(url)
+
+        if rsp.status_code != 200:
+            return
+
+        base64_data = base64.b64encode(rsp.content)
+        base64_data = base64_data.decode()
+
+        db_update['$set']['base64_data'] = base64_data
+
+    image_cache.update_one(db_filter, db_update, upsert=True)
 
 
 def get_image(cq_code) -> Optional[bytes]:
     cache = image_cache.find_one({'cq_code': cq_code})
     if not cache:
         return None
 
+    if 'base64_data' not in cache:
+        return None
+
     base64_data = cache['base64_data']
     return base64.b64decode(base64_data)
 
 
 def clear_image_cache(days: int = 5, times: int = 3):
-    idate = int(str(datetime.datetime.now() - datetime.timedelta(days=days)).replace('-', ''))
-    image_cache.delete_many({'date': {'$lt': idate}, 'ref_times': {'$lt': times}})
+    idate = int(
+        str((datetime.now() - timedelta(days=days)).date()).replace('-', ''))
+    image_cache.delete_many({'ref_times': {"$exists": False}})
+    image_cache.delete_many(
+        {'date': {'$lt': idate}, 'ref_times': {'$lt': times}})
+
+
+if __name__ == '__main__':
+    clear_image_cache()
diff --git a/src/plugins/repeater/model.py b/src/plugins/repeater/model.py
@@ -28,7 +28,7 @@
     print('TTS not available, error:', error)
     TTS_AVAIABLE = False
 
-mongo_client = pymongo.MongoClient('127.0.0.1', 27017, w=0,
+mongo_client = pymongo.MongoClient('127.0.0.1', 27017,
         unicode_decode_error_handler='ignore')
 
 mongo_db = mongo_client['PallasBot']

diff --git a/tools/clear_old_image.py b/tools/clear_old_image.py
@@ -1,7 +1,7 @@
 import pymongo
 import time
 
-mongo_client = pymongo.MongoClient('127.0.0.1', 27017, w=0,
+mongo_client = pymongo.MongoClient('127.0.0.1', 27017,
         unicode_decode_error_handler='ignore')
 
 mongo_db = mongo_client['PallasBot']