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 com.datastax.oss.driver.internal.core.loadbalancing;
  
  import com.datastax.oss.driver.api.core.ConsistencyLevel;
  import com.datastax.oss.driver.api.core.CqlIdentifier;
  import com.datastax.oss.driver.api.core.config.DefaultDriverOption;
  import com.datastax.oss.driver.api.core.config.DriverExecutionProfile;
  import com.datastax.oss.driver.api.core.context.DriverContext;
  import com.datastax.oss.driver.api.core.cql.Statement;
  import com.datastax.oss.driver.api.core.loadbalancing.LoadBalancingPolicy;
  import com.datastax.oss.driver.api.core.loadbalancing.NodeDistance;
  import com.datastax.oss.driver.api.core.loadbalancing.NodeDistanceEvaluator;
  import com.datastax.oss.driver.api.core.metadata.Node;
  import com.datastax.oss.driver.api.core.metadata.NodeState;
  import com.datastax.oss.driver.api.core.metadata.TokenMap;
  import com.datastax.oss.driver.api.core.metadata.token.Token;
  import com.datastax.oss.driver.api.core.session.Request;
  import com.datastax.oss.driver.api.core.session.Session;
  import com.datastax.oss.driver.internal.core.context.InternalDriverContext;
  import com.datastax.oss.driver.internal.core.loadbalancing.helper.DefaultNodeDistanceEvaluatorHelper;
  import com.datastax.oss.driver.internal.core.loadbalancing.helper.OptionalLocalDcHelper;
  import com.datastax.oss.driver.internal.core.loadbalancing.nodeset.DcAgnosticNodeSet;
  import com.datastax.oss.driver.internal.core.loadbalancing.nodeset.MultiDcNodeSet;
  import com.datastax.oss.driver.internal.core.loadbalancing.nodeset.NodeSet;
  import com.datastax.oss.driver.internal.core.loadbalancing.nodeset.SingleDcNodeSet;
  import com.datastax.oss.driver.internal.core.util.ArrayUtils;
  import com.datastax.oss.driver.internal.core.util.collection.CompositeQueryPlan;
  import com.datastax.oss.driver.internal.core.util.collection.LazyQueryPlan;
  import com.datastax.oss.driver.internal.core.util.collection.QueryPlan;
  import com.datastax.oss.driver.internal.core.util.collection.SimpleQueryPlan;
  import com.datastax.oss.driver.shaded.guava.common.base.Predicates;
  import com.datastax.oss.driver.shaded.guava.common.collect.ImmutableMap;
  import com.datastax.oss.driver.shaded.guava.common.collect.Lists;
  import com.datastax.oss.driver.shaded.guava.common.collect.Sets;
  import edu.umd.cs.findbugs.annotations.NonNull;
  import edu.umd.cs.findbugs.annotations.Nullable;
  import java.nio.ByteBuffer;
  import java.util.Collections;
  import java.util.LinkedHashSet;
  import java.util.List;
  import java.util.Map;
  import java.util.Objects;
  import java.util.Optional;
  import java.util.Queue;
  import java.util.Set;
  import java.util.UUID;
  import java.util.concurrent.atomic.AtomicInteger;
  import java.util.function.IntUnaryOperator;
  import net.jcip.annotations.ThreadSafe;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  /**
   * A basic implementation of {@link LoadBalancingPolicy} that can serve as a building block for more
   * advanced use cases.
   *
   * <p>To activate this policy, modify the {@code basic.load-balancing-policy} section in the driver
   * configuration, for example:
   *
   * <pre>
   * datastax-java-driver {
   *   basic.load-balancing-policy {
   *     class = BasicLoadBalancingPolicy
   *     local-datacenter = datacenter1 # optional
   *   }
   * }
   * </pre>
   *
   * See {@code reference.conf} (in the manual or core driver JAR) for more details.
   *
   * <p><b>Local datacenter</b>: This implementation will only define a local datacenter if it is
   * explicitly set either through configuration or programmatically; if the local datacenter is
   * unspecified, this implementation will effectively act as a datacenter-agnostic load balancing
   * policy and will consider all nodes in the cluster when creating query plans, regardless of their
   * datacenter.
   *
   * <p><b>Query plan</b>: This implementation prioritizes replica nodes over non-replica ones; if
   * more than one replica is available, the replicas will be shuffled. Non-replica nodes will be
   * included in a round-robin fashion. If the local datacenter is defined (see above), query plans
   * will only include local nodes, never remote ones; if it is unspecified however, query plans may
   * contain nodes from different datacenters.
   *
   * <p><b>This class is not recommended for normal users who should always prefer {@link
   * DefaultLoadBalancingPolicy}</b>.
   */
  @ThreadSafe
  public class BasicLoadBalancingPolicy implements LoadBalancingPolicy {
  
    private static final Logger LOG = LoggerFactory.getLogger(BasicLoadBalancingPolicy.class);
  
    protected static final IntUnaryOperator INCREMENT = i -> (i == Integer.MAX_VALUE) ? 0 : i + 1;
    private static final Object[] EMPTY_NODES = new Object[0];
  
    @NonNull protected final InternalDriverContext context;
    @NonNull protected final DriverExecutionProfile profile;
    @NonNull protected final String logPrefix;
  
    protected final AtomicInteger roundRobinAmount = new AtomicInteger();
  
    private final int maxNodesPerRemoteDc;
    private final boolean allowDcFailoverForLocalCl;
    private final ConsistencyLevel defaultConsistencyLevel;
  
    // private because they should be set in init() and never be modified after
    private volatile DistanceReporter distanceReporter;
    private volatile NodeDistanceEvaluator nodeDistanceEvaluator;
    private volatile String localDc;
    private volatile NodeSet liveNodes;
    private final LinkedHashSet<String> preferredRemoteDcs;
  
    public BasicLoadBalancingPolicy(@NonNull DriverContext context, @NonNull String profileName) {
      this.context = (InternalDriverContext) context;
      profile = context.getConfig().getProfile(profileName);
      logPrefix = context.getSessionName() + "|" + profileName;
      maxNodesPerRemoteDc =
          profile.getInt(DefaultDriverOption.LOAD_BALANCING_DC_FAILOVER_MAX_NODES_PER_REMOTE_DC);
      allowDcFailoverForLocalCl =
          profile.getBoolean(
              DefaultDriverOption.LOAD_BALANCING_DC_FAILOVER_ALLOW_FOR_LOCAL_CONSISTENCY_LEVELS);
      defaultConsistencyLevel =
          this.context
              .getConsistencyLevelRegistry()
              .nameToLevel(profile.getString(DefaultDriverOption.REQUEST_CONSISTENCY));
  
      preferredRemoteDcs =
          new LinkedHashSet<>(
              profile.getStringList(
                  DefaultDriverOption.LOAD_BALANCING_DC_FAILOVER_PREFERRED_REMOTE_DCS));
    }
  
    /**
     * Returns the local datacenter name, if known; empty otherwise.
     *
     * <p>When this method returns null, then datacenter awareness is completely disabled. All
     * non-ignored nodes will be considered "local" regardless of their actual datacenters, and will
     * have equal chances of being selected for query plans.
     *
     * <p>After the policy is {@linkplain #init(Map, DistanceReporter) initialized} this method will
     * return the local datacenter that was discovered by calling {@link #discoverLocalDc(Map)}.
     * Before initialization, this method always returns null.
     */
    @Nullable
    public String getLocalDatacenter() {
      return localDc;
    }
  
    @NonNull
    @Override
    public Map<String, ?> getStartupConfiguration() {
      ImmutableMap.Builder<String, Object> builder = new ImmutableMap.Builder<>();
      if (localDc != null) {
        builder.put("localDc", localDc);
      } else {
        // Local data center may not be discovered prior to connection pool initialization.
        // In such scenario, return configured local data center name.
        // Note that when using DC inferring load balancing policy, startup configuration
        // may not show local DC name, because it will be discovered only once control connection
        // is established and datacenter of contact points known.
        Optional<String> configuredDc =
            new OptionalLocalDcHelper(context, profile, logPrefix).configuredLocalDc();
        configuredDc.ifPresent(d -> builder.put("localDc", d));
      }
      if (!preferredRemoteDcs.isEmpty()) {
        builder.put("preferredRemoteDcs", preferredRemoteDcs);
      }
      if (allowDcFailoverForLocalCl) {
        builder.put("allowDcFailoverForLocalCl", allowDcFailoverForLocalCl);
      }
      if (maxNodesPerRemoteDc > 0) {
        builder.put("maxNodesPerRemoteDc", maxNodesPerRemoteDc);
      }
      return ImmutableMap.of(BasicLoadBalancingPolicy.class.getSimpleName(), builder.build());
    }
  
    /** @return The nodes currently considered as live. */
    protected NodeSet getLiveNodes() {
      return liveNodes;
    }
  
    @Override
    public void init(@NonNull Map<UUID, Node> nodes, @NonNull DistanceReporter distanceReporter) {
      this.distanceReporter = distanceReporter;
      localDc = discoverLocalDc(nodes).orElse(null);
      nodeDistanceEvaluator = createNodeDistanceEvaluator(localDc, nodes);
      liveNodes =
          localDc == null
              ? new DcAgnosticNodeSet()
              : maxNodesPerRemoteDc <= 0 ? new SingleDcNodeSet(localDc) : new MultiDcNodeSet();
      for (Node node : nodes.values()) {
        NodeDistance distance = computeNodeDistance(node);
        distanceReporter.setDistance(node, distance);
        if (distance != NodeDistance.IGNORED && node.getState() != NodeState.DOWN) {
          // This includes state == UNKNOWN. If the node turns out to be unreachable, this will be
          // detected when we try to open a pool to it, it will get marked down and this will be
          // signaled back to this policy, which will then remove it from the live set.
          liveNodes.add(node);
        }
      }
    }
  
    /**
     * Returns the local datacenter, if it can be discovered, or returns {@link Optional#empty empty}
     * otherwise.
     *
     * <p>This method is called only once, during {@linkplain LoadBalancingPolicy#init(Map,
     * LoadBalancingPolicy.DistanceReporter) initialization}.
     *
     * <p>Implementors may choose to throw {@link IllegalStateException} instead of returning {@link
     * Optional#empty empty}, if they require a local datacenter to be defined in order to operate
     * properly.
     *
     * <p>If this method returns empty, then datacenter awareness will be completely disabled. All
     * non-ignored nodes will be considered "local" regardless of their actual datacenters, and will
     * have equal chances of being selected for query plans.
     *
     * @param nodes All the nodes that were known to exist in the cluster (regardless of their state)
     *     when the load balancing policy was initialized. This argument is provided in case
     *     implementors need to inspect the cluster topology to discover the local datacenter.
     * @return The local datacenter, or {@link Optional#empty empty} if none found.
     * @throws IllegalStateException if the local datacenter could not be discovered, and this policy
     *     cannot operate without it.
     */
    @NonNull
    protected Optional<String> discoverLocalDc(@NonNull Map<UUID, Node> nodes) {
      return new OptionalLocalDcHelper(context, profile, logPrefix).discoverLocalDc(nodes);
    }
  
    /**
     * Creates a new node distance evaluator to use with this policy.
     *
     * <p>This method is called only once, during {@linkplain LoadBalancingPolicy#init(Map,
     * LoadBalancingPolicy.DistanceReporter) initialization}, and only after local datacenter
     * discovery has been attempted.
     *
     * @param localDc The local datacenter that was just discovered, or null if none found.
     * @param nodes All the nodes that were known to exist in the cluster (regardless of their state)
     *     when the load balancing policy was initialized. This argument is provided in case
     *     implementors need to inspect the cluster topology to create the evaluator.
     * @return the distance evaluator to use.
     */
    @NonNull
    protected NodeDistanceEvaluator createNodeDistanceEvaluator(
        @Nullable String localDc, @NonNull Map<UUID, Node> nodes) {
      return new DefaultNodeDistanceEvaluatorHelper(context, profile, logPrefix)
          .createNodeDistanceEvaluator(localDc, nodes);
    }
  
    @NonNull
    @Override
    public Queue<Node> newQueryPlan(@Nullable Request request, @Nullable Session session) {
      // Take a snapshot since the set is concurrent:
      Object[] currentNodes = liveNodes.dc(localDc).toArray();
  
      Set<Node> allReplicas = getReplicas(request, session);
      int replicaCount = 0; // in currentNodes
  
      if (!allReplicas.isEmpty()) {
        // Move replicas to the beginning
        for (int i = 0; i < currentNodes.length; i++) {
          Node node = (Node) currentNodes[i];
          if (allReplicas.contains(node)) {
            ArrayUtils.bubbleUp(currentNodes, i, replicaCount);
            replicaCount += 1;
          }
        }
  
        if (replicaCount > 1) {
          shuffleHead(currentNodes, replicaCount);
        }
      }
  
      LOG.trace("[{}] Prioritizing {} local replicas", logPrefix, replicaCount);
  
      // Round-robin the remaining nodes
      ArrayUtils.rotate(
          currentNodes,
          replicaCount,
          currentNodes.length - replicaCount,
          roundRobinAmount.getAndUpdate(INCREMENT));
  
      QueryPlan plan = currentNodes.length == 0 ? QueryPlan.EMPTY : new SimpleQueryPlan(currentNodes);
      return maybeAddDcFailover(request, plan);
    }
  
    @NonNull
    protected Set<Node> getReplicas(@Nullable Request request, @Nullable Session session) {
      if (request == null || session == null) {
        return Collections.emptySet();
      }
  
      Optional<TokenMap> maybeTokenMap = context.getMetadataManager().getMetadata().getTokenMap();
      if (!maybeTokenMap.isPresent()) {
        return Collections.emptySet();
      }
  
      // Note: we're on the hot path and the getXxx methods are potentially more than simple getters,
      // so we only call each method when strictly necessary (which is why the code below looks a bit
      // weird).
      CqlIdentifier keyspace;
      Token token;
      ByteBuffer key;
      try {
        keyspace = request.getKeyspace();
        if (keyspace == null) {
          keyspace = request.getRoutingKeyspace();
        }
        if (keyspace == null && session.getKeyspace().isPresent()) {
          keyspace = session.getKeyspace().get();
        }
        if (keyspace == null) {
          return Collections.emptySet();
        }
  
        token = request.getRoutingToken();
        key = (token == null) ? request.getRoutingKey() : null;
        if (token == null && key == null) {
          return Collections.emptySet();
        }
      } catch (Exception e) {
        // Protect against poorly-implemented Request instances
        LOG.error("Unexpected error while trying to compute query plan", e);
        return Collections.emptySet();
      }
  
      TokenMap tokenMap = maybeTokenMap.get();
      return token != null
          ? tokenMap.getReplicas(keyspace, token)
          : tokenMap.getReplicas(keyspace, key);
    }
  
    @NonNull
    protected Queue<Node> maybeAddDcFailover(@Nullable Request request, @NonNull Queue<Node> local) {
      if (maxNodesPerRemoteDc <= 0 || localDc == null) {
        return local;
      }
      if (!allowDcFailoverForLocalCl && request instanceof Statement) {
        Statement<?> statement = (Statement<?>) request;
        ConsistencyLevel consistency = statement.getConsistencyLevel();
        if (consistency == null) {
          consistency = defaultConsistencyLevel;
        }
        if (consistency.isDcLocal()) {
          return local;
        }
      }
      if (preferredRemoteDcs.isEmpty()) {
        return new CompositeQueryPlan(local, buildRemoteQueryPlanAll());
      }
      return new CompositeQueryPlan(local, buildRemoteQueryPlanPreferred());
    }
  
    private QueryPlan buildRemoteQueryPlanAll() {
  
      return new LazyQueryPlan() {
        @Override
        protected Object[] computeNodes() {
  
          Object[] remoteNodes =
              liveNodes.dcs().stream()
                  .filter(Predicates.not(Predicates.equalTo(localDc)))
                  .flatMap(dc -> liveNodes.dc(dc).stream().limit(maxNodesPerRemoteDc))
                  .toArray();
          if (remoteNodes.length == 0) {
            return EMPTY_NODES;
          }
          shuffleHead(remoteNodes, remoteNodes.length);
          return remoteNodes;
        }
      };
    }
  
    private QueryPlan buildRemoteQueryPlanPreferred() {
  
      Set<String> dcs = liveNodes.dcs();
      List<String> orderedDcs = Lists.newArrayListWithCapacity(dcs.size());
      orderedDcs.addAll(preferredRemoteDcs);
      orderedDcs.addAll(Sets.difference(dcs, preferredRemoteDcs));
  
      QueryPlan[] queryPlans =
          orderedDcs.stream()
              .filter(Predicates.not(Predicates.equalTo(localDc)))
              .map(
                  (dc) -> {
                    return new LazyQueryPlan() {
                      @Override
                      protected Object[] computeNodes() {
                        Object[] rv = liveNodes.dc(dc).stream().limit(maxNodesPerRemoteDc).toArray();
                        if (rv.length == 0) {
                          return EMPTY_NODES;
                        }
                        shuffleHead(rv, rv.length);
                        return rv;
                      }
                    };
                  })
              .toArray(QueryPlan[]::new);
  
      return new CompositeQueryPlan(queryPlans);
    }
  
    /** Exposed as a protected method so that it can be accessed by tests */
    protected void shuffleHead(Object[] currentNodes, int headLength) {
      ArrayUtils.shuffleHead(currentNodes, headLength);
    }
  
    @Override
    public void onAdd(@NonNull Node node) {
      NodeDistance distance = computeNodeDistance(node);
      // Setting to a non-ignored distance triggers the session to open a pool, which will in turn
      // set the node UP when the first channel gets opened, then #onUp will be called, and the
      // node will be eventually added to the live set.
      distanceReporter.setDistance(node, distance);
      LOG.debug("[{}] {} was added, setting distance to {}", logPrefix, node, distance);
    }
  
    @Override
    public void onUp(@NonNull Node node) {
      NodeDistance distance = computeNodeDistance(node);
      if (node.getDistance() != distance) {
        distanceReporter.setDistance(node, distance);
      }
      if (distance != NodeDistance.IGNORED && liveNodes.add(node)) {
        LOG.debug("[{}] {} came back UP, added to live set", logPrefix, node);
      }
    }
  
    @Override
    public void onDown(@NonNull Node node) {
      if (liveNodes.remove(node)) {
        LOG.debug("[{}] {} went DOWN, removed from live set", logPrefix, node);
      }
    }
  
    @Override
    public void onRemove(@NonNull Node node) {
      if (liveNodes.remove(node)) {
        LOG.debug("[{}] {} was removed, removed from live set", logPrefix, node);
      }
    }
  
    /**
     * Computes the distance of the given node.
     *
     * <p>This method is called during {@linkplain #init(Map, DistanceReporter) initialization}, when
     * a node {@linkplain #onAdd(Node) is added}, and when a node {@linkplain #onUp(Node) is back UP}.
     */
    protected NodeDistance computeNodeDistance(@NonNull Node node) {
      // We interrogate the custom evaluator every time since it could be dynamic
      // and change its verdict between two invocations of this method.
      NodeDistance distance = nodeDistanceEvaluator.evaluateDistance(node, localDc);
      if (distance != null) {
        return distance;
      }
      // no local DC defined: all nodes are considered LOCAL.
      if (localDc == null) {
        return NodeDistance.LOCAL;
      }
      // otherwise, the node is LOCAL if its datacenter is the local datacenter.
      if (Objects.equals(node.getDatacenter(), localDc)) {
        return NodeDistance.LOCAL;
      }
      // otherwise, the node will be either REMOTE or IGNORED, depending
      // on how many remote nodes we accept per DC.
      if (maxNodesPerRemoteDc > 0) {
        Object[] remoteNodes = liveNodes.dc(node.getDatacenter()).toArray();
        for (int i = 0; i < maxNodesPerRemoteDc; i++) {
          if (i == remoteNodes.length) {
            // there is still room for one more REMOTE node in this DC
            return NodeDistance.REMOTE;
          } else if (remoteNodes[i] == node) {
            return NodeDistance.REMOTE;
          }
        }
      }
      return NodeDistance.IGNORED;
    }
  
    @Override
    public void close() {
      // nothing to do
    }
  }
  