com.baidu.rigel.biplatform.tesseract.datasource.DynamicSqlDataSource.java Source code

Java tutorial

Introduction

Here is the source code for com.baidu.rigel.biplatform.tesseract.datasource.DynamicSqlDataSource.java

Source

/**
 * Copyright (c) 2014 Baidu, Inc. All Rights Reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *         http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
/**
 * 
 */
package com.baidu.rigel.biplatform.tesseract.datasource;

import java.sql.Connection;
import java.sql.PreparedStatement;
import java.sql.SQLException;
import java.util.Collections;
import java.util.HashMap;
import java.util.Iterator;
import java.util.LinkedHashSet;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;

import org.apache.commons.collections.MapUtils;
import org.apache.commons.lang.StringUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.baidu.rigel.biplatform.tesseract.datasource.impl.SqlDataSourceWrap;
import com.google.common.collect.Lists;
import com.google.common.collect.Maps;
import com.mchange.v2.c3p0.ComboPooledDataSource;

/**
 * ???
 * 
 * @author xiaoming.chen
 *
 */
public class DynamicSqlDataSource {

    /**
     * LOGGER
     */
    private static final Logger LOGGER = LoggerFactory.getLogger(DynamicSqlDataSource.class);

    /**
     * TEST_SQL ?SQL
     */
    private static final String TEST_SQL = "select 1";

    /**
     * validInterval 
     */
    private long validInterval = 10000L;

    /**
     * dataSources ??
     */
    private Map<String, SqlDataSourceWrap> dataSources;

    /**
     * executor ??
     */
    private ExecutorService executor;

    /**
     * validSql SQL
     */
    private String validSql = TEST_SQL;

    /**
     * failedMap DataSource
     */
    private final Map<String, SqlDataSourceWrap> failedMap = new ConcurrentHashMap<String, SqlDataSourceWrap>();

    /**
     * recoverHeartBeat ??
     */
    private DataSourceRecoverHeartBeat recoverHeartBeat;

    /**
     * accessValidDataSourceKeys ???KEY
     */
    private Set<String> accessValidDataSourceKeys = Collections.synchronizedSet(new LinkedHashSet<String>());

    /**
     * lastSuccessGetDataSourcTime ???
     */
    private long lastSuccessGetDataSourcTime;

    /**
     * construct with
     * 
     * @param dataSources ??
     */
    public DynamicSqlDataSource(Map<String, SqlDataSourceWrap> dataSources) {
        this.dataSources = Maps.newHashMap(dataSources);
        this.lastSuccessGetDataSourcTime = System.currentTimeMillis();
    }

    /**
     * construct with
     * 
     * @param validInterval ??1S
     * @param dataSources ??
     * @param validSql SQL?? select 1
     */
    public DynamicSqlDataSource(long validInterval, Map<String, SqlDataSourceWrap> dataSources, String validSql) {
        this(dataSources);
        this.validInterval = validInterval;
        this.validSql = validSql;
    }

    /**
     * ???
     * 
     * @throws Exception ?
     */
    public void destroy() throws Exception {
        shutdownHeartBean();
        if (MapUtils.isNotEmpty(dataSources)) {
            dataSources.forEach((key, ds) -> {
                if (ds.getDataSource() instanceof ComboPooledDataSource) {
                    ((ComboPooledDataSource) ds.getDataSource()).close();
                }
            });
        }
    }

    /**
     * ?????KEY??
     * 
     * @param dataSourceKey ??KEY
     * @return ??
     * @throws IllegalArgumentException ?
     */
    public synchronized SqlDataSourceWrap removeDataSourceByKey(String dataSourceKey) {
        if (StringUtils.isBlank(dataSourceKey)) {
            throw new IllegalArgumentException("can not remove datasource by null key");
        }
        accessValidDataSourceKeys.remove(dataSourceKey);
        failedMap.remove(dataSourceKey);
        return dataSources.remove(dataSourceKey);
    }

    /** 
     * clearDataSourceFailCount
     */
    public void clearDataSourceFailCount() {
        if (MapUtils.isNotEmpty(failedMap)) {
            final long current = System.currentTimeMillis();
            failedMap.forEach((key, ds) -> {
                if (ds.getFailCount() >= 5 && (ds.getFailTime() - current) / (36 * 2 * 1e5) > 1) {
                    ds.resetFailCount();
                }
            });
        }
    }

    /**
     * ??
     * 
     * @param key ??KEY
     * @param dataSource ??
     * @throws IllegalArgumentException ?
     */
    public synchronized void addDataSource(String key, SqlDataSourceWrap dataSource) {
        if (StringUtils.isBlank(key)) {
            throw new IllegalArgumentException("can not add datasource by null key");
        }
        if (dataSource == null) {
            throw new IllegalArgumentException("can not add datasource by null datasource");
        }
        dataSources.put(key, dataSource);
    }

    /**
     * ??DataSource
     * 
     * @return DataSource
     * @throw IllegalArgumentException ??
     */
    public synchronized SqlDataSourceWrap getDataSource() {
        String key = getDataSourceKey();
        SqlDataSourceWrap result = dataSources.get(key);
        Connection connection = null;
        try {
            connection = result.getConnection();
            LOGGER.info("validate datasource by key:" + key);
            validateConnection(connection);
            LOGGER.info("return datasource by key:" + key);
            // ?
            accessValidDataSourceKeys.remove(key);
            accessValidDataSourceKeys.add(key);
            return result;
        } catch (Exception e) {
            LOGGER.warn("datasource key:" + key + " is invalide,try another one!", e);
            failedMap.put(key, result);
            accessValidDataSourceKeys.remove(key);
            checkFailedDataSourceHeatBeat();
            return getDataSource();
        } finally {
            if (connection != null) {
                try {
                    connection.close();
                } catch (SQLException e) {
                    e.printStackTrace();
                }
            }
        }
    }

    /**
     * ?DataSourcekey???KEY ?KEY
     * 
     * @return ?DataSourcekey
     * @throw IllegalArgumentException ??
     */
    private String getDataSourceKey() {
        if (MapUtils.isEmpty(dataSources)) {
            throw new IllegalArgumentException("can not get datasource by empty datasources!");
        }
        List<String> allKeys = Lists.newArrayList(dataSources.keySet());
        LOGGER.debug("get " + allKeys.size() + " size avaliable datasources");
        if (!failedMap.isEmpty()) {
            Set<String> failedKeys = failedMap.keySet();
            LOGGER.warn("found failed datasource keys :" + failedKeys);
            allKeys.removeAll(failedKeys);
            accessValidDataSourceKeys.removeAll(failedKeys);
        }
        if (allKeys.size() > 0) {
            allKeys.removeAll(accessValidDataSourceKeys);
            if (allKeys.size() > 0) {
                LOGGER.debug("after filter access valid datasource," + allKeys);
                return allKeys.get(0);
            } else {
                LOGGER.debug("return first time access datasource key:" + accessValidDataSourceKeys);
                return (String) accessValidDataSourceKeys.toArray()[0];
            }
        }
        throw new IllegalArgumentException("can not get datasource key because all datasource is failed!");
    }

    /**
     * DataSource
     */
    private synchronized void checkFailedDataSourceHeatBeat() {
        if (recoverHeartBeat == null) {
            recoverHeartBeat = new DataSourceRecoverHeartBeat(this);
            if (executor == null) {
                executor = Executors.newFixedThreadPool(1);
            }
            executor.execute(recoverHeartBeat);
        } else {
            if (!recoverHeartBeat.isRuning()) {
                if (executor == null) {
                    executor = Executors.newFixedThreadPool(1);
                }
                executor.execute(recoverHeartBeat);
            }
        }

    }

    /**
     * connection
     * 
     * @param con connection
     * @throws SQLException connection
     */
    private void validateConnection(Connection con) throws SQLException {
        long current = System.currentTimeMillis();
        PreparedStatement stmt = con.prepareStatement(validSql); // test
                                                                 // connection
                                                                 // is ok
        stmt.executeQuery();
        stmt.close();
        LOGGER.debug("validate connection cost:" + (System.currentTimeMillis() - current));
    }

    /**
     * 
     */
    private synchronized void shutdownHeartBean() {
        if (recoverHeartBeat != null) {
            recoverHeartBeat.close();
        }
        if (executor != null) {
            executor.shutdown();
        }
    }

    /**
     * Connection status recover heart beat thread.
     * 
     * @author xiemalin
     *
     */
    private static class DataSourceRecoverHeartBeat implements Runnable {
        /**
         * Logger for this class
         */
        private static final Logger LOGGER = LoggerFactory.getLogger(DataSourceRecoverHeartBeat.class);

        /**
         * dynamicDataSource
         */
        private DynamicSqlDataSource dynamicDataSource;
        /**
         * runing 
         */
        private boolean runing;

        /**
         * close ?
         */
        private boolean close = false;

        /**
         * ?
         * 
         * @return ?
         */
        public boolean isRuning() {
            return runing;
        }

        /**
         * 
         */
        public void close() {
            close = true;
        }

        /**
         * construct with
         * 
         * @param dynamicDataSource ???
         */
        public DataSourceRecoverHeartBeat(DynamicSqlDataSource dynamicDataSource) {
            this.dynamicDataSource = dynamicDataSource;
        }

        public void run() {
            runing = true;
            while (!dynamicDataSource.failedMap.isEmpty() && !close) {

                // copy data
                Map<String, SqlDataSourceWrap> dataSourceMapCopy;
                dataSourceMapCopy = new HashMap<String, SqlDataSourceWrap>(dynamicDataSource.failedMap);

                Iterator<Entry<String, SqlDataSourceWrap>> iter;
                iter = dataSourceMapCopy.entrySet().iterator();
                while (iter.hasNext()) {
                    Entry<String, SqlDataSourceWrap> next = iter.next();
                    SqlDataSourceWrap ds = next.getValue();
                    String key = next.getKey();
                    if (ds.getFailCount() >= 10) {
                        LOGGER.warn("Datasource key {} has fail more than 5 time,skip.", key);
                        continue;
                    }
                    Connection con = null;
                    try {
                        con = ds.getConnection();
                        dynamicDataSource.validateConnection(con);
                        LOGGER.debug("Datasource key='" + key + "' valid ok.");
                        dynamicDataSource.failedMap.remove(key);
                        ds.resetFailCount();
                    } catch (SQLException e) {
                        LOGGER.warn("Datasource key='" + key + "' valid failed.");
                        ds.increaseFailCount();
                    } finally {
                        if (con != null) {
                            try {
                                con.close();
                            } catch (SQLException e) {
                                if (LOGGER.isDebugEnabled()) {
                                    LOGGER.debug(e.getMessage(), e);
                                }
                            }
                        }
                    }

                }

                try {
                    Thread.sleep(dynamicDataSource.validInterval);
                } catch (Exception e) {
                    // here we needn't care exception
                }
            }
            runing = false;
        }
    }

    /**
     * getter method for property validInterval
     * 
     * @return the validInterval
     */
    public long getValidInterval() {
        return validInterval;
    }

    /**
     * setter method for property validInterval
     * 
     * @param validInterval the validInterval to set
     */
    public void setValidInterval(long validInterval) {
        this.validInterval = validInterval;
    }

    /**
     * getter method for property validSql
     * 
     * @return the validSql
     */
    public String getValidSql() {
        return validSql;
    }

    /**
     * setter method for property validSql
     * 
     * @param validSql the validSql to set
     */
    public void setValidSql(String validSql) {
        this.validSql = validSql;
    }

    /**
     * get lastSuccessGetDataSourcTime
     * 
     * @return the lastSuccessGetDataSourcTime
     */
    public long getLastSuccessGetDataSourcTime() {
        return lastSuccessGetDataSourcTime;
    }

    /**
     * set lastSuccessGetDataSourcTime with lastSuccessGetDataSourcTime
     * 
     * @param lastSuccessGetDataSourcTime the lastSuccessGetDataSourcTime to set
     */
    public void setLastSuccessGetDataSourcTime(long lastSuccessGetDataSourcTime) {
        this.lastSuccessGetDataSourcTime = lastSuccessGetDataSourcTime;
    }

    /** 
     * ? dataSources 
     * @return the dataSources 
     */
    public Map<String, SqlDataSourceWrap> getDataSources() {

        return dataSources;
    }
}