Better lifecycles of tool daemons

Author: lukebemish

build.gradle

@@ -235,14 +235,14 @@ dependencies {
     compileOnly cLibs.bundles.compileonly
     annotationProcessor cLibs.bundles.annotationprocessor
 
-    implementation 'info.picocli:picocli:4.7.6'
-    implementation 'org.slf4j:slf4j-simple:2.0.13'
-    implementation 'org.ow2.asm:asm:9.7'
-    implementation 'dev.lukebemish.mirror.net.fabricmc:mapping-io:0.6.2'
+    implementation libs.picocli.core
+    implementation libs.slf4j.simple
+    implementation libs.asm
+    implementation libs.mappingio
     // For reading/writing parchment json files
-    implementation "org.parchmentmc:feather:1.1.0"
-    implementation "org.parchmentmc.feather:io-gson:1.1.0"
-    annotationProcessor 'info.picocli:picocli-codegen:4.7.6'
+    implementation libs.feather.core
+    implementation libs.feather.gson
+    annotationProcessor libs.picocli.codegen
     implementation(project(':')) {
         capabilities {
             requireCapability("dev.lukebemish:taskgraphrunner-model")
@@ -253,32 +253,32 @@ dependencies {
             requireCapability("dev.lukebemish:taskgraphrunner-signatures")
         }
     }
-    modelApi 'com.google.code.gson:gson:2.11.0'
+    modelApi libs.gson
     modelCompileOnlyApi cLibs.jspecify
     modelCompileOnly cLibs.bundles.compileonly
     modelAnnotationProcessor cLibs.bundles.annotationprocessor
 
-    signaturesApi 'org.ow2.asm:asm:9.7'
+    signaturesApi libs.asm
     signaturesCompileOnly cLibs.bundles.compileonly
     signaturesAnnotationProcessor cLibs.bundles.annotationprocessor
 
     executionCompileOnly cLibs.bundles.compileonly
     executionAnnotationProcessor cLibs.bundles.annotationprocessor
-    executionImplementation('dev.lukebemish:forkedtaskexecutor:0.1.1') {
+    executionImplementation(libs.forkedtaskexecutor) {
         capabilities {
             requireFeature('runner')
         }
     }
 
-    implementation 'dev.lukebemish:forkedtaskexecutor:0.1.1'
-    implementation('dev.lukebemish:forkedtaskexecutor:0.1.1') {
+    implementation libs.forkedtaskexecutor
+    implementation(libs.forkedtaskexecutor) {
         capabilities {
             requireFeature('runner')
         }
     }
-    modelApi 'dev.lukebemish:forkedtaskexecutor:0.1.1'
+    modelApi libs.forkedtaskexecutor
 
-    instrumentationImplementation 'org.ow2.asm:asm:9.7'
+    instrumentationImplementation libs.asm
 }
 
 // This keeps the same versions at runtime and compile time

gradle/libs.versions.toml

@@ -0,0 +1,26 @@
+[versions]
+
+forkedtaskexecutor = "0.1.3"
+picocli = "4.7.6"
+slf4j = "2.0.13"
+asm = "9.7"
+mappingio = "0.6.2"
+feather = "1.1.0"
+gson = "2.11.0"
+
+[libraries]
+
+forkedtaskexecutor = { group = "dev.lukebemish", name = "forkedtaskexecutor", version.ref = "forkedtaskexecutor" }
+picocli-core = { group = "info.picocli", name = "picocli", version.ref = "picocli" }
+picocli-codegen = { group = "info.picocli", name = "picocli-codegen", version.ref = "picocli" }
+
+slf4j-simple = { group = "org.slf4j", name = "slf4j-simple", version.ref = "slf4j" }
+
+asm = { group = "org.ow2.asm", name = "asm", version.ref = "asm" }
+
+mappingio = { group = "dev.lukebemish.mirror.net.fabricmc", name = "mapping-io", version.ref = "mappingio" }
+
+feather-core = { group = "org.parchmentmc", name = "feather", version.ref = "feather" }
+feather-gson = { group = "org.parchmentmc.feather", name = "io-gson", version.ref = "feather" }
+
+gson = { group = "com.google.code.gson", name = "gson", version.ref = "gson" }

src/execution/java/dev/lukebemish/taskgraphrunner/execution/ToolTask.java

@@ -17,6 +17,7 @@
 import java.net.URLClassLoader;
 import java.nio.charset.StandardCharsets;
 import java.nio.file.Paths;
+import java.time.Duration;
 import java.util.Arrays;
 import java.util.HashSet;
 import java.util.Set;
@@ -65,6 +66,35 @@ public PrintStream replaceSystemErr(PrintStream err) {
         return new PrintStream(OutputStream.nullOutputStream());
     }
 
+    @Override
+    public void setupLifecycleWatcher(Supplier<Integer> currentTasks, Supplier<Boolean> attemptShutdown) {
+        var thread = new Thread("TaskGraphRunner Tool Daemon Watcher") {
+            @Override
+            public void run() {
+                int emptyFor = 0;
+                while (true) {
+                    try {
+                        Thread.sleep(Duration.ofSeconds(1));
+                    } catch (InterruptedException e) {
+                        break;
+                    }
+                    if (currentTasks.get() == 0) {
+                        emptyFor++;
+                    } else {
+                        emptyFor = 0;
+                    }
+                    if (emptyFor > 10) {
+                        if (attemptShutdown.get()) {
+                            break;
+                        }
+                    }
+                }
+            }
+        };
+        thread.setDaemon(true);
+        thread.start();
+    }
+
     private record Target(URL[] urls, String mainClass) {}
 
     private int execute(String[] args, Supplier<Target> targetSupplier, String logFile) {

src/main/java/dev/lukebemish/taskgraphrunner/runtime/execution/ToolDaemonExecutor.java

@@ -6,6 +6,7 @@
 import dev.lukebemish.taskgraphrunner.runtime.Context;
 import dev.lukebemish.taskgraphrunner.runtime.util.FileUtils;
 import dev.lukebemish.taskgraphrunner.runtime.util.HashUtils;
+import org.jspecify.annotations.Nullable;
 import org.objectweb.asm.ClassReader;
 import org.objectweb.asm.ClassVisitor;
 import org.objectweb.asm.ClassWriter;
@@ -19,6 +20,7 @@
 import java.io.File;
 import java.io.IOException;
 import java.io.UncheckedIOException;
+import java.lang.ref.WeakReference;
 import java.nio.ByteBuffer;
 import java.nio.charset.StandardCharsets;
 import java.nio.file.Files;
@@ -31,7 +33,10 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
+import java.util.Queue;
 import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.ConcurrentLinkedQueue;
+import java.util.function.Consumer;
 import java.util.jar.JarInputStream;
 import java.util.jar.JarOutputStream;
 import java.util.stream.Collectors;
@@ -69,13 +74,13 @@ public static void execute(Collection<Path> classpath, String mainClass, Path lo
     private static final Map<String, ToolDaemonExecutor> CLASSPATH_INSTANCES = new ConcurrentHashMap<>();
 
     private final ForkedTaskExecutor executor;
-    private final Runnable onRemoval;
+    private final Consumer<ToolDaemonExecutor> onRemoval;
 
-    private ToolDaemonExecutor(Runnable onRemoval) throws IOException {
+    private ToolDaemonExecutor(Consumer<ToolDaemonExecutor> onRemoval) throws IOException {
         this(new Path[0], onRemoval);
     }
 
-    private ToolDaemonExecutor(Path[] classpath, Runnable onRemoval) throws IOException {
+    private ToolDaemonExecutor(Path[] classpath, Consumer<ToolDaemonExecutor> onRemoval) throws IOException {
         this.onRemoval = onRemoval;
         var workingDirectory = Files.createTempDirectory("taskgraphrunner");
         var tempFile = workingDirectory.resolve("execution-daemon.jar");
@@ -110,11 +115,14 @@ private ToolDaemonExecutor(Path[] classpath, Runnable onRemoval) throws IOExcept
             .addJvmOption("-cp")
             .addJvmOption("@"+ classpathFile.toAbsolutePath())
             .taskClass("dev.lukebemish.taskgraphrunner.execution.ToolTask")
+            .onShutdownRequest(() -> onRemoval.accept(this)) // When the daemon thinks it should be shut down, we stop any new requests from going to it.
             .build();
 
         this.executor = new ForkedTaskExecutor(spec);
     }
 
+    private static final Queue<WeakReference<ToolDaemonExecutor>> TO_CLOSE = new ConcurrentLinkedQueue<>();
+
     static {
         Runtime.getRuntime().addShutdownHook(new Thread(() -> {
             if (INSTANCE != null) {
@@ -123,30 +131,54 @@ private ToolDaemonExecutor(Path[] classpath, Runnable onRemoval) throws IOExcept
             }
             CLASSPATH_INSTANCES.values().forEach(ToolDaemonExecutor::close);
             CLASSPATH_INSTANCES.clear();
+            WeakReference<ToolDaemonExecutor> ref;
+            while ((ref = TO_CLOSE.poll()) != null) {
+                var executor = ref.get();
+                if (executor != null) {
+                    executor.close();
+                }
+            }
         }));
     }
 
     @Override
     public void close() {
-        onRemoval.run();
+        onRemoval.accept(this);
         executor.close();
     }
 
     private static synchronized ToolDaemonExecutor getInstance(Path[] classpath) {
-        var key = String.join(File.pathSeparator, Arrays.stream(classpath).map(it -> it.toAbsolutePath().toString()).toArray(CharSequence[]::new));
+        var key = key(classpath);
         return CLASSPATH_INSTANCES.computeIfAbsent(key, k -> {
             try {
-                return new ToolDaemonExecutor(classpath, () -> CLASSPATH_INSTANCES.remove(key));
+                return new ToolDaemonExecutor(classpath, it -> clearInstance(it, classpath));
             } catch (IOException e) {
                 throw new RuntimeException(e);
             }
         });
     }
 
+    private static String key(Path[] classpath) {
+        return String.join(File.pathSeparator, Arrays.stream(classpath).map(it -> it.toAbsolutePath().toString()).toArray(CharSequence[]::new));
+    }
+
+    private static synchronized void clearInstance(ToolDaemonExecutor it, Path @Nullable [] classpath) {
+        if (classpath == null && INSTANCE == it) {
+            var oldInstance = INSTANCE;
+            TO_CLOSE.add(new WeakReference<>(oldInstance));
+            INSTANCE = null;
+        } else {
+            var key = key(classpath);
+            if (CLASSPATH_INSTANCES.remove(key, it)) {
+                TO_CLOSE.add(new WeakReference<>(it));
+            }
+        }
+    }
+
     private static synchronized ToolDaemonExecutor getInstance() {
         if (INSTANCE == null) {
             try {
-                INSTANCE = new ToolDaemonExecutor(() -> INSTANCE = null);
+                INSTANCE = new ToolDaemonExecutor(it -> clearInstance(it, null));
             } catch (IOException e) {
                 throw new RuntimeException(e);
             }

src/model/java/dev/lukebemish/taskgraphrunner/model/conversion/NeoFormGenerator.java

@@ -322,7 +322,7 @@ public static Config convert(Path neoFormConfig, Value selfReference, Options op
                     }
                     List<Argument> args;
                     TaskModel tool;
-                    if (function.jvmArgs().isEmpty()) {
+                    if (!function.jvmArgs().isEmpty()) {
                         TaskModel.Tool toolModel;
                         tool = toolModel = new TaskModel.Tool(step.name(), List.of());
 
@@ -547,7 +547,9 @@ private static Argument processArgument(Map<String, Input> inputs, String arg, N
             var function = fullConfig.functions().get(step.type());
             if (isVineflower(function)) {
                 arg = arg.replace("TRACE", "WARN");
-            } else if (arg.startsWith("-Xmx")) {
+            }
+
+            if (arg.startsWith("-Xmx")) {
                 var rest = arg.substring(4);
                 var systemProp = "dev.lukebemish.taskgraphrunner."+step.type()+".maxHeap";
                 return new Argument.Untracked("-Xmx{}", new InputValue.DirectInput(new Value.SystemPropertyValue(systemProp, rest)));
@@ -560,6 +562,7 @@ private static Argument processArgument(Map<String, Input> inputs, String arg, N
                 var systemProp = "dev.lukebemish.taskgraphrunner."+step.type()+".maxThreads";
                 return new Argument.Untracked("-thr={}", new InputValue.DirectInput(new Value.SystemPropertyValue(systemProp, rest)));
             }
+
             return new Argument.ValueInput(null, new InputValue.DirectInput(new Value.DirectStringValue(arg)));
         }
     }