/*
 * 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 static com.datastax.driver.core.TestUtils.nonQuietClusterCloseOptions;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.times;
import static org.mockito.Mockito.verify;
import static org.scassandra.http.client.PrimingRequest.then;

import com.datastax.driver.core.exceptions.DriverException;
import com.datastax.driver.core.policies.ConstantSpeculativeExecutionPolicy;
import com.datastax.driver.core.policies.RetryPolicy;
import com.datastax.driver.core.policies.SpeculativeExecutionPolicy;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableMap;
import java.util.List;
import java.util.Map;
import org.mockito.Mockito;
import org.scassandra.http.client.Consistency;
import org.scassandra.http.client.PrimingRequest;
import org.scassandra.http.client.Result;
import org.testng.annotations.AfterMethod;
import org.testng.annotations.BeforeMethod;
import org.testng.annotations.Test;

public class SpeculativeExecutionTest {
  ScassandraCluster scassandras;

  Cluster cluster = null;
  SortingLoadBalancingPolicy loadBalancingPolicy;
  Metrics.Errors errors;
  Host host1, host2, host3;
  Session session;

  @BeforeMethod(groups = "short")
  public void beforeMethod() {
    scassandras = ScassandraCluster.builder().withNodes(3).build();
    scassandras.init();

    int speculativeExecutionDelay = 200;

    loadBalancingPolicy = new SortingLoadBalancingPolicy();
    cluster =
        Cluster.builder()
            .addContactPoints(scassandras.address(2).getAddress())
            .withPort(scassandras.getBinaryPort())
            .withLoadBalancingPolicy(loadBalancingPolicy)
            .withSpeculativeExecutionPolicy(
                new ConstantSpeculativeExecutionPolicy(speculativeExecutionDelay, 1))
            .withQueryOptions(new QueryOptions().setDefaultIdempotence(true))
            .withRetryPolicy(new CustomRetryPolicy())
            .withNettyOptions(nonQuietClusterCloseOptions)
            .build();

    session = cluster.connect();

    host1 = TestUtils.findHost(cluster, 1);
    host2 = TestUtils.findHost(cluster, 2);
    host3 = TestUtils.findHost(cluster, 3);

    errors = cluster.getMetrics().getErrorMetrics();
  }

  @AfterMethod(groups = "short", alwaysRun = true)
  public void afterMethod() {
    if (cluster != null) cluster.close();
    if (scassandras != null) scassandras.stop();
  }

  @Test(groups = "short")
  public void should_not_start_speculative_execution_if_first_execution_completes_successfully() {
    scassandras
        .node(1)
        .primingClient()
        .prime(
            PrimingRequest.queryBuilder()
                .withQuery("mock query")
                .withThen(then().withRows(row("result", "result1")))
                .build());

    long execStartCount = errors.getSpeculativeExecutions().getCount();

    ResultSet rs = session.execute("mock query");
    Row row = rs.one();

    assertThat(row.getString("result")).isEqualTo("result1");
    assertThat(errors.getSpeculativeExecutions().getCount()).isEqualTo(execStartCount);
    ExecutionInfo executionInfo = rs.getExecutionInfo();
    assertThat(executionInfo.getTriedHosts()).containsOnly(host1);
    assertThat(executionInfo.getQueriedHost()).isEqualTo(host1);
    assertThat(executionInfo.getSpeculativeExecutions()).isEqualTo(0);
    assertThat(executionInfo.getSuccessfulExecutionIndex()).isEqualTo(0);
  }

  @Test(groups = "short")
  public void
      should_not_start_speculative_execution_if_first_execution_retries_but_is_still_fast_enough() {
    // will retry once on this node:
    scassandras
        .node(1)
        .primingClient()
        .prime(
            PrimingRequest.queryBuilder()
                .withQuery("mock query")
                .withConsistency(Consistency.TWO)
                .withThen(then().withResult(Result.read_request_timeout))
                .build());

    scassandras
        .node(1)
        .primingClient()
        .prime(
            PrimingRequest.queryBuilder()
                .withQuery("mock query")
                .withConsistency(Consistency.ONE)
                .withThen(then().withRows(row("result", "result1")))
                .build());

    long execStartCount = errors.getSpeculativeExecutions().getCount();
    long retriesStartCount = errors.getRetriesOnUnavailable().getCount();

    SimpleStatement statement = new SimpleStatement("mock query");
    statement.setConsistencyLevel(ConsistencyLevel.TWO);
    ResultSet rs = session.execute(statement);
    Row row = rs.one();

    assertThat(row).isNotNull();
    assertThat(row.getString("result")).isEqualTo("result1");
    assertThat(errors.getSpeculativeExecutions().getCount()).isEqualTo(execStartCount);
    assertThat(errors.getRetriesOnReadTimeout().getCount()).isEqualTo(retriesStartCount + 1);
    ExecutionInfo executionInfo = rs.getExecutionInfo();
    assertThat(executionInfo.getTriedHosts()).containsOnly(host1);
    assertThat(executionInfo.getQueriedHost()).isEqualTo(host1);
    assertThat(executionInfo.getSpeculativeExecutions()).isEqualTo(0);
    assertThat(executionInfo.getSuccessfulExecutionIndex()).isEqualTo(0);
  }

  @Test(groups = "short")
  public void should_start_speculative_execution_if_first_execution_takes_too_long() {
    scassandras
        .node(1)
        .primingClient()
        .prime(
            PrimingRequest.queryBuilder()
                .withQuery("mock query")
                .withThen(then().withRows(row("result", "result1")).withFixedDelay(400L))
                .build());

    scassandras
        .node(2)
        .primingClient()
        .prime(
            PrimingRequest.queryBuilder()
                .withQuery("mock query")
                .withThen(then().withRows(row("result", "result2")))
                .build());
    long execStartCount = errors.getSpeculativeExecutions().getCount();

    ResultSet rs = session.execute("mock query");
    Row row = rs.one();

    assertThat(row.getString("result")).isEqualTo("result2");
    assertThat(errors.getSpeculativeExecutions().getCount()).isEqualTo(execStartCount + 1);
    ExecutionInfo executionInfo = rs.getExecutionInfo();
    // triedHosts does not contain host1 because the request to it had not completed yet
    assertThat(executionInfo.getTriedHosts()).containsOnly(host2);
    assertThat(executionInfo.getQueriedHost()).isEqualTo(host2);
    assertThat(executionInfo.getSpeculativeExecutions()).isEqualTo(1);
    assertThat(executionInfo.getSuccessfulExecutionIndex()).isEqualTo(1);
  }

  @Test(groups = "short")
  public void should_start_speculative_execution_on_multiple_hosts_with_zero_delay() {
    Cluster cluster =
        Cluster.builder()
            .addContactPoints(scassandras.address(2).getAddress())
            .withPort(scassandras.getBinaryPort())
            .withLoadBalancingPolicy(loadBalancingPolicy)
            .withSpeculativeExecutionPolicy(new ConstantSpeculativeExecutionPolicy(0, 2))
            .withQueryOptions(new QueryOptions().setDefaultIdempotence(true))
            .withRetryPolicy(new CustomRetryPolicy())
            .withNettyOptions(nonQuietClusterCloseOptions)
            .build();
    Session session = cluster.connect();
    host1 = TestUtils.findHost(cluster, 1);
    host2 = TestUtils.findHost(cluster, 2);
    host3 = TestUtils.findHost(cluster, 3);
    errors = cluster.getMetrics().getErrorMetrics();

    scassandras
        .node(1)
        .primingClient()
        .prime(
            PrimingRequest.queryBuilder()
                .withQuery("mock query")
                .withThen(then().withRows(row("result", "result1")).withFixedDelay(1000L))
                .build());

    scassandras
        .node(2)
        .primingClient()
        .prime(
            PrimingRequest.queryBuilder()
                .withQuery("mock query")
                .withThen(then().withRows(row("result", "result2")).withFixedDelay(1000L))
                .build());

    scassandras
        .node(3)
        .primingClient()
        .prime(
            PrimingRequest.queryBuilder()
                .withQuery("mock query")
                .withThen(then().withRows(row("result", "result3")))
                .build());
    long execStartCount = errors.getSpeculativeExecutions().getCount();

    ResultSet rs = session.execute("mock query");
    Row row = rs.one();

    assertThat(row.getString("result")).isEqualTo("result3");
    assertThat(errors.getSpeculativeExecutions().getCount()).isEqualTo(execStartCount + 2);
    ExecutionInfo executionInfo = rs.getExecutionInfo();
    // triedHosts does not contain host1 because the request to it had not completed yet
    assertThat(executionInfo.getTriedHosts()).containsOnly(host3);
    assertThat(executionInfo.getQueriedHost()).isEqualTo(host3);
    assertThat(executionInfo.getSpeculativeExecutions()).isEqualTo(2);
    assertThat(executionInfo.getSuccessfulExecutionIndex()).isEqualTo(2);
  }

  @Test(groups = "short")
  public void should_wait_until_all_executions_have_finished() {
    // Rely on read timeouts to trigger errors that cause an execution to move to the next node
    cluster.getConfiguration().getSocketOptions().setReadTimeoutMillis(1000);

    scassandras
        .node(1)
        .primingClient()
        // execution1 starts with host1, which will time out at t=1000
        .prime(
            PrimingRequest.queryBuilder()
                .withQuery("mock query")
                .withThen(then().withRows(row("result", "result1")).withFixedDelay(2000L))
                .build());
    // at t=1000, execution1 moves to host3, which eventually succeeds at t=1500
    scassandras
        .node(3)
        .primingClient()
        .prime(
            PrimingRequest.queryBuilder()
                .withQuery("mock query")
                .withThen(then().withRows(row("result", "result3")).withFixedDelay(500L))
                .build());
    // meanwhile, execution2 starts at t=200, using host2 which times out at t=1200
    // at that time, the query plan is empty so execution2 fails
    // The goal of this test is to check that execution2 does not fail the query, since execution1
    // is still running
    scassandras
        .node(2)
        .primingClient()
        .prime(
            PrimingRequest.queryBuilder()
                .withQuery("mock query")
                .withThen(then().withRows(row("result", "result2")).withFixedDelay(2000L))
                .build());
    long execStartCount = errors.getSpeculativeExecutions().getCount();

    ResultSet rs = session.execute("mock query");
    Row row = rs.one();

    assertThat(row.getString("result")).isEqualTo("result3");
    assertThat(errors.getSpeculativeExecutions().getCount()).isEqualTo(execStartCount + 1);
    ExecutionInfo executionInfo = rs.getExecutionInfo();
    assertThat(executionInfo.getTriedHosts()).containsOnly(host1, host2, host3);
    assertThat(executionInfo.getQueriedHost()).isEqualTo(host3);
    assertThat(executionInfo.getSpeculativeExecutions()).isEqualTo(1);
    assertThat(executionInfo.getSuccessfulExecutionIndex()).isEqualTo(0);
  }

  /**
   * Validates that when a Cluster is initialized that {@link
   * SpeculativeExecutionPolicy#init(Cluster)} is called and that when a Cluster is closed {@link
   * SpeculativeExecutionPolicy#close()} is called.
   *
   * @test_category queries:speculative_execution
   * @expected_result init and close are called on cluster init and close.
   * @jira_ticket JAVA-796
   * @since 2.0.11, 2.1.7, 2.2.1
   */
  @Test(groups = "short")
  public void should_init_and_close_policy_on_cluster() {
    SpeculativeExecutionPolicy mockPolicy = mock(SpeculativeExecutionPolicy.class);

    Cluster cluster =
        Cluster.builder()
            .addContactPoints(scassandras.address(2).getAddress())
            .withPort(scassandras.getBinaryPort())
            .withSpeculativeExecutionPolicy(mockPolicy)
            .build();

    verify(mockPolicy, times(0)).init(cluster);
    verify(mockPolicy, times(0)).close();

    try {
      cluster.init();
      verify(mockPolicy, times(1)).init(cluster);
    } finally {
      cluster.close();
      verify(mockPolicy, times(1)).close();
    }
  }

  /**
   * Validates that if a query gets speculatively re-executed, the second execution uses the same
   * default timestamp.
   *
   * @test_category tracing
   * @jira_ticket JAVA-724
   * @expected_result timestamp generator invoked only once for a query that caused two executions.
   */
  @Test(groups = "short")
  public void should_use_same_default_timestamp_for_all_executions() {
    TimestampGenerator timestampGenerator = Mockito.spy(ServerSideTimestampGenerator.INSTANCE);
    Cluster cluster =
        Cluster.builder()
            .addContactPoints(scassandras.address(2).getAddress())
            .withPort(scassandras.getBinaryPort())
            .withLoadBalancingPolicy(loadBalancingPolicy)
            .withTimestampGenerator(timestampGenerator)
            .withSpeculativeExecutionPolicy(new ConstantSpeculativeExecutionPolicy(1, 2))
            .withNettyOptions(nonQuietClusterCloseOptions)
            .build();

    try {
      scassandras
          .node(1)
          .primingClient()
          // execution1 starts with host1, which will time out at t=1000
          .prime(
              PrimingRequest.queryBuilder()
                  .withQuery("mock query")
                  .withThen(then().withRows(row("result", "result1")).withFixedDelay(2000L))
                  .build());

      Session session = cluster.connect();
      Metrics.Errors errors = cluster.getMetrics().getErrorMetrics();

      long execStartCount = errors.getSpeculativeExecutions().getCount();

      SimpleStatement statement = new SimpleStatement("mock query");
      statement.setIdempotent(true);
      session.execute(statement);

      // Should have only requested one timestamp.
      if (errors.getSpeculativeExecutions().getCount() == execStartCount + 1) {
        Mockito.verify(timestampGenerator, times(1)).next();
      }
    } finally {
      cluster.close();
    }
  }

  /**
   * Custom retry policy that retries at ONE on read timeout. This deals with the fact that
   * Scassandra only allows read timeouts with 0 replicas.
   */
  static class CustomRetryPolicy implements RetryPolicy {
    @Override
    public RetryDecision onReadTimeout(
        Statement statement,
        ConsistencyLevel cl,
        int requiredResponses,
        int receivedResponses,
        boolean dataRetrieved,
        int nbRetry) {
      if (nbRetry != 0) return RetryDecision.rethrow();
      return RetryDecision.retry(ConsistencyLevel.ONE);
    }

    @Override
    public RetryDecision onWriteTimeout(
        Statement statement,
        ConsistencyLevel cl,
        WriteType writeType,
        int requiredAcks,
        int receivedAcks,
        int nbRetry) {
      return RetryDecision.rethrow();
    }

    @Override
    public RetryDecision onUnavailable(
        Statement statement,
        ConsistencyLevel cl,
        int requiredReplica,
        int aliveReplica,
        int nbRetry) {
      return RetryDecision.rethrow();
    }

    @Override
    public RetryDecision onRequestError(
        Statement statement, ConsistencyLevel cl, DriverException e, int nbRetry) {
      return RetryDecision.tryNextHost(cl);
    }

    @Override
    public void init(Cluster cluster) {}

    @Override
    public void close() {}
  }

  private static List<Map<String, ?>> row(String key, String value) {
    return ImmutableList.<Map<String, ?>>of(ImmutableMap.of(key, value));
  }
}
