Reset desired balance (#94525)

This introduces an endpoint to reset the desired balance.
It could be used if computed balance diverged from the actual one a lot 
to start a new computation from the current state.

docs/reference/cluster.asciidoc

@@ -119,3 +119,5 @@ include::cluster/get-desired-nodes.asciidoc[]
 include::cluster/delete-desired-nodes.asciidoc[]
 
 include::cluster/get-desired-balance.asciidoc[]
+
+include::cluster/delete-desired-balance.asciidoc[]

docs/reference/cluster/delete-desired-balance.asciidoc

@@ -0,0 +1,18 @@
+[[delete-desired-balance]]
+=== Delete/reset desired balance API
+++++
+<titleabbrev>Delete/reset desired balance</titleabbrev>
+++++
+
+NOTE: {cloud-only}
+
+Exposes the desired balance and basic metrics.
+
+[[delete-desired-balance-request]]
+==== {api-request-title}
+
+[source,console]
+--------------------------------------------------
+DELETE /_internal/desired_balance
+--------------------------------------------------
+// TEST[skip:Can't reliably test desired balance]

qa/smoke-test-multinode/src/yamlRestTest/resources/rest-api-spec/test/smoke_test_multinode/30_desired_balance.yml

@@ -69,7 +69,7 @@ setup:
 
   - skip:
       version: " - 8.6.99"
-      reason: "Field added in in 8.7.0"
+      reason: "cluster_balance_stats added in in 8.7.0"
 
   - do:
       cluster.state: {}
@@ -117,7 +117,7 @@ setup:
 
   - skip:
       version: " - 8.7.99"
-      reason: "tier preference added in in 8.8.0"
+      reason: "tier_preference added in in 8.8.0"
 
   - do:
       indices.create:
@@ -142,7 +142,7 @@ setup:
 
   - skip:
       version: " - 8.7.99"
-      reason: "Field added in in 8.8.0"
+      reason: "cluster_info added in in 8.8.0"
 
   - do:
       _internal.get_desired_balance: { }

rest-api-spec/src/main/resources/rest-api-spec/api/_internal.delete_desired_balance.json

@@ -0,0 +1,23 @@
+{
+  "_internal.delete_desired_balance":{
+    "documentation":{
+      "url":"https://www.elastic.co/guide/en/elasticsearch/reference/current/delete-desired-balance.html",
+      "description": "This API is a diagnostics API and the output should not be relied upon for building applications."
+    },
+    "stability":"experimental",
+    "visibility":"private",
+    "headers":{
+      "accept": [ "application/json"]
+    },
+    "url":{
+      "paths":[
+        {
+          "path":"/_internal/desired_balance",
+          "methods":[
+            "DELETE"
+          ]
+        }
+      ]
+    }
+  }
+}

rest-api-spec/src/yamlRestTest/resources/rest-api-spec/test/cluster.desired_balance/10_basic.yml

@@ -5,7 +5,7 @@ setup:
       reason: "API added in in 8.6.0"
 
 ---
-"Test empty desired balance":
+"Test get empty desired balance":
 
   - do:
       _internal.get_desired_balance: { }
@@ -19,12 +19,51 @@ setup:
   - gte: { stats.reconciliation_time_in_millis: 0 }
   - match: { routing_table: {} }
 
+---
+"Test get desired balance for a single shard":
+  - do:
+      indices.create:
+        index: test
+        body:
+          settings:
+            number_of_shards: 1
+            number_of_replicas: 0
+
+  - do:
+      cluster.health:
+        index: test
+        wait_for_status: green
+
+  - do:
+      _internal.get_desired_balance: { }
+
+  - gte: { stats.computation_submitted: 0 }
+  - gte: { stats.computation_executed: 0 }
+  - gte: { stats.computation_converged: 0 }
+  - gte: { stats.computation_iterations: 0 }
+  - gte: { stats.computation_converged_index: 0 }
+  - gte: { stats.computation_time_in_millis: 0 }
+  - gte: { stats.reconciliation_time_in_millis: 0 }
+
+  - match: { routing_table.test.0.current.0.state: 'STARTED' }
+  - match: { routing_table.test.0.current.0.shard_id: 0 }
+  - match: { routing_table.test.0.current.0.index: test }
+  - is_true: 'routing_table.test.0.current.0.node_is_desired'
+  - is_false: 'routing_table.test.0.current.0.relocating_node'
+  - is_false: 'routing_table.test.0.current.0.relocating_node_is_desired'
+  - is_false: 'routing_table.test.0.current.0.forecast_write_load'
+  - is_false: 'routing_table.test.0.current.0.forecast_shard_size_in_bytes'
+  - match: { routing_table.test.0.desired.total: 1 }
+  - gte: { routing_table.test.0.desired.unassigned: 0 }
+  - gte: { routing_table.test.0.desired.ignored: 0 }
+  - is_true: 'routing_table.test.0.desired.node_ids'
+
 ---
 "Test cluster_balance_stats":
 
   - skip:
       version: " - 8.6.99"
-      reason: "Field added in in 8.7.0"
+      reason: "cluster_balance_stats added in in 8.7.0"
 
   - do:
       cluster.state: {}
@@ -72,7 +111,7 @@ setup:
 
   - skip:
       version: " - 8.7.99"
-      reason: "Field added in in 8.8.0"
+      reason: "cluster_info added in in 8.8.0"
 
   - do:
       _internal.get_desired_balance: { }
@@ -84,7 +123,7 @@ setup:
 
   - skip:
       version: " - 8.7.99"
-      reason: "Node ID added in in 8.8.0"
+      reason: "node_id and roles added in in 8.8.0"
 
   - do:
       cluster.state: {}
@@ -98,51 +137,12 @@ setup:
   - is_true: 'cluster_balance_stats.nodes.$node_name.node_id'
   - is_true: 'cluster_balance_stats.nodes.$node_name.roles'
 
----
-"Test get desired balance for single shard":
-  - do:
-      indices.create:
-        index: test
-        body:
-          settings:
-            number_of_shards: 1
-            number_of_replicas: 0
-
-  - do:
-      cluster.health:
-        index: test
-        wait_for_status: green
-
-  - do:
-      _internal.get_desired_balance: { }
-
-  - gte: { stats.computation_submitted: 0 }
-  - gte: { stats.computation_executed: 0 }
-  - gte: { stats.computation_converged: 0 }
-  - gte: { stats.computation_iterations: 0 }
-  - gte: { stats.computation_converged_index: 0 }
-  - gte: { stats.computation_time_in_millis: 0 }
-  - gte: { stats.reconciliation_time_in_millis: 0 }
-
-  - match: { routing_table.test.0.current.0.state: 'STARTED' }
-  - match: { routing_table.test.0.current.0.shard_id: 0 }
-  - match: { routing_table.test.0.current.0.index: test }
-  - is_true: 'routing_table.test.0.current.0.node_is_desired'
-  - is_false: 'routing_table.test.0.current.0.relocating_node'
-  - is_false: 'routing_table.test.0.current.0.relocating_node_is_desired'
-  - is_false: 'routing_table.test.0.current.0.forecast_write_load'
-  - is_false: 'routing_table.test.0.current.0.forecast_shard_size_in_bytes'
-  - match: { routing_table.test.0.desired.total: 1 }
-  - gte: { routing_table.test.0.desired.unassigned: 0 }
-  - gte: { routing_table.test.0.desired.ignored: 0 }
-  - is_true: 'routing_table.test.0.desired.node_ids'
-
 ---
 "Test tier_preference":
 
   - skip:
       version: " - 8.7.99"
-      reason: "tier preference added in in 8.8.0"
+      reason: "tier_preference added in in 8.8.0"
 
   - do:
       indices.create:
@@ -173,3 +173,13 @@ setup:
       _internal.get_desired_balance: { }
 
   - gte: { stats.computed_shard_movements: 0 }
+
+---
+"Test reset desired balance":
+
+  - skip:
+      version: " - 8.7.99"
+      reason: "reset API added in in 8.8.0"
+
+  - do:
+      _internal.delete_desired_balance: { }

server/src/main/java/org/elasticsearch/action/ActionModule.java

@@ -11,8 +11,10 @@ package org.elasticsearch.action;
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
 import org.elasticsearch.action.admin.cluster.allocation.ClusterAllocationExplainAction;
+import org.elasticsearch.action.admin.cluster.allocation.DeleteDesiredBalanceAction;
 import org.elasticsearch.action.admin.cluster.allocation.GetDesiredBalanceAction;
 import org.elasticsearch.action.admin.cluster.allocation.TransportClusterAllocationExplainAction;
+import org.elasticsearch.action.admin.cluster.allocation.TransportDeleteDesiredBalanceAction;
 import org.elasticsearch.action.admin.cluster.allocation.TransportGetDesiredBalanceAction;
 import org.elasticsearch.action.admin.cluster.configuration.AddVotingConfigExclusionsAction;
 import org.elasticsearch.action.admin.cluster.configuration.ClearVotingConfigExclusionsAction;
@@ -310,6 +312,7 @@ import org.elasticsearch.rest.action.admin.cluster.RestClusterStateAction;
 import org.elasticsearch.rest.action.admin.cluster.RestClusterStatsAction;
 import org.elasticsearch.rest.action.admin.cluster.RestClusterUpdateSettingsAction;
 import org.elasticsearch.rest.action.admin.cluster.RestCreateSnapshotAction;
+import org.elasticsearch.rest.action.admin.cluster.RestDeleteDesiredBalanceAction;
 import org.elasticsearch.rest.action.admin.cluster.RestDeleteDesiredNodesAction;
 import org.elasticsearch.rest.action.admin.cluster.RestDeleteRepositoryAction;
 import org.elasticsearch.rest.action.admin.cluster.RestDeleteSnapshotAction;
@@ -621,6 +624,7 @@ public class ActionModule extends AbstractModule {
         actions.register(ClearVotingConfigExclusionsAction.INSTANCE, TransportClearVotingConfigExclusionsAction.class);
         actions.register(ClusterAllocationExplainAction.INSTANCE, TransportClusterAllocationExplainAction.class);
         actions.register(GetDesiredBalanceAction.INSTANCE, TransportGetDesiredBalanceAction.class);
+        actions.register(DeleteDesiredBalanceAction.INSTANCE, TransportDeleteDesiredBalanceAction.class);
         actions.register(ClusterStatsAction.INSTANCE, TransportClusterStatsAction.class);
         actions.register(ClusterStateAction.INSTANCE, TransportClusterStateAction.class);
         actions.register(ClusterHealthAction.INSTANCE, TransportClusterHealthAction.class);
@@ -799,6 +803,7 @@ public class ActionModule extends AbstractModule {
         registerHandler.accept(new RestNodesHotThreadsAction());
         registerHandler.accept(new RestClusterAllocationExplainAction());
         registerHandler.accept(new RestGetDesiredBalanceAction());
+        registerHandler.accept(new RestDeleteDesiredBalanceAction());
         registerHandler.accept(new RestClusterStatsAction());
         registerHandler.accept(new RestClusterStateAction(settingsFilter, threadPool));
         registerHandler.accept(new RestClusterHealthAction());

server/src/main/java/org/elasticsearch/action/admin/cluster/allocation/DeleteDesiredBalanceAction.java

@@ -0,0 +1,22 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0 and the Server Side Public License, v 1; you may not use this file except
+ * in compliance with, at your election, the Elastic License 2.0 or the Server
+ * Side Public License, v 1.
+ */
+
+package org.elasticsearch.action.admin.cluster.allocation;
+
+import org.elasticsearch.action.ActionResponse;
+import org.elasticsearch.action.ActionType;
+
+public class DeleteDesiredBalanceAction extends ActionType<ActionResponse.Empty> {
+
+    public static final DeleteDesiredBalanceAction INSTANCE = new DeleteDesiredBalanceAction();
+    public static final String NAME = "cluster:admin/desired_balance/reset";
+
+    DeleteDesiredBalanceAction() {
+        super(NAME, in -> ActionResponse.Empty.INSTANCE);
+    }
+}

server/src/main/java/org/elasticsearch/action/admin/cluster/allocation/TransportDeleteDesiredBalanceAction.java

@@ -0,0 +1,129 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0 and the Server Side Public License, v 1; you may not use this file except
+ * in compliance with, at your election, the Elastic License 2.0 or the Server
+ * Side Public License, v 1.
+ */
+
+package org.elasticsearch.action.admin.cluster.allocation;
+
+import org.elasticsearch.ResourceNotFoundException;
+import org.elasticsearch.action.ActionListener;
+import org.elasticsearch.action.ActionResponse;
+import org.elasticsearch.action.support.ActionFilters;
+import org.elasticsearch.action.support.master.TransportMasterNodeAction;
+import org.elasticsearch.cluster.ClusterState;
+import org.elasticsearch.cluster.ClusterStateTaskExecutor;
+import org.elasticsearch.cluster.ClusterStateTaskListener;
+import org.elasticsearch.cluster.block.ClusterBlockException;
+import org.elasticsearch.cluster.block.ClusterBlockLevel;
+import org.elasticsearch.cluster.metadata.IndexNameExpressionResolver;
+import org.elasticsearch.cluster.routing.allocation.AllocationService;
+import org.elasticsearch.cluster.routing.allocation.allocator.AllocationActionMultiListener;
+import org.elasticsearch.cluster.routing.allocation.allocator.DesiredBalanceShardsAllocator;
+import org.elasticsearch.cluster.routing.allocation.allocator.ShardsAllocator;
+import org.elasticsearch.cluster.service.ClusterService;
+import org.elasticsearch.cluster.service.MasterServiceTaskQueue;
+import org.elasticsearch.common.Priority;
+import org.elasticsearch.common.inject.Inject;
+import org.elasticsearch.core.Nullable;
+import org.elasticsearch.tasks.Task;
+import org.elasticsearch.threadpool.ThreadPool;
+import org.elasticsearch.transport.TransportService;
+
+public class TransportDeleteDesiredBalanceAction extends TransportMasterNodeAction<DesiredBalanceRequest, ActionResponse.Empty> {
+
+    @Nullable
+    private final MasterServiceTaskQueue<ResetDesiredBalanceTask> resetDesiredBalanceTaskQueue;
+
+    @Inject
+    public TransportDeleteDesiredBalanceAction(
+        TransportService transportService,
+        ClusterService clusterService,
+        ThreadPool threadPool,
+        ActionFilters actionFilters,
+        IndexNameExpressionResolver indexNameExpressionResolver,
+        AllocationService allocationService,
+        ShardsAllocator shardsAllocator
+    ) {
+        super(
+            DeleteDesiredBalanceAction.NAME,
+            transportService,
+            clusterService,
+            threadPool,
+            actionFilters,
+            DesiredBalanceRequest::new,
+            indexNameExpressionResolver,
+            in -> ActionResponse.Empty.INSTANCE,
+            ThreadPool.Names.MANAGEMENT
+        );
+
+        this.resetDesiredBalanceTaskQueue = shardsAllocator instanceof DesiredBalanceShardsAllocator allocator
+            ? clusterService.createTaskQueue(
+                "reset-desired-balance",
+                Priority.NORMAL,
+                new ResetDesiredBalanceClusterExecutor(threadPool, allocationService, allocator)
+            )
+            : null;
+    }
+
+    public record ResetDesiredBalanceTask(ActionListener<Void> listener) implements ClusterStateTaskListener {
+
+        @Override
+        public void onFailure(Exception e) {
+            listener.onFailure(e);
+        }
+    }
+
+    private static final class ResetDesiredBalanceClusterExecutor implements ClusterStateTaskExecutor<ResetDesiredBalanceTask> {
+
+        private final ThreadPool threadPool;
+        private final AllocationService allocationService;
+        private final DesiredBalanceShardsAllocator desiredBalanceShardsAllocator;
+
+        ResetDesiredBalanceClusterExecutor(
+            ThreadPool threadPool,
+            AllocationService allocationService,
+            DesiredBalanceShardsAllocator desiredBalanceShardsAllocator
+        ) {
+            this.threadPool = threadPool;
+            this.allocationService = allocationService;
+            this.desiredBalanceShardsAllocator = desiredBalanceShardsAllocator;
+        }
+
+        @Override
+        public ClusterState execute(BatchExecutionContext<ResetDesiredBalanceTask> batchExecutionContext) throws InterruptedException {
+            var listener = new AllocationActionMultiListener<Void>(threadPool.getThreadContext());
+            var state = batchExecutionContext.initialState();
+            desiredBalanceShardsAllocator.resetDesiredBalance();
+            for (var taskContext : batchExecutionContext.taskContexts()) {
+                taskContext.success(() -> listener.delay(taskContext.getTask().listener()).onResponse(null));
+            }
+            return allocationService.reroute(state, "reset-desired-balance", listener.reroute());
+        }
+    }
+
+    @Override
+    protected void masterOperation(
+        Task task,
+        DesiredBalanceRequest request,
+        ClusterState state,
+        ActionListener<ActionResponse.Empty> listener
+    ) throws Exception {
+        if (resetDesiredBalanceTaskQueue == null) {
+            listener.onFailure(new ResourceNotFoundException("Desired balance allocator is not in use, no desired balance found"));
+            return;
+        }
+        resetDesiredBalanceTaskQueue.submitTask(
+            "reset-desired-balance",
+            new ResetDesiredBalanceTask(listener.map(ignored -> ActionResponse.Empty.INSTANCE)),
+            null
+        );
+    }
+
+    @Override
+    protected ClusterBlockException checkBlock(DesiredBalanceRequest request, ClusterState state) {
+        return state.blocks().globalBlockedException(ClusterBlockLevel.METADATA_READ);
+    }
+}

server/src/main/java/org/elasticsearch/cluster/routing/allocation/allocator/DesiredBalanceShardsAllocator.java

@@ -36,6 +36,7 @@ import org.elasticsearch.threadpool.ThreadPool;
 import java.util.ArrayList;
 import java.util.Comparator;
 import java.util.List;
+import java.util.Map;
 import java.util.Set;
 import java.util.concurrent.ConcurrentLinkedQueue;
 import java.util.concurrent.atomic.AtomicLong;
@@ -62,6 +63,7 @@ public class DesiredBalanceShardsAllocator implements ShardsAllocator {
     private final MasterServiceTaskQueue<ReconcileDesiredBalanceTask> masterServiceTaskQueue;
     private final NodeAllocationOrdering allocationOrdering = new NodeAllocationOrdering();
     private volatile DesiredBalance currentDesiredBalance = DesiredBalance.INITIAL;
+    private volatile boolean resetCurrentDesiredBalance = false;
 
     // stats
     protected final CounterMetric computationsSubmitted = new CounterMetric();
@@ -115,7 +117,7 @@ public class DesiredBalanceShardsAllocator implements ShardsAllocator {
                     cumulativeComputationTime,
                     () -> setCurrentDesiredBalance(
                         desiredBalanceComputer.compute(
-                            currentDesiredBalance,
+                            getInitialDesiredBalance(),
                             desiredBalanceInput,
                             pendingDesiredBalanceMoves,
                             this::isFresh
@@ -132,6 +134,16 @@ public class DesiredBalanceShardsAllocator implements ShardsAllocator {
                 }
             }
 
+            private DesiredBalance getInitialDesiredBalance() {
+                if (resetCurrentDesiredBalance) {
+                    logger.info("Resetting current desired balance");
+                    resetCurrentDesiredBalance = false;
+                    return new DesiredBalance(currentDesiredBalance.lastConvergedIndex(), Map.of());
+                } else {
+                    return currentDesiredBalance;
+                }
+            }
+
             @Override
             public String toString() {
                 return "DesiredBalanceShardsAllocator#updateDesiredBalanceAndReroute";
@@ -229,6 +241,10 @@ public class DesiredBalanceShardsAllocator implements ShardsAllocator {
         return currentDesiredBalance;
     }
 
+    public void resetDesiredBalance() {
+        resetCurrentDesiredBalance = true;
+    }
+
     public DesiredBalanceStats getStats() {
         return new DesiredBalanceStats(
             currentDesiredBalance.lastConvergedIndex(),

server/src/main/java/org/elasticsearch/rest/action/admin/cluster/RestDeleteDesiredBalanceAction.java

@@ -0,0 +1,41 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0 and the Server Side Public License, v 1; you may not use this file except
+ * in compliance with, at your election, the Elastic License 2.0 or the Server
+ * Side Public License, v 1.
+ */
+
+package org.elasticsearch.rest.action.admin.cluster;
+
+import org.elasticsearch.action.admin.cluster.allocation.DeleteDesiredBalanceAction;
+import org.elasticsearch.action.admin.cluster.allocation.DesiredBalanceRequest;
+import org.elasticsearch.client.internal.node.NodeClient;
+import org.elasticsearch.rest.BaseRestHandler;
+import org.elasticsearch.rest.RestRequest;
+import org.elasticsearch.rest.action.RestToXContentListener;
+
+import java.io.IOException;
+import java.util.List;
+
+public class RestDeleteDesiredBalanceAction extends BaseRestHandler {
+
+    @Override
+    public String getName() {
+        return "delete_desired_balance";
+    }
+
+    @Override
+    public List<Route> routes() {
+        return List.of(new Route(RestRequest.Method.DELETE, "_internal/desired_balance"));
+    }
+
+    @Override
+    protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient client) throws IOException {
+        return channel -> client.execute(
+            DeleteDesiredBalanceAction.INSTANCE,
+            new DesiredBalanceRequest(),
+            new RestToXContentListener<>(channel)
+        );
+    }
+}

server/src/test/java/org/elasticsearch/action/admin/cluster/allocation/TransportDeleteDesiredBalanceActionTests.java

@@ -0,0 +1,169 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0 and the Server Side Public License, v 1; you may not use this file except
+ * in compliance with, at your election, the Elastic License 2.0 or the Server
+ * Side Public License, v 1.
+ */
+
+package org.elasticsearch.action.admin.cluster.allocation;
+
+import org.elasticsearch.ResourceNotFoundException;
+import org.elasticsearch.Version;
+import org.elasticsearch.action.ActionResponse;
+import org.elasticsearch.action.support.ActionFilters;
+import org.elasticsearch.action.support.PlainActionFuture;
+import org.elasticsearch.cluster.ClusterName;
+import org.elasticsearch.cluster.ClusterState;
+import org.elasticsearch.cluster.ESAllocationTestCase;
+import org.elasticsearch.cluster.EmptyClusterInfoService;
+import org.elasticsearch.cluster.TestShardRoutingRoleStrategies;
+import org.elasticsearch.cluster.metadata.IndexMetadata;
+import org.elasticsearch.cluster.metadata.IndexNameExpressionResolver;
+import org.elasticsearch.cluster.metadata.Metadata;
+import org.elasticsearch.cluster.node.DiscoveryNodes;
+import org.elasticsearch.cluster.routing.RoutingTable;
+import org.elasticsearch.cluster.routing.allocation.AllocationService;
+import org.elasticsearch.cluster.routing.allocation.allocator.BalancedShardsAllocator;
+import org.elasticsearch.cluster.routing.allocation.allocator.DesiredBalance;
+import org.elasticsearch.cluster.routing.allocation.allocator.DesiredBalanceComputer;
+import org.elasticsearch.cluster.routing.allocation.allocator.DesiredBalanceInput;
+import org.elasticsearch.cluster.routing.allocation.allocator.DesiredBalanceShardsAllocator;
+import org.elasticsearch.cluster.routing.allocation.allocator.ShardsAllocator;
+import org.elasticsearch.cluster.routing.allocation.command.MoveAllocationCommand;
+import org.elasticsearch.cluster.service.ClusterService;
+import org.elasticsearch.common.UUIDs;
+import org.elasticsearch.common.settings.ClusterSettings;
+import org.elasticsearch.common.settings.Settings;
+import org.elasticsearch.index.shard.ShardId;
+import org.elasticsearch.tasks.Task;
+import org.elasticsearch.test.ClusterServiceUtils;
+import org.elasticsearch.test.gateway.TestGatewayAllocator;
+import org.elasticsearch.threadpool.TestThreadPool;
+import org.elasticsearch.threadpool.ThreadPool;
+import org.elasticsearch.transport.TransportService;
+
+import java.util.List;
+import java.util.Map;
+import java.util.Queue;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicReference;
+import java.util.function.Predicate;
+
+import static org.elasticsearch.cluster.metadata.IndexMetadata.SETTING_INDEX_VERSION_CREATED;
+import static org.elasticsearch.cluster.metadata.IndexMetadata.SETTING_NUMBER_OF_REPLICAS;
+import static org.elasticsearch.cluster.metadata.IndexMetadata.SETTING_NUMBER_OF_SHARDS;
+import static org.hamcrest.Matchers.anEmptyMap;
+import static org.hamcrest.Matchers.equalTo;
+import static org.hamcrest.Matchers.greaterThan;
+import static org.hamcrest.Matchers.not;
+import static org.hamcrest.Matchers.notNullValue;
+import static org.mockito.Mockito.mock;
+
+public class TransportDeleteDesiredBalanceActionTests extends ESAllocationTestCase {
+
+    public void testReturnsErrorIfAllocatorIsNotDesiredBalanced() throws Exception {
+
+        var listener = new PlainActionFuture<ActionResponse.Empty>();
+
+        new TransportDeleteDesiredBalanceAction(
+            mock(TransportService.class),
+            mock(ClusterService.class),
+            mock(ThreadPool.class),
+            mock(ActionFilters.class),
+            mock(IndexNameExpressionResolver.class),
+            mock(AllocationService.class),
+            mock(ShardsAllocator.class)
+        ).masterOperation(mock(Task.class), new DesiredBalanceRequest(), ClusterState.EMPTY_STATE, listener);
+
+        var exception = expectThrows(ResourceNotFoundException.class, listener::actionGet);
+        assertThat(exception.getMessage(), equalTo("Desired balance allocator is not in use, no desired balance found"));
+    }
+
+    public void testDeleteDesiredBalance() throws Exception {
+
+        var threadPool = new TestThreadPool(getTestName());
+
+        var shardId = new ShardId("test-index", UUIDs.randomBase64UUID(), 0);
+        var index = createIndex(shardId.getIndexName());
+        var clusterState = ClusterState.builder(ClusterName.DEFAULT)
+            .nodes(DiscoveryNodes.builder().add(newNode("master")).localNodeId("master").masterNodeId("master").build())
+            .metadata(Metadata.builder().put(index, false).build())
+            .routingTable(RoutingTable.builder(TestShardRoutingRoleStrategies.DEFAULT_ROLE_ONLY).addAsNew(index).build())
+            .build();
+
+        var clusterService = ClusterServiceUtils.createClusterService(clusterState, threadPool);
+
+        var settings = Settings.EMPTY;
+        var clusterSettings = ClusterSettings.createBuiltInClusterSettings(settings);
+
+        var delegate = new BalancedShardsAllocator();
+        var computer = new DesiredBalanceComputer(clusterSettings, threadPool, delegate) {
+
+            final AtomicReference<DesiredBalance> lastComputationInput = new AtomicReference<>();
+
+            @Override
+            public DesiredBalance compute(
+                DesiredBalance previousDesiredBalance,
+                DesiredBalanceInput desiredBalanceInput,
+                Queue<List<MoveAllocationCommand>> pendingDesiredBalanceMoves,
+                Predicate<DesiredBalanceInput> isFresh
+            ) {
+                lastComputationInput.set(previousDesiredBalance);
+                return super.compute(previousDesiredBalance, desiredBalanceInput, pendingDesiredBalanceMoves, isFresh);
+            }
+        };
+        var allocator = new DesiredBalanceShardsAllocator(delegate, threadPool, clusterService, computer, (state, action) -> state);
+        var allocationService = new MockAllocationService(
+            randomAllocationDeciders(settings, clusterSettings),
+            new TestGatewayAllocator(),
+            allocator,
+            EmptyClusterInfoService.INSTANCE,
+            SNAPSHOT_INFO_SERVICE_WITH_NO_SHARD_SIZES
+        );
+
+        PlainActionFuture.<Void, RuntimeException>get(
+            f -> allocationService.reroute(clusterState, "inital-allocate", f),
+            10,
+            TimeUnit.SECONDS
+        );
+
+        var balanceBeforeReset = allocator.getDesiredBalance();
+        assertThat(balanceBeforeReset.lastConvergedIndex(), greaterThan(DesiredBalance.INITIAL.lastConvergedIndex()));
+        assertThat(balanceBeforeReset.assignments(), not(anEmptyMap()));
+
+        var listener = new PlainActionFuture<ActionResponse.Empty>();
+
+        var action = new TransportDeleteDesiredBalanceAction(
+            mock(TransportService.class),
+            clusterService,
+            threadPool,
+            mock(ActionFilters.class),
+            mock(IndexNameExpressionResolver.class),
+            allocationService,
+            allocator
+        );
+
+        action.masterOperation(mock(Task.class), new DesiredBalanceRequest(), clusterState, listener);
+
+        try {
+            assertThat(listener.get(), notNullValue());
+            // resetting desired balance should trigger new computation with empty assignments
+            assertThat(computer.lastComputationInput.get(), equalTo(new DesiredBalance(balanceBeforeReset.lastConvergedIndex(), Map.of())));
+        } finally {
+            clusterService.close();
+            terminate(threadPool);
+        }
+    }
+
+    private static IndexMetadata createIndex(String name) {
+        return IndexMetadata.builder(name)
+            .settings(
+                Settings.builder()
+                    .put(SETTING_NUMBER_OF_SHARDS, 1)
+                    .put(SETTING_NUMBER_OF_REPLICAS, 0)
+                    .put(SETTING_INDEX_VERSION_CREATED.getKey(), Version.CURRENT)
+            )
+            .build();
+    }
+}

server/src/test/java/org/elasticsearch/cluster/routing/allocation/allocator/DesiredBalanceShardsAllocatorTests.java

@@ -11,16 +11,16 @@ package org.elasticsearch.cluster.routing.allocation.allocator;
 import org.apache.lucene.util.SetOnce;
 import org.elasticsearch.Version;
 import org.elasticsearch.action.ActionListener;
+import org.elasticsearch.action.support.PlainActionFuture;
 import org.elasticsearch.cluster.ClusterInfo;
 import org.elasticsearch.cluster.ClusterName;
 import org.elasticsearch.cluster.ClusterState;
 import org.elasticsearch.cluster.ClusterStateUpdateTask;
+import org.elasticsearch.cluster.ESAllocationTestCase;
 import org.elasticsearch.cluster.TestShardRoutingRoleStrategies;
 import org.elasticsearch.cluster.block.ClusterBlocks;
 import org.elasticsearch.cluster.metadata.IndexMetadata;
 import org.elasticsearch.cluster.metadata.Metadata;
-import org.elasticsearch.cluster.node.DiscoveryNode;
-import org.elasticsearch.cluster.node.DiscoveryNodeRole;
 import org.elasticsearch.cluster.node.DiscoveryNodes;
 import org.elasticsearch.cluster.routing.RoutingTable;
 import org.elasticsearch.cluster.routing.ShardRouting;
@@ -29,28 +29,30 @@ import org.elasticsearch.cluster.routing.allocation.AllocationService;
 import org.elasticsearch.cluster.routing.allocation.ExistingShardsAllocator;
 import org.elasticsearch.cluster.routing.allocation.RoutingAllocation;
 import org.elasticsearch.cluster.routing.allocation.ShardAllocationDecision;
+import org.elasticsearch.cluster.routing.allocation.allocator.DesiredBalanceShardsAllocator.DesiredBalanceReconcilerAction;
 import org.elasticsearch.cluster.routing.allocation.command.MoveAllocationCommand;
 import org.elasticsearch.cluster.routing.allocation.decider.AllocationDeciders;
 import org.elasticsearch.cluster.service.ClusterApplierService;
 import org.elasticsearch.cluster.service.ClusterService;
 import org.elasticsearch.cluster.service.FakeThreadPoolMasterService;
+import org.elasticsearch.common.UUIDs;
 import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.common.util.concurrent.DeterministicTaskQueue;
 import org.elasticsearch.common.util.concurrent.PrioritizedEsThreadPoolExecutor;
 import org.elasticsearch.gateway.GatewayAllocator;
+import org.elasticsearch.index.shard.ShardId;
 import org.elasticsearch.snapshots.SnapshotShardSizeInfo;
 import org.elasticsearch.test.ClusterServiceUtils;
-import org.elasticsearch.test.ESTestCase;
 import org.elasticsearch.threadpool.TestThreadPool;
 
 import java.util.List;
 import java.util.Map;
 import java.util.Queue;
-import java.util.Set;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.BiConsumer;
 import java.util.function.Consumer;
 import java.util.function.Predicate;
@@ -62,7 +64,7 @@ import static org.elasticsearch.common.settings.ClusterSettings.createBuiltInClu
 import static org.hamcrest.Matchers.equalTo;
 import static org.hamcrest.Matchers.hasItem;
 
-public class DesiredBalanceShardsAllocatorTests extends ESTestCase {
+public class DesiredBalanceShardsAllocatorTests extends ESAllocationTestCase {
 
     private static final String LOCAL_NODE_ID = "node-1";
     private static final String OTHER_NODE_ID = "node-2";
@@ -107,8 +109,8 @@ public class DesiredBalanceShardsAllocatorTests extends ESTestCase {
         var deterministicTaskQueue = new DeterministicTaskQueue();
         var threadPool = deterministicTaskQueue.getThreadPool();
 
-        var localNode = createDiscoveryNode(LOCAL_NODE_ID);
-        var otherNode = createDiscoveryNode(OTHER_NODE_ID);
+        var localNode = newNode(LOCAL_NODE_ID);
+        var otherNode = newNode(OTHER_NODE_ID);
         var initialState = ClusterState.builder(new ClusterName(ClusterServiceUtils.class.getSimpleName()))
             .nodes(DiscoveryNodes.builder().add(localNode).add(otherNode).localNodeId(localNode.getId()).masterNodeId(localNode.getId()))
             .blocks(ClusterBlocks.EMPTY_CLUSTER_BLOCK)
@@ -135,7 +137,7 @@ public class DesiredBalanceShardsAllocatorTests extends ESTestCase {
         clusterService.start();
 
         var allocationServiceRef = new SetOnce<AllocationService>();
-        var reconcileAction = new DesiredBalanceShardsAllocator.DesiredBalanceReconcilerAction() {
+        var reconcileAction = new DesiredBalanceReconcilerAction() {
             @Override
             public ClusterState apply(ClusterState clusterState, Consumer<RoutingAllocation> routingAllocationAction) {
                 return allocationServiceRef.get().executeWithRoutingAllocation(clusterState, "reconcile", routingAllocationAction);
@@ -200,7 +202,7 @@ public class DesiredBalanceShardsAllocatorTests extends ESTestCase {
         var listenersCalled = new CountDownLatch(2);
         var clusterStateUpdatesExecuted = new CountDownLatch(2);
 
-        var discoveryNode = createDiscoveryNode("node-0");
+        var discoveryNode = newNode("node-0");
         var initialState = ClusterState.builder(ClusterName.DEFAULT)
             .nodes(DiscoveryNodes.builder().add(discoveryNode).localNodeId(discoveryNode.getId()).masterNodeId(discoveryNode.getId()))
             .build();
@@ -208,7 +210,7 @@ public class DesiredBalanceShardsAllocatorTests extends ESTestCase {
         var threadPool = new TestThreadPool(getTestName());
         var clusterService = ClusterServiceUtils.createClusterService(initialState, threadPool);
         var allocationServiceRef = new SetOnce<AllocationService>();
-        var reconcileAction = new DesiredBalanceShardsAllocator.DesiredBalanceReconcilerAction() {
+        var reconcileAction = new DesiredBalanceReconcilerAction() {
             @Override
             public ClusterState apply(ClusterState clusterState, Consumer<RoutingAllocation> routingAllocationAction) {
                 reconciliations.incrementAndGet();
@@ -301,8 +303,8 @@ public class DesiredBalanceShardsAllocatorTests extends ESTestCase {
         var newMasterElected = new CountDownLatch(1);
         var clusterStateUpdatesExecuted = new CountDownLatch(1);
 
-        var node1 = createDiscoveryNode(LOCAL_NODE_ID);
-        var node2 = createDiscoveryNode(OTHER_NODE_ID);
+        var node1 = newNode(LOCAL_NODE_ID);
+        var node2 = newNode(OTHER_NODE_ID);
         var initial = ClusterState.builder(ClusterName.DEFAULT)
             .nodes(DiscoveryNodes.builder().add(node1).add(node2).localNodeId(node1.getId()).masterNodeId(node1.getId()))
             .build();
@@ -310,7 +312,7 @@ public class DesiredBalanceShardsAllocatorTests extends ESTestCase {
         var threadPool = new TestThreadPool(getTestName());
         var clusterService = ClusterServiceUtils.createClusterService(initial, threadPool);
         var allocationServiceRef = new SetOnce<AllocationService>();
-        var reconcileAction = new DesiredBalanceShardsAllocator.DesiredBalanceReconcilerAction() {
+        var reconcileAction = new DesiredBalanceReconcilerAction() {
             @Override
             public ClusterState apply(ClusterState clusterState, Consumer<RoutingAllocation> routingAllocationAction) {
                 return allocationServiceRef.get().executeWithRoutingAllocation(clusterState, "reconcile", routingAllocationAction);
@@ -391,15 +393,72 @@ public class DesiredBalanceShardsAllocatorTests extends ESTestCase {
         }
     }
 
-    private static DiscoveryNode createDiscoveryNode(String nodeId) {
-        return new DiscoveryNode(
-            nodeId,
-            nodeId,
-            buildNewFakeTransportAddress(),
-            Map.of(),
-            Set.of(DiscoveryNodeRole.MASTER_ROLE, DiscoveryNodeRole.DATA_ROLE),
-            Version.CURRENT
+    public void testResetDesiredBalance() {
+
+        var node1 = newNode(LOCAL_NODE_ID);
+        var node2 = newNode(OTHER_NODE_ID);
+
+        var shardId = new ShardId("test-index", UUIDs.randomBase64UUID(), 0);
+        var index = createIndex(shardId.getIndexName());
+        var clusterState = ClusterState.builder(ClusterName.DEFAULT)
+            .nodes(DiscoveryNodes.builder().add(node1).add(node2).localNodeId(node1.getId()).masterNodeId(node1.getId()))
+            .metadata(Metadata.builder().put(index, false).build())
+            .routingTable(RoutingTable.builder(TestShardRoutingRoleStrategies.DEFAULT_ROLE_ONLY).addAsNew(index).build())
+            .build();
+
+        var threadPool = new TestThreadPool(getTestName());
+        var clusterService = ClusterServiceUtils.createClusterService(clusterState, threadPool);
+
+        var delegateAllocator = createShardsAllocator();
+
+        var desiredBalanceComputer = new DesiredBalanceComputer(createBuiltInClusterSettings(), threadPool, delegateAllocator) {
+
+            final AtomicReference<DesiredBalance> lastComputationInput = new AtomicReference<>();
+
+            @Override
+            public DesiredBalance compute(
+                DesiredBalance previousDesiredBalance,
+                DesiredBalanceInput desiredBalanceInput,
+                Queue<List<MoveAllocationCommand>> pendingDesiredBalanceMoves,
+                Predicate<DesiredBalanceInput> isFresh
+            ) {
+                lastComputationInput.set(previousDesiredBalance);
+                return super.compute(previousDesiredBalance, desiredBalanceInput, pendingDesiredBalanceMoves, isFresh);
+            }
+        };
+
+        var desiredBalanceShardsAllocator = new DesiredBalanceShardsAllocator(
+            delegateAllocator,
+            threadPool,
+            clusterService,
+            desiredBalanceComputer,
+            (reconcilerClusterState, routingAllocationAction) -> reconcilerClusterState
         );
+
+        var service = createAllocationService(desiredBalanceShardsAllocator, createGatewayAllocator());
+
+        try {
+            // initial computation is based on DesiredBalance.INITIAL
+            rerouteAndWait(service, clusterState, "initial-allocation");
+            assertThat(desiredBalanceComputer.lastComputationInput.get(), equalTo(DesiredBalance.INITIAL));
+
+            // any next computation is based on current desired balance
+            var current = desiredBalanceShardsAllocator.getDesiredBalance();
+            rerouteAndWait(service, clusterState, "next-allocation");
+            assertThat(desiredBalanceComputer.lastComputationInput.get(), equalTo(current));
+
+            // when desired balance is resetted then computation is based on balance with no previous assignments
+            desiredBalanceShardsAllocator.resetDesiredBalance();
+            current = desiredBalanceShardsAllocator.getDesiredBalance();
+            rerouteAndWait(service, clusterState, "reset-desired-balance");
+            assertThat(
+                desiredBalanceComputer.lastComputationInput.get(),
+                equalTo(new DesiredBalance(current.lastConvergedIndex(), Map.of()))
+            );
+        } finally {
+            clusterService.close();
+            terminate(threadPool);
+        }
     }
 
     private static IndexMetadata createIndex(String name) {
@@ -475,4 +534,8 @@ public class DesiredBalanceShardsAllocatorTests extends ESTestCase {
             }
         };
     }
+
+    private static void rerouteAndWait(AllocationService service, ClusterState clusterState, String reason) {
+        PlainActionFuture.<Void, RuntimeException>get(f -> service.reroute(clusterState, reason, f), 10, TimeUnit.SECONDS);
+    }
 }

x-pack/docs/en/security/operator-privileges/operator-only-functionality.asciidoc

@@ -26,6 +26,7 @@ given {es} version.
 * <<get-desired-nodes>>
 * <<delete-desired-nodes>>
 * <<get-desired-balance>>
+* <<delete-desired-balance>>
 
 [[operator-only-dynamic-cluster-settings]]
 ==== Operator-only dynamic cluster settings

x-pack/plugin/security/src/main/java/org/elasticsearch/xpack/security/operator/OperatorOnlyRegistry.java

@@ -7,6 +7,7 @@
 
 package org.elasticsearch.xpack.security.operator;
 
+import org.elasticsearch.action.admin.cluster.allocation.DeleteDesiredBalanceAction;
 import org.elasticsearch.action.admin.cluster.allocation.GetDesiredBalanceAction;
 import org.elasticsearch.action.admin.cluster.configuration.AddVotingConfigExclusionsAction;
 import org.elasticsearch.action.admin.cluster.configuration.ClearVotingConfigExclusionsAction;
@@ -52,7 +53,8 @@ public class OperatorOnlyRegistry {
         DeleteDesiredNodesAction.NAME,
         GetDesiredNodesAction.NAME,
         UpdateDesiredNodesAction.NAME,
-        GetDesiredBalanceAction.NAME
+        GetDesiredBalanceAction.NAME,
+        DeleteDesiredBalanceAction.NAME
     );
 
     private final ClusterSettings clusterSettings;