Skip to content

Commit 7136a06

Browse files
DaveCTurnerDaveCTurner
authored
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.
1 parent 42c8402 commit 7136a06

File tree

4 files changed

+59
-33
lines changed

4 files changed

+59
-33
lines changed

‎docs/changelog/97775.yaml‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
pr: 97775
2+
summary: Fix NPE in Desired Balance API
3+
area: Allocation
4+
type: bug
5+
issues: []

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

Lines changed: 11 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -289,8 +289,18 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
289289

290290
public record ShardAssignmentView(Set<String> nodeIds, int total, int unassigned, int ignored) implements Writeable, ToXContentObject {
291291

292+
public static final ShardAssignmentView EMPTY = new ShardAssignmentView(Set.of(), 0, 0, 0);
293+
292294
public static ShardAssignmentView from(StreamInput in) throws IOException {
293-
return new ShardAssignmentView(in.readSet(StreamInput::readString), in.readVInt(), in.readVInt(), in.readVInt());
295+
final var nodeIds = in.readSet(StreamInput::readString);
296+
final var total = in.readVInt();
297+
final var unassigned = in.readVInt();
298+
final var ignored = in.readVInt();
299+
if (nodeIds.isEmpty() && total == 0 && unassigned == 0 && ignored == 0) {
300+
return EMPTY;
301+
} else {
302+
return new ShardAssignmentView(nodeIds, total, unassigned, ignored);
303+
}
294304
}
295305

296306
@Override

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

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -138,14 +138,14 @@ private Map<String, Map<Integer, DesiredBalanceResponse.DesiredShards>> createRo
138138
shardId,
139139
new DesiredBalanceResponse.DesiredShards(
140140
shardViews,
141-
shardAssignment != null
142-
? new DesiredBalanceResponse.ShardAssignmentView(
141+
shardAssignment == null
142+
? DesiredBalanceResponse.ShardAssignmentView.EMPTY
143+
: new DesiredBalanceResponse.ShardAssignmentView(
143144
shardAssignment.nodeIds(),
144145
shardAssignment.total(),
145146
shardAssignment.unassigned(),
146147
shardAssignment.ignored()
147148
)
148-
: null
149149
)
150150
);
151151
}

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

Lines changed: 40 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -9,8 +9,8 @@
99

1010
import org.elasticsearch.ResourceNotFoundException;
1111
import org.elasticsearch.Version;
12-
import org.elasticsearch.action.ActionListener;
1312
import org.elasticsearch.action.support.ActionFilters;
13+
import org.elasticsearch.action.support.PlainActionFuture;
1414
import org.elasticsearch.cluster.ClusterInfo;
1515
import org.elasticsearch.cluster.ClusterInfoService;
1616
import org.elasticsearch.cluster.ClusterName;
@@ -40,15 +40,17 @@
4040
import org.elasticsearch.index.shard.ShardId;
4141
import org.elasticsearch.rest.RestStatus;
4242
import org.elasticsearch.tasks.Task;
43+
import org.elasticsearch.tasks.TaskId;
44+
import org.elasticsearch.test.AbstractChunkedSerializingTestCase;
4345
import org.elasticsearch.threadpool.ThreadPool;
4446
import org.elasticsearch.transport.TransportService;
45-
import org.mockito.ArgumentCaptor;
4647

4748
import java.util.HashMap;
4849
import java.util.List;
4950
import java.util.Map;
5051
import java.util.Optional;
5152
import java.util.Set;
53+
import java.util.concurrent.TimeUnit;
5254
import java.util.stream.Collectors;
5355

5456
import static org.elasticsearch.cluster.ClusterModule.BALANCED_ALLOCATOR;
@@ -57,7 +59,6 @@
5759
import static org.hamcrest.Matchers.equalTo;
5860
import static org.hamcrest.Matchers.notNullValue;
5961
import static org.mockito.Mockito.mock;
60-
import static org.mockito.Mockito.verify;
6162
import static org.mockito.Mockito.when;
6263

6364
public class TransportGetDesiredBalanceActionTests extends ESAllocationTestCase {
@@ -74,13 +75,28 @@ public class TransportGetDesiredBalanceActionTests extends ESAllocationTestCase
7475
clusterInfoService,
7576
TEST_WRITE_LOAD_FORECASTER
7677
);
77-
@SuppressWarnings("unchecked")
78-
private final ActionListener<DesiredBalanceResponse> listener = mock(ActionListener.class);
78+
79+
private static DesiredBalanceResponse execute(TransportGetDesiredBalanceAction action, ClusterState clusterState) throws Exception {
80+
return PlainActionFuture.get(
81+
future -> action.masterOperation(
82+
new Task(1, "test", GetDesiredBalanceAction.NAME, "", TaskId.EMPTY_TASK_ID, Map.of()),
83+
new DesiredBalanceRequest(),
84+
clusterState,
85+
future
86+
),
87+
10,
88+
TimeUnit.SECONDS
89+
);
90+
}
91+
92+
private DesiredBalanceResponse executeAction(ClusterState clusterState) throws Exception {
93+
return execute(transportGetDesiredBalanceAction, clusterState);
94+
}
7995

8096
public void testReturnsErrorIfAllocatorIsNotDesiredBalanced() throws Exception {
8197
var clusterState = ClusterState.builder(ClusterName.DEFAULT).metadata(metadataWithConfiguredAllocator(BALANCED_ALLOCATOR)).build();
8298

83-
new TransportGetDesiredBalanceAction(
99+
final var action = new TransportGetDesiredBalanceAction(
84100
mock(TransportService.class),
85101
mock(ClusterService.class),
86102
mock(ThreadPool.class),
@@ -89,12 +105,9 @@ public void testReturnsErrorIfAllocatorIsNotDesiredBalanced() throws Exception {
89105
mock(ShardsAllocator.class),
90106
mock(ClusterInfoService.class),
91107
mock(WriteLoadForecaster.class)
92-
).masterOperation(mock(Task.class), mock(DesiredBalanceRequest.class), clusterState, listener);
93-
94-
ArgumentCaptor<ResourceNotFoundException> exceptionArgumentCaptor = ArgumentCaptor.forClass(ResourceNotFoundException.class);
95-
verify(listener).onFailure(exceptionArgumentCaptor.capture());
108+
);
96109

97-
final var exception = exceptionArgumentCaptor.getValue();
110+
final var exception = expectThrows(ResourceNotFoundException.class, () -> execute(action, clusterState));
98111
assertEquals("Desired balance allocator is not in use, no desired balance found", exception.getMessage());
99112
assertThat(exception.status(), equalTo(RestStatus.NOT_FOUND));
100113
}
@@ -104,12 +117,10 @@ public void testReturnsErrorIfDesiredBalanceIsNotAvailable() throws Exception {
104117
.metadata(metadataWithConfiguredAllocator(DESIRED_BALANCE_ALLOCATOR))
105118
.build();
106119

107-
transportGetDesiredBalanceAction.masterOperation(mock(Task.class), mock(DesiredBalanceRequest.class), clusterState, listener);
108-
109-
ArgumentCaptor<ResourceNotFoundException> exceptionArgumentCaptor = ArgumentCaptor.forClass(ResourceNotFoundException.class);
110-
verify(listener).onFailure(exceptionArgumentCaptor.capture());
111-
112-
assertEquals("Desired balance is not computed yet", exceptionArgumentCaptor.getValue().getMessage());
120+
assertEquals(
121+
"Desired balance is not computed yet",
122+
expectThrows(ResourceNotFoundException.class, () -> executeAction(clusterState)).getMessage()
123+
);
113124
}
114125

115126
public void testGetDesiredBalance() throws Exception {
@@ -220,15 +231,15 @@ public void testGetDesiredBalance() throws Exception {
220231
.routingTable(routingTable)
221232
.build();
222233

223-
transportGetDesiredBalanceAction.masterOperation(mock(Task.class), mock(DesiredBalanceRequest.class), clusterState, listener);
224-
225-
ArgumentCaptor<DesiredBalanceResponse> desiredBalanceResponseCaptor = ArgumentCaptor.forClass(DesiredBalanceResponse.class);
226-
verify(listener).onResponse(desiredBalanceResponseCaptor.capture());
227-
DesiredBalanceResponse desiredBalanceResponse = desiredBalanceResponseCaptor.getValue();
234+
final var desiredBalanceResponse = executeAction(clusterState);
228235
assertThat(desiredBalanceResponse.getStats(), equalTo(desiredBalanceStats));
229236
assertThat(desiredBalanceResponse.getClusterBalanceStats(), notNullValue());
230237
assertThat(desiredBalanceResponse.getClusterInfo(), equalTo(clusterInfo));
231238
assertEquals(indexShards.keySet(), desiredBalanceResponse.getRoutingTable().keySet());
239+
240+
assertEquals(desiredBalanceResponse, copyWriteable(desiredBalanceResponse, writableRegistry(), DesiredBalanceResponse::from));
241+
AbstractChunkedSerializingTestCase.assertChunkCount(desiredBalanceResponse, r -> 2 + r.getRoutingTable().size());
242+
232243
for (var e : desiredBalanceResponse.getRoutingTable().entrySet()) {
233244
String index = e.getKey();
234245
Map<Integer, DesiredBalanceResponse.DesiredShards> shardsMap = e.getValue();
@@ -267,14 +278,14 @@ public void testGetDesiredBalance() throws Exception {
267278
);
268279
assertEquals(indexMetadata.getTierPreference(), shardView.tierPreference());
269280
}
270-
Optional<ShardAssignment> shardAssignment = Optional.ofNullable(shardAssignments.get(indexShardRoutingTable.shardId()));
271-
if (shardAssignment.isPresent()) {
272-
assertEquals(shardAssignment.get().nodeIds(), desiredShard.desired().nodeIds());
273-
assertEquals(shardAssignment.get().total(), desiredShard.desired().total());
274-
assertEquals(shardAssignment.get().unassigned(), desiredShard.desired().unassigned());
275-
assertEquals(shardAssignment.get().ignored(), desiredShard.desired().ignored());
281+
final var shardAssignment = shardAssignments.get(indexShardRoutingTable.shardId());
282+
if (shardAssignment == null) {
283+
assertSame(desiredShard.desired(), DesiredBalanceResponse.ShardAssignmentView.EMPTY);
276284
} else {
277-
assertNull(desiredShard.desired());
285+
assertEquals(shardAssignment.nodeIds(), desiredShard.desired().nodeIds());
286+
assertEquals(shardAssignment.total(), desiredShard.desired().total());
287+
assertEquals(shardAssignment.unassigned(), desiredShard.desired().unassigned());
288+
assertEquals(shardAssignment.ignored(), desiredShard.desired().ignored());
278289
}
279290
}
280291
}

0 commit comments

Comments
 (0)