From f05ab3e23e4e29361b4913182e90ffd430e0f6ec Mon Sep 17 00:00:00 2001
From: Suomm <1474983351@qq.com>
Date: Sun, 25 Jun 2023 20:40:31 +0800
Subject: [PATCH] =?UTF-8?q?feat:=20=E7=BC=93=E5=AD=98=E6=96=B9=E6=B3=95?=
 =?UTF-8?q?=E5=AF=B9=E5=BA=94=E7=9A=84=E6=95=B0=E6=8D=AE=E6=BA=90=E3=80=82?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 .../datasource/DataSourceInterceptor.java     | 55 +++++++++++++------
 1 file changed, 37 insertions(+), 18 deletions(-)

diff --git a/mybatis-flex-spring/src/main/java/com/mybatisflex/spring/datasource/DataSourceInterceptor.java b/mybatis-flex-spring/src/main/java/com/mybatisflex/spring/datasource/DataSourceInterceptor.java
index 5face08d..f467f5f8 100644
--- a/mybatis-flex-spring/src/main/java/com/mybatisflex/spring/datasource/DataSourceInterceptor.java
+++ b/mybatis-flex-spring/src/main/java/com/mybatisflex/spring/datasource/DataSourceInterceptor.java
@@ -22,6 +22,12 @@ import com.mybatisflex.core.datasource.DataSourceKey;
 import com.mybatisflex.core.util.StringUtil;
 import org.aopalliance.intercept.MethodInterceptor;
 import org.aopalliance.intercept.MethodInvocation;
+import org.springframework.core.MethodClassKey;
+
+import java.lang.reflect.Method;
+import java.util.Map;
+import java.util.Objects;
+import java.util.concurrent.ConcurrentHashMap;
 
 /**
  * 多数据源切换拦截器。
@@ -39,7 +45,7 @@ public class DataSourceInterceptor implements MethodInterceptor {
             return invocation.proceed();
         }
 
-        dsKey = determineDataSourceKey(invocation);
+        dsKey = findDataSourceKey(invocation.getMethod(), invocation.getThis());
         if (StringUtil.isBlank(dsKey)) {
             return invocation.proceed();
         }
@@ -52,35 +58,48 @@ public class DataSourceInterceptor implements MethodInterceptor {
         }
     }
 
-    private String determineDataSourceKey(MethodInvocation invocation) {
+    /**
+     * 缓存方法对应的数据源。
+     */
+    private final Map<Object, String> dsCache = new ConcurrentHashMap<>();
+
+    private String findDataSourceKey(Method method, Object targetObject) {
+        Object cacheKey = new MethodClassKey(method, targetObject.getClass());
+        String dsKey = this.dsCache.get(cacheKey);
+        if (dsKey == null) {
+            dsKey = determineDataSourceKey(method, targetObject);
+            if (dsKey == null) {
+                dsKey = "";
+            }
+            this.dsCache.put(cacheKey, dsKey);
+        }
+        return dsKey;
+    }
+
+    private String determineDataSourceKey(Method method, Object targetObject) {
 
         // 方法上定义有 UseDataSource 注解
-        UseDataSource annotation = invocation.getMethod().getAnnotation(UseDataSource.class);
+        UseDataSource annotation = method.getAnnotation(UseDataSource.class);
         if (annotation != null) {
             return annotation.value();
         }
 
-        Object target = invocation.getThis();
+        // 类上定义有 UseDataSource 注解
+        Class<?> targetClass = targetObject.getClass();
+        annotation = targetClass.getAnnotation(UseDataSource.class);
+        if (annotation != null) {
+            return annotation.value();
+        }
 
-        if (target != null) {
-            // 类上定义有 UseDataSource 注解
-            Class<?> targetClass = target.getClass();
-            annotation = targetClass.getAnnotation(UseDataSource.class);
+        // 接口上定义有 UseDataSource 注解
+        Class<?>[] interfaces = targetClass.getInterfaces();
+        for (Class<?> anInterface : interfaces) {
+            annotation = anInterface.getAnnotation(UseDataSource.class);
             if (annotation != null) {
                 return annotation.value();
             }
-
-            // 接口上定义有 UseDataSource 注解
-            Class<?>[] interfaces = targetClass.getInterfaces();
-            for (Class<?> anInterface : interfaces) {
-                annotation = anInterface.getAnnotation(UseDataSource.class);
-                if (annotation != null) {
-                    return annotation.value();
-                }
-            }
         }
 
-
         return null;
     }