public static SqlTaskExecution createSqlTaskExecution(
TaskStateMachine taskStateMachine,
TaskContext taskContext,
SharedBuffer sharedBuffer,
PlanFragment fragment,
List<TaskSource> sources,
LocalExecutionPlanner planner,
TaskExecutor taskExecutor,
Executor notificationExecutor,
QueryMonitor queryMonitor) {
SqlTaskExecution task =
new SqlTaskExecution(
taskStateMachine,
taskContext,
sharedBuffer,
fragment,
planner,
taskExecutor,
queryMonitor,
notificationExecutor);
try (SetThreadName ignored = new SetThreadName("Task-%s", task.getTaskId())) {
task.start();
task.addSources(sources);
return task;
}
}
@Test
public void testBufferCloseOnFinish() throws Exception {
outputBuffers =
outputBuffers
.withBuffer("out", new UnpartitionedPagePartitionFunction())
.withNoMoreBufferIds();
taskExecution.addResultQueue(outputBuffers);
BufferResult bufferResult =
taskExecution.getResults(
"out", 0, new DataSize(1, Unit.MEGABYTE), new Duration(0, TimeUnit.MILLISECONDS));
assertFalse(bufferResult.isBufferClosed());
bufferResult =
taskExecution.getResults(
"out", 0, new DataSize(1, Unit.MEGABYTE), new Duration(0, TimeUnit.MILLISECONDS));
assertFalse(bufferResult.isBufferClosed());
taskExecution.addSources(
ImmutableList.of(new TaskSource(tableScanNodeId, ImmutableSet.<ScheduledSplit>of(), true)));
assertEquals(taskExecution.getTaskInfo(false).getState(), TaskState.FINISHED);
// buffer will be closed by cancel event (wait for 500 MS for event to fire)
bufferResult =
taskExecution.getResults(
"out", 0, new DataSize(1, Unit.MEGABYTE), new Duration(500, TimeUnit.MILLISECONDS));
assertTrue(bufferResult.isBufferClosed());
bufferResult =
taskExecution.getResults(
"out", 0, new DataSize(1, Unit.MEGABYTE), new Duration(500, TimeUnit.MILLISECONDS));
assertTrue(bufferResult.isBufferClosed());
}
public static SqlTaskExecution createSqlTaskExecution(
Session session,
TaskId taskId,
URI location,
PlanFragment fragment,
List<TaskSource> sources,
OutputBuffers outputBuffers,
LocalExecutionPlanner planner,
DataSize maxBufferSize,
TaskExecutor taskExecutor,
ExecutorService notificationExecutor,
DataSize maxTaskMemoryUsage,
DataSize operatorPreAllocatedMemory,
QueryMonitor queryMonitor,
boolean cpuTimerEnabled) {
SqlTaskExecution task =
new SqlTaskExecution(
session,
taskId,
location,
fragment,
outputBuffers,
planner,
maxBufferSize,
taskExecutor,
maxTaskMemoryUsage,
operatorPreAllocatedMemory,
queryMonitor,
notificationExecutor,
cpuTimerEnabled);
try (SetThreadName setThreadName = new SetThreadName("Task-%s", taskId)) {
task.start();
task.addSources(sources);
task.recordHeartbeat();
return task;
}
}