43
43
import org .elasticsearch .node .Node ;
44
44
import org .elasticsearch .snapshots .SnapshotShardSizeInfo ;
45
45
import org .elasticsearch .snapshots .SnapshotsInfoService ;
46
+ import org .elasticsearch .tasks .CancellableTask ;
47
+ import org .elasticsearch .tasks .TaskCancelHelper ;
48
+ import org .elasticsearch .tasks .TaskCancelledException ;
49
+ import org .elasticsearch .tasks .TaskId ;
46
50
import org .elasticsearch .test .ESTestCase ;
47
51
import org .elasticsearch .test .gateway .TestGatewayAllocator ;
48
52
import org .elasticsearch .xpack .core .ilm .ErrorStep ;
@@ -152,6 +156,7 @@ public void testEmptyCluster() {
152
156
ClusterState state = createTestClusterState (routingTable , List .of (), SingleNodeShutdownMetadata .Type .REMOVE );
153
157
154
158
ShutdownShardMigrationStatus status = TransportGetShutdownStatusAction .shardMigrationStatus (
159
+ new CancellableTask (1 , "direct" , GetShutdownStatusAction .NAME , "" , TaskId .EMPTY_TASK_ID , Map .of ()),
155
160
state ,
156
161
SHUTTING_DOWN_NODE_ID ,
157
162
SingleNodeShutdownMetadata .Type .REMOVE ,
@@ -182,6 +187,7 @@ public void testRestartAlwaysComplete() {
182
187
ClusterState state = createTestClusterState (routingTable .build (), List .of (imd ), SingleNodeShutdownMetadata .Type .RESTART );
183
188
184
189
ShutdownShardMigrationStatus status = TransportGetShutdownStatusAction .shardMigrationStatus (
190
+ new CancellableTask (1 , "direct" , GetShutdownStatusAction .NAME , "" , TaskId .EMPTY_TASK_ID , Map .of ()),
185
191
state ,
186
192
SHUTTING_DOWN_NODE_ID ,
187
193
SingleNodeShutdownMetadata .Type .RESTART ,
@@ -218,6 +224,7 @@ public void testComplete() {
218
224
ClusterState state = createTestClusterState (routingTable .build (), List .of (imd ), SingleNodeShutdownMetadata .Type .REMOVE );
219
225
220
226
ShutdownShardMigrationStatus status = TransportGetShutdownStatusAction .shardMigrationStatus (
227
+ new CancellableTask (1 , "direct" , GetShutdownStatusAction .NAME , "" , TaskId .EMPTY_TASK_ID , Map .of ()),
221
228
state ,
222
229
SHUTTING_DOWN_NODE_ID ,
223
230
SingleNodeShutdownMetadata .Type .REMOVE ,
@@ -231,6 +238,39 @@ public void testComplete() {
231
238
assertShardMigration (status , SingleNodeShutdownMetadata .Status .COMPLETE , 0 , nullValue ());
232
239
}
233
240
241
+ /**
242
+ * Ensures we check whether the task is cancelled during the computation
243
+ */
244
+ public void testCancelled () {
245
+ Index index = new Index (randomAlphaOfLength (5 ), randomAlphaOfLengthBetween (1 , 20 ));
246
+ IndexMetadata imd = generateIndexMetadata (index , 1 , 0 );
247
+ IndexRoutingTable indexRoutingTable = IndexRoutingTable .builder (index )
248
+ .addShard (TestShardRouting .newShardRouting (new ShardId (index , 0 ), SHUTTING_DOWN_NODE_ID , true , ShardRoutingState .STARTED ))
249
+ .build ();
250
+
251
+ RoutingTable .Builder routingTable = RoutingTable .builder ();
252
+ routingTable .add (indexRoutingTable );
253
+ ClusterState state = createTestClusterState (routingTable .build (), List .of (imd ), SingleNodeShutdownMetadata .Type .REMOVE );
254
+
255
+ final var task = new CancellableTask (1 , "direct" , GetShutdownStatusAction .NAME , "" , TaskId .EMPTY_TASK_ID , Map .of ());
256
+ TaskCancelHelper .cancel (task , "test" );
257
+
258
+ expectThrows (
259
+ TaskCancelledException .class ,
260
+ () -> TransportGetShutdownStatusAction .shardMigrationStatus (
261
+ task ,
262
+ state ,
263
+ SHUTTING_DOWN_NODE_ID ,
264
+ SingleNodeShutdownMetadata .Type .REMOVE ,
265
+ true ,
266
+ clusterInfoService ,
267
+ snapshotsInfoService ,
268
+ allocationService ,
269
+ allocationDeciders
270
+ )
271
+ );
272
+ }
273
+
234
274
/**
235
275
* Ensures that we properly detect "in progress" migrations while there are shards relocating off the node that's shutting down.
236
276
*/
@@ -263,6 +303,7 @@ public void testInProgressWithRelocatingShards() {
263
303
ClusterState state = createTestClusterState (routingTable .build (), List .of (imd ), SingleNodeShutdownMetadata .Type .REMOVE );
264
304
265
305
ShutdownShardMigrationStatus status = TransportGetShutdownStatusAction .shardMigrationStatus (
306
+ new CancellableTask (1 , "direct" , GetShutdownStatusAction .NAME , "" , TaskId .EMPTY_TASK_ID , Map .of ()),
266
307
state ,
267
308
SHUTTING_DOWN_NODE_ID ,
268
309
SingleNodeShutdownMetadata .Type .REMOVE ,
@@ -315,6 +356,7 @@ public void testInProgressWithShardsMovingBetweenOtherNodes() {
315
356
ClusterState state = createTestClusterState (routingTable .build (), List .of (imd ), SingleNodeShutdownMetadata .Type .REMOVE );
316
357
317
358
ShutdownShardMigrationStatus status = TransportGetShutdownStatusAction .shardMigrationStatus (
359
+ new CancellableTask (1 , "direct" , GetShutdownStatusAction .NAME , "" , TaskId .EMPTY_TASK_ID , Map .of ()),
318
360
state ,
319
361
SHUTTING_DOWN_NODE_ID ,
320
362
SingleNodeShutdownMetadata .Type .REMOVE ,
@@ -351,6 +393,7 @@ public void testStalled() {
351
393
ClusterState state = createTestClusterState (routingTable .build (), List .of (imd ), SingleNodeShutdownMetadata .Type .REMOVE );
352
394
353
395
ShutdownShardMigrationStatus status = TransportGetShutdownStatusAction .shardMigrationStatus (
396
+ new CancellableTask (1 , "direct" , GetShutdownStatusAction .NAME , "" , TaskId .EMPTY_TASK_ID , Map .of ()),
354
397
state ,
355
398
SHUTTING_DOWN_NODE_ID ,
356
399
SingleNodeShutdownMetadata .Type .REMOVE ,
@@ -441,6 +484,7 @@ public void testNotStalledIfAllShardsHaveACopyOnAnotherNode() {
441
484
ClusterState state = createTestClusterState (routingTable .build (), List .of (imd ), SingleNodeShutdownMetadata .Type .REMOVE );
442
485
443
486
ShutdownShardMigrationStatus status = TransportGetShutdownStatusAction .shardMigrationStatus (
487
+ new CancellableTask (1 , "direct" , GetShutdownStatusAction .NAME , "" , TaskId .EMPTY_TASK_ID , Map .of ()),
444
488
state ,
445
489
SHUTTING_DOWN_NODE_ID ,
446
490
SingleNodeShutdownMetadata .Type .REMOVE ,
@@ -473,6 +517,7 @@ public void testOnlyInitializingShardsRemaining() {
473
517
ClusterState state = createTestClusterState (routingTable .build (), List .of (imd ), SingleNodeShutdownMetadata .Type .REMOVE );
474
518
475
519
ShutdownShardMigrationStatus status = TransportGetShutdownStatusAction .shardMigrationStatus (
520
+ new CancellableTask (1 , "direct" , GetShutdownStatusAction .NAME , "" , TaskId .EMPTY_TASK_ID , Map .of ()),
476
521
state ,
477
522
SHUTTING_DOWN_NODE_ID ,
478
523
SingleNodeShutdownMetadata .Type .REMOVE ,
@@ -548,6 +593,7 @@ public void testNodeNotInCluster() {
548
593
.build ();
549
594
550
595
ShutdownShardMigrationStatus status = TransportGetShutdownStatusAction .shardMigrationStatus (
596
+ new CancellableTask (1 , "direct" , GetShutdownStatusAction .NAME , "" , TaskId .EMPTY_TASK_ID , Map .of ()),
551
597
state ,
552
598
bogusNodeId ,
553
599
SingleNodeShutdownMetadata .Type .REMOVE ,
@@ -642,6 +688,7 @@ private void checkStalledShardWithIlmState(
642
688
state = setIlmOperationMode (state , operationMode );
643
689
644
690
ShutdownShardMigrationStatus status = TransportGetShutdownStatusAction .shardMigrationStatus (
691
+ new CancellableTask (1 , "direct" , GetShutdownStatusAction .NAME , "" , TaskId .EMPTY_TASK_ID , Map .of ()),
645
692
state ,
646
693
SHUTTING_DOWN_NODE_ID ,
647
694
SingleNodeShutdownMetadata .Type .REMOVE ,
@@ -810,6 +857,7 @@ private ShutdownShardMigrationStatus getUnassignedShutdownStatus(Index index, In
810
857
);
811
858
812
859
return TransportGetShutdownStatusAction .shardMigrationStatus (
860
+ new CancellableTask (1 , "direct" , GetShutdownStatusAction .NAME , "" , TaskId .EMPTY_TASK_ID , Map .of ()),
813
861
state ,
814
862
SHUTTING_DOWN_NODE_ID ,
815
863
SingleNodeShutdownMetadata .Type .REMOVE ,
0 commit comments