/*
 * 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.ConditionChecker.check;
import static org.assertj.core.api.Assertions.assertThat;
import static org.mockito.Mockito.spy;
import static org.mockito.Mockito.timeout;
import static org.mockito.Mockito.verify;
import static org.testng.Assert.fail;

import com.datastax.driver.core.AbstractReconnectionHandler.HandlerFuture;
import com.datastax.driver.core.AbstractReconnectionHandlerTest.MockReconnectionWork.ReconnectBehavior;
import com.datastax.driver.core.exceptions.ConnectionException;
import com.datastax.driver.core.exceptions.UnsupportedProtocolVersionException;
import com.datastax.driver.core.policies.ReconnectionPolicy.ReconnectionSchedule;
import com.google.common.util.concurrent.Futures;
import com.google.common.util.concurrent.ListenableFuture;
import com.google.common.util.concurrent.SettableFuture;
import java.net.InetSocketAddress;
import java.util.concurrent.Callable;
import java.util.concurrent.CyclicBarrier;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicReference;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.testng.annotations.AfterMethod;
import org.testng.annotations.BeforeMethod;
import org.testng.annotations.Test;

public class AbstractReconnectionHandlerTest {
  private static final Logger logger =
      LoggerFactory.getLogger(AbstractReconnectionHandlerTest.class);

  ScheduledExecutorService executor;
  MockReconnectionSchedule schedule;
  MockReconnectionWork work;
  AtomicReference<ListenableFuture<?>> future = new AtomicReference<ListenableFuture<?>>();
  AbstractReconnectionHandler handler;
  Callable<Boolean> nextTryAssigned =
      new Callable<Boolean>() {
        @Override
        public Boolean call() throws Exception {
          return handler.handlerFuture.nextTry != null;
        }
      };

  @BeforeMethod(groups = {"unit", "short"})
  public void setup() {
    executor = spy(Executors.newScheduledThreadPool(2));
    schedule = new MockReconnectionSchedule();
    work = new MockReconnectionWork();
    future.set(null);
    handler =
        new AbstractReconnectionHandler("test", executor, schedule, future) {
          @Override
          protected Connection tryReconnect()
              throws ConnectionException, InterruptedException, UnsupportedProtocolVersionException,
                  ClusterNameMismatchException {
            return work.tryReconnect();
          }

          @Override
          protected void onReconnection(Connection connection) {
            work.onReconnection();
          }
        };
  }

  @AfterMethod(
      groups = {"unit", "short"},
      alwaysRun = true)
  public void tearDown() {
    if (future.get() != null) future.get().cancel(false);
    executor.shutdownNow();
  }

  @Test(groups = "unit")
  public void should_complete_if_first_reconnection_succeeds() {
    handler.start();

    assertThat(future.get()).isNotNull();
    assertThat(future.get().isDone()).isFalse();

    schedule.tick();
    work.nextReconnect = ReconnectBehavior.SUCCEED;
    work.tick();

    waitForCompletion();

    assertThat(work.success).isTrue();
    assertThat(work.tries).isEqualTo(1);
    assertThat(future.get()).isNull();
  }

  @Test(groups = "unit")
  public void should_retry_until_success() {
    handler.start();

    int simulatedErrors = 10;
    for (int i = 0; i < simulatedErrors; i++) {
      schedule.tick();
      work.nextReconnect = ReconnectBehavior.THROW_EXCEPTION;
      work.tick();
      assertThat(work.success).isFalse();
      assertThat(future.get().isDone()).isFalse();
    }

    schedule.tick();
    work.nextReconnect = ReconnectBehavior.SUCCEED;
    work.tick();

    waitForCompletion();

    assertThat(work.success).isTrue();
    assertThat(work.tries).isEqualTo(simulatedErrors + 1);
    assertThat(future.get()).isNull();
  }

  @Test(groups = "unit")
  public void should_stop_if_cancelled_before_first_attempt() {
    schedule.delay = 10 * 1000; // give ourselves time to cancel
    handler.start();
    schedule.tick();

    future.get().cancel(false);

    waitForCompletion();

    assertThat(work.success).isFalse();
    assertThat(work.tries).isEqualTo(0);
    assertThat(future.get().isCancelled()).isTrue();
  }

  @Test(groups = "short")
  public void should_stop_if_cancelled_between_attempts() {
    handler.start();

    // Wait for the initial schedule of a reconnect.
    verify(executor, timeout(10000)).schedule(handler, 0, TimeUnit.MILLISECONDS);

    // Force a failed reconnect.
    schedule.tick();
    work.nextReconnect = ReconnectBehavior.THROW_EXCEPTION;
    // Tick work, should trigger the barrier in tryReconnect.
    work.tick();

    // Tick schedule, should cause nextDelayMs to proceed, reconnect handler will call reschedule.
    schedule.delay = 3000;
    schedule.tick();

    // Ensure reconnect is scheduled (slight timing window after handling failed reconnect
    // and scheduling next reconnect).
    verify(executor, timeout(10000)).schedule(handler, schedule.delay, TimeUnit.MILLISECONDS);

    // Wait until nextTry is assigned after schedule completes.
    check().before(10000).that(nextTryAssigned).becomesTrue();

    future.get().cancel(false);

    // Should immediately return as the future was cancelled while the task was scheduled.
    waitForCompletion();

    assertThat(work.success).isFalse();
    // Should have had 1 failed attempt, no second attempt since cancelled.
    assertThat(work.tries).isEqualTo(1);

    // The future will be marked cancelled and thus not executed.
    ListenableFuture<?> currentAttempt = future.get();

    assertThat(currentAttempt).isInstanceOf(HandlerFuture.class);
    HandlerFuture handlerFuture = (HandlerFuture) currentAttempt;
    assertThat(handlerFuture.isCancelled());

    // The next try should also be cancelled.
    assertThat(handlerFuture.nextTry).isNotNull();
    assertThat(handlerFuture.nextTry.isCancelled());
  }

  @Test(groups = "unit")
  public void should_complete_if_cancelled_during_successful_reconnect()
      throws InterruptedException {
    handler.start();

    schedule.tick();
    work.nextReconnect = ReconnectBehavior.SUCCEED;

    // short pause to make sure we are in the middle of the handler's run method (it checks
    // if the future is cancelled at the beginning)
    TimeUnit.MILLISECONDS.sleep(100);
    // don't force interruption because that's what the production code does
    future.get().cancel(false);

    work.tick();

    waitForCompletion();

    assertThat(work.success).isTrue();
    assertThat(work.tries).isEqualTo(1);
  }

  @Test(groups = "unit")
  public void should_stop_if_cancelled_during_failed_reconnect() throws InterruptedException {
    handler.start();

    schedule.tick();
    work.nextReconnect = ReconnectBehavior.THROW_EXCEPTION;

    // short pause to make sure we are in the middle of the handler's run method (it checks
    // if the future is cancelled at the beginning)
    TimeUnit.MILLISECONDS.sleep(100);
    // don't force interruption because that's what the production code does
    future.get().cancel(false);

    work.tick();

    // Need to
    schedule.tick();

    waitForCompletion();

    assertThat(work.success).isFalse();
    assertThat(work.tries).isEqualTo(1);
  }

  @Test(groups = "unit")
  public void should_yield_to_another_running_handler() {
    // Set an uncompleted future, representing a running handler
    future.set(SettableFuture.create());

    handler.start();

    // Increase the delay to make sure that the first attempt does not start before the check
    // for cancellation (which would require calling work.tick())
    schedule.delay = 5000;
    schedule.tick();

    waitForCompletion();

    assertThat(work.success).isFalse();
  }

  /**
   * Note: a handler that succeeds immediately resets the future to null, so there is a very small
   * window of opportunity for this scenario. Therefore we consider that if we find a completed
   * future, the connection was successfully re-established a few milliseconds ago, so we don't
   * start another attempt.
   */
  @Test(groups = "unit")
  public void should_yield_to_another_handler_that_just_succeeded() {
    future.set(Futures.immediateCheckedFuture(null));

    handler.start();

    schedule.tick();

    waitForCompletion();

    assertThat(work.success).isFalse();
  }

  @Test(groups = "unit")
  public void should_run_if_another_handler_was_cancelled() {
    future.set(Futures.immediateCancelledFuture());

    handler.start();

    schedule.tick();
    work.nextReconnect = ReconnectBehavior.SUCCEED;
    work.tick();

    waitForCompletion();

    assertThat(work.success).isTrue();
    assertThat(work.tries).isEqualTo(1);
    assertThat(future.get()).isNull();
  }

  /**
   * A reconnection schedule that allows manually setting the delay.
   *
   * <p>To make testing easier, nextDelay blocks until tick() is called from the main thread.
   */
  static class MockReconnectionSchedule implements ReconnectionSchedule {
    volatile long delay;
    private final CyclicBarrier barrier = new CyclicBarrier(2);

    // Hack to work around the fact that the first call to nextDelayMs is synchronous
    private volatile boolean firstDelay = true;
    private volatile boolean firstTick = true;

    @Override
    public long nextDelayMs() {
      if (firstDelay) firstDelay = false;
      else {
        logger.debug("in schedule, waiting for tick from main thread");
        try {
          barrier.await(10, TimeUnit.SECONDS);
          logger.debug("in schedule, got tick from main thread, proceeding");
        } catch (Exception e) {
          fail("Error while waiting for tick", e);
        }
      }
      logger.debug("in schedule, returning {}", delay);
      return delay;
    }

    public void tick() {
      if (firstTick) firstTick = false;
      else {
        logger.debug("send tick to schedule");
        try {
          barrier.await(10, TimeUnit.SECONDS);
        } catch (Exception e) {
          fail("Error while sending tick, no thread was waiting", e);
        }
        barrier.reset();
      }
    }
  }

  /**
   * Simulates the work done by the overridable methods of the handler.
   *
   * <p>Allows choosing whether the next reconnect will succeed or throw an exception. To make
   * testing easier, tryReconnect blocks until tick() is called from the main thread.
   */
  static class MockReconnectionWork {
    enum ReconnectBehavior {
      SUCCEED,
      THROW_EXCEPTION
    }

    private final CyclicBarrier barrier = new CyclicBarrier(2);

    volatile ReconnectBehavior nextReconnect;

    volatile int tries = 0;
    volatile boolean success = false;

    protected Connection tryReconnect() throws ConnectionException {
      tries += 1;
      logger.debug("in reconnection work, wait for tick from main thread");
      try {
        barrier.await(60, TimeUnit.SECONDS);
        logger.debug("in reconnection work, got tick from main thread, proceeding");
      } catch (Exception e) {
        fail("Error while waiting for tick", e);
      }
      switch (nextReconnect) {
        case SUCCEED:
          logger.debug("simulate reconnection success");
          return null;
        case THROW_EXCEPTION:
          logger.debug("simulate reconnection error");
          throw new ConnectionException(
              EndPoints.forAddress(new InetSocketAddress(8888)),
              "Simulated exception from mock reconnection");
        default:
          throw new AssertionError();
      }
    }

    public void tick() {
      logger.debug("send tick to reconnection work");
      try {
        barrier.await(60, TimeUnit.SECONDS);
      } catch (Exception e) {
        fail("Error while sending tick, no thread was waiting", e);
      }
      barrier.reset();
    }

    protected void onReconnection() {
      success = true;
    }
  }

  private void waitForCompletion() {
    executor.shutdown();
    try {
      boolean shutdown = executor.awaitTermination(30, TimeUnit.SECONDS);
      if (!shutdown) fail("executor ran for longer than expected");
    } catch (InterruptedException e) {
      fail("Interrupted while waiting for executor to shutdown");
    }
  }
}
