优化数据源切换逻辑

hsweb-datasource/hsweb-datasource-api/src/main/java/org/hswebframework/web/datasource/AopDataSourceSwitcherAutoConfiguration.java

@@ -18,6 +18,7 @@ import java.lang.reflect.Method;
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.function.Consumer;
 import java.util.function.Function;
 
@@ -91,11 +92,12 @@ public class AopDataSourceSwitcherAutoConfiguration {
 
                 Consumer<MethodInterceptorContext> before = context -> {
                 };
-
+                AtomicBoolean dataSourceChanged = new AtomicBoolean(true);
                 if (matcher != null) {
                     before = before.andThen(context -> {
                         Strategy strategy = matcher.getStrategy(context);
                         if (strategy == null) {
+                            dataSourceChanged.set(false);
                             logger.warn("strategy matcher found:{}, but strategy is null!", matcher);
                         } else {
                             logger.debug("switch datasource.use strategy:{}", strategy);
@@ -117,8 +119,10 @@ public class AopDataSourceSwitcherAutoConfiguration {
                                         DataSourceHolder.switcher().use(id);
                                     }
                                 } catch (RuntimeException e) {
+                                    dataSourceChanged.set(false);
                                     throw e;
                                 } catch (Exception e) {
+                                    dataSourceChanged.set(false);
                                     throw new RuntimeException(e.getMessage(), e);
                                 }
                             }
@@ -142,7 +146,9 @@ public class AopDataSourceSwitcherAutoConfiguration {
                 try {
                     return methodInvocation.proceed();
                 } finally {
-                    DataSourceHolder.switcher().useLast();
+                    if (dataSourceChanged.get()) {
+                        DataSourceHolder.switcher().useLast();
+                    }
                     DataSourceHolder.tableSwitcher().reset();
                 }
             });
@@ -150,19 +156,21 @@ public class AopDataSourceSwitcherAutoConfiguration {
 
         @Override
         public boolean matches(Method method, Class<?> aClass) {
-            CacheKey key = new CacheKey(aClass, method);
+            Class<?> targetClass = ClassUtils.getUserClass(aClass);
+
+            CacheKey key = new CacheKey(targetClass, method);
             matchers.stream()
-                    .filter(matcher -> matcher.match(aClass, method))
+                    .filter(matcher -> matcher.match(targetClass, method))
                     .findFirst()
                     .ifPresent((matcher) -> cache.put(key, matcher));
 
-
             boolean datasourceMatched = cache.containsKey(key);
             boolean tableMatched = false;
             if (null != tableSwitcher) {
-                CachedTableSwitchStrategyMatcher.CacheKey tableCacheKey = new CachedTableSwitchStrategyMatcher.CacheKey(aClass, method);
+                CachedTableSwitchStrategyMatcher.CacheKey tableCacheKey = new CachedTableSwitchStrategyMatcher
+                        .CacheKey(targetClass, method);
                 tableSwitcher.stream()
-                        .filter(matcher -> matcher.match(aClass, method))
+                        .filter(matcher -> matcher.match(targetClass, method))
                         .findFirst()
                         .ifPresent((matcher) -> tableCache.put(tableCacheKey, matcher));
                 tableMatched = tableCache.containsKey(tableCacheKey);

hsweb-datasource/hsweb-datasource-api/src/main/java/org/hswebframework/web/datasource/strategy/CachedDataSourceSwitchStrategyMatcher.java

@@ -54,12 +54,14 @@ public abstract class CachedDataSourceSwitchStrategyMatcher implements DataSourc
                 return false;
             }
             CacheKey target = ((CacheKey) obj);
-            return target.target == this.target && target.method == method;
+            return target.target.getName().equals(this.target.getName())
+                    && target.method.getName().equals(method.getName())
+                    && target.method.getParameterCount() == method.getParameterCount();
         }
 
         public int hashCode() {
-            int result = this.target != null ? this.target.hashCode() : 0;
-            result = 31 * result + (this.method != null ? this.method.hashCode() : 0);
+            int result = this.target != null ? this.target.getName().hashCode() : 0;
+            result = 31 * result + (this.method != null ? this.method.getName().hashCode() : 0);
             return result;
         }
     }

hsweb-datasource/hsweb-datasource-api/src/test/java/org/hswebframework/web/datasource/switcher/DefaultDataSourceSwitcherTest.java

@@ -1,8 +1,11 @@
 package org.hswebframework.web.datasource.switcher;
 
+import lombok.SneakyThrows;
 import org.junit.Assert;
 import org.junit.Test;
 
+import java.lang.reflect.Method;
+
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertTrue;
 
@@ -56,5 +59,21 @@ public class DefaultDataSourceSwitcherTest {
         assertEquals(switcher.currentDataSourceId(), "test");
     }
 
+    @SneakyThrows
+    public static void main(String[] args) {
+        Method method = Test2.class.getMethod("test");
+
+        System.out.println(method);
+        System.out.println(method.getDeclaringClass());
+    }
+
+    public class Test2 extends TestClass {
 
+    }
+
+    public class TestClass {
+        public void test() {
+
+        }
+    }
 }

hsweb-datasource/hsweb-datasource-jta/src/test/java/org/hswebframework/web/datasource/jta/SimpleAtomikosTests.java

@@ -168,10 +168,20 @@ public class SimpleAtomikosTests extends AbstractTransactionalJUnit4SpringContex
         Thread.sleep(1000);
     }
 
+    public interface TestService {
+        List findAll() throws SQLException;
+    }
+
+    public static class AbstractTest implements TestService {
+        RDBDatabase database;
+
+        public List findAll() throws SQLException {
+            return database.getTable("s_user").createQuery().list();
+        }
+    }
 
     @Transactional
-    public static class DynDsTest {
-        private RDBDatabase database;
+    public static class DynDsTest extends AbstractTest {
 
         @Transactional(propagation = Propagation.NOT_SUPPORTED)
         public void testCreateTable() throws SQLException {
@@ -195,9 +205,6 @@ public class SimpleAtomikosTests extends AbstractTransactionalJUnit4SpringContex
             return database.getTable("s_user").createQuery().list();
         }
 
-        public List findAll() throws SQLException {
-            return database.getTable("s_user").createQuery().list();
-        }
 
         @UseDataSource("test_ds")
         public List query() throws SQLException {