Examples with InternalNode io.trino.metadata.InternalNode used on opensource projects

Example 21 with InternalNode

use of io.trino.metadata.InternalNode in project trino by trinodb.

the class TestNodeScheduler method testOptimizedLocalScheduling.

@Test
public void testOptimizedLocalScheduling() {
    InternalNode node1 = new InternalNode("node1", URI.create("http://10.0.0.1:11"), NodeVersion.UNKNOWN, false);
    nodeManager.addNode(CONNECTOR_ID, node1);
    InternalNode node2 = new InternalNode("node2", URI.create("http://10.0.0.1:12"), NodeVersion.UNKNOWN, false);
    nodeManager.addNode(CONNECTOR_ID, node2);
    Set<Split> splits = new LinkedHashSet<>();
    // 20 splits with node1 as local node to be assigned in the first iteration of computeAssignments
    for (int i = 0; i < 20; i++) {
        splits.add(new Split(CONNECTOR_ID, new TestSplitLocal(), Lifespan.taskWide()));
    }
    // computeAssignments just returns a mapping of nodes with splits to be assigned, it does not assign splits
    Multimap<InternalNode, Split> assignments1 = nodeSelector.computeAssignments(splits, ImmutableList.copyOf(taskMap.values())).getAssignments();
    // Check that all 20 splits are being assigned to node1 as optimized-local-scheduling is enabled
    assertEquals(assignments1.size(), 20);
    assertEquals(assignments1.keySet().size(), 2);
    assertTrue(assignments1.keySet().contains(node1));
    assertTrue(assignments1.keySet().contains(node2));
    // 19 splits with node2 as local node to be assigned in the first iteration of computeAssignments
    for (int i = 0; i < 19; i++) {
        splits.add(new Split(CONNECTOR_ID, new TestSplitRemote(HostAddress.fromString("10.0.0.1:12")), Lifespan.taskWide()));
    }
    Multimap<InternalNode, Split> assignments2 = nodeSelector.computeAssignments(splits, ImmutableList.copyOf(taskMap.values())).getAssignments();
    // Check that all 39 splits are being assigned (20 splits assigned to node1 and 19 splits assigned to node2)
    assertEquals(assignments2.size(), 39);
    assertEquals(assignments2.keySet().size(), 2);
    assertTrue(assignments2.keySet().contains(node1));
    assertTrue(assignments2.keySet().contains(node2));
    long node1Splits = assignments2.values().stream().map(Split::getConnectorSplit).filter(TestSplitLocal.class::isInstance).count();
    assertEquals(node1Splits, 20);
    long node2Splits = assignments2.values().stream().map(Split::getConnectorSplit).filter(TestSplitRemote.class::isInstance).count();
    assertEquals(node2Splits, 19);
    // 1 split with node1 as local node
    splits.add(new Split(CONNECTOR_ID, new TestSplitLocal(), Lifespan.taskWide()));
    // 1 split with node2 as local node
    splits.add(new Split(CONNECTOR_ID, new TestSplitRemote(HostAddress.fromString("10.0.0.1:12")), Lifespan.taskWide()));
    // splits now contains 41 splits : 21 with node1 as local node and 20 with node2 as local node
    Multimap<InternalNode, Split> assignments3 = nodeSelector.computeAssignments(splits, ImmutableList.copyOf(taskMap.values())).getAssignments();
    // Check that only 40 splits are being assigned as there is a single task
    assertEquals(assignments3.size(), 40);
    assertEquals(assignments3.keySet().size(), 2);
    assertTrue(assignments3.keySet().contains(node1));
    assertTrue(assignments3.keySet().contains(node2));
    // The first 20 splits have node1 as local, the next 19 have node2 as local, the 40th split has node1 as local and the 41st has node2 as local
    // If optimized-local-scheduling is disabled, the 41st split will be unassigned (the last slot in node2 will be taken up by the 40th split with node1 as local)
    // optimized-local-scheduling ensures that all splits that can be assigned locally will be assigned first
    node1Splits = assignments3.values().stream().map(Split::getConnectorSplit).filter(TestSplitLocal.class::isInstance).count();
    assertEquals(node1Splits, 20);
    node2Splits = assignments3.values().stream().map(Split::getConnectorSplit).filter(TestSplitRemote.class::isInstance).count();
    assertEquals(node2Splits, 20);
}

Example 22 with InternalNode

use of io.trino.metadata.InternalNode in project trino by trinodb.

the class TestNodeScheduler method testRedistributeSplit.

@Test
public void testRedistributeSplit() {
    InternalNode node1 = new InternalNode("node1", URI.create("http://10.0.0.1:11"), NodeVersion.UNKNOWN, false);
    nodeManager.addNode(CONNECTOR_ID, node1);
    InternalNode node2 = new InternalNode("node2", URI.create("http://10.0.0.1:12"), NodeVersion.UNKNOWN, false);
    nodeManager.addNode(CONNECTOR_ID, node2);
    Multimap<InternalNode, Split> assignment = HashMultimap.create();
    Set<Split> splitsAssignedToNode1 = new LinkedHashSet<>();
    // Node1 to be assigned 12 splits out of which 6 are local to it
    for (int i = 0; i < 6; i++) {
        splitsAssignedToNode1.add(new Split(CONNECTOR_ID, new TestSplitLocal(), Lifespan.taskWide()));
        splitsAssignedToNode1.add(new Split(CONNECTOR_ID, new TestSplitRemote(), Lifespan.taskWide()));
    }
    for (Split split : splitsAssignedToNode1) {
        assignment.put(node1, split);
    }
    Set<Split> splitsAssignedToNode2 = new LinkedHashSet<>();
    // Node2 to be assigned 10 splits
    for (int i = 0; i < 10; i++) {
        splitsAssignedToNode2.add(new Split(CONNECTOR_ID, new TestSplitRemote(), Lifespan.taskWide()));
    }
    for (Split split : splitsAssignedToNode2) {
        assignment.put(node2, split);
    }
    assertEquals(assignment.get(node1).size(), 12);
    assertEquals(assignment.get(node2).size(), 10);
    ImmutableSetMultimap.Builder<InetAddress, InternalNode> nodesByHost = ImmutableSetMultimap.builder();
    try {
        nodesByHost.put(node1.getInternalAddress(), node1);
        nodesByHost.put(node2.getInternalAddress(), node2);
    } catch (UnknownHostException e) {
        System.out.println("Could not convert the address");
    }
    // Redistribute 1 split from Node 1 to Node 2
    UniformNodeSelector.redistributeSplit(assignment, node1, node2, nodesByHost.build());
    assertEquals(assignment.get(node1).size(), 11);
    assertEquals(assignment.get(node2).size(), 11);
    Set<Split> redistributedSplit = Sets.difference(new HashSet<>(assignment.get(node2)), splitsAssignedToNode2);
    assertEquals(redistributedSplit.size(), 1);
    // Assert that the redistributed split is not a local split in Node 1. This test ensures that redistributeSingleSplit() prioritizes the transfer of a non-local split
    assertTrue(redistributedSplit.iterator().next().getConnectorSplit() instanceof TestSplitRemote);
}

Example 23 with InternalNode

use of io.trino.metadata.InternalNode in project trino by trinodb.

the class TestNodeScheduler method testMaxUnacknowledgedSplitsPerTask.

@Test
public void testMaxUnacknowledgedSplitsPerTask() {
    int maxUnacknowledgedSplitsPerTask = 5;
    nodeSelector = nodeScheduler.createNodeSelector(sessionWithMaxUnacknowledgedSplitsPerTask(maxUnacknowledgedSplitsPerTask), Optional.of(CONNECTOR_ID));
    setUpNodes();
    ImmutableList.Builder<Split> initialSplits = ImmutableList.builder();
    for (int i = 0; i < maxUnacknowledgedSplitsPerTask; i++) {
        initialSplits.add(new Split(CONNECTOR_ID, new TestSplitRemote(), Lifespan.taskWide()));
    }
    List<InternalNode> nodes = new ArrayList<>();
    List<MockRemoteTaskFactory.MockRemoteTask> tasks = new ArrayList<>();
    MockRemoteTaskFactory remoteTaskFactory = new MockRemoteTaskFactory(remoteTaskExecutor, remoteTaskScheduledExecutor);
    int counter = 1;
    for (InternalNode node : nodeManager.getActiveConnectorNodes(CONNECTOR_ID)) {
        // Max out number of unacknowledged splits on each task
        TaskId taskId = new TaskId(new StageId("test", 1), counter, 0);
        counter++;
        MockRemoteTaskFactory.MockRemoteTask remoteTask = remoteTaskFactory.createTableScanTask(taskId, node, initialSplits.build(), nodeTaskMap.createPartitionedSplitCountTracker(node, taskId));
        nodeTaskMap.addTask(node, remoteTask);
        remoteTask.setMaxUnacknowledgedSplits(maxUnacknowledgedSplitsPerTask);
        remoteTask.setUnacknowledgedSplits(maxUnacknowledgedSplitsPerTask);
        nodes.add(node);
        tasks.add(remoteTask);
    }
    // One split per node
    Set<Split> splits = new HashSet<>();
    for (int i = 0; i < nodes.size(); i++) {
        splits.add(new Split(CONNECTOR_ID, new TestSplitRemote(), Lifespan.taskWide()));
    }
    SplitPlacementResult splitPlacements = nodeSelector.computeAssignments(splits, ImmutableList.copyOf(tasks));
    // No splits should have been placed, max unacknowledged was already reached
    assertEquals(splitPlacements.getAssignments().size(), 0);
    // Unblock one task
    MockRemoteTaskFactory.MockRemoteTask taskOne = tasks.get(0);
    taskOne.finishSplits(1);
    taskOne.setUnacknowledgedSplits(taskOne.getUnacknowledgedPartitionedSplitCount() - 1);
    assertTrue(splitPlacements.getBlocked().isDone());
    // Attempt to schedule again, only the node with the unblocked task should be chosen
    splitPlacements = nodeSelector.computeAssignments(splits, ImmutableList.copyOf(tasks));
    assertEquals(splitPlacements.getAssignments().size(), 1);
    assertTrue(splitPlacements.getAssignments().keySet().contains(nodes.get(0)));
    // Make the first node appear to have no splits, unacknowledged splits alone should force the splits to be spread across nodes
    taskOne.clearSplits();
    // Give all tasks with room for 1 unacknowledged split
    tasks.forEach(task -> task.setUnacknowledgedSplits(maxUnacknowledgedSplitsPerTask - 1));
    splitPlacements = nodeSelector.computeAssignments(splits, ImmutableList.copyOf(tasks));
    // One split placed on each node
    assertEquals(splitPlacements.getAssignments().size(), nodes.size());
    assertTrue(splitPlacements.getAssignments().keySet().containsAll(nodes));
}

Example 24 with InternalNode

use of io.trino.metadata.InternalNode in project trino by trinodb.

the class TestNodeScheduler method testEmptyAssignmentWithFullNodes.

@Test
public void testEmptyAssignmentWithFullNodes() {
    InternalNode node1 = new InternalNode("node1", URI.create("http://10.0.0.1:11"), NodeVersion.UNKNOWN, false);
    nodeManager.addNode(CONNECTOR_ID, node1);
    InternalNode node2 = new InternalNode("node2", URI.create("http://10.0.0.1:12"), NodeVersion.UNKNOWN, false);
    nodeManager.addNode(CONNECTOR_ID, node2);
    Set<Split> splits = new LinkedHashSet<>();
    // 20 splits with node1 as local node to be assigned in the first iteration of computeAssignments
    for (int i = 0; i < (20 + 10 + 5) * 2; i++) {
        splits.add(new Split(CONNECTOR_ID, new TestSplitLocal(), Lifespan.taskWide()));
    }
    // computeAssignments just returns a mapping of nodes with splits to be assigned, it does not assign splits
    Multimap<InternalNode, Split> assignments1 = nodeSelector.computeAssignments(splits, ImmutableList.copyOf(taskMap.values())).getAssignments();
    assertEquals(assignments1.size(), 40);
    assertEquals(assignments1.keySet().size(), 2);
    assertEquals(assignments1.get(node1).size(), 20);
    assertEquals(assignments1.get(node2).size(), 20);
    MockRemoteTaskFactory remoteTaskFactory = new MockRemoteTaskFactory(remoteTaskExecutor, remoteTaskScheduledExecutor);
    int task = 0;
    for (InternalNode node : assignments1.keySet()) {
        TaskId taskId = new TaskId(new StageId("test", 1), task, 0);
        task++;
        MockRemoteTaskFactory.MockRemoteTask remoteTask = remoteTaskFactory.createTableScanTask(taskId, node, ImmutableList.copyOf(assignments1.get(node)), nodeTaskMap.createPartitionedSplitCountTracker(node, taskId));
        remoteTask.startSplits(20);
        nodeTaskMap.addTask(node, remoteTask);
        taskMap.put(node, remoteTask);
    }
    Set<Split> unassignedSplits = Sets.difference(splits, new HashSet<>(assignments1.values()));
    assertEquals(unassignedSplits.size(), 30);
    Multimap<InternalNode, Split> assignments2 = nodeSelector.computeAssignments(unassignedSplits, ImmutableList.copyOf(taskMap.values())).getAssignments();
    for (InternalNode node : assignments2.keySet()) {
        RemoteTask remoteTask = taskMap.get(node);
        remoteTask.addSplits(ImmutableMultimap.<PlanNodeId, Split>builder().putAll(new PlanNodeId("sourceId"), assignments2.get(node)).build());
    }
    unassignedSplits = Sets.difference(unassignedSplits, new HashSet<>(assignments2.values()));
    assertEquals(unassignedSplits.size(), 10);
    Multimap<InternalNode, Split> assignments3 = nodeSelector.computeAssignments(unassignedSplits, ImmutableList.copyOf(taskMap.values())).getAssignments();
    assertTrue(assignments3.isEmpty());
}

Example 25 with InternalNode

use of io.trino.metadata.InternalNode in project trino by trinodb.

the class TestFullNodeCapableNodeAllocator method testRemoveAssignedFullNode.

@Test(timeOut = TEST_TIMEOUT)
public void testRemoveAssignedFullNode() throws Exception {
    InMemoryNodeManager nodeManager = testingNodeManager(basicNodesMap(NODE_1, NODE_2));
    setupNodeAllocatorService(nodeManager, 1);
    try (NodeAllocator nodeAllocator = nodeAllocatorService.getNodeAllocator(Q1_SESSION)) {
        NodeAllocator.NodeLease sharedAcquire1 = nodeAllocator.acquire(NO_REQUIREMENTS);
        assertAcquired(sharedAcquire1);
        NodeAllocator.NodeLease sharedAcquire2 = nodeAllocator.acquire(NO_REQUIREMENTS);
        assertAcquired(sharedAcquire2);
        InternalNode nodeAcquired1 = sharedAcquire1.getNode().get().getNode();
        InternalNode nodeAcquired2 = sharedAcquire2.getNode().get().getNode();
        assertNotEquals(nodeAcquired1, nodeAcquired2);
        // try to acquire full node; should not happen
        NodeAllocator.NodeLease fullAcquire = nodeAllocator.acquire(FULL_NODE_REQUIREMENTS);
        assertNotAcquired(fullAcquire);
        Set<InternalNode> pendingFullNodes = nodeAllocatorService.getPendingFullNodes();
        InternalNode pendingFullNode = Iterables.getOnlyElement(pendingFullNodes);
        // remove assigned node and release shared allocation for it; full node acquire still should not be fulfilled
        nodeManager.removeNode(pendingFullNode);
        sharedAcquire1.release();
        assertNotAcquired(fullAcquire);
        // release remaining node in the cluster
        sharedAcquire2.release();
        // full node should be fulfilled now
        assertEventually(() -> {
            // we need to wait as pending acquires are processed asynchronously
            assertAcquired(fullAcquire, nodeAcquired2);
        });
    }
}