Presto UDF原理分析与浅实践

背景

这两天在研究Presto的UDF,发现Presto的UDF的方式很特别,后来自己思考了一下实现方式(没有看源码,有可能与Presto实现不同),现在把思路记录一下。

思路

Presto UDF 最后是将生成的Jar放到到plugin中,并且调用实现com.facebook.presto.spi.Plugin接口的类的getFunctions方法,那么我猜测的步骤应该如下
1.加载plugin文件夹下的依赖Jar
2.找到实现com.facebook.presto.spi.Plugin的类
3.运用反射获取实际逻辑的Class
下面我将主要的实现类写到下面,如有不足之处,还请大家指点

package com.levin.presto.plugin;

import com.facebook.presto.spi.Plugin;

import java.io.File;
import java.io.IOException;
import java.lang.annotation.Annotation;
import java.lang.reflect.InvocationTargetException;
import java.lang.reflect.Method;
import java.net.MalformedURLException;
import java.net.URL;
import java.net.URLClassLoader;
import java.util.*;
import java.util.jar.JarEntry;
import java.util.jar.JarFile;

/**
 * author Levin
 */
public class MyClassLoader {
    /**
     * 用来存放所有Class对象
     */
    private static Set<Class<?>> classes = new LinkedHashSet<Class<?>>();
    /**
     * 用来存放每个Class中的注解对象
     */
    private static Map<Class<?>, Annotation[]> classAnnotationMap = new HashMap<Class<?>, Annotation[]>();
    /**
     * 用来存放每个Class中方法的注解对象
     */
    private static Map<Class<?>, Map<Method, Annotation[]>> classMethodAnnoMap = new HashMap<Class<?>, Map<Method, Annotation[]>>();

    /**
     * 使用URLClassLoader加载制定lib下的所有Jar
     *
     * @param lib
     * @return
     * @throws MalformedURLException
     * @throws NoSuchMethodException
     * @throws InvocationTargetException
     * @throws IllegalAccessException
     */
    public static URLClassLoader libLoad(String lib) throws MalformedURLException, NoSuchMethodException, InvocationTargetException, IllegalAccessException {
        File file_directory = new File(lib);
        URLClassLoader loader = (URLClassLoader) ClassLoader.getSystemClassLoader();
        //使用反射调用addURL方法
        Method add = URLClassLoader.class.getDeclaredMethod("addURL", new Class[]{URL.class});
        add.setAccessible(true);
        if (!file_directory.exists()) {
            throw new RuntimeException("Jar不存在");
        }
        // load路径下的所有Jar
        for (File file : file_directory.listFiles()) {
            URL url = new URL("file:" + file.getAbsolutePath());
            add.invoke(loader, new Object[]{url});
        }
        return loader;
    }

    /**
     * 加载Jar包并且执行Jar包中某个Class的方法
     *
     * @param loader
     * @throws IOException
     * @throws ClassNotFoundException
     */
    public static void execJarClass(URLClassLoader loader, String jar_Path) throws IOException, ClassNotFoundException {
//        String jar_Path = "/home/levin/workspcae/java/plugin_test/target/plugin_test-1.0.jar";
        JarFile jarFile = new JarFile(new File(jar_Path));
        Enumeration<JarEntry> es = jarFile.entries();
        while (es.hasMoreElements()) {
            JarEntry jarEntry = (JarEntry) es.nextElement();
            String name = jarEntry.getName();
            // 只加载.class文件
            if (name != null && name.endsWith(".class")) {
                Class<?> c = loader.loadClass(name.replace("/", ".").substring(0, name.length() - 6));
                classes.add(c);
                Annotation[] classAnnos = c.getDeclaredAnnotations();
                classAnnotationMap.put(c, classAnnos);
                Method[] classMethods = c.getDeclaredMethods();
                Map<Method, Annotation[]> methodAnnoMap = new HashMap<Method, Annotation[]>();
                for (int i = 0; i < classMethods.length; i++) {
                    Annotation[] a = classMethods[i].getDeclaredAnnotations();
                    methodAnnoMap.put(classMethods[i], a);
                }
                classMethodAnnoMap.put(c, methodAnnoMap);
            }
        }
        System.out.println("添加 " + classes.size() + " 个Class文件");
    }

    public static void main(String[] args) throws IOException, NoSuchMethodException, IllegalAccessException, InvocationTargetException, ClassNotFoundException {
        String jar_Path = "/home/levin/workspcae/java/plugin_test/target/plugin_test-1.0.jar";
        URLClassLoader loader = libLoad("/home/levin/install-soft/bigdata/presto-server-0.187/lib");
        execJarClass(loader, jar_Path);
        Iterator class_iterator = classes.iterator();

        while (class_iterator.hasNext()) {
            Class c = (Class) class_iterator.next();
            if (Plugin.class.isAssignableFrom(c)) {
                System.out.println(c.getName().concat(" 类实现了 com.facebook.presto.spi.Plugin 接口"));
            } else {
                System.out.println(c.getName());
            }
        }
    }
}