@JsonCreator public static TaskId valueOf(String taskId) { return new TaskId(taskId); }
private Builder(Executor notificationExecutor, ScheduledExecutorService yieldExecutor, Session session) { this.notificationExecutor = notificationExecutor; this.yieldExecutor = yieldExecutor; this.session = session; this.taskStateMachine = new TaskStateMachine(new TaskId("query", 0, 0), notificationExecutor); }
public static TaskContext createTaskContext(DefaultQueryContext queryContext, Executor executor, Session session) { return createTaskContext(queryContext, session, new TaskStateMachine(new TaskId("query", 0, 0), executor)); }
public synchronized Optional<RemoteTask> scheduleTask(Node node, int partition, OptionalInt totalPartitions) { requireNonNull(node, "node is null"); if (stateMachine.getState().isDone()) { return Optional.empty(); } checkState(!splitsScheduled.get(), "scheduleTask can not be called once splits have been scheduled"); return Optional.of(scheduleTask(node, new TaskId(stateMachine.getStageId(), partition), ImmutableMultimap.of(), totalPartitions)); }
private void createTask(TaskSpecification specification) { if (specification.getType() == LEAF) { runningTasks.put(specification, new LeafTask( taskExecutor, specification, new TaskId(specification.getName(), 0, runningTasks.get(specification).size() + completedTasks.get(specification).size()))); } else { runningTasks.put(specification, new IntermediateTask( taskExecutor, specification, new TaskId(specification.getName(), 0, runningTasks.get(specification).size() + completedTasks.get(specification).size()))); } }
@Test public void testUpdateLevelWithCap() { MultilevelSplitQueue splitQueue = new MultilevelSplitQueue(2); TaskHandle handle0 = new TaskHandle(new TaskId("test0", 0, 0), splitQueue, () -> 1, 1, new Duration(1, SECONDS), OptionalInt.empty()); long quantaNanos = MINUTES.toNanos(10); handle0.addScheduledNanos(quantaNanos); long cappedNanos = Math.min(quantaNanos, LEVEL_CONTRIBUTION_CAP); for (int i = 0; i < (LEVEL_THRESHOLD_SECONDS.length - 1); i++) { long thisLevelTime = Math.min(SECONDS.toNanos(LEVEL_THRESHOLD_SECONDS[i + 1] - LEVEL_THRESHOLD_SECONDS[i]), cappedNanos); assertEquals(splitQueue.getLevelScheduledTime(i), thisLevelTime); cappedNanos -= thisLevelTime; } }
@Test public void testMultipleNodes() { FixedCountScheduler nodeScheduler = new FixedCountScheduler( (node, partition, totalPartitions) -> Optional.of(taskFactory.createTableScanTask( new TaskId("test", 1, 1), node, ImmutableList.of(), new PartitionedSplitCountTracker(delta -> {}))), generateRandomNodes(5)); ScheduleResult result = nodeScheduler.schedule(); assertTrue(result.isFinished()); assertTrue(result.getBlocked().isDone()); assertEquals(result.getNewTasks().size(), 5); assertEquals(result.getNewTasks().stream().map(RemoteTask::getNodeId).collect(toImmutableSet()).size(), 5); }
@Test public void testSingleNode() { FixedCountScheduler nodeScheduler = new FixedCountScheduler( (node, partition, totalPartitions) -> Optional.of(taskFactory.createTableScanTask( new TaskId("test", 1, 1), node, ImmutableList.of(), new PartitionedSplitCountTracker(delta -> {}))), generateRandomNodes(1)); ScheduleResult result = nodeScheduler.schedule(); assertTrue(result.isFinished()); assertTrue(result.getBlocked().isDone()); assertEquals(result.getNewTasks().size(), 1); assertTrue(result.getNewTasks().iterator().next().getNodeId().equals("other 0")); }
@Test public void testLevelContributionCap() { MultilevelSplitQueue splitQueue = new MultilevelSplitQueue(2); TaskHandle handle0 = new TaskHandle(new TaskId("test0", 0, 0), splitQueue, () -> 1, 1, new Duration(1, SECONDS), OptionalInt.empty()); TaskHandle handle1 = new TaskHandle(new TaskId("test1", 0, 0), splitQueue, () -> 1, 1, new Duration(1, SECONDS), OptionalInt.empty()); for (int i = 0; i < (LEVEL_THRESHOLD_SECONDS.length - 1); i++) { long levelAdvanceTime = SECONDS.toNanos(LEVEL_THRESHOLD_SECONDS[i + 1] - LEVEL_THRESHOLD_SECONDS[i]); handle0.addScheduledNanos(levelAdvanceTime); assertEquals(handle0.getPriority().getLevel(), i + 1); handle1.addScheduledNanos(levelAdvanceTime); assertEquals(handle1.getPriority().getLevel(), i + 1); assertEquals(splitQueue.getLevelScheduledTime(i), 2 * Math.min(levelAdvanceTime, LEVEL_CONTRIBUTION_CAP)); assertEquals(splitQueue.getLevelScheduledTime(i + 1), 0); } }
private RemoteTask createRemoteTask(HttpRemoteTaskFactory httpRemoteTaskFactory) { return httpRemoteTaskFactory.createRemoteTask( TEST_SESSION, new TaskId("test", 1, 2), new PrestoNode("node-id", URI.create("http://fake.invalid/"), new NodeVersion("version"), false), TaskTestUtils.PLAN_FRAGMENT, ImmutableMultimap.of(), OptionalInt.empty(), createInitialEmptyOutputBuffers(OutputBuffers.BufferType.BROADCAST), new NodeTaskMap.PartitionedSplitCountTracker(i -> {}), true); }
private OperatorContext createContexts(SqlTask sqlTask) { TaskContext taskContext = sqlTask.getQueryContext().addTaskContext(new TaskStateMachine(new TaskId("q", 1, 1), executor), session, false, false, OptionalInt.empty()); PipelineContext pipelineContext = taskContext.addPipelineContext(0, false, false, false); DriverContext driverContext = pipelineContext.addDriverContext(); OperatorContext operatorContext = driverContext.addOperatorContext(1, new PlanNodeId("na"), "na"); return operatorContext; }
@Test public void testTaskCompletion() throws Exception { MockRemoteTaskFactory remoteTaskFactory = new MockRemoteTaskFactory(remoteTaskExecutor, remoteTaskScheduledExecutor); Node chosenNode = Iterables.get(nodeManager.getActiveConnectorNodes(CONNECTOR_ID), 0); TaskId taskId = new TaskId("test", 1, 1); RemoteTask remoteTask = remoteTaskFactory.createTableScanTask( taskId, chosenNode, ImmutableList.of(new Split(CONNECTOR_ID, TestingTransactionHandle.create(), new TestSplitRemote())), nodeTaskMap.createPartitionedSplitCountTracker(chosenNode, taskId)); nodeTaskMap.addTask(chosenNode, remoteTask); assertEquals(nodeTaskMap.getPartitionedSplitsOnNode(chosenNode), 1); remoteTask.abort(); MILLISECONDS.sleep(100); // Sleep until cache expires assertEquals(nodeTaskMap.getPartitionedSplitsOnNode(chosenNode), 0); remoteTask.abort(); assertEquals(nodeTaskMap.getPartitionedSplitsOnNode(chosenNode), 0); }
@Test public void testMaxSplitsPerNode() { TestingTransactionHandle transactionHandle = TestingTransactionHandle.create(); Node newNode = new PrestoNode("other4", URI.create("http://127.0.0.1:14"), NodeVersion.UNKNOWN, false); nodeManager.addNode(CONNECTOR_ID, newNode); ImmutableList.Builder<Split> initialSplits = ImmutableList.builder(); for (int i = 0; i < 10; i++) { initialSplits.add(new Split(CONNECTOR_ID, transactionHandle, new TestSplitRemote())); } MockRemoteTaskFactory remoteTaskFactory = new MockRemoteTaskFactory(remoteTaskExecutor, remoteTaskScheduledExecutor); // Max out number of splits on node TaskId taskId1 = new TaskId("test", 1, 1); RemoteTask remoteTask1 = remoteTaskFactory.createTableScanTask(taskId1, newNode, initialSplits.build(), nodeTaskMap.createPartitionedSplitCountTracker(newNode, taskId1)); nodeTaskMap.addTask(newNode, remoteTask1); TaskId taskId2 = new TaskId("test", 1, 2); RemoteTask remoteTask2 = remoteTaskFactory.createTableScanTask(taskId2, newNode, initialSplits.build(), nodeTaskMap.createPartitionedSplitCountTracker(newNode, taskId2)); nodeTaskMap.addTask(newNode, remoteTask2); Set<Split> splits = new HashSet<>(); for (int i = 0; i < 5; i++) { splits.add(new Split(CONNECTOR_ID, transactionHandle, new TestSplitRemote())); } Multimap<Node, Split> assignments = nodeSelector.computeAssignments(splits, ImmutableList.copyOf(taskMap.values())).getAssignments(); // no split should be assigned to the newNode, as it already has maxNodeSplits assigned to it assertFalse(assignments.keySet().contains(newNode)); remoteTask1.abort(); remoteTask2.abort(); assertEquals(nodeTaskMap.getPartitionedSplitsOnNode(newNode), 0); }
@Test(invocationCount = 100) public void testQuantaFairness() { TestingTicker ticker = new TestingTicker(); TaskExecutor taskExecutor = new TaskExecutor(1, 2, 3, 4, ticker); taskExecutor.start(); ticker.increment(20, MILLISECONDS); try { TaskHandle shortQuantaTaskHandle = taskExecutor.addTask(new TaskId("shortQuanta", 0, 0), () -> 0, 10, new Duration(1, MILLISECONDS), OptionalInt.empty()); TaskHandle longQuantaTaskHandle = taskExecutor.addTask(new TaskId("longQuanta", 0, 0), () -> 0, 10, new Duration(1, MILLISECONDS), OptionalInt.empty()); Phaser globalPhaser = new Phaser(); TestingJob shortQuantaDriver = new TestingJob(ticker, new Phaser(), new Phaser(), globalPhaser, 10, 10); TestingJob longQuantaDriver = new TestingJob(ticker, new Phaser(), new Phaser(), globalPhaser, 10, 20); taskExecutor.enqueueSplits(shortQuantaTaskHandle, true, ImmutableList.of(shortQuantaDriver)); taskExecutor.enqueueSplits(longQuantaTaskHandle, true, ImmutableList.of(longQuantaDriver)); for (int i = 0; i < 11; i++) { globalPhaser.arriveAndAwaitAdvance(); } assertTrue(shortQuantaDriver.getCompletedPhases() >= 7 && shortQuantaDriver.getCompletedPhases() <= 8); assertTrue(longQuantaDriver.getCompletedPhases() >= 3 && longQuantaDriver.getCompletedPhases() <= 4); globalPhaser.arriveAndDeregister(); } finally { taskExecutor.stop(); } }
@Test public void testSplitCount() { MockRemoteTaskFactory remoteTaskFactory = new MockRemoteTaskFactory(remoteTaskExecutor, remoteTaskScheduledExecutor); Node chosenNode = Iterables.get(nodeManager.getActiveConnectorNodes(CONNECTOR_ID), 0); TaskId taskId1 = new TaskId("test", 1, 1); RemoteTask remoteTask1 = remoteTaskFactory.createTableScanTask(taskId1, chosenNode, ImmutableList.of( new Split(CONNECTOR_ID, TestingTransactionHandle.create(), new TestSplitRemote()), new Split(CONNECTOR_ID, TestingTransactionHandle.create(), new TestSplitRemote())), nodeTaskMap.createPartitionedSplitCountTracker(chosenNode, taskId1)); TaskId taskId2 = new TaskId("test", 1, 2); RemoteTask remoteTask2 = remoteTaskFactory.createTableScanTask( taskId2, chosenNode, ImmutableList.of(new Split(CONNECTOR_ID, TestingTransactionHandle.create(), new TestSplitRemote())), nodeTaskMap.createPartitionedSplitCountTracker(chosenNode, taskId2)); nodeTaskMap.addTask(chosenNode, remoteTask1); nodeTaskMap.addTask(chosenNode, remoteTask2); assertEquals(nodeTaskMap.getPartitionedSplitsOnNode(chosenNode), 3); remoteTask1.abort(); assertEquals(nodeTaskMap.getPartitionedSplitsOnNode(chosenNode), 1); remoteTask2.abort(); assertEquals(nodeTaskMap.getPartitionedSplitsOnNode(chosenNode), 0); }
@BeforeMethod public void setUpTest() { memoryPool = new MemoryPool(new MemoryPoolId("test"), memoryPoolSize); queryContext = new DefaultQueryContext( new QueryId("test_query"), queryMaxMemory, queryMaxTotalMemory, memoryPool, new TestingGcMonitor(), notificationExecutor, yieldExecutor, queryMaxSpillSize, spillSpaceTracker); taskContext = queryContext.addTaskContext( new TaskStateMachine(new TaskId("query", 0, 0), notificationExecutor), testSessionBuilder().build(), true, true, OptionalInt.empty()); pipelineContext = taskContext.addPipelineContext(0, true, true, false); driverContext = pipelineContext.addDriverContext(); operatorContext = driverContext.addOperatorContext(1, new PlanNodeId("a"), "test-operator"); }
@Test public void testTaskHandle() { TestingTicker ticker = new TestingTicker(); TaskExecutor taskExecutor = new TaskExecutor(4, 8, 3, 4, ticker); taskExecutor.start(); try { TaskId taskId = new TaskId("test", 0, 0); TaskHandle taskHandle = taskExecutor.addTask(taskId, () -> 0, 10, new Duration(1, MILLISECONDS), OptionalInt.empty()); Phaser beginPhase = new Phaser(); beginPhase.register(); Phaser verificationComplete = new Phaser(); verificationComplete.register(); TestingJob driver1 = new TestingJob(ticker, new Phaser(), beginPhase, verificationComplete, 10, 0); TestingJob driver2 = new TestingJob(ticker, new Phaser(), beginPhase, verificationComplete, 10, 0); // force enqueue a split taskExecutor.enqueueSplits(taskHandle, true, ImmutableList.of(driver1)); assertEquals(taskHandle.getRunningLeafSplits(), 0); // normal enqueue a split taskExecutor.enqueueSplits(taskHandle, false, ImmutableList.of(driver2)); assertEquals(taskHandle.getRunningLeafSplits(), 1); // let the split continue to run beginPhase.arriveAndDeregister(); verificationComplete.arriveAndDeregister(); } finally { taskExecutor.stop(); } }
private SqlTask newSqlTask() { TaskId taskId = new TaskId("query", 0, idGeneator.incrementAndGet()); URI location = URI.create("fake://task/" + taskId); return createSqlTask( taskId, location, "fake", new DefaultQueryContext(new QueryId("query"), new DataSize(1, MEGABYTE), new DataSize(2, MEGABYTE), memoryPool, new TestingGcMonitor(), executor, scheduledExecutor, new DataSize(1, GIGABYTE), spillSpaceTracker), sqlTaskExecutionFactory, executor, Functions.identity(), new DataSize(32, MEGABYTE), new CounterStat()); } }
@Test(invocationCount = 100) public void testLevelMovement() { TestingTicker ticker = new TestingTicker(); TaskExecutor taskExecutor = new TaskExecutor(2, 2, 3, 4, ticker); taskExecutor.start(); ticker.increment(20, MILLISECONDS); try { TaskHandle testTaskHandle = taskExecutor.addTask(new TaskId("test", 0, 0), () -> 0, 10, new Duration(1, MILLISECONDS), OptionalInt.empty()); Phaser globalPhaser = new Phaser(); globalPhaser.bulkRegister(3); int quantaTimeMills = 500; int phasesPerSecond = 1000 / quantaTimeMills; int totalPhases = LEVEL_THRESHOLD_SECONDS[LEVEL_THRESHOLD_SECONDS.length - 1] * phasesPerSecond; TestingJob driver1 = new TestingJob(ticker, globalPhaser, new Phaser(), new Phaser(), totalPhases, quantaTimeMills); TestingJob driver2 = new TestingJob(ticker, globalPhaser, new Phaser(), new Phaser(), totalPhases, quantaTimeMills); taskExecutor.enqueueSplits(testTaskHandle, true, ImmutableList.of(driver1, driver2)); int completedPhases = 0; for (int i = 0; i < (LEVEL_THRESHOLD_SECONDS.length - 1); i++) { for (; (completedPhases / phasesPerSecond) < LEVEL_THRESHOLD_SECONDS[i + 1]; completedPhases++) { globalPhaser.arriveAndAwaitAdvance(); } assertEquals(testTaskHandle.getPriority().getLevel(), i + 1); } globalPhaser.arriveAndDeregister(); } finally { taskExecutor.stop(); } }
public SqlTask createInitialTask() { TaskId taskId = new TaskId("query", 0, nextTaskId.incrementAndGet()); URI location = URI.create("fake://task/" + taskId); DefaultQueryContext queryContext = new DefaultQueryContext(new QueryId("query"), new DataSize(1, MEGABYTE), new DataSize(2, MEGABYTE), new MemoryPool(new MemoryPoolId("test"), new DataSize(1, GIGABYTE)), new TestingGcMonitor(), taskNotificationExecutor, driverYieldExecutor, new DataSize(1, MEGABYTE), new SpillSpaceTracker(new DataSize(1, GIGABYTE))); queryContext.addTaskContext(new TaskStateMachine(taskId, taskNotificationExecutor), testSessionBuilder().build(), false, false, OptionalInt.empty()); return createSqlTask( taskId, location, "fake", queryContext, sqlTaskExecutionFactory, taskNotificationExecutor, Functions.identity(), new DataSize(32, MEGABYTE), new CounterStat()); } }