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); }
@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); }
nodeTaskMap.addTask(node, remoteTask); taskMap.put(node, remoteTask);
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 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); }
nodeTaskMap.addTask(node, task);
@Test public void testTaskCompletion() throws Exception { MockRemoteTaskFactory remoteTaskFactory = new MockRemoteTaskFactory(remoteTaskExecutor); Node chosenNode = Iterables.get(nodeManager.getActiveDatasourceNodes("foo"), 0); TaskId taskId = new TaskId(new StageId("test", "1"), "1"); RemoteTask remoteTask = remoteTaskFactory.createTableScanTask( taskId, chosenNode, ImmutableList.of(new Split("foo", TestingTransactionHandle.create("foo"), 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() throws Exception { TestingTransactionHandle transactionHandle = TestingTransactionHandle.create("foo"); Node newNode = new PrestoNode("other4", URI.create("http://127.0.0.1:14"), NodeVersion.UNKNOWN); nodeManager.addNode("foo", newNode); ImmutableList.Builder<Split> initialSplits = ImmutableList.builder(); for (int i = 0; i < 10; i++) { initialSplits.add(new Split("foo", transactionHandle, new TestSplitRemote())); } MockRemoteTaskFactory remoteTaskFactory = new MockRemoteTaskFactory(remoteTaskExecutor); // Max out number of splits on node TaskId taskId1 = new TaskId(new StageId("test", "1"), "1"); RemoteTask remoteTask1 = remoteTaskFactory.createTableScanTask(taskId1, newNode, initialSplits.build(), nodeTaskMap.createPartitionedSplitCountTracker(newNode, taskId1)); nodeTaskMap.addTask(newNode, remoteTask1); TaskId taskId2 = new TaskId(new StageId("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("foo", transactionHandle, new TestSplitRemote())); } Multimap<Node, Split> assignments = nodeSelector.computeAssignments(splits, ImmutableList.copyOf(taskMap.values())); // 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.addTask(node, remoteTask); taskMap.put(node, remoteTask);
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 testSplitCount() throws Exception { MockRemoteTaskFactory remoteTaskFactory = new MockRemoteTaskFactory(remoteTaskExecutor); Node chosenNode = Iterables.get(nodeManager.getActiveDatasourceNodes("foo"), 0); TaskId taskId1 = new TaskId(new StageId("test", "1"), "1"); RemoteTask remoteTask1 = remoteTaskFactory.createTableScanTask(taskId1, chosenNode, ImmutableList.of(new Split("foo", TestingTransactionHandle.create("foo"), new TestSplitRemote()), new Split("bar", TestingTransactionHandle.create("bar"), new TestSplitRemote())), nodeTaskMap.createPartitionedSplitCountTracker(chosenNode, taskId1)); TaskId taskId2 = new TaskId(new StageId("test", "1"), "2"); RemoteTask remoteTask2 = remoteTaskFactory.createTableScanTask( taskId2, chosenNode, ImmutableList.of(new Split("foo2", TestingTransactionHandle.create("foo2"), 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); }