NIFI-8080: Compile Jython scripts before evaluating

This closes #4717

Signed-off-by: Mike Thomsen <mthomsen@apache.org>
This commit is contained in:
Matthew Burgess 2020-12-08 22:54:21 -05:00 committed by Mike Thomsen
parent cfbcecc4c6
commit 2c0671cf80
No known key found for this signature in database
GPG Key ID: 88511C3D4CAD246F
2 changed files with 44 additions and 12 deletions

View File

@ -20,15 +20,22 @@ import org.apache.nifi.logging.ComponentLog;
import org.apache.nifi.processors.script.ScriptEngineConfigurator;
import org.python.core.PyString;
import javax.script.Compilable;
import javax.script.CompiledScript;
import javax.script.ScriptEngine;
import javax.script.ScriptException;
import java.net.URL;
import java.util.Arrays;
import java.util.concurrent.atomic.AtomicReference;
import java.util.stream.Collectors;
/**
* A helper class to configure the Jython engine with any specific requirements
*/
public class JythonScriptEngineConfigurator implements ScriptEngineConfigurator {
private final AtomicReference<CompiledScript> compiledScriptRef = new AtomicReference<>();
@Override
public String getScriptEngineName() {
return "python";
@ -41,17 +48,9 @@ public class JythonScriptEngineConfigurator implements ScriptEngineConfigurator
}
@Override
public Object init(ScriptEngine engine, String[] modulePaths) throws ScriptException {
if (engine != null) {
// Need to import the module path inside the engine, in order to pick up
// other Python/Jython modules.
engine.eval("import sys");
if (modulePaths != null) {
for (String modulePath : modulePaths) {
engine.eval("sys.path.append(" + PyString.encode_UnicodeEscape(modulePath, true) + ")");
}
}
}
public Object init(ScriptEngine engine, String[] modulePaths) {
// Always compile when first run
compiledScriptRef.set(null);
return null;
}
@ -59,7 +58,17 @@ public class JythonScriptEngineConfigurator implements ScriptEngineConfigurator
public Object eval(ScriptEngine engine, String scriptBody, String[] modulePaths) throws ScriptException {
Object returnValue = null;
if (engine != null) {
returnValue = engine.eval(scriptBody);
final CompiledScript existing = compiledScriptRef.get();
if (existing == null) {
// Add prefix for import sys and all jython modules
String prefix = "import sys\n"
+ Arrays.stream(modulePaths).map((modulePath) -> "sys.path.append(" + PyString.encode_UnicodeEscape(modulePath, true) + ")")
.collect(Collectors.joining("\n"));
final CompiledScript compiled = ((Compilable) engine).compile(prefix + scriptBody);
compiledScriptRef.compareAndSet(null, compiled);
}
returnValue = compiledScriptRef.get().eval();
}
return returnValue;
}

View File

@ -19,6 +19,7 @@ package org.apache.nifi.processors.script;
import org.apache.nifi.script.ScriptingComponentUtils;
import org.apache.nifi.util.MockFlowFile;
import org.junit.Before;
import org.junit.Ignore;
import org.junit.Test;
import java.nio.charset.StandardCharsets;
@ -73,6 +74,28 @@ public class TestExecuteJython extends BaseScriptTest {
runner.assertValid();
runner.enqueue("test content".getBytes(StandardCharsets.UTF_8));
runner.run();
}
@Ignore("This is more of an integration test, can be run before and after changes to ExecuteScript to measure performance improvements")
@Test
public void testPerformance() throws Exception {
runner.setValidateExpressionUsage(false);
runner.setProperty(scriptingComponent.getScriptingComponentHelper().SCRIPT_ENGINE, "python");
runner.setProperty(ScriptingComponentUtils.SCRIPT_BODY,
"from org.apache.nifi.processors.script import ExecuteScript\n"
+ "flowFile = session.get()\n"
+ "flowFile = session.putAttribute(flowFile, \"from-content\", \"test content\")\n"
+ "session.transfer(flowFile, ExecuteScript.REL_SUCCESS)");
runner.assertValid();
final int ITERATIONS = 50000;
for (int i = 0; i < ITERATIONS; i++) {
runner.enqueue("test content".getBytes(StandardCharsets.UTF_8));
}
runner.run(ITERATIONS);
runner.assertAllFlowFilesTransferred(ExecuteScript.REL_SUCCESS, ITERATIONS);
final List<MockFlowFile> result = runner.getFlowFilesForRelationship(ExecuteScript.REL_SUCCESS);
result.get(0).assertAttributeEquals("from-content", "test content");
}
}