/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you 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.datastax.driver.core;

import static com.datastax.driver.core.Assertions.assertThat;
import static com.datastax.driver.core.ConditionChecker.check;
import static com.datastax.driver.core.PoolingOptions.NEW_CONNECTION_THRESHOLD_LOCAL_KEY;
import static com.google.common.collect.Lists.newArrayList;
import static java.util.concurrent.TimeUnit.MILLISECONDS;
import static java.util.concurrent.TimeUnit.SECONDS;
import static org.mockito.Mockito.after;
import static org.mockito.Mockito.any;
import static org.mockito.Mockito.doReturn;
import static org.mockito.Mockito.eq;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.never;
import static org.mockito.Mockito.reset;
import static org.mockito.Mockito.spy;
import static org.mockito.Mockito.timeout;
import static org.mockito.Mockito.times;
import static org.mockito.Mockito.verify;
import static org.scassandra.http.client.ClosedConnectionReport.CloseType.CLOSE;
import static org.scassandra.http.client.PrimingRequest.queryBuilder;
import static org.scassandra.http.client.PrimingRequest.then;
import static org.scassandra.http.client.Result.server_error;
import static org.testng.Assert.fail;

import com.codahale.metrics.Gauge;
import com.datastax.driver.core.exceptions.BusyConnectionException;
import com.datastax.driver.core.exceptions.BusyPoolException;
import com.datastax.driver.core.exceptions.ConnectionException;
import com.datastax.driver.core.exceptions.DriverException;
import com.datastax.driver.core.exceptions.NoHostAvailableException;
import com.datastax.driver.core.policies.ConstantReconnectionPolicy;
import com.google.common.base.Function;
import com.google.common.base.Predicate;
import com.google.common.base.Throwables;
import com.google.common.util.concurrent.ForwardingListeningExecutorService;
import com.google.common.util.concurrent.FutureCallback;
import com.google.common.util.concurrent.ListenableFuture;
import com.google.common.util.concurrent.ListeningExecutorService;
import com.google.common.util.concurrent.Uninterruptibles;
import java.net.InetSocketAddress;
import java.util.Collection;
import java.util.Collections;
import java.util.Iterator;
import java.util.List;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.Semaphore;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicReference;
import org.scassandra.cql.PrimitiveType;
import org.scassandra.http.client.PrimingRequest;
import org.testng.annotations.BeforeClass;
import org.testng.annotations.Test;

public class HostConnectionPoolTest extends ScassandraTestBase.PerClassCluster {

  static final Integer NEW_CONNECTION_THRESHOLD =
      PoolingOptions.DEFAULTS.get(ProtocolVersion.V1).get(NEW_CONNECTION_THRESHOLD_LOCAL_KEY);

  @BeforeClass(groups = {"short", "long"})
  public void reinitializeCluster() {
    // Don't use the provided cluster, each test will create its own instead.
    cluster.close();
  }

  /**
   * Ensure the given pool has the given size within 5 seconds.
   *
   * @param pool Pool to check.
   * @param expectedSize Expected size of pool.
   */
  private void assertPoolSize(HostConnectionPool pool, final int expectedSize) {
    check()
        .before(5, TimeUnit.SECONDS)
        .that(
            pool,
            new Predicate<HostConnectionPool>() {
              @Override
              public boolean apply(HostConnectionPool input) {
                return input.connections.size() == expectedSize;
              }
            })
        .becomesTrue();
  }

  private void assertBorrowedConnections(
      Iterable<MockRequest> requests, List<Connection> expectedConnections) {
    for (MockRequest request : requests) {
      assertThat(expectedConnections).contains(request.getConnection());
    }
  }

  private void assertBorrowedConnection(
      Iterable<MockRequest> requests, Connection expectedConnection) {
    assertBorrowedConnections(requests, Collections.singletonList(expectedConnection));
  }

  /**
   * Ensures that if a fixed-sized pool has filled its core connections and reached its maximum
   * number of enqueued requests, then borrowConnection will fail instead of creating a new
   * connection.
   *
   * @jira_ticket JAVA-419
   * @test_category connection:connection_pool
   * @since 2.0.10, 2.1.6
   */
  @Test(groups = "short")
  public void fixed_size_pool_should_fill_its_core_connections_and_queue_and_then_reject() {
    Cluster cluster = createClusterBuilder().build();
    List<MockRequest> allRequests = newArrayList();
    try {
      HostConnectionPool pool = createPool(cluster, 2, 2);
      int maxQueueSize = 256;

      assertThat(pool.connections.size()).isEqualTo(2);
      List<Connection> coreConnections = newArrayList(pool.connections);
      // fill connections
      List<MockRequest> requests = MockRequest.sendMany(2 * 128, pool);
      assertBorrowedConnections(requests, coreConnections);
      allRequests.addAll(requests);
      // fill queue
      allRequests.addAll(MockRequest.sendMany(maxQueueSize, pool, maxQueueSize));

      // add one more request, it should be rejected because the queue is full
      MockRequest failedBorrow = MockRequest.send(pool, maxQueueSize);
      try {
        failedBorrow.getConnection();
        fail("Expected a BusyPoolException");
      } catch (BusyPoolException e) {
        assertThat(e).hasMessageContaining("reached its max size");
      }
    } finally {
      MockRequest.completeAll(allRequests);
      cluster.close();
    }
  }

  /**
   * Ensures that if a fixed-sized pool has filled its core connections and reached a number of
   * requests to cause it to be enqueued, that if the request is not serviced within 100ms, a
   * BusyPoolException is raised with a timeout.
   *
   * @jira_ticket JAVA-1371
   * @test_category connection:connection_pool
   * @since 3.0.7 3.1.4 3.2.0
   */
  @Test(groups = "short")
  public void should_reject_if_enqueued_and_timeout_reached() {
    Cluster cluster = createClusterBuilder().build();
    List<MockRequest> allRequests = newArrayList();
    try {
      HostConnectionPool pool = createPool(cluster, 1, 1);
      List<MockRequest> requests = MockRequest.sendMany(128, pool);
      allRequests.addAll(requests);

      // pool is now full, this request will be enqueued
      MockRequest failedBorrow = MockRequest.send(pool, 100, 128);
      try {
        failedBorrow.getConnection();
        fail("Expected a BusyPoolException");
      } catch (BusyPoolException e) {
        assertThat(e).hasMessageContaining("timed out");
      }
    } finally {
      MockRequest.completeAll(allRequests);
      cluster.close();
    }
  }

  /**
   * Validates that if a borrow request is enqueued into a pool for a Host that is currently within
   * the window of reconnecting after an error that the future tied to that query times out after
   * {@link PoolingOptions#setPoolTimeoutMillis(int)}.
   *
   * <p>This primarily serves to demonstrate that the use case described in JAVA-1371 is fixed by
   * using {@link Session#execute(Statement)} shortly after a server error on a connection.
   *
   * @jira_ticket JAVA-1371
   * @test_category connection:connection_pool
   * @since 3.0.7 3.1.4 3.2.0
   */
  @Test(groups = "short")
  public void should_not_hang_when_executing_sync_queries() {
    primingClient.prime(
        queryBuilder()
            .withQuery("server_error query")
            .withThen(then().withResult(server_error))
            .build());

    Cluster cluster =
        createClusterBuilder()
            .withReconnectionPolicy(new ConstantReconnectionPolicy(10000))
            .build();
    // reduce timeout so test runs faster.
    cluster.getConfiguration().getPoolingOptions().setPoolTimeoutMillis(500);
    try {
      Session session = cluster.connect();
      try {
        session.execute("server_error query");
        fail("Exception expected");
      } catch (Exception e) {
        // error is expected in this case.
      }

      try {
        session.execute("this should not block indefinitely");
      } catch (NoHostAvailableException nhae) {
        // should raise a NHAE with a BusyPoolException.
        Collection<Throwable> errors = nhae.getErrors().values();
        assertThat(errors).hasSize(1);
        Throwable e = errors.iterator().next();
        assertThat(e).isInstanceOf(BusyPoolException.class);
        assertThat(e).hasMessageContaining("timed out");
      }
    } finally {
      cluster.close();
    }
  }

  /**
   * Ensures that any enqueued connection borrow requests are failed when their associated
   * connection pool closes.
   *
   * @jira_ticket JAVA-839
   * @test_category connection:connection_pool
   * @since 3.0.4, 3.1.1
   */
  @Test(groups = "short")
  public void requests_with_enqueued_borrow_requests_should_be_failed_when_pool_closes() {
    Cluster cluster = createClusterBuilder().build();
    List<MockRequest> requests = newArrayList();
    try {
      HostConnectionPool pool = createPool(cluster, 2, 2);
      int maxQueueSize = 256;

      assertThat(pool.connections.size()).isEqualTo(2);
      List<Connection> coreConnections = newArrayList(pool.connections);
      // fill connections
      requests = MockRequest.sendMany(2 * 128, pool);
      assertBorrowedConnections(requests, coreConnections);
      // fill queue
      List<MockRequest> queuedRequests = MockRequest.sendMany(maxQueueSize, pool, maxQueueSize);

      // Closing the pool should fail all queued connections.
      pool.closeAsync();

      for (MockRequest queuedRequest : queuedRequests) {
        // Future should be completed.
        assertThat(queuedRequest.connectionFuture.isDone()).isTrue();
        try {
          // Future should fail as result of the pool closing.
          queuedRequest.getConnection();
          fail("Expected a ConnectionException");
        } catch (ConnectionException e) {
          /*expected*/
        }
      }
    } finally {
      MockRequest.completeAll(requests);
      cluster.close();
    }
  }

  /**
   * Validates that if the keyspace tied to the Session's pool state is different than the keyspace
   * on the connection being used in dequeue that {@link Connection#setKeyspaceAsync(String)} is set
   * on that connection and that "USE keyspace" is only called once since setKeyspaceAsync should
   * not attempt setting the keyspace if there is already a request inflight that is doing this.
   *
   * @jira_ticket JAVA-839
   * @test_category connection:connection_pool
   * @since 3.0.4, 3.1.1
   */
  @Test(groups = "short")
  public void should_adjust_connection_keyspace_on_dequeue_if_pool_state_is_different()
      throws TimeoutException, ExecutionException {
    Cluster cluster = createClusterBuilder().build();
    List<MockRequest> requests = newArrayList();
    try {
      HostConnectionPool pool = createPool(cluster, 1, 1);
      int maxQueueSize = 256;

      assertThat(pool.connections.size()).isEqualTo(1);
      List<Connection> coreConnections = newArrayList(pool.connections);
      // fill connections
      requests = MockRequest.sendMany(128, pool);
      assertBorrowedConnections(requests, coreConnections);
      // fill queue
      List<MockRequest> queuedRequests = MockRequest.sendMany(maxQueueSize, pool, maxQueueSize);

      // Wait for connections to be borrowed before changing keyspace.
      for (MockRequest request : requests) {
        Uninterruptibles.getUninterruptibly(request.connectionFuture, 5, TimeUnit.SECONDS);
      }
      // Simulate change of keyspace on pool.  Prime a delay so existing requests can complete
      // beforehand.
      primingClient.prime(
          PrimingRequest.queryBuilder()
              .withQuery("USE \"newkeyspace\"")
              .withThen(PrimingRequest.then().withFixedDelay(2000L)));
      pool.manager.poolsState.setKeyspace("newkeyspace");

      // Complete all requests, this should cause dequeue on connection.
      MockRequest.completeAll(requests);

      // Check the status on queued request's connection futures.  We expect that dequeue should
      // be called when connection is released by previous requests completing, and that one set
      // keyspace attempt should be tried.

      int count = 0;
      for (MockRequest queuedRequest : queuedRequests) {
        try {
          Uninterruptibles.getUninterruptibly(queuedRequest.connectionFuture, 10, TimeUnit.SECONDS);
          count++;
        } catch (ExecutionException e) {
          // 128th request should timeout since all in flight requests are used.
          assertThat(e.getCause())
              .isInstanceOf(BusyPoolException.class)
              .hasMessageContaining("timed out after");
          assertThat(count).isEqualTo(128);
          break;
        }
      }

      // We should only have gotten one 'USE newkeyspace' query since Connection#setKeyspaceAsync
      // should only do
      // this once if there is already a request in flight.
      assertThat(activityClient.retrieveQueries())
          .extractingResultOf("getQuery")
          .containsOnlyOnce("USE \"newkeyspace\"");
    } finally {
      MockRequest.completeAll(requests);
      cluster.close();
    }
  }

  /**
   * Ensures that on borrowConnection if a set keyspace attempt is in progress on that connection
   * for a different keyspace than the pool state that the borrowConnection future returned is
   * failed.
   *
   * @jira_ticket JAVA-839
   * @test_category connection:connection_pool
   * @since 3.0.4, 3.1.1
   */
  @Test(groups = "short")
  public void
      should_fail_in_borrowConnection_when_setting_keyspace_and_another_set_keyspace_attempt_is_in_flight()
          throws TimeoutException {
    Cluster cluster = createClusterBuilder().build();
    try {
      HostConnectionPool pool = createPool(cluster, 1, 1);

      // Respond to setting as slowks very slowly.
      primingClient.prime(
          PrimingRequest.queryBuilder()
              .withQuery("USE \"slowks\"")
              .withThen(PrimingRequest.then().withFixedDelay(5000L)));

      Connection connection = pool.connections.get(0);

      connection.setKeyspaceAsync("slowks");

      // Simulate change of keyspace on pool.
      pool.manager.poolsState.setKeyspace("newks");

      MockRequest request = MockRequest.send(pool);

      try {
        Uninterruptibles.getUninterruptibly(request.connectionFuture, 5, TimeUnit.SECONDS);
        fail("Should have thrown exception");
      } catch (ExecutionException e) {
        assertThat(e.getCause()).isInstanceOf(DriverException.class);
        assertThat(e.getCause().getMessage())
            .contains(
                "Aborting attempt to set keyspace to 'newks' since there is already an in flight attempt to set keyspace to 'slowks'.");
      }
    } finally {
      cluster.close();
    }
  }

  /**
   * Ensures that while dequeuing borrow connection requests that if a set keyspace attempt is in
   * progress on that connection for a difference keyspace than the pool state that the future for
   * that borrow attempt is failed.
   *
   * @jira_ticket JAVA-839
   * @test_category connection:connection_pool
   * @since 3.0.4, 3.1.1
   */
  @Test(groups = "short")
  public void
      should_fail_in_dequeue_when_setting_keyspace_and_another_set_keyspace_attempt_is_in_flight()
          throws ExecutionException, TimeoutException {
    Cluster cluster = createClusterBuilder().build();
    List<MockRequest> requests = newArrayList();
    try {
      HostConnectionPool pool = createPool(cluster, 1, 1);
      // Limit requests per connection to 100 so we don't exhaust stream ids.
      cluster
          .getConfiguration()
          .getPoolingOptions()
          .setMaxRequestsPerConnection(HostDistance.LOCAL, 100);
      int maxQueueSize = 256;

      assertThat(pool.connections.size()).isEqualTo(1);
      List<Connection> coreConnections = newArrayList(pool.connections);

      // fill connections
      requests = MockRequest.sendMany(100, pool);
      assertBorrowedConnections(requests, coreConnections);

      // send a request that will be queued.
      MockRequest queuedRequest = MockRequest.send(pool, maxQueueSize);

      // Wait for connections to be borrowed before changing keyspace.
      for (MockRequest request : requests) {
        Uninterruptibles.getUninterruptibly(request.connectionFuture, 5, TimeUnit.SECONDS);
      }

      // Respond to setting as slowks very slowly.
      primingClient.prime(
          PrimingRequest.queryBuilder()
              .withQuery("USE \"slowks\"")
              .withThen(PrimingRequest.then().withFixedDelay(5000L)));
      Connection connection = pool.connections.get(0);
      connection.setKeyspaceAsync("slowks");

      // Simulate change of keyspace on pool.
      pool.manager.poolsState.setKeyspace("newkeyspace");

      // Complete all requests, this should cause dequeue on connection.
      MockRequest.completeAll(requests);

      try {
        Uninterruptibles.getUninterruptibly(queuedRequest.connectionFuture, 5, TimeUnit.SECONDS);
        fail("Should have thrown exception");
      } catch (ExecutionException e) {
        assertThat(e.getCause()).isInstanceOf(DriverException.class);
        assertThat(e.getCause().getMessage())
            .contains(
                "Aborting attempt to set keyspace to 'newkeyspace' since there is already an in flight attempt to set keyspace to 'slowks'.");
      }
    } finally {
      MockRequest.completeAll(requests);
      cluster.close();
    }
  }

  /**
   * Ensures that if a variable-sized pool has filled up to its maximum connections that
   * borrowConnection will return a failed future instead of creating a new connection.
   *
   * @jira_ticket JAVA-419
   * @test_category connection:connection_pool
   * @since 2.0.10, 2.1.6
   */
  @Test(groups = "short")
  public void variable_size_pool_should_fill_its_connections_and_then_reject() throws Exception {
    Cluster cluster = createClusterBuilder().build();
    List<MockRequest> allRequests = newArrayList();
    try {
      HostConnectionPool pool = createPool(cluster, 1, 2);
      Connection.Factory factory = spy(cluster.manager.connectionFactory);
      cluster.manager.connectionFactory = factory;

      assertThat(pool.connections.size()).isEqualTo(1);
      Connection coreConnection = pool.connections.get(0);

      // Fill enough connections to hit the threshold.
      List<MockRequest> requests = MockRequest.sendMany(NEW_CONNECTION_THRESHOLD, pool);
      assertBorrowedConnection(requests, coreConnection);
      allRequests.addAll(requests);
      allRequests.add(MockRequest.send(pool));

      // Allow time for new connection to be spawned.
      verify(factory, after(2000).times(1)).open(any(HostConnectionPool.class));
      assertPoolSize(pool, 2);

      // Borrow more and ensure the connection returned is a non-core connection.
      for (int i = 0; i < NEW_CONNECTION_THRESHOLD; i++) {
        MockRequest request = MockRequest.send(pool);
        assertThat(request.getConnection()).isNotEqualTo(coreConnection);
        allRequests.add(request);
      }

      // Fill remaining connections (28 + 28) - 1
      allRequests.addAll(MockRequest.sendMany(55, pool));

      MockRequest failedBorrow = MockRequest.send(pool);
      try {
        failedBorrow.getConnection();
        fail("Expected a BusyPoolException");
      } catch (BusyPoolException e) {
        /*expected*/
      }
    } finally {
      MockRequest.completeAll(allRequests);
      cluster.close();
    }
  }

  /**
   * Ensures that if the core connection pool is full that borrowConnection will create and use a
   * new connection.
   *
   * @jira_ticket JAVA-419
   * @test_category connection:connection_pool
   * @since 2.0.10, 2.1.6
   */
  @Test(groups = "short")
  public void should_add_extra_connection_when_core_full() throws Exception {
    Cluster cluster = createClusterBuilder().build();
    List<MockRequest> allRequests = newArrayList();
    try {
      HostConnectionPool pool = createPool(cluster, 1, 2);
      Connection.Factory factory = spy(cluster.manager.connectionFactory);
      cluster.manager.connectionFactory = factory;
      Connection core = pool.connections.get(0);

      // Fill core connection + 1
      List<MockRequest> requests = MockRequest.sendMany(NEW_CONNECTION_THRESHOLD, pool);
      assertBorrowedConnection(requests, core);
      allRequests.addAll(requests);
      allRequests.add(MockRequest.send(pool));

      // Reaching the threshold should have triggered the creation of an extra one
      verify(factory, after(2000).times(1)).open(any(HostConnectionPool.class));
      assertPoolSize(pool, 2);
    } finally {
      MockRequest.completeAll(allRequests);
      cluster.close();
    }
  }

  /**
   * Ensures that a trashed connection that has not been timed out should be resurrected into the
   * connection pool if borrowConnection is called and a new connection is needed.
   *
   * @jira_ticket JAVA-419
   * @test_category connection:connection_pool
   * @since 2.0.10, 2.1.6
   */
  @Test(groups = "long")
  public void should_resurrect_trashed_connection_within_idle_timeout() throws Exception {
    Cluster cluster =
        createClusterBuilder()
            .withPoolingOptions(new PoolingOptions().setIdleTimeoutSeconds(20))
            .build();
    List<MockRequest> allRequests = newArrayList();
    try {
      HostConnectionPool pool = createPool(cluster, 1, 2);
      Connection.Factory factory = spy(cluster.manager.connectionFactory);
      cluster.manager.connectionFactory = factory;
      Connection connection1 = pool.connections.get(0);

      List<MockRequest> requests = MockRequest.sendMany(NEW_CONNECTION_THRESHOLD, pool);
      assertBorrowedConnections(requests, Collections.singletonList(connection1));
      allRequests.addAll(requests);
      allRequests.add(MockRequest.send(pool));

      verify(factory, after(2000).times(1)).open(any(HostConnectionPool.class));
      assertPoolSize(pool, 2);
      Connection connection2 = pool.connections.get(1);

      assertThat(connection1.inFlight.get()).isEqualTo(101);
      assertThat(connection2.inFlight.get()).isEqualTo(0);

      // Go back under the capacity of 1 connection
      MockRequest.completeMany(51, allRequests);

      assertThat(connection1.inFlight.get()).isEqualTo(50);
      assertThat(connection2.inFlight.get()).isEqualTo(0);

      // Given enough time, one connection gets trashed (and the implementation picks the first one)
      Uninterruptibles.sleepUninterruptibly(20, TimeUnit.SECONDS);
      assertThat(pool.connections).containsExactly(connection2);
      assertThat(pool.trash).containsExactly(connection1);

      // Now borrow enough to go just under the 1 connection threshold
      allRequests.addAll(MockRequest.sendMany(50, pool));

      assertThat(pool.connections).containsExactly(connection2);
      assertThat(pool.trash).containsExactly(connection1);
      assertThat(connection1.inFlight.get()).isEqualTo(50);
      assertThat(connection2.inFlight.get()).isEqualTo(50);

      // Borrowing one more time should resurrect the trashed connection
      allRequests.addAll(MockRequest.sendMany(1, pool));
      verify(factory, after(2000).times(1)).open(any(HostConnectionPool.class));
      assertPoolSize(pool, 2);

      assertThat(pool.connections).containsExactly(connection2, connection1);
      assertThat(pool.trash).isEmpty();
      assertThat(connection1.inFlight.get()).isEqualTo(50);
      assertThat(connection2.inFlight.get()).isEqualTo(51);
    } finally {
      MockRequest.completeAll(allRequests);
      cluster.close();
    }
  }

  /**
   * Ensures that a trashed connection that has been timed out should not be resurrected into the
   * connection pool if borrowConnection is called and a new connection is needed.
   *
   * @jira_ticket JAVA-419
   * @test_category connection:connection_pool
   * @since 2.0.10, 2.1.6
   */
  @Test(groups = "long")
  public void should_not_resurrect_trashed_connection_after_idle_timeout() throws Exception {
    Cluster cluster =
        createClusterBuilder()
            .withPoolingOptions(new PoolingOptions().setIdleTimeoutSeconds(20))
            .build();
    List<MockRequest> allRequests = newArrayList();
    try {
      HostConnectionPool pool = createPool(cluster, 1, 2);
      Connection.Factory factory = spy(cluster.manager.connectionFactory);
      cluster.manager.connectionFactory = factory;
      Connection connection1 = pool.connections.get(0);

      List<MockRequest> requests = MockRequest.sendMany(NEW_CONNECTION_THRESHOLD, pool);
      assertBorrowedConnection(requests, connection1);
      allRequests.addAll(requests);
      allRequests.add(MockRequest.send(pool));

      verify(factory, after(2000).times(1)).open(any(HostConnectionPool.class));
      assertPoolSize(pool, 2);
      reset(factory);
      Connection connection2 = pool.connections.get(1);

      assertThat(connection1.inFlight.get()).isEqualTo(101);
      assertThat(connection2.inFlight.get()).isEqualTo(0);

      // Go back under the capacity of 1 connection
      MockRequest.completeMany(51, allRequests);

      assertThat(connection1.inFlight.get()).isEqualTo(50);
      assertThat(connection2.inFlight.get()).isEqualTo(0);

      // Given enough time, one connection gets trashed (and the implementation picks the first one)
      Uninterruptibles.sleepUninterruptibly(20, TimeUnit.SECONDS);
      assertThat(pool.connections).containsExactly(connection2);
      assertThat(pool.trash).containsExactly(connection1);

      // Return trashed connection down to 0 inFlight
      MockRequest.completeMany(50, allRequests);
      assertThat(connection1.inFlight.get()).isEqualTo(0);

      // Give enough time for trashed connection to be cleaned up from the trash:
      Uninterruptibles.sleepUninterruptibly(30, TimeUnit.SECONDS);
      assertThat(pool.connections).containsExactly(connection2);
      assertThat(pool.trash).isEmpty();
      assertThat(connection1.isClosed()).isTrue();

      // Fill the live connection to go over the threshold where a second one is needed
      requests = MockRequest.sendMany(NEW_CONNECTION_THRESHOLD, pool);
      assertBorrowedConnection(requests, connection2);
      allRequests.addAll(requests);
      allRequests.add(MockRequest.send(pool));
      assertThat(connection2.inFlight.get()).isEqualTo(101);
      verify(factory, after(2000).times(1)).open(any(HostConnectionPool.class));
      assertPoolSize(pool, 2);

      // Borrow again to get the new connection
      MockRequest request = MockRequest.send(pool);
      allRequests.add(request);
      assertThat(request.getConnection())
          .isNotEqualTo(connection2) // should not be the full connection
          .isNotEqualTo(connection1); // should not be the previously trashed one
    } finally {
      MockRequest.completeAll(allRequests);
      cluster.close();
    }
  }

  /**
   * Ensures that a trashed connection that has been timed out should not be closed until it has 0
   * in flight requests.
   *
   * @jira_ticket JAVA-419
   * @test_category connection:connection_pool
   * @since 2.0.10, 2.1.6
   */
  @Test(groups = "long")
  public void should_not_close_trashed_connection_until_no_in_flight() throws Exception {
    Cluster cluster =
        createClusterBuilder()
            .withPoolingOptions(new PoolingOptions().setIdleTimeoutSeconds(20))
            .build();
    List<MockRequest> allRequests = newArrayList();

    try {
      HostConnectionPool pool = createPool(cluster, 1, 2);
      Connection.Factory factory = spy(cluster.manager.connectionFactory);
      cluster.manager.connectionFactory = factory;
      Connection connection1 = pool.connections.get(0);

      // Fill core connection enough to trigger creation of another one
      List<MockRequest> requests = MockRequest.sendMany(NEW_CONNECTION_THRESHOLD, pool);
      assertBorrowedConnections(requests, Collections.singletonList(connection1));
      allRequests.addAll(requests);
      allRequests.add(MockRequest.send(pool));

      verify(factory, after(2000).times(1)).open(any(HostConnectionPool.class));
      assertThat(pool.connections).hasSize(2);

      // Return enough times to get back under the threshold where one connection is enough
      MockRequest.completeMany(50, allRequests);

      // Give enough time for one connection to be trashed. Due to the implementation, this will be
      // the first one.
      // It still has in-flight requests so should not get closed.
      Uninterruptibles.sleepUninterruptibly(30, TimeUnit.SECONDS);
      assertThat(pool.trash).containsExactly(connection1);
      assertThat(connection1.inFlight.get()).isEqualTo(51);
      assertThat(connection1.isClosed()).isFalse();

      // Consume all inFlight requests on the trashed connection.
      MockRequest.completeMany(51, allRequests);

      // Sleep enough time for the connection to be consider idled and closed.
      Uninterruptibles.sleepUninterruptibly(30, TimeUnit.SECONDS);

      // The connection should be now closed.
      // The trashed connection should be closed and not in the pool or trash.
      assertThat(connection1.isClosed()).isTrue();
      assertThat(pool.connections).doesNotContain(connection1);
      assertThat(pool.trash).doesNotContain(connection1);
    } finally {
      MockRequest.completeAll(allRequests);
      cluster.close();
    }
  }

  /**
   * Ensures that if a connection that has less than the minimum available stream ids is returned to
   * the pool that the connection is put in the trash.
   *
   * @jira_ticket JAVA-419
   * @test_category connection:connection_pool
   * @since 2.0.10, 2.1.6
   */
  @Test(groups = "short")
  public void should_trash_on_returning_connection_with_insufficient_streams() throws Exception {
    Cluster cluster = createClusterBuilder().build();
    List<MockRequest> allRequests = newArrayList();
    try {
      HostConnectionPool pool = createPool(cluster, 1, 2);
      Connection.Factory factory = spy(cluster.manager.connectionFactory);
      cluster.manager.connectionFactory = factory;
      Connection core = pool.connections.get(0);

      List<MockRequest> requests = MockRequest.sendMany(NEW_CONNECTION_THRESHOLD, pool);
      assertBorrowedConnections(requests, Collections.singletonList(core));
      allRequests.addAll(requests);
      allRequests.add(MockRequest.send(pool));

      verify(factory, after(2000).times(1)).open(any(HostConnectionPool.class));
      assertThat(pool.connections).hasSize(2);

      // Grab the new non-core connection and replace it with a spy.
      Connection extra1 = spy(pool.connections.get(1));
      pool.connections.set(1, extra1);

      // Borrow 10 times to ensure pool is utilized.
      allRequests.addAll(MockRequest.sendMany(10, pool));
      assertThat(pool.connections).hasSize(2);

      // stub the maxAvailableStreams method to return 0, indicating there are no remaining streams.
      // this should cause the connection to be replaced and trashed on returnConnection.
      doReturn(0).when(extra1).maxAvailableStreams();

      // On returning of the connection, should detect that there are no available streams and trash
      // it.
      assertThat(pool.trash).hasSize(0);
      pool.returnConnection(extra1, false);
      assertThat(pool.trash).hasSize(1);
    } finally {
      MockRequest.completeAll(allRequests);
      cluster.close();
    }
  }

  /**
   * Ensures that if a connection on a host is lost but other connections remain intact in the Pool
   * that the host is not marked down.
   *
   * @jira_ticket JAVA-544
   * @test_category connection:connection_pool
   * @since 2.0.11
   */
  @Test(groups = "short")
  public void should_keep_host_up_when_one_connection_lost() throws Exception {
    Cluster cluster = createClusterBuilder().build();
    try {
      HostConnectionPool pool = createPool(cluster, 2, 2);
      Connection core0 = pool.connections.get(0);
      Connection core1 = pool.connections.get(1);

      // Drop a connection and ensure the host stays up.
      currentClient.disableListener();
      currentClient.closeConnection(CLOSE, ((InetSocketAddress) core0.channel.localAddress()));
      Uninterruptibles.sleepUninterruptibly(100, TimeUnit.MILLISECONDS);

      // connection 0 should be down, while connection 1 and the Host should remain up.
      assertThat(core0.isClosed()).isTrue();
      assertThat(core1.isClosed()).isFalse();
      assertThat(pool.connections).doesNotContain(core0);
      assertThat(cluster).host(1).hasState(Host.State.UP);
      assertThat(cluster).hasOpenControlConnection();
    } finally {
      cluster.close();
    }
  }

  /**
   * Ensures that if all connections on a host are closed that the host is marked down and the
   * control connection is notified of that fact and re-established itself.
   *
   * @jira_ticket JAVA-544
   * @test_category connection:connection_pool
   * @since 2.0.11
   */
  @Test(groups = "short")
  public void should_mark_host_down_when_no_connections_remaining() throws Exception {
    int readTimeout = 1000;
    int reconnectInterval = 1000;
    Cluster cluster =
        this.createClusterBuilder()
            .withSocketOptions(
                new SocketOptions()
                    .setConnectTimeoutMillis(readTimeout)
                    .setReadTimeoutMillis(reconnectInterval))
            .withReconnectionPolicy(new ConstantReconnectionPolicy(1000))
            .build();
    try {
      cluster.init();

      Connection.Factory factory = spy(cluster.manager.connectionFactory);
      cluster.manager.connectionFactory = factory;

      HostConnectionPool pool = createPool(cluster, 8, 8);
      // copy list to track these connections.
      List<Connection> connections = newArrayList(pool.connections);

      reset(factory);

      // Drop all connections.
      currentClient.disableListener();
      currentClient.closeConnections(CLOSE);

      // The host should be marked down and the control connection closed.
      assertThat(cluster).host(1).goesDownWithin(10, TimeUnit.SECONDS);
      assertThat(cluster).hasClosedControlConnection();

      // Ensure all connections are closed.
      for (Connection connection : connections) {
        assertThat(connection.isClosed()).isTrue();
      }

      // Expect a reconnect attempt on host after reconnect interval
      // on behalf of the control connection.
      verify(factory, timeout(reconnectInterval * 2).atLeastOnce()).open(host);

      // Sleep for a bit to allow reconnect to fail.
      Uninterruptibles.sleepUninterruptibly(readTimeout * 2, TimeUnit.MILLISECONDS);

      // Ensure control connection is still closed.
      assertThat(cluster).hasClosedControlConnection();

      // Reenable connectivity.
      currentClient.enableListener();

      // Reconnect attempt should have been connected for control connection
      // and pool.
      // 2 attempts for connection.open (reconnect control connection and initial connection for
      // host state).
      verify(factory, after(reconnectInterval * 2).atLeast(2)).open(host);
      // 7 attempts for core connections after first initial connection.
      verify(factory, timeout(reconnectInterval * 2))
          .newConnections(any(HostConnectionPool.class), eq(7));

      // Wait some reasonable amount of time for connection to reestablish.
      Uninterruptibles.sleepUninterruptibly(readTimeout, TimeUnit.MILLISECONDS);

      // Control Connection should now be open.
      assertThat(cluster).hasOpenControlConnection();
      assertThat(cluster).host(1).hasState(Host.State.UP);
    } finally {
      cluster.close();
    }
  }

  /**
   * Ensures that if a connection on a host is lost that brings the number of active connections in
   * a pool under core connection count that up to core connections are re-established, but only
   * after the next reconnect schedule has elapsed.
   *
   * @jira_ticket JAVA-544
   * @test_category connection:connection_pool
   * @since 2.0.11
   */
  @Test(groups = "short")
  public void should_create_new_connections_when_connection_lost_and_under_core_connections()
      throws Exception {
    int readTimeout = 1000;
    int reconnectInterval = 1000;
    Cluster cluster =
        this.createClusterBuilder()
            .withSocketOptions(
                new SocketOptions()
                    .setConnectTimeoutMillis(readTimeout)
                    .setReadTimeoutMillis(reconnectInterval))
            .withReconnectionPolicy(new ConstantReconnectionPolicy(1000))
            .build();
    List<MockRequest> allRequests = newArrayList();
    try {
      cluster.init();

      Connection.Factory factory = spy(cluster.manager.connectionFactory);
      cluster.manager.connectionFactory = factory;

      TestExecutorService blockingExecutor =
          new TestExecutorService(cluster.manager.blockingExecutor);
      cluster.manager.blockingExecutor = blockingExecutor;

      HostConnectionPool pool = createPool(cluster, 3, 3);
      Connection core0 = pool.connections.get(0);
      Connection core1 = pool.connections.get(1);
      Connection core2 = pool.connections.get(2);

      // Drop two core connections.
      // Disable new connections initially and we'll eventually reenable it.
      currentClient.disableListener();
      currentClient.closeConnection(CLOSE, ((InetSocketAddress) core0.channel.localAddress()));
      currentClient.closeConnection(CLOSE, ((InetSocketAddress) core2.channel.localAddress()));
      Uninterruptibles.sleepUninterruptibly(100, TimeUnit.MILLISECONDS);

      // Since we have a connection left the host should remain up.
      assertThat(cluster).host(1).hasState(Host.State.UP);
      assertThat(pool.connections).hasSize(1);

      // The borrowed connection should be the open one.
      MockRequest request = MockRequest.send(pool);
      allRequests.add(request);
      assertThat(request.getConnection()).isEqualTo(core1);

      // Should not have tried to create a new core connection since reconnection time had not
      // elapsed.
      verify(factory, never()).open(any(HostConnectionPool.class));

      // Sleep to elapse the Reconnection Policy.
      Uninterruptibles.sleepUninterruptibly(reconnectInterval, TimeUnit.MILLISECONDS);

      // Attempt to borrow connection, this should trigger ensureCoreConnections thus spawning a new
      // connection.
      blockingExecutor.reset();
      request = MockRequest.send(pool);
      allRequests.add(request);
      assertThat(request.getConnection()).isEqualTo(core1);

      // Should have tried to open up to core connections as result of borrowing a connection past
      // reconnect time and not being at core.
      blockingExecutor.blockUntilNextTaskCompleted();
      verify(factory).open(any(HostConnectionPool.class));
      reset(factory);

      // Sleep for reconnect interval to allow reconnection time to elapse.
      Uninterruptibles.sleepUninterruptibly(
          (readTimeout + reconnectInterval) * 2, TimeUnit.MILLISECONDS);

      // Enable listening so new connections succeed.
      currentClient.enableListener();
      // Sleep to elapse the Reconnection Policy.
      Uninterruptibles.sleepUninterruptibly(reconnectInterval, TimeUnit.MILLISECONDS);

      // Try to borrow a connection, the pool should grow.
      blockingExecutor.reset();
      allRequests.add(MockRequest.send(pool));
      blockingExecutor.blockUntilNextTaskCompleted();
      verify(factory).open(any(HostConnectionPool.class));
      reset(factory);

      // Another core connection should be opened as result of another request to get us up to core
      // connections.
      blockingExecutor.reset();
      allRequests.add(MockRequest.send(pool));
      blockingExecutor.blockUntilNextTaskCompleted();
      verify(factory).open(any(HostConnectionPool.class));
      reset(factory);

      // Sending another request should not grow the pool any more, since we are now at core
      // connections.
      allRequests.add(MockRequest.send(pool));
      verify(factory, after((reconnectInterval + readTimeout) * 2).never())
          .open(any(HostConnectionPool.class));
    } finally {
      MockRequest.completeAll(allRequests);
      cluster.close();
    }
  }

  /**
   * Ensures that if a connection on a host is lost and the number of remaining connections is at
   * core connection count that no connections are re-established until after there are enough
   * inflight requests to justify creating one and the reconnection interval has elapsed.
   *
   * @jira_ticket JAVA-544
   * @test_category connection:connection_pool
   * @since 2.0.11
   */
  @Test(groups = "short")
  public void should_not_schedule_reconnect_when_connection_lost_and_at_core_connections()
      throws Exception {
    int readTimeout = 1000;
    int reconnectInterval = 1000;
    Cluster cluster =
        this.createClusterBuilder()
            .withSocketOptions(
                new SocketOptions()
                    .setConnectTimeoutMillis(readTimeout)
                    .setReadTimeoutMillis(reconnectInterval))
            .withReconnectionPolicy(new ConstantReconnectionPolicy(1000))
            .build();
    List<MockRequest> allRequests = newArrayList();
    try {
      cluster.init();

      Connection.Factory factory = spy(cluster.manager.connectionFactory);
      cluster.manager.connectionFactory = factory;

      HostConnectionPool pool = createPool(cluster, 1, 2);
      Connection core0 = pool.connections.get(0);

      // Create enough inFlight requests to spawn another connection.
      List<MockRequest> core0requests = newArrayList();
      for (int i = 0; i < 101; i++) {
        MockRequest request = MockRequest.send(pool);
        assertThat(request.getConnection()).isEqualTo(core0);
        core0requests.add(request);
      }

      // Pool should grow by 1.
      verify(factory, after(2000).times(1)).open(any(HostConnectionPool.class));
      assertThat(pool.connections).hasSize(2);

      // Reset factory mock as we'll be checking for new open() invokes later.
      reset(factory);

      // Grab the new non-core connection.
      Connection extra1 = pool.connections.get(1);

      // Drop a connection and disable listening.
      currentClient.closeConnection(CLOSE, ((InetSocketAddress) core0.channel.localAddress()));
      Uninterruptibles.sleepUninterruptibly(100, TimeUnit.MILLISECONDS);
      currentClient.disableListener();

      // Since core0 was closed, all of it's requests should have errored.
      for (MockRequest request : core0requests) {
        assertThat(request.state.get()).isEqualTo(MockRequest.State.FAILED);
      }

      assertThat(cluster).host(1).hasState(Host.State.UP);

      // Create enough inFlight requests to fill connection.
      List<MockRequest> requests = MockRequest.sendMany(100, pool);
      assertBorrowedConnections(requests, Collections.singletonList(extra1));
      allRequests.addAll(requests);
      assertThat(pool.connections).hasSize(1);

      // A new connection should never have been spawned since we didn't max out core.
      verify(factory, after(readTimeout).never()).open(any(HostConnectionPool.class));

      // Borrow another connection, since we exceed max another connection should be opened.
      MockRequest request = MockRequest.send(pool);
      allRequests.add(request);
      assertThat(request.getConnection()).isEqualTo(extra1);

      // After some time the a connection should attempt to be opened (but will fail).
      verify(factory, timeout(readTimeout)).open(any(HostConnectionPool.class));
      assertPoolSize(pool, 1);
      assertThat(pool.connections).hasSize(1);

      // Wait some reasonable amount of time for connection to reestablish then check pool size.
      Uninterruptibles.sleepUninterruptibly(readTimeout * 2, TimeUnit.MILLISECONDS);
      // Reconnecting failed since listening was enabled.
      assertPoolSize(pool, 1);

      // Re enable listening then wait for reconnect.
      currentClient.enableListener();
      Uninterruptibles.sleepUninterruptibly(reconnectInterval, TimeUnit.MILLISECONDS);

      // Borrow another connection, since we exceed max another connection should be opened.
      request = MockRequest.send(pool);
      allRequests.add(request);
      assertThat(request.getConnection()).isEqualTo(extra1);

      // Wait some reasonable amount of time for connection to reestablish then check pool size.
      Uninterruptibles.sleepUninterruptibly(readTimeout, TimeUnit.MILLISECONDS);
      // Reconnecting should have exceeded and pool will have grown.
      assertThat(pool.connections).hasSize(2);

      // Borrowed connection should be the newly spawned connection since the other one has some
      // inflight requests.
      request = MockRequest.send(pool);
      allRequests.add(request);
      assertThat(request.getConnection()).isNotEqualTo(core0).isNotEqualTo(extra1);
    } finally {
      MockRequest.completeAll(allRequests);
      cluster.close();
    }
  }

  /**
   * Ensures that if some connections fail on pool init that the host and subsequently the control
   * connection is not marked down. The test also ensures that when making requests on the pool that
   * connections are brought up to core.
   *
   * @jira_ticket JAVA-544
   * @test_category connection:connection_pool
   * @since 2.0.11
   */
  @Test(groups = "short")
  public void should_not_mark_host_down_if_some_connections_fail_on_init() throws Exception {
    int readTimeout = 1000;
    int reconnectInterval = 1000;
    Cluster cluster =
        this.createClusterBuilder()
            .withSocketOptions(
                new SocketOptions()
                    .setConnectTimeoutMillis(readTimeout)
                    .setReadTimeoutMillis(reconnectInterval))
            .withReconnectionPolicy(new ConstantReconnectionPolicy(1000))
            .build();
    List<MockRequest> allRequests = newArrayList();
    try {
      cluster.init();

      Connection.Factory factory = spy(cluster.manager.connectionFactory);
      cluster.manager.connectionFactory = factory;

      // Allow the first 4 connections to establish, but disable after that.
      currentClient.disableListener(4);
      HostConnectionPool pool = createPool(cluster, 8, 8);

      reset(factory);

      // Pool size should show all successful connections.
      assertThat(pool.connections).hasSize(4);

      // Control connection should remain up in addition to to host.
      assertThat(cluster).host(1).hasState(Host.State.UP);
      assertThat(cluster).hasOpenControlConnection();

      // Reenable listener, wait reconnectInterval and then try borrowing a connection.
      currentClient.enableListener();

      Uninterruptibles.sleepUninterruptibly(reconnectInterval, TimeUnit.MILLISECONDS);

      // Should open up to core connections, however it will only spawn up to 1 connection
      // per request, so we need to make enough requests to make up the deficit.  Additionally
      // we need to wait for connections to be established between requests for the pool
      // to spawn new connections (since it only allows one simultaneous creation).
      for (int i = 5; i <= 8; i++) {
        allRequests.add(MockRequest.send(pool));
        verify(factory, timeout(readTimeout)).open(any(HostConnectionPool.class));
        reset(factory);
        assertPoolSize(pool, i);
      }
    } finally {
      MockRequest.completeAll(allRequests);
      cluster.close();
    }
  }

  /**
   * Ensures that if all connections fail on pool init that the host and subsequently the control
   * connection is not marked down since the control connection is still active. The test also
   * ensures that borrow attempts on the pool fail if we are still in the reconnection window
   * according to the ConvictionPolicy.
   *
   * @jira_ticket JAVA-544
   * @test_category connection:connection_pool
   * @since 2.0.11
   */
  @Test(groups = "short")
  public void should_throw_exception_if_convicted_and_no_connections_available() {
    int readTimeout = 1000;
    int reconnectInterval = 1000;
    Cluster cluster =
        this.createClusterBuilder()
            .withSocketOptions(
                new SocketOptions()
                    .setConnectTimeoutMillis(readTimeout)
                    .setReadTimeoutMillis(reconnectInterval))
            .withReconnectionPolicy(new ConstantReconnectionPolicy(1000))
            .build();
    try {
      // Init cluster so control connection is created.
      cluster.init();
      assertThat(cluster).hasOpenControlConnection();

      Connection.Factory factory = spy(cluster.manager.connectionFactory);
      cluster.manager.connectionFactory = factory;

      // Disable listener so all connections on pool fail.
      currentClient.disableListener();
      HostConnectionPool pool = createPool(cluster, 8, 8);

      reset(factory);

      // Pool should be empty.
      assertThat(pool.connections).hasSize(0);

      // Control connection should stay up with the host.
      assertThat(cluster).host(1).hasState(Host.State.UP);
      assertThat(cluster).hasOpenControlConnection();

      MockRequest failedBorrow = MockRequest.send(pool);
      try {
        failedBorrow.getConnection();
        fail("Expected a BusyPoolException");
      } catch (BusyPoolException e) {
        /*expected*/
      }
    } finally {
      cluster.close();
    }
  }

  /**
   * Ensures that if all connections fail on pool init that the host and subsequently the control
   * connection is not marked down. The test also ensures that when making requests on the pool
   * after the conviction period that all core connections are created.
   *
   * @jira_ticket JAVA-544
   * @test_category connection:connection_pool
   * @since 2.0.11
   */
  @Test(groups = "short")
  public void should_wait_on_connection_if_not_convicted_and_no_connections_available()
      throws Exception {
    int readTimeout = 1000;
    int reconnectInterval = 1000;
    Cluster cluster =
        this.createClusterBuilder()
            .withSocketOptions(
                new SocketOptions()
                    .setConnectTimeoutMillis(readTimeout)
                    .setReadTimeoutMillis(reconnectInterval))
            .withReconnectionPolicy(new ConstantReconnectionPolicy(1000))
            .build();
    try {
      // Init cluster so control connection is created.
      cluster.init();
      assertThat(cluster).hasOpenControlConnection();

      Connection.Factory factory = spy(cluster.manager.connectionFactory);
      cluster.manager.connectionFactory = factory;

      // Disable listener so all connections on pool fail.
      currentClient.disableListener();
      HostConnectionPool pool = createPool(cluster, 8, 8);

      // Pool should be empty.
      assertThat(pool.connections).hasSize(0);

      // Control connection should stay up with the host.
      assertThat(cluster).host(1).hasState(Host.State.UP);
      assertThat(cluster).hasOpenControlConnection();

      currentClient.enableListener();

      // Wait for reconnectInterval so ConvictionPolicy allows connection to be created.
      Uninterruptibles.sleepUninterruptibly(reconnectInterval, TimeUnit.MILLISECONDS);

      reset(factory);
      MockRequest request = MockRequest.send(pool, 1);

      // Should create up to core connections.
      verify(factory, timeout(readTimeout * 8).times(8)).open(any(HostConnectionPool.class));
      assertPoolSize(pool, 8);

      request.simulateSuccessResponse();
    } finally {
      cluster.close();
    }
  }

  /**
   * Ensures that if a pool is created with zero core connections that when a request is first sent
   * that one and only one connection is created and that it waits on availability of that
   * connection and returns it.
   *
   * @jira_ticket JAVA-544
   * @test_category connection:connection_pool
   * @since 2.0.11
   */
  @Test(groups = "short")
  public void should_wait_on_connection_if_zero_core_connections() throws Exception {
    int readTimeout = 1000;
    int reconnectInterval = 1000;
    Cluster cluster =
        this.createClusterBuilder()
            .withSocketOptions(
                new SocketOptions()
                    .setConnectTimeoutMillis(readTimeout)
                    .setReadTimeoutMillis(reconnectInterval))
            .withReconnectionPolicy(new ConstantReconnectionPolicy(1000))
            .build();
    try {
      // Init cluster so control connection is created.
      cluster.init();
      assertThat(cluster).hasOpenControlConnection();

      Connection.Factory factory = spy(cluster.manager.connectionFactory);
      cluster.manager.connectionFactory = factory;

      HostConnectionPool pool = createPool(cluster, 0, 2);

      // Pool should be empty.
      assertThat(pool.connections).hasSize(0);

      // Control connection should stay up with the host.
      assertThat(cluster).host(1).hasState(Host.State.UP);
      assertThat(cluster).hasOpenControlConnection();

      reset(factory);
      MockRequest request = MockRequest.send(pool, 1);

      // Should create up to core connections.
      verify(factory, timeout(readTimeout).times(1)).open(any(HostConnectionPool.class));
      assertPoolSize(pool, 1);
      Uninterruptibles.getUninterruptibly(request.requestInitialized, 10, TimeUnit.SECONDS);
      request.simulateSuccessResponse();
    } finally {
      cluster.close();
    }
  }

  private HostConnectionPool createPool(Cluster cluster, int coreConnections, int maxConnections) {
    cluster
        .getConfiguration()
        .getPoolingOptions()
        .setNewConnectionThreshold(HostDistance.LOCAL, 100)
        .setMaxRequestsPerConnection(HostDistance.LOCAL, 128)
        .setMaxConnectionsPerHost(HostDistance.LOCAL, maxConnections)
        .setCoreConnectionsPerHost(HostDistance.LOCAL, coreConnections);
    Session session = cluster.connect();
    Host host = TestUtils.findHost(cluster, 1);

    // Replace the existing pool with a spy pool and return it.
    SessionManager sm = ((SessionManager) session);
    return sm.pools.get(host);
  }

  /**
   * This test uses a table named "Java349" with 1000 column and performs asynchronously 100k
   * insertions. While the insertions are being executed, the number of opened connection is
   * monitored.
   *
   * <p>If at anytime, the number of opened connections is negative, this test will fail.
   *
   * @jira_ticket JAVA-349
   * @test_category connection:connection_pool
   * @since 2.0.6, 2.1.1
   */
  @Test(groups = "long", enabled = false /* this test causes timeouts on Jenkins */)
  public void open_connections_metric_should_always_be_positive() throws InterruptedException {
    // Track progress in a dedicated thread
    int numberOfInserts = 100 * 1000;
    final CountDownLatch pendingInserts = new CountDownLatch(numberOfInserts);
    ExecutorService progressReportExecutor = Executors.newSingleThreadExecutor();
    final Runnable progressReporter =
        new Runnable() {
          @Override
          public void run() {
            pendingInserts.countDown();
          }
        };

    // Track opened connections in a dedicated thread every one second
    final AtomicBoolean negativeOpenConnectionCountSpotted = new AtomicBoolean(false);
    final Gauge<Integer> openConnections = cluster.getMetrics().getOpenConnections();
    ScheduledExecutorService openConnectionsWatcherExecutor = Executors.newScheduledThreadPool(1);
    final Runnable openConnectionsWatcher =
        new Runnable() {
          @Override
          public void run() {
            Integer value = openConnections.getValue();
            if (value < 0) {
              System.err.println("Negative value spotted for openConnection metric: " + value);
              negativeOpenConnectionCountSpotted.set(true);
            }
          }
        };
    openConnectionsWatcherExecutor.scheduleAtFixedRate(openConnectionsWatcher, 1, 1, SECONDS);

    // Insert 100k lines in a newly created 1k columns table
    PreparedStatement insertStatement = session.prepare(generateJava349InsertStatement());
    for (int key = 0; key < numberOfInserts; key++) {
      ResultSetFuture future = session.executeAsync(insertStatement.bind(key));
      future.addListener(progressReporter, progressReportExecutor);
    }

    // Wait for all inserts to happen and stop connections and progress tracking
    pendingInserts.await();
    openConnectionsWatcherExecutor.shutdownNow();
    progressReportExecutor.shutdownNow();

    if (negativeOpenConnectionCountSpotted.get()) {
      fail("Negative value spotted for open connection count");
    }
  }

  /**
   * Ensures that on a a host is bring up with a reused connection and zero core connections the
   * driver do not try to create any connections.
   *
   * @jira_ticket JAVA-1794
   * @test_category connection:connection_pool
   * @since 3.5.1
   */
  @Test(groups = "short")
  public void
      should_not_create_connections_if_zero_core_connections_and_reused_connection_on_reconnect()
          throws Exception {
    int reconnectInterval = 1000;

    Cluster cluster =
        this.createClusterBuilder()
            .withReconnectionPolicy(new ConstantReconnectionPolicy(1000))
            .build();

    Host.StateListener stateListener = mock(Host.StateListener.class);

    try {
      // Init cluster so control connection is created.
      cluster.init();
      assertThat(cluster).hasOpenControlConnection();

      Connection.Factory factory = spy(cluster.manager.connectionFactory);
      cluster.manager.connectionFactory = factory;

      HostConnectionPool pool = createPool(cluster, 0, 2);

      cluster.register(stateListener);
      cluster.manager.triggerOnDown(pool.host, true);

      Thread.sleep(reconnectInterval * 3);

      verify(stateListener, times(1)).onUp(pool.host);

      // Pool should be empty.
      assertThat(pool.connections).hasSize(0);

      // Control connection should stay up with the host.
      assertThat(cluster).host(1).hasState(Host.State.UP);
      assertThat(cluster).hasOpenControlConnection();

    } finally {
      cluster.close();
    }
  }

  private String generateJava349InsertStatement() {
    StringBuilder sb = new StringBuilder("INSERT INTO Java349 (mykey");
    for (int i = 0; i < 1000; i++) {
      sb.append(", column").append(i);
    }
    sb.append(") VALUES (?");
    for (int i = 0; i < 1000; i++) {
      sb.append(", ").append(i);
    }
    sb.append(");");

    PrimingRequest preparedStatementPrime =
        PrimingRequest.preparedStatementBuilder()
            .withQuery(sb.toString())
            .withThen(then().withVariableTypes(PrimitiveType.INT))
            .build();
    primingClient.prime(preparedStatementPrime);
    return sb.toString();
  }

  /**
   * Mock ResponseCallback that simulates the behavior of SpeculativeExecution (in terms of
   * borrowing/releasing connections).
   */
  static class MockRequest implements Connection.ResponseCallback {

    enum State {
      START,
      COMPLETED,
      FAILED,
      TIMED_OUT
    }

    final ListenableFuture<Connection> connectionFuture;

    final ListenableFuture<Connection.ResponseHandler> requestInitialized;

    private volatile Connection.ResponseHandler responseHandler;

    final AtomicReference<State> state = new AtomicReference<State>(State.START);

    static MockRequest send(HostConnectionPool pool) {
      return send(pool, 0);
    }

    static MockRequest send(HostConnectionPool pool, int maxQueueSize)
        throws ConnectionException, BusyConnectionException {
      return send(pool, 5000, maxQueueSize);
    }

    static MockRequest send(HostConnectionPool pool, int timeoutMillis, int maxQueueSize)
        throws ConnectionException, BusyConnectionException {
      return new MockRequest(pool, timeoutMillis, maxQueueSize);
    }

    private static List<MockRequest> sendMany(int count, HostConnectionPool pool)
        throws ConnectionException {
      return sendMany(count, pool, 0);
    }

    private static List<MockRequest> sendMany(int count, HostConnectionPool pool, int maxQueueSize)
        throws ConnectionException {
      List<MockRequest> requests = newArrayList();
      for (int i = 0; i < count; i++) {
        MockRequest request = send(pool, maxQueueSize);
        requests.add(request);
      }
      return requests;
    }

    /** Completes count requests by simulating a successful response. */
    private static void completeMany(int count, List<MockRequest> requests) {
      Iterator<MockRequest> requestIt = requests.iterator();
      for (int i = 0; i < count; i++) {
        if (requestIt.hasNext()) {
          MockRequest request = requestIt.next();
          request.simulateSuccessResponse();
          requestIt.remove();
        } else {
          break;
        }
      }
    }

    /** Completes all requests by simulating a successful response. */
    private static void completeAll(List<MockRequest> requests) {
      for (MockRequest request : requests) {
        request.simulateSuccessResponse();
      }
    }

    private MockRequest(HostConnectionPool pool, int timeoutMillis, int maxQueueSize)
        throws ConnectionException {
      this.connectionFuture = pool.borrowConnection(timeoutMillis, MILLISECONDS, maxQueueSize);
      requestInitialized =
          GuavaCompatibility.INSTANCE.transform(
              this.connectionFuture,
              new Function<Connection, Connection.ResponseHandler>() {
                @Override
                public Connection.ResponseHandler apply(Connection connection) {
                  MockRequest thisRequest = MockRequest.this;
                  thisRequest.responseHandler =
                      new Connection.ResponseHandler(connection, -1, thisRequest);
                  connection.dispatcher.add(thisRequest.responseHandler);
                  return responseHandler;
                }
              });
    }

    void simulateSuccessResponse() {
      onSet(getConnection(), null, 0, 0);
    }

    @SuppressWarnings("unused")
    void simulateErrorResponse() {
      onException(getConnection(), null, 0, 0);
    }

    @SuppressWarnings("unused")
    void simulateTimeout() {
      if (onTimeout(getConnection(), 0, 0)) responseHandler.cancelHandler();
    }

    Connection getConnection() {
      try {
        return Uninterruptibles.getUninterruptibly(connectionFuture, 500, MILLISECONDS);
      } catch (ExecutionException e) {
        throw Throwables.propagate(e.getCause());
      } catch (TimeoutException e) {
        fail("Timed out getting connection");
        return null; // never reached
      }
    }

    @Override
    public void onSet(
        Connection connection, Message.Response response, long latency, int retryCount) {
      assertThat(connectionFuture.isDone()).isTrue();
      try {
        assertThat(Uninterruptibles.getUninterruptibly(connectionFuture)).isNotNull();
      } catch (ExecutionException e) {
        Throwables.propagate(e.getCause());
      }
      if (state.compareAndSet(State.START, State.COMPLETED)) {
        connection.dispatcher.removeHandler(responseHandler, true);
        connection.release();
      }
    }

    @Override
    public void onException(
        Connection connection, Exception exception, long latency, int retryCount) {
      if (state.compareAndSet(State.START, State.FAILED)) {
        connection.dispatcher.removeHandler(responseHandler, true);
        connection.release();
      }
    }

    @Override
    public boolean onTimeout(Connection connection, long latency, int retryCount) {
      return state.compareAndSet(State.START, State.TIMED_OUT);
    }

    @Override
    public Message.Request request() {
      return null; // not needed for this test class
    }

    @Override
    public int retryCount() {
      return 0; // value not important for this test class
    }
  }

  static class TestExecutorService extends ForwardingListeningExecutorService {

    private final ListeningExecutorService delegate;

    private final Semaphore semaphore = new Semaphore(0);

    TestExecutorService(ListeningExecutorService delegate) {
      this.delegate = delegate;
    }

    @Override
    protected ListeningExecutorService delegate() {
      return delegate;
    }

    public void reset() {
      semaphore.drainPermits();
    }

    public void blockUntilNextTaskCompleted() throws InterruptedException {
      semaphore.tryAcquire(1, 1, TimeUnit.MINUTES);
    }

    @Override
    public ListenableFuture<?> submit(Runnable task) {
      ListenableFuture<?> future = super.submit(task);
      GuavaCompatibility.INSTANCE.addCallback(
          future,
          new FutureCallback<Object>() {
            @Override
            public void onSuccess(Object result) {
              semaphore.release(1);
            }

            @Override
            public void onFailure(Throwable t) {
              semaphore.release(1);
            }
          });
      return future;
    }
  }
}
