Metrics based custom scheduler plugin

presto-main/src/main/java/com/facebook/presto/server/PluginClassLoader.java

@@ -26,7 +26,7 @@
 
 import static java.util.Objects.requireNonNull;
 
-class PluginClassLoader
+public class PluginClassLoader
         extends URLClassLoader
 {
     private static final ClassLoader PLATFORM_CLASS_LOADER = findPlatformClassLoader();

presto-main/src/main/java/com/facebook/presto/server/PluginDiscovery.java

@@ -38,7 +38,7 @@
 import static java.nio.file.Files.walkFileTree;
 
 // This is a hack for development and does not support nested classes.
-final class PluginDiscovery
+public final class PluginDiscovery
 {
     private static final String CLASS_FILE_SUFFIX = ".class";
 

presto-router-custom-scheduler/README.txt

@@ -0,0 +1,17 @@
+# Presto Router Custom Scheduler Plugin
+This package implements third party custom schedulers to be used by Presto router
+
+## Adding the custom scheduler plugin to presto router
+The custom plugin jar that is built is placed in the presto-router/plugin/custom-scheduler directory,
+to be picked up by the presto-router
+The scheduler name has to be set to CUSTOM_PLUGIN_SCHEDULER in the router-config.json
+    "scheduler": "CUSTOM_PLUGIN_SCHEDULER"
+The name of the custom scheduler factory needs to be set in the property file router-scheduler.properties in presto-router/etc/router-config
+
+## 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 - Custom scheduler implementing the scheduling logic for clusters. More similar classes can be added implementing specific custom scheduling logic.
+    This class implements the interface com.facebook.presto.spi.router.Scheduler

presto-router-custom-scheduler/pom.xml

@@ -0,0 +1,86 @@
+<?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.291-SNAPSHOT</version>
+    </parent>
+
+    <groupId>com.facebook.presto.router</groupId>
+    <artifactId>scheduler</artifactId>
+    <version>0.291-SNAPSHOT</version>
+    <packaging>jar</packaging>
+    <name>scheduler</name>
+    <description>SPI plugin project for presto router custom scheduler</description>
+
+    <dependencies>
+        <dependency>
+            <groupId>com.facebook.presto</groupId>
+            <artifactId>presto-spi</artifactId>
+        </dependency>
+
+        <dependency>
+            <groupId>com.facebook.airlift</groupId>
+            <artifactId>log</artifactId>
+            <version>0.215</version>
+        </dependency>
+
+        <!-- test dependencies-->
+        <dependency>
+            <groupId>org.testng</groupId>
+            <artifactId>testng</artifactId>
+            <version>7.10.2</version>
+            <scope>test</scope>
+        </dependency>
+        <dependency>
+            <groupId>org.junit.jupiter</groupId>
+            <artifactId>junit-jupiter-engine</artifactId>
+            <version>5.9.1</version> <!-- Use the latest stable version -->
+            <scope>test</scope>
+        </dependency>
+    </dependencies>
+
+    <!-- Build configurations -->
+    <build>
+        <plugins>
+            <plugin>
+                <groupId>org.apache.maven.plugins</groupId>
+                <artifactId>maven-compiler-plugin</artifactId>
+                <version>3.8.1</version>
+                <configuration>
+                    <source>11</source> 
+                    <target>11</target>
+                </configuration>
+            </plugin>
+
+            <plugin>
+                <groupId>org.apache.maven.plugins</groupId>
+                <artifactId>maven-assembly-plugin</artifactId>
+                <version>3.3.0</version>
+                <executions>
+                    <execution>
+                        <id>make-assembly</id>
+                        <phase>package</phase>
+                        <goals>
+                            <goal>single</goal>
+                        </goals>
+                        <configuration>
+                            <descriptorRefs>
+                                <descriptorRef>jar-with-dependencies</descriptorRef>
+                            </descriptorRefs>
+                            <archive>
+                                <manifest>
+                                    <addClasspath>true</addClasspath>
+                                    <mainClass>com.facebook.presto.router.scheduler.RouterSchedulerPlugin</mainClass>
+                                </manifest>
+                            </archive>
+                        </configuration>
+                    </execution>
+                </executions>
+            </plugin>
+        </plugins>
+    </build>
+
+</project>

presto-router-custom-scheduler/src/main/java/com/facebook/presto/router/scheduler/CustomSchedulerConstants.java

@@ -0,0 +1,23 @@
+/*
+ * 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;
+
+public class CustomSchedulerConstants
+{
+    private CustomSchedulerConstants()
+    {
+    }
+
+    public static final String METRICS_BASED = "metricsBased";
+}

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

@@ -0,0 +1,86 @@
+/*
+ * 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.airlift.log.Logger;
+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 java.util.PriorityQueue;
+import java.util.stream.Collectors;
+
+//Metrics based scheduler uses coordinator and/or worker metrics for scheduling decisions
+public class MetricsBasedScheduler
+        implements Scheduler
+{
+    private static final Logger log = Logger.get(MetricsBasedScheduler.class);
+    Map<URI, ClusterInfo> clusterInfos;
+
+    //Min heap based priority queue. Cluster with lowest number of queries will be on top of the queue
+    PriorityQueue<ClusterQueryInfo> clusterQueue = new PriorityQueue<ClusterQueryInfo>((x, y) -> (int) (x.totalQueries - y.totalQueries));
+
+    private List<URI> candidates;
+
+    private class ClusterQueryInfo
+    {
+        URI clusterUri;
+        long runningQueries;
+        long queuedQueries;
+        long totalQueries;
+
+        ClusterQueryInfo(URI clusterUri, long runningQueries, long queuedQueries)
+        {
+            this.clusterUri = clusterUri;
+            this.runningQueries = runningQueries;
+            this.queuedQueries = queuedQueries;
+            totalQueries = (long) (runningQueries + queuedQueries);
+            log.info("Cluster URI : " + clusterUri + ", runningQueries : " + runningQueries + ", queuedQueries : " + queuedQueries + ", totalQueries : " + totalQueries);
+        }
+    }
+
+    @Override
+    public Optional<URI> getDestination(String user, String query)
+    {
+        try {
+            if (clusterInfos != null && clusterInfos.size() > 0) {
+                clusterQueue.clear();
+                clusterQueue.addAll(clusterInfos.keySet().stream()
+                        .map(uri -> new ClusterQueryInfo(uri, clusterInfos.get(uri).getRunningQueries(), clusterInfos.get(uri).getQueuedQueries()))
+                        .collect(Collectors.toList()));
+                return Optional.of(clusterQueue.poll().clusterUri);
+            }
+            return Optional.empty();
+        }
+        catch (IllegalArgumentException e) {
+            log.warn(e, "Error getting destination for user " + user);
+            return Optional.empty();
+        }
+    }
+
+    @Override
+    public void setCandidates(List<URI> candidates)
+    {
+        this.candidates = candidates;
+    }
+
+    @Override
+    public void setClusterInfos(Map<URI, ClusterInfo> clusterInfos)
+    {
+        this.clusterInfos = clusterInfos;
+    }
+}

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

@@ -0,0 +1,35 @@
+/*
+ * 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
+{
+    @Override
+    public String getName()
+    {
+        return CustomSchedulerConstants.METRICS_BASED;
+    }
+
+    @Override
+    public Scheduler create(Map<String, String> config)
+    {
+        return new MetricsBasedScheduler();
+    }
+}

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

@@ -0,0 +1,32 @@
+/*
+ * 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 java.util.ArrayList;
+import java.util.List;
+
+public class RouterSchedulerPlugin
+        implements RouterPlugin
+{
+    @Override
+    public Iterable<SchedulerFactory> getSchedulerFactories()
+    {
+        List<SchedulerFactory> schedulerFactories = new ArrayList<>();
+        schedulerFactories.add(new MetricsBasedSchedulerFactory());
+        return schedulerFactories;
+    }
+}

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

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

presto-router-custom-scheduler/src/modernizer/violations.xml

@@ -0,0 +1,45 @@
+<?xml version="1.0"?>
+<modernizer>
+    <violation>
+        <name>java/lang/Class.newInstance:()Ljava/lang/Object;</name>
+        <version>1.1</version>
+        <comment>Prefer Class.getConstructor().newInstance()</comment>
+    </violation>
+
+    <violation>
+        <name>java/lang/String.toLowerCase:()Ljava/lang/String;</name>
+        <version>1.1</version>
+        <comment>Prefer String.toLowerCase(java.util.Locale)</comment>
+    </violation>
+
+    <violation>
+        <name>com/google/common/primitives/Ints.checkedCast:(J)I</name>
+        <version>1.8</version>
+        <comment>Prefer Math.toIntExact(long)</comment>
+    </violation>
+
+    <violation>
+        <name>org/testng/Assert.assertEquals:(Ljava/lang/Iterable;Ljava/lang/Iterable;)V</name>
+        <version>1.8</version>
+        <comment>Use com.facebook.presto.testing.assertions.Assert.assertEquals due to TestNG #543</comment>
+    </violation>
+
+    <violation>
+        <name>org/testng/Assert.assertEquals:(Ljava/lang/Iterable;Ljava/lang/Iterable;Ljava/lang/String;)V</name>
+        <version>1.8</version>
+        <comment>Use com.facebook.presto.testing.assertions.Assert.assertEquals due to TestNG #543</comment>
+    </violation>
+
+    <violation>
+        <name>java/util/TimeZone.getTimeZone:(Ljava/lang/String;)Ljava/util/TimeZone;</name>
+        <version>1.8</version>
+        <comment>Avoid TimeZone.getTimeZone as it returns GMT for a zone not supported by the JVM. Use TimeZone.getTimeZone(ZoneId.of(..)) instead, or TimeZone.getTimeZone(..., false).</comment>
+    </violation>
+
+    <violation>
+        <name>org/joda/time/DateTimeZone.toTimeZone:()Ljava/util/TimeZone;</name>
+        <version>1.8</version>
+        <comment>Avoid DateTimeZone.toTimeZone as it returns GMT for a zone not supported by the JVM. Use TimeZone.getTimeZone(ZoneId.of(dtz.getId())) instead.</comment>
+    </violation>
+
+</modernizer>