From a5cf2edda31a240839011e0dc1d60ebe39bdd180 Mon Sep 17 00:00:00 2001 From: Anshu Agarwal Date: Thu, 3 Nov 2022 23:09:43 +0530 Subject: [PATCH] Weighted routing/backport2.4 (#5057) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * Weighted round-robin scheduling policy for shard coordination traffic… (#4241) * Add PUT api to update shard routing weights (#4272) * Add GET api to get shard routing weights (#4275) * Fix weighted routing metadata deserialization error during node restart (#4691) * Delete API for weighted round robin search routing (#4917 * Mark apis experimental Signed-off-by: Anshu Agarwal --- CHANGELOG.md | 5 + .../client/RestHighLevelClientTests.java | 5 +- .../api/cluster.delete_weighted_routing.json | 19 + .../api/cluster.get_weighted_routing.json | 25 ++ .../api/cluster.put_weighted_routing.json | 25 ++ .../cluster/routing/WeightedRoutingIT.java | 351 ++++++++++++++++++ .../search/SearchWeightedRoutingIT.java | 161 ++++++++ .../org/opensearch/action/ActionModule.java | 16 + .../ClusterDeleteWeightedRoutingAction.java | 25 ++ .../ClusterDeleteWeightedRoutingRequest.java | 44 +++ ...erDeleteWeightedRoutingRequestBuilder.java | 27 ++ .../ClusterDeleteWeightedRoutingResponse.java | 37 ++ .../TransportDeleteWeightedRoutingAction.java | 86 +++++ .../routing/weighted/delete/package-info.java | 10 + .../get/ClusterGetWeightedRoutingAction.java | 25 ++ .../get/ClusterGetWeightedRoutingRequest.java | 62 ++++ ...usterGetWeightedRoutingRequestBuilder.java | 37 ++ .../ClusterGetWeightedRoutingResponse.java | 125 +++++++ .../TransportGetWeightedRoutingAction.java | 111 ++++++ .../routing/weighted/get/package-info.java | 10 + .../put/ClusterAddWeightedRoutingAction.java | 26 ++ .../put/ClusterPutWeightedRoutingRequest.java | 173 +++++++++ ...usterPutWeightedRoutingRequestBuilder.java | 33 ++ .../ClusterPutWeightedRoutingResponse.java | 29 ++ .../TransportAddWeightedRoutingAction.java | 97 +++++ .../routing/weighted/put/package-info.java | 10 + .../opensearch/client/ClusterAdminClient.java | 55 +++ .../java/org/opensearch/client/Requests.java | 30 ++ .../client/support/AbstractClient.java | 66 ++++ .../org/opensearch/cluster/ClusterModule.java | 9 + .../opensearch/cluster/metadata/Metadata.java | 8 + .../metadata/WeightedRoutingMetadata.java | 165 ++++++++ .../routing/IndexShardRoutingTable.java | 134 +++++++ .../cluster/routing/OperationRouting.java | 46 ++- .../cluster/routing/WeightedRoundRobin.java | 106 ++++++ .../cluster/routing/WeightedRouting.java | 75 ++++ .../routing/WeightedRoutingService.java | 157 ++++++++ .../common/settings/ClusterSettings.java | 1 + ...estClusterDeleteWeightedRoutingAction.java | 53 +++ .../RestClusterGetWeightedRoutingAction.java | 52 +++ .../RestClusterPutWeightedRoutingAction.java | 58 +++ ...ClusterPutWeightedRoutingRequestTests.java | 65 ++++ ...ClusterGetWeightedRoutingRequestTests.java | 39 ++ ...lusterGetWeightedRoutingResponseTests.java | 38 ++ ...ransportGetWeightedRoutingActionTests.java | 262 +++++++++++++ .../WeightedRoutingMetadataTests.java | 36 ++ .../routing/OperationRoutingTests.java | 318 ++++++++++++++++ .../routing/WeightedRoundRobinTests.java | 151 ++++++++ .../routing/WeightedRoutingServiceTests.java | 279 ++++++++++++++ .../structure/RoutingIteratorTests.java | 210 +++++++++++ ...tClusterAddWeightedRoutingActionTests.java | 76 ++++ 51 files changed, 4054 insertions(+), 9 deletions(-) create mode 100644 rest-api-spec/src/main/resources/rest-api-spec/api/cluster.delete_weighted_routing.json create mode 100644 rest-api-spec/src/main/resources/rest-api-spec/api/cluster.get_weighted_routing.json create mode 100644 rest-api-spec/src/main/resources/rest-api-spec/api/cluster.put_weighted_routing.json create mode 100644 server/src/internalClusterTest/java/org/opensearch/cluster/routing/WeightedRoutingIT.java create mode 100644 server/src/internalClusterTest/java/org/opensearch/search/SearchWeightedRoutingIT.java create mode 100644 server/src/main/java/org/opensearch/action/admin/cluster/shards/routing/weighted/delete/ClusterDeleteWeightedRoutingAction.java create mode 100644 server/src/main/java/org/opensearch/action/admin/cluster/shards/routing/weighted/delete/ClusterDeleteWeightedRoutingRequest.java create mode 100644 server/src/main/java/org/opensearch/action/admin/cluster/shards/routing/weighted/delete/ClusterDeleteWeightedRoutingRequestBuilder.java create mode 100644 server/src/main/java/org/opensearch/action/admin/cluster/shards/routing/weighted/delete/ClusterDeleteWeightedRoutingResponse.java create mode 100644 server/src/main/java/org/opensearch/action/admin/cluster/shards/routing/weighted/delete/TransportDeleteWeightedRoutingAction.java create mode 100644 server/src/main/java/org/opensearch/action/admin/cluster/shards/routing/weighted/delete/package-info.java create mode 100644 server/src/main/java/org/opensearch/action/admin/cluster/shards/routing/weighted/get/ClusterGetWeightedRoutingAction.java create mode 100644 server/src/main/java/org/opensearch/action/admin/cluster/shards/routing/weighted/get/ClusterGetWeightedRoutingRequest.java create mode 100644 server/src/main/java/org/opensearch/action/admin/cluster/shards/routing/weighted/get/ClusterGetWeightedRoutingRequestBuilder.java create mode 100644 server/src/main/java/org/opensearch/action/admin/cluster/shards/routing/weighted/get/ClusterGetWeightedRoutingResponse.java create mode 100644 server/src/main/java/org/opensearch/action/admin/cluster/shards/routing/weighted/get/TransportGetWeightedRoutingAction.java create mode 100644 server/src/main/java/org/opensearch/action/admin/cluster/shards/routing/weighted/get/package-info.java create mode 100644 server/src/main/java/org/opensearch/action/admin/cluster/shards/routing/weighted/put/ClusterAddWeightedRoutingAction.java create mode 100644 server/src/main/java/org/opensearch/action/admin/cluster/shards/routing/weighted/put/ClusterPutWeightedRoutingRequest.java create mode 100644 server/src/main/java/org/opensearch/action/admin/cluster/shards/routing/weighted/put/ClusterPutWeightedRoutingRequestBuilder.java create mode 100644 server/src/main/java/org/opensearch/action/admin/cluster/shards/routing/weighted/put/ClusterPutWeightedRoutingResponse.java create mode 100644 server/src/main/java/org/opensearch/action/admin/cluster/shards/routing/weighted/put/TransportAddWeightedRoutingAction.java create mode 100644 server/src/main/java/org/opensearch/action/admin/cluster/shards/routing/weighted/put/package-info.java create mode 100644 server/src/main/java/org/opensearch/cluster/metadata/WeightedRoutingMetadata.java create mode 100644 server/src/main/java/org/opensearch/cluster/routing/WeightedRoundRobin.java create mode 100644 server/src/main/java/org/opensearch/cluster/routing/WeightedRouting.java create mode 100644 server/src/main/java/org/opensearch/cluster/routing/WeightedRoutingService.java create mode 100644 server/src/main/java/org/opensearch/rest/action/admin/cluster/RestClusterDeleteWeightedRoutingAction.java create mode 100644 server/src/main/java/org/opensearch/rest/action/admin/cluster/RestClusterGetWeightedRoutingAction.java create mode 100644 server/src/main/java/org/opensearch/rest/action/admin/cluster/RestClusterPutWeightedRoutingAction.java create mode 100644 server/src/test/java/org/opensearch/action/admin/cluster/shards/routing/weighted/put/ClusterPutWeightedRoutingRequestTests.java create mode 100644 server/src/test/java/org/opensearch/cluster/action/shard/routing/weighted/get/ClusterGetWeightedRoutingRequestTests.java create mode 100644 server/src/test/java/org/opensearch/cluster/action/shard/routing/weighted/get/ClusterGetWeightedRoutingResponseTests.java create mode 100644 server/src/test/java/org/opensearch/cluster/action/shard/routing/weighted/get/TransportGetWeightedRoutingActionTests.java create mode 100644 server/src/test/java/org/opensearch/cluster/metadata/WeightedRoutingMetadataTests.java create mode 100644 server/src/test/java/org/opensearch/cluster/routing/WeightedRoundRobinTests.java create mode 100644 server/src/test/java/org/opensearch/cluster/routing/WeightedRoutingServiceTests.java create mode 100644 server/src/test/java/org/opensearch/rest/action/admin/cluster/RestClusterAddWeightedRoutingActionTests.java diff --git a/CHANGELOG.md b/CHANGELOG.md index 7008fbae1c0ff..31ade664fb1ac 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -50,6 +50,7 @@ Inspired from [Keep a Changelog](https://keepachangelog.com/en/1.0.0/) ### Changed - 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)) +- Weighted round-robin scheduling policy for shard coordination traffic ([#4241](https://github.com/opensearch-project/OpenSearch/pull/4241)) - Add index specific setting for remote repository ([#4253](https://github.com/opensearch-project/OpenSearch/pull/4253)) - [Segment Replication] Update replicas to commit SegmentInfos instead of relying on SIS files from primary shards. ([#4402](https://github.com/opensearch-project/OpenSearch/pull/4402)) - Change the version to remove deprecated code of adding node name into log pattern of log4j property file ([#4569](https://github.com/opensearch-project/OpenSearch/pull/4569)) @@ -58,6 +59,10 @@ Inspired from [Keep a Changelog](https://keepachangelog.com/en/1.0.0/) - Plugin ZIP publication groupId value is configurable ([#4156](https://github.com/opensearch-project/OpenSearch/pull/4156)) - Further simplification of the ZIP publication implementation ([#4360](https://github.com/opensearch-project/OpenSearch/pull/4360)) - [Remote Store] Change behaviour in replica recovery for remote translog enabled indices ([#4318](https://github.com/opensearch-project/OpenSearch/pull/4318)) +- PUT api for weighted shard routing ([#4272](https://github.com/opensearch-project/OpenSearch/pull/4272)) +- GET api for weighted shard routing([#4275](https://github.com/opensearch-project/OpenSearch/pull/4275/)) +- Delete api for weighted shard routing([#4400](https://github.com/opensearch-project/OpenSearch/pull/4400/)) +- Fix weighted routing metadata deserialization error on process restart ([#4691](https://github.com/opensearch-project/OpenSearch/pull/4691)) ### Deprecated ### Removed diff --git a/client/rest-high-level/src/test/java/org/opensearch/client/RestHighLevelClientTests.java b/client/rest-high-level/src/test/java/org/opensearch/client/RestHighLevelClientTests.java index 8945856519412..5003d4aaeeb53 100644 --- a/client/rest-high-level/src/test/java/org/opensearch/client/RestHighLevelClientTests.java +++ b/client/rest-high-level/src/test/java/org/opensearch/client/RestHighLevelClientTests.java @@ -888,7 +888,10 @@ public void testApiNamingConventions() throws Exception { "nodes.usage", "nodes.reload_secure_settings", "search_shards", - "remote_store.restore", }; + "remote_store.restore", + "cluster.put_weighted_routing", + "cluster.get_weighted_routing", + "cluster.delete_weighted_routing", }; List booleanReturnMethods = Arrays.asList("security.enable_user", "security.disable_user", "security.change_password"); Set deprecatedMethods = new HashSet<>(); deprecatedMethods.add("indices.force_merge"); diff --git a/rest-api-spec/src/main/resources/rest-api-spec/api/cluster.delete_weighted_routing.json b/rest-api-spec/src/main/resources/rest-api-spec/api/cluster.delete_weighted_routing.json new file mode 100644 index 0000000000000..274c32f2a91be --- /dev/null +++ b/rest-api-spec/src/main/resources/rest-api-spec/api/cluster.delete_weighted_routing.json @@ -0,0 +1,19 @@ +{ + "cluster.delete_weighted_routing": { + "documentation": { + "url": "https://opensearch.org/docs/latest/opensearch/rest-api/weighted-routing/delete", + "description": "Delete weighted shard routing weights" + }, + "stability": "experimental", + "url": { + "paths": [ + { + "path": "/_cluster/routing/awareness/weights", + "methods": [ + "DELETE" + ] + } + ] + } + } +} diff --git a/rest-api-spec/src/main/resources/rest-api-spec/api/cluster.get_weighted_routing.json b/rest-api-spec/src/main/resources/rest-api-spec/api/cluster.get_weighted_routing.json new file mode 100644 index 0000000000000..95e776d2ffb8b --- /dev/null +++ b/rest-api-spec/src/main/resources/rest-api-spec/api/cluster.get_weighted_routing.json @@ -0,0 +1,25 @@ +{ + "cluster.get_weighted_routing": { + "documentation": { + "url": "https://opensearch.org/docs/latest/opensearch/rest-api/weighted-routing/get", + "description": "Fetches weighted shard routing weights" + }, + "stability": "experimental", + "url": { + "paths": [ + { + "path": "/_cluster/routing/awareness/{attribute}/weights", + "methods": [ + "GET" + ], + "parts": { + "attribute": { + "type": "string", + "description": "Awareness attribute name" + } + } + } + ] + } + } +} diff --git a/rest-api-spec/src/main/resources/rest-api-spec/api/cluster.put_weighted_routing.json b/rest-api-spec/src/main/resources/rest-api-spec/api/cluster.put_weighted_routing.json new file mode 100644 index 0000000000000..b8ce13ab33621 --- /dev/null +++ b/rest-api-spec/src/main/resources/rest-api-spec/api/cluster.put_weighted_routing.json @@ -0,0 +1,25 @@ +{ + "cluster.put_weighted_routing": { + "documentation": { + "url": "https://opensearch.org/docs/latest/opensearch/rest-api/weighted-routing/put", + "description": "Updates weighted shard routing weights" + }, + "stability": "experimental", + "url": { + "paths": [ + { + "path": "/_cluster/routing/awareness/{attribute}/weights", + "methods": [ + "PUT" + ], + "parts": { + "attribute": { + "type": "string", + "description": "Awareness attribute name" + } + } + } + ] + } + } +} diff --git a/server/src/internalClusterTest/java/org/opensearch/cluster/routing/WeightedRoutingIT.java b/server/src/internalClusterTest/java/org/opensearch/cluster/routing/WeightedRoutingIT.java new file mode 100644 index 0000000000000..bba07d878a42c --- /dev/null +++ b/server/src/internalClusterTest/java/org/opensearch/cluster/routing/WeightedRoutingIT.java @@ -0,0 +1,351 @@ +/* + * 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.routing; + +import org.opensearch.action.admin.cluster.health.ClusterHealthResponse; +import org.opensearch.action.admin.cluster.shards.routing.weighted.delete.ClusterDeleteWeightedRoutingResponse; +import org.opensearch.action.admin.cluster.shards.routing.weighted.get.ClusterGetWeightedRoutingResponse; +import org.opensearch.action.admin.cluster.shards.routing.weighted.put.ClusterPutWeightedRoutingResponse; +import org.opensearch.common.settings.Settings; +import org.opensearch.test.OpenSearchIntegTestCase; + +import java.io.IOException; +import java.util.List; +import java.util.Map; + +import static org.hamcrest.Matchers.equalTo; + +@OpenSearchIntegTestCase.ClusterScope(scope = OpenSearchIntegTestCase.Scope.TEST, numDataNodes = 0, minNumDataNodes = 3) +public class WeightedRoutingIT extends OpenSearchIntegTestCase { + + public void testPutWeightedRouting() { + Settings commonSettings = Settings.builder() + .put("cluster.routing.allocation.awareness.attributes", "zone") + .put("cluster.routing.allocation.awareness.force.zone.values", "a,b,c") + .build(); + + logger.info("--> starting 6 nodes on different zones"); + int nodeCountPerAZ = 2; + + logger.info("--> starting a dedicated cluster manager node"); + internalCluster().startClusterManagerOnlyNode(Settings.builder().put(commonSettings).build()); + + logger.info("--> starting 1 nodes on zones 'a' & 'b' & 'c'"); + List nodes_in_zone_a = internalCluster().startDataOnlyNodes( + nodeCountPerAZ, + Settings.builder().put(commonSettings).put("node.attr.zone", "a").build() + ); + List nodes_in_zone_b = internalCluster().startDataOnlyNodes( + nodeCountPerAZ, + Settings.builder().put(commonSettings).put("node.attr.zone", "b").build() + ); + List nodes_in_zone_c = internalCluster().startDataOnlyNodes( + nodeCountPerAZ, + Settings.builder().put(commonSettings).put("node.attr.zone", "c").build() + ); + + logger.info("--> waiting for nodes to form a cluster"); + ClusterHealthResponse health = client().admin().cluster().prepareHealth().setWaitForNodes("7").execute().actionGet(); + assertThat(health.isTimedOut(), equalTo(false)); + + ensureGreen(); + + logger.info("--> setting shard routing weights for weighted round robin"); + Map weights = Map.of("a", 1.0, "b", 2.0, "c", 3.0); + WeightedRouting weightedRouting = new WeightedRouting("zone", weights); + + ClusterPutWeightedRoutingResponse response = client().admin() + .cluster() + .prepareWeightedRouting() + .setWeightedRouting(weightedRouting) + .get(); + assertEquals(response.isAcknowledged(), true); + + // put call made on a data node in zone a + response = internalCluster().client(randomFrom(nodes_in_zone_a.get(0), nodes_in_zone_a.get(1))) + .admin() + .cluster() + .prepareWeightedRouting() + .setWeightedRouting(weightedRouting) + .get(); + assertEquals(response.isAcknowledged(), true); + } + + public void testPutWeightedRouting_InvalidAwarenessAttribute() { + Settings commonSettings = Settings.builder() + .put("cluster.routing.allocation.awareness.attributes", "zone") + .put("cluster.routing.allocation.awareness.force.zone.values", "a,b,c") + .build(); + + internalCluster().startNodes( + Settings.builder().put(commonSettings).put("node.attr.zone", "a").build(), + Settings.builder().put(commonSettings).put("node.attr.zone", "b").build(), + Settings.builder().put(commonSettings).put("node.attr.zone", "c").build() + ); + + logger.info("--> waiting for nodes to form a cluster"); + ClusterHealthResponse health = client().admin().cluster().prepareHealth().setWaitForNodes("3").execute().actionGet(); + assertThat(health.isTimedOut(), equalTo(false)); + + ensureGreen(); + + logger.info("--> setting shard routing weights for weighted round robin"); + Map weights = Map.of("a", 1.0, "b", 2.0, "c", 3.0); + WeightedRouting weightedRouting = new WeightedRouting("zone1", weights); + + assertThrows( + IllegalArgumentException.class, + () -> client().admin().cluster().prepareWeightedRouting().setWeightedRouting(weightedRouting).get() + ); + } + + public void testPutWeightedRouting_MoreThanOneZoneHasZeroWeight() { + Settings commonSettings = Settings.builder() + .put("cluster.routing.allocation.awareness.attributes", "zone") + .put("cluster.routing.allocation.awareness.force.zone.values", "a,b,c") + .build(); + + internalCluster().startNodes( + Settings.builder().put(commonSettings).put("node.attr.zone", "a").build(), + Settings.builder().put(commonSettings).put("node.attr.zone", "b").build(), + Settings.builder().put(commonSettings).put("node.attr.zone", "c").build() + ); + + logger.info("--> waiting for nodes to form a cluster"); + ClusterHealthResponse health = client().admin().cluster().prepareHealth().setWaitForNodes("3").execute().actionGet(); + assertThat(health.isTimedOut(), equalTo(false)); + + ensureGreen(); + + logger.info("--> setting shard routing weights for weighted round robin"); + Map weights = Map.of("a", 1.0, "b", 0.0, "c", 0.0); + WeightedRouting weightedRouting = new WeightedRouting("zone1", weights); + + assertThrows( + IllegalArgumentException.class, + () -> client().admin().cluster().prepareWeightedRouting().setWeightedRouting(weightedRouting).get() + ); + } + + public void testGetWeightedRouting_WeightsNotSet() { + Settings commonSettings = Settings.builder() + .put("cluster.routing.allocation.awareness.attributes", "zone") + .put("cluster.routing.allocation.awareness.force.zone.values", "a,b,c") + .build(); + + internalCluster().startNodes( + Settings.builder().put(commonSettings).put("node.attr.zone", "a").build(), + Settings.builder().put(commonSettings).put("node.attr.zone", "b").build(), + Settings.builder().put(commonSettings).put("node.attr.zone", "c").build() + ); + + logger.info("--> waiting for nodes to form a cluster"); + ClusterHealthResponse health = client().admin().cluster().prepareHealth().setWaitForNodes("3").execute().actionGet(); + assertThat(health.isTimedOut(), equalTo(false)); + + ensureGreen(); + + ClusterGetWeightedRoutingResponse weightedRoutingResponse = client().admin() + .cluster() + .prepareGetWeightedRouting() + .setAwarenessAttribute("zone") + .get(); + assertNull(weightedRoutingResponse.weights()); + } + + public void testGetWeightedRouting_WeightsAreSet() throws IOException { + + Settings commonSettings = Settings.builder() + .put("cluster.routing.allocation.awareness.attributes", "zone") + .put("cluster.routing.allocation.awareness.force.zone.values", "a,b,c") + .build(); + + int nodeCountPerAZ = 2; + + logger.info("--> starting a dedicated cluster manager node"); + internalCluster().startClusterManagerOnlyNode(Settings.builder().put(commonSettings).build()); + + logger.info("--> starting 2 nodes on zones 'a' & 'b' & 'c'"); + List nodes_in_zone_a = internalCluster().startDataOnlyNodes( + nodeCountPerAZ, + Settings.builder().put(commonSettings).put("node.attr.zone", "a").build() + ); + List nodes_in_zone_b = internalCluster().startDataOnlyNodes( + nodeCountPerAZ, + Settings.builder().put(commonSettings).put("node.attr.zone", "b").build() + ); + List nodes_in_zone_c = internalCluster().startDataOnlyNodes( + nodeCountPerAZ, + Settings.builder().put(commonSettings).put("node.attr.zone", "c").build() + ); + + logger.info("--> waiting for nodes to form a cluster"); + ClusterHealthResponse health = client().admin().cluster().prepareHealth().setWaitForNodes("7").execute().actionGet(); + assertThat(health.isTimedOut(), equalTo(false)); + + ensureGreen(); + + logger.info("--> setting shard routing weights for weighted round robin"); + Map weights = Map.of("a", 1.0, "b", 2.0, "c", 3.0); + WeightedRouting weightedRouting = new WeightedRouting("zone", weights); + // put api call to set weights + ClusterPutWeightedRoutingResponse response = client().admin() + .cluster() + .prepareWeightedRouting() + .setWeightedRouting(weightedRouting) + .get(); + assertEquals(response.isAcknowledged(), true); + + // get api call to fetch weights + ClusterGetWeightedRoutingResponse weightedRoutingResponse = client().admin() + .cluster() + .prepareGetWeightedRouting() + .setAwarenessAttribute("zone") + .get(); + assertEquals(weightedRouting, weightedRoutingResponse.weights()); + + // get api to fetch local node weight for a node in zone a + weightedRoutingResponse = internalCluster().client(randomFrom(nodes_in_zone_a.get(0), nodes_in_zone_a.get(1))) + .admin() + .cluster() + .prepareGetWeightedRouting() + .setAwarenessAttribute("zone") + .setRequestLocal(true) + .get(); + assertEquals(weightedRouting, weightedRoutingResponse.weights()); + assertEquals("1.0", weightedRoutingResponse.getLocalNodeWeight()); + + // get api to fetch local node weight for a node in zone b + weightedRoutingResponse = internalCluster().client(randomFrom(nodes_in_zone_b.get(0), nodes_in_zone_b.get(1))) + .admin() + .cluster() + .prepareGetWeightedRouting() + .setAwarenessAttribute("zone") + .setRequestLocal(true) + .get(); + assertEquals(weightedRouting, weightedRoutingResponse.weights()); + assertEquals("2.0", weightedRoutingResponse.getLocalNodeWeight()); + + // get api to fetch local node weight for a node in zone c + weightedRoutingResponse = internalCluster().client(randomFrom(nodes_in_zone_c.get(0), nodes_in_zone_c.get(1))) + .admin() + .cluster() + .prepareGetWeightedRouting() + .setAwarenessAttribute("zone") + .setRequestLocal(true) + .get(); + assertEquals(weightedRouting, weightedRoutingResponse.weights()); + assertEquals("3.0", weightedRoutingResponse.getLocalNodeWeight()); + } + + public void testWeightedRoutingMetadataOnOSProcessRestart() throws Exception { + Settings commonSettings = Settings.builder() + .put("cluster.routing.allocation.awareness.attributes", "zone") + .put("cluster.routing.allocation.awareness.force.zone.values", "a,b,c") + .build(); + + internalCluster().startNodes( + Settings.builder().put(commonSettings).put("node.attr.zone", "a").build(), + Settings.builder().put(commonSettings).put("node.attr.zone", "b").build(), + Settings.builder().put(commonSettings).put("node.attr.zone", "c").build() + ); + + logger.info("--> waiting for nodes to form a cluster"); + ClusterHealthResponse health = client().admin().cluster().prepareHealth().setWaitForNodes("3").execute().actionGet(); + assertThat(health.isTimedOut(), equalTo(false)); + + ensureGreen(); + + logger.info("--> setting shard routing weights for weighted round robin"); + Map weights = Map.of("a", 1.0, "b", 2.0, "c", 3.0); + WeightedRouting weightedRouting = new WeightedRouting("zone", weights); + // put api call to set weights + ClusterPutWeightedRoutingResponse response = client().admin() + .cluster() + .prepareWeightedRouting() + .setWeightedRouting(weightedRouting) + .get(); + assertEquals(response.isAcknowledged(), true); + + ensureStableCluster(3); + + // routing weights are set in cluster metadata + assertNotNull(internalCluster().clusterService().state().metadata().weightedRoutingMetadata()); + + ensureGreen(); + + // Restart a random data node and check that OS process comes healthy + internalCluster().restartRandomDataNode(); + ensureGreen(); + assertNotNull(internalCluster().clusterService().state().metadata().weightedRoutingMetadata()); + } + + public void testDeleteWeightedRouting_WeightsNotSet() { + Settings commonSettings = Settings.builder() + .put("cluster.routing.allocation.awareness.attributes", "zone") + .put("cluster.routing.allocation.awareness.force.zone.values", "a,b,c") + .build(); + + internalCluster().startNodes( + Settings.builder().put(commonSettings).put("node.attr.zone", "a").build(), + Settings.builder().put(commonSettings).put("node.attr.zone", "b").build(), + Settings.builder().put(commonSettings).put("node.attr.zone", "c").build() + ); + + logger.info("--> waiting for nodes to form a cluster"); + ClusterHealthResponse health = client().admin().cluster().prepareHealth().setWaitForNodes("3").execute().actionGet(); + assertThat(health.isTimedOut(), equalTo(false)); + + ensureGreen(); + + assertNull(internalCluster().clusterService().state().metadata().weightedRoutingMetadata()); + + // delete weighted routing metadata + ClusterDeleteWeightedRoutingResponse deleteResponse = client().admin().cluster().prepareDeleteWeightedRouting().get(); + assertTrue(deleteResponse.isAcknowledged()); + assertNull(internalCluster().clusterService().state().metadata().weightedRoutingMetadata()); + } + + public void testDeleteWeightedRouting_WeightsAreSet() { + Settings commonSettings = Settings.builder() + .put("cluster.routing.allocation.awareness.attributes", "zone") + .put("cluster.routing.allocation.awareness.force.zone.values", "a,b,c") + .build(); + + internalCluster().startNodes( + Settings.builder().put(commonSettings).put("node.attr.zone", "a").build(), + Settings.builder().put(commonSettings).put("node.attr.zone", "b").build(), + Settings.builder().put(commonSettings).put("node.attr.zone", "c").build() + ); + + logger.info("--> waiting for nodes to form a cluster"); + ClusterHealthResponse health = client().admin().cluster().prepareHealth().setWaitForNodes("3").execute().actionGet(); + assertThat(health.isTimedOut(), equalTo(false)); + + ensureGreen(); + + logger.info("--> setting shard routing weights for weighted round robin"); + Map weights = Map.of("a", 1.0, "b", 2.0, "c", 3.0); + WeightedRouting weightedRouting = new WeightedRouting("zone", weights); + + // put api call to set weights + ClusterPutWeightedRoutingResponse response = client().admin() + .cluster() + .prepareWeightedRouting() + .setWeightedRouting(weightedRouting) + .get(); + assertEquals(response.isAcknowledged(), true); + assertNotNull(internalCluster().clusterService().state().metadata().weightedRoutingMetadata()); + + // delete weighted routing metadata + ClusterDeleteWeightedRoutingResponse deleteResponse = client().admin().cluster().prepareDeleteWeightedRouting().get(); + assertTrue(deleteResponse.isAcknowledged()); + assertNull(internalCluster().clusterService().state().metadata().weightedRoutingMetadata()); + } +} diff --git a/server/src/internalClusterTest/java/org/opensearch/search/SearchWeightedRoutingIT.java b/server/src/internalClusterTest/java/org/opensearch/search/SearchWeightedRoutingIT.java new file mode 100644 index 0000000000000..097775b7ab4ac --- /dev/null +++ b/server/src/internalClusterTest/java/org/opensearch/search/SearchWeightedRoutingIT.java @@ -0,0 +1,161 @@ +/* + * 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.search; + +import org.junit.Assert; +import org.opensearch.action.admin.cluster.health.ClusterHealthResponse; +import org.opensearch.action.admin.cluster.node.stats.NodeStats; +import org.opensearch.action.admin.cluster.node.stats.NodesStatsResponse; +import org.opensearch.action.admin.cluster.shards.routing.weighted.delete.ClusterDeleteWeightedRoutingResponse; +import org.opensearch.action.admin.cluster.shards.routing.weighted.put.ClusterPutWeightedRoutingResponse; +import org.opensearch.action.search.SearchResponse; +import org.opensearch.cluster.node.DiscoveryNode; +import org.opensearch.cluster.node.DiscoveryNodes; +import org.opensearch.cluster.routing.WeightedRouting; +import org.opensearch.cluster.routing.allocation.decider.AwarenessAllocationDecider; +import org.opensearch.common.settings.Settings; +import org.opensearch.index.query.QueryBuilders; +import org.opensearch.index.search.stats.SearchStats; +import org.opensearch.test.OpenSearchIntegTestCase; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Set; + +import static org.hamcrest.Matchers.equalTo; +import static org.hamcrest.Matchers.lessThanOrEqualTo; +import static org.opensearch.test.hamcrest.OpenSearchAssertions.assertAcked; + +@OpenSearchIntegTestCase.ClusterScope(scope = OpenSearchIntegTestCase.Scope.TEST, numDataNodes = 0, minNumDataNodes = 3) +public class SearchWeightedRoutingIT extends OpenSearchIntegTestCase { + @Override + protected int numberOfReplicas() { + return 2; + } + + public void testSearchWithWRRShardRouting() throws IOException { + Settings commonSettings = Settings.builder() + .put(AwarenessAllocationDecider.CLUSTER_ROUTING_ALLOCATION_AWARENESS_FORCE_GROUP_SETTING.getKey() + "zone" + ".values", "a,b,c") + .put(AwarenessAllocationDecider.CLUSTER_ROUTING_ALLOCATION_AWARENESS_ATTRIBUTE_SETTING.getKey(), "zone") + .build(); + + logger.info("--> starting 6 nodes on different zones"); + List nodes = internalCluster().startNodes( + Settings.builder().put(commonSettings).put("node.attr.zone", "a").build(), + Settings.builder().put(commonSettings).put("node.attr.zone", "b").build(), + Settings.builder().put(commonSettings).put("node.attr.zone", "b").build(), + Settings.builder().put(commonSettings).put("node.attr.zone", "a").build(), + Settings.builder().put(commonSettings).put("node.attr.zone", "c").build(), + Settings.builder().put(commonSettings).put("node.attr.zone", "c").build() + ); + + String A_0 = nodes.get(0); + String B_0 = nodes.get(1); + String B_1 = nodes.get(2); + String A_1 = nodes.get(3); + String C_0 = nodes.get(4); + String C_1 = nodes.get(5); + + logger.info("--> waiting for nodes to form a cluster"); + ClusterHealthResponse health = client().admin().cluster().prepareHealth().setWaitForNodes("6").execute().actionGet(); + assertThat(health.isTimedOut(), equalTo(false)); + + assertAcked( + prepareCreate("test").setSettings(Settings.builder().put("index.number_of_shards", 10).put("index.number_of_replicas", 2)) + ); + ensureGreen(); + logger.info("--> creating indices for test"); + for (int i = 0; i < 100; i++) { + client().prepareIndex("test_" + i).setId("" + i).setSource("field_" + i, "value_" + i).get(); + } + + logger.info("--> setting shard routing weights for weighted round robin"); + Map weights = Map.of("a", 1.0, "b", 1.0, "c", 0.0); + WeightedRouting weightedRouting = new WeightedRouting("zone", weights); + + ClusterPutWeightedRoutingResponse response = client().admin() + .cluster() + .prepareWeightedRouting() + .setWeightedRouting(weightedRouting) + .get(); + assertEquals(response.isAcknowledged(), true); + + Set hitNodes = new HashSet<>(); + // making search requests + for (int i = 0; i < 50; i++) { + SearchResponse searchResponse = internalCluster().client(randomFrom(A_0, A_1, B_0, B_1)) + .prepareSearch() + .setQuery(QueryBuilders.matchAllQuery()) + .get(); + assertEquals(searchResponse.getFailedShards(), 0); + for (int j = 0; j < searchResponse.getHits().getHits().length; j++) { + hitNodes.add(searchResponse.getHits().getAt(j).getShard().getNodeId()); + } + } + // search should not go to nodes in zone c + assertThat(hitNodes.size(), lessThanOrEqualTo(4)); + DiscoveryNodes dataNodes = internalCluster().clusterService().state().nodes(); + List nodeIdsFromZoneWithWeightZero = new ArrayList<>(); + for (DiscoveryNode node : dataNodes) { + if (node.getAttributes().get("zone").equals("c")) { + nodeIdsFromZoneWithWeightZero.add(node.getId()); + } + } + for (String nodeId : nodeIdsFromZoneWithWeightZero) { + assertFalse(hitNodes.contains(nodeId)); + } + + NodesStatsResponse nodeStats = client().admin().cluster().prepareNodesStats().execute().actionGet(); + for (NodeStats stat : nodeStats.getNodes()) { + SearchStats.Stats searchStats = stat.getIndices().getSearch().getTotal(); + if (stat.getNode().getAttributes().get("zone").equals("c")) { + assertEquals(0, searchStats.getQueryCount()); + assertEquals(0, searchStats.getFetchCount()); + + } else { + Assert.assertTrue(searchStats.getQueryCount() > 0L); + Assert.assertTrue(searchStats.getFetchCount() > 0L); + } + } + + logger.info("--> deleted shard routing weights for weighted round robin"); + + ClusterDeleteWeightedRoutingResponse deleteResponse = client().admin().cluster().prepareDeleteWeightedRouting().get(); + assertEquals(deleteResponse.isAcknowledged(), true); + + hitNodes = new HashSet<>(); + // making search requests + for (int i = 0; i < 100; i++) { + SearchResponse searchResponse = internalCluster().client(randomFrom(A_0, A_1, B_0, B_1)) + .prepareSearch() + .setQuery(QueryBuilders.matchAllQuery()) + .get(); + assertEquals(searchResponse.getFailedShards(), 0); + for (int j = 0; j < searchResponse.getHits().getHits().length; j++) { + hitNodes.add(searchResponse.getHits().getAt(j).getShard().getNodeId()); + } + } + + // Check shard routing requests hit data nodes in zone c + for (String nodeId : nodeIdsFromZoneWithWeightZero) { + assertFalse(!hitNodes.contains(nodeId)); + } + nodeStats = client().admin().cluster().prepareNodesStats().execute().actionGet(); + + for (NodeStats stat : nodeStats.getNodes()) { + SearchStats.Stats searchStats = stat.getIndices().getSearch().getTotal(); + Assert.assertTrue(searchStats.getQueryCount() > 0L); + Assert.assertTrue(searchStats.getFetchCount() > 0L); + } + } + +} diff --git a/server/src/main/java/org/opensearch/action/ActionModule.java b/server/src/main/java/org/opensearch/action/ActionModule.java index b8cbba930951a..89eb82de97da8 100644 --- a/server/src/main/java/org/opensearch/action/ActionModule.java +++ b/server/src/main/java/org/opensearch/action/ActionModule.java @@ -79,6 +79,12 @@ import org.opensearch.action.admin.cluster.settings.TransportClusterUpdateSettingsAction; import org.opensearch.action.admin.cluster.shards.ClusterSearchShardsAction; import org.opensearch.action.admin.cluster.shards.TransportClusterSearchShardsAction; +import org.opensearch.action.admin.cluster.shards.routing.weighted.delete.ClusterDeleteWeightedRoutingAction; +import org.opensearch.action.admin.cluster.shards.routing.weighted.delete.TransportDeleteWeightedRoutingAction; +import org.opensearch.action.admin.cluster.shards.routing.weighted.get.ClusterGetWeightedRoutingAction; +import org.opensearch.action.admin.cluster.shards.routing.weighted.get.TransportGetWeightedRoutingAction; +import org.opensearch.action.admin.cluster.shards.routing.weighted.put.ClusterAddWeightedRoutingAction; +import org.opensearch.action.admin.cluster.shards.routing.weighted.put.TransportAddWeightedRoutingAction; import org.opensearch.action.admin.cluster.snapshots.clone.CloneSnapshotAction; import org.opensearch.action.admin.cluster.snapshots.clone.TransportCloneSnapshotAction; import org.opensearch.action.admin.cluster.snapshots.create.CreateSnapshotAction; @@ -292,8 +298,11 @@ import org.opensearch.rest.action.admin.cluster.RestClearVotingConfigExclusionsAction; import org.opensearch.rest.action.admin.cluster.RestCloneSnapshotAction; import org.opensearch.rest.action.admin.cluster.RestClusterAllocationExplainAction; +import org.opensearch.rest.action.admin.cluster.RestClusterDeleteWeightedRoutingAction; import org.opensearch.rest.action.admin.cluster.RestClusterGetSettingsAction; +import org.opensearch.rest.action.admin.cluster.RestClusterGetWeightedRoutingAction; import org.opensearch.rest.action.admin.cluster.RestClusterHealthAction; +import org.opensearch.rest.action.admin.cluster.RestClusterPutWeightedRoutingAction; import org.opensearch.rest.action.admin.cluster.RestClusterRerouteAction; import org.opensearch.rest.action.admin.cluster.RestClusterSearchShardsAction; import org.opensearch.rest.action.admin.cluster.RestClusterStateAction; @@ -563,6 +572,9 @@ public void reg actions.register(RestoreSnapshotAction.INSTANCE, TransportRestoreSnapshotAction.class); actions.register(SnapshotsStatusAction.INSTANCE, TransportSnapshotsStatusAction.class); + actions.register(ClusterAddWeightedRoutingAction.INSTANCE, TransportAddWeightedRoutingAction.class); + actions.register(ClusterGetWeightedRoutingAction.INSTANCE, TransportGetWeightedRoutingAction.class); + actions.register(ClusterDeleteWeightedRoutingAction.INSTANCE, TransportDeleteWeightedRoutingAction.class); actions.register(IndicesStatsAction.INSTANCE, TransportIndicesStatsAction.class); actions.register(IndicesSegmentsAction.INSTANCE, TransportIndicesSegmentsAction.class); actions.register(IndicesShardStoresAction.INSTANCE, TransportIndicesShardStoresAction.class); @@ -745,6 +757,10 @@ public void initRestHandlers(Supplier nodesInCluster) { registerHandler.accept(new RestOpenIndexAction()); registerHandler.accept(new RestAddIndexBlockAction()); + registerHandler.accept(new RestClusterPutWeightedRoutingAction()); + registerHandler.accept(new RestClusterGetWeightedRoutingAction()); + registerHandler.accept(new RestClusterDeleteWeightedRoutingAction()); + registerHandler.accept(new RestUpdateSettingsAction()); registerHandler.accept(new RestGetSettingsAction()); diff --git a/server/src/main/java/org/opensearch/action/admin/cluster/shards/routing/weighted/delete/ClusterDeleteWeightedRoutingAction.java b/server/src/main/java/org/opensearch/action/admin/cluster/shards/routing/weighted/delete/ClusterDeleteWeightedRoutingAction.java new file mode 100644 index 0000000000000..aa438cd31b934 --- /dev/null +++ b/server/src/main/java/org/opensearch/action/admin/cluster/shards/routing/weighted/delete/ClusterDeleteWeightedRoutingAction.java @@ -0,0 +1,25 @@ +/* + * 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.action.admin.cluster.shards.routing.weighted.delete; + +import org.opensearch.action.ActionType; + +/** + * Action to delete weights for weighted round-robin shard routing policy. + * + * @opensearch.internal + */ +public class ClusterDeleteWeightedRoutingAction extends ActionType { + public static final ClusterDeleteWeightedRoutingAction INSTANCE = new ClusterDeleteWeightedRoutingAction(); + public static final String NAME = "cluster:admin/routing/awareness/weights/delete"; + + private ClusterDeleteWeightedRoutingAction() { + super(NAME, ClusterDeleteWeightedRoutingResponse::new); + } +} diff --git a/server/src/main/java/org/opensearch/action/admin/cluster/shards/routing/weighted/delete/ClusterDeleteWeightedRoutingRequest.java b/server/src/main/java/org/opensearch/action/admin/cluster/shards/routing/weighted/delete/ClusterDeleteWeightedRoutingRequest.java new file mode 100644 index 0000000000000..71eab8ff35a2d --- /dev/null +++ b/server/src/main/java/org/opensearch/action/admin/cluster/shards/routing/weighted/delete/ClusterDeleteWeightedRoutingRequest.java @@ -0,0 +1,44 @@ +/* + * 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.action.admin.cluster.shards.routing.weighted.delete; + +import org.opensearch.action.ActionRequestValidationException; +import org.opensearch.action.support.clustermanager.ClusterManagerNodeRequest; +import org.opensearch.common.io.stream.StreamInput; +import org.opensearch.common.io.stream.StreamOutput; + +import java.io.IOException; + +/** + * Request to delete weights for weighted round-robin shard routing policy. + * + * @opensearch.internal + */ +public class ClusterDeleteWeightedRoutingRequest extends ClusterManagerNodeRequest { + public ClusterDeleteWeightedRoutingRequest() {} + + public ClusterDeleteWeightedRoutingRequest(StreamInput in) throws IOException { + super(in); + } + + @Override + public ActionRequestValidationException validate() { + return null; + } + + @Override + public void writeTo(StreamOutput out) throws IOException { + super.writeTo(out); + } + + @Override + public String toString() { + return "ClusterDeleteWeightedRoutingRequest"; + } +} diff --git a/server/src/main/java/org/opensearch/action/admin/cluster/shards/routing/weighted/delete/ClusterDeleteWeightedRoutingRequestBuilder.java b/server/src/main/java/org/opensearch/action/admin/cluster/shards/routing/weighted/delete/ClusterDeleteWeightedRoutingRequestBuilder.java new file mode 100644 index 0000000000000..19976ac6b07aa --- /dev/null +++ b/server/src/main/java/org/opensearch/action/admin/cluster/shards/routing/weighted/delete/ClusterDeleteWeightedRoutingRequestBuilder.java @@ -0,0 +1,27 @@ +/* + * 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.action.admin.cluster.shards.routing.weighted.delete; + +import org.opensearch.action.support.clustermanager.ClusterManagerNodeOperationRequestBuilder; +import org.opensearch.client.OpenSearchClient; + +/** + * Request builder to delete weights for weighted round-robin shard routing policy. + * + * @opensearch.internal + */ +public class ClusterDeleteWeightedRoutingRequestBuilder extends ClusterManagerNodeOperationRequestBuilder< + ClusterDeleteWeightedRoutingRequest, + ClusterDeleteWeightedRoutingResponse, + ClusterDeleteWeightedRoutingRequestBuilder> { + + public ClusterDeleteWeightedRoutingRequestBuilder(OpenSearchClient client, ClusterDeleteWeightedRoutingAction action) { + super(client, action, new ClusterDeleteWeightedRoutingRequest()); + } +} diff --git a/server/src/main/java/org/opensearch/action/admin/cluster/shards/routing/weighted/delete/ClusterDeleteWeightedRoutingResponse.java b/server/src/main/java/org/opensearch/action/admin/cluster/shards/routing/weighted/delete/ClusterDeleteWeightedRoutingResponse.java new file mode 100644 index 0000000000000..b98ac6c0c55be --- /dev/null +++ b/server/src/main/java/org/opensearch/action/admin/cluster/shards/routing/weighted/delete/ClusterDeleteWeightedRoutingResponse.java @@ -0,0 +1,37 @@ +/* + * 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.action.admin.cluster.shards.routing.weighted.delete; + +import org.opensearch.action.support.master.AcknowledgedResponse; +import org.opensearch.common.io.stream.StreamInput; +import org.opensearch.common.io.stream.StreamOutput; + +import java.io.IOException; + +/** + * Response from deleting weights for weighted round-robin search routing policy. + * + * @opensearch.internal + */ +public class ClusterDeleteWeightedRoutingResponse extends AcknowledgedResponse { + + ClusterDeleteWeightedRoutingResponse(StreamInput in) throws IOException { + super(in); + } + + public ClusterDeleteWeightedRoutingResponse(boolean acknowledged) { + super(acknowledged); + } + + @Override + public void writeTo(StreamOutput out) throws IOException { + super.writeTo(out); + + } +} diff --git a/server/src/main/java/org/opensearch/action/admin/cluster/shards/routing/weighted/delete/TransportDeleteWeightedRoutingAction.java b/server/src/main/java/org/opensearch/action/admin/cluster/shards/routing/weighted/delete/TransportDeleteWeightedRoutingAction.java new file mode 100644 index 0000000000000..8f88d8af71b70 --- /dev/null +++ b/server/src/main/java/org/opensearch/action/admin/cluster/shards/routing/weighted/delete/TransportDeleteWeightedRoutingAction.java @@ -0,0 +1,86 @@ +/* + * 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.action.admin.cluster.shards.routing.weighted.delete; + +import org.apache.logging.log4j.LogManager; +import org.apache.logging.log4j.Logger; +import org.opensearch.action.ActionListener; +import org.opensearch.action.support.ActionFilters; +import org.opensearch.action.support.clustermanager.TransportClusterManagerNodeAction; +import org.opensearch.cluster.ClusterState; +import org.opensearch.cluster.block.ClusterBlockException; +import org.opensearch.cluster.block.ClusterBlockLevel; +import org.opensearch.cluster.metadata.IndexNameExpressionResolver; +import org.opensearch.cluster.routing.WeightedRoutingService; +import org.opensearch.cluster.service.ClusterService; +import org.opensearch.common.inject.Inject; +import org.opensearch.common.io.stream.StreamInput; +import org.opensearch.threadpool.ThreadPool; +import org.opensearch.transport.TransportService; + +import java.io.IOException; + +/** + * Transport action for deleting weights for weighted round-robin search routing policy + * + * @opensearch.internal + */ +public class TransportDeleteWeightedRoutingAction extends TransportClusterManagerNodeAction< + ClusterDeleteWeightedRoutingRequest, + ClusterDeleteWeightedRoutingResponse> { + + private static final Logger logger = LogManager.getLogger(TransportDeleteWeightedRoutingAction.class); + + private final WeightedRoutingService weightedRoutingService; + + @Inject + public TransportDeleteWeightedRoutingAction( + TransportService transportService, + ClusterService clusterService, + WeightedRoutingService weightedRoutingService, + ThreadPool threadPool, + ActionFilters actionFilters, + IndexNameExpressionResolver indexNameExpressionResolver + ) { + super( + ClusterDeleteWeightedRoutingAction.NAME, + transportService, + clusterService, + threadPool, + actionFilters, + ClusterDeleteWeightedRoutingRequest::new, + indexNameExpressionResolver + ); + this.weightedRoutingService = weightedRoutingService; + } + + @Override + protected String executor() { + return ThreadPool.Names.SAME; + } + + @Override + protected ClusterDeleteWeightedRoutingResponse read(StreamInput in) throws IOException { + return new ClusterDeleteWeightedRoutingResponse(in); + } + + @Override + protected ClusterBlockException checkBlock(ClusterDeleteWeightedRoutingRequest request, ClusterState state) { + return state.blocks().globalBlockedException(ClusterBlockLevel.METADATA_WRITE); + } + + @Override + protected void clusterManagerOperation( + ClusterDeleteWeightedRoutingRequest request, + ClusterState state, + ActionListener listener + ) throws Exception { + weightedRoutingService.deleteWeightedRoutingMetadata(request, listener); + } +} diff --git a/server/src/main/java/org/opensearch/action/admin/cluster/shards/routing/weighted/delete/package-info.java b/server/src/main/java/org/opensearch/action/admin/cluster/shards/routing/weighted/delete/package-info.java new file mode 100644 index 0000000000000..d24c88ec674f3 --- /dev/null +++ b/server/src/main/java/org/opensearch/action/admin/cluster/shards/routing/weighted/delete/package-info.java @@ -0,0 +1,10 @@ +/* + * 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. + */ + +/** delete weighted-round robin shard routing weights. */ +package org.opensearch.action.admin.cluster.shards.routing.weighted.delete; diff --git a/server/src/main/java/org/opensearch/action/admin/cluster/shards/routing/weighted/get/ClusterGetWeightedRoutingAction.java b/server/src/main/java/org/opensearch/action/admin/cluster/shards/routing/weighted/get/ClusterGetWeightedRoutingAction.java new file mode 100644 index 0000000000000..7662b7cc6fcc8 --- /dev/null +++ b/server/src/main/java/org/opensearch/action/admin/cluster/shards/routing/weighted/get/ClusterGetWeightedRoutingAction.java @@ -0,0 +1,25 @@ +/* + * 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.action.admin.cluster.shards.routing.weighted.get; + +import org.opensearch.action.ActionType; + +/** + * Action to get weights for weighted round-robin search routing policy. + * + * @opensearch.internal + */ +public class ClusterGetWeightedRoutingAction extends ActionType { + public static final ClusterGetWeightedRoutingAction INSTANCE = new ClusterGetWeightedRoutingAction(); + public static final String NAME = "cluster:admin/routing/awareness/weights/get"; + + private ClusterGetWeightedRoutingAction() { + super(NAME, ClusterGetWeightedRoutingResponse::new); + } +} diff --git a/server/src/main/java/org/opensearch/action/admin/cluster/shards/routing/weighted/get/ClusterGetWeightedRoutingRequest.java b/server/src/main/java/org/opensearch/action/admin/cluster/shards/routing/weighted/get/ClusterGetWeightedRoutingRequest.java new file mode 100644 index 0000000000000..aaa000baa95f3 --- /dev/null +++ b/server/src/main/java/org/opensearch/action/admin/cluster/shards/routing/weighted/get/ClusterGetWeightedRoutingRequest.java @@ -0,0 +1,62 @@ +/* + * 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.action.admin.cluster.shards.routing.weighted.get; + +import org.opensearch.action.ActionRequestValidationException; +import org.opensearch.action.support.clustermanager.ClusterManagerNodeReadRequest; +import org.opensearch.common.io.stream.StreamInput; +import org.opensearch.common.io.stream.StreamOutput; + +import java.io.IOException; + +import static org.opensearch.action.ValidateActions.addValidationError; + +/** + * Request to get weights for weighted round-robin search routing policy. + * + * @opensearch.internal + */ +public class ClusterGetWeightedRoutingRequest extends ClusterManagerNodeReadRequest { + + String awarenessAttribute; + + public String getAwarenessAttribute() { + return awarenessAttribute; + } + + public void setAwarenessAttribute(String awarenessAttribute) { + this.awarenessAttribute = awarenessAttribute; + } + + public ClusterGetWeightedRoutingRequest(String awarenessAttribute) { + this.awarenessAttribute = awarenessAttribute; + } + + public ClusterGetWeightedRoutingRequest() {} + + public ClusterGetWeightedRoutingRequest(StreamInput in) throws IOException { + super(in); + awarenessAttribute = in.readString(); + } + + @Override + public void writeTo(StreamOutput out) throws IOException { + super.writeTo(out); + out.writeString(awarenessAttribute); + } + + @Override + public ActionRequestValidationException validate() { + ActionRequestValidationException validationException = null; + if (awarenessAttribute == null || awarenessAttribute.isEmpty()) { + validationException = addValidationError("Awareness attribute is missing", validationException); + } + return validationException; + } +} diff --git a/server/src/main/java/org/opensearch/action/admin/cluster/shards/routing/weighted/get/ClusterGetWeightedRoutingRequestBuilder.java b/server/src/main/java/org/opensearch/action/admin/cluster/shards/routing/weighted/get/ClusterGetWeightedRoutingRequestBuilder.java new file mode 100644 index 0000000000000..82f4c1106461d --- /dev/null +++ b/server/src/main/java/org/opensearch/action/admin/cluster/shards/routing/weighted/get/ClusterGetWeightedRoutingRequestBuilder.java @@ -0,0 +1,37 @@ +/* + * 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.action.admin.cluster.shards.routing.weighted.get; + +import org.opensearch.action.support.clustermanager.ClusterManagerNodeReadOperationRequestBuilder; +import org.opensearch.client.OpenSearchClient; + +/** + * Request builder to get weights for weighted round-robin search routing policy. + * + * @opensearch.internal + */ +public class ClusterGetWeightedRoutingRequestBuilder extends ClusterManagerNodeReadOperationRequestBuilder< + ClusterGetWeightedRoutingRequest, + ClusterGetWeightedRoutingResponse, + ClusterGetWeightedRoutingRequestBuilder> { + + public ClusterGetWeightedRoutingRequestBuilder(OpenSearchClient client, ClusterGetWeightedRoutingAction action) { + super(client, action, new ClusterGetWeightedRoutingRequest()); + } + + public ClusterGetWeightedRoutingRequestBuilder setRequestLocal(boolean local) { + request.local(local); + return this; + } + + public ClusterGetWeightedRoutingRequestBuilder setAwarenessAttribute(String attribute) { + request.setAwarenessAttribute(attribute); + return this; + } +} diff --git a/server/src/main/java/org/opensearch/action/admin/cluster/shards/routing/weighted/get/ClusterGetWeightedRoutingResponse.java b/server/src/main/java/org/opensearch/action/admin/cluster/shards/routing/weighted/get/ClusterGetWeightedRoutingResponse.java new file mode 100644 index 0000000000000..bb77576b63d20 --- /dev/null +++ b/server/src/main/java/org/opensearch/action/admin/cluster/shards/routing/weighted/get/ClusterGetWeightedRoutingResponse.java @@ -0,0 +1,125 @@ +/* + * 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.action.admin.cluster.shards.routing.weighted.get; + +import org.opensearch.OpenSearchParseException; +import org.opensearch.action.ActionResponse; + +import org.opensearch.cluster.routing.WeightedRouting; +import org.opensearch.common.io.stream.StreamInput; +import org.opensearch.common.io.stream.StreamOutput; +import org.opensearch.common.xcontent.ToXContentObject; +import org.opensearch.common.xcontent.XContentBuilder; +import org.opensearch.common.xcontent.XContentParser; + +import java.io.IOException; +import java.util.HashMap; +import java.util.Map; +import java.util.Objects; + +import static org.opensearch.common.xcontent.XContentParserUtils.ensureExpectedToken; + +/** + * Response from fetching weights for weighted round-robin search routing policy. + * + * @opensearch.internal + */ +public class ClusterGetWeightedRoutingResponse extends ActionResponse implements ToXContentObject { + private WeightedRouting weightedRouting; + private String localNodeWeight; + private static final String NODE_WEIGHT = "node_weight"; + + public String getLocalNodeWeight() { + return localNodeWeight; + } + + ClusterGetWeightedRoutingResponse() { + this.weightedRouting = null; + } + + public ClusterGetWeightedRoutingResponse(String localNodeWeight, WeightedRouting weightedRouting) { + this.localNodeWeight = localNodeWeight; + this.weightedRouting = weightedRouting; + } + + ClusterGetWeightedRoutingResponse(StreamInput in) throws IOException { + if (in.available() != 0) { + this.weightedRouting = new WeightedRouting(in); + } + } + + /** + * List of weights to return + * + * @return list or weights + */ + public WeightedRouting weights() { + return this.weightedRouting; + } + + @Override + public void writeTo(StreamOutput out) throws IOException { + if (weightedRouting != null) { + weightedRouting.writeTo(out); + } + } + + @Override + public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException { + builder.startObject(); + if (this.weightedRouting != null) { + for (Map.Entry entry : weightedRouting.weights().entrySet()) { + builder.field(entry.getKey(), entry.getValue().toString()); + } + if (localNodeWeight != null) { + builder.field(NODE_WEIGHT, localNodeWeight); + } + } + builder.endObject(); + return builder; + } + + public static ClusterGetWeightedRoutingResponse fromXContent(XContentParser parser) throws IOException { + ensureExpectedToken(XContentParser.Token.START_OBJECT, parser.nextToken(), parser); + XContentParser.Token token; + String attrKey = null, attrValue = null; + String localNodeWeight = null; + Map weights = new HashMap<>(); + + 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(); + if (attrKey != null && attrKey.equals(NODE_WEIGHT)) { + localNodeWeight = attrValue; + } else if (attrKey != null) { + weights.put(attrKey, Double.parseDouble(attrValue)); + } + } else { + throw new OpenSearchParseException("failed to parse weighted routing response"); + } + } + WeightedRouting weightedRouting = new WeightedRouting("", weights); + return new ClusterGetWeightedRoutingResponse(localNodeWeight, weightedRouting); + } + + @Override + public boolean equals(Object o) { + if (this == o) return true; + if (o == null || getClass() != o.getClass()) return false; + ClusterGetWeightedRoutingResponse that = (ClusterGetWeightedRoutingResponse) o; + return weightedRouting.equals(that.weightedRouting) && localNodeWeight.equals(that.localNodeWeight); + } + + @Override + public int hashCode() { + return Objects.hash(weightedRouting, localNodeWeight); + } +} diff --git a/server/src/main/java/org/opensearch/action/admin/cluster/shards/routing/weighted/get/TransportGetWeightedRoutingAction.java b/server/src/main/java/org/opensearch/action/admin/cluster/shards/routing/weighted/get/TransportGetWeightedRoutingAction.java new file mode 100644 index 0000000000000..9421967a5df26 --- /dev/null +++ b/server/src/main/java/org/opensearch/action/admin/cluster/shards/routing/weighted/get/TransportGetWeightedRoutingAction.java @@ -0,0 +1,111 @@ +/* + * 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.action.admin.cluster.shards.routing.weighted.get; + +import org.opensearch.action.ActionListener; +import org.apache.logging.log4j.LogManager; +import org.apache.logging.log4j.Logger; + +import org.opensearch.action.support.ActionFilters; +import org.opensearch.action.support.clustermanager.TransportClusterManagerNodeReadAction; +import org.opensearch.cluster.ClusterState; +import org.opensearch.cluster.block.ClusterBlockException; +import org.opensearch.cluster.block.ClusterBlockLevel; +import org.opensearch.cluster.metadata.IndexNameExpressionResolver; + +import org.opensearch.cluster.metadata.WeightedRoutingMetadata; +import org.opensearch.cluster.node.DiscoveryNode; +import org.opensearch.cluster.routing.WeightedRouting; +import org.opensearch.cluster.routing.WeightedRoutingService; +import org.opensearch.cluster.service.ClusterService; +import org.opensearch.common.inject.Inject; +import org.opensearch.common.io.stream.StreamInput; + +import org.opensearch.threadpool.ThreadPool; +import org.opensearch.transport.TransportService; + +import java.io.IOException; + +/** + * Transport action for getting weights for weighted round-robin search routing policy + * + * @opensearch.internal + */ +public class TransportGetWeightedRoutingAction extends TransportClusterManagerNodeReadAction< + ClusterGetWeightedRoutingRequest, + ClusterGetWeightedRoutingResponse> { + private static final Logger logger = LogManager.getLogger(TransportGetWeightedRoutingAction.class); + private final WeightedRoutingService weightedRoutingService; + + @Inject + public TransportGetWeightedRoutingAction( + TransportService transportService, + ClusterService clusterService, + WeightedRoutingService weightedRoutingService, + ThreadPool threadPool, + ActionFilters actionFilters, + IndexNameExpressionResolver indexNameExpressionResolver + ) { + super( + ClusterGetWeightedRoutingAction.NAME, + transportService, + clusterService, + threadPool, + actionFilters, + ClusterGetWeightedRoutingRequest::new, + indexNameExpressionResolver + ); + this.weightedRoutingService = weightedRoutingService; + } + + @Override + protected String executor() { + return ThreadPool.Names.SAME; + } + + @Override + protected ClusterGetWeightedRoutingResponse read(StreamInput in) throws IOException { + return new ClusterGetWeightedRoutingResponse(in); + } + + @Override + protected ClusterBlockException checkBlock(ClusterGetWeightedRoutingRequest request, ClusterState state) { + return state.blocks().globalBlockedException(ClusterBlockLevel.METADATA_READ); + } + + @Override + protected void clusterManagerOperation( + final ClusterGetWeightedRoutingRequest request, + ClusterState state, + final ActionListener listener + ) { + try { + weightedRoutingService.verifyAwarenessAttribute(request.getAwarenessAttribute()); + WeightedRoutingMetadata weightedRoutingMetadata = state.metadata().custom(WeightedRoutingMetadata.TYPE); + ClusterGetWeightedRoutingResponse clusterGetWeightedRoutingResponse = new ClusterGetWeightedRoutingResponse(); + String weight = null; + if (weightedRoutingMetadata != null && weightedRoutingMetadata.getWeightedRouting() != null) { + WeightedRouting weightedRouting = weightedRoutingMetadata.getWeightedRouting(); + if (request.local()) { + DiscoveryNode localNode = state.getNodes().getLocalNode(); + if (localNode.getAttributes().get(request.getAwarenessAttribute()) != null) { + String attrVal = localNode.getAttributes().get(request.getAwarenessAttribute()); + if (weightedRouting.weights().containsKey(attrVal)) { + weight = weightedRouting.weights().get(attrVal).toString(); + } + } + } + clusterGetWeightedRoutingResponse = new ClusterGetWeightedRoutingResponse(weight, weightedRouting); + } + listener.onResponse(clusterGetWeightedRoutingResponse); + } catch (Exception ex) { + listener.onFailure(ex); + } + } +} diff --git a/server/src/main/java/org/opensearch/action/admin/cluster/shards/routing/weighted/get/package-info.java b/server/src/main/java/org/opensearch/action/admin/cluster/shards/routing/weighted/get/package-info.java new file mode 100644 index 0000000000000..45e5b32b72e50 --- /dev/null +++ b/server/src/main/java/org/opensearch/action/admin/cluster/shards/routing/weighted/get/package-info.java @@ -0,0 +1,10 @@ +/* + * 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. + */ + +/** get weighted-round robin shard routing weights. */ +package org.opensearch.action.admin.cluster.shards.routing.weighted.get; diff --git a/server/src/main/java/org/opensearch/action/admin/cluster/shards/routing/weighted/put/ClusterAddWeightedRoutingAction.java b/server/src/main/java/org/opensearch/action/admin/cluster/shards/routing/weighted/put/ClusterAddWeightedRoutingAction.java new file mode 100644 index 0000000000000..65c5ccca71461 --- /dev/null +++ b/server/src/main/java/org/opensearch/action/admin/cluster/shards/routing/weighted/put/ClusterAddWeightedRoutingAction.java @@ -0,0 +1,26 @@ +/* + * 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.action.admin.cluster.shards.routing.weighted.put; + +import org.opensearch.action.ActionType; + +/** + * Action to update weights for weighted round-robin shard routing policy. + * + * @opensearch.internal + */ +public final class ClusterAddWeightedRoutingAction extends ActionType { + + public static final ClusterAddWeightedRoutingAction INSTANCE = new ClusterAddWeightedRoutingAction(); + public static final String NAME = "cluster:admin/routing/awareness/weights/put"; + + private ClusterAddWeightedRoutingAction() { + super(NAME, ClusterPutWeightedRoutingResponse::new); + } +} diff --git a/server/src/main/java/org/opensearch/action/admin/cluster/shards/routing/weighted/put/ClusterPutWeightedRoutingRequest.java b/server/src/main/java/org/opensearch/action/admin/cluster/shards/routing/weighted/put/ClusterPutWeightedRoutingRequest.java new file mode 100644 index 0000000000000..af229fb12b4f0 --- /dev/null +++ b/server/src/main/java/org/opensearch/action/admin/cluster/shards/routing/weighted/put/ClusterPutWeightedRoutingRequest.java @@ -0,0 +1,173 @@ +/* + * 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.action.admin.cluster.shards.routing.weighted.put; + +import org.apache.logging.log4j.LogManager; +import org.apache.logging.log4j.Logger; +import org.opensearch.OpenSearchGenerationException; +import org.opensearch.OpenSearchParseException; +import org.opensearch.action.ActionRequestValidationException; +import org.opensearch.action.support.clustermanager.ClusterManagerNodeRequest; +import org.opensearch.cluster.routing.WeightedRouting; +import org.opensearch.common.bytes.BytesReference; +import org.opensearch.common.io.stream.StreamInput; +import org.opensearch.common.io.stream.StreamOutput; +import org.opensearch.common.xcontent.DeprecationHandler; +import org.opensearch.common.xcontent.NamedXContentRegistry; +import org.opensearch.common.xcontent.XContentBuilder; +import org.opensearch.common.xcontent.XContentFactory; +import org.opensearch.common.xcontent.XContentHelper; +import org.opensearch.common.xcontent.XContentParser; +import org.opensearch.common.xcontent.XContentType; + +import java.io.IOException; +import java.util.HashMap; +import java.util.Locale; +import java.util.Map; + +import static org.opensearch.action.ValidateActions.addValidationError; + +/** + * Request to update weights for weighted round-robin shard routing policy. + * + * @opensearch.internal + */ +public class ClusterPutWeightedRoutingRequest extends ClusterManagerNodeRequest { + private static final Logger logger = LogManager.getLogger(ClusterPutWeightedRoutingRequest.class); + + private WeightedRouting weightedRouting; + private String attributeName; + + public ClusterPutWeightedRoutingRequest() {} + + public WeightedRouting getWeightedRouting() { + return weightedRouting; + } + + public ClusterPutWeightedRoutingRequest setWeightedRouting(WeightedRouting weightedRouting) { + this.weightedRouting = weightedRouting; + return this; + } + + public void attributeName(String attributeName) { + this.attributeName = attributeName; + } + + public ClusterPutWeightedRoutingRequest(StreamInput in) throws IOException { + super(in); + weightedRouting = new WeightedRouting(in); + } + + public ClusterPutWeightedRoutingRequest(String attributeName) { + this.attributeName = attributeName; + } + + public void setWeightedRouting(Map source) { + try { + if (source.isEmpty()) { + throw new OpenSearchParseException(("Empty request body")); + } + XContentBuilder builder = XContentFactory.jsonBuilder(); + builder.map(source); + setWeightedRouting(BytesReference.bytes(builder), builder.contentType()); + } catch (IOException e) { + throw new OpenSearchGenerationException("Failed to generate [" + source + "]", e); + } + } + + public void setWeightedRouting(BytesReference source, XContentType contentType) { + try ( + XContentParser parser = XContentHelper.createParser( + NamedXContentRegistry.EMPTY, + DeprecationHandler.THROW_UNSUPPORTED_OPERATION, + source, + contentType + ) + ) { + String attrValue = null; + Map weights = new HashMap<>(); + Double attrWeight = null; + XContentParser.Token token; + // move to the first alias + parser.nextToken(); + while ((token = parser.nextToken()) != XContentParser.Token.END_OBJECT) { + if (token == XContentParser.Token.FIELD_NAME) { + attrValue = parser.currentName(); + } else if (token == XContentParser.Token.VALUE_STRING) { + attrWeight = Double.parseDouble(parser.text()); + weights.put(attrValue, attrWeight); + } else { + throw new OpenSearchParseException( + "failed to parse weighted routing request attribute [{}], " + "unknown type", + attrWeight + ); + } + } + this.weightedRouting = new WeightedRouting(this.attributeName, weights); + } catch (IOException e) { + logger.error("error while parsing put for weighted routing request object", e); + } + } + + @Override + public ActionRequestValidationException validate() { + ActionRequestValidationException validationException = null; + if (weightedRouting == null) { + validationException = addValidationError("Weighted routing request object is null", validationException); + } + if (weightedRouting.attributeName() == null || weightedRouting.attributeName().isEmpty()) { + validationException = addValidationError("Attribute name is missing", validationException); + } + if (weightedRouting.weights() == null || weightedRouting.weights().isEmpty()) { + validationException = addValidationError("Weights are missing", validationException); + } + int countValueWithZeroWeights = 0; + double weight; + try { + for (Object value : weightedRouting.weights().values()) { + if (value == null) { + validationException = addValidationError(("Weight is null"), validationException); + } else { + weight = Double.parseDouble(value.toString()); + countValueWithZeroWeights = (weight == 0) ? countValueWithZeroWeights + 1 : countValueWithZeroWeights; + } + } + } catch (NumberFormatException e) { + validationException = addValidationError(("Weight is not a number"), validationException); + } + if (countValueWithZeroWeights > 1) { + validationException = addValidationError( + (String.format(Locale.ROOT, "More than one [%d] value has weight set as 0", countValueWithZeroWeights)), + validationException + ); + } + return validationException; + } + + /** + * @param source weights definition from request body + * @return this request + */ + public ClusterPutWeightedRoutingRequest source(Map source) { + setWeightedRouting(source); + return this; + } + + @Override + public void writeTo(StreamOutput out) throws IOException { + super.writeTo(out); + weightedRouting.writeTo(out); + } + + @Override + public String toString() { + return "ClusterPutWeightedRoutingRequest{" + "weightedRouting= " + weightedRouting.toString() + "}"; + } + +} diff --git a/server/src/main/java/org/opensearch/action/admin/cluster/shards/routing/weighted/put/ClusterPutWeightedRoutingRequestBuilder.java b/server/src/main/java/org/opensearch/action/admin/cluster/shards/routing/weighted/put/ClusterPutWeightedRoutingRequestBuilder.java new file mode 100644 index 0000000000000..b437f4c54d8d6 --- /dev/null +++ b/server/src/main/java/org/opensearch/action/admin/cluster/shards/routing/weighted/put/ClusterPutWeightedRoutingRequestBuilder.java @@ -0,0 +1,33 @@ +/* + * 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.action.admin.cluster.shards.routing.weighted.put; + +import org.opensearch.action.support.clustermanager.ClusterManagerNodeOperationRequestBuilder; +import org.opensearch.client.OpenSearchClient; +import org.opensearch.cluster.routing.WeightedRouting; + +/** + * Request builder to update weights for weighted round-robin shard routing policy. + * + * @opensearch.internal + */ +public class ClusterPutWeightedRoutingRequestBuilder extends ClusterManagerNodeOperationRequestBuilder< + ClusterPutWeightedRoutingRequest, + ClusterPutWeightedRoutingResponse, + ClusterPutWeightedRoutingRequestBuilder> { + public ClusterPutWeightedRoutingRequestBuilder(OpenSearchClient client, ClusterAddWeightedRoutingAction action) { + super(client, action, new ClusterPutWeightedRoutingRequest()); + } + + public ClusterPutWeightedRoutingRequestBuilder setWeightedRouting(WeightedRouting weightedRouting) { + request.setWeightedRouting(weightedRouting); + return this; + } + +} diff --git a/server/src/main/java/org/opensearch/action/admin/cluster/shards/routing/weighted/put/ClusterPutWeightedRoutingResponse.java b/server/src/main/java/org/opensearch/action/admin/cluster/shards/routing/weighted/put/ClusterPutWeightedRoutingResponse.java new file mode 100644 index 0000000000000..b0154aceef0c2 --- /dev/null +++ b/server/src/main/java/org/opensearch/action/admin/cluster/shards/routing/weighted/put/ClusterPutWeightedRoutingResponse.java @@ -0,0 +1,29 @@ +/* + * 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.action.admin.cluster.shards.routing.weighted.put; + +import org.opensearch.action.support.master.AcknowledgedResponse; +import org.opensearch.common.io.stream.StreamInput; + +import java.io.IOException; + +/** + * Response from updating weights for weighted round-robin search routing policy. + * + * @opensearch.internal + */ +public class ClusterPutWeightedRoutingResponse extends AcknowledgedResponse { + public ClusterPutWeightedRoutingResponse(boolean acknowledged) { + super(acknowledged); + } + + public ClusterPutWeightedRoutingResponse(StreamInput in) throws IOException { + super(in); + } +} diff --git a/server/src/main/java/org/opensearch/action/admin/cluster/shards/routing/weighted/put/TransportAddWeightedRoutingAction.java b/server/src/main/java/org/opensearch/action/admin/cluster/shards/routing/weighted/put/TransportAddWeightedRoutingAction.java new file mode 100644 index 0000000000000..249e313c1f53b --- /dev/null +++ b/server/src/main/java/org/opensearch/action/admin/cluster/shards/routing/weighted/put/TransportAddWeightedRoutingAction.java @@ -0,0 +1,97 @@ +/* + * 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.action.admin.cluster.shards.routing.weighted.put; + +import org.opensearch.action.ActionListener; +import org.opensearch.action.ActionRequestValidationException; +import org.opensearch.action.support.ActionFilters; +import org.opensearch.action.support.clustermanager.TransportClusterManagerNodeAction; +import org.opensearch.cluster.ClusterState; +import org.opensearch.cluster.block.ClusterBlockException; +import org.opensearch.cluster.block.ClusterBlockLevel; +import org.opensearch.cluster.metadata.IndexNameExpressionResolver; +import org.opensearch.cluster.routing.WeightedRoutingService; +import org.opensearch.cluster.service.ClusterService; +import org.opensearch.common.inject.Inject; +import org.opensearch.common.io.stream.StreamInput; +import org.opensearch.threadpool.ThreadPool; +import org.opensearch.transport.TransportService; + +import java.io.IOException; + +/** + * Transport action for updating weights for weighted round-robin search routing policy + * + * @opensearch.internal + */ +public class TransportAddWeightedRoutingAction extends TransportClusterManagerNodeAction< + ClusterPutWeightedRoutingRequest, + ClusterPutWeightedRoutingResponse> { + + private final WeightedRoutingService weightedRoutingService; + + @Inject + public TransportAddWeightedRoutingAction( + TransportService transportService, + ClusterService clusterService, + WeightedRoutingService weightedRoutingService, + ThreadPool threadPool, + ActionFilters actionFilters, + IndexNameExpressionResolver indexNameExpressionResolver + ) { + super( + ClusterAddWeightedRoutingAction.NAME, + transportService, + clusterService, + threadPool, + actionFilters, + ClusterPutWeightedRoutingRequest::new, + indexNameExpressionResolver + ); + this.weightedRoutingService = weightedRoutingService; + } + + @Override + protected String executor() { + return ThreadPool.Names.SAME; + } + + @Override + protected ClusterPutWeightedRoutingResponse read(StreamInput in) throws IOException { + return new ClusterPutWeightedRoutingResponse(in); + } + + @Override + protected void clusterManagerOperation( + ClusterPutWeightedRoutingRequest request, + ClusterState state, + ActionListener listener + ) throws Exception { + try { + weightedRoutingService.verifyAwarenessAttribute(request.getWeightedRouting().attributeName()); + } catch (ActionRequestValidationException ex) { + listener.onFailure(ex); + return; + } + weightedRoutingService.registerWeightedRoutingMetadata( + request, + ActionListener.delegateFailure( + listener, + (delegatedListener, response) -> { + delegatedListener.onResponse(new ClusterPutWeightedRoutingResponse(response.isAcknowledged())); + } + ) + ); + } + + @Override + protected ClusterBlockException checkBlock(ClusterPutWeightedRoutingRequest request, ClusterState state) { + return state.blocks().globalBlockedException(ClusterBlockLevel.METADATA_WRITE); + } +} diff --git a/server/src/main/java/org/opensearch/action/admin/cluster/shards/routing/weighted/put/package-info.java b/server/src/main/java/org/opensearch/action/admin/cluster/shards/routing/weighted/put/package-info.java new file mode 100644 index 0000000000000..4f18b220cd343 --- /dev/null +++ b/server/src/main/java/org/opensearch/action/admin/cluster/shards/routing/weighted/put/package-info.java @@ -0,0 +1,10 @@ +/* + * 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. + */ + +/** add/update weighted-round robin shard routing weights. */ +package org.opensearch.action.admin.cluster.shards.routing.weighted.put; diff --git a/server/src/main/java/org/opensearch/client/ClusterAdminClient.java b/server/src/main/java/org/opensearch/client/ClusterAdminClient.java index 7a7b98bf724f6..47fb5bd8d4600 100644 --- a/server/src/main/java/org/opensearch/client/ClusterAdminClient.java +++ b/server/src/main/java/org/opensearch/client/ClusterAdminClient.java @@ -86,6 +86,15 @@ import org.opensearch.action.admin.cluster.shards.ClusterSearchShardsRequest; import org.opensearch.action.admin.cluster.shards.ClusterSearchShardsRequestBuilder; import org.opensearch.action.admin.cluster.shards.ClusterSearchShardsResponse; +import org.opensearch.action.admin.cluster.shards.routing.weighted.delete.ClusterDeleteWeightedRoutingRequest; +import org.opensearch.action.admin.cluster.shards.routing.weighted.delete.ClusterDeleteWeightedRoutingRequestBuilder; +import org.opensearch.action.admin.cluster.shards.routing.weighted.delete.ClusterDeleteWeightedRoutingResponse; +import org.opensearch.action.admin.cluster.shards.routing.weighted.get.ClusterGetWeightedRoutingRequest; +import org.opensearch.action.admin.cluster.shards.routing.weighted.get.ClusterGetWeightedRoutingRequestBuilder; +import org.opensearch.action.admin.cluster.shards.routing.weighted.get.ClusterGetWeightedRoutingResponse; +import org.opensearch.action.admin.cluster.shards.routing.weighted.put.ClusterPutWeightedRoutingRequest; +import org.opensearch.action.admin.cluster.shards.routing.weighted.put.ClusterPutWeightedRoutingRequestBuilder; +import org.opensearch.action.admin.cluster.shards.routing.weighted.put.ClusterPutWeightedRoutingResponse; import org.opensearch.action.admin.cluster.snapshots.clone.CloneSnapshotRequest; import org.opensearch.action.admin.cluster.snapshots.clone.CloneSnapshotRequestBuilder; import org.opensearch.action.admin.cluster.snapshots.create.CreateSnapshotRequest; @@ -791,4 +800,50 @@ public interface ClusterAdminClient extends OpenSearchClient { * Delete specified dangling indices. */ ActionFuture deleteDanglingIndex(DeleteDanglingIndexRequest request); + + /** + * Updates weights for weighted round-robin search routing policy. + */ + ActionFuture putWeightedRouting(ClusterPutWeightedRoutingRequest request); + + /** + * Updates weights for weighted round-robin search routing policy. + */ + void putWeightedRouting(ClusterPutWeightedRoutingRequest request, ActionListener listener); + + /** + * Updates weights for weighted round-robin search routing policy. + */ + ClusterPutWeightedRoutingRequestBuilder prepareWeightedRouting(); + + /** + * Gets weights for weighted round-robin search routing policy. + */ + ActionFuture getWeightedRouting(ClusterGetWeightedRoutingRequest request); + + /** + * Gets weights for weighted round-robin search routing policy. + */ + void getWeightedRouting(ClusterGetWeightedRoutingRequest request, ActionListener listener); + + /** + * Gets weights for weighted round-robin search routing policy. + */ + ClusterGetWeightedRoutingRequestBuilder prepareGetWeightedRouting(); + + /** + * Deletes weights for weighted round-robin search routing policy. + */ + ActionFuture deleteWeightedRouting(ClusterDeleteWeightedRoutingRequest request); + + /** + * Deletes weights for weighted round-robin search routing policy. + */ + void deleteWeightedRouting(ClusterDeleteWeightedRoutingRequest request, ActionListener listener); + + /** + * Deletes weights for weighted round-robin search routing policy. + */ + ClusterDeleteWeightedRoutingRequestBuilder prepareDeleteWeightedRouting(); + } diff --git a/server/src/main/java/org/opensearch/client/Requests.java b/server/src/main/java/org/opensearch/client/Requests.java index b04de7830a780..107cfe5f04d06 100644 --- a/server/src/main/java/org/opensearch/client/Requests.java +++ b/server/src/main/java/org/opensearch/client/Requests.java @@ -47,6 +47,9 @@ import org.opensearch.action.admin.cluster.reroute.ClusterRerouteRequest; import org.opensearch.action.admin.cluster.settings.ClusterUpdateSettingsRequest; import org.opensearch.action.admin.cluster.shards.ClusterSearchShardsRequest; +import org.opensearch.action.admin.cluster.shards.routing.weighted.get.ClusterGetWeightedRoutingRequest; +import org.opensearch.action.admin.cluster.shards.routing.weighted.put.ClusterPutWeightedRoutingRequest; +import org.opensearch.action.admin.cluster.shards.routing.weighted.delete.ClusterDeleteWeightedRoutingRequest; import org.opensearch.action.admin.cluster.snapshots.create.CreateSnapshotRequest; import org.opensearch.action.admin.cluster.snapshots.delete.DeleteSnapshotRequest; import org.opensearch.action.admin.cluster.snapshots.get.GetSnapshotsRequest; @@ -548,4 +551,31 @@ public static DeleteSnapshotRequest deleteSnapshotRequest(String repository, Str public static SnapshotsStatusRequest snapshotsStatusRequest(String repository) { return new SnapshotsStatusRequest(repository); } + + /** + * Updates weights for weighted round-robin search routing policy + * + * @return update weight request + */ + public static ClusterPutWeightedRoutingRequest putWeightedRoutingRequest(String attributeName) { + return new ClusterPutWeightedRoutingRequest(attributeName); + } + + /** + * Gets weights for weighted round-robin search routing policy + * + * @return get weight request + */ + public static ClusterGetWeightedRoutingRequest getWeightedRoutingRequest(String attributeName) { + return new ClusterGetWeightedRoutingRequest(attributeName); + } + + /** + * Deletes weights for weighted round-robin search routing policy + * + * @return delete weight request + */ + public static ClusterDeleteWeightedRoutingRequest deleteWeightedRoutingRequest() { + return new ClusterDeleteWeightedRoutingRequest(); + } } diff --git a/server/src/main/java/org/opensearch/client/support/AbstractClient.java b/server/src/main/java/org/opensearch/client/support/AbstractClient.java index 21cd01bf65a45..80d04156f475b 100644 --- a/server/src/main/java/org/opensearch/client/support/AbstractClient.java +++ b/server/src/main/java/org/opensearch/client/support/AbstractClient.java @@ -110,6 +110,18 @@ import org.opensearch.action.admin.cluster.shards.ClusterSearchShardsRequest; import org.opensearch.action.admin.cluster.shards.ClusterSearchShardsRequestBuilder; import org.opensearch.action.admin.cluster.shards.ClusterSearchShardsResponse; +import org.opensearch.action.admin.cluster.shards.routing.weighted.delete.ClusterDeleteWeightedRoutingAction; +import org.opensearch.action.admin.cluster.shards.routing.weighted.delete.ClusterDeleteWeightedRoutingRequest; +import org.opensearch.action.admin.cluster.shards.routing.weighted.delete.ClusterDeleteWeightedRoutingRequestBuilder; +import org.opensearch.action.admin.cluster.shards.routing.weighted.delete.ClusterDeleteWeightedRoutingResponse; +import org.opensearch.action.admin.cluster.shards.routing.weighted.get.ClusterGetWeightedRoutingAction; +import org.opensearch.action.admin.cluster.shards.routing.weighted.get.ClusterGetWeightedRoutingRequest; +import org.opensearch.action.admin.cluster.shards.routing.weighted.get.ClusterGetWeightedRoutingRequestBuilder; +import org.opensearch.action.admin.cluster.shards.routing.weighted.get.ClusterGetWeightedRoutingResponse; +import org.opensearch.action.admin.cluster.shards.routing.weighted.put.ClusterAddWeightedRoutingAction; +import org.opensearch.action.admin.cluster.shards.routing.weighted.put.ClusterPutWeightedRoutingRequest; +import org.opensearch.action.admin.cluster.shards.routing.weighted.put.ClusterPutWeightedRoutingRequestBuilder; +import org.opensearch.action.admin.cluster.shards.routing.weighted.put.ClusterPutWeightedRoutingResponse; import org.opensearch.action.admin.cluster.snapshots.clone.CloneSnapshotAction; import org.opensearch.action.admin.cluster.snapshots.clone.CloneSnapshotRequest; import org.opensearch.action.admin.cluster.snapshots.clone.CloneSnapshotRequestBuilder; @@ -1272,6 +1284,60 @@ public ActionFuture deleteDanglingIndex(DeleteDanglingInde return execute(DeleteDanglingIndexAction.INSTANCE, request); } + @Override + public ActionFuture putWeightedRouting(ClusterPutWeightedRoutingRequest request) { + return execute(ClusterAddWeightedRoutingAction.INSTANCE, request); + } + + @Override + public void putWeightedRouting( + ClusterPutWeightedRoutingRequest request, + ActionListener listener + ) { + execute(ClusterAddWeightedRoutingAction.INSTANCE, request, listener); + } + + @Override + public ClusterPutWeightedRoutingRequestBuilder prepareWeightedRouting() { + return new ClusterPutWeightedRoutingRequestBuilder(this, ClusterAddWeightedRoutingAction.INSTANCE); + } + + @Override + public ActionFuture getWeightedRouting(ClusterGetWeightedRoutingRequest request) { + return execute(ClusterGetWeightedRoutingAction.INSTANCE, request); + } + + @Override + public void getWeightedRouting( + ClusterGetWeightedRoutingRequest request, + ActionListener listener + ) { + execute(ClusterGetWeightedRoutingAction.INSTANCE, request, listener); + } + + @Override + public ClusterGetWeightedRoutingRequestBuilder prepareGetWeightedRouting() { + return new ClusterGetWeightedRoutingRequestBuilder(this, ClusterGetWeightedRoutingAction.INSTANCE); + } + + @Override + public ActionFuture deleteWeightedRouting(ClusterDeleteWeightedRoutingRequest request) { + return execute(ClusterDeleteWeightedRoutingAction.INSTANCE, request); + } + + @Override + public void deleteWeightedRouting( + ClusterDeleteWeightedRoutingRequest request, + ActionListener listener + ) { + execute(ClusterDeleteWeightedRoutingAction.INSTANCE, request, listener); + } + + @Override + public ClusterDeleteWeightedRoutingRequestBuilder prepareDeleteWeightedRouting() { + return new ClusterDeleteWeightedRoutingRequestBuilder(this, ClusterDeleteWeightedRoutingAction.INSTANCE); + } + @Override public void deleteDanglingIndex(DeleteDanglingIndexRequest request, ActionListener listener) { execute(DeleteDanglingIndexAction.INSTANCE, request, listener); diff --git a/server/src/main/java/org/opensearch/cluster/ClusterModule.java b/server/src/main/java/org/opensearch/cluster/ClusterModule.java index bce92e66509d6..85f49b19bc967 100644 --- a/server/src/main/java/org/opensearch/cluster/ClusterModule.java +++ b/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.WeightedRoutingMetadata; import org.opensearch.cluster.routing.DelayedAllocationService; import org.opensearch.cluster.routing.allocation.AllocationService; import org.opensearch.cluster.routing.allocation.ExistingShardsAllocator; @@ -192,6 +193,7 @@ public static List getNamedWriteables() { ComposableIndexTemplateMetadata::readDiffFrom ); registerMetadataCustom(entries, DataStreamMetadata.TYPE, DataStreamMetadata::new, DataStreamMetadata::readDiffFrom); + registerMetadataCustom(entries, WeightedRoutingMetadata.TYPE, WeightedRoutingMetadata::new, WeightedRoutingMetadata::readDiffFrom); // Task Status (not Diffable) entries.add(new Entry(Task.Status.class, PersistentTasksNodeService.Status.NAME, PersistentTasksNodeService.Status::new)); return entries; @@ -275,6 +277,13 @@ public static List getNamedXWriteables() { DataStreamMetadata::fromXContent ) ); + entries.add( + new NamedXContentRegistry.Entry( + Metadata.Custom.class, + new ParseField(WeightedRoutingMetadata.TYPE), + WeightedRoutingMetadata::fromXContent + ) + ); return entries; } diff --git a/server/src/main/java/org/opensearch/cluster/metadata/Metadata.java b/server/src/main/java/org/opensearch/cluster/metadata/Metadata.java index 5f7e98e9e1199..086865d2170c3 100644 --- a/server/src/main/java/org/opensearch/cluster/metadata/Metadata.java +++ b/server/src/main/java/org/opensearch/cluster/metadata/Metadata.java @@ -810,6 +810,14 @@ public IndexGraveyard indexGraveyard() { return custom(IndexGraveyard.TYPE); } + /** + * * + * @return The weighted routing metadata for search requests + */ + public WeightedRoutingMetadata weightedRoutingMetadata() { + return custom(WeightedRoutingMetadata.TYPE); + } + public T custom(String type) { return (T) customs.get(type); } diff --git a/server/src/main/java/org/opensearch/cluster/metadata/WeightedRoutingMetadata.java b/server/src/main/java/org/opensearch/cluster/metadata/WeightedRoutingMetadata.java new file mode 100644 index 0000000000000..07cdc949c4529 --- /dev/null +++ b/server/src/main/java/org/opensearch/cluster/metadata/WeightedRoutingMetadata.java @@ -0,0 +1,165 @@ +/* + * 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.WeightedRouting; +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 routing + * + * @opensearch.internal + */ +public class WeightedRoutingMetadata extends AbstractNamedDiffable implements Metadata.Custom { + private static final Logger logger = LogManager.getLogger(WeightedRoutingMetadata.class); + public static final String TYPE = "weighted_shard_routing"; + public static final String AWARENESS = "awareness"; + private WeightedRouting weightedRouting; + + public WeightedRouting getWeightedRouting() { + return weightedRouting; + } + + public WeightedRoutingMetadata setWeightedRouting(WeightedRouting weightedRouting) { + this.weightedRouting = weightedRouting; + return this; + } + + public WeightedRoutingMetadata(StreamInput in) throws IOException { + if (in.available() != 0) { + this.weightedRouting = new WeightedRouting(in); + } + } + + public WeightedRoutingMetadata(WeightedRouting weightedRouting) { + this.weightedRouting = weightedRouting; + } + + @Override + public EnumSet context() { + return Metadata.API_AND_GATEWAY; + } + + @Override + public String getWriteableName() { + return TYPE; + } + + @Override + public Version getMinimalSupportedVersion() { + return Version.V_2_4_0; + } + + @Override + public void writeTo(StreamOutput out) throws IOException { + if (weightedRouting != null) { + weightedRouting.writeTo(out); + } + } + + public static NamedDiff readDiffFrom(StreamInput in) throws IOException { + return readDiffFrom(Metadata.Custom.class, TYPE, in); + } + + public static WeightedRoutingMetadata fromXContent(XContentParser parser) throws IOException { + String attrKey = null; + Double attrValue; + String attributeName = null; + Map weights = new HashMap<>(); + WeightedRouting weightedRouting = null; + XContentParser.Token token; + 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 weighted routing 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 weighted routing 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_NUMBER) { + attrValue = Double.parseDouble(parser.text()); + weights.put(attrKey, attrValue); + } else { + throw new OpenSearchParseException( + "failed to parse weighted routing metadata attribute " + "[{}], unknown type", + attributeName + ); + } + } + } + } + } + weightedRouting = new WeightedRouting(attributeName, weights); + return new WeightedRoutingMetadata(weightedRouting); + } + + @Override + public boolean equals(Object o) { + if (this == o) return true; + if (o == null || getClass() != o.getClass()) return false; + WeightedRoutingMetadata that = (WeightedRoutingMetadata) o; + return weightedRouting.equals(that.weightedRouting); + } + + @Override + public int hashCode() { + return weightedRouting.hashCode(); + } + + @Override + public XContentBuilder toXContent(XContentBuilder builder, ToXContent.Params params) throws IOException { + toXContent(weightedRouting, builder); + return builder; + } + + public static void toXContent(WeightedRouting weightedRouting, XContentBuilder builder) throws IOException { + builder.startObject(AWARENESS); + builder.startObject(weightedRouting.attributeName()); + for (Map.Entry entry : weightedRouting.weights().entrySet()) { + builder.field(entry.getKey(), entry.getValue()); + } + builder.endObject(); + builder.endObject(); + } + + @Override + public String toString() { + return Strings.toString(this); + } +} diff --git a/server/src/main/java/org/opensearch/cluster/routing/IndexShardRoutingTable.java b/server/src/main/java/org/opensearch/cluster/routing/IndexShardRoutingTable.java index d4597f47d9a6c..9026e7068e9fe 100644 --- a/server/src/main/java/org/opensearch/cluster/routing/IndexShardRoutingTable.java +++ b/server/src/main/java/org/opensearch/cluster/routing/IndexShardRoutingTable.java @@ -85,6 +85,9 @@ public class IndexShardRoutingTable implements Iterable { private volatile Map activeShardsByAttributes = emptyMap(); private volatile Map initializingShardsByAttributes = emptyMap(); private final Object shardsByAttributeMutex = new Object(); + private final Object shardsByWeightMutex = new Object(); + private volatile Map> activeShardsByWeight = emptyMap(); + private volatile Map> initializingShardsByWeight = emptyMap(); /** * The initializing list, including ones that are initializing on a target node because of relocation. @@ -233,6 +236,10 @@ public List assignedShards() { return this.assignedShards; } + public Map> getActiveShardsByWeight() { + return activeShardsByWeight; + } + public ShardIterator shardsRandomIt() { return new PlainShardIterator(shardId, shuffler.shuffle(shards)); } @@ -292,6 +299,73 @@ 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 weightedRouting 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 activeInitializingShardsWeightedIt(WeightedRouting weightedRouting, DiscoveryNodes nodes, double defaultWeight) { + final int seed = shuffler.nextSeed(); + List ordered = new ArrayList<>(); + List orderedActiveShards = getActiveShardsByWeight(weightedRouting, nodes, defaultWeight); + ordered.addAll(shuffler.shuffle(orderedActiveShards, seed)); + if (!allInitializingShards.isEmpty()) { + List orderedInitializingShards = getInitializingShardsByWeight(weightedRouting, nodes, defaultWeight); + ordered.addAll(orderedInitializingShards); + } + return new PlainShardIterator(shardId, ordered); + } + + /** + * Returns a list containing shard routings ordered using weighted round-robin scheduling. + */ + private List shardsOrderedByWeight( + List shards, + WeightedRouting weightedRouting, + DiscoveryNodes nodes, + double defaultWeight + ) { + WeightedRoundRobin weightedRoundRobin = new WeightedRoundRobin<>( + calculateShardWeight(shards, weightedRouting, nodes, defaultWeight) + ); + List> shardsOrderedbyWeight = weightedRoundRobin.orderEntities(); + List orderedShardRouting = new ArrayList<>(activeShards.size()); + if (shardsOrderedbyWeight != null) { + for (WeightedRoundRobin.Entity shardRouting : shardsOrderedbyWeight) { + orderedShardRouting.add(shardRouting.getTarget()); + } + } + return orderedShardRouting; + } + + /** + * Returns a list containing shard routing and associated weight. This function iterates through all the shards and + * uses weighted routing to find weight for the corresponding shard. This is fed to weighted round-robin scheduling + * to order shards by weight. + */ + private List> calculateShardWeight( + List shards, + WeightedRouting weightedRouting, + DiscoveryNodes nodes, + double defaultWeight + ) { + List> shardsWithWeights = new ArrayList<>(); + for (ShardRouting shard : shards) { + DiscoveryNode node = nodes.get(shard.currentNodeId()); + if (node != null) { + String attVal = node.getAttributes().get(weightedRouting.attributeName()); + // If weight for a zone is not defined, considering it as 1 by default + Double weight = weightedRouting.weights().getOrDefault(attVal, defaultWeight); + shardsWithWeights.add(new WeightedRoundRobin.Entity<>(weight, shard)); + } + } + return shardsWithWeights; + } + private static Set getAllNodeIds(final List shards) { final Set nodeIds = new HashSet<>(); for (ShardRouting shard : shards) { @@ -698,6 +772,66 @@ public int shardsMatchingPredicateCount(Predicate predicate) { return count; } + /** + * Key for WeightedRouting Shard Iterator + * + * @opensearch.internal + */ + public static class WeightedRoutingKey { + private final WeightedRouting weightedRouting; + + public WeightedRoutingKey(WeightedRouting weightedRouting) { + this.weightedRouting = weightedRouting; + } + + @Override + public boolean equals(Object o) { + if (this == o) return true; + if (o == null || getClass() != o.getClass()) return false; + WeightedRoutingKey key = (WeightedRoutingKey) o; + if (!weightedRouting.equals(key.weightedRouting)) return false; + return true; + } + + @Override + public int hashCode() { + int result = weightedRouting.hashCode(); + return result; + } + } + + /** + * * + * Gets active shard routing from memory if available, else calculates and put it in memory. + */ + private List getActiveShardsByWeight(WeightedRouting weightedRouting, DiscoveryNodes nodes, double defaultWeight) { + WeightedRoutingKey key = new WeightedRoutingKey(weightedRouting); + List shardRoutings = activeShardsByWeight.get(key); + if (shardRoutings == null) { + synchronized (shardsByWeightMutex) { + shardRoutings = shardsOrderedByWeight(activeShards, weightedRouting, nodes, defaultWeight); + activeShardsByWeight = new MapBuilder().put(key, shardRoutings).immutableMap(); + } + } + return shardRoutings; + } + + /** + * * + * Gets initializing shard routing from memory if available, else calculates and put it in memory. + */ + private List getInitializingShardsByWeight(WeightedRouting weightedRouting, DiscoveryNodes nodes, double defaultWeight) { + WeightedRoutingKey key = new WeightedRoutingKey(weightedRouting); + List shardRoutings = initializingShardsByWeight.get(key); + if (shardRoutings == null) { + synchronized (shardsByWeightMutex) { + shardRoutings = shardsOrderedByWeight(activeShards, weightedRouting, nodes, defaultWeight); + initializingShardsByWeight = new MapBuilder().put(key, shardRoutings).immutableMap(); + } + } + return shardRoutings; + } + /** * Builder of an index shard routing table. * diff --git a/server/src/main/java/org/opensearch/cluster/routing/OperationRouting.java b/server/src/main/java/org/opensearch/cluster/routing/OperationRouting.java index 30f6408c19783..9026da667ccb0 100644 --- a/server/src/main/java/org/opensearch/cluster/routing/OperationRouting.java +++ b/server/src/main/java/org/opensearch/cluster/routing/OperationRouting.java @@ -34,6 +34,7 @@ import org.opensearch.cluster.ClusterState; import org.opensearch.cluster.metadata.IndexMetadata; +import org.opensearch.cluster.metadata.WeightedRoutingMetadata; import org.opensearch.cluster.node.DiscoveryNodes; import org.opensearch.cluster.routing.allocation.decider.AwarenessAllocationDecider; import org.opensearch.common.Nullable; @@ -75,9 +76,17 @@ public class OperationRouting { Setting.Property.Dynamic, Setting.Property.NodeScope ); + public static final Setting WEIGHTED_ROUTING_DEFAULT_WEIGHT = Setting.doubleSetting( + "cluster.routing.weighted.default_weight", + 1.0, + 1.0, + Setting.Property.Dynamic, + Setting.Property.NodeScope + ); private volatile List awarenessAttributes; private volatile boolean useAdaptiveReplicaSelection; private volatile boolean ignoreAwarenessAttr; + private volatile double weightedRoutingDefaultWeight; public OperationRouting(Settings settings, ClusterSettings clusterSettings) { // whether to ignore awareness attributes when routing requests @@ -88,8 +97,10 @@ public OperationRouting(Settings settings, ClusterSettings clusterSettings) { this::setAwarenessAttributes ); this.useAdaptiveReplicaSelection = USE_ADAPTIVE_REPLICA_SELECTION_SETTING.get(settings); + this.weightedRoutingDefaultWeight = WEIGHTED_ROUTING_DEFAULT_WEIGHT.get(settings); clusterSettings.addSettingsUpdateConsumer(USE_ADAPTIVE_REPLICA_SELECTION_SETTING, this::setUseAdaptiveReplicaSelection); clusterSettings.addSettingsUpdateConsumer(IGNORE_AWARENESS_ATTRIBUTES_SETTING, this::setIgnoreAwarenessAttributes); + clusterSettings.addSettingsUpdateConsumer(WEIGHTED_ROUTING_DEFAULT_WEIGHT, this::setWeightedRoutingDefaultWeight); } void setUseAdaptiveReplicaSelection(boolean useAdaptiveReplicaSelection) { @@ -100,6 +111,10 @@ void setIgnoreAwarenessAttributes(boolean ignoreAwarenessAttributes) { this.ignoreAwarenessAttr = ignoreAwarenessAttributes; } + void setWeightedRoutingDefaultWeight(double weightedRoutingDefaultWeight) { + this.weightedRoutingDefaultWeight = weightedRoutingDefaultWeight; + } + public boolean isIgnoreAwarenessAttr() { return ignoreAwarenessAttr; } @@ -116,6 +131,10 @@ public boolean ignoreAwarenessAttributes() { return this.awarenessAttributes.isEmpty() || this.ignoreAwarenessAttr; } + public double getWeightedRoutingDefaultWeight() { + return this.weightedRoutingDefaultWeight; + } + public ShardIterator indexShards(ClusterState clusterState, String index, String id, @Nullable String routing) { return shards(clusterState, index, id, routing).shardsIt(); } @@ -133,7 +152,8 @@ public ShardIterator getShards( clusterState.nodes(), preference, null, - null + null, + clusterState.getMetadata().weightedRoutingMetadata() ); } @@ -145,7 +165,8 @@ public ShardIterator getShards(ClusterState clusterState, String index, int shar clusterState.nodes(), preference, null, - null + null, + clusterState.metadata().weightedRoutingMetadata() ); } @@ -175,7 +196,8 @@ public GroupShardsIterator searchShards( clusterState.nodes(), preference, collectorService, - nodeCounts + nodeCounts, + clusterState.metadata().weightedRoutingMetadata() ); if (iterator != null) { set.add(iterator); @@ -225,10 +247,11 @@ private ShardIterator preferenceActiveShardIterator( DiscoveryNodes nodes, @Nullable String preference, @Nullable ResponseCollectorService collectorService, - @Nullable Map nodeCounts + @Nullable Map nodeCounts, + @Nullable WeightedRoutingMetadata weightedRoutingMetadata ) { if (preference == null || preference.isEmpty()) { - return shardRoutings(indexShard, nodes, collectorService, nodeCounts); + return shardRoutings(indexShard, nodes, collectorService, nodeCounts, weightedRoutingMetadata); } if (preference.charAt(0) == '_') { Preference preferenceType = Preference.parse(preference); @@ -255,7 +278,7 @@ private ShardIterator preferenceActiveShardIterator( } // no more preference if (index == -1 || index == preference.length() - 1) { - return shardRoutings(indexShard, nodes, collectorService, nodeCounts); + return shardRoutings(indexShard, nodes, collectorService, nodeCounts, weightedRoutingMetadata); } else { // update the preference and continue preference = preference.substring(index + 1); @@ -298,9 +321,16 @@ private ShardIterator shardRoutings( IndexShardRoutingTable indexShard, DiscoveryNodes nodes, @Nullable ResponseCollectorService collectorService, - @Nullable Map nodeCounts + @Nullable Map nodeCounts, + @Nullable WeightedRoutingMetadata weightedRoutingMetadata ) { - if (ignoreAwarenessAttributes()) { + if (weightedRoutingMetadata != null) { + return indexShard.activeInitializingShardsWeightedIt( + weightedRoutingMetadata.getWeightedRouting(), + nodes, + getWeightedRoutingDefaultWeight() + ); + } else if (ignoreAwarenessAttributes()) { if (useAdaptiveReplicaSelection) { return indexShard.activeInitializingShardsRankedIt(collectorService, nodeCounts); } else { diff --git a/server/src/main/java/org/opensearch/cluster/routing/WeightedRoundRobin.java b/server/src/main/java/org/opensearch/cluster/routing/WeightedRoundRobin.java new file mode 100644 index 0000000000000..15d437db9c8ff --- /dev/null +++ b/server/src/main/java/org/opensearch/cluster/routing/WeightedRoundRobin.java @@ -0,0 +1,106 @@ +/* + * 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.routing; + +import java.util.ArrayList; +import java.util.List; + +/** + * Weighted Round Robin Scheduling policy + * + */ +public class WeightedRoundRobin { + + private List> entities; + + public WeightedRoundRobin(List> entities) { + this.entities = entities; + } + + /** + * * + * @return list of entities that is ordered using weighted round-robin scheduling + * http://kb.linuxvirtualserver.org/wiki/Weighted_Round-Robin_Scheduling + */ + public List> orderEntities() { + int lastSelectedEntity = -1; + int size = entities.size(); + double currentWeight = 0; + List> orderedWeight = new ArrayList<>(); + if (size == 0) { + return null; + } + // Find maximum weight and greatest common divisor of weight across all entities + double maxWeight = 0; + double sumWeight = 0; + Double gcd = null; + for (WeightedRoundRobin.Entity entity : entities) { + maxWeight = Math.max(maxWeight, entity.getWeight()); + gcd = (gcd == null) ? entity.getWeight() : gcd(gcd, entity.getWeight()); + sumWeight += entity.getWeight() > 0 ? entity.getWeight() : 0; + } + int count = 0; + while (count < sumWeight) { + lastSelectedEntity = (lastSelectedEntity + 1) % size; + if (lastSelectedEntity == 0) { + currentWeight = currentWeight - gcd; + if (currentWeight <= 0) { + currentWeight = maxWeight; + if (currentWeight == 0) { + return orderedWeight; + } + } + } + if (entities.get(lastSelectedEntity).getWeight() >= currentWeight) { + orderedWeight.add(entities.get(lastSelectedEntity)); + count++; + } + } + return orderedWeight; + } + + /** + * Return greatest common divisor for two integers + * https://en.wikipedia.org/wiki/Greatest_common_divisor#Using_Euclid.27s_algorithm + * + * @param a first number + * @param b second number + * @return greatest common divisor + */ + private double gcd(double a, double b) { + return (b == 0) ? a : gcd(b, a % b); + } + + static final class Entity { + private double weight; + private T target; + + public Entity(double weight, T target) { + this.weight = weight; + this.target = target; + } + + public T getTarget() { + return this.target; + } + + public void setTarget(T target) { + this.target = target; + } + + public double getWeight() { + return this.weight; + } + + public void setWeight(double weight) { + this.weight = weight; + } + } + +} diff --git a/server/src/main/java/org/opensearch/cluster/routing/WeightedRouting.java b/server/src/main/java/org/opensearch/cluster/routing/WeightedRouting.java new file mode 100644 index 0000000000000..df2d8d595eaab --- /dev/null +++ b/server/src/main/java/org/opensearch/cluster/routing/WeightedRouting.java @@ -0,0 +1,75 @@ +/* + * 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.routing; + +import org.opensearch.common.io.stream.StreamInput; +import org.opensearch.common.io.stream.StreamOutput; +import org.opensearch.common.io.stream.Writeable; + +import java.io.IOException; +import java.util.Map; +import java.util.Objects; + +/** + * Entity for Weighted Round Robin weights + * + * @opensearch.internal + */ +public class WeightedRouting implements Writeable { + private String attributeName; + private Map weights; + + public WeightedRouting(String attributeName, Map weights) { + this.attributeName = attributeName; + this.weights = weights; + } + + public WeightedRouting(WeightedRouting weightedRouting) { + this.attributeName = weightedRouting.attributeName(); + this.weights = weightedRouting.weights; + } + + public WeightedRouting(StreamInput in) throws IOException { + attributeName = in.readString(); + weights = (Map) in.readGenericValue(); + } + + @Override + public void writeTo(StreamOutput out) throws IOException { + out.writeString(attributeName); + out.writeGenericValue(weights); + } + + @Override + public boolean equals(Object o) { + if (this == o) return true; + if (o == null || getClass() != o.getClass()) return false; + WeightedRouting that = (WeightedRouting) o; + if (!attributeName.equals(that.attributeName)) return false; + return weights.equals(that.weights); + } + + @Override + public int hashCode() { + return Objects.hash(attributeName, weights); + } + + @Override + public String toString() { + return "WeightedRouting{" + attributeName + "}{" + weights().toString() + "}"; + } + + public Map weights() { + return this.weights; + } + + public String attributeName() { + return this.attributeName; + } +} diff --git a/server/src/main/java/org/opensearch/cluster/routing/WeightedRoutingService.java b/server/src/main/java/org/opensearch/cluster/routing/WeightedRoutingService.java new file mode 100644 index 0000000000000..54f2c81aea384 --- /dev/null +++ b/server/src/main/java/org/opensearch/cluster/routing/WeightedRoutingService.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.routing; + +import org.apache.logging.log4j.LogManager; +import org.apache.logging.log4j.Logger; +import org.apache.logging.log4j.message.ParameterizedMessage; +import org.opensearch.action.ActionListener; +import org.opensearch.action.admin.cluster.shards.routing.weighted.delete.ClusterDeleteWeightedRoutingRequest; +import org.opensearch.action.ActionRequestValidationException; +import org.opensearch.action.admin.cluster.shards.routing.weighted.delete.ClusterDeleteWeightedRoutingResponse; +import org.opensearch.action.admin.cluster.shards.routing.weighted.put.ClusterPutWeightedRoutingRequest; +import org.opensearch.cluster.ClusterState; +import org.opensearch.cluster.ClusterStateUpdateTask; +import org.opensearch.cluster.ack.ClusterStateUpdateResponse; +import org.opensearch.cluster.metadata.Metadata; +import org.opensearch.cluster.metadata.WeightedRoutingMetadata; +import org.opensearch.cluster.routing.allocation.decider.AwarenessAllocationDecider; +import org.opensearch.cluster.service.ClusterService; +import org.opensearch.common.Priority; +import org.opensearch.common.inject.Inject; + +import org.opensearch.common.settings.ClusterSettings; +import org.opensearch.common.settings.Settings; +import org.opensearch.threadpool.ThreadPool; + +import java.util.List; +import java.util.Locale; + +import static org.opensearch.action.ValidateActions.addValidationError; + +/** + * * Service responsible for updating cluster state metadata with weighted routing weights + */ +public class WeightedRoutingService { + private static final Logger logger = LogManager.getLogger(WeightedRoutingService.class); + private final ClusterService clusterService; + private final ThreadPool threadPool; + private volatile List awarenessAttributes; + + @Inject + public WeightedRoutingService( + ClusterService clusterService, + ThreadPool threadPool, + Settings settings, + ClusterSettings clusterSettings + ) { + this.clusterService = clusterService; + this.threadPool = threadPool; + this.awarenessAttributes = AwarenessAllocationDecider.CLUSTER_ROUTING_ALLOCATION_AWARENESS_ATTRIBUTE_SETTING.get(settings); + clusterSettings.addSettingsUpdateConsumer( + AwarenessAllocationDecider.CLUSTER_ROUTING_ALLOCATION_AWARENESS_ATTRIBUTE_SETTING, + this::setAwarenessAttributes + ); + } + + public void registerWeightedRoutingMetadata( + final ClusterPutWeightedRoutingRequest request, + final ActionListener listener + ) { + final WeightedRoutingMetadata newWeightedRoutingMetadata = new WeightedRoutingMetadata(request.getWeightedRouting()); + clusterService.submitStateUpdateTask("update_weighted_routing", new ClusterStateUpdateTask(Priority.URGENT) { + @Override + public ClusterState execute(ClusterState currentState) { + Metadata metadata = currentState.metadata(); + Metadata.Builder mdBuilder = Metadata.builder(currentState.metadata()); + WeightedRoutingMetadata weightedRoutingMetadata = metadata.custom(WeightedRoutingMetadata.TYPE); + if (weightedRoutingMetadata == null) { + logger.info("put weighted routing weights in metadata [{}]", request.getWeightedRouting()); + weightedRoutingMetadata = new WeightedRoutingMetadata(request.getWeightedRouting()); + } else { + if (!checkIfSameWeightsInMetadata(newWeightedRoutingMetadata, weightedRoutingMetadata)) { + logger.info("updated weighted routing weights [{}] in metadata", request.getWeightedRouting()); + weightedRoutingMetadata = new WeightedRoutingMetadata(newWeightedRoutingMetadata.getWeightedRouting()); + } else { + return currentState; + } + } + mdBuilder.putCustom(WeightedRoutingMetadata.TYPE, weightedRoutingMetadata); + logger.info("building cluster state with weighted routing weights [{}]", request.getWeightedRouting()); + return ClusterState.builder(currentState).metadata(mdBuilder).build(); + } + + @Override + public void onFailure(String source, Exception e) { + logger.warn(() -> new ParameterizedMessage("failed to update cluster state for weighted routing weights [{}]", e)); + listener.onFailure(e); + } + + @Override + public void clusterStateProcessed(String source, ClusterState oldState, ClusterState newState) { + logger.debug("cluster weighted routing weights metadata change is processed by all the nodes"); + listener.onResponse(new ClusterStateUpdateResponse(true)); + } + }); + } + + private boolean checkIfSameWeightsInMetadata( + WeightedRoutingMetadata newWeightedRoutingMetadata, + WeightedRoutingMetadata oldWeightedRoutingMetadata + ) { + return newWeightedRoutingMetadata.getWeightedRouting().equals(oldWeightedRoutingMetadata.getWeightedRouting()); + } + + public void deleteWeightedRoutingMetadata( + final ClusterDeleteWeightedRoutingRequest request, + final ActionListener listener + ) { + clusterService.submitStateUpdateTask("delete_weighted_routing", new ClusterStateUpdateTask(Priority.URGENT) { + @Override + public ClusterState execute(ClusterState currentState) { + logger.info("Deleting weighted routing metadata from the cluster state"); + Metadata.Builder mdBuilder = Metadata.builder(currentState.metadata()); + mdBuilder.removeCustom(WeightedRoutingMetadata.TYPE); + return ClusterState.builder(currentState).metadata(mdBuilder).build(); + } + + @Override + public void onFailure(String source, Exception e) { + logger.error("failed to remove weighted routing metadata from cluster state", e); + listener.onFailure(e); + } + + @Override + public void clusterStateProcessed(String source, ClusterState oldState, ClusterState newState) { + logger.debug("cluster weighted routing metadata change is processed by all the nodes"); + assert newState.metadata().weightedRoutingMetadata() == null; + listener.onResponse(new ClusterDeleteWeightedRoutingResponse(true)); + } + }); + } + + List getAwarenessAttributes() { + return awarenessAttributes; + } + + private void setAwarenessAttributes(List awarenessAttributes) { + this.awarenessAttributes = awarenessAttributes; + } + + public void verifyAwarenessAttribute(String attributeName) { + if (getAwarenessAttributes().contains(attributeName) == false) { + ActionRequestValidationException validationException = null; + validationException = addValidationError( + String.format(Locale.ROOT, "invalid awareness attribute %s requested for updating weighted routing", attributeName), + validationException + ); + throw validationException; + } + } +} diff --git a/server/src/main/java/org/opensearch/common/settings/ClusterSettings.java b/server/src/main/java/org/opensearch/common/settings/ClusterSettings.java index 003dfeb11504b..9d07e306cfa09 100644 --- a/server/src/main/java/org/opensearch/common/settings/ClusterSettings.java +++ b/server/src/main/java/org/opensearch/common/settings/ClusterSettings.java @@ -537,6 +537,7 @@ public void apply(Settings value, Settings current, Settings previous) { Node.BREAKER_TYPE_KEY, OperationRouting.USE_ADAPTIVE_REPLICA_SELECTION_SETTING, OperationRouting.IGNORE_AWARENESS_ATTRIBUTES_SETTING, + OperationRouting.WEIGHTED_ROUTING_DEFAULT_WEIGHT, IndexGraveyard.SETTING_MAX_TOMBSTONES, PersistentTasksClusterService.CLUSTER_TASKS_ALLOCATION_RECHECK_INTERVAL_SETTING, EnableAssignmentDecider.CLUSTER_TASKS_ALLOCATION_ENABLE_SETTING, diff --git a/server/src/main/java/org/opensearch/rest/action/admin/cluster/RestClusterDeleteWeightedRoutingAction.java b/server/src/main/java/org/opensearch/rest/action/admin/cluster/RestClusterDeleteWeightedRoutingAction.java new file mode 100644 index 0000000000000..9742cc373d520 --- /dev/null +++ b/server/src/main/java/org/opensearch/rest/action/admin/cluster/RestClusterDeleteWeightedRoutingAction.java @@ -0,0 +1,53 @@ +/* + * 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.rest.action.admin.cluster; + +import org.apache.logging.log4j.LogManager; +import org.apache.logging.log4j.Logger; +import org.opensearch.action.admin.cluster.shards.routing.weighted.delete.ClusterDeleteWeightedRoutingRequest; +import org.opensearch.client.Requests; +import org.opensearch.client.node.NodeClient; +import org.opensearch.rest.BaseRestHandler; +import org.opensearch.rest.RestRequest; +import org.opensearch.rest.action.RestToXContentListener; + +import java.io.IOException; +import java.util.List; + +import static java.util.Collections.singletonList; +import static org.opensearch.rest.RestRequest.Method.DELETE; + +/** + * Delete Weighted Round Robin based shard routing weights + * + * @opensearch.api + * + */ +public class RestClusterDeleteWeightedRoutingAction extends BaseRestHandler { + + private static final Logger logger = LogManager.getLogger(RestClusterDeleteWeightedRoutingAction.class); + + @Override + public List routes() { + return singletonList(new Route(DELETE, "/_cluster/routing/awareness/weights")); + } + + @Override + public String getName() { + return "delete_weighted_routing_action"; + } + + @Override + protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient client) throws IOException { + ClusterDeleteWeightedRoutingRequest clusterDeleteWeightedRoutingRequest = Requests.deleteWeightedRoutingRequest(); + return channel -> client.admin() + .cluster() + .deleteWeightedRouting(clusterDeleteWeightedRoutingRequest, new RestToXContentListener<>(channel)); + } +} diff --git a/server/src/main/java/org/opensearch/rest/action/admin/cluster/RestClusterGetWeightedRoutingAction.java b/server/src/main/java/org/opensearch/rest/action/admin/cluster/RestClusterGetWeightedRoutingAction.java new file mode 100644 index 0000000000000..7c9d1190f0b1d --- /dev/null +++ b/server/src/main/java/org/opensearch/rest/action/admin/cluster/RestClusterGetWeightedRoutingAction.java @@ -0,0 +1,52 @@ +/* + * 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.rest.action.admin.cluster; + +import org.apache.logging.log4j.LogManager; +import org.apache.logging.log4j.Logger; +import org.opensearch.action.admin.cluster.shards.routing.weighted.get.ClusterGetWeightedRoutingRequest; +import org.opensearch.client.Requests; +import org.opensearch.client.node.NodeClient; +import org.opensearch.rest.BaseRestHandler; +import org.opensearch.rest.RestRequest; +import org.opensearch.rest.action.RestToXContentListener; + +import java.io.IOException; +import java.util.List; + +import static java.util.Collections.singletonList; +import static org.opensearch.rest.RestRequest.Method.GET; + +/** + * Fetch Weighted Round Robin based shard routing weights + * + * @opensearch.api + * + */ +public class RestClusterGetWeightedRoutingAction extends BaseRestHandler { + + private static final Logger logger = LogManager.getLogger(RestClusterGetWeightedRoutingAction.class); + + @Override + public List routes() { + return singletonList(new Route(GET, "/_cluster/routing/awareness/{attribute}/weights")); + } + + @Override + public String getName() { + return "get_weighted_routing_action"; + } + + @Override + protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient client) throws IOException { + ClusterGetWeightedRoutingRequest getWeightedRoutingRequest = Requests.getWeightedRoutingRequest(request.param("attribute")); + getWeightedRoutingRequest.local(request.paramAsBoolean("local", getWeightedRoutingRequest.local())); + return channel -> client.admin().cluster().getWeightedRouting(getWeightedRoutingRequest, new RestToXContentListener<>(channel)); + } +} diff --git a/server/src/main/java/org/opensearch/rest/action/admin/cluster/RestClusterPutWeightedRoutingAction.java b/server/src/main/java/org/opensearch/rest/action/admin/cluster/RestClusterPutWeightedRoutingAction.java new file mode 100644 index 0000000000000..1cf44e665cf84 --- /dev/null +++ b/server/src/main/java/org/opensearch/rest/action/admin/cluster/RestClusterPutWeightedRoutingAction.java @@ -0,0 +1,58 @@ +/* + * 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.rest.action.admin.cluster; + +import org.apache.logging.log4j.LogManager; +import org.apache.logging.log4j.Logger; +import org.opensearch.action.admin.cluster.shards.routing.weighted.put.ClusterPutWeightedRoutingRequest; +import org.opensearch.client.Requests; +import org.opensearch.client.node.NodeClient; +import org.opensearch.rest.BaseRestHandler; +import org.opensearch.rest.RestRequest; +import org.opensearch.rest.action.RestToXContentListener; + +import java.io.IOException; +import java.util.List; + +import static java.util.Collections.singletonList; +import static org.opensearch.rest.RestRequest.Method.PUT; + +/** + * Update Weighted Round Robin based shard routing weights + * + * @opensearch.api + * + */ +public class RestClusterPutWeightedRoutingAction extends BaseRestHandler { + + private static final Logger logger = LogManager.getLogger(RestClusterPutWeightedRoutingAction.class); + + @Override + public List routes() { + return singletonList(new Route(PUT, "/_cluster/routing/awareness/{attribute}/weights")); + } + + @Override + public String getName() { + return "put_weighted_routing_action"; + } + + @Override + protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient client) throws IOException { + ClusterPutWeightedRoutingRequest putWeightedRoutingRequest = createRequest(request); + return channel -> client.admin().cluster().putWeightedRouting(putWeightedRoutingRequest, new RestToXContentListener<>(channel)); + } + + public static ClusterPutWeightedRoutingRequest createRequest(RestRequest request) throws IOException { + ClusterPutWeightedRoutingRequest putWeightedRoutingRequest = Requests.putWeightedRoutingRequest(request.param("attribute")); + request.applyContentParser(p -> putWeightedRoutingRequest.source(p.mapStrings())); + return putWeightedRoutingRequest; + } + +} diff --git a/server/src/test/java/org/opensearch/action/admin/cluster/shards/routing/weighted/put/ClusterPutWeightedRoutingRequestTests.java b/server/src/test/java/org/opensearch/action/admin/cluster/shards/routing/weighted/put/ClusterPutWeightedRoutingRequestTests.java new file mode 100644 index 0000000000000..186e7e8638f17 --- /dev/null +++ b/server/src/test/java/org/opensearch/action/admin/cluster/shards/routing/weighted/put/ClusterPutWeightedRoutingRequestTests.java @@ -0,0 +1,65 @@ +/* + * 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.action.admin.cluster.shards.routing.weighted.put; + +import org.opensearch.action.ActionRequestValidationException; +import org.opensearch.cluster.routing.WeightedRouting; +import org.opensearch.common.bytes.BytesArray; +import org.opensearch.common.xcontent.XContentType; +import org.opensearch.test.OpenSearchTestCase; + +import java.util.Map; + +public class ClusterPutWeightedRoutingRequestTests extends OpenSearchTestCase { + + public void testSetWeightedRoutingWeight() { + String reqString = "{\"us-east-1c\" : \"0\", \"us-east-1b\":\"1\",\"us-east-1a\":\"1\"}"; + ClusterPutWeightedRoutingRequest request = new ClusterPutWeightedRoutingRequest("zone"); + Map weights = Map.of("us-east-1a", 1.0, "us-east-1b", 1.0, "us-east-1c", 0.0); + WeightedRouting weightedRouting = new WeightedRouting("zone", weights); + request.setWeightedRouting(new BytesArray(reqString), XContentType.JSON); + assertEquals(request.getWeightedRouting(), weightedRouting); + } + + public void testValidate_ValuesAreProper() { + String reqString = "{\"us-east-1c\" : \"1\", \"us-east-1b\":\"0\",\"us-east-1a\":\"1\"}"; + ClusterPutWeightedRoutingRequest request = new ClusterPutWeightedRoutingRequest("zone"); + request.setWeightedRouting(new BytesArray(reqString), XContentType.JSON); + ActionRequestValidationException actionRequestValidationException = request.validate(); + assertNull(actionRequestValidationException); + } + + public void testValidate_TwoZonesWithZeroWeight() { + String reqString = "{\"us-east-1c\" : \"0\", \"us-east-1b\":\"0\",\"us-east-1a\":\"1\"}"; + ClusterPutWeightedRoutingRequest request = new ClusterPutWeightedRoutingRequest("zone"); + request.setWeightedRouting(new BytesArray(reqString), XContentType.JSON); + ActionRequestValidationException actionRequestValidationException = request.validate(); + assertNotNull(actionRequestValidationException); + assertTrue(actionRequestValidationException.getMessage().contains("More than one [2] value has weight set as " + "0")); + } + + public void testValidate_MissingWeights() { + String reqString = "{}"; + ClusterPutWeightedRoutingRequest request = new ClusterPutWeightedRoutingRequest("zone"); + request.setWeightedRouting(new BytesArray(reqString), XContentType.JSON); + ActionRequestValidationException actionRequestValidationException = request.validate(); + assertNotNull(actionRequestValidationException); + assertTrue(actionRequestValidationException.getMessage().contains("Weights are missing")); + } + + public void testValidate_AttributeMissing() { + String reqString = "{\"us-east-1c\" : \"0\", \"us-east-1b\":\"1\",\"us-east-1a\":\"1\"}"; + ClusterPutWeightedRoutingRequest request = new ClusterPutWeightedRoutingRequest(); + request.setWeightedRouting(new BytesArray(reqString), XContentType.JSON); + ActionRequestValidationException actionRequestValidationException = request.validate(); + assertNotNull(actionRequestValidationException); + assertTrue(actionRequestValidationException.getMessage().contains("Attribute name is missing")); + } + +} diff --git a/server/src/test/java/org/opensearch/cluster/action/shard/routing/weighted/get/ClusterGetWeightedRoutingRequestTests.java b/server/src/test/java/org/opensearch/cluster/action/shard/routing/weighted/get/ClusterGetWeightedRoutingRequestTests.java new file mode 100644 index 0000000000000..0a4dad4cbc597 --- /dev/null +++ b/server/src/test/java/org/opensearch/cluster/action/shard/routing/weighted/get/ClusterGetWeightedRoutingRequestTests.java @@ -0,0 +1,39 @@ +/* + * 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.action.shard.routing.weighted.get; + +import org.opensearch.action.ActionRequestValidationException; +import org.opensearch.action.admin.cluster.shards.routing.weighted.get.ClusterGetWeightedRoutingRequest; +import org.opensearch.test.OpenSearchTestCase; + +public class ClusterGetWeightedRoutingRequestTests extends OpenSearchTestCase { + + public void testValidate_AwarenessAttributeIsSet() { + ClusterGetWeightedRoutingRequest request = new ClusterGetWeightedRoutingRequest(); + request.setAwarenessAttribute("zone"); + ActionRequestValidationException actionRequestValidationException = request.validate(); + assertNull(actionRequestValidationException); + } + + public void testValidate_AwarenessAttributeNotSet() { + ClusterGetWeightedRoutingRequest request = new ClusterGetWeightedRoutingRequest(); + ActionRequestValidationException actionRequestValidationException = request.validate(); + assertNotNull(actionRequestValidationException); + assertTrue(actionRequestValidationException.getMessage().contains("Awareness attribute is missing")); + } + + public void testValidate_AwarenessAttributeIsEmpty() { + ClusterGetWeightedRoutingRequest request = new ClusterGetWeightedRoutingRequest(); + request.setAwarenessAttribute(""); + ActionRequestValidationException actionRequestValidationException = request.validate(); + assertNotNull(actionRequestValidationException); + assertTrue(actionRequestValidationException.getMessage().contains("Awareness attribute is missing")); + } + +} diff --git a/server/src/test/java/org/opensearch/cluster/action/shard/routing/weighted/get/ClusterGetWeightedRoutingResponseTests.java b/server/src/test/java/org/opensearch/cluster/action/shard/routing/weighted/get/ClusterGetWeightedRoutingResponseTests.java new file mode 100644 index 0000000000000..e9add55ca774b --- /dev/null +++ b/server/src/test/java/org/opensearch/cluster/action/shard/routing/weighted/get/ClusterGetWeightedRoutingResponseTests.java @@ -0,0 +1,38 @@ +/* + * 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.action.shard.routing.weighted.get; + +import org.opensearch.action.admin.cluster.shards.routing.weighted.get.ClusterGetWeightedRoutingResponse; +import org.opensearch.cluster.routing.WeightedRouting; +import org.opensearch.common.xcontent.XContentParser; +import org.opensearch.test.AbstractXContentTestCase; + +import java.io.IOException; +import java.util.Map; + +public class ClusterGetWeightedRoutingResponseTests extends AbstractXContentTestCase { + @Override + protected ClusterGetWeightedRoutingResponse createTestInstance() { + Map weights = Map.of("zone_A", 1.0, "zone_B", 0.0, "zone_C", 1.0); + WeightedRouting weightedRouting = new WeightedRouting("", weights); + ClusterGetWeightedRoutingResponse response = new ClusterGetWeightedRoutingResponse("1", weightedRouting); + return response; + } + + @Override + protected ClusterGetWeightedRoutingResponse doParseInstance(XContentParser parser) throws IOException { + return ClusterGetWeightedRoutingResponse.fromXContent(parser); + } + + @Override + protected boolean supportsUnknownFields() { + return false; + } + +} diff --git a/server/src/test/java/org/opensearch/cluster/action/shard/routing/weighted/get/TransportGetWeightedRoutingActionTests.java b/server/src/test/java/org/opensearch/cluster/action/shard/routing/weighted/get/TransportGetWeightedRoutingActionTests.java new file mode 100644 index 0000000000000..f28e932e068ac --- /dev/null +++ b/server/src/test/java/org/opensearch/cluster/action/shard/routing/weighted/get/TransportGetWeightedRoutingActionTests.java @@ -0,0 +1,262 @@ +/* + * 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.action.shard.routing.weighted.get; + +import org.junit.After; +import org.junit.Before; +import org.opensearch.Version; +import org.opensearch.action.admin.cluster.shards.routing.weighted.get.ClusterGetWeightedRoutingAction; +import org.opensearch.action.admin.cluster.shards.routing.weighted.get.ClusterGetWeightedRoutingRequestBuilder; +import org.opensearch.action.admin.cluster.shards.routing.weighted.get.ClusterGetWeightedRoutingResponse; +import org.opensearch.action.admin.cluster.shards.routing.weighted.get.TransportGetWeightedRoutingAction; +import org.opensearch.action.support.ActionFilters; +import org.opensearch.action.support.ActionTestUtils; +import org.opensearch.client.node.NodeClient; +import org.opensearch.cluster.ClusterName; +import org.opensearch.cluster.ClusterState; +import org.opensearch.cluster.metadata.IndexNameExpressionResolver; +import org.opensearch.cluster.metadata.Metadata; +import org.opensearch.cluster.metadata.WeightedRoutingMetadata; +import org.opensearch.cluster.node.DiscoveryNode; +import org.opensearch.cluster.node.DiscoveryNodeRole; +import org.opensearch.cluster.node.DiscoveryNodes; +import org.opensearch.cluster.routing.WeightedRouting; +import org.opensearch.cluster.routing.WeightedRoutingService; +import org.opensearch.cluster.routing.allocation.decider.AwarenessAllocationDecider; +import org.opensearch.cluster.service.ClusterService; +import org.opensearch.common.settings.ClusterSettings; +import org.opensearch.common.settings.Settings; +import org.opensearch.test.ClusterServiceUtils; +import org.opensearch.test.OpenSearchTestCase; +import org.opensearch.test.transport.MockTransport; +import org.opensearch.threadpool.TestThreadPool; +import org.opensearch.threadpool.ThreadPool; +import org.opensearch.transport.TransportService; + +import java.util.Collections; +import java.util.HashSet; +import java.util.Map; +import java.util.Set; + +import static java.util.Collections.emptySet; +import static org.mockito.Mockito.mock; + +public class TransportGetWeightedRoutingActionTests extends OpenSearchTestCase { + + private ThreadPool threadPool; + private ClusterService clusterService; + private TransportService transportService; + private WeightedRoutingService weightedRoutingService; + private TransportGetWeightedRoutingAction transportGetWeightedRoutingAction; + private ClusterSettings clusterSettings; + NodeClient client; + + final private static Set CLUSTER_MANAGER_ROLE = Collections.unmodifiableSet( + new HashSet<>(Collections.singletonList(DiscoveryNodeRole.CLUSTER_MANAGER_ROLE)) + ); + + final private static Set DATA_ROLE = Collections.unmodifiableSet( + new HashSet<>(Collections.singletonList(DiscoveryNodeRole.DATA_ROLE)) + ); + + @Override + public void setUp() throws Exception { + super.setUp(); + threadPool = new TestThreadPool("test", Settings.EMPTY); + clusterService = ClusterServiceUtils.createClusterService(threadPool); + } + + @Before + public void setUpService() { + ClusterState clusterState = ClusterState.builder(new ClusterName("test")).build(); + clusterState = addClusterManagerNodes(clusterState); + clusterState = addDataNodes(clusterState); + clusterState = setLocalNode(clusterState, "nodeA1"); + + ClusterState.Builder builder = ClusterState.builder(clusterState); + ClusterServiceUtils.setState(clusterService, builder); + + final MockTransport transport = new MockTransport(); + transportService = transport.createTransportService( + Settings.EMPTY, + threadPool, + TransportService.NOOP_TRANSPORT_INTERCEPTOR, + boundTransportAddress -> clusterService.state().nodes().get("nodes1"), + null, + Collections.emptySet() + + ); + + Settings.Builder settingsBuilder = Settings.builder() + .put(AwarenessAllocationDecider.CLUSTER_ROUTING_ALLOCATION_AWARENESS_ATTRIBUTE_SETTING.getKey(), "zone"); + + clusterSettings = new ClusterSettings(settingsBuilder.build(), ClusterSettings.BUILT_IN_CLUSTER_SETTINGS); + transportService.start(); + transportService.acceptIncomingRequests(); + + this.weightedRoutingService = new WeightedRoutingService(clusterService, threadPool, settingsBuilder.build(), clusterSettings); + + this.transportGetWeightedRoutingAction = new TransportGetWeightedRoutingAction( + transportService, + clusterService, + weightedRoutingService, + threadPool, + new ActionFilters(emptySet()), + mock(IndexNameExpressionResolver.class) + ); + client = new NodeClient(Settings.EMPTY, threadPool); + } + + private ClusterState addDataNodes(ClusterState clusterState) { + clusterState = addDataNodeForAZone(clusterState, "zone_A", "nodeA1", "nodeA2", "nodeA3"); + clusterState = addDataNodeForAZone(clusterState, "zone_B", "nodeB1", "nodeB2", "nodeB3"); + clusterState = addDataNodeForAZone(clusterState, "zone_C", "nodeC1", "nodeC2", "nodeC3"); + return clusterState; + } + + private ClusterState addClusterManagerNodes(ClusterState clusterState) { + clusterState = addClusterManagerNodeForAZone(clusterState, "zone_A", "nodeMA"); + clusterState = addClusterManagerNodeForAZone(clusterState, "zone_B", "nodeMB"); + clusterState = addClusterManagerNodeForAZone(clusterState, "zone_C", "nodeMC"); + return clusterState; + } + + private ClusterState addDataNodeForAZone(ClusterState clusterState, String zone, String... nodeIds) { + DiscoveryNodes.Builder nodeBuilder = DiscoveryNodes.builder(clusterState.nodes()); + org.opensearch.common.collect.List.of(nodeIds) + .forEach( + nodeId -> nodeBuilder.add( + new DiscoveryNode( + nodeId, + buildNewFakeTransportAddress(), + Collections.singletonMap("zone", zone), + DATA_ROLE, + Version.CURRENT + ) + ) + ); + clusterState = ClusterState.builder(clusterState).nodes(nodeBuilder).build(); + return clusterState; + } + + private ClusterState addClusterManagerNodeForAZone(ClusterState clusterState, String zone, String... nodeIds) { + + DiscoveryNodes.Builder nodeBuilder = DiscoveryNodes.builder(clusterState.nodes()); + org.opensearch.common.collect.List.of(nodeIds) + .forEach( + nodeId -> nodeBuilder.add( + new DiscoveryNode( + nodeId, + buildNewFakeTransportAddress(), + Collections.singletonMap("zone", zone), + CLUSTER_MANAGER_ROLE, + Version.CURRENT + ) + ) + ); + clusterState = ClusterState.builder(clusterState).nodes(nodeBuilder).build(); + return clusterState; + } + + private ClusterState setLocalNode(ClusterState clusterState, String nodeId) { + DiscoveryNodes.Builder nodeBuilder = DiscoveryNodes.builder(clusterState.nodes()); + nodeBuilder.localNodeId(nodeId); + nodeBuilder.clusterManagerNodeId(nodeId); + clusterState = ClusterState.builder(clusterState).nodes(nodeBuilder).build(); + return clusterState; + } + + private ClusterState setWeightedRoutingWeights(ClusterState clusterState, Map weights) { + WeightedRouting weightedRouting = new WeightedRouting("zone", weights); + WeightedRoutingMetadata weightedRoutingMetadata = new WeightedRoutingMetadata(weightedRouting); + Metadata.Builder metadataBuilder = Metadata.builder(clusterState.metadata()); + metadataBuilder.putCustom(WeightedRoutingMetadata.TYPE, weightedRoutingMetadata); + clusterState = ClusterState.builder(clusterState).metadata(metadataBuilder).build(); + return clusterState; + } + + public void testGetWeightedRouting_WeightsNotSetInMetadata() { + + final ClusterGetWeightedRoutingRequestBuilder request = new ClusterGetWeightedRoutingRequestBuilder( + client, + ClusterGetWeightedRoutingAction.INSTANCE + ); + request.setAwarenessAttribute("zone"); + ClusterState state = clusterService.state(); + + ClusterGetWeightedRoutingResponse response = ActionTestUtils.executeBlocking(transportGetWeightedRoutingAction, request.request()); + assertEquals(response.getLocalNodeWeight(), null); + assertEquals(response.weights(), null); + } + + public void testGetWeightedRouting_WeightsSetInMetadata() { + ClusterGetWeightedRoutingRequestBuilder request = new ClusterGetWeightedRoutingRequestBuilder( + client, + ClusterGetWeightedRoutingAction.INSTANCE + ); + request.setAwarenessAttribute("zone"); + + ClusterState state = clusterService.state(); + state = setLocalNode(state, "nodeB1"); + Map weights = Map.of("zone_A", 1.0, "zone_B", 0.0, "zone_C", 1.0); + state = setWeightedRoutingWeights(state, weights); + ClusterState.Builder builder = ClusterState.builder(state); + ClusterServiceUtils.setState(clusterService, builder); + + ClusterGetWeightedRoutingResponse response = ActionTestUtils.executeBlocking(transportGetWeightedRoutingAction, request.request()); + assertEquals(weights, response.weights().weights()); + } + + public void testGetWeightedRoutingLocalWeight_WeightsSetInMetadata() { + + ClusterGetWeightedRoutingRequestBuilder request = new ClusterGetWeightedRoutingRequestBuilder( + client, + ClusterGetWeightedRoutingAction.INSTANCE + ); + + request.setRequestLocal(true); + request.setAwarenessAttribute("zone"); + + ClusterState state = clusterService.state(); + state = setLocalNode(state, "nodeB1"); + Map weights = Map.of("zone_A", 1.0, "zone_B", 0.0, "zone_C", 1.0); + state = setWeightedRoutingWeights(state, weights); + ClusterState.Builder builder = ClusterState.builder(state); + ClusterServiceUtils.setState(clusterService, builder); + + ClusterGetWeightedRoutingResponse response = ActionTestUtils.executeBlocking(transportGetWeightedRoutingAction, request.request()); + assertEquals("0.0", response.getLocalNodeWeight()); + } + + public void testGetWeightedRoutingLocalWeight_WeightsNotSetInMetadata() { + + ClusterGetWeightedRoutingRequestBuilder request = new ClusterGetWeightedRoutingRequestBuilder( + client, + ClusterGetWeightedRoutingAction.INSTANCE + ); + + request.setRequestLocal(true); + request.setAwarenessAttribute("zone"); + + ClusterState state = clusterService.state(); + state = setLocalNode(state, "nodeB1"); + ClusterState.Builder builder = ClusterState.builder(state); + ClusterServiceUtils.setState(clusterService, builder); + + ClusterGetWeightedRoutingResponse response = ActionTestUtils.executeBlocking(transportGetWeightedRoutingAction, request.request()); + assertEquals(null, response.getLocalNodeWeight()); + } + + @After + public void shutdown() { + clusterService.stop(); + threadPool.shutdown(); + } + +} diff --git a/server/src/test/java/org/opensearch/cluster/metadata/WeightedRoutingMetadataTests.java b/server/src/test/java/org/opensearch/cluster/metadata/WeightedRoutingMetadataTests.java new file mode 100644 index 0000000000000..a0a9d2bd9586b --- /dev/null +++ b/server/src/test/java/org/opensearch/cluster/metadata/WeightedRoutingMetadataTests.java @@ -0,0 +1,36 @@ +/* + * 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.opensearch.cluster.routing.WeightedRouting; +import org.opensearch.common.xcontent.XContentParser; +import org.opensearch.test.AbstractXContentTestCase; + +import java.io.IOException; +import java.util.Map; + +public class WeightedRoutingMetadataTests extends AbstractXContentTestCase { + @Override + protected WeightedRoutingMetadata createTestInstance() { + Map weights = Map.of("a", 1.0, "b", 1.0, "c", 0.0); + WeightedRouting weightedRouting = new WeightedRouting("zone", weights); + WeightedRoutingMetadata weightedRoutingMetadata = new WeightedRoutingMetadata(weightedRouting); + return weightedRoutingMetadata; + } + + @Override + protected WeightedRoutingMetadata doParseInstance(XContentParser parser) throws IOException { + return WeightedRoutingMetadata.fromXContent(parser); + } + + @Override + protected boolean supportsUnknownFields() { + return false; + } +} diff --git a/server/src/test/java/org/opensearch/cluster/routing/OperationRoutingTests.java b/server/src/test/java/org/opensearch/cluster/routing/OperationRoutingTests.java index 8bf2b1626292a..87cab4a006a63 100644 --- a/server/src/test/java/org/opensearch/cluster/routing/OperationRoutingTests.java +++ b/server/src/test/java/org/opensearch/cluster/routing/OperationRoutingTests.java @@ -36,6 +36,7 @@ import org.opensearch.cluster.ClusterState; import org.opensearch.cluster.metadata.IndexMetadata; import org.opensearch.cluster.metadata.Metadata; +import org.opensearch.cluster.metadata.WeightedRoutingMetadata; import org.opensearch.cluster.node.DiscoveryNode; import org.opensearch.cluster.node.DiscoveryNodeRole; import org.opensearch.cluster.routing.allocation.decider.AwarenessAllocationDecider; @@ -759,6 +760,232 @@ public void testAdaptiveReplicaSelectionWithZoneAwarenessIgnored() throws Except terminate(threadPool); } + private ClusterState clusterStateForWeightedRouting(String[] indexNames, int numShards, int numReplicas) { + DiscoveryNode[] allNodes = setUpNodesForWeightedRouting(); + ClusterState state = ClusterStateCreationUtils.state(allNodes[0], allNodes[6], allNodes); + + Map> discoveryNodeMap = new HashMap<>(); + List nodesZoneA = new ArrayList<>(); + nodesZoneA.add(allNodes[0]); + nodesZoneA.add(allNodes[1]); + + List nodesZoneB = new ArrayList<>(); + nodesZoneB.add(allNodes[2]); + nodesZoneB.add(allNodes[3]); + + List nodesZoneC = new ArrayList<>(); + nodesZoneC.add(allNodes[4]); + nodesZoneC.add(allNodes[5]); + discoveryNodeMap.put("a", nodesZoneA); + discoveryNodeMap.put("b", nodesZoneB); + discoveryNodeMap.put("c", nodesZoneC); + + // Updating cluster state with node, index and shard details + state = updateStatetoTestWeightedRouting(indexNames, numShards, numReplicas, state, discoveryNodeMap); + + return state; + + } + + private ClusterState setWeightedRoutingWeights(ClusterState clusterState, Map weights) { + WeightedRouting weightedRouting = new WeightedRouting("zone", weights); + WeightedRoutingMetadata weightedRoutingMetadata = new WeightedRoutingMetadata(weightedRouting); + Metadata.Builder metadataBuilder = Metadata.builder(clusterState.metadata()); + metadataBuilder.putCustom(WeightedRoutingMetadata.TYPE, weightedRoutingMetadata); + clusterState = ClusterState.builder(clusterState).metadata(metadataBuilder).build(); + return clusterState; + } + + public void testWeightedOperationRouting() throws Exception { + final int numIndices = 2; + final int numShards = 3; + final int numReplicas = 2; + // setting up indices + final String[] indexNames = new String[numIndices]; + for (int i = 0; i < numIndices; i++) { + indexNames[i] = "test" + i; + } + ClusterService clusterService = null; + TestThreadPool threadPool = null; + try { + ClusterState state = clusterStateForWeightedRouting(indexNames, numShards, numReplicas); + + Settings setting = Settings.builder().put("cluster.routing.allocation.awareness.attributes", "zone").build(); + + threadPool = new TestThreadPool("testThatOnlyNodesSupport"); + clusterService = ClusterServiceUtils.createClusterService(threadPool); + + OperationRouting opRouting = new OperationRouting( + setting, + new ClusterSettings(Settings.EMPTY, ClusterSettings.BUILT_IN_CLUSTER_SETTINGS) + ); + assertTrue(opRouting.ignoreAwarenessAttributes()); + Set selectedNodes = new HashSet<>(); + ResponseCollectorService collector = new ResponseCollectorService(clusterService); + Map outstandingRequests = new HashMap<>(); + + // Setting up weights for weighted round-robin in cluster state + Map weights = Map.of("a", 1.0, "b", 1.0, "c", 0.0); + state = setWeightedRoutingWeights(state, weights); + + ClusterState.Builder builder = ClusterState.builder(state); + ClusterServiceUtils.setState(clusterService, builder); + + // search shards call + GroupShardsIterator groupIterator = opRouting.searchShards( + state, + indexNames, + null, + null, + collector, + outstandingRequests + + ); + + for (ShardIterator it : groupIterator) { + List shardRoutings = Collections.singletonList(it.nextOrNull()); + for (ShardRouting shardRouting : shardRoutings) { + selectedNodes.add(shardRouting.currentNodeId()); + } + } + // tests no shards are assigned to nodes in zone c + for (String nodeID : selectedNodes) { + // No shards are assigned to nodes in zone c since its weight is 0 + assertFalse(nodeID.contains("c")); + } + + selectedNodes = new HashSet<>(); + setting = Settings.builder().put("cluster.routing.allocation.awareness.attributes", "zone").build(); + + // Updating weighted round robin weights in cluster state + weights = Map.of("a", 1.0, "b", 0.0, "c", 1.0); + state = setWeightedRoutingWeights(state, weights); + + builder = ClusterState.builder(state); + ClusterServiceUtils.setState(clusterService, builder); + + opRouting = new OperationRouting(setting, new ClusterSettings(Settings.EMPTY, ClusterSettings.BUILT_IN_CLUSTER_SETTINGS)); + + // search shards call + groupIterator = opRouting.searchShards(state, indexNames, null, null, collector, outstandingRequests); + + for (ShardIterator it : groupIterator) { + List shardRoutings = Collections.singletonList(it.nextOrNull()); + for (ShardRouting shardRouting : shardRoutings) { + selectedNodes.add(shardRouting.currentNodeId()); + } + } + // tests that no shards are assigned to zone with weight zero + for (String nodeID : selectedNodes) { + // No shards are assigned to nodes in zone b since its weight is 0 + assertFalse(nodeID.contains("b")); + } + } finally { + IOUtils.close(clusterService); + terminate(threadPool); + } + } + + public void testWeightedOperationRoutingWeightUndefinedForOneZone() throws Exception { + final int numIndices = 2; + final int numShards = 3; + final int numReplicas = 2; + // setting up indices + final String[] indexNames = new String[numIndices]; + for (int i = 0; i < numIndices; i++) { + indexNames[i] = "test" + i; + } + + ClusterService clusterService = null; + TestThreadPool threadPool = null; + try { + ClusterState state = clusterStateForWeightedRouting(indexNames, numShards, numReplicas); + + Settings setting = Settings.builder().put("cluster.routing.allocation.awareness.attributes", "zone").build(); + + threadPool = new TestThreadPool("testThatOnlyNodesSupport"); + clusterService = ClusterServiceUtils.createClusterService(threadPool); + + OperationRouting opRouting = new OperationRouting( + setting, + new ClusterSettings(Settings.EMPTY, ClusterSettings.BUILT_IN_CLUSTER_SETTINGS) + ); + assertTrue(opRouting.ignoreAwarenessAttributes()); + Set selectedNodes = new HashSet<>(); + ResponseCollectorService collector = new ResponseCollectorService(clusterService); + Map outstandingRequests = new HashMap<>(); + + // Setting up weights for weighted round-robin in cluster state, weight for nodes in zone b is not set + Map weights = Map.of("a", 1.0, "c", 0.0); + state = setWeightedRoutingWeights(state, weights); + ClusterServiceUtils.setState(clusterService, ClusterState.builder(state)); + + // search shards call + GroupShardsIterator groupIterator = opRouting.searchShards( + state, + indexNames, + null, + null, + collector, + outstandingRequests + + ); + + for (ShardIterator it : groupIterator) { + List shardRoutings = Collections.singletonList(it.nextOrNull()); + for (ShardRouting shardRouting : shardRoutings) { + selectedNodes.add(shardRouting.currentNodeId()); + } + } + boolean weighAwayNodesInUndefinedZone = true; + // tests no shards are assigned to nodes in zone c + // tests shards are assigned to nodes in zone b + for (String nodeID : selectedNodes) { + // shard from nodes in zone c is not selected since its weight is 0 + assertFalse(nodeID.contains("c")); + if (nodeID.contains("b")) { + weighAwayNodesInUndefinedZone = false; + } + } + assertFalse(weighAwayNodesInUndefinedZone); + + selectedNodes = new HashSet<>(); + setting = Settings.builder().put("cluster.routing.allocation.awareness.attributes", "zone").build(); + + // Updating weighted round robin weights in cluster state + weights = Map.of("a", 0.0, "b", 1.0); + + state = setWeightedRoutingWeights(state, weights); + ClusterServiceUtils.setState(clusterService, ClusterState.builder(state)); + + opRouting = new OperationRouting(setting, new ClusterSettings(Settings.EMPTY, ClusterSettings.BUILT_IN_CLUSTER_SETTINGS)); + + // search shards call + groupIterator = opRouting.searchShards(state, indexNames, null, null, collector, outstandingRequests); + + for (ShardIterator it : groupIterator) { + List shardRoutings = Collections.singletonList(it.nextOrNull()); + for (ShardRouting shardRouting : shardRoutings) { + selectedNodes.add(shardRouting.currentNodeId()); + } + } + // tests that no shards are assigned to zone with weight zero + // tests shards are assigned to nodes in zone c + weighAwayNodesInUndefinedZone = true; + for (String nodeID : selectedNodes) { + // shard from nodes in zone a is not selected since its weight is 0 + assertFalse(nodeID.contains("a")); + if (nodeID.contains("c")) { + weighAwayNodesInUndefinedZone = false; + } + } + assertFalse(weighAwayNodesInUndefinedZone); + } finally { + IOUtils.close(clusterService); + terminate(threadPool); + } + } + private DiscoveryNode[] setupNodes() { // Sets up two data nodes in zone-a and one data node in zone-b List zones = Arrays.asList("a", "a", "b"); @@ -785,6 +1012,32 @@ private DiscoveryNode[] setupNodes() { return allNodes; } + private DiscoveryNode[] setUpNodesForWeightedRouting() { + List zones = Arrays.asList("a", "a", "b", "b", "c", "c"); + DiscoveryNode[] allNodes = new DiscoveryNode[7]; + int i = 0; + for (String zone : zones) { + DiscoveryNode node = new DiscoveryNode( + "node_" + zone + "_" + i, + buildNewFakeTransportAddress(), + singletonMap("zone", zone), + Collections.singleton(DiscoveryNodeRole.DATA_ROLE), + Version.CURRENT + ); + allNodes[i++] = node; + } + + DiscoveryNode clusterManager = new DiscoveryNode( + "cluster-manager", + buildNewFakeTransportAddress(), + Collections.emptyMap(), + Collections.singleton(DiscoveryNodeRole.CLUSTER_MANAGER_ROLE), + Version.CURRENT + ); + allNodes[i] = clusterManager; + return allNodes; + } + public void testAllocationAwarenessDeprecation() { OperationRouting routing = new OperationRouting( Settings.builder() @@ -841,4 +1094,69 @@ private ClusterState updateStatetoTestARS( clusterState.routingTable(routingTableBuilder.build()); return clusterState.build(); } + + private ClusterState updateStatetoTestWeightedRouting( + String[] indices, + int numberOfShards, + int numberOfReplicas, + ClusterState state, + Map> discoveryNodeMap + ) { + RoutingTable.Builder routingTableBuilder = RoutingTable.builder(); + Metadata.Builder metadataBuilder = Metadata.builder(); + ClusterState.Builder clusterState = ClusterState.builder(state); + List nodesZoneA = discoveryNodeMap.get("a"); + List nodesZoneB = discoveryNodeMap.get("b"); + List nodesZoneC = discoveryNodeMap.get("c"); + for (String index : indices) { + IndexMetadata indexMetadata = IndexMetadata.builder(index) + .settings( + Settings.builder() + .put(SETTING_VERSION_CREATED, Version.CURRENT) + .put(SETTING_NUMBER_OF_SHARDS, numberOfShards) + .put(SETTING_NUMBER_OF_REPLICAS, numberOfReplicas) + .put(SETTING_CREATION_DATE, System.currentTimeMillis()) + ) + .build(); + metadataBuilder.put(indexMetadata, false).generateClusterUuidIfNeeded(); + IndexRoutingTable.Builder indexRoutingTableBuilder = IndexRoutingTable.builder(indexMetadata.getIndex()); + for (int i = 0; i < numberOfShards; i++) { + final ShardId shardId = new ShardId(index, "_na_", i); + IndexShardRoutingTable.Builder indexShardRoutingBuilder = new IndexShardRoutingTable.Builder(shardId); + // Assign all the primary shards on nodes in zone-a (node_a0 or node_a1) + indexShardRoutingBuilder.addShard( + TestShardRouting.newShardRouting( + index, + i, + nodesZoneA.get(randomInt(nodesZoneA.size() - 1)).getId(), + null, + true, + ShardRoutingState.STARTED + ) + ); + for (int replica = 0; replica < numberOfReplicas; replica++) { + // Assign all the replicas on nodes in zone-b (node_b2) + String nodeId = ""; + if (replica == 0) { + nodeId = nodesZoneB.get(randomInt(nodesZoneB.size() - 1)).getId(); + } else { + nodeId = nodesZoneC.get(randomInt(nodesZoneC.size() - 1)).getId(); + } + indexShardRoutingBuilder.addShard( + TestShardRouting.newShardRouting(index, i, nodeId, null, false, ShardRoutingState.STARTED) + ); + } + indexRoutingTableBuilder.addIndexShard(indexShardRoutingBuilder.build()); + } + routingTableBuilder.add(indexRoutingTableBuilder.build()); + } + // add weighted routing weights in metadata + Map weights = Map.of("a", 1.0, "b", 1.0, "c", 0.0); + WeightedRouting weightedRouting = new WeightedRouting("zone", weights); + WeightedRoutingMetadata weightedRoutingMetadata = new WeightedRoutingMetadata(weightedRouting); + metadataBuilder.putCustom(WeightedRoutingMetadata.TYPE, weightedRoutingMetadata); + clusterState.metadata(metadataBuilder); + clusterState.routingTable(routingTableBuilder.build()); + return clusterState.build(); + } } diff --git a/server/src/test/java/org/opensearch/cluster/routing/WeightedRoundRobinTests.java b/server/src/test/java/org/opensearch/cluster/routing/WeightedRoundRobinTests.java new file mode 100644 index 0000000000000..5f62d30486e86 --- /dev/null +++ b/server/src/test/java/org/opensearch/cluster/routing/WeightedRoundRobinTests.java @@ -0,0 +1,151 @@ +/* + * 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.routing; + +import org.opensearch.test.OpenSearchTestCase; + +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; + +public class WeightedRoundRobinTests extends OpenSearchTestCase { + + public void testWeightedRoundRobinOrder() { + // weights set as A:4, B:3, C:2 + List> entity = new ArrayList>(); + entity.add(new WeightedRoundRobin.Entity<>(4, "A")); + entity.add(new WeightedRoundRobin.Entity<>(3, "B")); + entity.add(new WeightedRoundRobin.Entity<>(2, "C")); + WeightedRoundRobin weightedRoundRobin = new WeightedRoundRobin(entity); + List> orderedEntities = weightedRoundRobin.orderEntities(); + List expectedOrdering = Arrays.asList("A", "A", "B", "A", "B", "C", "A", "B", "C"); + List actualOrdering = new ArrayList<>(); + for (WeightedRoundRobin.Entity en : orderedEntities) { + actualOrdering.add(en.getTarget()); + } + assertEquals(expectedOrdering, actualOrdering); + + // weights set as A:1, B:1, C:0 + entity = new ArrayList>(); + entity.add(new WeightedRoundRobin.Entity<>(1, "A")); + entity.add(new WeightedRoundRobin.Entity<>(1, "B")); + entity.add(new WeightedRoundRobin.Entity<>(0, "C")); + weightedRoundRobin = new WeightedRoundRobin(entity); + orderedEntities = weightedRoundRobin.orderEntities(); + expectedOrdering = Arrays.asList("A", "B"); + actualOrdering = new ArrayList<>(); + for (WeightedRoundRobin.Entity en : orderedEntities) { + actualOrdering.add(en.getTarget()); + } + assertEquals(expectedOrdering, actualOrdering); + + // weights set as A:0, B:0, C:0 + entity = new ArrayList>(); + entity.add(new WeightedRoundRobin.Entity<>(0, "A")); + entity.add(new WeightedRoundRobin.Entity<>(0, "B")); + entity.add(new WeightedRoundRobin.Entity<>(0, "C")); + weightedRoundRobin = new WeightedRoundRobin(entity); + orderedEntities = weightedRoundRobin.orderEntities(); + expectedOrdering = Arrays.asList(); + actualOrdering = new ArrayList<>(); + for (WeightedRoundRobin.Entity en : orderedEntities) { + actualOrdering.add(en.getTarget()); + } + assertEquals(expectedOrdering, actualOrdering); + + // weights set as A:-1, B:0, C:1 + entity = new ArrayList>(); + entity.add(new WeightedRoundRobin.Entity<>(-1, "A")); + entity.add(new WeightedRoundRobin.Entity<>(0, "B")); + entity.add(new WeightedRoundRobin.Entity<>(1, "C")); + weightedRoundRobin = new WeightedRoundRobin(entity); + orderedEntities = weightedRoundRobin.orderEntities(); + expectedOrdering = Arrays.asList("C"); + actualOrdering = new ArrayList<>(); + for (WeightedRoundRobin.Entity en : orderedEntities) { + actualOrdering.add(en.getTarget()); + } + assertEquals(expectedOrdering, actualOrdering); + + // weights set as A:-1, B:3, C:0, D:10 + entity = new ArrayList>(); + entity.add(new WeightedRoundRobin.Entity<>(-1, "A")); + entity.add(new WeightedRoundRobin.Entity<>(3, "B")); + entity.add(new WeightedRoundRobin.Entity<>(0, "C")); + entity.add(new WeightedRoundRobin.Entity<>(10, "D")); + weightedRoundRobin = new WeightedRoundRobin(entity); + orderedEntities = weightedRoundRobin.orderEntities(); + expectedOrdering = Arrays.asList("B", "D", "B", "D", "B", "D", "D", "D", "D", "D", "D", "D", "D"); + actualOrdering = new ArrayList<>(); + for (WeightedRoundRobin.Entity en : orderedEntities) { + actualOrdering.add(en.getTarget()); + } + assertEquals(expectedOrdering, actualOrdering); + + // weights set as A:-1, B:3, C:0, D:10000 + entity = new ArrayList>(); + entity.add(new WeightedRoundRobin.Entity<>(-1, "A")); + entity.add(new WeightedRoundRobin.Entity<>(3, "B")); + entity.add(new WeightedRoundRobin.Entity<>(0, "C")); + entity.add(new WeightedRoundRobin.Entity<>(10000, "D")); + weightedRoundRobin = new WeightedRoundRobin(entity); + orderedEntities = weightedRoundRobin.orderEntities(); + assertEquals(10003, orderedEntities.size()); + // Count of D's + int countD = 0; + // Count of B's + int countB = 0; + for (WeightedRoundRobin.Entity en : orderedEntities) { + if (en.getTarget().equals("D")) { + countD++; + } else if (en.getTarget().equals("B")) { + countB++; + } + } + assertEquals(3, countB); + assertEquals(10000, countD); + + // weights set C:0 + entity = new ArrayList>(); + entity.add(new WeightedRoundRobin.Entity<>(0, "C")); + weightedRoundRobin = new WeightedRoundRobin(entity); + orderedEntities = weightedRoundRobin.orderEntities(); + expectedOrdering = Arrays.asList(); + actualOrdering = new ArrayList<>(); + for (WeightedRoundRobin.Entity en : orderedEntities) { + actualOrdering.add(en.getTarget()); + } + assertEquals(expectedOrdering, actualOrdering); + + // weights set C:1 + entity = new ArrayList>(); + entity.add(new WeightedRoundRobin.Entity<>(1, "C")); + weightedRoundRobin = new WeightedRoundRobin(entity); + orderedEntities = weightedRoundRobin.orderEntities(); + expectedOrdering = Arrays.asList("C"); + actualOrdering = new ArrayList<>(); + for (WeightedRoundRobin.Entity en : orderedEntities) { + actualOrdering.add(en.getTarget()); + } + assertEquals(expectedOrdering, actualOrdering); + + // weights set C:2 + entity = new ArrayList>(); + entity.add(new WeightedRoundRobin.Entity<>(2, "C")); + weightedRoundRobin = new WeightedRoundRobin(entity); + orderedEntities = weightedRoundRobin.orderEntities(); + expectedOrdering = Arrays.asList("C", "C"); + actualOrdering = new ArrayList<>(); + for (WeightedRoundRobin.Entity en : orderedEntities) { + actualOrdering.add(en.getTarget()); + } + assertEquals(expectedOrdering, actualOrdering); + } + +} diff --git a/server/src/test/java/org/opensearch/cluster/routing/WeightedRoutingServiceTests.java b/server/src/test/java/org/opensearch/cluster/routing/WeightedRoutingServiceTests.java new file mode 100644 index 0000000000000..fc5d46ef84c79 --- /dev/null +++ b/server/src/test/java/org/opensearch/cluster/routing/WeightedRoutingServiceTests.java @@ -0,0 +1,279 @@ +/* + * 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.routing; + +import org.junit.After; +import org.junit.Before; +import org.opensearch.Version; +import org.opensearch.action.ActionListener; +import org.opensearch.action.admin.cluster.shards.routing.weighted.delete.ClusterDeleteWeightedRoutingRequest; +import org.opensearch.action.ActionRequestValidationException; +import org.opensearch.action.admin.cluster.shards.routing.weighted.delete.ClusterDeleteWeightedRoutingResponse; +import org.opensearch.action.admin.cluster.shards.routing.weighted.put.ClusterAddWeightedRoutingAction; +import org.opensearch.action.admin.cluster.shards.routing.weighted.put.ClusterPutWeightedRoutingRequestBuilder; +import org.opensearch.client.node.NodeClient; +import org.opensearch.cluster.ClusterName; +import org.opensearch.cluster.ClusterState; +import org.opensearch.cluster.ack.ClusterStateUpdateResponse; +import org.opensearch.cluster.metadata.Metadata; +import org.opensearch.cluster.metadata.WeightedRoutingMetadata; +import org.opensearch.cluster.node.DiscoveryNode; +import org.opensearch.cluster.node.DiscoveryNodeRole; +import org.opensearch.cluster.node.DiscoveryNodes; +import org.opensearch.cluster.routing.allocation.decider.AwarenessAllocationDecider; +import org.opensearch.cluster.service.ClusterService; +import org.opensearch.common.settings.ClusterSettings; +import org.opensearch.common.settings.Settings; +import org.opensearch.test.ClusterServiceUtils; +import org.opensearch.test.OpenSearchTestCase; +import org.opensearch.test.transport.MockTransport; +import org.opensearch.threadpool.TestThreadPool; +import org.opensearch.threadpool.ThreadPool; +import org.opensearch.transport.TransportService; + +import java.util.Collections; +import java.util.HashSet; +import java.util.Map; +import java.util.Set; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; + +public class WeightedRoutingServiceTests extends OpenSearchTestCase { + private ThreadPool threadPool; + private ClusterService clusterService; + private TransportService transportService; + private WeightedRoutingService weightedRoutingService; + private ClusterSettings clusterSettings; + NodeClient client; + + final private static Set CLUSTER_MANAGER_ROLE = Collections.unmodifiableSet( + new HashSet<>(Collections.singletonList(DiscoveryNodeRole.CLUSTER_MANAGER_ROLE)) + ); + + final private static Set DATA_ROLE = Collections.unmodifiableSet( + new HashSet<>(Collections.singletonList(DiscoveryNodeRole.DATA_ROLE)) + ); + + @Override + public void setUp() throws Exception { + super.setUp(); + threadPool = new TestThreadPool("test", Settings.EMPTY); + clusterService = ClusterServiceUtils.createClusterService(threadPool); + } + + @Before + public void setUpService() { + ClusterState clusterState = ClusterState.builder(new ClusterName("test")).build(); + clusterState = addClusterManagerNodes(clusterState); + clusterState = addDataNodes(clusterState); + clusterState = setLocalNode(clusterState, "nodeA1"); + + ClusterState.Builder builder = ClusterState.builder(clusterState); + ClusterServiceUtils.setState(clusterService, builder); + + final MockTransport transport = new MockTransport(); + transportService = transport.createTransportService( + Settings.EMPTY, + threadPool, + TransportService.NOOP_TRANSPORT_INTERCEPTOR, + boundTransportAddress -> clusterService.state().nodes().get("nodes1"), + null, + Collections.emptySet() + + ); + + Settings.Builder settingsBuilder = Settings.builder() + .put(AwarenessAllocationDecider.CLUSTER_ROUTING_ALLOCATION_AWARENESS_ATTRIBUTE_SETTING.getKey(), "zone"); + + clusterSettings = new ClusterSettings(settingsBuilder.build(), ClusterSettings.BUILT_IN_CLUSTER_SETTINGS); + transportService.start(); + transportService.acceptIncomingRequests(); + + this.weightedRoutingService = new WeightedRoutingService(clusterService, threadPool, settingsBuilder.build(), clusterSettings); + client = new NodeClient(Settings.EMPTY, threadPool); + } + + @After + public void shutdown() { + clusterService.stop(); + threadPool.shutdown(); + } + + private ClusterState addDataNodes(ClusterState clusterState) { + clusterState = addDataNodeForAZone(clusterState, "zone_A", "nodeA1", "nodeA2", "nodeA3"); + clusterState = addDataNodeForAZone(clusterState, "zone_B", "nodeB1", "nodeB2", "nodeB3"); + clusterState = addDataNodeForAZone(clusterState, "zone_C", "nodeC1", "nodeC2", "nodeC3"); + return clusterState; + } + + private ClusterState addClusterManagerNodes(ClusterState clusterState) { + clusterState = addClusterManagerNodeForAZone(clusterState, "zone_A", "nodeMA"); + clusterState = addClusterManagerNodeForAZone(clusterState, "zone_B", "nodeMB"); + clusterState = addClusterManagerNodeForAZone(clusterState, "zone_C", "nodeMC"); + return clusterState; + } + + private ClusterState addDataNodeForAZone(ClusterState clusterState, String zone, String... nodeIds) { + DiscoveryNodes.Builder nodeBuilder = DiscoveryNodes.builder(clusterState.nodes()); + org.opensearch.common.collect.List.of(nodeIds) + .forEach( + nodeId -> nodeBuilder.add( + new DiscoveryNode( + nodeId, + buildNewFakeTransportAddress(), + Collections.singletonMap("zone", zone), + DATA_ROLE, + Version.CURRENT + ) + ) + ); + clusterState = ClusterState.builder(clusterState).nodes(nodeBuilder).build(); + return clusterState; + } + + private ClusterState addClusterManagerNodeForAZone(ClusterState clusterState, String zone, String... nodeIds) { + + DiscoveryNodes.Builder nodeBuilder = DiscoveryNodes.builder(clusterState.nodes()); + org.opensearch.common.collect.List.of(nodeIds) + .forEach( + nodeId -> nodeBuilder.add( + new DiscoveryNode( + nodeId, + buildNewFakeTransportAddress(), + Collections.singletonMap("zone", zone), + CLUSTER_MANAGER_ROLE, + Version.CURRENT + ) + ) + ); + clusterState = ClusterState.builder(clusterState).nodes(nodeBuilder).build(); + return clusterState; + } + + private ClusterState setLocalNode(ClusterState clusterState, String nodeId) { + DiscoveryNodes.Builder nodeBuilder = DiscoveryNodes.builder(clusterState.nodes()); + nodeBuilder.localNodeId(nodeId); + nodeBuilder.clusterManagerNodeId(nodeId); + clusterState = ClusterState.builder(clusterState).nodes(nodeBuilder).build(); + return clusterState; + } + + private ClusterState setWeightedRoutingWeights(ClusterState clusterState, Map weights) { + WeightedRouting weightedRouting = new WeightedRouting("zone", weights); + WeightedRoutingMetadata weightedRoutingMetadata = new WeightedRoutingMetadata(weightedRouting); + Metadata.Builder metadataBuilder = Metadata.builder(clusterState.metadata()); + metadataBuilder.putCustom(WeightedRoutingMetadata.TYPE, weightedRoutingMetadata); + clusterState = ClusterState.builder(clusterState).metadata(metadataBuilder).build(); + return clusterState; + } + + public void testRegisterWeightedRoutingMetadataWithChangedWeights() throws InterruptedException { + Map weights = Map.of("zone_A", 1.0, "zone_B", 1.0, "zone_C", 1.0); + ClusterState state = clusterService.state(); + state = setWeightedRoutingWeights(state, weights); + ClusterState.Builder builder = ClusterState.builder(state); + ClusterServiceUtils.setState(clusterService, builder); + + ClusterPutWeightedRoutingRequestBuilder request = new ClusterPutWeightedRoutingRequestBuilder( + client, + ClusterAddWeightedRoutingAction.INSTANCE + ); + WeightedRouting updatedWeightedRouting = new WeightedRouting("zone", Map.of("zone_A", 1.0, "zone_B", 0.0, "zone_C", 0.0)); + request.setWeightedRouting(updatedWeightedRouting); + final CountDownLatch countDownLatch = new CountDownLatch(1); + ActionListener listener = new ActionListener() { + @Override + public void onResponse(ClusterStateUpdateResponse clusterStateUpdateResponse) { + assertTrue(clusterStateUpdateResponse.isAcknowledged()); + assertEquals(updatedWeightedRouting, clusterService.state().metadata().weightedRoutingMetadata().getWeightedRouting()); + countDownLatch.countDown(); + } + + @Override + public void onFailure(Exception e) { + fail("request should not fail"); + } + }; + weightedRoutingService.registerWeightedRoutingMetadata(request.request(), listener); + assertTrue(countDownLatch.await(30, TimeUnit.SECONDS)); + } + + public void testRegisterWeightedRoutingMetadataWithSameWeights() throws InterruptedException { + Map weights = Map.of("zone_A", 1.0, "zone_B", 1.0, "zone_C", 1.0); + ClusterState state = clusterService.state(); + state = setWeightedRoutingWeights(state, weights); + ClusterState.Builder builder = ClusterState.builder(state); + ClusterServiceUtils.setState(clusterService, builder); + + ClusterPutWeightedRoutingRequestBuilder request = new ClusterPutWeightedRoutingRequestBuilder( + client, + ClusterAddWeightedRoutingAction.INSTANCE + ); + WeightedRouting updatedWeightedRouting = new WeightedRouting("zone", weights); + request.setWeightedRouting(updatedWeightedRouting); + final CountDownLatch countDownLatch = new CountDownLatch(1); + ActionListener listener = new ActionListener() { + @Override + public void onResponse(ClusterStateUpdateResponse clusterStateUpdateResponse) { + assertTrue(clusterStateUpdateResponse.isAcknowledged()); + assertEquals(updatedWeightedRouting, clusterService.state().metadata().weightedRoutingMetadata().getWeightedRouting()); + countDownLatch.countDown(); + } + + @Override + public void onFailure(Exception e) { + fail("request should not fail"); + } + }; + weightedRoutingService.registerWeightedRoutingMetadata(request.request(), listener); + assertTrue(countDownLatch.await(30, TimeUnit.SECONDS)); + } + + public void testDeleteWeightedRoutingMetadata() throws InterruptedException { + Map weights = Map.of("zone_A", 1.0, "zone_B", 1.0, "zone_C", 1.0); + ClusterState state = clusterService.state(); + state = setWeightedRoutingWeights(state, weights); + ClusterState.Builder builder = ClusterState.builder(state); + ClusterServiceUtils.setState(clusterService, builder); + + ClusterDeleteWeightedRoutingRequest clusterDeleteWeightedRoutingRequest = new ClusterDeleteWeightedRoutingRequest(); + final CountDownLatch countDownLatch = new CountDownLatch(1); + ActionListener listener = new ActionListener() { + @Override + public void onResponse(ClusterDeleteWeightedRoutingResponse clusterDeleteWeightedRoutingResponse) { + assertTrue(clusterDeleteWeightedRoutingResponse.isAcknowledged()); + assertNull(clusterService.state().metadata().weightedRoutingMetadata()); + countDownLatch.countDown(); + } + + @Override + public void onFailure(Exception e) { + fail("on failure shouldn't have been called"); + } + }; + weightedRoutingService.deleteWeightedRoutingMetadata(clusterDeleteWeightedRoutingRequest, listener); + assertTrue(countDownLatch.await(30, TimeUnit.SECONDS)); + } + + public void testVerifyAwarenessAttribute_InvalidAttributeName() { + assertThrows( + "invalid awareness attribute %s requested for updating weighted routing", + ActionRequestValidationException.class, + () -> weightedRoutingService.verifyAwarenessAttribute("zone2") + ); + } + + public void testVerifyAwarenessAttribute_ValidAttributeName() { + try { + weightedRoutingService.verifyAwarenessAttribute("zone"); + } catch (Exception e) { + fail("verify awareness attribute should not fail"); + } + } +} diff --git a/server/src/test/java/org/opensearch/cluster/structure/RoutingIteratorTests.java b/server/src/test/java/org/opensearch/cluster/structure/RoutingIteratorTests.java index 68ad47fa1bbc9..8f5aa1b764551 100644 --- a/server/src/test/java/org/opensearch/cluster/structure/RoutingIteratorTests.java +++ b/server/src/test/java/org/opensearch/cluster/structure/RoutingIteratorTests.java @@ -40,6 +40,7 @@ import org.opensearch.cluster.metadata.Metadata; import org.opensearch.cluster.node.DiscoveryNodes; import org.opensearch.cluster.routing.GroupShardsIterator; +import org.opensearch.cluster.routing.IndexShardRoutingTable; import org.opensearch.cluster.routing.OperationRouting; import org.opensearch.cluster.routing.PlainShardIterator; import org.opensearch.cluster.routing.RotationShardShuffler; @@ -48,11 +49,15 @@ import org.opensearch.cluster.routing.ShardRouting; import org.opensearch.cluster.routing.ShardShuffler; import org.opensearch.cluster.routing.ShardsIterator; +import org.opensearch.cluster.routing.WeightedRouting; import org.opensearch.cluster.routing.allocation.AllocationService; import org.opensearch.cluster.routing.allocation.decider.ClusterRebalanceAllocationDecider; +import org.opensearch.cluster.service.ClusterService; import org.opensearch.common.settings.ClusterSettings; import org.opensearch.common.settings.Settings; import org.opensearch.index.shard.ShardId; +import org.opensearch.test.ClusterServiceUtils; +import org.opensearch.threadpool.TestThreadPool; import java.util.Arrays; import java.util.Collections; @@ -497,4 +502,209 @@ public void testReplicaShardPreferenceIters() throws Exception { } } + public void testWeightedRoutingWithDifferentWeights() { + TestThreadPool threadPool = null; + try { + Settings.Builder settings = Settings.builder() + .put("cluster.routing.allocation.node_concurrent_recoveries", 10) + .put("cluster.routing.allocation.awareness.attributes", "zone"); + AllocationService strategy = createAllocationService(settings.build()); + + Metadata metadata = Metadata.builder() + .put(IndexMetadata.builder("test").settings(settings(Version.CURRENT)).numberOfShards(1).numberOfReplicas(2)) + .build(); + + RoutingTable routingTable = RoutingTable.builder().addAsNew(metadata.index("test")).build(); + + ClusterState clusterState = ClusterState.builder(ClusterName.CLUSTER_NAME_SETTING.getDefault(Settings.EMPTY)) + .metadata(metadata) + .routingTable(routingTable) + .build(); + + threadPool = new TestThreadPool("testThatOnlyNodesSupport"); + ClusterService clusterService = ClusterServiceUtils.createClusterService(threadPool); + + Map node1Attributes = new HashMap<>(); + node1Attributes.put("zone", "zone1"); + Map node2Attributes = new HashMap<>(); + node2Attributes.put("zone", "zone2"); + Map node3Attributes = new HashMap<>(); + node3Attributes.put("zone", "zone3"); + clusterState = ClusterState.builder(clusterState) + .nodes( + DiscoveryNodes.builder() + .add(newNode("node1", unmodifiableMap(node1Attributes))) + .add(newNode("node2", unmodifiableMap(node2Attributes))) + .add(newNode("node3", unmodifiableMap(node3Attributes))) + .localNodeId("node1") + ) + .build(); + clusterState = strategy.reroute(clusterState, "reroute"); + + clusterState = startInitializingShardsAndReroute(strategy, clusterState); + clusterState = startInitializingShardsAndReroute(strategy, clusterState); + + Map weights = Map.of("zone1", 1.0, "zone2", 1.0, "zone3", 0.0); + WeightedRouting weightedRouting = new WeightedRouting("zone", weights); + + ShardIterator shardIterator = clusterState.routingTable() + .index("test") + .shard(0) + .activeInitializingShardsWeightedIt(weightedRouting, clusterState.nodes(), 1); + + assertEquals(2, shardIterator.size()); + ShardRouting shardRouting; + while (shardIterator.remaining() > 0) { + shardRouting = shardIterator.nextOrNull(); + assertNotNull(shardRouting); + assertFalse(Arrays.asList("node3").contains(shardRouting.currentNodeId())); + } + + weights = Map.of("zone1", 1.0, "zone2", 1.0, "zone3", 1.0); + weightedRouting = new WeightedRouting("zone", weights); + shardIterator = clusterState.routingTable() + .index("test") + .shard(0) + .activeInitializingShardsWeightedIt(weightedRouting, clusterState.nodes(), 1); + assertEquals(3, shardIterator.size()); + + weights = Map.of("zone1", -1.0, "zone2", 0.0, "zone3", 1.0); + weightedRouting = new WeightedRouting("zone", weights); + shardIterator = clusterState.routingTable() + .index("test") + .shard(0) + .activeInitializingShardsWeightedIt(weightedRouting, clusterState.nodes(), 1); + assertEquals(1, shardIterator.size()); + while (shardIterator.remaining() > 0) { + shardRouting = shardIterator.nextOrNull(); + assertNotNull(shardRouting); + assertFalse(Arrays.asList("node2", "node1").contains(shardRouting.currentNodeId())); + } + + weights = Map.of("zone1", 0.0, "zone2", 0.0, "zone3", 0.0); + weightedRouting = new WeightedRouting("zone", weights); + shardIterator = clusterState.routingTable() + .index("test") + .shard(0) + .activeInitializingShardsWeightedIt(weightedRouting, clusterState.nodes(), 1); + assertEquals(0, shardIterator.size()); + } finally { + terminate(threadPool); + } + } + + public void testWeightedRoutingInMemoryStore() { + TestThreadPool threadPool = null; + try { + Settings.Builder settings = Settings.builder() + .put("cluster.routing.allocation.node_concurrent_recoveries", 10) + .put("cluster.routing.allocation.awareness.attributes", "zone"); + AllocationService strategy = createAllocationService(settings.build()); + + Metadata metadata = Metadata.builder() + .put(IndexMetadata.builder("test").settings(settings(Version.CURRENT)).numberOfShards(1).numberOfReplicas(2)) + .build(); + + RoutingTable routingTable = RoutingTable.builder().addAsNew(metadata.index("test")).build(); + + ClusterState clusterState = ClusterState.builder(ClusterName.CLUSTER_NAME_SETTING.getDefault(Settings.EMPTY)) + .metadata(metadata) + .routingTable(routingTable) + .build(); + + threadPool = new TestThreadPool("testThatOnlyNodesSupport"); + ClusterService clusterService = ClusterServiceUtils.createClusterService(threadPool); + + Map node1Attributes = new HashMap<>(); + node1Attributes.put("zone", "zone1"); + Map node2Attributes = new HashMap<>(); + node2Attributes.put("zone", "zone2"); + Map node3Attributes = new HashMap<>(); + node3Attributes.put("zone", "zone3"); + clusterState = ClusterState.builder(clusterState) + .nodes( + DiscoveryNodes.builder() + .add(newNode("node1", unmodifiableMap(node1Attributes))) + .add(newNode("node2", unmodifiableMap(node2Attributes))) + .add(newNode("node3", unmodifiableMap(node3Attributes))) + .localNodeId("node1") + ) + .build(); + clusterState = strategy.reroute(clusterState, "reroute"); + + clusterState = startInitializingShardsAndReroute(strategy, clusterState); + clusterState = startInitializingShardsAndReroute(strategy, clusterState); + + Map weights = Map.of("zone1", 1.0, "zone2", 1.0, "zone3", 0.0); + WeightedRouting weightedRouting = new WeightedRouting("zone", weights); + + IndexShardRoutingTable indexShardRoutingTable = clusterState.routingTable().index("test").shard(0); + + assertNull( + indexShardRoutingTable.getActiveShardsByWeight().get(new IndexShardRoutingTable.WeightedRoutingKey(weightedRouting)) + ); + ShardIterator shardIterator = clusterState.routingTable() + .index("test") + .shard(0) + .activeInitializingShardsWeightedIt(weightedRouting, clusterState.nodes(), 1); + assertEquals(2, shardIterator.size()); + ShardRouting shardRouting; + while (shardIterator.remaining() > 0) { + shardRouting = shardIterator.nextOrNull(); + assertNotNull(shardRouting); + assertFalse(Arrays.asList("node3").contains(shardRouting.currentNodeId())); + } + + // Make iterator call with same WeightedRouting instance + assertNotNull( + indexShardRoutingTable.getActiveShardsByWeight().get(new IndexShardRoutingTable.WeightedRoutingKey(weightedRouting)) + ); + shardIterator = clusterState.routingTable() + .index("test") + .shard(0) + .activeInitializingShardsWeightedIt(weightedRouting, clusterState.nodes(), 1); + assertEquals(2, shardIterator.size()); + while (shardIterator.remaining() > 0) { + shardRouting = shardIterator.nextOrNull(); + assertNotNull(shardRouting); + assertFalse(Arrays.asList("node3").contains(shardRouting.currentNodeId())); + } + + // Make iterator call with new instance of WeightedRouting but same weights + Map weights1 = Map.of("zone1", 1.0, "zone2", 1.0, "zone3", 0.0); + weightedRouting = new WeightedRouting("zone", weights1); + assertNotNull( + indexShardRoutingTable.getActiveShardsByWeight().get(new IndexShardRoutingTable.WeightedRoutingKey(weightedRouting)) + ); + shardIterator = clusterState.routingTable() + .index("test") + .shard(0) + .activeInitializingShardsWeightedIt(weightedRouting, clusterState.nodes(), 1); + assertEquals(2, shardIterator.size()); + while (shardIterator.remaining() > 0) { + shardRouting = shardIterator.nextOrNull(); + assertNotNull(shardRouting); + assertFalse(Arrays.asList("node3").contains(shardRouting.currentNodeId())); + } + + // Make iterator call with different weights + Map weights2 = Map.of("zone1", 1.0, "zone2", 0.0, "zone3", 1.0); + weightedRouting = new WeightedRouting("zone", weights2); + assertNull( + indexShardRoutingTable.getActiveShardsByWeight().get(new IndexShardRoutingTable.WeightedRoutingKey(weightedRouting)) + ); + shardIterator = clusterState.routingTable() + .index("test") + .shard(0) + .activeInitializingShardsWeightedIt(weightedRouting, clusterState.nodes(), 1); + assertEquals(2, shardIterator.size()); + while (shardIterator.remaining() > 0) { + shardRouting = shardIterator.nextOrNull(); + assertNotNull(shardRouting); + assertFalse(Arrays.asList("node2").contains(shardRouting.currentNodeId())); + } + } finally { + terminate(threadPool); + } + } } diff --git a/server/src/test/java/org/opensearch/rest/action/admin/cluster/RestClusterAddWeightedRoutingActionTests.java b/server/src/test/java/org/opensearch/rest/action/admin/cluster/RestClusterAddWeightedRoutingActionTests.java new file mode 100644 index 0000000000000..a4cd6224217b7 --- /dev/null +++ b/server/src/test/java/org/opensearch/rest/action/admin/cluster/RestClusterAddWeightedRoutingActionTests.java @@ -0,0 +1,76 @@ +/* + * 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.rest.action.admin.cluster; + +import com.fasterxml.jackson.core.JsonParseException; +import org.junit.Before; +import org.opensearch.OpenSearchParseException; +import org.opensearch.action.admin.cluster.shards.routing.weighted.put.ClusterPutWeightedRoutingRequest; +import org.opensearch.common.bytes.BytesArray; +import org.opensearch.common.xcontent.XContentType; +import org.opensearch.rest.RestRequest; +import org.opensearch.test.rest.FakeRestRequest; +import org.opensearch.test.rest.RestActionTestCase; + +import java.io.IOException; +import java.util.HashMap; +import java.util.Map; + +import static java.util.Collections.singletonMap; + +public class RestClusterAddWeightedRoutingActionTests extends RestActionTestCase { + private RestClusterPutWeightedRoutingAction action; + + @Before + public void setupAction() { + action = new RestClusterPutWeightedRoutingAction(); + controller().registerHandler(action); + } + + public void testCreateRequest_SupportedRequestBody() throws IOException { + String req = "{\"us-east-1c\" : \"1\", \"us-east-1d\":\"1.0\", \"us-east-1a\":\"0.0\"}"; + RestRequest restRequest = buildRestRequest(req); + ClusterPutWeightedRoutingRequest clusterPutWeightedRoutingRequest = RestClusterPutWeightedRoutingAction.createRequest(restRequest); + assertEquals("zone", clusterPutWeightedRoutingRequest.getWeightedRouting().attributeName()); + assertNotNull(clusterPutWeightedRoutingRequest.getWeightedRouting().weights()); + assertEquals("1.0", clusterPutWeightedRoutingRequest.getWeightedRouting().weights().get("us-east-1c").toString()); + assertEquals("1.0", clusterPutWeightedRoutingRequest.getWeightedRouting().weights().get("us-east-1d").toString()); + assertEquals("0.0", clusterPutWeightedRoutingRequest.getWeightedRouting().weights().get("us-east-1a").toString()); + } + + public void testCreateRequest_UnsupportedRequestBody() throws IOException { + Map params = new HashMap<>(); + String req = "[\"us-east-1c\" : \"1\", \"us-east-1d\":\"1\", \"us-east-1a\":\"0\"]"; + RestRequest restRequest = buildRestRequest(req); + assertThrows(OpenSearchParseException.class, () -> RestClusterPutWeightedRoutingAction.createRequest(restRequest)); + } + + public void testCreateRequest_MalformedRequestBody() throws IOException { + Map params = new HashMap<>(); + + String req = "{\"us-east-1c\" : \1\", \"us-east-1d\":\"1\", \"us-east-1a\":\"0\"}"; + RestRequest restRequest = buildRestRequest(req); + assertThrows(JsonParseException.class, () -> RestClusterPutWeightedRoutingAction.createRequest(restRequest)); + } + + public void testCreateRequest_EmptyRequestBody() throws IOException { + String req = "{}"; + RestRequest restRequest = buildRestRequest(req); + assertThrows(OpenSearchParseException.class, () -> RestClusterPutWeightedRoutingAction.createRequest(restRequest)); + } + + private RestRequest buildRestRequest(String content) { + return new FakeRestRequest.Builder(xContentRegistry()).withMethod(RestRequest.Method.PUT) + .withPath("/_cluster/routing/awareness/zone/weights") + .withParams(singletonMap("attribute", "zone")) + .withContent(new BytesArray(content), XContentType.JSON) + .build(); + } + +}