diff --git a/recaf-core/src/main/java/software/coley/recaf/services/plugin/CdiClassAllocator.java b/recaf-core/src/main/java/software/coley/recaf/services/plugin/CdiClassAllocator.java index 02cb822b5..cc4d328fd 100644 --- a/recaf-core/src/main/java/software/coley/recaf/services/plugin/CdiClassAllocator.java +++ b/recaf-core/src/main/java/software/coley/recaf/services/plugin/CdiClassAllocator.java @@ -6,8 +6,10 @@ import jakarta.enterprise.inject.spi.*; import jakarta.inject.Inject; -import java.util.IdentityHashMap; import java.util.Map; +import java.util.WeakHashMap; +import java.util.concurrent.locks.Lock; +import java.util.concurrent.locks.ReentrantLock; /** * Allocator instance that ties into the CDI container. @@ -16,7 +18,8 @@ */ @ApplicationScoped public class CdiClassAllocator implements ClassAllocator { - private final Map, Bean> classBeanMap = new IdentityHashMap<>(); + private final Map, Bean> classBeanMap = new WeakHashMap<>(); + private final Lock lock = new ReentrantLock(); private final BeanManager beanManager; @Inject @@ -28,13 +31,18 @@ public CdiClassAllocator(@Nonnull BeanManager beanManager) { @Override @SuppressWarnings("unchecked") public T instance(@Nonnull Class cls) throws AllocationException { + lock.lock(); try { // Create bean Bean bean = (Bean) classBeanMap.computeIfAbsent(cls, c -> { - AnnotatedType annotatedClass = beanManager.createAnnotatedType(cls); + // TODO something here is bugged + // bean.create(creationalContext).getClass().getClassLoader() == cls.getClassLoader() + // - Evaluates to false + // - In AnnotatedTypeIdentifier#forBackedAnnotatedType the ClassLoader is not considered, just the class's name + AnnotatedType annotatedClass = beanManager.createAnnotatedType((Class) c); BeanAttributes attributes = beanManager.createBeanAttributes(annotatedClass); InjectionTargetFactory factory = beanManager.getInjectionTargetFactory(annotatedClass); - return beanManager.createBean(attributes, cls, factory); + return beanManager.createBean(attributes, (Class) c, factory); }); CreationalContext creationalContext = beanManager.createCreationalContext(bean); @@ -42,6 +50,8 @@ public T instance(@Nonnull Class cls) throws AllocationException { return bean.create(creationalContext); } catch (Throwable t) { throw new AllocationException(cls, t); + } finally { + lock.unlock(); } } } diff --git a/recaf-core/src/main/java/software/coley/recaf/services/script/CancellationSingleton.java b/recaf-core/src/main/java/software/coley/recaf/services/script/CancellationSingleton.java new file mode 100644 index 000000000..adab09912 --- /dev/null +++ b/recaf-core/src/main/java/software/coley/recaf/services/script/CancellationSingleton.java @@ -0,0 +1,44 @@ +package software.coley.recaf.services.script; + +import software.coley.recaf.util.CancelSignal; + +/** + * Cancellation singleton. + *

+ * Injected into generated scripts. + * Do not use in normal code. + * + * @author xDark + * @see InsertCancelSignalVisitor + * @see GenerateResult#requestStop() + */ +public final class CancellationSingleton { + private static volatile boolean shouldStop; + + private CancellationSingleton() {} + + /** + * Schedules cancellation of the script. + */ + public static void stop() { + shouldStop = true; + } + + /** + * Clears any prior cancellation request. + */ + public static void reset() { + shouldStop = false; + } + + /** + * Polls for cancellation. + * + * @throws CancelSignal + * If cancellation has been requested. + */ + public static void poll() { + if (shouldStop) + throw CancelSignal.get(); + } +} diff --git a/recaf-core/src/main/java/software/coley/recaf/services/script/GenerateResult.java b/recaf-core/src/main/java/software/coley/recaf/services/script/GenerateResult.java index a364d604d..f074bcbfe 100644 --- a/recaf-core/src/main/java/software/coley/recaf/services/script/GenerateResult.java +++ b/recaf-core/src/main/java/software/coley/recaf/services/script/GenerateResult.java @@ -4,6 +4,7 @@ import jakarta.annotation.Nullable; import software.coley.recaf.services.compile.CompilerDiagnostic; +import java.lang.reflect.InvocationTargetException; import java.util.List; /** @@ -23,4 +24,42 @@ public record GenerateResult(@Nullable Class cls, @Nonnull List + * Cancellation is scoped to the generated class loader that owns {@link #cls()}. + * Callers that wish to stop running scripts must use {@link ScriptEngine#run(GenerateResult)} + * and track the returned instance for stopping. + * + * @throws IllegalStateException + * If something went wrong. + */ + public void requestStop() { + invokeCancellationSingleton("stop"); + } + + /** + * Clears any prior request to stop the script. If the generation failed, this method will do nothing. + * + * @throws IllegalStateException + * If something went wrong. + */ + public void resetStop() { + invokeCancellationSingleton("reset"); + } + + private void invokeCancellationSingleton(@Nonnull String methodName) { + Class cls = this.cls; + if (cls == null) + return; + try { + Class cancellationSingleton = cls.getClassLoader().loadClass(CancellationSingleton.class.getName()); + cancellationSingleton.getDeclaredMethod(methodName).invoke(null); + } catch (InvocationTargetException ex) { + throw new IllegalStateException(ex.getTargetException()); + } catch (Exception ex) { + throw new IllegalStateException(ex); + } + } +} diff --git a/recaf-core/src/main/java/software/coley/recaf/services/script/InsertCancelSignalVisitor.java b/recaf-core/src/main/java/software/coley/recaf/services/script/InsertCancelSignalVisitor.java new file mode 100644 index 000000000..062c5c78b --- /dev/null +++ b/recaf-core/src/main/java/software/coley/recaf/services/script/InsertCancelSignalVisitor.java @@ -0,0 +1,265 @@ +package software.coley.recaf.services.script; + +import jakarta.annotation.Nonnull; +import jakarta.annotation.Nullable; +import org.objectweb.asm.ClassVisitor; +import org.objectweb.asm.MethodVisitor; +import org.objectweb.asm.Type; +import org.objectweb.asm.tree.AbstractInsnNode; +import org.objectweb.asm.tree.InsnList; +import org.objectweb.asm.tree.InsnNode; +import org.objectweb.asm.tree.JumpInsnNode; +import org.objectweb.asm.tree.LabelNode; +import org.objectweb.asm.tree.LookupSwitchInsnNode; +import org.objectweb.asm.tree.MethodInsnNode; +import org.objectweb.asm.tree.MethodNode; +import org.objectweb.asm.tree.TableSwitchInsnNode; +import org.objectweb.asm.tree.TryCatchBlockNode; +import software.coley.recaf.RecafConstants; +import software.coley.recaf.util.CancelSignal; + +import java.util.ArrayList; +import java.util.Collections; +import java.util.IdentityHashMap; +import java.util.List; +import java.util.Map; +import java.util.Set; + +import static org.objectweb.asm.Opcodes.*; + +/** + * Class visitor that inserts calls to {@link CancellationSingleton#poll()} at the start of loops. + *

+ * This does not provide a perfect guarantee of cancellation, but it is a best-effort + * to allow scripts to be stopped in a timely manner. + * + * @author xDark + */ +final class InsertCancelSignalVisitor extends ClassVisitor { + private static final String SINGLETON_TYPE = Type.getInternalName(CancellationSingleton.class); + private static final String CANCEL_SIGNAL_TYPE = Type.getInternalName(CancelSignal.class); + private static final String POLL = "poll"; + private static final String POLL_DESC = "()V"; + + /** + * @param cv + * Parent visitor to delegate to. + */ + InsertCancelSignalVisitor(@Nonnull ClassVisitor cv) { + super(RecafConstants.getAsmVersion(), cv); + } + + @Override + public MethodVisitor visitMethod(int access, String name, String descriptor, String signature, + String[] exceptions) { + MethodVisitor delegate = super.visitMethod(access, name, descriptor, signature, exceptions); + return new MethodNode(api, access, name, descriptor, signature, exceptions) { + @Override + public void visitEnd() { + insertPolls(this); + insertCancelSignalRethrows(this); + accept(delegate); + } + }; + } + + /** + * Inserts calls to {@link CancellationSingleton#poll()} at the start of loops. + * + * @param method + * Method to insert into. + */ + private static void insertPolls(@Nonnull MethodNode method) { + // Collect label positions for loop detection. + Map positions = new IdentityHashMap<>(); + int position = 0; + for (AbstractInsnNode insn : method.instructions) { + if (insn instanceof LabelNode label) + positions.put(label, position); + position++; + } + + // Collect loop heading labels. + List loopHeaders = new ArrayList<>(); + Set addedLoopHeaders = Collections.newSetFromMap(new IdentityHashMap<>()); + position = 0; + for (AbstractInsnNode insn : method.instructions) { + if (insn instanceof JumpInsnNode jumpInsn) { + if (isLoopback(positions, position, jumpInsn.label) && addedLoopHeaders.add(jumpInsn.label)) + loopHeaders.add(jumpInsn.label); + } else if (insn instanceof LookupSwitchInsnNode switchInsn) { + addLoopbackTarget(positions, position, switchInsn.dflt, loopHeaders, addedLoopHeaders); + addLoopbackTargets(positions, position, switchInsn.labels, loopHeaders, addedLoopHeaders); + } else if (insn instanceof TableSwitchInsnNode switchInsn) { + addLoopbackTarget(positions, position, switchInsn.dflt, loopHeaders, addedLoopHeaders); + addLoopbackTargets(positions, position, switchInsn.labels, loopHeaders, addedLoopHeaders); + } + position++; + } + + // Insert polls at loop headers. + for (LabelNode loopHeader : loopHeaders) + method.instructions.insert(loopHeaderInsertionPoint(loopHeader), newPoll()); + } + + /** + * Adds loopback targets from a list of labels. + * + * @param positions + * Label positions in the method. + * @param sourcePosition + * Position of the jump instruction. + * @param labels + * List of jump target labels. + * @param loopHeaders + * List to add detected loop headers to. + * @param addedLoopHeaders + * Set to track already added loop headers and prevent duplicates. + */ + private static void addLoopbackTargets(@Nonnull Map positions, int sourcePosition, + @Nullable List labels, + @Nonnull List loopHeaders, + @Nonnull Set addedLoopHeaders) { + if (labels == null) + return; + for (LabelNode label : labels) + addLoopbackTarget(positions, sourcePosition, label, loopHeaders, addedLoopHeaders); + } + + /** + * Adds a loopback target if it is a loop header. + * + * @param positions + * Label positions in the method. + * @param sourcePosition + * Position of the jump instruction. + * @param label + * Jump target label. + * @param loopHeaders + * List to add detected loop headers to. + * @param addedLoopHeaders + * Set to track already added loop headers and prevent duplicates. + */ + private static void addLoopbackTarget(@Nonnull Map positions, int sourcePosition, + @Nonnull LabelNode label, + @Nonnull List loopHeaders, + @Nonnull Set addedLoopHeaders) { + if (isLoopback(positions, sourcePosition, label) && addedLoopHeaders.add(label)) + loopHeaders.add(label); + } + + /** + * Detects if a jump is a loopback edge. + * + * @param positions + * Label positions in the method. + * @param sourcePosition + * Position of the jump instruction. + * @param target + * Jump target label. + * + * @return {@code true} if the jump is a loopback edge, {@code false} otherwise. + */ + private static boolean isLoopback(@Nonnull Map positions, int sourcePosition, + @Nonnull LabelNode target) { + Integer targetPosition = positions.get(target); + return targetPosition != null && targetPosition < sourcePosition; + } + + /** + * Inserts rethrow logic for cancel signals at the start of catch blocks that catch them. + * + * @param method + * Method to insert into. + */ + private static void insertCancelSignalRethrows(@Nonnull MethodNode method) { + // No catch blocks exist, so no rethrowing needed. + if (method.tryCatchBlocks == null) + return; + + // Insert rethrow logic at the start of any catch block that catches cancel signals. + for (TryCatchBlockNode tryCatchBlock : method.tryCatchBlocks) { + if (!catchesCancelSignal(tryCatchBlock.type)) + continue; + AbstractInsnNode firstRealInsn = nextRealInstruction(tryCatchBlock.handler); + method.instructions.insertBefore(firstRealInsn, newCancelSignalRethrow()); + } + } + + /** + * @param type + * Exception handler type. + * + * @return {@code true} if the catch block catches {@link CancelSignal}. + */ + private static boolean catchesCancelSignal(@Nullable String type) { + return type == null + || "java/lang/Throwable".equals(type) + || "java/lang/Error".equals(type) + || CANCEL_SIGNAL_TYPE.equals(type); + } + + /** + * Finds the appropriate insertion point for a poll check at the start of a loop header. + * + * @param label + * Loop header label. + * + * @return Instruction to use as the insertion point for a poll check. + */ + @Nonnull + private static AbstractInsnNode loopHeaderInsertionPoint(@Nonnull LabelNode label) { + AbstractInsnNode prev = label; + AbstractInsnNode next; + while ((next = prev.getNext()) != null) { + if (next.getOpcode() >= 0) + break; + prev = next; + } + return prev; + } + + /** + * Finds the first real instruction after a label, skipping over any labels, line numbers, or frames. + * + * @param label + * Label to start from. + * + * @return Next real instruction after the label. + */ + @Nonnull + private static AbstractInsnNode nextRealInstruction(@Nonnull LabelNode label) { + AbstractInsnNode next = label; + while ((next = next.getNext()) != null) + if (next.getOpcode() >= 0) + return next; + + // Should never happen in practice. Would imply execution falls off the end of the method. + return label; + } + + /** + * @return Block to call to {@link CancellationSingleton#poll()}. + */ + @Nonnull + private static InsnList newPoll() { + InsnList instructions = new InsnList(); + instructions.add(new MethodInsnNode(INVOKESTATIC, SINGLETON_TYPE, POLL, POLL_DESC, false)); + return instructions; + } + + /** + * @return Block to rethrow an exception if it is an instance of {@link CancelSignal}. + */ + @Nonnull + private static InsnList newCancelSignalRethrow() { + InsnList instructions = new InsnList(); + LabelNode continueLabel = new LabelNode(); + instructions.add(new InsnNode(DUP)); + instructions.add(new org.objectweb.asm.tree.TypeInsnNode(INSTANCEOF, CANCEL_SIGNAL_TYPE)); + instructions.add(new JumpInsnNode(IFEQ, continueLabel)); + instructions.add(new InsnNode(ATHROW)); + instructions.add(continueLabel); + return instructions; + } +} diff --git a/recaf-core/src/main/java/software/coley/recaf/services/script/JavacScriptEngine.java b/recaf-core/src/main/java/software/coley/recaf/services/script/JavacScriptEngine.java index d9cf45ba0..4c6238fbe 100644 --- a/recaf-core/src/main/java/software/coley/recaf/services/script/JavacScriptEngine.java +++ b/recaf-core/src/main/java/software/coley/recaf/services/script/JavacScriptEngine.java @@ -3,19 +3,35 @@ import jakarta.annotation.Nonnull; import jakarta.enterprise.context.ApplicationScoped; import jakarta.inject.Inject; +import org.objectweb.asm.ClassReader; +import org.objectweb.asm.ClassWriter; +import org.objectweb.asm.Type; import regexodus.Matcher; import software.coley.recaf.analytics.logging.DebuggingLogger; import software.coley.recaf.analytics.logging.Logging; -import software.coley.recaf.services.compile.*; +import software.coley.recaf.services.compile.CompilerDiagnostic; +import software.coley.recaf.services.compile.CompilerResult; +import software.coley.recaf.services.compile.JavacArguments; +import software.coley.recaf.services.compile.JavacArgumentsBuilder; +import software.coley.recaf.services.compile.JavacCompiler; import software.coley.recaf.services.plugin.CdiClassAllocator; -import software.coley.recaf.util.ClassDefiner; +import software.coley.recaf.util.CancelSignal; import software.coley.recaf.util.ReflectUtil; import software.coley.recaf.util.RegexUtil; import software.coley.recaf.util.StringUtil; import software.coley.recaf.util.threading.ThreadPoolFactory; +import java.io.IOException; +import java.io.InputStream; +import java.io.UncheckedIOException; +import java.lang.reflect.InvocationTargetException; import java.lang.reflect.Method; -import java.util.*; +import java.util.Arrays; +import java.util.HashMap; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Set; import java.util.concurrent.CompletableFuture; import java.util.concurrent.ExecutorService; import java.util.stream.Collectors; @@ -73,8 +89,9 @@ public class JavacScriptEngine implements ScriptEngine { "jakarta.inject.*", "org.slf4j.Logger" ); - private final Map generateResultMap = new HashMap<>(); - private final ExecutorService compileAndRunPool = ThreadPoolFactory.newSingleThreadExecutor("script-loader"); + private final Map scriptResultMap = new HashMap<>(); + private final ExecutorService compilePool = ThreadPoolFactory.newSingleThreadExecutor("script-compiler"); + private final ExecutorService runPool = ThreadPoolFactory.newCachedThreadPool("script-runner"); private final JavacCompiler compiler; private final CdiClassAllocator allocator; private final ScriptEngineConfig config; @@ -101,28 +118,31 @@ public ScriptEngineConfig getServiceConfig() { @Nonnull @Override public CompletableFuture run(@Nonnull String script) { - return CompletableFuture.supplyAsync(() -> handleExecute(script), compileAndRunPool); + return compile(script).thenCompose(this::run); + } + + @Nonnull + @Override + public CompletableFuture run(@Nonnull GenerateResult result) { + return CompletableFuture.supplyAsync(() -> handleExecute(result), runPool); } @Nonnull @Override public CompletableFuture compile(@Nonnull String scriptSource) { - return CompletableFuture.supplyAsync(() -> generate(scriptSource), compileAndRunPool); + return CompletableFuture.supplyAsync(() -> generate(scriptSource), compilePool); } /** - * Compiles and executes the script. - * If the same script has already been compiled previously, the prior class reference will be used - * to reduce duplicate compilations. + * Executes the generated script. * - * @param script - * Script to execute. + * @param result + * Compiled script to execute. * * @return Result of script execution. */ @Nonnull - private ScriptResult handleExecute(@Nonnull String script) { - GenerateResult result = generate(script); + private ScriptResult handleExecute(@Nonnull GenerateResult result) { if (result.cls() != null) { try { logger.debugging(l -> l.info("Allocating script instance")); @@ -132,9 +152,23 @@ private ScriptResult handleExecute(@Nonnull String script) { run.invoke(instance); logger.debugging(l -> l.info("Successfully ran script")); return new ScriptResult(result.diagnostics()); + } catch (InvocationTargetException ex) { + Throwable target = ex.getTargetException(); + if (target instanceof CancelSignal) { + logger.debugging(l -> l.info("Cancelled script")); + return ScriptResult.cancelled(result.diagnostics()); + } + logger.error("Failed to execute script", target); + return new ScriptResult(result.diagnostics(), target); } catch (Exception ex) { logger.error("Failed to execute script", ex); return new ScriptResult(result.diagnostics(), ex); + } finally { + try { + result.resetStop(); + } catch (IllegalStateException ex) { + logger.error("Failed to reset script cancellation state", ex); + } } } else { logger.error("Failed to compile script"); @@ -155,16 +189,17 @@ private ScriptResult handleExecute(@Nonnull String script) { * * @return Compiler result wrapper containing the loaded class reference. */ + @Nonnull private GenerateResult generate(@Nonnull String script) { int hash = script.hashCode(); GenerateResult result; if (RegexUtil.matchesAny(PATTERN_CLASS_NAME, script)) { logger.debugging(l -> l.info("Compiling script as class")); - result = generateResultMap.computeIfAbsent(hash, n -> generateStandardClass(script)); + result = scriptResultMap.computeIfAbsent(hash, n -> generateStandardClass(script).generateResult()); } else { logger.debugging(l -> l.info("Compiling script as function")); String className = "Script" + Math.abs(hash); - result = generateResultMap.computeIfAbsent(hash, n -> generateScriptClass(className, script)); + result = scriptResultMap.computeIfAbsent(hash, n -> generateScriptClass(className, script).generateResult()); } return result; } @@ -179,7 +214,7 @@ private GenerateResult generate(@Nonnull String script) { * @return Compiler result wrapper containing the loaded class reference. */ @Nonnull - private GenerateResult generateStandardClass(@Nonnull String source) { + private ScriptTemplate generateStandardClass(@Nonnull String source) { String originalSource = source; // Extract package name @@ -212,7 +247,7 @@ private GenerateResult generateStandardClass(@Nonnull String source) { source = source.replace(" " + originalName + "(", " " + modifiedName + "("); source = source.replace("\t" + originalName + "(", "\t" + modifiedName + "("); } else { - return new GenerateResult(null, List.of( + return new ScriptTemplate.Failed(List.of( new CompilerDiagnostic(-1, -1, 0, "Could not determine name of class", CompilerDiagnostic.Level.ERROR))); } @@ -230,10 +265,10 @@ private GenerateResult generateStandardClass(@Nonnull String source) { * @param script * Initial source of the script. * - * @return Compiler result wrapper containing the loaded class reference. + * @return Script wrapper containing the loaded class reference. */ @Nonnull - private GenerateResult generateScriptClass(@Nonnull String className, @Nonnull String script) { + private ScriptTemplate generateScriptClass(@Nonnull String className, @Nonnull String script) { String originalSource = script; Set imports = new HashSet<>(DEFAULT_IMPORTS); Matcher matcher = RegexUtil.getMatcher(PATTERN_IMPORT, script); @@ -252,7 +287,7 @@ private GenerateResult generateScriptClass(@Nonnull String className, @Nonnull S "@Dependent public class " + className + " implements Runnable, Opcodes { " + "private static final Logger log = Logging.get(\"script\"); " + "Workspace workspace; " + - "@Inject " + className +"(Workspace workspace) { this.workspace = workspace; } " + + "@Inject " + className + "(Workspace workspace) { this.workspace = workspace; } " + "public void run() {\n" + script + "\n" + "}" + "}"); for (String imp : imports) @@ -264,6 +299,24 @@ private GenerateResult generateScriptClass(@Nonnull String className, @Nonnull S return generate(className, originalSource, code.toString()); } + @Nonnull + private byte[] postProcessClass(@Nonnull byte[] classBytes) { + ClassReader cr = new ClassReader(classBytes); + ClassWriter cw = new ClassWriter(cr, ClassWriter.COMPUTE_FRAMES | ClassWriter.COMPUTE_MAXS); + cr.accept(new InsertCancelSignalVisitor(cw), ClassReader.EXPAND_FRAMES); + return cw.toByteArray(); + } + + private void injectClasses(@Nonnull Map classMap) { + for (Class c : List.of(CancellationSingleton.class)) { + try (InputStream in = c.getClassLoader().getResourceAsStream(Type.getInternalName(c) + ".class")) { + classMap.put(c.getName(), in.readAllBytes()); + } catch (IOException ex) { + throw new UncheckedIOException(ex); + } + } + } + /** * @param className * Name of the script class. @@ -272,29 +325,25 @@ private GenerateResult generateScriptClass(@Nonnull String className, @Nonnull S * @param compileSource * Full source of the script to pass to the compiler. * - * @return Compiler result wrapper containing the loaded class reference. + * @return Script wrapper containing the loaded class reference. */ @Nonnull - private GenerateResult generate(@Nonnull String className, - @Nonnull String originalSource, - @Nonnull String compileSource) { + private ScriptTemplate generate(@Nonnull String className, + @Nonnull String originalSource, + @Nonnull String compileSource) { JavacArguments args = new JavacArgumentsBuilder() .withClassName(className) .withClassSource(compileSource) .build(); CompilerResult result = compiler.compile(args, null, null); + List diagnostics = mapDiagnostics(originalSource, compileSource, result.getDiagnostics()); if (result.wasSuccess()) { - try { - Map classes = result.getCompilations().entrySet().stream() - .collect(Collectors.toMap(e -> e.getKey().replace('/', '.'), Map.Entry::getValue)); - ClassDefiner definer = new ClassDefiner(classes); - Class cls = definer.findClass(className.replace('/', '.')); - return new GenerateResult(cls, mapDiagnostics(originalSource, compileSource, result.getDiagnostics())); - } catch (Exception ex) { - logger.error("Failed to define generated script class", ex); - } + Map classes = result.getCompilations().entrySet().stream() + .collect(Collectors.toMap(e -> e.getKey().replace('/', '.'), e -> postProcessClass(e.getValue()))); + injectClasses(classes); + return new ScriptTemplate.Generated(className.replace('/', '.'), Map.copyOf(classes), diagnostics); } - return new GenerateResult(null, mapDiagnostics(originalSource, compileSource, result.getDiagnostics())); + return new ScriptTemplate.Failed(diagnostics); } /** @@ -308,8 +357,8 @@ private GenerateResult generate(@Nonnull String className, * @return List of updated diagnostics. */ private List mapDiagnostics(@Nonnull String originalSource, - @Nonnull String compileSource, - @Nonnull List diagnostics) { + @Nonnull String compileSource, + @Nonnull List diagnostics) { int syntheticLineCount = StringUtil.count("\n", StringUtil.cutOffAtFirst(compileSource, originalSource)); return diagnostics.stream() diff --git a/recaf-core/src/main/java/software/coley/recaf/services/script/ScriptEngine.java b/recaf-core/src/main/java/software/coley/recaf/services/script/ScriptEngine.java index b27561979..d826d22c8 100644 --- a/recaf-core/src/main/java/software/coley/recaf/services/script/ScriptEngine.java +++ b/recaf-core/src/main/java/software/coley/recaf/services/script/ScriptEngine.java @@ -22,6 +22,17 @@ public interface ScriptEngine extends Service { @Nonnull CompletableFuture run(@Nonnull String scriptSource); + /** + * @param result + * Compiled script to execute. + * Use this overload when the caller needs to keep the {@link GenerateResult} + * as a cancellation handle for the running script. + * + * @return Future of script execution. + */ + @Nonnull + CompletableFuture run(@Nonnull GenerateResult result); + /** * @param scriptSource * Script source to compile. diff --git a/recaf-core/src/main/java/software/coley/recaf/services/script/ScriptResult.java b/recaf-core/src/main/java/software/coley/recaf/services/script/ScriptResult.java index c8f3236a1..55113ffa8 100644 --- a/recaf-core/src/main/java/software/coley/recaf/services/script/ScriptResult.java +++ b/recaf-core/src/main/java/software/coley/recaf/services/script/ScriptResult.java @@ -7,13 +7,16 @@ import java.util.List; /** - * Wrapper for results of {@link ScriptEngine#run(String)} calls. + * Wrapper for results of script execution calls. * * @author Matt Coley + * @see ScriptEngine#run(String) + * @see ScriptEngine#run(GenerateResult) */ public class ScriptResult { private final List diagnostics; private final Throwable throwable; + private final boolean cancelled; /** * @param diagnostics @@ -30,15 +33,39 @@ public ScriptResult(@Nonnull List diagnostics) { * Runtime error value. */ public ScriptResult(@Nonnull List diagnostics, @Nullable Throwable throwable) { + this(diagnostics, throwable, false); + } + + /** + * @param diagnostics + * Compiler error list. + * @param throwable + * Runtime error value. + * @param cancelled + * Flag indicating cancellation was requested and observed. + */ + public ScriptResult(@Nonnull List diagnostics, @Nullable Throwable throwable, boolean cancelled) { this.diagnostics = diagnostics; this.throwable = throwable; + this.cancelled = cancelled; + } + + /** + * @param diagnostics + * Compiler error list. + * + * @return Cancelled script result. + */ + @Nonnull + public static ScriptResult cancelled(@Nonnull List diagnostics) { + return new ScriptResult(diagnostics, null, true); } /** * @return {@code true} when there were no compiler or runtime errors. */ public boolean wasSuccess() { - return !wasCompileFailure() && !wasRuntimeError(); + return !wasCompileFailure() && !wasRuntimeError() && !wasCancelled(); } /** @@ -55,6 +82,13 @@ public boolean wasRuntimeError() { return throwable != null; } + /** + * @return {@code true} when script execution was cancelled. + */ + public boolean wasCancelled() { + return cancelled; + } + /** * @return List of compiler diagnostics. */ diff --git a/recaf-core/src/main/java/software/coley/recaf/services/script/ScriptTemplate.java b/recaf-core/src/main/java/software/coley/recaf/services/script/ScriptTemplate.java new file mode 100644 index 000000000..df86e66a7 --- /dev/null +++ b/recaf-core/src/main/java/software/coley/recaf/services/script/ScriptTemplate.java @@ -0,0 +1,70 @@ +package software.coley.recaf.services.script; + +import jakarta.annotation.Nonnull; +import software.coley.recaf.analytics.logging.DebuggingLogger; +import software.coley.recaf.analytics.logging.Logging; +import software.coley.recaf.services.compile.CompilerDiagnostic; +import software.coley.recaf.util.ClassDefiner; + +import java.util.List; +import java.util.Map; + +/** + * Result of script template generation. + * + * @author xDark + */ +sealed interface ScriptTemplate { + + /** + * @return Result of the generation process. + */ + @Nonnull + GenerateResult generateResult(); + + /** + * Generation succeeded. A class was generated. + * + * @param className + * Name of the generated class. + * @param classMap + * Map of class names to bytecode for all generated classes. + * The main script class is included in this map. + * @param diagnostics + * Compiler diagnostics from the generation process. + * May include warnings or informational messages. + */ + record Generated(@Nonnull String className, + @Nonnull Map classMap, + @Nonnull List diagnostics) implements ScriptTemplate { + private static final DebuggingLogger logger = Logging.get(Generated.class); + + @Nonnull + @Override + public GenerateResult generateResult() { + try { + ClassDefiner definer = new ClassDefiner(classMap); + Class cls = definer.findClass(className); + return new GenerateResult(cls, diagnostics); + } catch (Exception ex) { + logger.error("Failed to define generated script class", ex); + } + return new GenerateResult(null, diagnostics); + } + } + + /** + * Generation failed. No class was generated. + * + * @param diagnostics + * Compiler diagnostics from the failed generation. + */ + record Failed(@Nonnull List diagnostics) implements ScriptTemplate { + + @Nonnull + @Override + public GenerateResult generateResult() { + return new GenerateResult(null, diagnostics); + } + } +} diff --git a/recaf-core/src/main/java/software/coley/recaf/util/ClassDefiner.java b/recaf-core/src/main/java/software/coley/recaf/util/ClassDefiner.java index a04eb3493..5067e1d4c 100644 --- a/recaf-core/src/main/java/software/coley/recaf/util/ClassDefiner.java +++ b/recaf-core/src/main/java/software/coley/recaf/util/ClassDefiner.java @@ -32,6 +32,26 @@ public ClassDefiner(@Nonnull Map classes) { this.classes = classes; } + @Override + protected final Class loadClass(String name, boolean resolve) throws ClassNotFoundException { + // We must override this to prevent parent-first delegation. + Class c; + if ((c = findLoadedClass(name)) != null) { + return c; + } + synchronized (getClassLoadingLock(name)) { + if ((c = findLoadedClass(name)) != null) { + return c; + } + try { + c = findClass(name); + } catch (ClassNotFoundException e) { + c = super.loadClass(name, resolve); + } + return c; + } + } + @Override public final Class findClass(String name) throws ClassNotFoundException { byte[] bytecode = classes.get(name); diff --git a/recaf-core/src/test/java/software/coley/recaf/services/script/JavacScriptEngineTest.java b/recaf-core/src/test/java/software/coley/recaf/services/script/JavacScriptEngineTest.java index 2a3e5f8a3..d6f5d9092 100644 --- a/recaf-core/src/test/java/software/coley/recaf/services/script/JavacScriptEngineTest.java +++ b/recaf-core/src/test/java/software/coley/recaf/services/script/JavacScriptEngineTest.java @@ -3,20 +3,30 @@ import org.junit.jupiter.api.BeforeAll; import org.junit.jupiter.api.Nested; import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.parallel.Execution; +import org.junit.jupiter.api.parallel.ExecutionMode; import software.coley.recaf.services.compile.CompilerDiagnostic; import software.coley.recaf.test.TestBase; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CountDownLatch; import java.util.concurrent.ExecutionException; import java.util.concurrent.TimeUnit; import java.util.concurrent.TimeoutException; +import java.util.concurrent.atomic.AtomicInteger; import static org.junit.jupiter.api.Assertions.*; /** * Tests for {@link JavacScriptEngine} */ +@Execution(ExecutionMode.SAME_THREAD) public class JavacScriptEngineTest extends TestBase { static JavacScriptEngine engine; + public static CountDownLatch concurrentStarted; + public static CountDownLatch concurrentRelease; + public static CountDownLatch restartStarted; + public static AtomicInteger restartCounter = new AtomicInteger(); @BeforeAll static void setup() { @@ -30,6 +40,129 @@ void testHelloWorld() { assertSuccess("System.out.println(\"hello\");"); } + @Test + void testCancellation() throws ExecutionException, InterruptedException, TimeoutException { + GenerateResult generated = engine.compile(""" + long value = 0; + while (true) { + value++; + } + """).get(5, TimeUnit.SECONDS); + assertTrue(generated.wasSuccess()); + + CompletableFuture run = engine.run(generated); + generated.requestStop(); + + ScriptResult result = run.get(5, TimeUnit.SECONDS); + assertTrue(result.wasCancelled()); + assertFalse(result.wasSuccess()); + assertFalse(result.wasCompileFailure()); + assertFalse(result.wasRuntimeError()); + assertNull(result.getRuntimeThrowable()); + } + + @Test + void testRunAfterCancellation() throws ExecutionException, InterruptedException, TimeoutException { + restartStarted = new CountDownLatch(1); + restartCounter.set(0); + GenerateResult generated = engine.compile(""" + software.coley.recaf.services.script.JavacScriptEngineTest.restartStarted.countDown(); + for (int i = 0; i < 10; i++) { + software.coley.recaf.services.script.JavacScriptEngineTest.restartCounter.incrementAndGet(); + try { + Thread.sleep(10); + } catch (Throwable t) {} + } + """).get(5, TimeUnit.SECONDS); + assertTrue(generated.wasSuccess()); + + CompletableFuture stoppedRun = engine.run(generated); + assertTrue(restartStarted.await(1, TimeUnit.SECONDS)); + generated.requestStop(); + assertTrue(stoppedRun.get(5, TimeUnit.SECONDS).wasCancelled()); + + restartStarted = new CountDownLatch(1); + restartCounter.set(0); + ScriptResult secondRun = engine.run(generated).get(5, TimeUnit.SECONDS); + assertTrue(secondRun.wasSuccess()); + assertEquals(10, restartCounter.get()); + } + + @Test + void testStopAfterRestart() throws ExecutionException, InterruptedException, TimeoutException { + testStopAfterRestart(false); + } + + @Test + void testStopAfterRecompileRestart() throws ExecutionException, InterruptedException, TimeoutException { + testStopAfterRestart(true); + } + + void testStopAfterRestart(boolean recompile) throws ExecutionException, InterruptedException, TimeoutException { + String script = """ + software.coley.recaf.services.script.JavacScriptEngineTest.restartStarted.countDown(); + for (int i = 0; i < 1000; i++) { + software.coley.recaf.services.script.JavacScriptEngineTest.restartCounter.incrementAndGet(); + try { + Thread.sleep(10); + } catch (Throwable t) {} + } + """; + GenerateResult generated = engine.compile(script).get(5, TimeUnit.SECONDS); + assertTrue(generated.wasSuccess()); + + restartStarted = new CountDownLatch(1); + restartCounter.set(0); + CompletableFuture firstRun = engine.run(generated); + assertTrue(restartStarted.await(1, TimeUnit.SECONDS)); + awaitCounter(2); + generated.requestStop(); + assertTrue(firstRun.get(5, TimeUnit.SECONDS).wasCancelled()); + + if (recompile) + generated = engine.compile(script).get(5, TimeUnit.SECONDS); + + restartStarted = new CountDownLatch(1); + restartCounter.set(0); + CompletableFuture secondRun = engine.run(generated); + assertTrue(restartStarted.await(1, TimeUnit.SECONDS)); + awaitCounter(2); + generated.requestStop(); + + ScriptResult result = secondRun.get(5, TimeUnit.SECONDS); + assertTrue(result.wasCancelled()); + assertTrue(restartCounter.get() < 1000); + } + + @Test + void testConcurrentRuns() throws ExecutionException, InterruptedException, TimeoutException { + concurrentStarted = new CountDownLatch(2); + concurrentRelease = new CountDownLatch(1); + String template = """ + software.coley.recaf.services.script.JavacScriptEngineTest.concurrentStarted.countDown(); + try { + software.coley.recaf.services.script.JavacScriptEngineTest.concurrentRelease.await(5, java.util.concurrent.TimeUnit.SECONDS); + } catch (InterruptedException ex) { + throw new RuntimeException(ex); + } + """; + GenerateResult first = engine.compile(template + "\n// first").get(5, TimeUnit.SECONDS); + GenerateResult second = engine.compile(template + "\n// second").get(5, TimeUnit.SECONDS); + assertTrue(first.wasSuccess()); + assertTrue(second.wasSuccess()); + + CompletableFuture firstRun = engine.run(first); + CompletableFuture secondRun = engine.run(second); + try { + assertTrue(concurrentStarted.await(1, TimeUnit.SECONDS), "Scripts did not run concurrently"); + } finally { + concurrentRelease.countDown(); + } + + assertTrue(firstRun.get(5, TimeUnit.SECONDS).wasSuccess()); + assertTrue(secondRun.get(5, TimeUnit.SECONDS).wasSuccess()); + } + @Test void repeatedDefinitions() { String propertyName = "test-xyz"; @@ -154,4 +287,11 @@ static void assertSuccess(String code) { fail(ex); } } + + static void awaitCounter(int min) throws InterruptedException { + long end = System.currentTimeMillis() + TimeUnit.SECONDS.toMillis(1); + while (restartCounter.get() < min && System.currentTimeMillis() < end) + Thread.sleep(10); + assertTrue(restartCounter.get() >= min); + } } diff --git a/recaf-ui/src/main/java/software/coley/recaf/services/script/ScriptRunController.java b/recaf-ui/src/main/java/software/coley/recaf/services/script/ScriptRunController.java new file mode 100644 index 000000000..1449fc442 --- /dev/null +++ b/recaf-ui/src/main/java/software/coley/recaf/services/script/ScriptRunController.java @@ -0,0 +1,261 @@ +package software.coley.recaf.services.script; + +import jakarta.annotation.Nonnull; +import jakarta.enterprise.context.ApplicationScoped; +import jakarta.inject.Inject; +import javafx.beans.property.ReadOnlyBooleanProperty; +import javafx.beans.property.ReadOnlyBooleanWrapper; +import javafx.beans.property.ReadOnlyIntegerProperty; +import javafx.beans.property.ReadOnlyIntegerWrapper; +import javafx.beans.property.ReadOnlyLongProperty; +import javafx.beans.property.ReadOnlyLongWrapper; +import org.slf4j.Logger; +import software.coley.recaf.analytics.logging.Logging; +import software.coley.recaf.util.FxThreadUtil; + +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.concurrent.CompletableFuture; + +/** + * Coordinates script execution state across the UI. + * + * @author Matt Coley + */ +@ApplicationScoped +public class ScriptRunController { + private static final Logger logger = Logging.get(ScriptRunController.class); + private final Object lock = new Object(); + private final Map activeExecutions = new HashMap<>(); + private final ReadOnlyIntegerWrapper activeRunCount = new ReadOnlyIntegerWrapper(); + /** + * In order to support cancellation we currently limit script execution to one at a time via the UI. + * Supporting concurrent runs would require each active control to own and present its own {@link GenerateResult} + * and completion state. If there is a demand for this later we can look into refactoring to support it. + */ + private final ReadOnlyBooleanWrapper running = new ReadOnlyBooleanWrapper(); + private final ReadOnlyLongWrapper executionStateVersion = new ReadOnlyLongWrapper(); + private final ScriptEngine engine; + + @Inject + public ScriptRunController(@Nonnull ScriptEngine engine) { + this.engine = engine; + } + + /** + * @return Observable count of currently active script runs. + */ + @Nonnull + public ReadOnlyIntegerProperty activeRunCountProperty() { + return activeRunCount.getReadOnlyProperty(); + } + + /** + * @return Observable flag for any script execution state. + */ + @Nonnull + public ReadOnlyBooleanProperty runningProperty() { + return running.getReadOnlyProperty(); + } + + /** + * @return Observable version incremented whenever execution state may have changed. + */ + @Nonnull + public ReadOnlyLongProperty executionStateVersionProperty() { + return executionStateVersion.getReadOnlyProperty(); + } + + /** + * @return {@code true} when any script is currently compiling or running. + */ + public boolean isRunning() { + synchronized (lock) { + return !activeExecutions.isEmpty(); + } + } + + /** + * @param key + * Script identity. + * + * @return {@code true} when the script identified by the key is currently compiling or running. + */ + public boolean isRunning(@Nonnull Object key) { + synchronized (lock) { + return activeExecutions.containsKey(key); + } + } + + /** + * Starts a script keyed by its source text. + * + * @param source + * Script source to compile and execute. + * + * @return Future of script execution. + */ + @Nonnull + public CompletableFuture start(@Nonnull String source) { + return start(source, source); + } + + /** + * Starts a script if no other script is already active. + *

+ * This preserves the single active cancellation handle expected by the current UI. + * + * @param key + * Script identity. + * @param source + * Script source to compile and execute. + * + * @return Future of script execution. + */ + @Nonnull + public CompletableFuture start(@Nonnull Object key, @Nonnull String source) { + // Setup tracking for this script execution. + ActiveExecution execution = new ActiveExecution(); + synchronized (lock) { + if (!activeExecutions.isEmpty()) + return CompletableFuture.failedFuture(new IllegalStateException("A script is already running")); + activeExecutions.put(key, execution); + } + updateExecutionState(); + + // Compile and run the script. + engine.compile(source).whenComplete((generated, compileError) -> { + // If compilation failed, complete the future and clear the active execution. + if (compileError != null) { + execution.future.completeExceptionally(compileError); + clearActive(key, execution); + return; + } + + // Check if the execution was cancelled while compiling. + boolean stopRequested; + synchronized (lock) { + if (activeExecutions.get(key) != execution) { + stopRequested = true; + } else { + execution.generateResult = generated; + stopRequested = execution.stopRequested; + } + } + if (stopRequested) + requestStop(generated); + + // Run the script and complete the future when done, then clear the active execution + engine.run(generated).whenComplete((result, runError) -> { + if (runError != null) + execution.future.completeExceptionally(runError); + else + execution.future.complete(result); + clearActive(key, execution); + }); + }); + + return execution.future; + } + + /** + * Requests cancellation of all active scripts. + */ + public void requestStop() { + for (ActiveExecution execution : activeExecutions()) { + GenerateResult generated; + synchronized (lock) { + execution.stopRequested = true; + generated = execution.generateResult; + } + if (generated != null) + requestStop(generated); + } + } + + /** + * Requests cancellation of a script. + * + * @param key + * Script identity. + */ + public void requestStop(@Nonnull Object key) { + GenerateResult generated; + synchronized (lock) { + ActiveExecution execution = activeExecutions.get(key); + if (execution == null) + return; + execution.stopRequested = true; + generated = execution.generateResult; + } + if (generated != null) + requestStop(generated); + } + + /** + * @param generated + * Script generation result to request cancellation of. + * + * @see GenerateResult#requestStop() + */ + private void requestStop(@Nonnull GenerateResult generated) { + try { + generated.requestStop(); + } catch (IllegalStateException ex) { + logger.error("Failed to request script cancellation", ex); + } + } + + /** + * @return List of active executions. + */ + @Nonnull + private List activeExecutions() { + synchronized (lock) { + return new ArrayList<>(activeExecutions.values()); + } + } + + /** + * Clears an execution if it is still active. + * + * @param key + * Script key. + * @param execution + * Active execution to clear. + */ + private void clearActive(@Nonnull Object key, @Nonnull ActiveExecution execution) { + synchronized (lock) { + if (activeExecutions.get(key) != execution) + return; + activeExecutions.remove(key); + } + updateExecutionState(); + } + + /** + * Updates observable execution state properties on the FX thread. + */ + private void updateExecutionState() { + FxThreadUtil.run(() -> { + int activeCount; + synchronized (lock) { + activeCount = activeExecutions.size(); + } + activeRunCount.set(activeCount); + running.set(activeCount > 0); + executionStateVersion.set(executionStateVersion.get() + 1); + }); + } + + /** + * State of an active script execution. + */ + private static class ActiveExecution { + private final CompletableFuture future = new CompletableFuture<>(); + private GenerateResult generateResult; + private boolean stopRequested; + } +} diff --git a/recaf-ui/src/main/java/software/coley/recaf/ui/control/ActionButton.java b/recaf-ui/src/main/java/software/coley/recaf/ui/control/ActionButton.java index 8a6f93ac1..37f278f45 100644 --- a/recaf-ui/src/main/java/software/coley/recaf/ui/control/ActionButton.java +++ b/recaf-ui/src/main/java/software/coley/recaf/ui/control/ActionButton.java @@ -22,6 +22,11 @@ public class ActionButton extends Button implements Tooltipable { private static final ExecutorService service = ThreadPoolFactory.newSingleThreadExecutor("async-button-action"); + /** + * Set nothing initially. + */ + public ActionButton() {} + /** * @param text * Button display text. @@ -153,7 +158,7 @@ public ActionButton width(double width) { return this; } - private static void wrap(ActionEvent e, Runnable action) { + protected static void wrap(ActionEvent e, Runnable action) { // This stops the input from 'bleeding' through to parent control handlers. // - Useful for when the button is used in 'x.setGraphic(button)' scenarios e.consume(); diff --git a/recaf-ui/src/main/java/software/coley/recaf/ui/menubar/ScriptMenu.java b/recaf-ui/src/main/java/software/coley/recaf/ui/menubar/ScriptMenu.java index 8513ed63d..1dd27176e 100644 --- a/recaf-ui/src/main/java/software/coley/recaf/ui/menubar/ScriptMenu.java +++ b/recaf-ui/src/main/java/software/coley/recaf/ui/menubar/ScriptMenu.java @@ -8,9 +8,9 @@ import org.kordamp.ikonli.carbonicons.CarbonIcons; import software.coley.observables.ObservableBoolean; import software.coley.observables.ObservableCollection; -import software.coley.recaf.services.script.ScriptEngine; import software.coley.recaf.services.script.ScriptFile; import software.coley.recaf.services.script.ScriptManager; +import software.coley.recaf.services.script.ScriptRunController; import software.coley.recaf.services.window.WindowManager; import software.coley.recaf.ui.control.FontIconView; import software.coley.recaf.ui.pane.ScriptManagerPane; @@ -34,14 +34,14 @@ public class ScriptMenu extends Menu { private final ObservableCollection> scriptFiles; private final ObservableBoolean hasScripts; private final ScriptManagerPane scriptManagerPane; - private final ScriptEngine engine; + private final ScriptRunController scriptRunController; @Inject - public ScriptMenu(WindowManager windowManager, ScriptManager scriptManager, ScriptEngine engine, + public ScriptMenu(WindowManager windowManager, ScriptManager scriptManager, ScriptRunController scriptRunController, ScriptManagerPane scriptManagerPane) { this.windowManager = windowManager; this.scriptManagerPane = scriptManagerPane; - this.engine = engine; + this.scriptRunController = scriptRunController; scriptFiles = scriptManager.getScriptFiles(); hasScripts = scriptFiles.mapBoolean(List::isEmpty).negated(); @@ -54,6 +54,7 @@ public ScriptMenu(WindowManager windowManager, ScriptManager scriptManager, Scri getItems().add(action("menu.scripting.new", CarbonIcons.ADD_ALT, this::newScript)); hasScripts.addChangeListener((ob, old, cur) -> refreshList()); + scriptRunController.executionStateVersionProperty().addListener((ob, old, cur) -> refreshList()); refreshList(); } @@ -65,8 +66,16 @@ private void refreshList() { menuScripts.getItems().clear(); if (hasScripts.getValue()) { for (ScriptFile scriptFile : new TreeSet<>(scriptFiles)) { - menuScripts.getItems().add(actionLiteral(scriptFile.name(), CarbonIcons.PLAY_FILLED, - () -> scriptFile.execute(engine))); + // Menu entries do not own result presentation, so any entry can cancel the single active run. + boolean running = scriptRunController.isRunning(); + menuScripts.getItems().add(actionLiteral(scriptFile.name(), + running ? CarbonIcons.STOP_FILLED : CarbonIcons.PLAY_FILLED, + () -> { + if (scriptRunController.isRunning()) + scriptRunController.requestStop(); + else + scriptRunController.start(scriptFile.path(), scriptFile.source()); + })); } } else { MenuItem item = new MenuItem(); diff --git a/recaf-ui/src/main/java/software/coley/recaf/ui/pane/ScriptManagerPane.java b/recaf-ui/src/main/java/software/coley/recaf/ui/pane/ScriptManagerPane.java index 486d82863..4f857e114 100644 --- a/recaf-ui/src/main/java/software/coley/recaf/ui/pane/ScriptManagerPane.java +++ b/recaf-ui/src/main/java/software/coley/recaf/ui/pane/ScriptManagerPane.java @@ -3,6 +3,7 @@ import atlantafx.base.controls.Popover; import atlantafx.base.theme.Styles; import jakarta.annotation.Nonnull; +import jakarta.annotation.Nullable; import jakarta.enterprise.context.Dependent; import jakarta.enterprise.inject.Instance; import jakarta.inject.Inject; @@ -33,6 +34,7 @@ import software.coley.recaf.services.script.ScriptFile; import software.coley.recaf.services.script.ScriptManager; import software.coley.recaf.services.script.ScriptManagerConfig; +import software.coley.recaf.services.script.ScriptRunController; import software.coley.recaf.services.window.WindowFactory; import software.coley.recaf.ui.config.KeybindingConfig; import software.coley.recaf.ui.control.ActionButton; @@ -69,6 +71,7 @@ public class ScriptManagerPane extends BorderPane { private final ScriptManager scriptManager; private final ScriptManagerConfig config; private final ScriptEngine engine; + private final ScriptRunController scriptRunController; private final FileTypeSyntaxAssociationService languageAssociation; private final WindowFactory windowFactory; private final RecafDirectoriesConfig directories; @@ -79,6 +82,7 @@ public class ScriptManagerPane extends BorderPane { public ScriptManagerPane(@Nonnull ScriptManagerConfig config, @Nonnull ScriptManager scriptManager, @Nonnull ScriptEngine engine, + @Nonnull ScriptRunController scriptRunController, @Nonnull FileTypeSyntaxAssociationService languageAssociation, @Nonnull WindowFactory windowFactory, @Nonnull RecafDirectoriesConfig directories, @@ -88,6 +92,7 @@ public ScriptManagerPane(@Nonnull ScriptManagerConfig config, this.scriptManager = scriptManager; this.config = config; this.engine = engine; + this.scriptRunController = scriptRunController; this.languageAssociation = languageAssociation; this.directories = directories; this.keys = keys; @@ -269,44 +274,19 @@ private void save() { * Editor component to call {@link ScriptEngine#run(String)}. */ private class RunScriptComponent extends ActionButton { - private RunScriptComponent() { - super(new FontIconView(CarbonIcons.PLAY_FILLED, Color.LIME), Lang.getBinding("menu.scripting.execute"), - () -> { - problemTracking.removeByPhase(ProblemPhase.BUILD); - engine.run(editor.getText()).whenCompleteAsync((result, throwable) -> { - if (result != null && result.wasSuccess()) { - // Don't care about compilation, just wanted to validate it was valid semantics. - Animations.animateSuccess(editor, 1000); - } else { - // Handle compile-result failure, or uncaught thrown exception. - if (result != null) { - for (CompilerDiagnostic diagnostic : result.getCompileDiagnostics()) - problemTracking.addItem(Problem.fromDiagnostic(diagnostic)); - - // Display runtime error if given. - Throwable runtimeThrowable = result.getRuntimeThrowable(); - if (runtimeThrowable != null) { - Label traceString = new Label(StringUtil.traceToString(runtimeThrowable)); - traceString.setGraphic(new FontIconView(CarbonIcons.ERROR, Color.RED)); - - Popover popover = new Popover(); - popover.setArrowLocation(Popover.ArrowLocation.BOTTOM_RIGHT); - popover.setContentNode(traceString); - - // Hack to get self - ObservableList children = editor.getPrimaryStack().getChildrenUnmodifiable(); - popover.show(children.get(children.size() - 1)); - } - } else { - logger.error("Compilation encountered an error", throwable); - } - Animations.animateFailure(editor, 1000); - } + private final Node GRAPHIC_EXEC = new FontIconView(CarbonIcons.PLAY_FILLED, Color.LIME); + private final Node GRAPHIC_STOP = new FontIconView(CarbonIcons.STOP_FILLED, Color.RED); - // Redraw paragraph graphics to update things like in-line problem graphics. - editor.redrawParagraphGraphics(); - }, FxThreadUtil.executor()); - }); + private RunScriptComponent() { + setOnAction(e -> wrap(e, () -> { + if (scriptRunController.isRunning(getScriptKey())) + stop(); + else if (!scriptRunController.isRunning()) + execute(); + })); + scriptRunController.executionStateVersionProperty() + .addListener((ob, old, cur) -> updateDisplay()); + updateDisplay(); // Layout tweaks StackPane.setAlignment(this, Pos.BOTTOM_RIGHT); @@ -314,6 +294,64 @@ private RunScriptComponent() { editor.getVerticalScrollbar().visibleProperty() .addListener((ob, old, cur) -> ScrollbarPaddingUtil.handleScrollbarVisibility(this, cur)); } + + private void execute() { + problemTracking.removeByPhase(ProblemPhase.BUILD); + scriptRunController.start(getScriptKey(), editor.getText()).whenCompleteAsync((result, error) -> { + if (result != null && result.wasCancelled()) + return; + + if (result != null && result.wasSuccess()) { + // Don't care about compilation, just wanted to validate it was valid semantics. + Animations.animateSuccess(editor, 1000); + } else { + // Handle compile-result failure, or uncaught thrown exception. + if (result != null) { + for (CompilerDiagnostic diagnostic : result.getCompileDiagnostics()) + problemTracking.addItem(Problem.fromDiagnostic(diagnostic)); + + // Display runtime error if given. + Throwable runtimeThrowable = result.getRuntimeThrowable(); + if (runtimeThrowable != null) { + Label traceString = new Label(StringUtil.traceToString(runtimeThrowable)); + traceString.setGraphic(new FontIconView(CarbonIcons.ERROR, Color.RED)); + + Popover popover = new Popover(); + popover.setArrowLocation(Popover.ArrowLocation.BOTTOM_RIGHT); + popover.setContentNode(traceString); + + // Hack to get self + ObservableList children = editor.getPrimaryStack().getChildrenUnmodifiable(); + popover.show(children.getLast()); + } + } else { + logger.error("Compilation encountered an error", error); + } + Animations.animateFailure(editor, 1000); + } + + // Redraw paragraph graphics to update things like in-line problem graphics. + editor.redrawParagraphGraphics(); + }, FxThreadUtil.executor()); + } + + private void stop() { + scriptRunController.requestStop(getScriptKey()); + } + + @Nonnull + private Object getScriptKey() { + Path path = scriptPath; + return path == null ? ScriptEditor.this : path; + } + + private void updateDisplay() { + boolean activeScript = scriptRunController.isRunning(getScriptKey()); + boolean otherScriptActive = scriptRunController.isRunning() && !activeScript; + setGraphic(activeScript ? GRAPHIC_STOP : GRAPHIC_EXEC); + setText(Lang.get(activeScript ? "menu.scripting.stop" : "menu.scripting.execute")); + setDisable(otherScriptActive); + } } } @@ -360,16 +398,32 @@ protected String computeValue() { ScriptEntry entry = this; - Button executeButton = new ActionButton(CarbonIcons.PLAY_FILLED_ALT, getBinding("menu.scripting.execute"), () -> { - script.execute(engine) - .whenComplete((result, error) -> { - if (result != null && result.wasSuccess()) { - Animations.animateSuccess(entry, 1000); - } else { - Animations.animateFailure(entry, 1000); - } - }); + Path scriptKey = script.path(); + Button executeButton = new ActionButton(); + executeButton.setOnAction(e -> { + e.consume(); + + // The controller intentionally allows one active run atm, so only the owning row can stop it. + if (scriptRunController.isRunning(scriptKey)) { + scriptRunController.requestStop(scriptKey); + } else if (!scriptRunController.isRunning()) { + scriptRunController.start(scriptKey, script.source()) + .whenCompleteAsync((result, error) -> { + if (result != null && result.wasCancelled()) + return; + if (result != null && result.wasSuccess()) { + Animations.animateSuccess(entry, 1000); + } else { + if (error != null) + logger.error("Script execution encountered an error", error); + Animations.animateFailure(entry, 1000); + } + }, FxThreadUtil.executor()); + } }); + scriptRunController.executionStateVersionProperty() + .addListener((ob, old, cur) -> updateExecuteButton(executeButton, scriptKey)); + updateExecuteButton(executeButton, scriptKey); executeButton.setAlignment(Pos.CENTER_LEFT); executeButton.setPrefSize(130, 30); @@ -388,6 +442,16 @@ protected String computeValue() { prefWidthProperty().bind(widthProperty()); } + private void updateExecuteButton(@Nonnull Button button, @Nonnull Object scriptKey) { + boolean activeScript = scriptRunController.isRunning(scriptKey); + boolean otherScriptActive = scriptRunController.isRunning() && !activeScript; + // Non-owning rows remain execute buttons, but are disabled while another script owns the run handle. + button.setGraphic(new FontIconView(activeScript ? CarbonIcons.STOP_FILLED : CarbonIcons.PLAY_FILLED_ALT, + activeScript ? Color.RED : Color.LIME)); + button.setText(Lang.get(activeScript ? "menu.scripting.stop" : "menu.scripting.execute")); + button.setDisable(otherScriptActive); + } + /** * Used to display bullet point format. * diff --git a/recaf-ui/src/main/resources/translations/en_US.lang b/recaf-ui/src/main/resources/translations/en_US.lang index faea842fc..01548a177 100644 --- a/recaf-ui/src/main/resources/translations/en_US.lang +++ b/recaf-ui/src/main/resources/translations/en_US.lang @@ -97,6 +97,7 @@ menu.scripting.edit=Edit menu.scripting.browse=Browse scripts menu.scripting.save=Save script menu.scripting.execute=Execute +menu.scripting.stop=Stop menu.scripting.editor=Script editor menu.scripting.author=Author menu.scripting.version=Version diff --git a/recaf-ui/src/test/java/software/coley/recaf/services/script/ScriptRunControllerTest.java b/recaf-ui/src/test/java/software/coley/recaf/services/script/ScriptRunControllerTest.java new file mode 100644 index 000000000..655aafd31 --- /dev/null +++ b/recaf-ui/src/test/java/software/coley/recaf/services/script/ScriptRunControllerTest.java @@ -0,0 +1,104 @@ +package software.coley.recaf.services.script; + +import jakarta.annotation.Nonnull; +import org.junit.jupiter.api.BeforeAll; +import org.junit.jupiter.api.Test; +import software.coley.recaf.test.TestBase; + +import java.nio.file.Path; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.TimeUnit; + +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertTrue; + +/** + * Tests for {@link ScriptRunController}. + */ +class ScriptRunControllerTest extends TestBase { + static ScriptRunController controller; + + @BeforeAll + static void setup() { + controller = new ScriptRunController(recaf.get(ScriptEngine.class)); + } + + @Test + void testStopAfterRestart() throws Exception { + Path key = Path.of("test-script"); + String counterProperty = "recaf.script-run-controller.counter"; + String script = loopPrintingScript(counterProperty); + + // Start the script that writes to the given property. + System.clearProperty(counterProperty); + CompletableFuture firstRun = controller.start(key, script); + + // Wait for the script to start and write to the property a few times. + awaitCounter(counterProperty, 2); + + // Verify the script is running and then request it to stop. + assertTrue(controller.isRunning(key)); + controller.requestStop(); + + // Wait for the script to acknowledge the stop request and verify it was canceled. + assertTrue(firstRun.get(5, TimeUnit.SECONDS).wasCancelled()); + assertFalse(controller.isRunning(key)); + + // Start the script again and verify it can run after being stopped. + System.clearProperty(counterProperty); + CompletableFuture secondRun = controller.start(key, script); + awaitCounter(counterProperty, 2); + assertTrue(controller.isRunning(key)); + controller.requestStop(); + + // Wait for the script to acknowledge the stop request and verify it was canceled. + ScriptResult secondResult = secondRun.get(5, TimeUnit.SECONDS); + assertTrue(secondResult.wasCancelled()); + assertFalse(controller.isRunning(key)); + assertTrue(Integer.parseInt(System.getProperty(counterProperty, "100")) < 100); + } + + @Test + void testOnlyOneScriptRunsAtATime() throws Exception { + Path firstKey = Path.of("first-script"); + Path secondKey = Path.of("second-script"); + String counterProperty = "recaf.script-run-controller.concurrent"; + String script = loopPrintingScript(counterProperty); + + // Start the first script that writes to the given property. + System.clearProperty(counterProperty); + CompletableFuture firstRun = controller.start(firstKey, script); + awaitCounter(counterProperty, 2); + + // Attempt to start a second script while the first is still running. This should fail. + CompletableFuture secondRun = controller.start(secondKey, + "System.setProperty(\"recaf.script-run-controller.second\", \"ran\");"); + assertTrue(secondRun.isCompletedExceptionally()); + assertFalse(controller.isRunning(secondKey)); + + // Request the first script to stop and verify it was canceled. + controller.requestStop(); + assertTrue(firstRun.get(5, TimeUnit.SECONDS).wasCancelled()); + assertFalse(controller.isRunning()); + } + + @Nonnull + private static String loopPrintingScript(String counterProperty) { + return """ + for (int i = 0; i < 100; i++) { + System.setProperty("%s", String.valueOf(i)); + try { + Thread.sleep(10); + } catch (Throwable t) {} + } + """.formatted(counterProperty); + } + + @SuppressWarnings("all") + private static void awaitCounter(String property, int min) throws InterruptedException { + long end = System.currentTimeMillis() + TimeUnit.SECONDS.toMillis(1); + while (Integer.parseInt(System.getProperty(property, "-1")) < min && System.currentTimeMillis() < end) + Thread.sleep(10); + assertTrue(Integer.parseInt(System.getProperty(property, "-1")) >= min); + } +}