com.datastax.driver.core.policies.TokenAwarePolicy Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of cassandra-driver Show documentation
Shaded version of DataStax Java Driver for Apache Cassandra
There is a newer version: 3.6.0-1
/*
 *      Copyright (C) 2012-2015 DataStax Inc.
 *
 *   Licensed 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 com.datastax.driver.core.*;
import com.datastax.driver.$internal.com.google.common.collect.AbstractIterator;
import com.datastax.driver.$internal.com.google.common.collect.Lists;

import java.nio.ByteBuffer;
import java.util.*;

/**
 * A wrapper load balancing policy that add token awareness to a child policy.
 * 
 * This policy encapsulates another policy. The resulting policy works in
 * the following way:
 * 

 * the {@code distance} method is inherited from the child policy.
 * the iterator return by the {@code newQueryPlan} method will first
 * return the {@code LOCAL} replicas for the query (based on {@link Statement#getRoutingKey})
 * if possible (i.e. if the query {@code getRoutingKey} method
 * doesn't return {@code null} and if {@link Metadata#getReplicas}
 * returns a non empty set of replicas for that partition key). If no
 * local replica can be either found or successfully contacted, the rest
 * of the query plan will fallback to one of the child policy.
 * 
 * 
 * Do note that only replica for which the child policy {@code distance}
 * method returns {@code HostDistance.LOCAL} will be considered having
 * priority. For example, if you wrap {@link DCAwareRoundRobinPolicy} with this
 * token aware policy, replicas from remote data centers may only be
 * returned after all the host of the local data center.
 */
public class TokenAwarePolicy implements ChainableLoadBalancingPolicy {

    private final LoadBalancingPolicy childPolicy;
    private final boolean shuffleReplicas;
    private volatile Metadata clusterMetadata;
    private volatile ProtocolVersion protocolVersion;
    private volatile CodecRegistry codecRegistry;

    /**
     * Creates a new {@code TokenAware} policy.
     *
     * @param childPolicy     the load balancing policy to wrap with token awareness.
     * @param shuffleReplicas whether to shuffle the replicas returned by {@code getRoutingKey}.
     *                        Note that setting this parameter to {@code true} might decrease the
     *                        effectiveness of caching (especially at consistency level ONE), since
     *                        the same row will be retrieved from any replica (instead of only the
     *                        "primary" replica without shuffling).
     *                        On the other hand, shuffling will better distribute writes, and can
     *                        alleviate hotspots caused by "fat" partitions.
     */
    public TokenAwarePolicy(LoadBalancingPolicy childPolicy, boolean shuffleReplicas) {
        this.childPolicy = childPolicy;
        this.shuffleReplicas = shuffleReplicas;
    }

    /**
     * Creates a new {@code TokenAware} policy with shuffling of replicas.
     *
     * @param childPolicy the load balancing policy to wrap with token
     *                    awareness.
     * @see #TokenAwarePolicy(LoadBalancingPolicy, boolean)
     */
    public TokenAwarePolicy(LoadBalancingPolicy childPolicy) {
        this(childPolicy, true);
    }

    @Override
    public LoadBalancingPolicy getChildPolicy() {
        return childPolicy;
    }

    @Override
    public void init(Cluster cluster, Collection hosts) {
        clusterMetadata = cluster.getMetadata();
        protocolVersion = cluster.getConfiguration().getProtocolOptions().getProtocolVersion();
        codecRegistry = cluster.getConfiguration().getCodecRegistry();
        childPolicy.init(cluster, hosts);
    }

    /**
     * Return the HostDistance for the provided host.
     *
     * @param host the host of which to return the distance of.
     * @return the HostDistance to {@code host} as returned by the wrapped policy.
     */
    @Override
    public HostDistance distance(Host host) {
        return childPolicy.distance(host);
    }

    /**
     * Returns the hosts to use for a new query.
     * 
     * The returned plan will first return replicas (whose {@code HostDistance}
     * for the child policy is {@code LOCAL}) for the query if it can determine
     * them (i.e. mainly if {@code statement.getRoutingKey()} is not {@code null}).
     * Following what it will return the plan of the child policy.
     *
     * @param statement the query for which to build the plan.
     * @return the new query plan.
     */
    @Override
    public Iterator newQueryPlan(final String loggedKeyspace, final Statement statement) {

        ByteBuffer partitionKey = statement.getRoutingKey(protocolVersion, codecRegistry);
        String keyspace = statement.getKeyspace();
        if (keyspace == null)
            keyspace = loggedKeyspace;

        if (partitionKey == null || keyspace == null)
            return childPolicy.newQueryPlan(keyspace, statement);

        final Set replicas = clusterMetadata.getReplicas(Metadata.quote(keyspace), partitionKey);
        if (replicas.isEmpty())
            return childPolicy.newQueryPlan(loggedKeyspace, statement);

        final Iterator iter;
        if (shuffleReplicas) {
            List l = Lists.newArrayList(replicas);
            Collections.shuffle(l);
            iter = l.iterator();
        } else {
            iter = replicas.iterator();
        }

        return new AbstractIterator() {

            private Iterator childIterator;

            @Override
            protected Host computeNext() {
                while (iter.hasNext()) {
                    Host host = iter.next();
                    if (host.isUp() && childPolicy.distance(host) == HostDistance.LOCAL)
                        return host;
                }

                if (childIterator == null)
                    childIterator = childPolicy.newQueryPlan(loggedKeyspace, statement);

                while (childIterator.hasNext()) {
                    Host host = childIterator.next();
                    // Skip it if it was already a local replica
                    if (!replicas.contains(host) || childPolicy.distance(host) != HostDistance.LOCAL)
                        return host;
                }
                return endOfData();
            }
        };
    }

    @Override
    public void onUp(Host host) {
        childPolicy.onUp(host);
    }

    @Override
    public void onDown(Host host) {
        childPolicy.onDown(host);
    }

    @Override
    public void onAdd(Host host) {
        childPolicy.onAdd(host);
    }

    @Override
    public void onRemove(Host host) {
        childPolicy.onRemove(host);
    }

    @Override
    public void close() {
        childPolicy.close();
    }
}