/*
 * 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 com.datastax.driver.core.exceptions.NoHostAvailableException;
import com.datastax.driver.core.policies.ConstantReconnectionPolicy;
import com.datastax.driver.core.utils.SocketChannelMonitor;
import com.google.common.util.concurrent.FutureCallback;
import com.google.common.util.concurrent.ThreadFactoryBuilder;
import com.google.common.util.concurrent.Uninterruptibles;
import io.netty.channel.socket.SocketChannel;
import java.util.Collection;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.Semaphore;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicInteger;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.testng.annotations.Test;

@CCMConfig(numberOfNodes = 3)
public class TimeoutStressTest extends CCMTestsSupport {

  static final Logger logger = LoggerFactory.getLogger(TimeoutStressTest.class);

  // Maximum number of concurrent queries running at a given time.
  static final int CONCURRENT_QUERIES = 25;

  // How long the test should run for, may want to consider running for longer periods to time to
  // check for leaks
  // that could occur over very tiny timing windows.
  static final long DURATION = 60000;

  // Configured read timeout - this may need to be tuned to the host system running the test.
  static final int READ_TIMEOUT_IN_MS = 50;

  // Configured connection timeout - this may need to be tuned to the host system running the test.
  static final int CONNECTION_TIMEOUT_IN_MS = 20;

  private static AtomicInteger executedQueries = new AtomicInteger(0);

  private SocketChannelMonitor channelMonitor;

  @Override
  public Cluster.Builder createClusterBuilder() {
    channelMonitor = register(new SocketChannelMonitor());
    PoolingOptions poolingOptions =
        new PoolingOptions().setConnectionsPerHost(HostDistance.LOCAL, 8, 8);
    return super.createClusterBuilder()
        .withPoolingOptions(poolingOptions)
        .withNettyOptions(channelMonitor.nettyOptions())
        .withReconnectionPolicy(new ConstantReconnectionPolicy(1000));
  }

  @Override
  public void onTestContextInitialized() {
    execute(
        "create table record (\n"
            + "  name text,\n"
            + "  phone text,\n"
            + "  value text,\n"
            + "  PRIMARY KEY (name, phone)\n"
            + ")");
  }

  /**
   * Validates that under extreme timeout conditions the driver is able to properly maintain
   * connection pools in addition to not leaking connections.
   *
   * <p>
   *
   * <p>Does the following:
   *
   * <ol>
   *   <li>Creates a table and loads 30k rows in a single partition.
   *   <li>Sets the connection and read timeout {@link SocketOptions} to very low values.
   *   <li>Spawns workers that concurrently execute queries.
   *   <li>For some duration, repeatedly measures number of open socket connections and warns if
   *       exceeded.
   *   <li>After a duration, resets {@link SocketOptions} to defaults.
   *   <li>Wait for 20 seconds for reaper to remove old connections and restore pools.
   *   <li>Ensure pools are restored.
   *   <li>Shutdown session and ensure that there remains only 1 open connection.
   * </ol>
   *
   * @test_category connection:connection_pool
   * @expected_result no connections leak and all host pools are maintained.
   * @jira_ticket JAVA-692
   * @since 2.0.10, 2.1.6
   */
  @Test(groups = "stress")
  public void host_state_should_be_maintained_with_timeouts() throws Exception {
    insertRecords();
    session().close();

    // Set very low timeouts.
    cluster()
        .getConfiguration()
        .getSocketOptions()
        .setConnectTimeoutMillis(CONNECTION_TIMEOUT_IN_MS);
    cluster().getConfiguration().getSocketOptions().setReadTimeoutMillis(READ_TIMEOUT_IN_MS);
    Session newSession = cluster().connect(keyspace);
    PreparedStatement statement =
        newSession.prepare("select * from record where name=? limit 1000;");

    int workers = Runtime.getRuntime().availableProcessors();
    ExecutorService workerPool =
        Executors.newFixedThreadPool(
            workers,
            new ThreadFactoryBuilder()
                .setNameFormat("timeout-stress-test-worker-%d")
                .setDaemon(true)
                .build());

    AtomicBoolean stopped = new AtomicBoolean(false);

    // Ensure that we never exceed MaxConnectionsPerHost * nodes + 1 control connection.
    int maxConnections =
        TestUtils.numberOfLocalCoreConnections(cluster()) * getContactPoints().size() + 1;

    try {
      Semaphore concurrentQueries = new Semaphore(CONCURRENT_QUERIES);
      for (int i = 0; i < workers; i++) {
        workerPool.submit(
            new TimeoutStressWorker(newSession, statement, concurrentQueries, stopped));
      }

      long startTime = System.currentTimeMillis();
      while (System.currentTimeMillis() - startTime < DURATION) {
        Uninterruptibles.sleepUninterruptibly(1, TimeUnit.SECONDS);
        channelMonitor.report();
        // Some connections that are being closed may have had active requests which are delegated
        // to the
        // reaper for cleanup later.
        Collection<SocketChannel> openChannels =
            channelMonitor.openChannels(getContactPointsWithPorts());

        // Ensure that we don't exceed maximum connections.  Log as warning as there will be a bit
        // of a timing
        // factor between retrieving open connections and checking the reaper.
        if (openChannels.size() > maxConnections) {
          logger.warn(
              "{} of open channels: {} exceeds maximum expected: {}.  "
                  + "This could be because there are connections to be cleaned up in the reaper.",
              openChannels.size(),
              maxConnections,
              openChannels);
        }
      }
    } finally {
      stopped.set(true);

      // Reset socket timeouts to allow pool to recover.
      cluster()
          .getConfiguration()
          .getSocketOptions()
          .setConnectTimeoutMillis(SocketOptions.DEFAULT_CONNECT_TIMEOUT_MILLIS);
      cluster()
          .getConfiguration()
          .getSocketOptions()
          .setReadTimeoutMillis(SocketOptions.DEFAULT_READ_TIMEOUT_MILLIS);

      logger.debug(
          "Sleeping 20 seconds to allow connection reaper to clean up connections "
              + "and for the pools to recover.");
      Uninterruptibles.sleepUninterruptibly(20, TimeUnit.SECONDS);

      Collection<SocketChannel> openChannels =
          channelMonitor.openChannels(getContactPointsWithPorts());
      assertThat(openChannels.size())
          .as("Number of open connections does not meet expected: %s", openChannels)
          .isLessThanOrEqualTo(maxConnections);

      // Each host should be in an up state.
      assertThat(cluster()).host(1).comesUpWithin(0, TimeUnit.SECONDS);
      assertThat(cluster()).host(2).comesUpWithin(0, TimeUnit.SECONDS);
      assertThat(cluster()).host(3).comesUpWithin(0, TimeUnit.SECONDS);

      newSession.close();

      openChannels = channelMonitor.openChannels(getContactPointsWithPorts());
      assertThat(openChannels.size())
          .as("Number of open connections does not meet expected: %s", openChannels)
          .isEqualTo(1);

      workerPool.shutdown();
    }
  }

  private void insertRecords() {
    int records = 30000;
    PreparedStatement insertStmt =
        session().prepare("insert into record (name, phone, value) values (?, ?, ?)");

    for (int i = 0; i < records; i++) {
      if (i % 1000 == 0) logger.debug("Inserting record {}.", i);
      session().execute(insertStmt.bind("0", Integer.toString(i), "test"));
    }
    logger.debug("Inserts complete.");
  }

  public static class TimeoutStressWorker implements Runnable {

    private final Semaphore concurrentQueries;
    private final AtomicBoolean stopped;
    private final Session session;
    private final PreparedStatement statement;

    public TimeoutStressWorker(
        Session session,
        PreparedStatement statement,
        Semaphore concurrentQueries,
        AtomicBoolean stopped) {
      this.session = session;
      this.statement = statement;
      this.concurrentQueries = concurrentQueries;
      this.stopped = stopped;
    }

    @Override
    public void run() {

      while (!stopped.get()) {
        try {
          concurrentQueries.acquire();
          ResultSetFuture future = session.executeAsync(statement.bind("0"));
          GuavaCompatibility.INSTANCE.addCallback(
              future,
              new FutureCallback<ResultSet>() {

                @Override
                public void onSuccess(ResultSet result) {
                  concurrentQueries.release();
                  if (executedQueries.incrementAndGet() % 1000 == 0)
                    logger.debug(
                        "Successfully executed {}.  rows: {}",
                        executedQueries.get(),
                        result.getAvailableWithoutFetching());
                }

                @Override
                public void onFailure(Throwable t) {
                  concurrentQueries.release();
                  if (t instanceof NoHostAvailableException) {
                    // logger.error("NHAE: {}", t.getMessage());
                  } else {
                    // logger.error("Exception", t);
                  }
                }
              });
        } catch (Exception e) {
          logger.error("Failure while submitting query.", e);
        }
      }
    }
  }
}
