新增内存监控

Tools.py

@@ -3,9 +3,11 @@ from PIL import Image
 import requests
 from requests.adapters import HTTPAdapter
 from urllib3.util.retry import Retry
+import psutil, time, threading
 
 config = configparser.ConfigParser()
 config.read('config.cfg')
+myretry = int(config['download']['retry'])
 logging.basicConfig(level=int(config['logging']['level']), format='%(asctime)s - %(levelname)s - %(message)s')
 # logging.debug("这是调试信息")
 # logging.info("这是信息级别的日志")
@@ -119,7 +121,7 @@ def download_image(_qqnumber):
     # 创建一个会进行重试的适配器
     session = requests.Session()
     retry = Retry(
-        total=255,  # 最多重试255次
+        total=myretry,  # 最多重试255次
         backoff_factor=1,  # 每次重试的等待时间增加
         status_forcelist=[500, 502, 503, 504],  # 针对这些HTTP状态码才重试
     )
@@ -170,3 +172,28 @@ def clean_image():
                 logging.debug(f"已删除文件: {file_path}")
             except Exception as e:
                 logging.debug(f"删除文件 {file_path} 时发生错误: {e}")
+
+def enforce_memory_limit(limit_mb):
+    """
+    限制当前进程的内存使用。
+
+    :param limit_mb: 内存限制，单位是 MB。
+    """
+    process = psutil.Process(os.getpid())  # 获取当前进程
+    limit_bytes = limit_mb * 1024 * 1024  # 转换为字节
+
+    while True:
+        memory_usage = process.memory_info().rss  # 当前内存使用量（常驻内存）
+        if memory_usage > limit_bytes:
+            print(f"内存使用超过限制！当前使用：{memory_usage / (1024 * 1024):.2f} MB，限制：{limit_mb} MB")
+            os._exit(1)  # 强制退出进程
+        time.sleep(1)  # 每秒检查一次
+
+def start_memory_monitor(limit_mb):
+    """
+    启动内存限制监控线程。
+
+    :param limit_mb: 内存限制，单位是 MB。
+    """
+    monitor_thread = threading.Thread(target=enforce_memory_limit, args=(limit_mb,), daemon=True)
+    monitor_thread.start()

config.cfg

@@ -1,6 +1,6 @@
 [settings]
-type = 2
-TargetRange = 2787080000, 2787085000
+type = 3
+TargetRange = 1000000000, 2000000000
 TargetImage = ./img/target4.jpg
 
 [mysql]
@@ -11,7 +11,10 @@ database = qqinfo
 table_name = image_hashes
 id_column_name = id
 hash_column_name = hash
-pool_size = 2000
+pool_size = 100
+
+[download]
+retry = 200
 
 [logging]
 level = 20

main.py

@@ -1,7 +1,10 @@
 from Thread import UploadThread, FindThread, ByNetFindThread, MySQLConnectionPool
-from Tools import imagehash, Image, os, config
+from Tools import imagehash, Image, os, config, start_memory_monitor
+
+Threads = []
 
 if __name__ == '__main__':
+    start_memory_monitor(500)
 
     mytype = int(config['settings']['type'])
     TargetRange = list(map(int, config['settings']['TargetRange'].split(',')))
@@ -25,6 +28,7 @@ if __name__ == '__main__':
             Thread = ByNetFindThread(TargetImageHash, baseArr)
             Thread.start()
             baseArr += 1
+            Threads.append(Thread)
 
     elif mytype == 3:
         pool = MySQLConnectionPool()
@@ -32,3 +36,7 @@ if __name__ == '__main__':
             Thread = UploadThread(baseArr, pool)
             Thread.start()
             baseArr += 1
+            Threads.append(Thread)
+
+    for Thread in Threads:
+        Thread.join()