ScriptEngine.java

package com.thinkbiganalytics.spark.repl;

/*-
 * #%L
 * thinkbig-commons-spark-repl
 * %%
 * Copyright (C) 2017 ThinkBig Analytics
 * %%
 * 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.
 * #L%
 */

import com.google.common.base.Charsets;
import com.google.common.base.Throwables;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.Maps;
import com.thinkbiganalytics.spark.util.ArrayUtils;

import org.apache.commons.io.output.ByteArrayOutputStream;
import org.apache.spark.SparkContext;
import org.apache.spark.sql.SQLContext;
import org.apache.spark.sql.hive.HiveContext;

import java.io.PrintWriter;
import java.io.UnsupportedEncodingException;
import java.util.List;
import java.util.Map;
import java.util.concurrent.atomic.AtomicReference;

import javax.annotation.Nonnull;
import javax.annotation.Nullable;
import javax.script.ScriptException;

import scala.tools.nsc.interpreter.NamedParam;

/**
 * Interface for an interpreter that compiles and evaluates Scala code containing a Spark job.
 *
 * <p>Scripts may access a {@link SparkContext} through the {@code sc} variable and a {@link SQLContext} through the
 * {@code sqlContext} variable.</p>
 *
 * <p>This class is <i>thread-safe</i> and ensures that only one script </p>
 */
public abstract class ScriptEngine {

    /**
     * End of line character
     */
    private static final byte[] END_LINE = new byte[]{'\n'};

    /**
     * Label used by the compiler to indicate a compile error
     */
    private static final byte[] LABEL = "<console>".getBytes(Charsets.UTF_8);

    /**
     * Separator between label, line number, and error message
     */
    private static final byte[] SEPARATOR = new byte[]{':'};

    /**
     * Exception thrown by the last script
     */
    @Nonnull
    private final AtomicReference<Throwable> exception = new AtomicReference<>();

    /**
     * Compiler output stream for capturing compile errors
     */
    @Nonnull
    private final ByteArrayOutputStream out = new ByteArrayOutputStream();

    /**
     * Result of the last script
     */
    @Nonnull
    private final AtomicReference<Object> result = new AtomicReference<>();
    /**
     * Map of variable names to values for bindings
     */
    @Nonnull
    private final Map<String, Object> values = Maps.newHashMap();
    /**
     * Spark context
     */
    @Nullable
    private SparkContext sparkContext;
    /**
     * Spark SQL context
     */
    @Nullable
    private SQLContext sqlContext;

    /**
     * Executes the specified script.
     *
     * @param script the script to be executed
     * @return the value returned from the script
     * @throws ScriptException if an error occurs in the script
     */
    @Nullable
    public synchronized Object eval(@Nonnull final String script) throws ScriptException {
        List<NamedParam> bindings = ImmutableList.of();
        return eval(script, bindings);
    }

    /**
     * Executes the specified script with the given bindings.
     *
     * @param script   the script to be executed
     * @param bindings the variable bindings to be accessible to the script
     * @return the value returned from the script
     * @throws ScriptException if an error occurs in the script
     */
    @Nullable
    public synchronized Object eval(@Nonnull final String script, @Nonnull final List<NamedParam> bindings)
        throws ScriptException {
        // Define class containing script
        final StringBuilder cls = new StringBuilder();
        cls.append("class Script (engine: com.thinkbiganalytics.spark.repl.ScriptEngine)");
        cls.append("    extends com.thinkbiganalytics.spark.repl.Script (engine) {\n");
        cls.append("  override def eval (): Any = {\n");
        cls.append(script);
        cls.append("  }\n");

        // Add bindings to class
        this.values.clear();

        for (NamedParam param : bindings) {
            cls.append("  def ");
            cls.append(param.name());
            cls.append(" (): ");
            cls.append(param.tpe());
            cls.append(" = getValue(\"");
            cls.append(param.name());
            cls.append("\")\n");
            this.values.put(param.name(), param.value());
        }

        cls.append("}\n");

        // Instantiate class
        cls.append("new Script(engine).run()\n");

        // Execute script
        this.out.reset();

        execute(cls.toString());

        // Check for exception and return result
        checkCompileError();
        checkRuntimeError();

        return this.result.get();
    }

    /**
     * Gets the class loader used by the interpreter.
     *
     * @return the class loader
     */
    @Nonnull
    public abstract ClassLoader getClassLoader();

    /**
     * Gets the {@code SparkContext} available to scripts as {@code sc}.
     *
     * @return the Spark context
     */
    @Nonnull
    public SparkContext getSparkContext() {
        if (this.sparkContext == null) {
            this.sparkContext = createSparkContext();
        }
        return this.sparkContext;
    }

    /**
     * Gets the {@code SQLContext} available to scripts as {@code sqlContext}.
     *
     * @return the SQL context
     */
    @Nonnull
    public SQLContext getSQLContext() {
        if (this.sqlContext == null) {
            this.sqlContext = new HiveContext(getSparkContext());
        }
        return this.sqlContext;
    }

    /**
     * Creates the {@code SparkContext} that will be available to scripts as {@code sc}.
     *
     * @return the Spark context
     */
    @Nonnull
    protected abstract SparkContext createSparkContext();

    /**
     * Executes the specified script.
     *
     * @param script the script to be executed
     * @throws ScriptException if an error occurs in the script
     */
    protected abstract void execute(@Nonnull final String script) throws ScriptException;

    /**
     * Gets the writer for capturing compile errors.
     *
     * @return the compiler output stream
     */
    protected PrintWriter getPrintWriter() {
        return new PrintWriter(this.out);
    }

    /**
     * Resets the engine state so the {@link SparkContext} can be recreated.
     */
    protected void reset() {
        // Stop Spark
        if (sparkContext != null && !sparkContext.isStopped()) {
            sparkContext.stop();
        }

        // Clear instance variables
        exception.set(null);
        out.reset();
        result.set(null);
        sparkContext = null;
        sqlContext = null;
    }

    /**
     * Gets the value of the specified binding.
     *
     * @param name the name of the binding
     * @return the value of the binding
     */
    @Nullable
    Object getValue(@Nonnull final String name) {
        return this.values.get(name);
    }

    /**
     * Sets the runtime exception for the current script.
     *
     * @param t the exception
     */
    void setException(@Nonnull final Throwable t) {
        this.exception.set(t);
    }

    /**
     * Sets the result of the current script.
     *
     * @param result the result
     */
    void setResult(@Nullable final Object result) {
        this.exception.set(null);
        this.result.set(result);
    }

    /**
     * Checks the output stream for a compile error.
     *
     * @throws ScriptException if a compile error is found
     */
    private void checkCompileError() throws ScriptException {
        byte[] outBytes = this.out.toByteArray();

        // Look for label
        int labelIndex = ArrayUtils.indexOf(outBytes, 0, outBytes.length, LABEL, 0, LABEL.length, 0);
        if (labelIndex == -1) {
            return;
        }

        // Look for start and end of message
        int lineIndex = labelIndex + LABEL.length + 1;
        int msgStart = ArrayUtils.indexOf(outBytes, 0, outBytes.length, SEPARATOR, 0, SEPARATOR.length, lineIndex) + 2;
        int msgEnd = ArrayUtils.indexOf(outBytes, 0, outBytes.length, END_LINE, 0, END_LINE.length, msgStart);

        // Throw exception
        int line;
        String message;

        try {
            line = Integer.parseInt(new String(outBytes, lineIndex, msgStart - lineIndex - 2));
            message = new String(outBytes, msgStart, msgEnd - msgStart, "UTF-8");
        } catch (UnsupportedEncodingException e) {
            throw new IllegalStateException(e);
        }

        throw new ScriptException(message, "<console>", line);
    }

    /**
     * Checks for a runtime exception.
     *
     * @throws ScriptException if an exception is found
     */
    private void checkRuntimeError() throws ScriptException {
        Throwable exception = this.exception.get();

        if (exception != null) {
            Throwables.propagateIfPossible(exception, ScriptException.class);

            if (exception instanceof Exception) {
                throw new ScriptException((Exception) exception);
            } else {
                throw new ScriptException(exception.getMessage());
            }
        }
    }
}