diff --git a/presto-main/src/main/java/com/facebook/presto/connector/ConnectorAwareNodeManager.java b/presto-main/src/main/java/com/facebook/presto/connector/ConnectorAwareNodeManager.java index 6764405dca877..3506d55393f21 100644 --- a/presto-main/src/main/java/com/facebook/presto/connector/ConnectorAwareNodeManager.java +++ b/presto-main/src/main/java/com/facebook/presto/connector/ConnectorAwareNodeManager.java @@ -13,14 +13,19 @@ */ package com.facebook.presto.connector; +import com.facebook.presto.metadata.InternalNode; import com.facebook.presto.metadata.InternalNodeManager; import com.facebook.presto.spi.ConnectorId; import com.facebook.presto.spi.Node; import com.facebook.presto.spi.NodeManager; +import com.facebook.presto.spi.PrestoException; import com.google.common.collect.ImmutableSet; import java.util.Set; +import static com.facebook.presto.spi.StandardErrorCode.NO_CPP_SIDECARS; +import static com.facebook.presto.spi.StandardErrorCode.TOO_MANY_SIDECARS; +import static com.google.common.collect.Iterables.getOnlyElement; import static java.util.Objects.requireNonNull; public class ConnectorAwareNodeManager @@ -58,6 +63,19 @@ public Node getCurrentNode() return nodeManager.getCurrentNode(); } + @Override + public Node getSidecarNode() + { + Set coordinatorSidecars = nodeManager.getCoordinatorSidecars(); + if (coordinatorSidecars.isEmpty()) { + throw new PrestoException(NO_CPP_SIDECARS, "Expected exactly one coordinator sidecar, but found none"); + } + if (coordinatorSidecars.size() > 1) { + throw new PrestoException(TOO_MANY_SIDECARS, "Expected exactly one coordinator sidecar, but found " + coordinatorSidecars.size()); + } + return getOnlyElement(coordinatorSidecars); + } + @Override public String getEnvironment() { diff --git a/presto-main/src/main/java/com/facebook/presto/nodeManager/PluginNodeManager.java b/presto-main/src/main/java/com/facebook/presto/nodeManager/PluginNodeManager.java new file mode 100644 index 0000000000000..34c2a1a3ccaab --- /dev/null +++ b/presto-main/src/main/java/com/facebook/presto/nodeManager/PluginNodeManager.java @@ -0,0 +1,97 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.facebook.presto.nodeManager; + +import com.facebook.presto.metadata.InternalNode; +import com.facebook.presto.metadata.InternalNodeManager; +import com.facebook.presto.spi.Node; +import com.facebook.presto.spi.NodeManager; +import com.facebook.presto.spi.PrestoException; +import com.google.common.collect.ImmutableSet; +import com.google.inject.Inject; + +import java.util.Set; + +import static com.facebook.presto.spi.StandardErrorCode.NO_CPP_SIDECARS; +import static com.facebook.presto.spi.StandardErrorCode.TOO_MANY_SIDECARS; +import static com.google.common.collect.ImmutableSet.toImmutableSet; +import static com.google.common.collect.Iterables.getOnlyElement; +import static java.util.Objects.requireNonNull; + +/** + * This class simplifies managing Presto's cluster nodes, + * focusing on active workers and coordinators without tying to specific connectors. + */ +public class PluginNodeManager + implements NodeManager +{ + private final InternalNodeManager nodeManager; + private final String environment; + + @Inject + public PluginNodeManager(InternalNodeManager nodeManager) + { + this.nodeManager = nodeManager; + this.environment = "test"; + } + + public PluginNodeManager(InternalNodeManager nodeManager, String environment) + { + this.nodeManager = requireNonNull(nodeManager, "nodeManager is null"); + this.environment = requireNonNull(environment, "environment is null"); + } + + @Override + public Set getAllNodes() + { + return ImmutableSet.builder() + .addAll(getWorkerNodes()) + .addAll(nodeManager.getCoordinators()) + .build(); + } + + @Override + public Set getWorkerNodes() + { + //Retrieves all active worker nodes, excluding coordinators, resource managers, and catalog servers. + return nodeManager.getAllNodes().getActiveNodes().stream() + .filter(node -> !node.isResourceManager() && !node.isCoordinator() && !node.isCatalogServer()) + .collect(toImmutableSet()); + } + + @Override + public Node getCurrentNode() + { + return nodeManager.getCurrentNode(); + } + + @Override + public Node getSidecarNode() + { + Set coordinatorSidecars = nodeManager.getCoordinatorSidecars(); + if (coordinatorSidecars.isEmpty()) { + throw new PrestoException(NO_CPP_SIDECARS, "Expected exactly one coordinator sidecar, but found none"); + } + if (coordinatorSidecars.size() > 1) { + throw new PrestoException(TOO_MANY_SIDECARS, "Expected exactly one coordinator sidecar, but found " + coordinatorSidecars.size()); + } + return getOnlyElement(coordinatorSidecars); + } + + @Override + public String getEnvironment() + { + return environment; + } +} diff --git a/presto-main/src/main/java/com/facebook/presto/testing/TestingNodeManager.java b/presto-main/src/main/java/com/facebook/presto/testing/TestingNodeManager.java index d76239c726078..5b2eab906325b 100644 --- a/presto-main/src/main/java/com/facebook/presto/testing/TestingNodeManager.java +++ b/presto-main/src/main/java/com/facebook/presto/testing/TestingNodeManager.java @@ -92,6 +92,12 @@ public Node getCurrentNode() return localNode; } + @Override + public Node getSidecarNode() + { + return localNode; + } + @Override public String getEnvironment() { diff --git a/presto-main/src/test/java/com/facebook/presto/nodeManager/TestPluginNodeManager.java b/presto-main/src/test/java/com/facebook/presto/nodeManager/TestPluginNodeManager.java new file mode 100644 index 0000000000000..aa5ef70e7f3f1 --- /dev/null +++ b/presto-main/src/test/java/com/facebook/presto/nodeManager/TestPluginNodeManager.java @@ -0,0 +1,94 @@ +/* + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.facebook.presto.nodeManager; + +import com.facebook.presto.client.NodeVersion; +import com.facebook.presto.metadata.InMemoryNodeManager; +import com.facebook.presto.metadata.InternalNode; +import com.facebook.presto.spi.ConnectorId; +import com.facebook.presto.spi.Node; +import org.testng.annotations.BeforeClass; +import org.testng.annotations.Test; + +import java.net.URI; +import java.util.Arrays; +import java.util.Set; + +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertNotNull; +import static org.testng.Assert.assertTrue; + +public class TestPluginNodeManager +{ + private InMemoryNodeManager inMemoryNodeManager; + private PluginNodeManager pluginNodeManager; + + @BeforeClass + public void setUp() + { + // Initialize the InMemoryNodeManager and PluginNodeManager before each test. + inMemoryNodeManager = new InMemoryNodeManager(); + pluginNodeManager = new PluginNodeManager(inMemoryNodeManager, "test-env"); + } + + @Test + public void testGetAllNodes() + { + ConnectorId connectorId = new ConnectorId("test-connector"); + InternalNode activeNode1 = new InternalNode("activeNode1", URI.create("http://example1.com"), new NodeVersion("1"), false); + InternalNode activeNode2 = new InternalNode("activeNode2", URI.create("http://example2.com"), new NodeVersion("1"), false); + InternalNode coordinatorNode = new InternalNode("coordinatorNode", URI.create("http://example3.com"), new NodeVersion("1"), true); + + inMemoryNodeManager.addNode(connectorId, activeNode1); + inMemoryNodeManager.addNode(connectorId, activeNode2); + inMemoryNodeManager.addNode(connectorId, coordinatorNode); + + Set allNodes = pluginNodeManager.getAllNodes(); + // The expected count is 4, considering two active nodes, one coordinator, and one local node added by InMemoryNodeManager by default. + assertEquals(4, allNodes.size()); + assertTrue(allNodes.containsAll(Arrays.asList(activeNode1, activeNode2, coordinatorNode))); + } + + @Test + public void testGetWorkerNodes() + { + ConnectorId connectorId = new ConnectorId("test-connector"); + InternalNode activeNode1 = new InternalNode("activeNode1", URI.create("http://example1.com"), new NodeVersion("1"), false); + InternalNode activeNode2 = new InternalNode("activeNode2", URI.create("http://example2.com"), new NodeVersion("1"), false); + + inMemoryNodeManager.addNode(connectorId, activeNode1); + inMemoryNodeManager.addNode(connectorId, activeNode2); + + Set workerNodes = pluginNodeManager.getWorkerNodes(); + // Expected count is 3, accounting for two explicitly added active nodes and one local node. + assertEquals(3, workerNodes.size()); + assertTrue(workerNodes.containsAll(Arrays.asList(activeNode1, activeNode2))); + } + + @Test + public void testGetEnvironment() + { + // Validate that the PluginNodeManager correctly returns the environment string set during initialization. + assertEquals("test-env", pluginNodeManager.getEnvironment()); + } + + @Test + public void testGetCurrentNode() + { + Node currentNode = pluginNodeManager.getCurrentNode(); + assertNotNull(currentNode); + // Validate that the current node is not null and its identifier matches the expected local node identifier. + assertEquals("local", currentNode.getNodeIdentifier()); + } +} diff --git a/presto-spi/src/main/java/com/facebook/presto/spi/NodeManager.java b/presto-spi/src/main/java/com/facebook/presto/spi/NodeManager.java index 92b84d0187fb3..602205786503d 100644 --- a/presto-spi/src/main/java/com/facebook/presto/spi/NodeManager.java +++ b/presto-spi/src/main/java/com/facebook/presto/spi/NodeManager.java @@ -25,6 +25,8 @@ public interface NodeManager Node getCurrentNode(); + Node getSidecarNode(); + String getEnvironment(); default Set getRequiredWorkerNodes()