Fix NPE in Desired Balance API (#97775) (#97802)

If a shard has no desired assignment then we cannot send the desired balance over the wire or render it as JSON because the `ShardAssignmentView` is `null`. This commit replaces the spurious `null` with an empty object, and cleans up the tests a little to remove some unnecessary mocking.

Author: DaveCTurner

docs/changelog/97775.yaml

@@ -0,0 +1,5 @@
+pr: 97775
+summary: Fix NPE in Desired Balance API
+area: Allocation
+type: bug
+issues: []

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

@@ -289,8 +289,18 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
 
  public record ShardAssignmentView(Set<String> nodeIds, int total, int unassigned, int ignored) implements Writeable, ToXContentObject {
 
+ public static final ShardAssignmentView EMPTY = new ShardAssignmentView(Set.of(), 0, 0, 0);
+
  public static ShardAssignmentView from(StreamInput in) throws IOException {
- return new ShardAssignmentView(in.readSet(StreamInput::readString), in.readVInt(), in.readVInt(), in.readVInt());
+ final var nodeIds = in.readSet(StreamInput::readString);
+ final var total = in.readVInt();
+ final var unassigned = in.readVInt();
+ final var ignored = in.readVInt();
+ if (nodeIds.isEmpty() && total == 0 && unassigned == 0 && ignored == 0) {
+ return EMPTY;
+ } else {
+ return new ShardAssignmentView(nodeIds, total, unassigned, ignored);
+ }
  }
 
  @Override

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

@@ -138,14 +138,14 @@ private Map<String, Map<Integer, DesiredBalanceResponse.DesiredShards>> createRo
  shardId,
  new DesiredBalanceResponse.DesiredShards(
  shardViews,
- shardAssignment != null
- ? new DesiredBalanceResponse.ShardAssignmentView(
+ shardAssignment == null
+ ? DesiredBalanceResponse.ShardAssignmentView.EMPTY
+ : new DesiredBalanceResponse.ShardAssignmentView(
  shardAssignment.nodeIds(),
  shardAssignment.total(),
  shardAssignment.unassigned(),
  shardAssignment.ignored()
  )
- : null
  )
  );
  }

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

@@ -9,8 +9,8 @@
 
 import org.elasticsearch.ResourceNotFoundException;
 import org.elasticsearch.Version;
-import org.elasticsearch.action.ActionListener;
 import org.elasticsearch.action.support.ActionFilters;
+import org.elasticsearch.action.support.PlainActionFuture;
 import org.elasticsearch.cluster.ClusterInfo;
 import org.elasticsearch.cluster.ClusterInfoService;
 import org.elasticsearch.cluster.ClusterName;
@@ -40,15 +40,17 @@
 import org.elasticsearch.index.shard.ShardId;
 import org.elasticsearch.rest.RestStatus;
 import org.elasticsearch.tasks.Task;
+import org.elasticsearch.tasks.TaskId;
+import org.elasticsearch.test.AbstractChunkedSerializingTestCase;
 import org.elasticsearch.threadpool.ThreadPool;
 import org.elasticsearch.transport.TransportService;
-import org.mockito.ArgumentCaptor;
 
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
 import java.util.Set;
+import java.util.concurrent.TimeUnit;
 import java.util.stream.Collectors;
 
 import static org.elasticsearch.cluster.ClusterModule.BALANCED_ALLOCATOR;
@@ -57,7 +59,6 @@
 import static org.hamcrest.Matchers.equalTo;
 import static org.hamcrest.Matchers.notNullValue;
 import static org.mockito.Mockito.mock;
-import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
 public class TransportGetDesiredBalanceActionTests extends ESAllocationTestCase {
@@ -74,13 +75,28 @@ public class TransportGetDesiredBalanceActionTests extends ESAllocationTestCase
  clusterInfoService,
  TEST_WRITE_LOAD_FORECASTER
  );
- @SuppressWarnings("unchecked")
- private final ActionListener<DesiredBalanceResponse> listener = mock(ActionListener.class);
+
+ private static DesiredBalanceResponse execute(TransportGetDesiredBalanceAction action, ClusterState clusterState) throws Exception {
+ return PlainActionFuture.get(
+ future -> action.masterOperation(
+ new Task(1, "test", GetDesiredBalanceAction.NAME, "", TaskId.EMPTY_TASK_ID, Map.of()),
+ new DesiredBalanceRequest(),
+ clusterState,
+ future
+ ),
+ 10,
+ TimeUnit.SECONDS
+ );
+ }
+
+ private DesiredBalanceResponse executeAction(ClusterState clusterState) throws Exception {
+ return execute(transportGetDesiredBalanceAction, clusterState);
+ }
 
  public void testReturnsErrorIfAllocatorIsNotDesiredBalanced() throws Exception {
  var clusterState = ClusterState.builder(ClusterName.DEFAULT).metadata(metadataWithConfiguredAllocator(BALANCED_ALLOCATOR)).build();
 
- new TransportGetDesiredBalanceAction(
+ final var action = new TransportGetDesiredBalanceAction(
  mock(TransportService.class),
  mock(ClusterService.class),
  mock(ThreadPool.class),
@@ -89,12 +105,9 @@ public void testReturnsErrorIfAllocatorIsNotDesiredBalanced() throws Exception {
  mock(ShardsAllocator.class),
  mock(ClusterInfoService.class),
  mock(WriteLoadForecaster.class)
- ).masterOperation(mock(Task.class), mock(DesiredBalanceRequest.class), clusterState, listener);
-
- ArgumentCaptor<ResourceNotFoundException> exceptionArgumentCaptor = ArgumentCaptor.forClass(ResourceNotFoundException.class);
- verify(listener).onFailure(exceptionArgumentCaptor.capture());
+ );
 
- final var exception = exceptionArgumentCaptor.getValue();
+ final var exception = expectThrows(ResourceNotFoundException.class, () -> execute(action, clusterState));
  assertEquals("Desired balance allocator is not in use, no desired balance found", exception.getMessage());
  assertThat(exception.status(), equalTo(RestStatus.NOT_FOUND));
  }
@@ -104,12 +117,10 @@ public void testReturnsErrorIfDesiredBalanceIsNotAvailable() throws Exception {
  .metadata(metadataWithConfiguredAllocator(DESIRED_BALANCE_ALLOCATOR))
  .build();
 
- transportGetDesiredBalanceAction.masterOperation(mock(Task.class), mock(DesiredBalanceRequest.class), clusterState, listener);
-
- ArgumentCaptor<ResourceNotFoundException> exceptionArgumentCaptor = ArgumentCaptor.forClass(ResourceNotFoundException.class);
- verify(listener).onFailure(exceptionArgumentCaptor.capture());
-
- assertEquals("Desired balance is not computed yet", exceptionArgumentCaptor.getValue().getMessage());
+ assertEquals(
+ "Desired balance is not computed yet",
+ expectThrows(ResourceNotFoundException.class, () -> executeAction(clusterState)).getMessage()
+ );
  }
 
  public void testGetDesiredBalance() throws Exception {
@@ -220,15 +231,15 @@ public void testGetDesiredBalance() throws Exception {
  .routingTable(routingTable)
  .build();
 
- transportGetDesiredBalanceAction.masterOperation(mock(Task.class), mock(DesiredBalanceRequest.class), clusterState, listener);
-
- ArgumentCaptor<DesiredBalanceResponse> desiredBalanceResponseCaptor = ArgumentCaptor.forClass(DesiredBalanceResponse.class);
- verify(listener).onResponse(desiredBalanceResponseCaptor.capture());
- DesiredBalanceResponse desiredBalanceResponse = desiredBalanceResponseCaptor.getValue();
+ final var desiredBalanceResponse = executeAction(clusterState);
  assertThat(desiredBalanceResponse.getStats(), equalTo(desiredBalanceStats));
  assertThat(desiredBalanceResponse.getClusterBalanceStats(), notNullValue());
  assertThat(desiredBalanceResponse.getClusterInfo(), equalTo(clusterInfo));
  assertEquals(indexShards.keySet(), desiredBalanceResponse.getRoutingTable().keySet());
+
+ assertEquals(desiredBalanceResponse, copyWriteable(desiredBalanceResponse, writableRegistry(), DesiredBalanceResponse::from));
+ AbstractChunkedSerializingTestCase.assertChunkCount(desiredBalanceResponse, r -> 2 + r.getRoutingTable().size());
+
  for (var e : desiredBalanceResponse.getRoutingTable().entrySet()) {
  String index = e.getKey();
  Map<Integer, DesiredBalanceResponse.DesiredShards> shardsMap = e.getValue();
@@ -267,14 +278,14 @@ public void testGetDesiredBalance() throws Exception {
  );
  assertEquals(indexMetadata.getTierPreference(), shardView.tierPreference());
  }
- Optional<ShardAssignment> shardAssignment = Optional.ofNullable(shardAssignments.get(indexShardRoutingTable.shardId()));
- if (shardAssignment.isPresent()) {
- assertEquals(shardAssignment.get().nodeIds(), desiredShard.desired().nodeIds());
- assertEquals(shardAssignment.get().total(), desiredShard.desired().total());
- assertEquals(shardAssignment.get().unassigned(), desiredShard.desired().unassigned());
- assertEquals(shardAssignment.get().ignored(), desiredShard.desired().ignored());
+ final var shardAssignment = shardAssignments.get(indexShardRoutingTable.shardId());
+ if (shardAssignment == null) {
+ assertSame(desiredShard.desired(), DesiredBalanceResponse.ShardAssignmentView.EMPTY);
  } else {
- assertNull(desiredShard.desired());
+ assertEquals(shardAssignment.nodeIds(), desiredShard.desired().nodeIds());
+ assertEquals(shardAssignment.total(), desiredShard.desired().total());
+ assertEquals(shardAssignment.unassigned(), desiredShard.desired().unassigned());
+ assertEquals(shardAssignment.ignored(), desiredShard.desired().ignored());
  }
  }
  }