@Test
public void testCancel() throws Exception {
SqlTask sqlTask = createInitialTask();
TaskInfo taskInfo =
sqlTask.updateTask(
SESSION, PLAN_FRAGMENT, ImmutableList.<TaskSource>of(), INITIAL_EMPTY_OUTPUT_BUFFERS);
assertEquals(taskInfo.getState(), TaskState.RUNNING);
assertNull(taskInfo.getStats().getEndTime());
taskInfo = sqlTask.getTaskInfo();
assertEquals(taskInfo.getState(), TaskState.RUNNING);
assertNull(taskInfo.getStats().getEndTime());
taskInfo = sqlTask.cancel();
assertEquals(taskInfo.getState(), TaskState.CANCELED);
assertNotNull(taskInfo.getStats().getEndTime());
taskInfo = sqlTask.getTaskInfo();
assertEquals(taskInfo.getState(), TaskState.CANCELED);
assertNotNull(taskInfo.getStats().getEndTime());
}
@Test
public void testBufferCloseOnCancel() throws Exception {
SqlTask sqlTask = createInitialTask();
updateTask(
sqlTask,
EMPTY_SOURCES,
INITIAL_EMPTY_OUTPUT_BUFFERS.withBuffer("out", new UnpartitionedPagePartitionFunction()));
ListenableFuture<BufferResult> bufferResult =
sqlTask.getTaskResults("out", 0, new DataSize(1, MEGABYTE));
assertFalse(bufferResult.isDone());
sqlTask.cancel();
assertEquals(sqlTask.getTaskInfo().getState(), TaskState.CANCELED);
// buffer will be closed by cancel event.. the event is async so wait a bit for event to
// propagate
assertTrue(bufferResult.get(200, MILLISECONDS).isBufferClosed());
bufferResult = sqlTask.getTaskResults("out", 0, new DataSize(1, MEGABYTE));
assertTrue(bufferResult.isDone());
assertTrue(bufferResult.get().isBufferClosed());
}