gistfile1.java

import com.codahale.metrics.Gauge;
import com.codahale.metrics.Metric;
import com.codahale.metrics.MetricRegistry;
import com.sensus.common.database.hikari.SingletonMetricsRegistry;
import com.sensus.util.timing.Timeout;
import com.sms.db.manager.DatabaseState;
import com.sms.db.manager.DatabaseType;
import com.sms.db.manager.DbConnector;
import com.zaxxer.hikari.HikariConfig;
import com.zaxxer.hikari.HikariDataSource;
import org.junit.After;
import org.junit.Before;
import org.junit.Test;
import org.springframework.jdbc.datasource.DriverManagerDataSource;

import javax.sql.DataSource;
import java.sql.Connection;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.concurrent.atomic.AtomicBoolean;


public class HikariPoolTest {

    protected DataSource source;
    protected DatabaseType databaseType = DatabaseType.MSSQL;

    @Before
    public void testSetup() throws Exception {
        source = DbConnector.getInstance().fromBackup(databaseType, DatabaseState.CLEAN);
    }

    @After
    public void testTearDown() throws Exception {
        DbConnector.getInstance().returnDatabase(source);
    }

    @Test
    public void testHikariPoolSize() throws Exception {

        final List<Integer> threadCounts = Arrays.asList(1, 8, 16, 32);
        final List<Integer> operationTime = Arrays.asList(0, 1, 10, 100, 1000);
        final List<Integer> idleTimes = Arrays.asList(0, 1, 10, 100);

//        final List<Integer> threadCounts = Arrays.asList(8);
//        final List<Integer> operationTime = Arrays.asList(1, 100);
//        final List<Integer> idleTimes = Arrays.asList(0);

        List<Counts> counts = new ArrayList<>();

        for (Integer threadCount : threadCounts) {
            for (Integer opTime : operationTime) {
                for (Integer idleTime : idleTimes) {
                    counts.add(collectSteadyState(opTime, idleTime, threadCount, 10000));
                }
            }
        }

        System.out.println("Overall");
        for (Counts count : counts) {
            System.out.println(count);
        }
    }

    private Counts collectSteadyState(final int operationTime, final int idleTime, int threadCount, int testTime) throws InterruptedException {
        final String poolName = "test" + operationTime + "id" + idleTime;
        DriverManagerDataSource ds = (DriverManagerDataSource) source;
        // we don't need to copy DriverClassName since that is actually just a hint to call Class.forName() with
        // the driver name so it registers itself for the prefix "jtds:" etc.
        HikariConfig config = new HikariConfig();
        config.setJdbcUrl(ds.getUrl());
        config.setUsername(ds.getUsername());
        config.setPassword(ds.getPassword());
        config.setConnectionTestQuery("select 1 from dual");
        config.setPoolName(poolName);
        config.setRegisterMbeans(true);
        config.setMetricRegistry(SingletonMetricsRegistry.getInstance().getRegistry());
        // we want to minimize how many connections to start so we don't overload the database server on restart
        config.setMinimumIdle(2);
        config.setMaximumPoolSize(100);
        final HikariDataSource hikariDataSource = new HikariDataSource(config);

        List<Thread> threads = new ArrayList<>();

        final AtomicBoolean running = new AtomicBoolean(true);

        for (int i = 0; i < threadCount; i++) {
            threads.add(new Thread(new Runnable() {
                @Override
                public void run() {
                    try {
                        while (running.get()) {
                            final Connection connection = hikariDataSource.getConnection();
                            if (operationTime > 0) {
                                Thread.sleep(operationTime);
                            }
                            connection.close();
                            if (idleTime > 0) {
                                Thread.sleep(idleTime);
                            }
                        }
                    } catch (Exception e) {
                        e.printStackTrace();
                    }
                }
            }));
        }

        for (Thread thread : threads) {
            thread.start();
        }

        int maxActive = 0;
        int maxTotal = 0;

        Timeout testRunTime = new Timeout(testTime);
        while (!testRunTime.isExpired()) {
            Thread.sleep(100);
            final Counts counts = getCounts(hikariDataSource, poolName);
            maxActive = Math.max(maxActive, counts.active);
            maxTotal = Math.max(maxTotal, counts.total);
            System.out.println(counts);
        }

        running.set(false);
        Thread.sleep(operationTime * 2);

        for (Thread thread : threads) {
            thread.interrupt();
            thread.join();
        }
        hikariDataSource.close();
        return new Counts(maxActive, maxTotal, "opTime=" + operationTime + "idleTime=" + idleTime + " threadCount=" + threadCount);
    }

    private Counts getCounts(HikariDataSource hikariDataSource, String poolName1) {
        final MetricRegistry metricRegistry = (MetricRegistry) hikariDataSource.getMetricRegistry();

        final Metric activeMetric = metricRegistry.getMetrics().get(poolName1 + ".pool.ActiveConnections");
        final Integer active = (Integer) ((Gauge) activeMetric).getValue();
        final Metric totalMetric = metricRegistry.getMetrics().get(poolName1 + ".pool.TotalConnections");
        final Integer total = (Integer) ((Gauge) totalMetric).getValue();
        return new Counts(active, total, "Test");
    }

    private static class Counts {
        private final String name;
        private final int active;
        private final int total;

        public Counts(int active, int total, String name) {
            this.active = active;
            this.total = total;
            this.name = name;
        }

        @Override
        public String toString() {
            return name + " active=" + active +
                    "  total=" + total;
        }
    }
}