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

Example 1 with InMemoryNodeManager

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

the class TestTableScanNodePartitioning method createLocalQueryRunner.

@Override
protected LocalQueryRunner createLocalQueryRunner() {
    Session.SessionBuilder sessionBuilder = testSessionBuilder().setCatalog(MOCK_CATALOG).setSchema(TEST_SCHEMA);
    LocalQueryRunner queryRunner = LocalQueryRunner.builder(sessionBuilder.build()).withNodeCountForStats(10).build();
    queryRunner.createCatalog(MOCK_CATALOG, createMockFactory(), ImmutableMap.of());
    queryRunner.getNodePartitioningManager().addPartitioningProvider(new CatalogName(MOCK_CATALOG), new TestPartitioningProvider(new InMemoryNodeManager()));
    return queryRunner;
}

Example 2 with InMemoryNodeManager

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

the class TaskTestUtils method createTestingPlanner.

public static LocalExecutionPlanner createTestingPlanner() {
    PageSourceManager pageSourceManager = new PageSourceManager();
    pageSourceManager.addConnectorPageSourceProvider(CONNECTOR_ID, new TestingPageSourceProvider());
    // we don't start the finalizer so nothing will be collected, which is ok for a test
    FinalizerService finalizerService = new FinalizerService();
    BlockTypeOperators blockTypeOperators = new BlockTypeOperators(PLANNER_CONTEXT.getTypeOperators());
    NodeScheduler nodeScheduler = new NodeScheduler(new UniformNodeSelectorFactory(new InMemoryNodeManager(), new NodeSchedulerConfig().setIncludeCoordinator(true), new NodeTaskMap(finalizerService)));
    NodePartitioningManager nodePartitioningManager = new NodePartitioningManager(nodeScheduler, blockTypeOperators);
    PageFunctionCompiler pageFunctionCompiler = new PageFunctionCompiler(PLANNER_CONTEXT.getFunctionManager(), 0);
    return new LocalExecutionPlanner(PLANNER_CONTEXT, createTestingTypeAnalyzer(PLANNER_CONTEXT), Optional.empty(), pageSourceManager, new IndexManager(), nodePartitioningManager, new PageSinkManager(), new MockDirectExchangeClientSupplier(), new ExpressionCompiler(PLANNER_CONTEXT.getFunctionManager(), pageFunctionCompiler), pageFunctionCompiler, new JoinFilterFunctionCompiler(PLANNER_CONTEXT.getFunctionManager()), new IndexJoinLookupStats(), new TaskManagerConfig(), new GenericSpillerFactory((types, spillContext, memoryContext) -> {
        throw new UnsupportedOperationException();
    }), (types, spillContext, memoryContext) -> {
        throw new UnsupportedOperationException();
    }, (types, partitionFunction, spillContext, memoryContext) -> {
        throw new UnsupportedOperationException();
    }, new PagesIndex.TestingFactory(false), new JoinCompiler(PLANNER_CONTEXT.getTypeOperators()), new TrinoOperatorFactories(), new OrderingCompiler(PLANNER_CONTEXT.getTypeOperators()), new DynamicFilterConfig(), blockTypeOperators, new TableExecuteContextManager(), new ExchangeManagerRegistry(new ExchangeHandleResolver()));
}

Example 3 with InMemoryNodeManager

use of io.trino.metadata.InMemoryNodeManager 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);
        });
    }
}

Example 4 with InMemoryNodeManager

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

the class TestFullNodeCapableNodeAllocator method testNoFullNodeAvailable.

@Test(timeOut = TEST_TIMEOUT)
public void testNoFullNodeAvailable() throws Exception {
    InMemoryNodeManager nodeManager = testingNodeManager(nodesMapBuilder().put(NODE_1, ImmutableList.of(CATALOG_2)).buildOrThrow());
    setupNodeAllocatorService(nodeManager, 100);
    try (NodeAllocator nodeAllocator = nodeAllocatorService.getNodeAllocator(Q1_SESSION)) {
        // request a full node with specific catalog (not present)
        assertThatThrownBy(() -> nodeAllocator.acquire(FULL_NODE_CATALOG_1_REQUIREMENTS)).hasMessage("No nodes available to run query");
        // add node with specific catalog
        addNode(nodeManager, NODE_2, CATALOG_1);
        // we should be able to acquire the node now
        NodeAllocator.NodeLease acquire1 = nodeAllocator.acquire(FULL_NODE_CATALOG_1_REQUIREMENTS);
        assertAcquired(acquire1, NODE_2);
        // acquiring one more should block (all nodes with catalog already used)
        NodeAllocator.NodeLease acquire2 = nodeAllocator.acquire(FULL_NODE_CATALOG_1_REQUIREMENTS);
        assertNotAcquired(acquire2);
        // remove node with catalog
        nodeManager.removeNode(NODE_2);
        // TODO: make FullNodeCapableNodeAllocatorService react on node removed automatically
        nodeAllocatorService.wakeupProcessPendingAcquires();
        // pending acquire2 should be completed now but with an exception
        assertEventually(() -> {
            assertFalse(acquire2.getNode().isCancelled());
            assertTrue(acquire2.getNode().isDone());
            assertThatThrownBy(() -> getFutureValue(acquire2.getNode())).hasMessage("No nodes available to run query");
        });
    }
}

Example 5 with InMemoryNodeManager

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

the class TestFullNodeCapableNodeAllocator method testAllocateFullSimple.

private void testAllocateFullSimple(NodeRequirements fullNodeRequirements) throws Exception {
    InMemoryNodeManager nodeManager = testingNodeManager(basicNodesMap(NODE_1, NODE_2));
    setupNodeAllocatorService(nodeManager, 3);
    try (NodeAllocator nodeAllocator = nodeAllocatorService.getNodeAllocator(Q1_SESSION)) {
        // allocate 2 full nodes should not block
        NodeAllocator.NodeLease acquire1 = nodeAllocator.acquire(fullNodeRequirements);
        assertAcquired(acquire1);
        NodeAllocator.NodeLease acquire2 = nodeAllocator.acquire(fullNodeRequirements);
        assertAcquired(acquire2);
        // trying to allocate third full node should block
        NodeAllocator.NodeLease acquire3 = nodeAllocator.acquire(fullNodeRequirements);
        assertNotAcquired(acquire3);
        // third acquisition should unblock if one of old ones is released
        acquire1.release();
        assertEventually(() -> {
            assertAcquired(acquire3);
            assertEquals(acquire3.getNode().get().getNode(), acquire1.getNode().get().getNode());
        });
        // both nodes are used exclusively so we should no be able to acquire shared node
        NodeAllocator.NodeLease acquire4 = nodeAllocator.acquire(NO_REQUIREMENTS);
        assertNotAcquired(acquire4);
        // shared acquisition should unblock if one of full ones is released
        acquire2.release();
        assertEventually(() -> {
            assertAcquired(acquire4);
            assertEquals(acquire4.getNode().get().getNode(), acquire2.getNode().get().getNode());
        });
        // shared acquisition should block full acquisition
        NodeAllocator.NodeLease acquire5 = nodeAllocator.acquire(fullNodeRequirements);
        assertNotAcquired(acquire5);
        // and when shared acquisition is gone full node should be acquired
        acquire4.release();
        assertEventually(() -> {
            assertAcquired(acquire5);
            assertEquals(acquire5.getNode().get().getNode(), acquire4.getNode().get().getNode());
        });
    }
}