/*
 * 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 org.testng.Assert.assertEquals;
import static org.testng.Assert.fail;

import com.datastax.driver.core.utils.SocketChannelMonitor;
import com.google.common.collect.Lists;
import com.google.common.util.concurrent.Uninterruptibles;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.Callable;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.Future;
import java.util.concurrent.TimeUnit;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.testng.annotations.AfterMethod;
import org.testng.annotations.Test;

public class ClusterStressTest extends CCMTestsSupport {

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

  private ExecutorService executorService;

  public ClusterStressTest() {
    // 8 threads should be enough so that we stress the driver and not the OS thread scheduler
    executorService = Executors.newFixedThreadPool(8);
  }

  @Test(groups = "long")
  public void clusters_should_not_leak_connections() {
    int numberOfClusters = 10;
    int numberOfIterations = 500;
    try {
      for (int i = 1; i < numberOfIterations; i++) {
        logger.info("On iteration {}/{}.", i, numberOfIterations);
        logger.info("Creating {} clusters", numberOfClusters);
        List<CreateClusterAndCheckConnections> actions =
            waitForCreates(createClustersConcurrently(numberOfClusters));
        waitForCloses(closeClustersConcurrently(actions));
        if (logger.isDebugEnabled())
          logger.debug(
              "# {} threads currently running", Thread.getAllStackTraces().keySet().size());
      }
    } finally {
      logger.info("Sleeping 60 seconds to free TCP resources");
      Uninterruptibles.sleepUninterruptibly(60, TimeUnit.SECONDS);
    }
  }

  @AfterMethod(groups = "long", alwaysRun = true)
  public void shutdown() throws Exception {
    executorService.shutdown();
    try {
      boolean shutdown = executorService.awaitTermination(30, TimeUnit.SECONDS);
      if (!shutdown) fail("executor ran for longer than expected");
    } catch (InterruptedException e) {
      fail("Interrupted while waiting for executor to shutdown");
    } finally {
      executorService = null;
      System.gc();
    }
  }

  private List<Future<CreateClusterAndCheckConnections>> createClustersConcurrently(
      int numberOfClusters) {
    final CountDownLatch countDownLatch = new CountDownLatch(1);
    return createClustersConcurrently(numberOfClusters, countDownLatch);
  }

  private List<Future<CreateClusterAndCheckConnections>> createClustersConcurrently(
      int numberOfClusters, CountDownLatch countDownLatch) {
    List<Future<CreateClusterAndCheckConnections>> clusterFutures =
        Lists.newArrayListWithCapacity(numberOfClusters);
    for (int i = 0; i < numberOfClusters; i++) {
      clusterFutures.add(
          executorService.submit(new CreateClusterAndCheckConnections(countDownLatch)));
    }
    countDownLatch.countDown();
    return clusterFutures;
  }

  private List<Future<Void>> closeClustersConcurrently(
      List<CreateClusterAndCheckConnections> actions) {
    final CountDownLatch countDownLatch = new CountDownLatch(1);
    return closeClustersConcurrently(actions, countDownLatch);
  }

  private List<Future<Void>> closeClustersConcurrently(
      List<CreateClusterAndCheckConnections> actions, CountDownLatch startSignal) {
    List<Future<Void>> closeFutures = Lists.newArrayListWithCapacity(actions.size());
    for (CreateClusterAndCheckConnections action : actions) {
      closeFutures.add(
          executorService.submit(
              new CloseCluster(action.cluster, action.channelMonitor, startSignal)));
    }
    startSignal.countDown();
    return closeFutures;
  }

  private List<CreateClusterAndCheckConnections> waitForCreates(
      List<Future<CreateClusterAndCheckConnections>> futures) {
    List<CreateClusterAndCheckConnections> actions = Lists.newArrayListWithCapacity(futures.size());
    // If an error occurs, we will abort the test, but we still want to close all the clusters
    // that were opened successfully, so we iterate over the whole list no matter what.
    AssertionError error = null;
    for (Future<CreateClusterAndCheckConnections> future : futures) {
      try {
        actions.add(future.get());
      } catch (InterruptedException e) {
        Thread.currentThread().interrupt();
        if (error == null)
          error = assertionError("Interrupted while waiting for future creation", e);
      } catch (ExecutionException e) {
        if (error == null) {
          Throwable cause = e.getCause();
          if (cause instanceof AssertionError) error = (AssertionError) cause;
          else error = assertionError("Error while creating a cluster", cause);
        }
      }
    }
    if (error != null) {
      for (CreateClusterAndCheckConnections action : actions) action.cluster.close();
      throw error;
    } else return actions;
  }

  private List<Void> waitForCloses(List<Future<Void>> futures) {
    List<Void> result = new ArrayList<Void>(futures.size());
    AssertionError error = null;
    for (Future<Void> future : futures) {
      try {
        result.add(future.get());
      } catch (InterruptedException e) {
        Thread.currentThread().interrupt();
        if (error == null) error = assertionError("Interrupted while waiting for future", e);
      } catch (ExecutionException e) {
        if (error == null) {
          Throwable cause = e.getCause();
          if (cause instanceof AssertionError) error = (AssertionError) cause;
          else error = assertionError("Error while closing a cluster", cause);
        }
      }
    }
    if (error != null) throw error;
    else return result;
  }

  private class CreateClusterAndCheckConnections
      implements Callable<CreateClusterAndCheckConnections> {
    private final CountDownLatch startSignal;
    private Cluster cluster;
    private final SocketChannelMonitor channelMonitor = new SocketChannelMonitor();

    CreateClusterAndCheckConnections(CountDownLatch startSignal) {
      this.startSignal = startSignal;
      this.cluster =
          createClusterBuilder()
              .withPoolingOptions(
                  new PoolingOptions().setCoreConnectionsPerHost(HostDistance.LOCAL, 1))
              .withNettyOptions(channelMonitor.nettyOptions())
              .build();
    }

    @Override
    public CreateClusterAndCheckConnections call() throws Exception {
      startSignal.await();

      try {
        // There should be 1 control connection after initializing.
        cluster.init();
        assertEquals(cluster.manager.sessions.size(), 0);
        assertEquals((int) cluster.getMetrics().getOpenConnections().getValue(), 1);
        assertEquals(channelMonitor.openChannels(getContactPointsWithPorts()).size(), 1);

        // The first session initializes the cluster and its control connection
        Session session = cluster.connect();
        assertEquals(cluster.manager.sessions.size(), 1);
        assertEquals(
            (int) cluster.getMetrics().getOpenConnections().getValue(),
            1 + TestUtils.numberOfLocalCoreConnections(cluster));
        assertEquals(
            channelMonitor.openChannels(getContactPointsWithPorts()).size(),
            1 + TestUtils.numberOfLocalCoreConnections(cluster));

        // Closing the session keeps the control connection opened
        session.close();
        assertEquals(cluster.manager.sessions.size(), 0);
        assertEquals((int) cluster.getMetrics().getOpenConnections().getValue(), 1);
        assertEquals(channelMonitor.openChannels(getContactPointsWithPorts()).size(), 1);

        return this;
      } catch (AssertionError e) {
        // If an assertion fails, close the cluster now, because it's the last time we
        // have a reference to it.
        cluster.close();
        cluster = null;
        throw e;
      } finally {
        channelMonitor.stop();
      }
    }
  }

  private class CloseCluster implements Callable<Void> {
    private Cluster cluster;
    private SocketChannelMonitor channelMonitor;
    private final CountDownLatch startSignal;

    CloseCluster(Cluster cluster, SocketChannelMonitor channelMonitor, CountDownLatch startSignal) {
      this.cluster = cluster;
      this.channelMonitor = channelMonitor;
      this.startSignal = startSignal;
    }

    @Override
    public Void call() throws Exception {
      startSignal.await();
      try {
        cluster.close();
        assertEquals(cluster.manager.sessions.size(), 0);
        assertEquals(channelMonitor.openChannels(getContactPointsWithPorts()).size(), 0);
      } finally {
        channelMonitor.stop();
        cluster = null;
        channelMonitor = null;
      }
      return null;
    }
  }

  private static AssertionError assertionError(String message, Throwable cause) {
    AssertionError error = new AssertionError(message);
    error.initCause(cause);
    return error;
  }
}
