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.List;
  
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  import accord.utils.Invariants;
  import org.apache.cassandra.distributed.Cluster;
  import org.apache.cassandra.distributed.api.ConsistencyLevel;
  import org.apache.cassandra.distributed.api.ICluster;
  import org.apache.cassandra.distributed.api.ICoordinator;
  import org.apache.cassandra.distributed.api.IInstance;
  import org.apache.cassandra.harry.SchemaSpec;
  import org.apache.cassandra.harry.op.Visit;
  import org.apache.cassandra.harry.op.Operations;
  import org.apache.cassandra.harry.model.Model;
  import org.apache.cassandra.harry.model.QuiescentChecker;
  import org.apache.cassandra.harry.model.TokenPlacementModel;
  import org.apache.cassandra.harry.util.ByteUtils;
  import org.apache.cassandra.harry.util.TokenUtil;
  
  /**
   * Executes all modifications queries with NODE_LOCAL. Executes all validation queries
   * with NODE_LOCAL _on each node_.
   */
  public class RingAwareInJvmDTestVisitExecutor extends InJvmDTestVisitExecutor
  {
      private static final Logger logger = LoggerFactory.getLogger(InJvmDTestVisitExecutor.class);
  
      private final TokenPlacementModel.ReplicationFactor rf;
  
      private RingAwareInJvmDTestVisitExecutor(SchemaSpec schema,
                                               DataTracker dataTracker,
                                               Model model,
                                               ICluster<?> cluster,
                                               NodeSelector nodeSelector,
                                               PageSizeSelector pageSizeSelector,
                                               RetryPolicy retryPolicy,
                                               ConsistencyLevelSelector consistencyLevel,
                                               TokenPlacementModel.ReplicationFactor rf,
                                               QueryBuildingVisitExecutor.WrapQueries wrapQueries)
      {
          super(schema, dataTracker, model, cluster, nodeSelector, pageSizeSelector, retryPolicy, consistencyLevel, QueryBuildingVisitExecutor.WrapQueries.UNLOGGED_BATCH);
          this.rf = rf;
      }
  
      protected TokenPlacementModel.ReplicatedRanges getRing()
      {
          IInstance node = ((Cluster)cluster).firstAlive();
          ICoordinator coordinator = node.coordinator();
          List<TokenPlacementModel.Node> other = TokenPlacementModel.peerStateToNodes(coordinator.execute("select peer, tokens, data_center, rack from system.peers", ConsistencyLevel.ONE));
          List<TokenPlacementModel.Node> self = TokenPlacementModel.peerStateToNodes(coordinator.execute("select broadcast_address, tokens, data_center, rack from system.local", ConsistencyLevel.ONE));
          List<TokenPlacementModel.Node> all = new ArrayList<>();
          all.addAll(self);
          all.addAll(other);
          all.sort(TokenPlacementModel.Node::compareTo);
          return rf.replicate(all);
      }
  
      public List<TokenPlacementModel.Replica> getReplicasFor(long pd)
      {
          return getRing().replicasFor(token(pd));
      }
  
      protected long token(long pd)
      {
          return TokenUtil.token(ByteUtils.compose(ByteUtils.objectsToBytes(schema.valueGenerators.pkGen().inflate(pd))));
      }
  
      @Override
      protected void executeValidatingVisit(Visit visit, List<Operations.SelectStatement> selects, CompiledStatement statement)
      {
          try
          {
              for (TokenPlacementModel.Replica replica : getReplicasFor(selects.get(0).pd))
              {
                  IInstance instance = cluster
                                       .stream()
                                       .filter((n) -> n.config().broadcastAddress().toString().contains(replica.node().id()))
                                       .findFirst()
                                       .get();
                  ConsistencyLevel consistencyLevel = this.consistencyLevel.consistencyLevel(visit);
                  int pageSize = PageSizeSelector.NO_PAGING;
                  if (consistencyLevel != ConsistencyLevel.NODE_LOCAL)
                      pageSize = pageSizeSelector.pages(visit);
                  List<ResultSetRow> resultSetRow = executeWithResult(visit, instance.config().num(), pageSize, statement, consistencyLevel);
                  model.validate(selects.get(0), resultSetRow);
              }
          }
          catch (Throwable t)
          {
              throw new AssertionError(String.format("Caught an exception while validating %s:\n%s", selects.get(0), statement), t);
          }
      }
  
      @Override
      protected void executeWithoutResult(Visit visit, CompiledStatement statement)
      {
          try
          {
              Invariants.require(visit.visitedPartitions.size() == 1,
                                    "Ring aware executor can only read and write one partition at a time");
              for (TokenPlacementModel.Replica replica : getReplicasFor(visit.visitedPartitions.iterator().next().longValue()))
              {
                  IInstance instance = cluster
                                       .stream()
                                       .filter((n) -> n.config().broadcastAddress().toString().contains(replica.node().id()))
                                       .findFirst()
                                       .get();
                  executeWithoutResult(visit, instance.config().num(), statement, consistencyLevel.consistencyLevel(visit));
              }
          }
          catch (Throwable t)
          {
              throw new AssertionError(String.format("Caught an exception while validating %s", statement), t);
          }
      }
  
      public static Builder builder()
      {
          return new Builder();
      }
  
      public static class Builder extends InJvmDTestVisitExecutor.Builder
      {
          protected TokenPlacementModel.ReplicationFactor rf;
  
          public Builder replicationFactor(TokenPlacementModel.ReplicationFactor rf)
          {
              this.rf = rf;
              return this;
          }
  
          @Override
          public Builder wrapQueries(QueryBuildingVisitExecutor.WrapQueries wrapQueries)
          {
              super.wrapQueries(wrapQueries);
              return this;
          }
  
          @Override
          public Builder consistencyLevel(ConsistencyLevel consistencyLevel)
          {
              super.consistencyLevel(consistencyLevel);
              return this;
          }
  
          @Override
          public Builder nodeSelector(NodeSelector nodeSelector)
          {
              super.nodeSelector(nodeSelector);
              return this;
          }
  
          @Override
          public Builder pageSizeSelector(PageSizeSelector pageSizeSelector)
          {
              super.pageSizeSelector(pageSizeSelector);
              return this;
          }
  
          @Override
          public Builder retryPolicy(RetryPolicy retryPolicy)
          {
              super.retryPolicy(retryPolicy);
              return this;
          }
  
          @Override
          protected void setDefaults(SchemaSpec schema, ICluster<?> cluster)
          {
              super.setDefaults(schema, cluster);
              if (this.rf == null)
              {
                  this.rf = new TokenPlacementModel.SimpleReplicationFactor(Math.min(3, cluster.size()));
              }
          }
  
          public RingAwareInJvmDTestVisitExecutor build(SchemaSpec schema, Model.Replay replay, ICluster<?> cluster)
          {
              DataTracker tracker = new DataTracker.SequentialDataTracker();
              Model model = new QuiescentChecker(schema.valueGenerators, tracker, replay);
              return build(schema, tracker, model, cluster);
          }
  
          public RingAwareInJvmDTestVisitExecutor build(SchemaSpec schema, DataTracker tracker, Model model, ICluster<?> cluster)
          {
              setDefaults(schema, cluster);
              return new RingAwareInJvmDTestVisitExecutor(schema, tracker, model, cluster,
                                                          nodeSelector, pageSizeSelector, retryPolicy, consistencyLevel, rf, wrapQueries);
          }
      }
  }