Weighted round-robin scheduling policy for shard coordination traffic…

CHANGELOG.md

@@ -15,6 +15,7 @@ Inspired from [Keep a Changelog](https://keepachangelog.com/en/1.0.0/)
 - Dependency updates (httpcore, mockito, slf4j, httpasyncclient, commons-codec) ([#4308](https://github.com/opensearch-project/OpenSearch/pull/4308))
 - Use RemoteSegmentStoreDirectory instead of RemoteDirectory ([#4240](https://github.com/opensearch-project/OpenSearch/pull/4240))
 - Plugin ZIP publication groupId value is configurable ([#4156](https://github.com/opensearch-project/OpenSearch/pull/4156))
+- Weighted round-robin scheduling policy for shard coordination traffic ([#4241](https://github.com/opensearch-project/OpenSearch/pull/4241))
 
 ### Deprecated
 

server/src/main/java/org/opensearch/cluster/ClusterModule.java

@@ -48,6 +48,7 @@
 import org.opensearch.cluster.metadata.MetadataMappingService;
 import org.opensearch.cluster.metadata.MetadataUpdateSettingsService;
 import org.opensearch.cluster.metadata.RepositoriesMetadata;
+import org.opensearch.cluster.metadata.WeightedRoundRobinRoutingMetadata;
 import org.opensearch.cluster.routing.DelayedAllocationService;
 import org.opensearch.cluster.routing.allocation.AllocationService;
 import org.opensearch.cluster.routing.allocation.ExistingShardsAllocator;
@@ -191,6 +192,12 @@ public static List<Entry> getNamedWriteables() {
             ComposableIndexTemplateMetadata::readDiffFrom
         );
         registerMetadataCustom(entries, DataStreamMetadata.TYPE, DataStreamMetadata::new, DataStreamMetadata::readDiffFrom);
+        registerMetadataCustom(
+            entries,
+            WeightedRoundRobinRoutingMetadata.TYPE,
+            WeightedRoundRobinRoutingMetadata::new,
+            WeightedRoundRobinRoutingMetadata::readDiffFrom
+        );
         // Task Status (not Diffable)
         entries.add(new Entry(Task.Status.class, PersistentTasksNodeService.Status.NAME, PersistentTasksNodeService.Status::new));
         return entries;
@@ -274,6 +281,13 @@ public static List<NamedXContentRegistry.Entry> getNamedXWriteables() {
                 DataStreamMetadata::fromXContent
             )
         );
+        entries.add(
+            new NamedXContentRegistry.Entry(
+                Metadata.Custom.class,
+                new ParseField(WeightedRoundRobinRoutingMetadata.TYPE),
+                WeightedRoundRobinRoutingMetadata::fromXContent
+            )
+        );
         return entries;
     }
 

server/src/main/java/org/opensearch/cluster/metadata/WeightedRoundRobinRoutingMetadata.java

@@ -0,0 +1,157 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+package org.opensearch.cluster.metadata;
+
+import org.apache.logging.log4j.LogManager;
+import org.apache.logging.log4j.Logger;
+import org.opensearch.OpenSearchParseException;
+import org.opensearch.Version;
+import org.opensearch.cluster.AbstractNamedDiffable;
+import org.opensearch.cluster.NamedDiff;
+import org.opensearch.cluster.routing.WRRWeights;
+import org.opensearch.common.Strings;
+import org.opensearch.common.io.stream.StreamInput;
+import org.opensearch.common.io.stream.StreamOutput;
+import org.opensearch.common.xcontent.ToXContent;
+import org.opensearch.common.xcontent.XContentBuilder;
+import org.opensearch.common.xcontent.XContentParser;
+
+import java.io.IOException;
+import java.util.EnumSet;
+import java.util.HashMap;
+import java.util.Map;
+
+/**
+ * Contains metadata for weighted round-robin shard routing weights
+ *
+ * @opensearch.internal
+ */
+public class WeightedRoundRobinRoutingMetadata extends AbstractNamedDiffable<Metadata.Custom> implements Metadata.Custom {
+    private static final Logger logger = LogManager.getLogger(WeightedRoundRobinRoutingMetadata.class);
+    public static final String TYPE = "wrr_shard_routing";
+    private WRRWeights wrrWeight;
+
+    public WRRWeights getWrrWeight() {
+        return wrrWeight;
+    }
+
+    public WeightedRoundRobinRoutingMetadata setWrrWeight(WRRWeights wrrWeight) {
+        this.wrrWeight = wrrWeight;
+        return this;
+    }
+
+    public WeightedRoundRobinRoutingMetadata(StreamInput in) throws IOException {
+        this.wrrWeight = new WRRWeights(in);
+    }
+
+    public WeightedRoundRobinRoutingMetadata(WRRWeights wrrWeight) {
+        this.wrrWeight = wrrWeight;
+    }
+
+    @Override
+    public EnumSet<Metadata.XContentContext> context() {
+        return Metadata.API_AND_GATEWAY;
+    }
+
+    @Override
+    public String getWriteableName() {
+        return TYPE;
+    }
+
+    @Override
+    public Version getMinimalSupportedVersion() {
+        return Version.V_2_3_0;
+
+    }
+
+    @Override
+    public void writeTo(StreamOutput out) throws IOException {
+        wrrWeight.writeTo(out);
+    }
+
+    public static NamedDiff<Metadata.Custom> readDiffFrom(StreamInput in) throws IOException {
+        return readDiffFrom(Metadata.Custom.class, TYPE, in);
+    }
+
+    public static WeightedRoundRobinRoutingMetadata fromXContent(XContentParser parser) throws IOException {
+        String attrKey = null;
+        Object attrValue;
+        String attributeName = null;
+        Map<String, Object> weights = new HashMap<>();
+        WRRWeights wrrWeight = null;
+        XContentParser.Token token;
+        // move to the first alias
+        parser.nextToken();
+        String awarenessField = null;
+
+        while ((token = parser.nextToken()) != XContentParser.Token.END_OBJECT) {
+            if (token == XContentParser.Token.FIELD_NAME) {
+                awarenessField = parser.currentName();
+                if (parser.nextToken() != XContentParser.Token.START_OBJECT) {
+                    throw new OpenSearchParseException("failed to parse wrr metadata  [{}], expected object", awarenessField);
+                }
+                while ((token = parser.nextToken()) != XContentParser.Token.END_OBJECT) {
+                    attributeName = parser.currentName();
+                    if (parser.nextToken() != XContentParser.Token.START_OBJECT) {
+                        throw new OpenSearchParseException("failed to parse wrr metadata  [{}], expected object", attributeName);
+                    }
+                    while ((token = parser.nextToken()) != XContentParser.Token.END_OBJECT) {
+                        if (token == XContentParser.Token.FIELD_NAME) {
+                            attrKey = parser.currentName();
+                        } else if (token == XContentParser.Token.VALUE_STRING) {
+                            attrValue = parser.text();
+                            weights.put(attrKey, attrValue);
+                        } else {
+                            throw new OpenSearchParseException("failed to parse wrr metadata attribute [{}], unknown type", attributeName);
+                        }
+                    }
+                }
+            } else {
+                throw new OpenSearchParseException("failed to parse wrr metadata attribute [{}]", attributeName);
+            }
+        }
+        wrrWeight = new WRRWeights(attributeName, weights);
+        return new WeightedRoundRobinRoutingMetadata(wrrWeight);
+    }
+
+    @Override
+    public boolean equals(Object o) {
+        if (this == o) return true;
+        if (o == null || getClass() != o.getClass()) return false;
+        WeightedRoundRobinRoutingMetadata that = (WeightedRoundRobinRoutingMetadata) o;
+        return wrrWeight.equals(that.wrrWeight);
+    }
+
+    @Override
+    public int hashCode() {
+        return wrrWeight.hashCode();
+    }
+
+    @Override
+    public XContentBuilder toXContent(XContentBuilder builder, ToXContent.Params params) throws IOException {
+        toXContent(wrrWeight, builder);
+        return builder;
+    }
+
+    public static void toXContent(WRRWeights wrrWeight, XContentBuilder builder) throws IOException {
+        builder.startObject("awareness");
+        builder.startObject(wrrWeight.attributeName());
+        for (Map.Entry<String, Object> entry : wrrWeight.weights().entrySet()) {
+            builder.field(entry.getKey(), entry.getValue());
+        }
+        builder.endObject();
+        builder.endObject();
+    }
+
+    @Override
+    public String toString() {
+        return Strings.toString(this);
+    }
+
+}

server/src/main/java/org/opensearch/cluster/routing/IndexShardRoutingTable.java

@@ -292,6 +292,77 @@ public ShardIterator activeInitializingShardsRankedIt(
         return new PlainShardIterator(shardId, ordered);
     }
 
+    /**
+     * Returns an iterator over active and initializing shards, shards are ordered by weighted round-robin scheduling
+     * policy.
+     * @param wrrWeight Weighted round-robin weight entity
+     * @param nodes discovered nodes in the cluster
+     * @return an iterator over active and initializing shards, ordered by weighted round-robin
+     * scheduling policy. Making sure that initializing shards are the last to iterate through.
+     */
+    public ShardIterator activeInitializingShardsWRR(WRRWeights wrrWeight, DiscoveryNodes nodes, WRRShardsCache cache) {
+        final int seed = shuffler.nextSeed();
+        List<ShardRouting> ordered = new ArrayList<>(activeShards.size() + allInitializingShards.size());
+        List<ShardRouting> orderedActiveShards;
+        if (cache.getCache().get(new WRRShardsCache.Key(shardId)) != null) {
+            orderedActiveShards = cache.getCache().get(new WRRShardsCache.Key(shardId));
+        } else {
+            orderedActiveShards = getShardsWRR(activeShards, wrrWeight, nodes);
+            cache.getCache().put(new WRRShardsCache.Key(shardId), orderedActiveShards);
+        }
+        ordered.addAll(orderedActiveShards);
+
+        if (!allInitializingShards.isEmpty()) {
+            List<ShardRouting> orderedInitializingShards = getShardsWRR(allInitializingShards, wrrWeight, nodes);
+            ordered.addAll(orderedInitializingShards);
+        }
+        return new PlainShardIterator(shardId, ordered);
+    }
+
+    /**
+     *
+     * @param shards shards to be ordered using weighted round-robin scheduling policy
+     * @param wrrWeight weights to be considered for routing
+     * @param nodes discovered nodes in the cluster
+     * @return list of shards ordered using weighted round-robin scheduling.
+     */
+    private List<ShardRouting> getShardsWRR(List<ShardRouting> shards, WRRWeights wrrWeight, DiscoveryNodes nodes) {
+        List<WeightedRoundRobin.Entity<ShardRouting>> weightedShards = calculateShardWeight(shards, wrrWeight, nodes);
+        WeightedRoundRobin<ShardRouting> wrr = new WeightedRoundRobin<>(weightedShards);
+        List<WeightedRoundRobin.Entity<ShardRouting>> wrrOrderedActiveShards = wrr.orderEntities();
+        List<ShardRouting> orderedActiveShards = new ArrayList<>(activeShards.size());
+        for (WeightedRoundRobin.Entity<ShardRouting> shardRouting : wrrOrderedActiveShards) {
+            orderedActiveShards.add(shardRouting.getTarget());
+        }
+        return orderedActiveShards;
+    }
+
+    /**
+     *
+     * @param shards associate weights to shards
+     * @param wrrWeight weights to be used for association
+     * @param nodes discovered nodes in the cluster
+     * @return list of entity containing shard routing and associated weight.
+     */
+    private List<WeightedRoundRobin.Entity<ShardRouting>> calculateShardWeight(
+        List<ShardRouting> shards,
+        WRRWeights wrrWeight,
+        DiscoveryNodes nodes
+    ) {
+        List<WeightedRoundRobin.Entity<ShardRouting>> weightedShards = new ArrayList<>();
+        for (ShardRouting shard : shards) {
+            DiscoveryNode node = nodes.get(shard.currentNodeId());
+            String attVal = node.getAttributes().get(wrrWeight.attributeName());
+            // If weight for a zone is not defined, not considering shards from that zone
+            if (wrrWeight.weights().get(attVal) == null) {
+                continue;
+            }
+            Double weight = Double.parseDouble(wrrWeight.weights().get(attVal).toString());
+            weightedShards.add(new WeightedRoundRobin.Entity<>(weight, shard));
+        }
+        return weightedShards;
+    }
+
     private static Set<String> getAllNodeIds(final List<ShardRouting> shards) {
         final Set<String> nodeIds = new HashSet<>();
         for (ShardRouting shard : shards) {

server/src/main/java/org/opensearch/cluster/routing/OperationRouting.java

@@ -34,8 +34,10 @@
 
 import org.opensearch.cluster.ClusterState;
 import org.opensearch.cluster.metadata.IndexMetadata;
+import org.opensearch.cluster.metadata.WeightedRoundRobinRoutingMetadata;
 import org.opensearch.cluster.node.DiscoveryNodes;
 import org.opensearch.cluster.routing.allocation.decider.AwarenessAllocationDecider;
+import org.opensearch.cluster.service.ClusterService;
 import org.opensearch.common.Nullable;
 import org.opensearch.common.Strings;
 import org.opensearch.common.settings.ClusterSettings;
@@ -79,6 +81,33 @@ public class OperationRouting {
     private volatile boolean useAdaptiveReplicaSelection;
     private volatile boolean ignoreAwarenessAttr;
 
+    // reads value from cluster setting
+    private volatile boolean useWeightedRoundRobin;
+    /**
+     * Reads value from cluster setting and cluster state to determine if weighted round-robin
+     * search routing is enabled
+     * This is true if useWeightedRoundRobin=true and weights are set in cluster metadata.
+     */
+    private volatile boolean isWeightedRoundRobinEnabled;
+
+    private volatile WRRWeights wrrWeights;
+
+    public WRRShardsCache getWrrShardsCache() {
+        return wrrShardsCache;
+    }
+
+    private WRRShardsCache wrrShardsCache;
+
+    public ClusterService getClusterService() {
+        return clusterService;
+    }
+
+    public void setClusterService(ClusterService clusterService) {
+        this.clusterService = clusterService;
+    }
+
+    private ClusterService clusterService;
+
     public OperationRouting(Settings settings, ClusterSettings clusterSettings) {
         // whether to ignore awareness attributes when routing requests
         this.ignoreAwarenessAttr = clusterSettings.get(IGNORE_AWARENESS_ATTRIBUTES_SETTING);
@@ -90,6 +119,7 @@ public OperationRouting(Settings settings, ClusterSettings clusterSettings) {
         this.useAdaptiveReplicaSelection = USE_ADAPTIVE_REPLICA_SELECTION_SETTING.get(settings);
         clusterSettings.addSettingsUpdateConsumer(USE_ADAPTIVE_REPLICA_SELECTION_SETTING, this::setUseAdaptiveReplicaSelection);
         clusterSettings.addSettingsUpdateConsumer(IGNORE_AWARENESS_ATTRIBUTES_SETTING, this::setIgnoreAwarenessAttributes);
+
     }
 
     void setUseAdaptiveReplicaSelection(boolean useAdaptiveReplicaSelection) {
@@ -116,6 +146,10 @@ public boolean ignoreAwarenessAttributes() {
         return this.awarenessAttributes.isEmpty() || this.ignoreAwarenessAttr;
     }
 
+    public WRRWeights getWrrWeights() {
+        return wrrWeights;
+    }
+
     public ShardIterator indexShards(ClusterState clusterState, String index, String id, @Nullable String routing) {
         return shards(clusterState, index, id, routing).shardsIt();
     }
@@ -139,6 +173,7 @@ public ShardIterator getShards(
 
     public ShardIterator getShards(ClusterState clusterState, String index, int shardId, @Nullable String preference) {
         final IndexShardRoutingTable indexShard = clusterState.getRoutingTable().shardRoutingTable(index, shardId);
+        setWeightedRoundRobinAttributes(clusterState, getClusterService());
         return preferenceActiveShardIterator(
             indexShard,
             clusterState.nodes().getLocalNodeId(),
@@ -168,6 +203,7 @@ public GroupShardsIterator<ShardIterator> searchShards(
     ) {
         final Set<IndexShardRoutingTable> shards = computeTargetedShards(clusterState, concreteIndices, routing);
         final Set<ShardIterator> set = new HashSet<>(shards.size());
+        setWeightedRoundRobinAttributes(clusterState, getClusterService());
         for (IndexShardRoutingTable shard : shards) {
             ShardIterator iterator = preferenceActiveShardIterator(
                 shard,
@@ -184,6 +220,20 @@ public GroupShardsIterator<ShardIterator> searchShards(
         return GroupShardsIterator.sortAndCreate(new ArrayList<>(set));
     }
 
+    private void setWeightedRoundRobinAttributes(ClusterState clusterState, ClusterService clusterService) {
+        WeightedRoundRobinRoutingMetadata weightedRoundRobinRoutingMetadata = clusterState.metadata()
+            .custom(WeightedRoundRobinRoutingMetadata.TYPE);
+        this.isWeightedRoundRobinEnabled = weightedRoundRobinRoutingMetadata != null ? true : false;
+        if (this.isWeightedRoundRobinEnabled) {
+            this.wrrWeights = weightedRoundRobinRoutingMetadata.getWrrWeight();
+            this.wrrShardsCache = getWrrShardsCache() != null ? getWrrShardsCache() : new WRRShardsCache(clusterService);
+        }
+    }
+
+    private boolean isWeightedRoundRobinEnabled() {
+        return isWeightedRoundRobinEnabled;
+    }
+
     public static ShardIterator getShards(ClusterState clusterState, ShardId shardId) {
         final IndexShardRoutingTable shard = clusterState.routingTable().shardRoutingTable(shardId);
         return shard.activeInitializingShardsRandomIt();
@@ -227,6 +277,7 @@ private ShardIterator preferenceActiveShardIterator(
         @Nullable ResponseCollectorService collectorService,
         @Nullable Map<String, Long> nodeCounts
     ) {
+
         if (preference == null || preference.isEmpty()) {
             return shardRoutings(indexShard, nodes, collectorService, nodeCounts);
         }
@@ -300,7 +351,9 @@ private ShardIterator shardRoutings(
         @Nullable ResponseCollectorService collectorService,
         @Nullable Map<String, Long> nodeCounts
     ) {
-        if (ignoreAwarenessAttributes()) {
+        if (isWeightedRoundRobinEnabled()) {
+            return indexShard.activeInitializingShardsWRR(getWrrWeights(), nodes, wrrShardsCache);
+        } else if (ignoreAwarenessAttributes()) {
             if (useAdaptiveReplicaSelection) {
                 return indexShard.activeInitializingShardsRankedIt(collectorService, nodeCounts);
             } else {