动态代码生成技术在 Presto 中使用简介

谁践踏了优雅 2022-09-16 00:20 232阅读 0赞

在《[ASM 与 Presto 动态代码生成简介][ASM _ Presto]

## EmbedVersion ##

我们往 Presto 提交 SQL 查询以及 TaskExecutor 启动 TaskRunner 执行 Task 的时候都会使用到 EmbedVersion 类里面的 embedVersion 方法。embedVersion 方法其实就是初始化一个 Runnable 实例，比如启动 TaskRunner 的代码片段如下：

executor.execute(embedVersion.embedVersion(new TaskRunner()));

其中 TaskRunner 就是实现 Runnable 接口的。EmbedVersion 的 embedVersion 方法实现如下：

public Runnable embedVersion(Runnable runnable)
    {
        requireNonNull(runnable, "runnable is null");
        try {
            return (Runnable) runnableConstructor.invoke(runnable);
        }
        catch (Throwable throwable) {
            throwIfUnchecked(throwable);
            throw new RuntimeException(throwable);
        }
    }

其中 runnableConstructor 就是使用 ASM 进行代码生成的类，实现如下：

// 这里定义了一个类，类名大概为 Presto_null__testversion____20211011_105831_1，
    // 它的父类是 Object，并实现了 Runnable 接口。
    ClassDefinition classDefinition = new ClassDefinition(
        a(PUBLIC, FINAL),
        makeClassName(baseClassName(serverConfig)),
        type(Object.class),
        type(Runnable.class));
    // 定义了一个名为 runnable 的局部变量，类型为 Runnable
    FieldDefinition field = classDefinition.declareField(a(PRIVATE), "runnable", Runnable.class);
    Parameter parameter = arg("runnable", type(Runnable.class));
    // 定义了这个类的构造函数，参数为 runnable，参数类型为 Runnable
    MethodDefinition constructor = classDefinition.declareConstructor(a(PUBLIC), parameter);
    // 构造方法里面其实就是把参数 runnable 的值赋值给局部变量 runnable
    constructor.getBody()
        .comment("super(runnable);")
        .append(constructor.getThis())
        .invokeConstructor(Object.class)
        .append(constructor.getThis())
        .append(parameter)
        .putField(field)
        .ret();
    // 定义了一个名为 run 的方法，事实上就是实现 Runnable 接口里面的 run 方法
    MethodDefinition run = classDefinition.declareMethod(a(PUBLIC), "run", type(void.class));
    // run 里面其实就是调用局部变量 runnable 的 run 方法
    run.getBody()
        .comment("runnable.run();")
        .append(run.getThis())
        .getField(field)
        .invokeInterface(Runnable.class, "run", void.class)
        .ret();
    // 定义这个类，并加载到 ClassLoader 中
    Class<? extends Runnable> generatedClass = defineClass(classDefinition, Runnable.class, ImmutableMap.of(), getClass().getClassLoader());
    this.runnableConstructor = constructorMethodHandle(generatedClass, Runnable.class);

上面是 Presto 操作 Java 字节码并动态生成了一个类，其生成的类大概如下面所示：

package com.facebook.presto.$gen;
    public final class Presto_null__testversion____20211011_105831_1 implements Runnable {
        private Runnable runnable;
        public Presto_null__testversion____20211011_105831_1(Runnable runnable) {
            this.runnable = runnable;
        }
        public void run() {
            this.runnable.run();
        }
    }

看起来内容其实很简单。EmbedVersion 类算是 Presto 里面动态代码生成最简单的例子了。

## concat 函数实现 ##

下面我们来看下稍微复杂的，也就是 Presto 里面内置函数的实现。Presto 的内置函数的实现很多也是用到代码生成技术，比如 map\_filter、transform\_keys 以及 transform\_values 等。我们这里也举一个比较简单的例子，也就是 concat 函数的实现。比如下面的 SQL 查询：

select concat(o_orderstatus, o_orderpriority) from orders limit 10;

在 Presto 里面，concat 函数的实现就是通过代码生成进行的，其实现代码可以参见 com.facebook.presto.operator.scalar.ConcatFunction。Presto 接收到上面的 SQL 查询后，会在 Coordinator 端进行解析，并生成相应的 Tasks，提交给 Worker 执行。在 Worker 端，执行 Task 的时候，会调用 LocalExecutionPlanner 的 plan 方法生成 LocalExecutionPlan 其实就是本地可执行的计划，在 plan 方法里面会调用 com.facebook.presto.sql.planner.LocalExecutionPlanner.Visitor 对 Coordinator 传过来的 PlanNode 进行变量生成 PhysicalOperation。在我们的例子中，会在 com.facebook.presto.sql.planner.LocalExecutionPlanner.Visitor\#visitScanFilterAndProject 里面对 concat(o\_orderstatus, o\_orderpriority) 进行代码生成，最终调用到 com.facebook.presto.operator.scalar.ConcatFunction 的 generateConcat 方法，其就是 Presto 的 concat 函数实现逻辑，如下：

// arity 代表 Concat 函数输入参数的个数
    private static Class<?> generateConcat(TypeSignature type, int arity)
    {
        checkCondition(arity <= 254, NOT_SUPPORTED, "Too many arguments for string concatenation");
        // 定义动态代码生成的类名，生成的类名大概是 varchar_concat2ScalarFunction_20211011_062900_3 样子的
        ClassDefinition definition = new ClassDefinition(
            a(PUBLIC, FINAL),
            makeClassName(type.getBase() + "_concat" + arity + "ScalarFunction"),
            type(Object.class));
        // 生成类的构造函数，这里是使用 private 修饰的
        // Generate constructor
        definition.declareDefaultConstructor(a(PRIVATE));
        // Generate concat()
        // 定义 concat 函数的参数，比如 arg0、arg1；类型是 Slice
        List<Parameter> parameters = IntStream.range(0, arity)
            .mapToObj(i -> arg("arg" + i, Slice.class))
            .collect(toImmutableList());
        // 定义一个名为 concat 的函数，它的修饰符是 public static，
        // 返回类型是 Slice，输入参数是上面定义的 arg0、arg1 等。
        MethodDefinition method = definition.declareMethod(a(PUBLIC, STATIC), "concat", type(Slice.class), parameters);
        Scope scope = method.getScope();
        BytecodeBlock body = method.getBody();
        // 定义一个名为 length 的局部变量，类型为 int
        Variable length = scope.declareVariable(int.class, "length");
        // length 变量初始化为0
        body.append(length.set(constantInt(0)));
        // 下面是计算 concat 函数每个参数的长度（其实就是调用 string 的 length 方法）
        // 然后再把得到的字符串长度加到 length 里面，并赋值给 length
        for (int i = 0; i < arity; ++i) {
            body.append(length.set(generateCheckedAdd(length, parameters.get(i).invoke("length", int.class))));
        }
        // 定义一个名为 result 的局部变量，类型为 Slice
        Variable result = scope.declareVariable(Slice.class, "result");
        // 调用 Slices 的 allocate 方法分配出长度为 length 空间的 Slice 对象，并赋值给 result
        body.append(result.set(invokeStatic(Slices.class, "allocate", Slice.class, length)));
        // 定义一个名为 position 的局部变量，类型为 int，赋值为 0
        Variable position = scope.declareVariable(int.class, "position");
        body.append(position.set(constantInt(0)));
        // 下面是循环调用 result 的 setBytes 方法，并分别把 arg0、arg1 里面的内容放到 result 里面去
        // 最后计算 arg0 或 arg1 字符串长度再加上 position 的值，结果再赋值给 position
        for (int i = 0; i < arity; ++i) {
            body.append(result.invoke("setBytes", void.class, position, parameters.get(i)));
            body.append(position.set(add(position, parameters.get(i).invoke("length", int.class))));
        }
        // 返回 result
        body.getVariable(result)
            .retObject();
        // 定义生成的类，并把它加载打破 DynamicClassLoader 里面去
        return defineClass(definition, Object.class, ImmutableMap.of(), new DynamicClassLoader(ConcatFunction.class.getClassLoader()));
    }
    private static BytecodeExpression generateCheckedAdd(BytecodeExpression x, BytecodeExpression y)
    {
        // 调用 ConcatFunction 类里面的 checkedAdd 静态方法
        return invokeStatic(ConcatFunction.class, "checkedAdd", int.class, x, y);
    }
    @UsedByGeneratedCode
    public static int checkedAdd(int x, int y)
    {
        try {
            return addExact(x, y);
        }
        catch (ArithmeticException e) {
            throw new PrestoException(INVALID_FUNCTION_ARGUMENT, "Concatenated string is too large");
        }
    }

为了方便理解，我对 generateConcat 方法的实现进行了注释，应该很好理解。为了性能问题，最终生成的函数会进行缓存，下一次再调用 concat 函数，只要函数签名一样，就不用再一次进行 concat 代码的生成。比如我们前面的例子是对两个字符串进行合并（函数签名为 presto.default.concat(varchar,varchar):varchar ），如果下一次还是调用这个函数就不用再进行代码生成了。但是如果下一次是对三个字符串进行合并，还是要进行一次代码生成的。

到这里，大家可能还是不太明白 Presto 代码生成到底生成了什么东西。这里我就进一步介绍一下。如果运行我们上面的 SQL 查询，Presto 生成的 concat 实现大概如下面所示：

package com.facebook.presto.$gen;
    import com.facebook.presto.operator.scalar.ConcatFunction;
    import io.airlift.slice.Slice;
    import io.airlift.slice.Slices;
    public final class varchar_concat2ScalarFunction_20211011_062900_3 {
        private varchar_concat2ScalarFunction_20211011_062900_3() {
        }
        public static Slice concat(Slice arg0, Slice arg1) {
            int length = 0;
            int length = ConcatFunction.checkedAdd(length, arg0.length());
            length = ConcatFunction.checkedAdd(length, arg1.length());
            Slice result = Slices.allocate(length);
            int position = 0;
            result.setBytes(position, arg0);
            int position = position + arg0.length();
            result.setBytes(position, arg1);
            int var10000 = position + arg1.length();
            return result;
        }
    }

注意，Presto 里面生成的是 Java 字节码，这里只是为了说明的方便，给出了 Java 源代码。可以看到，最终生成的代码其实很好理解。Presto 里面对两个字符串进行 concat 其实就是执行上面的代码片段。

## 总结 ##

本文通过两个例子简单的介绍了 Presto 里面动态代码生成的用法，通过上面两个例子，相信大家应该能够大概了解 Presto 里面的动态代码生成的技术和用法。后面有空我会介绍一下表达式这块代码生成和整个 page 处理逻辑是怎么联系到一起的，敬请关注。

[ASM _ Presto]: https://blog.csdn.net/junerli/article/details/112687385