Custom scheduler plugin support

pom.xml

@@ -214,6 +214,7 @@
         <module>presto-native-sidecar-plugin</module>
         <module>presto-base-arrow-flight</module>
         <module>presto-function-server</module>
+        <module>presto-router-example-plugin-scheduler</module>
     </modules>
 
     <dependencyManagement>

presto-main-base/src/main/java/com/facebook/presto/server/PluginInstaller.java

@@ -15,10 +15,13 @@
 
 import com.facebook.presto.spi.CoordinatorPlugin;
 import com.facebook.presto.spi.Plugin;
+import com.facebook.presto.spi.RouterPlugin;
 
 public interface PluginInstaller
 {
     void installPlugin(Plugin plugin);
 
     void installCoordinatorPlugin(CoordinatorPlugin plugin);
+
+    default void installRouterPlugin(RouterPlugin plugin){};
 }

presto-main-base/src/main/java/com/facebook/presto/server/PluginManagerUtil.java

@@ -17,6 +17,7 @@
 import com.facebook.presto.metadata.Metadata;
 import com.facebook.presto.spi.CoordinatorPlugin;
 import com.facebook.presto.spi.Plugin;
+import com.facebook.presto.spi.RouterPlugin;
 import com.facebook.presto.spi.classloader.ThreadContextClassLoader;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.Ordering;
@@ -139,6 +140,7 @@ public static void loadPlugin(
                 pluginServicesFile,
                 parent);
         try (ThreadContextClassLoader ignored = new ThreadContextClassLoader(pluginClassLoader)) {
+            loadPlugin(pluginClassLoader, RouterPlugin.class, pluginInstaller);
             loadPlugin(pluginClassLoader, CoordinatorPlugin.class, pluginInstaller);
             loadPlugin(pluginClassLoader, Plugin.class, pluginInstaller);
         }
@@ -165,6 +167,9 @@ public static void loadPlugin(
             else if (plugin instanceof CoordinatorPlugin) {
                 pluginInstaller.installCoordinatorPlugin((CoordinatorPlugin) plugin);
             }
+            else if (plugin instanceof RouterPlugin) {
+                pluginInstaller.installRouterPlugin((RouterPlugin) plugin);
+            }
             else {
                 log.warn("Unknown plugin type: %s", plugin.getClass().getName());
             }

presto-router-example-plugin-scheduler/README.md

@@ -0,0 +1,22 @@
+# Example - Presto Router Custom Scheduler Plugin
+This package contains an example of a custom scheduler plugin.
+
+## Add the Custom Scheduler Plugin to the Presto Router
+Place the plugin jar and all the dependent jars for the plugin in the `plugin` directory relative to the Presto install directory.
+
+The configuration file for this plugin must be `etc/router-config/router-scheduler.properties`.  
+In `router-scheduler.properties`, set the name of the custom scheduler factory.  
+For example, use the following line to set the custom scheduler factory name to `metricsBased`.  
+``router-scheduler.name=metricsBased``
+
+The scheduler name must be set to `CUSTOM_PLUGIN_SCHEDULER` in `etc/router-config.json`.  
+    ``scheduler``: ``CUSTOM_PLUGIN_SCHEDULER``
+
+## Main Classes:
+* `RouterSchedulerPlugin` - Custom Scheduler Plugin class to be loaded by the Router plugin manager.  
+  This class implements the interface `com.facebook.presto.spi.RouterPlugin`.
+* `MetricsBasedSchedulerFactory` - Factory for creating specific custom scheduler.  
+  This class implements the interface `com.facebook.presto.spi.SchedulerFactory`.
+* `MetricsBasedScheduler` - Example custom scheduler implementing the scheduling logic for clusters.  
+  This class implements the interface `com.facebook.presto.spi.router.Scheduler`.  
+  Similar classes can be added implementing specific custom scheduling logic.

presto-router-example-plugin-scheduler/pom.xml

@@ -0,0 +1,71 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<project xmlns="http://maven.apache.org/POM/4.0.0" xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd">
+    <modelVersion>4.0.0</modelVersion>
+
+    <parent>
+        <groupId>com.facebook.presto</groupId>
+        <artifactId>presto-root</artifactId>
+        <version>0.293-SNAPSHOT</version>
+    </parent>
+
+    <groupId>com.facebook.presto.router</groupId>
+    <artifactId>presto-router-example-plugin-scheduler</artifactId>
+    <version>0.293-SNAPSHOT</version>
+    <packaging>jar</packaging>
+    <name>presto-router-example-plugin-scheduler</name>
+    <description>Presto-router example custom plugin scheduler</description>
+
+    <properties>
+        <air.main.basedir>${project.parent.basedir}</air.main.basedir>
+    </properties>
+
+    <dependencies>
+        <dependency>
+            <groupId>com.facebook.presto</groupId>
+            <artifactId>presto-spi</artifactId>
+            <scope>provided</scope>
+        </dependency>
+
+        <dependency>
+            <groupId>com.google.guava</groupId>
+            <artifactId>guava</artifactId>
+        </dependency>
+
+        <!-- test dependencies-->
+        <dependency>
+            <groupId>org.testng</groupId>
+            <artifactId>testng</artifactId>
+            <scope>test</scope>
+        </dependency>
+    </dependencies>
+
+    <!-- Build configurations -->
+    <build>
+        <plugins>
+            <plugin>
+                <groupId>com.facebook.presto</groupId>
+                <artifactId>presto-maven-plugin</artifactId>
+                <executions>
+                    <execution>
+                        <phase>validate</phase>
+                        <goals>
+                            <goal>check-spi-dependencies</goal>
+                        </goals>
+                    </execution>
+                </executions>
+            </plugin>
+            <plugin>
+                <groupId>ca.vanzyl.provisio.maven.plugins</groupId>
+                <artifactId>provisio-maven-plugin</artifactId>
+                <executions>
+                    <execution>
+                        <phase>package</phase>
+                        <goals>
+                            <goal>provision</goal>
+                        </goals>
+                    </execution>
+                </executions>
+            </plugin>
+        </plugins>
+    </build>
+</project>

presto-router-example-plugin-scheduler/src/main/java/com/facebook/presto/router/scheduler/MetricsBasedScheduler.java

@@ -0,0 +1,63 @@
+/*
+ * 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.router.scheduler;
+
+import com.facebook.presto.spi.router.ClusterInfo;
+import com.facebook.presto.spi.router.Scheduler;
+
+import java.net.URI;
+import java.util.List;
+import java.util.Map;
+import java.util.Optional;
+
+import static java.util.Comparator.comparingLong;
+
+/**
+ * Metrics based scheduler uses coordinator and/or worker metrics for scheduling decisions.
+ */
+public class MetricsBasedScheduler
+        implements Scheduler
+{
+    private Map<URI, ClusterInfo> clusterInfos;
+
+    @Override
+    public Optional<URI> getDestination(String user)
+    {
+        return Optional.empty();
+    }
+
+    /**
+     * Returns the destination cluster URI with the fewest number of active and queued queries.
+     */
+    @Override
+    public Optional<URI> getDestination(String user, String query)
+    {
+        if (clusterInfos != null && !clusterInfos.isEmpty()) {
+            //Cluster with the fewest number of queries will be returned
+            return clusterInfos.entrySet().stream()
+                    .min(comparingLong(entry -> entry.getValue().getRunningQueries() + entry.getValue().getQueuedQueries()))
+                    .map(Map.Entry::getKey);
+        }
+        return Optional.empty();
+    }
+
+    @Override
+    public void setCandidates(List<URI> candidates) {}
+
+    @Override
+    public void setClusterInfos(Map<URI, ClusterInfo> clusterInfos)
+    {
+        this.clusterInfos = clusterInfos;
+    }
+}

presto-router-example-plugin-scheduler/src/main/java/com/facebook/presto/router/scheduler/MetricsBasedSchedulerFactory.java

@@ -0,0 +1,37 @@
+/*
+ * 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.router.scheduler;
+
+import com.facebook.presto.spi.router.Scheduler;
+import com.facebook.presto.spi.router.SchedulerFactory;
+
+import java.util.Map;
+
+public class MetricsBasedSchedulerFactory
+        implements SchedulerFactory
+{
+    public static final String METRICS_BASED = "metricsBased";
+
+    @Override
+    public String getName()
+    {
+        return METRICS_BASED;
+    }
+
+    @Override
+    public Scheduler create(Map<String, String> config)
+    {
+        return new MetricsBasedScheduler();
+    }
+}

presto-router-example-plugin-scheduler/src/main/java/com/facebook/presto/router/scheduler/RouterSchedulerPlugin.java

@@ -0,0 +1,28 @@
+/*
+ * 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.router.scheduler;
+
+import com.facebook.presto.spi.RouterPlugin;
+import com.facebook.presto.spi.router.SchedulerFactory;
+import com.google.common.collect.ImmutableList;
+
+public class RouterSchedulerPlugin
+        implements RouterPlugin
+{
+    @Override
+    public Iterable<SchedulerFactory> getSchedulerFactories()
+    {
+        return ImmutableList.of(new MetricsBasedSchedulerFactory());
+    }
+}

presto-router-example-plugin-scheduler/src/main/provisio/plugin.xml

@@ -0,0 +1,4 @@
+<assembly>
+    <artifactSet to="/" ref="runtime.classpath" />
+    <archive name="${project.artifactId}-${project.version}.zip" />
+</assembly>

presto-router-example-plugin-scheduler/src/main/resources/META-INF/services/com.facebook.presto.spi.RouterPlugin

@@ -0,0 +1 @@
+com.facebook.presto.router.scheduler.RouterSchedulerPlugin

presto-router-example-plugin-scheduler/src/test/java/com/facebook/presto/router/scheduler/TestMetricsBasedScheduler.java

@@ -0,0 +1,109 @@
+/*
+ * 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.router.scheduler;
+
+import com.facebook.presto.spi.router.ClusterInfo;
+import com.facebook.presto.spi.router.Scheduler;
+import org.testng.annotations.Test;
+
+import java.net.URI;
+import java.util.HashMap;
+import java.util.Map;
+
+import static org.testng.Assert.assertEquals;
+
+public class TestMetricsBasedScheduler
+{
+    private final Map<URI, ClusterInfo> clusterInfos = new HashMap<>();
+
+    private static class MockRemoteClusterInfo
+            implements ClusterInfo
+    {
+        private final long runningQueries;
+        private final long queuedQueries;
+
+        MockRemoteClusterInfo(long runningQueries, long queuedQueries)
+        {
+            this.runningQueries = runningQueries;
+            this.queuedQueries = queuedQueries;
+        }
+
+        @Override
+        public long getRunningQueries()
+        {
+            return runningQueries;
+        }
+
+        @Override
+        public long getQueuedQueries()
+        {
+            return queuedQueries;
+        }
+
+        @Override
+        public long getBlockedQueries()
+        {
+            return 0;
+        }
+
+        @Override
+        public long getActiveWorkers()
+        {
+            return 0;
+        }
+
+        @Override
+        public long getRunningDrivers()
+        {
+            return 0;
+        }
+    }
+
+    @Test
+    public void testMetricsBasedScheduler()
+            throws Exception
+    {
+        Scheduler scheduler = new MetricsBasedScheduler();
+
+        URI uri1 = new URI("192.168.0.1");
+        URI uri2 = new URI("192.168.0.2");
+        URI uri3 = new URI("192.168.0.3");
+
+        clusterInfos.put(uri1, new MockRemoteClusterInfo(10, 10));
+        clusterInfos.put(uri2, new MockRemoteClusterInfo(20, 20));
+        clusterInfos.put(uri3, new MockRemoteClusterInfo(30, 30));
+        scheduler.setClusterInfos(clusterInfos);
+        URI target = scheduler.getDestination("test", null).orElseThrow(AssertionError::new);
+        assertEquals(target, uri1);
+
+        clusterInfos.put(uri1, new MockRemoteClusterInfo(20, 20));
+        clusterInfos.put(uri2, new MockRemoteClusterInfo(10, 10));
+        clusterInfos.put(uri3, new MockRemoteClusterInfo(30, 30));
+        scheduler.setClusterInfos(clusterInfos);
+        target = scheduler.getDestination("test", null).orElseThrow(AssertionError::new);
+        assertEquals(target, uri2);
+
+        clusterInfos.put(uri1, new MockRemoteClusterInfo(20, 20));
+        clusterInfos.put(uri2, new MockRemoteClusterInfo(30, 30));
+        clusterInfos.put(uri3, new MockRemoteClusterInfo(10, 10));
+        scheduler.setClusterInfos(clusterInfos);
+        target = scheduler.getDestination("test", null).orElseThrow(AssertionError::new);
+        assertEquals(target, uri3);
+
+        scheduler.setClusterInfos(new HashMap<>());
+        target = scheduler.getDestination("test", null).orElse(new URI("invalid"));
+        assertEquals(target, new URI("invalid"));
+    }
+}

presto-router/etc/router-config/router-scheduler.properties

@@ -0,0 +1 @@
+router-scheduler.name=metricsBased