Completed multiple and dynamic datasource and load balance

Author: 何惠民

src/main/java/cn/com/hellowood/dynamicdatasource/DynamicDataSourceApplication.java

@@ -2,9 +2,8 @@
 
 import org.springframework.boot.SpringApplication;
 import org.springframework.boot.autoconfigure.SpringBootApplication;
-import org.springframework.boot.autoconfigure.jdbc.DataSourceAutoConfiguration;
 
-@SpringBootApplication(exclude = {DataSourceAutoConfiguration.class})
+@SpringBootApplication
 public class DynamicDataSourceApplication {
 
     public static void main(String[] args) {

src/main/java/cn/com/hellowood/dynamicdatasource/common/CommonConstant.java

@@ -9,16 +9,22 @@
  */
 public class CommonConstant {
 
-    //Request result message
+    /**
+     * Request result message
+     */
     public static final String DEFAULT_SUCCESS_MESSAGE = "success";
     public static final String DEFAULT_FAIL_MESSAGE = "fail";
     public static final String NO_RESULT_MESSAGE = "no result";
 
-    //Operation status
+    /**
+     * Operation status
+     */
     public static final String SUCCESS = "SUCCESS";
     public static final String ERROR = "ERROR";
 
-    //Error or exception message
+    /**
+     * Error or exception message
+     */
     public static final String DB_ERROR_MESSAGE = "Database Error";
     public static final String SERVER_ERROR_MESSAGE = "Server Error";
 }

src/main/java/cn/com/hellowood/dynamicdatasource/common/DataSourceKey.java

@@ -0,0 +1,27 @@
+package cn.com.hellowood.dynamicdatasource.common;
+
+/**
+ * The enum Data source key.
+ *
+ * @author HelloWood
+ * @date 2017-08-15 14:26
+ * @Email hellowoodes@gmail.com
+ */
+public enum DataSourceKey {
+    /**
+     * Master data source key.
+     */
+    master,
+    /**
+     * Slave alpha data source key.
+     */
+    slaveAlpha,
+    /**
+     * Slave beta data source key.
+     */
+    slaveBeta,
+    /**
+     * Slave gamma data source key.
+     */
+    slaveGamma
+}

src/main/java/cn/com/hellowood/dynamicdatasource/common/ResponseUtil.java

@@ -75,8 +75,7 @@ public static CommonResponse generateResponse(Object data) {
         } else {
             commonResponse
                     .setCode(ResponseCode.SUCCESS)
-                    .setMessage(CommonConstant.NO_RESULT_MESSAGE)
-                    .setData(data);
+                    .setMessage(CommonConstant.NO_RESULT_MESSAGE);
 
         }
         return commonResponse;

src/main/java/cn/com/hellowood/dynamicdatasource/configuration/CustomHandlerExceptionResolver.java

@@ -32,16 +32,17 @@ public ModelAndView resolveException(HttpServletRequest request, HttpServletResp
         CommonResponse commonResponse = new CommonResponse();
         if (handler instanceof HandlerMethod) {
             HandlerMethod handlerMethod = (HandlerMethod) handler;
-
-            if (ex instanceof ServiceException) {//Service exception,handler exception from service
+            //Service exception,handler exception from service
+            if (ex instanceof ServiceException) {
                 commonResponse.setCode(ResponseCode.SUCCESS).setMessage(ex.getMessage());
                 logger.warn(ex.getMessage());
             } else {
-
-                if (ex instanceof DataAccessException) {//DB exception
+                //DB exception
+                if (ex instanceof DataAccessException) {
                     commonResponse.setCode(ResponseCode.INTERNAL_SERVER_ERROR)
                             .setMessage(CommonConstant.DB_ERROR_MESSAGE);
-                } else {//Others exception
+                } else {
+                    //Others exception
                     commonResponse.setCode(ResponseCode.INTERNAL_SERVER_ERROR)
                             .setMessage(CommonConstant.SERVER_ERROR_MESSAGE);
                 }

src/main/java/cn/com/hellowood/dynamicdatasource/configuration/DataSourceConfigurer.java

@@ -1,10 +1,14 @@
 package cn.com.hellowood.dynamicdatasource.configuration;
 
+import cn.com.hellowood.dynamicdatasource.common.DataSourceKey;
 import org.mybatis.spring.SqlSessionFactoryBean;
 import org.springframework.boot.autoconfigure.jdbc.DataSourceBuilder;
 import org.springframework.boot.context.properties.ConfigurationProperties;
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Configuration;
+import org.springframework.context.annotation.Primary;
+import org.springframework.jdbc.datasource.DataSourceTransactionManager;
+import org.springframework.transaction.PlatformTransactionManager;
 
 import javax.sql.DataSource;
 import java.util.HashMap;
@@ -26,19 +30,42 @@ public class DataSourceConfigurer {
      * @return data source
      */
     @Bean("master")
+    @Primary
     @ConfigurationProperties(prefix = "application.server.db.master")
     public DataSource master() {
         return DataSourceBuilder.create().build();
     }
 
     /**
-     * slave DataSource
+     * Slave alpha data source.
      *
-     * @return data source
+     * @return the data source
+     */
+    @Bean("slaveAlpha")
+    @ConfigurationProperties(prefix = "application.server.db.slave-alpha")
+    public DataSource slaveAlpha() {
+        return DataSourceBuilder.create().build();
+    }
+
+    /**
+     * Slave beta data source.
+     *
+     * @return the data source
      */
-    @Bean("slave")
-    @ConfigurationProperties(prefix = "application.server.db.slave")
-    public DataSource slave() {
+    @Bean("slaveBeta")
+    @ConfigurationProperties(prefix = "application.server.db.slave-beta")
+    public DataSource slaveBeta() {
+        return DataSourceBuilder.create().build();
+    }
+
+    /**
+     * Slave gamma data source.
+     *
+     * @return the data source
+     */
+    @Bean("slaveGamma")
+    @ConfigurationProperties(prefix = "application.server.db.slave-gamma")
+    public DataSource slaveGamma() {
         return DataSourceBuilder.create().build();
     }
 
@@ -50,9 +77,11 @@ public DataSource slave() {
     @Bean("dynamicDataSource")
     public DataSource dynamicDataSource() {
         DynamicRoutingDataSource dynamicRoutingDataSource = new DynamicRoutingDataSource();
-        Map<Object, Object> dataSourceMap = new HashMap<>(2);
-        dataSourceMap.put("master", master());
-        dataSourceMap.put("slave", slave());
+        Map<Object, Object> dataSourceMap = new HashMap<>(4);
+        dataSourceMap.put(DataSourceKey.master.name(), master());
+        dataSourceMap.put(DataSourceKey.slaveAlpha.name(), slaveAlpha());
+        dataSourceMap.put(DataSourceKey.slaveBeta.name(), slaveBeta());
+        dataSourceMap.put(DataSourceKey.slaveGamma.name(), slaveGamma());
 
         // Set master datasource as default
         dynamicRoutingDataSource.setDefaultTargetDataSource(master());
@@ -61,6 +90,10 @@ public DataSource dynamicDataSource() {
 
         // To put datasource keys into DataSourceContextHolder to judge if the datasource is exist
         DynamicDataSourceContextHolder.dataSourceKeys.addAll(dataSourceMap.keySet());
+
+        // To put slave datasource keys into DataSourceContextHolder to load balance
+        DynamicDataSourceContextHolder.slaveDataSourceKeys.addAll(dataSourceMap.keySet());
+        DynamicDataSourceContextHolder.slaveDataSourceKeys.remove(DataSourceKey.master.name());
         return dynamicRoutingDataSource;
     }
 
@@ -83,5 +116,15 @@ public SqlSessionFactoryBean sqlSessionFactoryBean() {
         sqlSessionFactoryBean.setDataSource(dynamicDataSource());
         return sqlSessionFactoryBean;
     }
+
+    /**
+     * Transaction manager platform transaction manager.
+     *
+     * @return the platform transaction manager
+     */
+    @Bean
+    public PlatformTransactionManager transactionManager() {
+        return new DataSourceTransactionManager(dynamicDataSource());
+    }
 }
 

src/main/java/cn/com/hellowood/dynamicdatasource/configuration/DynamicDataSourceAspect.java

@@ -4,9 +4,9 @@
 import org.aspectj.lang.annotation.After;
 import org.aspectj.lang.annotation.Aspect;
 import org.aspectj.lang.annotation.Before;
+import org.aspectj.lang.annotation.Pointcut;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
-import org.springframework.core.annotation.Order;
 import org.springframework.stereotype.Component;
 
 /**
@@ -17,39 +17,60 @@
  * @email hellowoodes@gmail.com
  */
 @Aspect
-@Order(-1) // To ensure execute before @Transactional
 @Component
 public class DynamicDataSourceAspect {
     private static final Logger logger = LoggerFactory.getLogger(DynamicDataSourceAspect.class);
 
+    private final String[] QUERY_PREFIX = {"select"};
+
+    /**
+     * Dao aspect.
+     */
+    @Pointcut("execution( * cn.com.hellowood.dynamicdatasource.mapper.*.*(..))")
+    public void daoAspect() {
+    }
+
     /**
      * Switch DataSource
      *
-     * @param point
-     * @param targetDataSource
+     * @param point the point
      */
-    @Before("@annotation(targetDataSource))")
-    public void switchDataSource(JoinPoint point, TargetDataSource targetDataSource) {
-        if (!DynamicDataSourceContextHolder.containDataSourceKey(targetDataSource.value())) {
-            logger.error("DataSource [{}] doesn't exist, use default DataSource [{}]", targetDataSource.value());
-        } else {
-            DynamicDataSourceContextHolder.setDataSourceKey(targetDataSource.value());
-            logger.info("Switch DataSource to [{}] in Method [{}]",
+    @Before("daoAspect()")
+    public void switchDataSource(JoinPoint point) {
+        Boolean isQueryMethod = isQueryMethod(point.getSignature().getName());
+        if (isQueryMethod) {
+            DynamicDataSourceContextHolder.useSlaveDataSource();
+            logger.debug("Switch DataSource to [{}] in Method [{}]",
                     DynamicDataSourceContextHolder.getDataSourceKey(), point.getSignature());
         }
     }
 
     /**
      * Restore DataSource
      *
-     * @param point
-     * @param targetDataSource
+     * @param point the point
      */
-    @After("@annotation(targetDataSource))")
-    public void restoreDataSource(JoinPoint point, TargetDataSource targetDataSource) {
+    @After("daoAspect())")
+    public void restoreDataSource(JoinPoint point) {
         DynamicDataSourceContextHolder.clearDataSourceKey();
-        logger.info("Restore DataSource to [{}] in Method [{}]",
+        logger.debug("Restore DataSource to [{}] in Method [{}]",
                 DynamicDataSourceContextHolder.getDataSourceKey(), point.getSignature());
     }
 
+
+    /**
+     * Judge if method start with query prefix
+     *
+     * @param methodName
+     * @return
+     */
+    private Boolean isQueryMethod(String methodName) {
+        for (String prefix : QUERY_PREFIX) {
+            if (methodName.startsWith(prefix)) {
+                return true;
+            }
+        }
+        return false;
+    }
+
 }

src/main/java/cn/com/hellowood/dynamicdatasource/configuration/DynamicDataSourceContextHolder.java

@@ -1,8 +1,14 @@
 package cn.com.hellowood.dynamicdatasource.configuration;
 
 
+import cn.com.hellowood.dynamicdatasource.common.DataSourceKey;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
 import java.util.ArrayList;
 import java.util.List;
+import java.util.concurrent.locks.Lock;
+import java.util.concurrent.locks.ReentrantLock;
 
 /**
  * Multiple DataSource Context Holder
@@ -13,30 +19,61 @@
  */
 public class DynamicDataSourceContextHolder {
 
-//    private static final Logger logger = LoggerFactory.getLogger(DynamicDataSourceContextHolder.class);
+    private static final Logger logger = LoggerFactory.getLogger(DynamicDataSourceContextHolder.class);
+
+    private static Lock lock = new ReentrantLock();
+
+    private static int counter = 0;
 
     /**
      * Maintain variable for every thread, to avoid effect other thread
      */
-    private static final ThreadLocal<String> contextHolder = new ThreadLocal<String>() {
-        @Override
-        protected String initialValue() {
-            return "master";
-        }
-    };
+    private static final ThreadLocal<String> CONTEXT_HOLDER = ThreadLocal.withInitial(DataSourceKey.master::name);
+
 
     /**
      * All DataSource List
      */
     public static List<Object> dataSourceKeys = new ArrayList<>();
 
+    /**
+     * The constant slaveDataSourceKeys.
+     */
+    public static List<Object> slaveDataSourceKeys = new ArrayList<>();
+
     /**
      * To switch DataSource
      *
      * @param key the key
      */
     public static void setDataSourceKey(String key) {
-        contextHolder.set(key);
+        CONTEXT_HOLDER.set(key);
+    }
+
+    /**
+     * Use master data source.
+     */
+    public static void useMasterDataSource() {
+        CONTEXT_HOLDER.set(DataSourceKey.master.name());
+    }
+
+    /**
+     * Use slave data source.
+     */
+    public static void useSlaveDataSource() {
+        lock.lock();
+
+        try {
+            int datasourceKeyIndex = counter % slaveDataSourceKeys.size();
+            CONTEXT_HOLDER.set(String.valueOf(slaveDataSourceKeys.get(datasourceKeyIndex)));
+            counter++;
+        } catch (Exception e) {
+            logger.error("Switch slave datasource failed, error message is {}", e.getMessage());
+            useMasterDataSource();
+            e.printStackTrace();
+        } finally {
+            lock.unlock();
+        }
     }
 
     /**
@@ -45,14 +82,14 @@ public static void setDataSourceKey(String key) {
      * @return data source key
      */
     public static String getDataSourceKey() {
-        return contextHolder.get();
+        return CONTEXT_HOLDER.get();
     }
 
     /**
      * To set DataSource as default
      */
     public static void clearDataSourceKey() {
-        contextHolder.remove();
+        CONTEXT_HOLDER.remove();
     }
 
     /**