/*
 * 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.ListenableFuture;
import com.google.common.util.concurrent.ListeningExecutorService;
import com.google.common.util.concurrent.MoreExecutors;
import com.google.common.util.concurrent.Uninterruptibles;
import java.util.Iterator;
import java.util.List;
import java.util.Stack;
import java.util.concurrent.Callable;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutionException;
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;

@CCMConfig(dirtiesContext = true)
public class SessionStressTest extends CCMTestsSupport {

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

  private ListeningExecutorService executorService;

  private Cluster stressCluster;

  private final SocketChannelMonitor channelMonitor = new SocketChannelMonitor();

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

  @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();
    }
  }

  /**
   * Stress test on opening/closing sessions.
   *
   * <p>This test opens and closes {@code Session} in a multithreaded environment and makes sure
   * that there is not connection leak. More specifically, this test performs the following steps:
   *
   * <p>
   *
   * <ul>
   *   <li>Open 2000 {@code Session} concurrently
   *   <li>Verify that 2000 sessions are reported as open by the {@code Cluster}
   *   <li>Verify that 4001 connections are reported as open by the {@code Cluster}
   *   <li>Close 1000 {@code Session} concurrently
   *   <li>Verify that 1000 sessions are reported as open by the {@code Cluster}
   *   <li>Verify that 2001 connections are reported as open by the {@code Cluster}
   *   <li>Open concurrently 1000 {@code Session} while 1000 other {@code Session} are closed
   *       concurrently
   *   <li>Verify that 1000 sessions are reported as open by the {@code Cluster}
   *   <li>Verify that 2001 connections are reported as open by the {@code Cluster}
   *   <li>Close 1000 {@code Session} concurrently
   *   <li>Verify that 0 sessions are reported as open by the {@code Cluster}
   *   <li>Verify that 1 connection is reported as open by the {@code Cluster}
   * </ul>
   *
   * <p>This test is linked to JAVA-432.
   */
  @Test(groups = "long")
  public void sessions_should_not_leak_connections() {
    // override inherited field with a new cluster object and ensure 0 sessions and connections
    channelMonitor.reportAtFixedInterval(1, TimeUnit.SECONDS);
    stressCluster =
        createClusterBuilder()
            .withPoolingOptions(
                new PoolingOptions().setCoreConnectionsPerHost(HostDistance.LOCAL, 1))
            .withNettyOptions(channelMonitor.nettyOptions())
            .build();

    try {
      stressCluster.init();

      // The cluster has been initialized, we should have 1 connection.
      assertEquals(stressCluster.manager.sessions.size(), 0);
      assertEquals((int) stressCluster.getMetrics().getOpenConnections().getValue(), 1);

      // The first session initializes the cluster and its control connection
      // This is a local cluster so we also have 2 connections per session
      Session session = stressCluster.connect();
      assertEquals(stressCluster.manager.sessions.size(), 1);
      int coreConnections = TestUtils.numberOfLocalCoreConnections(stressCluster);
      assertEquals(
          (int) stressCluster.getMetrics().getOpenConnections().getValue(), 1 + coreConnections);
      assertEquals(
          channelMonitor.openChannels(getContactPointsWithPorts()).size(), 1 + coreConnections);

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

      int nbOfSessions = 2000;
      int halfOfTheSessions = nbOfSessions / 2;
      int nbOfIterations = 5;
      int sleepTime = 20;

      for (int iteration = 1; iteration <= nbOfIterations; iteration++) {
        logger.info("On iteration {}/{}.", iteration, nbOfIterations);
        logger.info("Creating {} sessions.", nbOfSessions);
        waitFor(openSessionsConcurrently(nbOfSessions));

        // We should see the exact number of opened sessions
        // Since we have 2 connections per session, we should see 2 * sessions + control connection
        assertEquals(stressCluster.manager.sessions.size(), nbOfSessions);
        assertEquals(
            (int) stressCluster.getMetrics().getOpenConnections().getValue(),
            coreConnections * nbOfSessions + 1);
        assertEquals(
            channelMonitor.openChannels(getContactPointsWithPorts()).size(),
            coreConnections * nbOfSessions + 1);

        // Close half of the sessions asynchronously
        logger.info("Closing {}/{} sessions.", halfOfTheSessions, nbOfSessions);
        waitFor(closeSessionsConcurrently(halfOfTheSessions));

        // Check that we have the right number of sessions and connections
        assertEquals(stressCluster.manager.sessions.size(), halfOfTheSessions);
        assertEquals(
            (int) stressCluster.getMetrics().getOpenConnections().getValue(),
            coreConnections * (nbOfSessions / 2) + 1);
        assertEquals(
            channelMonitor.openChannels(getContactPointsWithPorts()).size(),
            coreConnections * (nbOfSessions / 2) + 1);

        // Close and open the same number of sessions concurrently
        logger.info("Closing and Opening {} sessions concurrently.", halfOfTheSessions);
        CountDownLatch startSignal = new CountDownLatch(2);
        List<ListenableFuture<Session>> openSessionFutures =
            openSessionsConcurrently(halfOfTheSessions, startSignal);
        List<ListenableFuture<Void>> closeSessionsFutures =
            closeSessionsConcurrently(halfOfTheSessions, startSignal);
        startSignal.countDown();
        waitFor(openSessionFutures);
        waitFor(closeSessionsFutures);

        // Check that we have the same number of sessions and connections
        assertEquals(stressCluster.manager.sessions.size(), halfOfTheSessions);
        assertEquals(
            (int) stressCluster.getMetrics().getOpenConnections().getValue(),
            coreConnections * (nbOfSessions / 2) + 1);
        assertEquals(
            channelMonitor.openChannels(getContactPointsWithPorts()).size(),
            coreConnections * (nbOfSessions / 2) + 1);

        // Close the remaining sessions
        logger.info("Closing remaining {} sessions.", halfOfTheSessions);
        waitFor(closeSessionsConcurrently(halfOfTheSessions));

        // Check that we have a clean state
        assertEquals(stressCluster.manager.sessions.size(), 0);
        assertEquals((int) stressCluster.getMetrics().getOpenConnections().getValue(), 1);
        assertEquals(channelMonitor.openChannels(getContactPointsWithPorts()).size(), 1);

        // On OSX, the TCP connections are released after 15s by default (sysctl -a
        // net.inet.tcp.msl)
        logger.info(
            "Sleeping {} seconds so that TCP connections are released by the OS", sleepTime);
        Uninterruptibles.sleepUninterruptibly(sleepTime, TimeUnit.SECONDS);
      }
    } finally {
      stressCluster.close();
      stressCluster = null;

      // Ensure no channels remain open.
      assertEquals(channelMonitor.openChannels(getContactPointsWithPorts()).size(), 0);

      channelMonitor.stop();
      channelMonitor.report();

      logger.info("Sleeping 60 extra seconds");
      Uninterruptibles.sleepUninterruptibly(60, TimeUnit.SECONDS);
    }
  }

  private List<ListenableFuture<Session>> openSessionsConcurrently(int iterations) {
    final CountDownLatch countDownLatch = new CountDownLatch(1);
    return openSessionsConcurrently(iterations, countDownLatch);
  }

  private List<ListenableFuture<Session>> openSessionsConcurrently(
      int iterations, CountDownLatch countDownLatch) {
    // Open new sessions once all tasks have been created
    List<ListenableFuture<Session>> sessionFutures = Lists.newArrayListWithCapacity(iterations);
    for (int i = 0; i < iterations; i++) {
      sessionFutures.add(executorService.submit(new OpenSession(countDownLatch)));
    }
    countDownLatch.countDown();
    return sessionFutures;
  }

  private List<ListenableFuture<Void>> closeSessionsConcurrently(int iterations) {
    final CountDownLatch countDownLatch = new CountDownLatch(1);
    return closeSessionsConcurrently(iterations, countDownLatch);
  }

  private List<ListenableFuture<Void>> closeSessionsConcurrently(
      int iterations, CountDownLatch countDownLatch) {
    // Get a reference to every session we want to close
    Stack<Session> sessionsToClose = new Stack<Session>();
    Iterator<? extends Session> iterator = stressCluster.manager.sessions.iterator();
    for (int i = 0; i < iterations; i++) {
      sessionsToClose.push(iterator.next());
    }

    // Close sessions asynchronously once all tasks have been created
    List<ListenableFuture<CloseFuture>> closeFutures = Lists.newArrayListWithCapacity(iterations);
    for (int i = 0; i < iterations; i++) {
      closeFutures.add(
          executorService.submit(new CloseSession(sessionsToClose.pop(), countDownLatch)));
    }
    countDownLatch.countDown();

    // Immediately wait for CloseFutures, this should be very quick since all this work does is call
    // closeAsync.
    List<ListenableFuture<Void>> futures = Lists.newArrayListWithCapacity(iterations);
    for (ListenableFuture<CloseFuture> closeFuture : closeFutures) {
      try {
        futures.add(closeFuture.get());
      } catch (Exception e) {
        logger.error("Got interrupted exception while waiting on closeFuture.", e);
      }
    }
    return futures;
  }

  private <E> void waitFor(List<ListenableFuture<E>> futures) {
    for (Future<E> future : futures) {
      try {
        future.get();
      } catch (InterruptedException e) {
        throw new RuntimeException("Interrupted while waiting for future", e);
      } catch (ExecutionException e) {
        e.printStackTrace();
        fail(e.getMessage());
      }
    }
  }

  private class OpenSession implements Callable<Session> {
    private final CountDownLatch startSignal;

    OpenSession(CountDownLatch startSignal) {
      this.startSignal = startSignal;
    }

    @Override
    public Session call() throws Exception {
      startSignal.await();
      return stressCluster.connect();
    }
  }

  private static class CloseSession implements Callable<CloseFuture> {
    private Session session;
    private final CountDownLatch startSignal;

    CloseSession(Session session, CountDownLatch startSignal) {
      this.session = session;
      this.startSignal = startSignal;
    }

    @Override
    public CloseFuture call() throws Exception {
      startSignal.await();
      try {
        return session.closeAsync();
      } finally {
        session = null;
      }
    }
  }
}
