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

Example 21 with InMemoryNodeManager

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

the class TestSourcePartitionedScheduler method testWorkerBalancedSplitAssignment.

@Test
public void testWorkerBalancedSplitAssignment() {
    // use private node manager so we can add a node later
    InMemoryNodeManager nodeManager = new InMemoryNodeManager();
    nodeManager.addNode(CONNECTOR_ID, new InternalNode("other1", URI.create("http://127.0.0.1:11"), NodeVersion.UNKNOWN, false), new InternalNode("other2", URI.create("http://127.0.0.1:12"), NodeVersion.UNKNOWN, false), new InternalNode("other3", URI.create("http://127.0.0.1:13"), NodeVersion.UNKNOWN, false));
    NodeTaskMap nodeTaskMap = new NodeTaskMap(finalizerService);
    // Schedule 15 splits - there are 3 nodes, each node should get 5 splits
    PlanFragment firstPlan = createFragment();
    StageExecution firstStage = createStageExecution(firstPlan, nodeTaskMap);
    StageScheduler firstScheduler = getSourcePartitionedScheduler(createFixedSplitSource(15, TestingSplit::createRemoteSplit), firstStage, nodeManager, nodeTaskMap, 200, NODE);
    ScheduleResult scheduleResult = firstScheduler.schedule();
    assertEffectivelyFinished(scheduleResult, firstScheduler);
    assertTrue(scheduleResult.getBlocked().isDone());
    assertEquals(scheduleResult.getNewTasks().size(), 3);
    assertEquals(firstStage.getAllTasks().size(), 3);
    for (RemoteTask remoteTask : firstStage.getAllTasks()) {
        PartitionedSplitsInfo splitsInfo = remoteTask.getPartitionedSplitsInfo();
        assertEquals(splitsInfo.getCount(), 5);
    }
    // Add new node
    InternalNode additionalNode = new InternalNode("other4", URI.create("http://127.0.0.1:14"), NodeVersion.UNKNOWN, false);
    nodeManager.addNode(CONNECTOR_ID, additionalNode);
    // Schedule 5 splits in another query. Since the new node does not have any splits, all 5 splits are assigned to the new node
    PlanFragment secondPlan = createFragment();
    StageExecution secondStage = createStageExecution(secondPlan, nodeTaskMap);
    StageScheduler secondScheduler = getSourcePartitionedScheduler(createFixedSplitSource(5, TestingSplit::createRemoteSplit), secondStage, nodeManager, nodeTaskMap, 200, NODE);
    scheduleResult = secondScheduler.schedule();
    assertEffectivelyFinished(scheduleResult, secondScheduler);
    assertTrue(scheduleResult.getBlocked().isDone());
    assertEquals(scheduleResult.getNewTasks().size(), 1);
    assertEquals(secondStage.getAllTasks().size(), 1);
    RemoteTask task = secondStage.getAllTasks().get(0);
    assertEquals(task.getPartitionedSplitsInfo().getCount(), 5);
    firstStage.abort();
    secondStage.abort();
}

Example 22 with InMemoryNodeManager

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

the class TestFullNodeCapableNodeAllocator method testAllocateFullNodeReleaseBeforeAcquiredWaitingOnOtherNodesUsed.

@Test(timeOut = TEST_TIMEOUT)
public void testAllocateFullNodeReleaseBeforeAcquiredWaitingOnOtherNodesUsed() throws Exception {
    InMemoryNodeManager nodeManager = testingNodeManager(basicNodesMap(NODE_1));
    setupNodeAllocatorService(nodeManager, 100);
    try (NodeAllocator nodeAllocator = nodeAllocatorService.getNodeAllocator(Q1_SESSION)) {
        // allocate NODE_1 in shared mode
        NodeAllocator.NodeLease acquire1 = nodeAllocator.acquire(NO_REQUIREMENTS);
        assertAcquired(acquire1, NODE_1);
        // add one more node
        addNode(nodeManager, NODE_2);
        // first full allocation should not block
        NodeAllocator.NodeLease acquire2 = nodeAllocator.acquire(FULL_NODE_REQUIREMENTS);
        assertAcquired(acquire2, NODE_2);
        // next two should block (all nodes used)
        NodeAllocator.NodeLease acquire3 = nodeAllocator.acquire(FULL_NODE_REQUIREMENTS);
        assertNotAcquired(acquire3);
        NodeAllocator.NodeLease acquire4 = nodeAllocator.acquire(FULL_NODE_REQUIREMENTS);
        assertNotAcquired(acquire4);
        // releasing a blocked one should not unblock anything
        acquire3.release();
        assertNotAcquired(acquire4);
        // releasing node acquired in shared move one should unblock one which is still blocked
        acquire1.release();
        assertEventually(() -> assertAcquired(acquire4, NODE_1));
    }
}

Example 23 with InMemoryNodeManager

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

the class TestFullNodeCapableNodeAllocator method testAllocateFullWithCatalogRequirements.

@Test(timeOut = TEST_TIMEOUT)
public void testAllocateFullWithCatalogRequirements() throws Exception {
    InMemoryNodeManager nodeManager = testingNodeManager(nodesMapBuilder().put(NODE_1, ImmutableList.of(CATALOG_1)).put(NODE_2, ImmutableList.of(CATALOG_1)).put(NODE_3, ImmutableList.of(CATALOG_2)).buildOrThrow());
    setupNodeAllocatorService(nodeManager, 2);
    try (NodeAllocator nodeAllocator = nodeAllocatorService.getNodeAllocator(Q1_SESSION)) {
        // we have 3 nodes available and per-query limit set to 2 but only 1 node that exposes CATALOG_2
        NodeAllocator.NodeLease acquire1 = nodeAllocator.acquire(FULL_NODE_CATALOG_2_REQUIREMENTS);
        assertAcquired(acquire1);
        NodeAllocator.NodeLease acquire2 = nodeAllocator.acquire(FULL_NODE_CATALOG_2_REQUIREMENTS);
        assertNotAcquired(acquire2);
        // releasing CATALOG_2 node allows pending lease to acquire it
        acquire1.release();
        assertEventually(() -> assertAcquired(acquire2));
    }
}

Example 24 with InMemoryNodeManager

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

the class TestFullNodeCapableNodeAllocator method testAllocateFullReleaseBeforeAcquired.

@Test(timeOut = TEST_TIMEOUT)
public void testAllocateFullReleaseBeforeAcquired() throws Exception {
    InMemoryNodeManager nodeManager = testingNodeManager(basicNodesMap(NODE_1));
    setupNodeAllocatorService(nodeManager, 1);
    try (NodeAllocator nodeAllocator = nodeAllocatorService.getNodeAllocator(Q1_SESSION)) {
        // first allocation should not block
        NodeAllocator.NodeLease acquire1 = nodeAllocator.acquire(FULL_NODE_REQUIREMENTS);
        assertAcquired(acquire1, NODE_1);
        // another two should block
        NodeAllocator.NodeLease acquire2 = nodeAllocator.acquire(FULL_NODE_REQUIREMENTS);
        assertNotAcquired(acquire2);
        NodeAllocator.NodeLease acquire3 = nodeAllocator.acquire(FULL_NODE_REQUIREMENTS);
        assertNotAcquired(acquire3);
        // releasing a blocked one should not unblock anything
        acquire2.release();
        assertNotAcquired(acquire3);
        // releasing one acquired one should unblock one which is still blocked
        acquire1.release();
        assertEventually(() -> assertAcquired(acquire3, NODE_1));
    }
}

Example 25 with InMemoryNodeManager

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

the class TestFullNodeCapableNodeAllocator method testAllocateFullWithQueryLimit.

@Test(timeOut = TEST_TIMEOUT)
public void testAllocateFullWithQueryLimit() throws Exception {
    InMemoryNodeManager nodeManager = testingNodeManager(basicNodesMap(NODE_1, NODE_2, NODE_3));
    setupNodeAllocatorService(nodeManager, 2);
    try (NodeAllocator q1NodeAllocator = nodeAllocatorService.getNodeAllocator(Q1_SESSION);
        NodeAllocator q2NodeAllocator = nodeAllocatorService.getNodeAllocator(Q2_SESSION)) {
        // allocate 2 full nodes for Q1 should not block
        NodeAllocator.NodeLease q1Acquire1 = q1NodeAllocator.acquire(FULL_NODE_REQUIREMENTS);
        assertAcquired(q1Acquire1);
        NodeAllocator.NodeLease q1Acquire2 = q1NodeAllocator.acquire(FULL_NODE_REQUIREMENTS);
        assertAcquired(q1Acquire2);
        // third allocation for Q1 should block even though we have 3 nodes available
        NodeAllocator.NodeLease q1Acquire3 = q1NodeAllocator.acquire(FULL_NODE_REQUIREMENTS);
        assertNotAcquired(q1Acquire3);
        // we should still be able to acquire full node for another query
        NodeAllocator.NodeLease q2Acquire1 = q2NodeAllocator.acquire(FULL_NODE_REQUIREMENTS);
        assertAcquired(q2Acquire1);
        // when we release one of the nodes for Q1 pending q1Acquire3 should unblock
        q1Acquire1.release();
        assertEventually(() -> {
            assertAcquired(q1Acquire3);
            assertEquals(q1Acquire3.getNode().get().getNode(), q1Acquire1.getNode().get().getNode());
        });
    }
}