org.apache.solr.cloud.ExclusiveSliceProperty.java Source code

Java tutorial

Introduction

Here is the source code for org.apache.solr.cloud.ExclusiveSliceProperty.java

Source

/*
 * 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.solr.cloud;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.ListIterator;
import java.util.Locale;
import java.util.Map;
import java.util.Random;
import java.util.Set;

import org.apache.commons.lang.StringUtils;
import org.apache.solr.cloud.overseer.ClusterStateMutator;
import org.apache.solr.cloud.overseer.CollectionMutator;
import org.apache.solr.cloud.overseer.SliceMutator;
import org.apache.solr.common.SolrException;
import org.apache.solr.common.cloud.ClusterState;
import org.apache.solr.common.cloud.DocCollection;
import org.apache.solr.common.cloud.Replica;
import org.apache.solr.common.cloud.Slice;
import org.apache.solr.common.cloud.ZkNodeProps;
import org.apache.solr.common.cloud.ZkStateReader;

import static org.apache.solr.cloud.OverseerCollectionMessageHandler.ONLY_ACTIVE_NODES;
import static org.apache.solr.cloud.OverseerCollectionMessageHandler.SHARD_UNIQUE;
import static org.apache.solr.common.params.CollectionParams.CollectionAction.BALANCESHARDUNIQUE;

// Class to encapsulate processing replica properties that have at most one replica hosting a property per slice.
class ExclusiveSliceProperty {
    private ClusterState clusterState;
    private final boolean onlyActiveNodes;
    private final String property;
    private final DocCollection collection;
    private final String collectionName;

    // Key structure. For each node, list all replicas on it regardless of whether they have the property or not.
    private final Map<String, List<SliceReplica>> nodesHostingReplicas = new HashMap<>();
    // Key structure. For each node, a list of the replicas _currently_ hosting the property.
    private final Map<String, List<SliceReplica>> nodesHostingProp = new HashMap<>();
    Set<String> shardsNeedingHosts = new HashSet<>();
    Map<String, Slice> changedSlices = new HashMap<>(); // Work on copies rather than the underlying cluster state.

    private int origMaxPropPerNode = 0;
    private int origModulo = 0;
    private int tmpMaxPropPerNode = 0;
    private int tmpModulo = 0;
    Random rand = new Random();

    private int assigned = 0;

    ExclusiveSliceProperty(ClusterState clusterState, ZkNodeProps message) {
        this.clusterState = clusterState;
        String tmp = message.getStr(ZkStateReader.PROPERTY_PROP);
        if (StringUtils.startsWith(tmp, OverseerCollectionMessageHandler.COLL_PROP_PREFIX) == false) {
            tmp = OverseerCollectionMessageHandler.COLL_PROP_PREFIX + tmp;
        }
        this.property = tmp.toLowerCase(Locale.ROOT);
        collectionName = message.getStr(ZkStateReader.COLLECTION_PROP);

        if (StringUtils.isBlank(collectionName) || StringUtils.isBlank(property)) {
            throw new SolrException(SolrException.ErrorCode.BAD_REQUEST,
                    "Overseer '" + message.getStr(Overseer.QUEUE_OPERATION) + "'  requires both the '"
                            + ZkStateReader.COLLECTION_PROP + "' and '" + ZkStateReader.PROPERTY_PROP
                            + "' parameters. No action taken ");
        }

        Boolean shardUnique = Boolean.parseBoolean(message.getStr(SHARD_UNIQUE));
        if (shardUnique == false && SliceMutator.SLICE_UNIQUE_BOOLEAN_PROPERTIES.contains(this.property) == false) {
            throw new SolrException(SolrException.ErrorCode.BAD_REQUEST,
                    "Balancing properties amongst replicas in a slice requires that"
                            + " the property be a pre-defined property (e.g. 'preferredLeader') or that 'shardUnique' be set to 'true' "
                            + " Property: " + this.property + " shardUnique: " + Boolean.toString(shardUnique));
        }

        collection = clusterState.getCollection(collectionName);
        if (collection == null) {
            throw new SolrException(SolrException.ErrorCode.BAD_REQUEST,
                    "Could not find collection ' " + collectionName + "' for overseer operation '"
                            + message.getStr(Overseer.QUEUE_OPERATION) + "'. No action taken.");
        }
        onlyActiveNodes = Boolean.parseBoolean(message.getStr(ONLY_ACTIVE_NODES, "true"));
    }

    DocCollection getDocCollection() {
        return collection;
    }

    private boolean isActive(Replica replica) {
        return replica.getState() == Replica.State.ACTIVE;
    }

    // Collect a list of all the nodes that _can_ host the indicated property. Along the way, also collect any of
    // the replicas on that node that _already_ host the property as well as any slices that do _not_ have the
    // property hosted.
    //
    // Return true if anything node needs it's property reassigned. False if the property is already balanced for
    // the collection.

    private boolean collectCurrentPropStats() {
        int maxAssigned = 0;
        // Get a list of potential replicas that can host the property _and_ their counts
        // Move any obvious entries to a list of replicas to change the property on
        Set<String> allHosts = new HashSet<>();
        for (Slice slice : collection.getSlices()) {
            boolean sliceHasProp = false;
            for (Replica replica : slice.getReplicas()) {
                if (onlyActiveNodes && isActive(replica) == false) {
                    if (StringUtils.isNotBlank(replica.getStr(property))) {
                        removeProp(slice, replica.getName()); // Note, we won't be committing this to ZK until later.
                    }
                    continue;
                }
                allHosts.add(replica.getNodeName());
                String nodeName = replica.getNodeName();
                if (StringUtils.isNotBlank(replica.getStr(property))) {
                    if (sliceHasProp) {
                        throw new SolrException(SolrException.ErrorCode.BAD_REQUEST,
                                "'" + BALANCESHARDUNIQUE
                                        + "' should only be called for properties that have at most one member "
                                        + "in any slice with the property set. No action taken.");
                    }
                    if (nodesHostingProp.containsKey(nodeName) == false) {
                        nodesHostingProp.put(nodeName, new ArrayList<>());
                    }
                    nodesHostingProp.get(nodeName).add(new SliceReplica(slice, replica));
                    ++assigned;
                    maxAssigned = Math.max(maxAssigned, nodesHostingProp.get(nodeName).size());
                    sliceHasProp = true;
                }
                if (nodesHostingReplicas.containsKey(nodeName) == false) {
                    nodesHostingReplicas.put(nodeName, new ArrayList<>());
                }
                nodesHostingReplicas.get(nodeName).add(new SliceReplica(slice, replica));
            }
        }

        // If the total number of already-hosted properties assigned to nodes
        // that have potential to host leaders is equal to the slice count _AND_ none of the current nodes has more than
        // the max number of properties, there's nothing to do.
        origMaxPropPerNode = collection.getSlices().size() / allHosts.size();

        // Some nodes can have one more of the proeprty if the numbers aren't exactly even.
        origModulo = collection.getSlices().size() % allHosts.size();
        if (origModulo > 0) {
            origMaxPropPerNode++; // have to have some nodes with 1 more property.
        }

        // We can say for sure that we need to rebalance if we don't have as many assigned properties as slices.
        if (assigned != collection.getSlices().size()) {
            return true;
        }

        // Make sure there are no more slices at the limit than the "leftovers"
        // Let's say there's 7 slices and 3 nodes. We need to distribute the property as 3 on node1, 2 on node2 and 2 on node3
        // (3, 2, 2) We need to be careful to not distribute them as 3, 3, 1. that's what this check is all about.
        int counter = origModulo;
        for (List<SliceReplica> list : nodesHostingProp.values()) {
            if (list.size() == origMaxPropPerNode)
                --counter;
        }
        if (counter == 0)
            return false; // nodes with 1 extra leader are exactly the needed number

        return true;
    }

    private void removeSliceAlreadyHostedFromPossibles(String sliceName) {
        for (Map.Entry<String, List<SliceReplica>> entReplica : nodesHostingReplicas.entrySet()) {

            ListIterator<SliceReplica> iter = entReplica.getValue().listIterator();
            while (iter.hasNext()) {
                SliceReplica sr = iter.next();
                if (sr.slice.getName().equals(sliceName))
                    iter.remove();
            }
        }
    }

    private void balanceUnassignedReplicas() {
        tmpMaxPropPerNode = origMaxPropPerNode; // A bit clumsy, but don't want to duplicate code.
        tmpModulo = origModulo;

        // Get the nodeName and shardName for the node that has the least room for this

        while (shardsNeedingHosts.size() > 0) {
            String nodeName = "";
            int minSize = Integer.MAX_VALUE;
            SliceReplica srToChange = null;
            for (String slice : shardsNeedingHosts) {
                for (Map.Entry<String, List<SliceReplica>> ent : nodesHostingReplicas.entrySet()) {
                    // A little tricky. If we don't set this to something below, then it means all possible places to
                    // put this property are full up, so just put it somewhere.
                    if (srToChange == null && ent.getValue().size() > 0) {
                        srToChange = ent.getValue().get(0);
                    }
                    ListIterator<SliceReplica> iter = ent.getValue().listIterator();
                    while (iter.hasNext()) {
                        SliceReplica sr = iter.next();
                        if (StringUtils.equals(slice, sr.slice.getName()) == false) {
                            continue;
                        }
                        if (nodesHostingProp.containsKey(ent.getKey()) == false) {
                            nodesHostingProp.put(ent.getKey(), new ArrayList<SliceReplica>());
                        }
                        if (minSize > nodesHostingReplicas.get(ent.getKey()).size()
                                && nodesHostingProp.get(ent.getKey()).size() < tmpMaxPropPerNode) {
                            minSize = nodesHostingReplicas.get(ent.getKey()).size();
                            srToChange = sr;
                            nodeName = ent.getKey();
                        }
                    }
                }
            }
            // Now, you have a slice and node to put it on
            shardsNeedingHosts.remove(srToChange.slice.getName());
            if (nodesHostingProp.containsKey(nodeName) == false) {
                nodesHostingProp.put(nodeName, new ArrayList<SliceReplica>());
            }
            nodesHostingProp.get(nodeName).add(srToChange);
            adjustLimits(nodesHostingProp.get(nodeName));
            removeSliceAlreadyHostedFromPossibles(srToChange.slice.getName());
            addProp(srToChange.slice, srToChange.replica.getName());
        }
    }

    // Adjust the min/max counts per allowed per node. Special handling here for dealing with the fact
    // that no node should have more than 1 more replica with this property than any other.
    private void adjustLimits(List<SliceReplica> changeList) {
        if (changeList.size() == tmpMaxPropPerNode) {
            if (tmpModulo < 0)
                return;

            --tmpModulo;
            if (tmpModulo == 0) {
                --tmpMaxPropPerNode;
                --tmpModulo; // Prevent dropping tmpMaxPropPerNode again.
            }
        }
    }

    // Go through the list of presently-hosted properties and remove any that have too many replicas that host the property
    private void removeOverallocatedReplicas() {
        tmpMaxPropPerNode = origMaxPropPerNode; // A bit clumsy, but don't want to duplicate code.
        tmpModulo = origModulo;

        for (Map.Entry<String, List<SliceReplica>> ent : nodesHostingProp.entrySet()) {
            while (ent.getValue().size() > tmpMaxPropPerNode) { // remove delta nodes
                ent.getValue().remove(rand.nextInt(ent.getValue().size()));
            }
            adjustLimits(ent.getValue());
        }
    }

    private void removeProp(Slice origSlice, String replicaName) {
        getReplicaFromChanged(origSlice, replicaName).getProperties().remove(property);
    }

    private void addProp(Slice origSlice, String replicaName) {
        getReplicaFromChanged(origSlice, replicaName).getProperties().put(property, "true");
    }

    // Just a place to encapsulate the fact that we need to have new slices (copy) to update before we
    // put this all in the cluster state.
    private Replica getReplicaFromChanged(Slice origSlice, String replicaName) {
        Slice newSlice = changedSlices.get(origSlice.getName());
        Replica replica;
        if (newSlice != null) {
            replica = newSlice.getReplica(replicaName);
        } else {
            newSlice = new Slice(origSlice.getName(), origSlice.getReplicasCopy(), origSlice.shallowCopy());
            changedSlices.put(origSlice.getName(), newSlice);
            replica = newSlice.getReplica(replicaName);
        }
        if (replica == null) {
            throw new SolrException(SolrException.ErrorCode.INVALID_STATE, "Should have been able to find replica '"
                    + replicaName + "' in slice '" + origSlice.getName() + "'. No action taken");
        }
        return replica;

    }
    // Main entry point for carrying out the action. Returns "true" if we have actually moved properties around.

    boolean balanceProperty() {
        if (collectCurrentPropStats() == false) {
            return false;
        }

        // we have two lists based on nodeName
        // 1> all the nodes that _could_ host a property for the slice
        // 2> all the nodes that _currently_ host a property for the slice.

        // So, remove a replica from the nodes that have too many
        removeOverallocatedReplicas();

        // prune replicas belonging to a slice that have the property currently assigned from the list of replicas
        // that could host the property.
        for (Map.Entry<String, List<SliceReplica>> entProp : nodesHostingProp.entrySet()) {
            for (SliceReplica srHosting : entProp.getValue()) {
                removeSliceAlreadyHostedFromPossibles(srHosting.slice.getName());
            }
        }

        // Assemble the list of slices that do not have any replica hosting the property:
        for (Map.Entry<String, List<SliceReplica>> ent : nodesHostingReplicas.entrySet()) {
            ListIterator<SliceReplica> iter = ent.getValue().listIterator();
            while (iter.hasNext()) {
                SliceReplica sr = iter.next();
                shardsNeedingHosts.add(sr.slice.getName());
            }
        }

        // At this point, nodesHostingProp should contain _only_ lists of replicas that belong to slices that do _not_
        // have any replica hosting the property. So let's assign them.

        balanceUnassignedReplicas();
        for (Slice newSlice : changedSlices.values()) {
            DocCollection docCollection = CollectionMutator.updateSlice(collectionName,
                    clusterState.getCollection(collectionName), newSlice);
            clusterState = ClusterStateMutator.newState(clusterState, collectionName, docCollection);
        }
        return true;
    }

    private static class SliceReplica {
        Slice slice;
        Replica replica;

        SliceReplica(Slice slice, Replica replica) {
            this.slice = slice;
            this.replica = replica;
        }
    }
}