@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); }
private NodeTasks createOrGetNodeTasks(Node node) { NodeTasks nodeTasks = nodeTasksMap.get(node); if (nodeTasks == null) { nodeTasks = addNodeTask(node); } return nodeTasks; }
public int getPartitionedSplitsOnNode(Node node) { return createOrGetNodeTasks(node).getPartitionedSplitCount(); }
totalPartitions, outputBuffers, nodeTaskMap.createPartitionedSplitCountTracker(node, taskId), summarizeTaskInfo); nodeTaskMap.addTask(node, task);
NodeTaskMap nodeTaskMap = new NodeTaskMap(finalizerService); MockRemoteTaskFactory.MockRemoteTask remoteTask = remoteTaskFactory.createTableScanTask(taskId, node, initialSplits.build(), nodeTaskMap.createPartitionedSplitCountTracker(node, taskId)); nodeTaskMap.addTask(node, remoteTask); taskMap.put(node, remoteTask);
@Test public void testScheduleNoSplits() { StageExecutionPlan plan = createPlan(createFixedSplitSource(0, TestingSplit::createRemoteSplit)); NodeTaskMap nodeTaskMap = new NodeTaskMap(finalizerService); SqlStageExecution stage = createSqlStageExecution(plan, nodeTaskMap); StageScheduler scheduler = getSourcePartitionedScheduler(plan, stage, nodeManager, nodeTaskMap, 1); ScheduleResult scheduleResult = scheduler.schedule(); assertEquals(scheduleResult.getNewTasks().size(), 1); assertEffectivelyFinished(scheduleResult, scheduler); stage.abort(); }
NodeTaskMap nodeTaskMap = new NodeTaskMap(finalizerService); InMemoryNodeManager nodeManager = new InMemoryNodeManager(); TaskId taskId = new TaskId("test", 1, task); task++; MockRemoteTaskFactory.MockRemoteTask remoteTask = remoteTaskFactory.createTableScanTask(taskId, node, ImmutableList.copyOf(assignments.get(node)), nodeTaskMap.createPartitionedSplitCountTracker(node, taskId)); remoteTask.startSplits(25); nodeTaskMap.addTask(node, remoteTask); taskMap.put(node, remoteTask);
@Test public void testScheduleSlowSplitSource() { QueuedSplitSource queuedSplitSource = new QueuedSplitSource(TestingSplit::createRemoteSplit); StageExecutionPlan plan = createPlan(queuedSplitSource); NodeTaskMap nodeTaskMap = new NodeTaskMap(finalizerService); SqlStageExecution stage = createSqlStageExecution(plan, nodeTaskMap); StageScheduler scheduler = getSourcePartitionedScheduler(plan, stage, nodeManager, nodeTaskMap, 1); // schedule with no splits - will block ScheduleResult scheduleResult = scheduler.schedule(); assertFalse(scheduleResult.isFinished()); assertFalse(scheduleResult.getBlocked().isDone()); assertEquals(scheduleResult.getNewTasks().size(), 0); assertEquals(stage.getAllTasks().size(), 0); queuedSplitSource.addSplits(1); assertTrue(scheduleResult.getBlocked().isDone()); }
initialSplits.build(), outputBuffers.get(), nodeTaskMap.createPartitionedSplitCountTracker(node, taskId)); nodeTaskMap.addTask(node, task);
@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); }
NodeTaskMap nodeTaskMap = new NodeTaskMap(finalizerService); MockRemoteTaskFactory.MockRemoteTask remoteTask = remoteTaskFactory.createTableScanTask(taskId, node, initialSplits.build(), nodeTaskMap.createPartitionedSplitCountTracker(node, taskId)); nodeTaskMap.addTask(node, remoteTask); taskMap.put(node, remoteTask);
@BeforeMethod public void setUp() { finalizerService = new FinalizerService(); nodeTaskMap = new NodeTaskMap(finalizerService); nodeManager = new InMemoryNodeManager(); ImmutableList.Builder<Node> nodeBuilder = ImmutableList.builder(); nodeBuilder.add(new PrestoNode("other1", URI.create("http://127.0.0.1:11"), NodeVersion.UNKNOWN, false)); nodeBuilder.add(new PrestoNode("other2", URI.create("http://127.0.0.1:12"), NodeVersion.UNKNOWN, false)); nodeBuilder.add(new PrestoNode("other3", URI.create("http://127.0.0.1:13"), NodeVersion.UNKNOWN, false)); ImmutableList<Node> nodes = nodeBuilder.build(); nodeManager.addNode(CONNECTOR_ID, nodes); NodeSchedulerConfig nodeSchedulerConfig = new NodeSchedulerConfig() .setMaxSplitsPerNode(20) .setIncludeCoordinator(false) .setMaxPendingSplitsPerTask(10); NodeScheduler nodeScheduler = new NodeScheduler(new LegacyNetworkTopology(), nodeManager, nodeSchedulerConfig, nodeTaskMap); // contents of taskMap indicate the node-task map for the current stage taskMap = new HashMap<>(); nodeSelector = nodeScheduler.createNodeSelector(CONNECTOR_ID); remoteTaskExecutor = newCachedThreadPool(daemonThreadsNamed("remoteTaskExecutor-%s")); remoteTaskScheduledExecutor = newScheduledThreadPool(2, daemonThreadsNamed("remoteTaskScheduledExecutor-%s")); finalizerService.start(); }
public void addTask(Node node, RemoteTask task) { createOrGetNodeTasks(node).addTask(task); }
private NodeTasks createOrGetNodeTasks(Node node) { NodeTasks nodeTasks = nodeTasksMap.get(node); if (nodeTasks == null) { nodeTasks = addNodeTask(node); } return nodeTasks; }
RemoteTask remoteTask = remoteTaskFactory.createTableScanTask(taskId, node, initialSplits.build(), nodeTaskMap.createPartitionedSplitCountTracker(node, taskId)); nodeTaskMap.addTask(node, remoteTask); tasks.add(remoteTask); RemoteTask newRemoteTask = remoteTaskFactory.createTableScanTask(taskId, newNode, initialSplits.build(), nodeTaskMap.createPartitionedSplitCountTracker(newNode, taskId)); nodeTaskMap.addTask(newNode, newRemoteTask); tasks.add(newRemoteTask); task.abort(); assertEquals(nodeTaskMap.getPartitionedSplitsOnNode(newNode), 0);
NodeTaskMap nodeTaskMap = new NodeTaskMap(finalizerService); InMemoryNodeManager nodeManager = new InMemoryNodeManager(); TaskId taskId = new TaskId(new StageId("test", "1"), String.valueOf(task)); task++; MockRemoteTaskFactory.MockRemoteTask remoteTask = remoteTaskFactory.createTableScanTask(taskId, node, ImmutableList.copyOf(assignments.get(node)), nodeTaskMap.createPartitionedSplitCountTracker(node, taskId)); remoteTask.startSplits(20); nodeTaskMap.addTask(node, remoteTask); taskMap.put(node, remoteTask);
@Test public void testScheduleSplitsBatched() { StageExecutionPlan plan = createPlan(createFixedSplitSource(60, TestingSplit::createRemoteSplit)); NodeTaskMap nodeTaskMap = new NodeTaskMap(finalizerService); SqlStageExecution stage = createSqlStageExecution(plan, nodeTaskMap); StageScheduler scheduler = getSourcePartitionedScheduler(plan, stage, nodeManager, nodeTaskMap, 7); for (int i = 0; i <= (60 / 7); i++) { ScheduleResult scheduleResult = scheduler.schedule(); // finishes when last split is fetched if (i == (60 / 7)) { assertEffectivelyFinished(scheduleResult, scheduler); } else { assertFalse(scheduleResult.isFinished()); } // never blocks assertTrue(scheduleResult.getBlocked().isDone()); // first three splits create new tasks assertEquals(scheduleResult.getNewTasks().size(), i == 0 ? 3 : 0); assertEquals(stage.getAllTasks().size(), 3); assertPartitionedSplitCount(stage, min((i + 1) * 7, 60)); } for (RemoteTask remoteTask : stage.getAllTasks()) { assertEquals(remoteTask.getPartitionedSplitCount(), 20); } stage.abort(); }
public PartitionedSplitCountTracker createPartitionedSplitCountTracker(Node node, TaskId taskId) { return createOrGetNodeTasks(node).createPartitionedSplitCountTracker(taskId); }
@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); }
@Test public void testScheduleSplitsOneAtATime() { StageExecutionPlan plan = createPlan(createFixedSplitSource(60, TestingSplit::createRemoteSplit)); NodeTaskMap nodeTaskMap = new NodeTaskMap(finalizerService); SqlStageExecution stage = createSqlStageExecution(plan, nodeTaskMap); StageScheduler scheduler = getSourcePartitionedScheduler(plan, stage, nodeManager, nodeTaskMap, 1); for (int i = 0; i < 60; i++) { ScheduleResult scheduleResult = scheduler.schedule(); // only finishes when last split is fetched if (i == 59) { assertEffectivelyFinished(scheduleResult, scheduler); } else { assertFalse(scheduleResult.isFinished()); } // never blocks assertTrue(scheduleResult.getBlocked().isDone()); // first three splits create new tasks assertEquals(scheduleResult.getNewTasks().size(), i < 3 ? 1 : 0); assertEquals(stage.getAllTasks().size(), i < 3 ? i + 1 : 3); assertPartitionedSplitCount(stage, min(i + 1, 60)); } for (RemoteTask remoteTask : stage.getAllTasks()) { assertEquals(remoteTask.getPartitionedSplitCount(), 20); } stage.abort(); }