|
13 | 13 | import org.apache.lucene.util.Constants;
|
14 | 14 | import org.opensearch.ExceptionsHelper;
|
15 | 15 | import org.opensearch.action.admin.cluster.node.tasks.cancel.CancelTasksRequest;
|
| 16 | +import org.opensearch.action.admin.cluster.node.tasks.get.GetTaskRequest; |
| 17 | +import org.opensearch.action.admin.cluster.node.tasks.get.GetTaskResponse; |
16 | 18 | import org.opensearch.action.admin.cluster.node.tasks.list.ListTasksRequest;
|
17 | 19 | import org.opensearch.action.admin.cluster.node.tasks.list.ListTasksResponse;
|
18 | 20 | import org.opensearch.action.support.ActionTestUtils;
|
@@ -563,8 +565,57 @@ public void testOnDemandRefreshWhileFetchingTasks() throws InterruptedException
|
563 | 565 |
|
564 | 566 | assertNotNull(taskInfo.getResourceStats());
|
565 | 567 | assertNotNull(taskInfo.getResourceStats().getResourceUsageInfo());
|
566 |
| - assertTrue(taskInfo.getResourceStats().getResourceUsageInfo().get("total") instanceof TaskResourceUsage); |
567 |
| - TaskResourceUsage taskResourceUsage = (TaskResourceUsage) taskInfo.getResourceStats().getResourceUsageInfo().get("total"); |
| 568 | + assertNotNull(taskInfo.getResourceStats().getResourceUsageInfo().get("total")); |
| 569 | + TaskResourceUsage taskResourceUsage = taskInfo.getResourceStats().getResourceUsageInfo().get("total"); |
| 570 | + assertCPUTime(taskResourceUsage.getCpuTimeInNanos()); |
| 571 | + assertTrue(taskResourceUsage.getMemoryInBytes() > 0); |
| 572 | + }; |
| 573 | + |
| 574 | + taskTestContext.operationFinishedValidator = (task, threadId) -> { assertEquals(0, resourceTasks.size()); }; |
| 575 | + |
| 576 | + startResourceAwareNodesAction(testNodes[0], false, taskTestContext, new ActionListener<NodesResponse>() { |
| 577 | + @Override |
| 578 | + public void onResponse(NodesResponse listTasksResponse) { |
| 579 | + responseReference.set(listTasksResponse); |
| 580 | + taskTestContext.requestCompleteLatch.countDown(); |
| 581 | + } |
| 582 | + |
| 583 | + @Override |
| 584 | + public void onFailure(Exception e) { |
| 585 | + throwableReference.set(e); |
| 586 | + taskTestContext.requestCompleteLatch.countDown(); |
| 587 | + } |
| 588 | + }); |
| 589 | + |
| 590 | + // Waiting for whole request to complete and return successfully till client |
| 591 | + taskTestContext.requestCompleteLatch.await(); |
| 592 | + |
| 593 | + assertTasksRequestFinishedSuccessfully(responseReference.get(), throwableReference.get()); |
| 594 | + } |
| 595 | + |
| 596 | + public void testOnDemandRefreshWhileGetTask() throws InterruptedException { |
| 597 | + setup(true, false); |
| 598 | + |
| 599 | + final AtomicReference<Throwable> throwableReference = new AtomicReference<>(); |
| 600 | + final AtomicReference<NodesResponse> responseReference = new AtomicReference<>(); |
| 601 | + |
| 602 | + TaskTestContext taskTestContext = new TaskTestContext(); |
| 603 | + |
| 604 | + Map<Long, Task> resourceTasks = testNodes[0].taskResourceTrackingService.getResourceAwareTasks(); |
| 605 | + |
| 606 | + taskTestContext.operationStartValidator = (task, threadId) -> { |
| 607 | + assertFalse(resourceTasks.isEmpty()); |
| 608 | + GetTaskResponse getTaskResponse = ActionTestUtils.executeBlocking( |
| 609 | + testNodes[0].transportGetTaskAction, |
| 610 | + new GetTaskRequest().setTaskId(new TaskId(testNodes[0].getNodeId(), new ArrayList<>(resourceTasks.values()).get(0).getId())) |
| 611 | + ); |
| 612 | + |
| 613 | + TaskInfo taskInfo = getTaskResponse.getTask().getTask(); |
| 614 | + |
| 615 | + assertNotNull(taskInfo.getResourceStats()); |
| 616 | + assertNotNull(taskInfo.getResourceStats().getResourceUsageInfo()); |
| 617 | + assertNotNull(taskInfo.getResourceStats().getResourceUsageInfo().get("total")); |
| 618 | + TaskResourceUsage taskResourceUsage = taskInfo.getResourceStats().getResourceUsageInfo().get("total"); |
568 | 619 | assertCPUTime(taskResourceUsage.getCpuTimeInNanos());
|
569 | 620 | assertTrue(taskResourceUsage.getMemoryInBytes() > 0);
|
570 | 621 | };
|
|
0 commit comments