<commit_before>package com.tinkerpop.gremlin.groovy.engine;

import org.apache.commons.lang3.concurrent.BasicThreadFactory;
import org.javatuples.Pair;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.script.Bindings;
import javax.script.ScriptException;
import javax.script.SimpleBindings;
import java.io.File;
import java.io.FileReader;
import java.io.IOException;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.ScheduledFuture;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.function.BiConsumer;
import java.util.function.Consumer;

/**
 * Execute Gremlin scripts against a {@code ScriptEngine} instance.
 *
 * @author <NAME> (http://stephen.genoprime.com)
 */
public class GremlinExecutor {
	private static final Logger logger = LoggerFactory.getLogger(GremlinExecutor.class);

	/**
	 * {@link ScriptEngines} instance to evaluate Gremlin script requests.
	 */
	private ScriptEngines sharedScriptEngines;

	private final Map<String, EngineSettings> settings;
	private final long scriptEvaluationTimeout;
	private final Bindings globalBindings;
	private final List<List<String>> use;
	private final ExecutorService executorService;
	private final ScheduledExecutorService scheduledExecutorService;
	private final Consumer<Bindings> beforeEval;
	private final Consumer<Bindings> afterSuccess;
	private final Consumer<Bindings> afterTimeout;
	private final BiConsumer<Bindings,Exception> afterFailure;

	private GremlinExecutor(final Map<String, EngineSettings> settings, final List<List<String>> use,
							final long scriptEvaluationTimeout, final Bindings globalBindings,
							final ExecutorService executorService, final ScheduledExecutorService scheduledExecutorService,
							final Consumer<Bindings> beforeEval, final Consumer<Bindings> afterSuccess,
							final Consumer<Bindings> afterTimeout, final BiConsumer<Bindings,Exception> afterFailure) {
		if (logger.isDebugEnabled())
			logger.debug("Initializing GremlinExecutor.  This should not happen more than once.");

		this.executorService = executorService;
		this.scheduledExecutorService = scheduledExecutorService;
		this.beforeEval = beforeEval;
		this.afterSuccess = afterSuccess;
		this.afterTimeout = afterTimeout;
		this.afterFailure = afterFailure;
		this.use = use;
		this.settings = settings;
		this.scriptEvaluationTimeout = scriptEvaluationTimeout;
		this.globalBindings = globalBindings;
		sharedScriptEngines = createScriptEngines();
	}

	public CompletableFuture<Object> eval(final String script) {
		return eval(script, Optional.empty(), new SimpleBindings());
	}

	public CompletableFuture<Object> eval(final String script, final Bindings boundVars) {
		return eval(script, Optional.empty(), boundVars);
	}

	public CompletableFuture<Object> eval(final String script, final Map<String,Object> boundVars) {
		return eval(script, Optional.empty(), new SimpleBindings(boundVars));
	}

	public CompletableFuture<Object> eval(final String script, final Optional<String> language, final Bindings boundVars) {
		final Bindings bindings = new SimpleBindings();
		bindings.putAll(this.globalBindings);
		bindings.putAll(boundVars);

		final String lang = language.orElse("gremlin-groovy");

		if (logger.isDebugEnabled()) logger.debug("Preparing to evaluate script - {} - in thread [{}]", script, Thread.currentThread().getName());

		// select the gremlin threadpool to execute the script evaluation in
		final AtomicBoolean abort = new AtomicBoolean(false);
		final CompletableFuture<Object> future = CompletableFuture.supplyAsync(() -> {
			try {
				if (logger.isDebugEnabled()) logger.debug("Evaluating script - {} - in thread [{}]", script, Thread.currentThread().getName());

				beforeEval.accept(bindings);
				final Object o = sharedScriptEngines.eval(script, bindings, lang);

				if (abort.get())
					afterTimeout.accept(bindings);
				else
					afterSuccess.accept(bindings);

				return o;
			} catch (Exception ex) {
				afterFailure.accept(bindings, ex);
				throw new RuntimeException(ex);
			}
		}, executorService);

		scheduleTimeout(future, script, abort);

		return future;
	}

	private void scheduleTimeout(final CompletableFuture<Object> evaluationFuture, final String script, final AtomicBoolean abort) {
		if (scriptEvaluationTimeout > 0) {
			// Schedule a timeout in the io threadpool for future execution - killing an eval is cheap
			final ScheduledFuture<?> sf = scheduledExecutorService.schedule(() -> {
				if (logger.isDebugEnabled())
					logger.info("Timing out script - {} - in thread [{}]", script, Thread.currentThread().getName());

				if (!evaluationFuture.isDone()) {<commit_msg>Add method to get ScriptEngine from the executor.<commit_after>package com.tinkerpop.gremlin.groovy.engine;

import org.apache.commons.lang3.concurrent.BasicThreadFactory;
import org.javatuples.Pair;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.script.Bindings;
import javax.script.ScriptException;
import javax.script.SimpleBindings;
import java.io.File;
import java.io.FileReader;
import java.io.IOException;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.ScheduledFuture;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.function.BiConsumer;
import java.util.function.Consumer;

/**
 * Execute Gremlin scripts against a {@code ScriptEngine} instance.  It is designed to host any JSR-223 enabled
 * {@code ScriptEngine} and assumes such engines are designed to be thread-safe in the evaluation.  Script evaluation
 * functions return a {@link java.util.concurrent.CompletableFuture} where scripts may timeout if their evaluation
 * takes too long.
 *
 * @author <NAME> (http://stephen.genoprime.com)
 */
public class GremlinExecutor {
	private static final Logger logger = LoggerFactory.getLogger(GremlinExecutor.class);

	/**
	 * {@link ScriptEngines} instance to evaluate Gremlin script requests.
	 */
	private ScriptEngines scriptEngines;

	private final Map<String, EngineSettings> settings;
	private final long scriptEvaluationTimeout;
	private final Bindings globalBindings;
	private final List<List<String>> use;
	private final ExecutorService executorService;
	private final ScheduledExecutorService scheduledExecutorService;
	private final Consumer<Bindings> beforeEval;
	private final Consumer<Bindings> afterSuccess;
	private final Consumer<Bindings> afterTimeout;
	private final BiConsumer<Bindings,Exception> afterFailure;

	private GremlinExecutor(final Map<String, EngineSettings> settings, final List<List<String>> use,
							final long scriptEvaluationTimeout, final Bindings globalBindings,
							final ExecutorService executorService, final ScheduledExecutorService scheduledExecutorService,
							final Consumer<Bindings> beforeEval, final Consumer<Bindings> afterSuccess,
							final Consumer<Bindings> afterTimeout, final BiConsumer<Bindings,Exception> afterFailure) {
		if (logger.isDebugEnabled())
			logger.debug("Initializing GremlinExecutor.  This should not happen more than once.");

		this.executorService = executorService;
		this.scheduledExecutorService = scheduledExecutorService;
		this.beforeEval = beforeEval;
		this.afterSuccess = afterSuccess;
		this.afterTimeout = afterTimeout;
		this.afterFailure = afterFailure;
		this.use = use;
		this.settings = settings;
		this.scriptEvaluationTimeout = scriptEvaluationTimeout;
		this.globalBindings = globalBindings;
		this.scriptEngines = createScriptEngines();
	}

	public CompletableFuture<Object> eval(final String script) {
		return eval(script, Optional.empty(), new SimpleBindings());
	}

	public CompletableFuture<Object> eval(final String script, final Bindings boundVars) {
		return eval(script, Optional.empty(), boundVars);
	}

	public CompletableFuture<Object> eval(final String script, final Map<String,Object> boundVars) {
		return eval(script, Optional.empty(), new SimpleBindings(boundVars));
	}

	public CompletableFuture<Object> eval(final String script, final Optional<String> language, final Bindings boundVars) {
		final Bindings bindings = new SimpleBindings();
		bindings.putAll(this.globalBindings);
		bindings.putAll(boundVars);

		final String lang = language.orElse("gremlin-groovy");

		if (logger.isDebugEnabled()) logger.debug("Preparing to evaluate script - {} - in thread [{}]", script, Thread.currentThread().getName());

		// select the gremlin threadpool to execute the script evaluation in
		final AtomicBoolean abort = new AtomicBoolean(false);
		final CompletableFuture<Object> future = CompletableFuture.supplyAsync(() -> {
			try {
				if (logger.isDebugEnabled()) logger.debug("Evaluating script - {} - in thread [{}]", script, Thread.currentThread().getName());

				beforeEval.accept(bindings);
				final Object o = scriptEngines.eval(script, bindings, lang);

				if (abort.get())
					afterTimeout.accept(bindings);
				else
					afterSuccess.accept(bindings);

				return o;
			} catch (Exception ex) {
				afterFailure.accept(bindings, ex);
				throw new RuntimeException(ex);
			}
		}, executorService);

		scheduleTimeout(future, script, abort);

		return future;
	}

	public ScriptEngines getScriptEngines() {
		return this.scriptEngines;
	}

	private void scheduleTimeout(final CompletableFuture<Object> evaluationFuture, final String script, final AtomicBoolean abort) {
		if (scriptEvaluationTimeout > 0) {
			// Schedule a timeout in the io threadpool for future execution - killing an eval is cheap
			final ScheduledFuture<?> sf = scheduledExecutorService.schedule(() -> {
				if (logger.isDebugEnabled())
					logger.info("Timing out script - {} - in thread [{}]", script, Thread.currentThread().getName());

				if (!evaluationFuture.isDone()) {
{'fraction_non_alphanumeric': 0.09808224406068124, 'fraction_numerical': 0.001049518175746589, 'mean_word_length': 5.490402476780186, 'pattern_counts': {'":': 0, '<': 51, '<?xml version=': 0, '>': 57, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6276334', 'n_tokens_mistral': 3133, 'n_tokens_neox': 2832, 'n_words': 767}