Add sidecars to NodeManager and add PluginNodeManager

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<InternalNode> 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()
     {

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<Node> getAllNodes()
+    {
+        return ImmutableSet.<Node>builder()
+                .addAll(getWorkerNodes())
+                .addAll(nodeManager.getCoordinators())
+                .build();
+    }
+
+    @Override
+    public Set<Node> 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<InternalNode> 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;
+    }
+}

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()
     {

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<Node> 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<Node> 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());
+    }
+}

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<Node> getRequiredWorkerNodes()