content: /*
   * 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 org.apache.cassandra.harry.execution;
  
  import java.util.ArrayList;
  import java.util.Arrays;
  import java.util.Iterator;
  import java.util.List;
  import java.util.function.Function;
  
  import com.google.common.base.Throwables;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  import accord.utils.Invariants;
  import org.apache.cassandra.distributed.api.ConsistencyLevel;
  import org.apache.cassandra.distributed.api.ICluster;
  import org.apache.cassandra.exceptions.RequestTimeoutException;
  import org.apache.cassandra.harry.ColumnSpec;
  import org.apache.cassandra.harry.MagicConstants;
  import org.apache.cassandra.harry.SchemaSpec;
  import org.apache.cassandra.harry.model.Model;
  import org.apache.cassandra.harry.model.QuiescentChecker;
  import org.apache.cassandra.harry.op.Operations;
  import org.apache.cassandra.harry.op.Visit;
  import org.apache.cassandra.utils.AssertionUtils;
  
  import static org.apache.cassandra.harry.MagicConstants.LTS_UNKNOWN;
  import static org.apache.cassandra.harry.MagicConstants.NIL_DESCR;
  import static org.apache.cassandra.harry.MagicConstants.NIL_KEY;
  import static org.apache.cassandra.harry.MagicConstants.UNKNOWN_DESCR;
  import static org.apache.cassandra.harry.MagicConstants.UNSET_DESCR;
  import static org.apache.cassandra.harry.execution.QueryBuildingVisitExecutor.*;
  
  public class InJvmDTestVisitExecutor extends CQLVisitExecutor
  {
      private static final Logger logger = LoggerFactory.getLogger(InJvmDTestVisitExecutor.class);
  
      protected final ICluster<?> cluster;
      protected final ConsistencyLevelSelector consistencyLevel;
  
      protected final NodeSelector nodeSelector;
      protected final PageSizeSelector pageSizeSelector;
      protected final RetryPolicy retryPolicy;
  
      protected InJvmDTestVisitExecutor(SchemaSpec schema,
                                        DataTracker dataTracker,
                                        Model model,
                                        ICluster<?> cluster,
  
                                        NodeSelector nodeSelector,
                                        PageSizeSelector pageSizeSelector,
                                        RetryPolicy retryPolicy,
                                        ConsistencyLevelSelector consistencyLevel,
                                        WrapQueries wrapQueries)
      {
          super(schema, dataTracker, model, new QueryBuildingVisitExecutor(schema, wrapQueries));
          this.cluster = cluster;
          this.consistencyLevel = consistencyLevel;
  
          this.nodeSelector = nodeSelector;
          this.pageSizeSelector = pageSizeSelector;
          this.retryPolicy = retryPolicy;
      }
  
      @Override
      protected List<ResultSetRow> executeWithResult(Visit visit, CompiledStatement statement)
      {
          while (true)
          {
              try
              {
                  ConsistencyLevel consistencyLevel = this.consistencyLevel.consistencyLevel(visit);
                  int pageSize = PageSizeSelector.NO_PAGING;
                  if (consistencyLevel != ConsistencyLevel.NODE_LOCAL)
                      pageSize = pageSizeSelector.pages(visit);
                  return executeWithResult(visit, nodeSelector.select(visit.lts), pageSize, statement, consistencyLevel);
              }
              catch (Throwable t)
              {
                  if (retryPolicy.retry(t))
                      continue;
                  throw t;
              }
          }
      }
  
      protected List<ResultSetRow> executeWithResult(Visit visit, int node, int pageSize, CompiledStatement statement, ConsistencyLevel consistencyLevel)
      {
          Invariants.require(visit.operations.length == 1);
          Object[][] rows;
          if (consistencyLevel == ConsistencyLevel.NODE_LOCAL)
              rows = cluster.get(node).executeInternal(statement.cql(), statement.bindings());
          else
          {
              if (pageSize == PageSizeSelector.NO_PAGING)
                  rows = cluster.coordinator(node).execute(statement.cql(), consistencyLevel, statement.bindings());
              else
                  rows = iterToArr(cluster.coordinator(node)
                                          .executeWithPaging(statement.cql(), consistencyLevel, pageSize, statement.bindings()));
          }
  
          if (logger.isTraceEnabled())
              logger.trace("{} returned {} results", statement, rows.length);
  
          return rowsToResultSet(schema, (Operations.SelectStatement) visit.operations[0], rows);
      }
  
      protected static Object[][] iterToArr(Iterator<Object[]> iter)
      {
          List<Object[]> tmp = new ArrayList<>();
          while (iter.hasNext())
              tmp.add(iter.next());
          return tmp.toArray(new Object[tmp.size()][]);
      }
  
      @Override
      protected void executeWithoutResult(Visit visit, CompiledStatement statement)
      {
          while (true)
          {
              try
              {
                  ConsistencyLevel consistencyLevel = this.consistencyLevel.consistencyLevel(visit);
                  executeWithoutResult(visit, nodeSelector.select(visit.lts), statement, consistencyLevel);
                  return;
              }
              catch (Throwable t)
              {
                  if (retryPolicy.retry(t))
                      continue;
                  throw t;
              }
          }
      }
  
      protected void executeWithoutResult(Visit visit, int node, CompiledStatement statement, ConsistencyLevel consistencyLevel)
      {
          if (consistencyLevel == ConsistencyLevel.NODE_LOCAL)
              cluster.get(node).executeInternal(statement.cql(), statement.bindings());
          else
              cluster.coordinator(node).execute(statement.cql(), consistencyLevel, statement.bindings());
      }
  
      public static List<ResultSetRow> rowsToResultSet(SchemaSpec schema, Operations.SelectStatement select, Object[][] result)
      {
          List<ResultSetRow> rs = new ArrayList<>();
          for (Object[] res : result)
              rs.add(rowToResultSet(schema, select, res));
          return rs;
      }
  
      public static ResultSetRow rowToResultSet(SchemaSpec schema, Operations.SelectStatement select, Object[] result)
      {
          long[] staticColumns = new long[schema.staticColumns.size()];
          long[] regularColumns = new long[schema.regularColumns.size()];
          Arrays.fill(staticColumns, UNKNOWN_DESCR);
          Arrays.fill(regularColumns, UNKNOWN_DESCR);
          // TODO: with timestamps
          long[] staticLts = LTS_UNKNOWN;
          long[] regularLts = LTS_UNKNOWN;
  
          long pd = UNKNOWN_DESCR;
          Operations.Selection selection = Operations.Selection.fromBitSet(select.selection(), schema);
          if (selection.selectsAllOf(schema.partitionKeys))
          {
              Object[] partitionKey = new Object[schema.partitionKeys.size()];
              for (int i = 0; i < schema.partitionKeys.size(); i++)
                  partitionKey[i] = result[selection.indexOf(schema.partitionKeys.get(i))];
  
              pd = schema.valueGenerators.pkGen().deflate(partitionKey);
          }
  
          // Deflate logic for clustering key is a bit more involved, since CK can be nil in case of a single static row.
          long cd = UNKNOWN_DESCR;
          if (selection.selectsAllOf(schema.clusteringKeys))
          {
              Object[] clusteringKey = new Object[schema.clusteringKeys.size()];
              for (int i = 0; i < schema.clusteringKeys.size(); i++)
              {
                  Object v = result[selection.indexOf(schema.clusteringKeys.get(i))];
                  if (v == null)
                  {
                      for (int j = 0; j < schema.clusteringKeys.size(); j++)
                      {
                          Invariants.require(result[selection.indexOf(schema.clusteringKeys.get(j))] == null,
                                                "All elements of clustering key should have been null");
                      }
                      clusteringKey = NIL_KEY;
                      break;
                  }
                  clusteringKey[i] = v;
              }
  
              // Clusterings can not be set to nil, so if we do not see, we assume it is unset
              if (clusteringKey == NIL_KEY)
                  cd = UNSET_DESCR;
              else
                  cd = schema.valueGenerators.ckGen().deflate(clusteringKey);
          }
  
          for (int i = 0; i < schema.regularColumns.size(); i++)
          {
              ColumnSpec<?> column = schema.regularColumns.get(i);
              if (selection.selects(column))
              {
                  Object v = result[selection.indexOf(schema.regularColumns.get(i))];
                  if (v == null)
                      regularColumns[i] = NIL_DESCR;
                  else
                      regularColumns[i] = schema.valueGenerators.regularColumnGen(i).deflate(v);
              }
              else
              {
                  regularLts[i] = UNKNOWN_DESCR;
              }
          }
  
          for (int i = 0; i < schema.staticColumns.size(); i++)
          {
              ColumnSpec<?> column = schema.staticColumns.get(i);
              if (selection.selects(column))
              {
                  Object v = result[selection.indexOf(schema.staticColumns.get(i))];
                  if (v == null)
                      staticColumns[i] = NIL_DESCR;
                  else
                      staticColumns[i] = schema.valueGenerators.staticColumnGen(i).deflate(v);
              }
              else
              {
                  staticColumns[i] = UNKNOWN_DESCR;
              }
          }
  
          if (selection.includeTimestamps())
          {
              long[] slts = new long[schema.staticColumns.size()];
              Arrays.fill(slts, MagicConstants.NO_TIMESTAMP);
              for (int i = 0, sltsBase = schema.allColumnInSelectOrder.size(); i < slts.length && sltsBase + i < result.length; i++)
              {
                  Object v = result[schema.allColumnInSelectOrder.size() + i];
                  if (v != null)
                      slts[i] = (long) v;
              }
  
              long[] lts = new long[schema.regularColumns.size()];
              Arrays.fill(lts, MagicConstants.NO_TIMESTAMP);
              for (int i = 0, ltsBase = schema.allColumnInSelectOrder.size() + slts.length; i < lts.length && ltsBase + i < result.length; i++)
              {
                  Object v = result[ltsBase + i];
                  if (v != null)
                      lts[i] = (long) v;
              }
          }
  
          return new ResultSetRow(pd, cd, staticColumns, staticLts, regularColumns, regularLts);
      }
  
      public interface NodeSelector
      {
          int select(long lts);
      }
  
      public interface PageSizeSelector
      {
          int NO_PAGING = -1;
          int pages(Visit visit);
      }
  
      public interface ConsistencyLevelSelector
      {
          ConsistencyLevel consistencyLevel(Visit visit);
      }
  
      public interface RetryPolicy
      {
          RetryPolicy RETRY_ON_TIMEOUT = (t) -> {
              return t.getMessage() != null &&
                     (t.getMessage().contains("timed out") ||
                      AssertionUtils.isInstanceof(RequestTimeoutException.class)
                                    .matches(Throwables.getRootCause(t)));
          };
          RetryPolicy NO_RETRY = (t) -> false;
          boolean retry(Throwable t);
      }
  
      public static Builder builder()
      {
          return new Builder();
      }
  
      public static class Builder
      {
          protected ConsistencyLevelSelector consistencyLevel = v -> ConsistencyLevel.QUORUM;
  
          protected NodeSelector nodeSelector = null;
          protected PageSizeSelector pageSizeSelector = (lts) -> 1;
          protected RetryPolicy retryPolicy = RetryPolicy.RETRY_ON_TIMEOUT;
          protected WrapQueries wrapQueries = WrapQueries.UNLOGGED_BATCH;
  
          public Builder wrapQueries(WrapQueries wrapQueries)
          {
              this.wrapQueries = wrapQueries;
              return this;
          }
  
          public Builder consistencyLevel(ConsistencyLevel consistencyLevel)
          {
              this.consistencyLevel = v -> consistencyLevel;
              return this;
          }
  
          public Builder consistencyLevel(ConsistencyLevelSelector consistencyLevel)
          {
              this.consistencyLevel = consistencyLevel;
              return this;
          }
  
          public Builder nodeSelector(NodeSelector nodeSelector)
          {
              this.nodeSelector = nodeSelector;
              return this;
          }
  
          public Builder pageSizeSelector(PageSizeSelector pageSizeSelector)
          {
              this.pageSizeSelector = pageSizeSelector;
              return this;
          }
  
          public Builder retryPolicy(RetryPolicy retryPolicy)
          {
              this.retryPolicy = retryPolicy;
              return this;
          }
  
          protected void setDefaults(SchemaSpec schema, ICluster<?> cluster)
          {
              if (nodeSelector == null)
              {
                  this.nodeSelector = new NodeSelector()
                  {
                      final int nodes = cluster.size();
                      long counter = 0;
  
                      @Override
                      public int select(long lts)
                      {
                          return (int) (counter++ % nodes + 1);
                      }
                  };
              }
          }
  
          public InJvmDTestVisitExecutor build(SchemaSpec schema, Model.Replay replay, ICluster<?> cluster)
          {
              setDefaults(schema, cluster);
              DataTracker tracker = new DataTracker.SequentialDataTracker();
              Model model = new QuiescentChecker(schema.valueGenerators, tracker, replay);
              return new InJvmDTestVisitExecutor(schema, tracker, model, cluster,
                                                 nodeSelector, pageSizeSelector, retryPolicy, consistencyLevel, wrapQueries);
          }
  
          public InJvmDTestVisitExecutor build(SchemaSpec schema, ICluster<?> cluster, Function<Builder, InJvmDTestVisitExecutor> overrides)
          {
              setDefaults(schema, cluster);
              return overrides.apply(this);
          }
  
          /**
           * WARNING: highly experimental
           */
          public InJvmDTestVisitExecutor doubleWriting(SchemaSpec schema, Model.Replay replay, ICluster<?> cluster, String secondTable)
          {
              setDefaults(schema, cluster);
              DataTracker tracker = new DataTracker.SequentialDataTracker();
              Model model = new QuiescentChecker(schema.valueGenerators, tracker, replay);
              return new InJvmDTestVisitExecutor(schema, tracker, model, cluster,
                                                 nodeSelector, pageSizeSelector, retryPolicy, consistencyLevel, wrapQueries)
              {
                  @Override
                  protected List<ResultSetRow> executeWithResult(Visit visit, int node, int pageSize, CompiledStatement statement, ConsistencyLevel consistencyLevel)
                  {
                      List<ResultSetRow> rows = super.executeWithResult(visit, node, pageSize, statement, consistencyLevel);
                      List<ResultSetRow> secondOpinion = super.executeWithResult(visit, node, pageSize, statement.withSchema(schema.keyspace, schema.table,
                                                                                                                             schema.keyspace, secondTable),
                                                                                 consistencyLevel);
                      if (!rows.equals(secondOpinion))
                      {
                          logger.debug("Second opinion: ");
                          for (ResultSetRow resultSetRow : secondOpinion)
                              logger.debug(resultSetRow.toString(schema.valueGenerators));
                      }
                      return rows;
                  }
  
                  @Override
                  protected void executeWithoutResult(Visit visit, int node, CompiledStatement statement, ConsistencyLevel consistencyLevel)
                  {
                      super.executeWithoutResult(visit, node, statement, consistencyLevel);
                      super.executeWithoutResult(visit, node, statement.withSchema(schema.keyspace, schema.table,
                                                                                   schema.keyspace, secondTable),
                                                 consistencyLevel);
                  }
              };
          }
      }
  
  }