/*
 * 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.policies;

import static com.datastax.driver.core.TestUtils.nonQuietClusterCloseOptions;
import static java.util.concurrent.TimeUnit.SECONDS;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.when;
import static org.scassandra.http.client.PrimingRequest.queryBuilder;
import static org.scassandra.http.client.PrimingRequest.then;
import static org.scassandra.http.client.Result.already_exists;
import static org.scassandra.http.client.Result.invalid;
import static org.scassandra.http.client.Result.read_request_timeout;
import static org.scassandra.http.client.Result.success;
import static org.scassandra.http.client.Result.syntax_error;
import static org.scassandra.http.client.Result.unauthorized;
import static org.scassandra.http.client.Result.unavailable;
import static org.scassandra.http.client.Result.write_request_timeout;

import com.datastax.driver.core.Cluster;
import com.datastax.driver.core.ConditionChecker;
import com.datastax.driver.core.Host;
import com.datastax.driver.core.LatencyTracker;
import com.datastax.driver.core.QueryTracker;
import com.datastax.driver.core.ScassandraCluster;
import com.datastax.driver.core.Session;
import com.datastax.driver.core.SortingLoadBalancingPolicy;
import com.datastax.driver.core.Statement;
import com.datastax.driver.core.exceptions.AlreadyExistsException;
import com.datastax.driver.core.exceptions.InvalidQueryException;
import com.datastax.driver.core.exceptions.ReadTimeoutException;
import com.datastax.driver.core.exceptions.SyntaxError;
import com.datastax.driver.core.exceptions.UnauthorizedException;
import com.datastax.driver.core.exceptions.UnavailableException;
import com.datastax.driver.core.exceptions.WriteTimeoutException;
import java.util.concurrent.Callable;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;
import org.scassandra.http.client.Result;
import org.testng.annotations.AfterMethod;
import org.testng.annotations.BeforeMethod;
import org.testng.annotations.Test;

public class ErrorAwarePolicyIntegrationTest {

  private QueryTracker queryTracker;
  private Clock clock;
  private ScassandraCluster sCluster;
  private AtomicInteger errorCounter;
  private LatencyTracker latencyTracker;

  @BeforeMethod(groups = "short")
  public void setUp() {
    queryTracker = new QueryTracker();
    clock = mock(Clock.class);
    sCluster = ScassandraCluster.builder().withNodes(2).build();
    sCluster.init();
    errorCounter = new AtomicInteger(0);
    latencyTracker =
        new LatencyTracker() {
          @Override
          public void update(
              Host host, Statement statement, Exception exception, long newLatencyNanos) {
            if (exception != null) errorCounter.incrementAndGet();
          }

          @Override
          public void onRegister(Cluster cluster) {}

          @Override
          public void onUnregister(Cluster cluster) {}
        };
    // By default node 1 should always fail, 2 succeed.
    prime(1, unauthorized);
    prime(2, success);
  }

  @AfterMethod(groups = "short")
  public void tearDown() {
    sCluster.stop();
  }

  private Cluster.Builder builder(LoadBalancingPolicy lbp) {
    return Cluster.builder()
        .withNettyOptions(nonQuietClusterCloseOptions)
        .addContactPoints(sCluster.address(1).getAddress())
        .withPort(sCluster.getBinaryPort())
        .withLoadBalancingPolicy(lbp);
  }

  private void prime(int node, Result result) {
    sCluster
        .node(node)
        .primingClient()
        .prime(
            queryBuilder()
                .withQuery(QueryTracker.QUERY)
                .withThen(then().withResult(result))
                .build());
  }

  /**
   * Checks that {@link LatencyTracker#update} was called at least expectedCount times within 5
   * seconds.
   *
   * <p>Note that the usefulness of this is dependent on the {@link
   * ErrorAwarePolicy.PerHostErrorTracker} being invoked before the latency tracker being used to
   * track error invocations in this test. The existing implementation seems to invoke update on
   * latency trackers in order that they are registered.
   *
   * @param expectedCount Expected number of errors to have been invoked.
   */
  private void awaitTrackerUpdate(final int expectedCount) {
    ConditionChecker.check()
        .every(10)
        .that(
            new Callable<Boolean>() {
              @Override
              public Boolean call() throws Exception {
                return errorCounter.get() >= expectedCount;
              }
            })
        .before(5000)
        .becomesTrue();
  }

  private void setTime(long time, TimeUnit timeUnit) {
    when(clock.nanoTime()).thenReturn(TimeUnit.NANOSECONDS.convert(time, timeUnit));
  }

  /**
   * Validates that {@link ErrorAwarePolicy} properly excludes a host after the maximum number of
   * errors is exceeded.
   *
   * <p>This test configures a maximum of 1 error per minute and executes 2 failing queries against
   * host1 during the first 5 simulated seconds. host1 should be excluded as soon as the rolling
   * count updates over the next 5-second interval.
   *
   * <p>It then makes another query and ensures it is executed against host2 and that the response
   * was successful.
   *
   * @jira_ticket JAVA-1055
   * @test_category load_balancing:error_aware
   * @since 3.1.0
   */
  @Test(groups = "short")
  public void should_exclude_host_after_reaching_maximum_errors() throws InterruptedException {
    LoadBalancingPolicy lbp =
        ErrorAwarePolicy.builder(new SortingLoadBalancingPolicy())
            .withMaxErrorsPerMinute(1)
            .withClock(clock)
            .build();

    Cluster cluster = builder(lbp).build();

    try {
      Session session = cluster.connect();
      cluster.register(latencyTracker);

      setTime(0, SECONDS);
      // Make 2 queries producing a count higher than the threshold
      queryTracker.query(session, 2, UnauthorizedException.class, sCluster.address(1));
      awaitTrackerUpdate(2);

      // Advance time so that RollingCount ticks and updates its count.
      setTime(5, SECONDS);

      // The next query should succeed and hit node 2 since node 1 is now ignored.
      queryTracker.query(session, 1, sCluster.address(2));
    } finally {
      cluster.close();
    }
  }

  /**
   * Validates that {@link ErrorAwarePolicy} will include a previously excluded host after the
   * configured retry period has elapsed.
   *
   * <p>The test executes queries with error to get host1 excluded. It then executes queries over 70
   * simulated seconds and then executes another query after this time has elapsed and ensures that
   * the next query execution uses host1.
   *
   * @jira_ticket JAVA-1055
   * @test_category load_balancing:error_aware
   * @since 3.1.0
   */
  @Test(groups = "short")
  public void should_resurrect_host_after_retry_period() throws InterruptedException {
    LoadBalancingPolicy lbp =
        ErrorAwarePolicy.builder(new SortingLoadBalancingPolicy())
            .withMaxErrorsPerMinute(1)
            .withRetryPeriod(70, SECONDS)
            .withClock(clock)
            .build();

    Cluster cluster = builder(lbp).build();
    try {
      Session session = cluster.connect();
      cluster.register(latencyTracker);

      setTime(0, SECONDS);
      // Make 2 queries producing a count higher than the threshold
      queryTracker.query(session, 2, UnauthorizedException.class, sCluster.address(1));
      awaitTrackerUpdate(2);

      // Advance time so that RollingCount ticks and updates its count.
      setTime(5, SECONDS);

      // Execute some queries, these should all succeed and hit host2 since host1 is excluded.
      queryTracker.query(session, 5, sCluster.address(2));

      // Advance time after the retry period
      setTime(75, SECONDS);

      // At this the load balancing policy should resurrect node 1 which will be used and fail.
      queryTracker.query(session, 1, UnauthorizedException.class, sCluster.address(1));

    } finally {
      cluster.close();
    }
  }

  /**
   * Validates that {@link ErrorAwarePolicy} will not penalize errors that are not considered in the
   * default {@link ErrorAwarePolicy.ErrorFilter} implementation.
   *
   * <p>Executes 10 queries with each error type and ensures that host1 is used each time, verifying
   * that it was never excluded.
   *
   * @jira_ticket JAVA-1055
   * @test_category load_balancing:error_aware
   * @since 3.1.0
   */
  @Test(groups = "short")
  public void should_not_penalize_default_ignored_exceptions() throws InterruptedException {
    LoadBalancingPolicy lbp =
        ErrorAwarePolicy.builder(new SortingLoadBalancingPolicy())
            .withMaxErrorsPerMinute(1)
            .withClock(clock)
            .build();

    // Use fall through retry policy so other hosts aren't tried.
    Cluster cluster = builder(lbp).withRetryPolicy(FallthroughRetryPolicy.INSTANCE).build();
    try {
      Session session = cluster.connect();
      cluster.register(latencyTracker);

      setTime(0, SECONDS);
      // TODO: Add Read and Write Failure, FunctionExecution exception when Scassandra supports v4.
      prime(1, read_request_timeout);
      queryTracker.query(session, 10, ReadTimeoutException.class, sCluster.address(1));
      awaitTrackerUpdate(10);

      setTime(5, SECONDS);

      prime(1, write_request_timeout);
      queryTracker.query(session, 10, WriteTimeoutException.class, sCluster.address(1));
      awaitTrackerUpdate(20);

      setTime(10, SECONDS);

      prime(1, unavailable);
      queryTracker.query(session, 10, UnavailableException.class, sCluster.address(1));
      awaitTrackerUpdate(30);

      setTime(15, SECONDS);

      prime(1, already_exists);
      queryTracker.query(session, 10, AlreadyExistsException.class, sCluster.address(1));
      awaitTrackerUpdate(40);

      setTime(20, SECONDS);

      prime(1, invalid);
      queryTracker.query(session, 10, InvalidQueryException.class, sCluster.address(1));
      awaitTrackerUpdate(50);

      setTime(25, SECONDS);

      prime(1, syntax_error);
      queryTracker.query(session, 10, SyntaxError.class, sCluster.address(1));
      awaitTrackerUpdate(60);

      setTime(30, SECONDS);

      // ensure host1 still used after another tick.
      queryTracker.query(session, 10, SyntaxError.class, sCluster.address(1));
    } finally {
      cluster.close();
    }
  }

  /**
   * Validates that {@link ErrorAwarePolicy} will regard a custom {@link
   * ErrorAwarePolicy.ErrorFilter} by only penalizing a node when it produces exceptions that
   * evaluate to true in the filter implementation.
   *
   * <p>It first executes 10 queries with an error type that is not considered and verify that host1
   * is never excluded.
   *
   * <p>It then executes queries with an error type that is considered and verifies that host1 is
   * then excluded and host2 is used instead.
   *
   * @jira_ticket JAVA-1055
   * @test_category load_balancing:error_aware
   * @since 3.1.0
   */
  @Test(groups = "short")
  public void should_only_consider_exceptions_based_on_errors_filter() throws InterruptedException {
    ErrorAwarePolicy.ErrorFilter iqeOnlyFilter =
        new ErrorAwarePolicy.ErrorFilter() {
          @Override
          public boolean shouldConsiderError(Exception e, Host host, Statement statement) {
            return e.getClass().isAssignableFrom(InvalidQueryException.class);
          }
        };

    LoadBalancingPolicy lbp =
        ErrorAwarePolicy.builder(new SortingLoadBalancingPolicy())
            .withMaxErrorsPerMinute(1)
            .withClock(clock)
            .withErrorsFilter(iqeOnlyFilter)
            .build();

    // Use fall through retry policy so other hosts aren't tried.
    Cluster cluster = builder(lbp).withRetryPolicy(FallthroughRetryPolicy.INSTANCE).build();
    try {
      Session session = cluster.connect();
      cluster.register(latencyTracker);

      setTime(0, SECONDS);
      // UnauthorizedException evaluates to false in the filter, so it should not be considered.
      prime(1, unauthorized);
      queryTracker.query(session, 10, UnauthorizedException.class, sCluster.address(1));
      awaitTrackerUpdate(10);

      setTime(5, SECONDS);
      // should still query host1
      queryTracker.query(session, 1, UnauthorizedException.class, sCluster.address(1));

      // InvalidQueryException evaluates to true, so it *should* be considered increment the count
      prime(1, invalid);
      queryTracker.query(session, 2, InvalidQueryException.class, sCluster.address(1));
      awaitTrackerUpdate(13);

      // Advance time so that the rolling count updates the next time we query.
      // The first errors that were considered were at t = 5 seconds and this is when the rolling
      // count was
      // initialized, so we want to be at t > 5 + 5 for the rolling count to update
      setTime(10, SECONDS);

      // The next query should succeed and hit node 2 since node 1 is now ignored.
      queryTracker.query(session, 1, sCluster.address(2));
    } finally {
      cluster.close();
    }
  }

  /**
   * Validates that an {@link ErrorAwarePolicy} configured with its defaults behaves as documented,
   * that being that the maximum number of errors is 1 and the retry period is 120 seconds.
   *
   * @jira_ticket JAVA-1055
   * @test_category load_balancing:error_aware
   * @since 3.1.0
   */
  @Test(groups = "short")
  public void should_regard_defaults() throws InterruptedException {
    LoadBalancingPolicy lbp =
        ErrorAwarePolicy.builder(new SortingLoadBalancingPolicy()).withClock(clock).build();

    Cluster cluster = builder(lbp).build();
    try {
      Session session = cluster.connect();
      cluster.register(latencyTracker);

      setTime(0, SECONDS);
      // Make 2 queries producing a count higher than the threshold
      queryTracker.query(session, 2, UnauthorizedException.class, sCluster.address(1));
      awaitTrackerUpdate(2);

      // Advance time so the rolling count ticks, next query should go to host 2
      setTime(5, SECONDS);
      queryTracker.query(session, 5, sCluster.address(2));

      // Advance clock 30 seconds, this is within the retry period so host 1 should still be
      // excluded.
      setTime(35, SECONDS);
      queryTracker.query(session, 5, sCluster.address(2));

      // At this point 120 seconds have elapsed, the load balancing policy should see that the retry
      // period has
      // elapsed and resurrect node 1 which will be used and fail.
      setTime(125, SECONDS);
      queryTracker.query(session, 1, UnauthorizedException.class, sCluster.address(1));

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