/*
 * Copyright (C) 2013 Brett Wooldridge
 *
 * 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.zaxxer.hikari.pool;

import static com.zaxxer.hikari.pool.TestElf.newHikariConfig;
import static com.zaxxer.hikari.pool.TestElf.newHikariDataSource;
import static com.zaxxer.hikari.pool.TestElf.getPool;
import static com.zaxxer.hikari.pool.TestElf.setConfigUnitTest;
import static com.zaxxer.hikari.pool.TestElf.setSlf4jLogLevel;
import static com.zaxxer.hikari.pool.TestElf.setSlf4jTargetStream;
import static com.zaxxer.hikari.util.UtilityElf.quietlySleep;
import static java.util.concurrent.TimeUnit.SECONDS;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertFalse;
import static org.junit.Assert.assertNotNull;
import static org.junit.Assert.assertNotSame;
import static org.junit.Assert.assertNull;
import static org.junit.Assert.assertSame;
import static org.junit.Assert.assertTrue;
import static org.junit.Assert.fail;

import java.sql.Connection;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.sql.SQLTransientConnectionException;
import java.sql.Statement;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicReference;

import org.apache.logging.log4j.Level;
import org.junit.After;
import org.junit.Before;
import org.junit.Test;

import com.zaxxer.hikari.HikariConfig;
import com.zaxxer.hikari.HikariDataSource;
import com.zaxxer.hikari.mocks.StubConnection;
import com.zaxxer.hikari.mocks.StubDataSource;
import com.zaxxer.hikari.mocks.StubStatement;
import com.zaxxer.hikari.pool.HikariPool.PoolInitializationException;

/**
 * @author Brett Wooldridge
 */
public class TestConnections
{
   @Before
   public void before()
   {
      setSlf4jTargetStream(HikariPool.class, System.err);
      setSlf4jLogLevel(HikariPool.class, Level.DEBUG);
      setSlf4jLogLevel(PoolBase.class, Level.DEBUG);
   }

   @After
   public void after()
   {
      System.getProperties().remove("com.zaxxer.hikari.housekeeping.periodMs");
      setSlf4jLogLevel(HikariPool.class, Level.WARN);
      setSlf4jLogLevel(PoolBase.class, Level.WARN);
   }

   @Test
   public void testCreate() throws SQLException
   {
      HikariConfig config = newHikariConfig();
      config.setMinimumIdle(1);
      config.setMaximumPoolSize(1);
      config.setConnectionTestQuery("VALUES 1");
      config.setConnectionInitSql("SELECT 1");
      config.setReadOnly(true);
      config.setConnectionTimeout(2500);
      config.setLeakDetectionThreshold(TimeUnit.SECONDS.toMillis(30));
      config.setDataSourceClassName("com.zaxxer.hikari.mocks.StubDataSource");

      try (HikariDataSource ds = new HikariDataSource(config)) {
         ds.setLoginTimeout(10);
         assertSame(10, ds.getLoginTimeout());

         HikariPool pool = getPool(ds);
         ds.getConnection().close();
         assertSame("Total connections not as expected", 1, pool.getTotalConnections());
         assertSame("Idle connections not as expected", 1, pool.getIdleConnections());

         try (Connection connection = ds.getConnection();
              PreparedStatement statement = connection.prepareStatement("SELECT * FROM device WHERE device_id=?")) {

            assertNotNull(connection);
            assertNotNull(statement);
   
            assertSame("Total connections not as expected", 1, pool.getTotalConnections());
            assertSame("Idle connections not as expected", 0, pool.getIdleConnections());
   
            statement.setInt(1, 0);

            try (ResultSet resultSet = statement.executeQuery()) {
               assertNotNull(resultSet);
      
               assertFalse(resultSet.next());
            }   
         }

         assertSame("Total connections not as expected", 1, pool.getTotalConnections());
         assertSame("Idle connections not as expected", 1, pool.getIdleConnections());
      }
   }

   @Test
   public void testMaxLifetime() throws Exception
   {
      HikariConfig config = newHikariConfig();
      config.setMinimumIdle(0);
      config.setMaximumPoolSize(1);
      config.setConnectionTimeout(2500);
      config.setConnectionTestQuery("VALUES 1");
      config.setInitializationFailTimeout(Long.MAX_VALUE);
      config.setDataSourceClassName("com.zaxxer.hikari.mocks.StubDataSource");

      System.setProperty("com.zaxxer.hikari.housekeeping.periodMs", "100");

      setConfigUnitTest(true);
      try (HikariDataSource ds = new HikariDataSource(config)) {
         System.clearProperty("com.zaxxer.hikari.housekeeping.periodMs");

         ds.setMaxLifetime(700);

         HikariPool pool = getPool(ds);

         assertSame("Total connections not as expected", 0, pool.getTotalConnections());
         assertSame("Idle connections not as expected", 0, pool.getIdleConnections());

         Connection unwrap;
         Connection unwrap2;
         try (Connection connection = ds.getConnection()) {
            unwrap = connection.unwrap(Connection.class);
            assertNotNull(connection);
   
            assertSame("Second total connections not as expected", 1, pool.getTotalConnections());
            assertSame("Second idle connections not as expected", 0, pool.getIdleConnections());
         }

         assertSame("Idle connections not as expected", 1, pool.getIdleConnections());

         try (Connection connection = ds.getConnection()) {
            unwrap2 = connection.unwrap(Connection.class);
            assertSame(unwrap, unwrap2);
            assertSame("Second total connections not as expected", 1, pool.getTotalConnections());
            assertSame("Second idle connections not as expected", 0, pool.getIdleConnections());
         }

         quietlySleep(TimeUnit.SECONDS.toMillis(2));

         try (Connection connection = ds.getConnection()) {
            unwrap2 = connection.unwrap(Connection.class);
            assertNotSame("Expected a different connection", unwrap, unwrap2);
         }

         assertSame("Post total connections not as expected", 1, pool.getTotalConnections());
         assertSame("Post idle connections not as expected", 1, pool.getIdleConnections());
      }
      finally {
         setConfigUnitTest(false);
      }
   }

   @Test
   public void testMaxLifetime2() throws Exception
   {
      HikariConfig config = newHikariConfig();
      config.setMinimumIdle(0);
      config.setMaximumPoolSize(1);
      config.setConnectionTimeout(2500);
      config.setConnectionTestQuery("VALUES 1");
      config.setDataSourceClassName("com.zaxxer.hikari.mocks.StubDataSource");

      System.setProperty("com.zaxxer.hikari.housekeeping.periodMs", "100");

      setConfigUnitTest(true);
      try (HikariDataSource ds = new HikariDataSource(config)) {
         ds.setMaxLifetime(700);

         HikariPool pool = getPool(ds);
         assertSame("Total connections not as expected", 0, pool.getTotalConnections());
         assertSame("Idle connections not as expected", 0, pool.getIdleConnections());

         Connection unwrap;
         Connection unwrap2;
         try (Connection connection = ds.getConnection()) {
            unwrap = connection.unwrap(Connection.class);
            assertNotNull(connection);
   
            assertSame("Second total connections not as expected", 1, pool.getTotalConnections());
            assertSame("Second idle connections not as expected", 0, pool.getIdleConnections());
         }

         assertSame("Idle connections not as expected", 1, pool.getIdleConnections());

         try (Connection connection = ds.getConnection()) {
            unwrap2 = connection.unwrap(Connection.class);
            assertSame(unwrap, unwrap2);
            assertSame("Second total connections not as expected", 1, pool.getTotalConnections());
            assertSame("Second idle connections not as expected", 0, pool.getIdleConnections());
         }

         quietlySleep(800);

         try (Connection connection = ds.getConnection()) {
            unwrap2 = connection.unwrap(Connection.class);
            assertNotSame("Expected a different connection", unwrap, unwrap2);
         }

         assertSame("Post total connections not as expected", 1, pool.getTotalConnections());
         assertSame("Post idle connections not as expected", 1, pool.getIdleConnections());
      }
      finally {
         setConfigUnitTest(false);
      }
   }

   @Test
   public void testDoubleClose() throws Exception
   {
      HikariConfig config = newHikariConfig();
      config.setMinimumIdle(1);
      config.setMaximumPoolSize(1);
      config.setConnectionTimeout(2500);
      config.setConnectionTestQuery("VALUES 1");
      config.setDataSourceClassName("com.zaxxer.hikari.mocks.StubDataSource");

      try (HikariDataSource ds = new HikariDataSource(config);
         Connection connection = ds.getConnection()) {
         connection.close();

         // should no-op
         connection.abort(null);

         assertTrue("Connection should have closed", connection.isClosed());
         assertFalse("Connection should have closed", connection.isValid(5));
         assertTrue("Expected to contain ClosedConnection, but was " + connection, connection.toString().contains("ClosedConnection"));
      }
   }

   @Test
   public void testEviction() throws Exception
   {
      HikariConfig config = newHikariConfig();
      config.setMinimumIdle(0);
      config.setMaximumPoolSize(5);
      config.setConnectionTimeout(2500);
      config.setConnectionTestQuery("VALUES 1");
      config.setDataSourceClassName("com.zaxxer.hikari.mocks.StubDataSource");

      try (HikariDataSource ds = new HikariDataSource(config)) {
         Connection connection = ds.getConnection();

         HikariPool pool = getPool(ds);
         assertEquals(1, pool.getTotalConnections());
         ds.evictConnection(connection);
         assertEquals(0, pool.getTotalConnections());
      }
   }

   @Test
   public void testBackfill() throws Exception
   {
      HikariConfig config = newHikariConfig();
      config.setMinimumIdle(1);
      config.setMaximumPoolSize(4);
      config.setConnectionTimeout(1000);
      config.setInitializationFailTimeout(Long.MAX_VALUE);
      config.setConnectionTestQuery("VALUES 1");
      config.setDataSourceClassName("com.zaxxer.hikari.mocks.StubDataSource");

      StubConnection.slowCreate = true;
      try (HikariDataSource ds = new HikariDataSource(config)) {

         HikariPool pool = getPool(ds);
         quietlySleep(1250);

         assertSame("Total connections not as expected", 1, pool.getTotalConnections());
         assertSame("Idle connections not as expected", 1, pool.getIdleConnections());

         // This will take the pool down to zero
         try (Connection connection = ds.getConnection()) {
            assertNotNull(connection);

            assertSame("Total connections not as expected", 1, pool.getTotalConnections());
            assertSame("Idle connections not as expected", 0, pool.getIdleConnections());

            PreparedStatement statement = connection.prepareStatement("SELECT some, thing FROM somewhere WHERE something=?");
            assertNotNull(statement);

            ResultSet resultSet = statement.executeQuery();
            assertNotNull(resultSet);

            try {
               statement.getMaxFieldSize();
               fail();
            }
            catch (Exception e) {
               assertSame(SQLException.class, e.getClass());
            }

            pool.logPoolState("testBackfill() before close...");

            // The connection will be ejected from the pool here
         }

         assertSame("Total connections not as expected", 0, pool.getTotalConnections());

         pool.logPoolState("testBackfill() after close...");

         quietlySleep(1250);

         assertSame("Total connections not as expected", 1, pool.getTotalConnections());
         assertSame("Idle connections not as expected", 1, pool.getIdleConnections());
      }
      finally {
         StubConnection.slowCreate = false;
      }
   }

   @Test
   public void testMaximumPoolLimit() throws Exception
   {
      HikariConfig config = newHikariConfig();
      config.setMinimumIdle(1);
      config.setMaximumPoolSize(4);
      config.setConnectionTimeout(20000);
      config.setInitializationFailTimeout(0);
      config.setConnectionTestQuery("VALUES 1");
      config.setDataSourceClassName("com.zaxxer.hikari.mocks.StubDataSource");

      final AtomicReference<Exception> ref = new AtomicReference<>();

      StubConnection.count.set(0); // reset counter

      try (final HikariDataSource ds = new HikariDataSource(config)) {

         final HikariPool pool = getPool(ds);

         Thread[] threads = new Thread[20];
         for (int i = 0; i < threads.length; i++) {
            threads[i] = new Thread(new Runnable() {
               @Override
               public void run()
               {
                  try {
                     pool.logPoolState("Before acquire ");
                     try (Connection connection = ds.getConnection()) {
                        pool.logPoolState("After  acquire ");
                        quietlySleep(500);
                     }
                  }
                  catch (Exception e) {
                     ref.set(e);
                  }
               }
            });
         }

         for (int i = 0; i < threads.length; i++) {
            threads[i].start();
         }

         for (int i = 0; i < threads.length; i++) {
            threads[i].join();
         }

         pool.logPoolState("before check ");
         assertNull((ref.get() != null ? ref.get().toString() : ""), ref.get());
         assertSame("StubConnection count not as expected", 4, StubConnection.count.get());
      }
   }

   @Test
   public void testOldDriver() throws Exception
   {
      HikariConfig config = newHikariConfig();
      config.setMinimumIdle(1);
      config.setMaximumPoolSize(1);
      config.setConnectionTimeout(2500);
      config.setConnectionTestQuery("VALUES 1");
      config.setDataSourceClassName("com.zaxxer.hikari.mocks.StubDataSource");

      StubConnection.oldDriver = true;
      StubStatement.oldDriver = true;
      try (HikariDataSource ds = new HikariDataSource(config)) {
         quietlySleep(500);

         try (Connection connection = ds.getConnection()) {
            // close
         }

         quietlySleep(500);
         try (Connection connection = ds.getConnection()) {
            // close
         }
      }
      finally {
         StubConnection.oldDriver = false;
         StubStatement.oldDriver = false;
      }
   }

   @Test
   public void testSuspendResume() throws Exception
   {
      HikariConfig config = newHikariConfig();
      config.setMinimumIdle(3);
      config.setMaximumPoolSize(3);
      config.setConnectionTimeout(2500);
      config.setAllowPoolSuspension(true);
      config.setConnectionTestQuery("VALUES 1");
      config.setDataSourceClassName("com.zaxxer.hikari.mocks.StubDataSource");

      try (final HikariDataSource ds = new HikariDataSource(config)) {
         HikariPool pool = getPool(ds);
         while (pool.getTotalConnections() < 3) {
            quietlySleep(50);
         }

         Thread t = new Thread(new Runnable() {
            @Override
            public void run()
            {
               try {
                  ds.getConnection();
                  ds.getConnection();
               }
               catch (Exception e) {
                  fail();
               }
            }
         });

         try (Connection c3 = ds.getConnection()) {
            assertEquals(2, pool.getIdleConnections());
   
            pool.suspendPool();
            t.start();
   
            quietlySleep(500);
            assertEquals(2, pool.getIdleConnections());
         }
         assertEquals(3, pool.getIdleConnections());
         pool.resumePool();
         quietlySleep(500);
         assertEquals(1, pool.getIdleConnections());
      }
   }

   @Test
   public void testInitializationFailure1() throws SQLException
   {
      StubDataSource stubDataSource = new StubDataSource();
      stubDataSource.setThrowException(new SQLException("Connection refused"));

      try (HikariDataSource ds = newHikariDataSource()) {
         ds.setMinimumIdle(1);
         ds.setMaximumPoolSize(1);
         ds.setConnectionTimeout(2500);
         ds.setConnectionTestQuery("VALUES 1");
         ds.setDataSource(stubDataSource);

         try (Connection c = ds.getConnection()) {
            fail("Initialization should have failed");
         }
         catch (SQLException e) {
            // passed
         }
      }
   }

   @Test
   public void testInitializationFailure2() throws SQLException
   {
      StubDataSource stubDataSource = new StubDataSource();
      stubDataSource.setThrowException(new SQLException("Connection refused"));

      HikariConfig config = newHikariConfig();
      config.setMinimumIdle(1);
      config.setConnectionTestQuery("VALUES 1");
      config.setDataSource(stubDataSource);

      try (HikariDataSource ds = new HikariDataSource(config);
           Connection c = ds.getConnection()) {
         fail("Initialization should have failed");
      }
      catch (PoolInitializationException e) {
         // passed
      }
   }

   @Test
   public void testInvalidConnectionTestQuery()
   {
      class BadConnection extends StubConnection {
         /** {@inheritDoc} */
         @Override
         public Statement createStatement() throws SQLException
         {
            throw new SQLException("Simulated exception in createStatement()");
         }
      }

      StubDataSource stubDataSource = new StubDataSource() {
         /** {@inheritDoc} */
         @Override
         public Connection getConnection() throws SQLException
         {
            return new BadConnection();
         }
      };

      HikariConfig config = newHikariConfig();
      config.setMinimumIdle(1);
      config.setMaximumPoolSize(2);
      config.setConnectionTimeout(TimeUnit.SECONDS.toMillis(3));
      config.setConnectionTestQuery("VALUES 1");
      config.setInitializationFailTimeout(TimeUnit.SECONDS.toMillis(2));
      config.setDataSource(stubDataSource);

      try (HikariDataSource ds = new HikariDataSource(config)) {
         try (Connection c = ds.getConnection()) {
            fail("getConnection() should have failed");
         }
         catch (SQLException e) {
            assertSame("Simulated exception in createStatement()", e.getNextException().getMessage());
         }
      }
      catch (PoolInitializationException e) {
         assertSame("Simulated exception in createStatement()", e.getCause().getMessage());
      }
      
      config.setInitializationFailTimeout(0);
      try (HikariDataSource ds = new HikariDataSource(config)) {
         fail("Initialization should have failed");
      }
      catch (PoolInitializationException e) {
         // passed
      }
   }

   @Test
   public void testPopulationSlowAcquisition() throws InterruptedException, SQLException
   {
      HikariConfig config = newHikariConfig();
      config.setMaximumPoolSize(20);
      config.setConnectionTestQuery("VALUES 1");
      config.setDataSourceClassName("com.zaxxer.hikari.mocks.StubDataSource");

      System.setProperty("com.zaxxer.hikari.housekeeping.periodMs", "1000");

      StubConnection.slowCreate = true;
      try (HikariDataSource ds = new HikariDataSource(config)) {
         System.clearProperty("com.zaxxer.hikari.housekeeping.periodMs");

         ds.setIdleTimeout(3000);

         SECONDS.sleep(2);

         HikariPool pool = getPool(ds);
         assertSame("Total connections not as expected", 2, pool.getTotalConnections());
         assertSame("Idle connections not as expected", 2, pool.getIdleConnections());

         try (Connection connection = ds.getConnection()) {
            assertNotNull(connection);
   
            SECONDS.sleep(20);
   
            assertSame("Second total connections not as expected", 20, pool.getTotalConnections());
            assertSame("Second idle connections not as expected", 19, pool.getIdleConnections());
         }

         assertSame("Idle connections not as expected", 20, pool.getIdleConnections());

         SECONDS.sleep(5);

         assertSame("Third total connections not as expected", 20, pool.getTotalConnections());
         assertSame("Third idle connections not as expected", 20, pool.getIdleConnections());
      }
      finally {
         StubConnection.slowCreate = false;
      }
   }

   @Test
   public void testMinimumIdleZero() throws SQLException
   {
      HikariConfig config = newHikariConfig();
      config.setMinimumIdle(0);
      config.setMaximumPoolSize(5);
      config.setConnectionTimeout(1000L);
      config.setConnectionTestQuery("VALUES 1");
      config.setDataSourceClassName("com.zaxxer.hikari.mocks.StubDataSource");

      try (HikariDataSource ds = new HikariDataSource(config);
           Connection connection = ds.getConnection()) {
         // passed
      }
      catch (SQLTransientConnectionException sqle) {
         fail("Failed to obtain connection");
      }
   }
}