From 5f19924f9f93d325181c2a1c5537c73e44f505f0 Mon Sep 17 00:00:00 2001
From: shaohuzhang1 <80892890+shaohuzhang1@users.noreply.github.com>
Date: Mon, 9 Jun 2025 21:09:51 +0800
Subject: [PATCH] perf: model manage (#3226)

---
 apps/common/config/embedding_config.py | 50 +++++++++++++++++---------
 apps/common/utils/rsa_util.py          | 11 +++---
 2 files changed, 37 insertions(+), 24 deletions(-)

diff --git a/apps/common/config/embedding_config.py b/apps/common/config/embedding_config.py
index 50ed5e23..6d6bfde9 100644
--- a/apps/common/config/embedding_config.py
+++ b/apps/common/config/embedding_config.py
@@ -6,40 +6,56 @@
     @date：2023/10/23 16:03
     @desc:
 """
+
 import threading
 import time
 
 from common.cache.mem_cache import MemCache
 
-lock = threading.Lock()
+_lock = threading.Lock()
+locks = {}
 
 
 class ModelManage:
     cache = MemCache('model', {})
     up_clear_time = time.time()
 
+    @staticmethod
+    def _get_lock(_id):
+        lock = locks.get(_id)
+        if lock is None:
+            with _lock:
+                lock = locks.get(_id)
+                if lock is None:
+                    lock = threading.Lock()
+                    locks[_id] = lock
+
+        return lock
+
     @staticmethod
     def get_model(_id, get_model):
-        # 获取锁
-        lock.acquire()
-        try:
-            model_instance = ModelManage.cache.get(_id)
-            if model_instance is None or not model_instance.is_cache_model():
+        model_instance = ModelManage.cache.get(_id)
+        if model_instance is None:
+            lock = ModelManage._get_lock(_id)
+            with lock:
+                model_instance = ModelManage.cache.get(_id)
+                if model_instance is None:
+                    model_instance = get_model(_id)
+                    ModelManage.cache.set(_id, model_instance, timeout=60 * 60 * 8)
+        else:
+            if model_instance.is_cache_model():
+                ModelManage.cache.touch(_id, timeout=60 * 60 * 8)
+            else:
                 model_instance = get_model(_id)
-                ModelManage.cache.set(_id, model_instance, timeout=60 * 30)
-                return model_instance
-            # 续期
-            ModelManage.cache.touch(_id, timeout=60 * 30)
-            ModelManage.clear_timeout_cache()
-            return model_instance
-        finally:
-            # 释放锁
-            lock.release()
+                ModelManage.cache.set(_id, model_instance, timeout=60 * 60 * 8)
+        ModelManage.clear_timeout_cache()
+        return model_instance
 
     @staticmethod
     def clear_timeout_cache():
-        if time.time() - ModelManage.up_clear_time > 60:
-            ModelManage.cache.clear_timeout_data()
+        if time.time() - ModelManage.up_clear_time > 60 * 60:
+            threading.Thread(target=lambda: ModelManage.cache.clear_timeout_data()).start()
+            ModelManage.up_clear_time = time.time()
 
     @staticmethod
     def delete_key(_id):
diff --git a/apps/common/utils/rsa_util.py b/apps/common/utils/rsa_util.py
index 3e20f594..c47663fd 100644
--- a/apps/common/utils/rsa_util.py
+++ b/apps/common/utils/rsa_util.py
@@ -40,15 +40,12 @@ def generate():
 def get_key_pair():
     rsa_value = rsa_cache.get(cache_key)
     if rsa_value is None:
-        lock.acquire()
-        rsa_value = rsa_cache.get(cache_key)
-        if rsa_value is not None:
-            return rsa_value
-        try:
+        with lock:
+            rsa_value = rsa_cache.get(cache_key)
+            if rsa_value is not None:
+                return rsa_value
             rsa_value = get_key_pair_by_sql()
             rsa_cache.set(cache_key, rsa_value)
-        finally:
-            lock.release()
     return rsa_value