Move CDS warmup to the CLI directly

The Java class data sharing archive can now be created by the CLI
using the warmup command. This makes it possible to build the CDS
archive during an install step so that the archive uses the right
paths. Creating the archive before distributing the CLI results in
the archive using absolute paths pointing to the where the CLI was
built rather than where it is installed on end user machines.

smithy-cli/build.gradle

@@ -187,26 +187,10 @@ runtime {
     }
 }
 
-// First, call validate with no args and create a class list to use application class data sharing.
-tasks.register("_createClassList", Exec) {
-    environment("SMITHY_OPTS", "-XX:DumpLoadedClassList=${project.buildDir}/image/smithy-cli-${imageOs}/lib/smithy.lst")
-    environment("SMITHY_WARMUP_INTERNAL_ONLY", "true")
-    commandLine("$smithyBinary", "warmup", "--discover")
-}
-
-// Next, actually dump out the archive of the collected classes. This is platform specific,
-// so it can only be done for the current OS+architecture.
-tasks.register("_dumpArchive", Exec) {
-    environment("SMITHY_OPTS", "-Xshare:dump -XX:SharedArchiveFile=${project.buildDir}/image/smithy-cli-${imageOs}/lib/smithy.jsa -XX:SharedClassListFile=${project.buildDir}/image/smithy-cli-${imageOs}/lib/smithy.lst")
-    environment("SMITHY_WARMUP_INTERNAL_ONLY", "true")
-    commandLine("$smithyBinary", "warmup", "--discover")
-}
-
-// Can't do CDS if the OS and architecture is not one of our targets.
-if (!imageOs.isEmpty()) {
-    tasks["_dumpArchive"].dependsOn("_createClassList")
-    tasks["runtime"].finalizedBy("_dumpArchive")
+tasks.register("optimize", Exec) {
+    commandLine("$smithyBinary", "warmup")
 }
+tasks["optimize"].dependsOn("runtime")
 
 // Always shadow the JAR and replace the JAR by the shadowed JAR.
 tasks['jar'].finalizedBy("shadowJar")

smithy-cli/src/it/java/software/amazon/smithy/cli/WarmupTest.java

@@ -3,23 +3,26 @@
 import static org.hamcrest.MatcherAssert.assertThat;
 import static org.hamcrest.Matchers.containsString;
 import static org.hamcrest.Matchers.equalTo;
-import static org.hamcrest.Matchers.not;
 
+import java.io.IOException;
+import java.nio.file.Files;
+import java.nio.file.Path;
 import org.junit.jupiter.api.Test;
 import software.amazon.smithy.utils.ListUtils;
 
 public class WarmupTest {
     @Test
     public void providingNoInputPrintsHelpExits0() {
         IntegUtils.run("simple-config-sources", ListUtils.of("--help"), result -> {
-            assertThat(result.getOutput(), not(containsString("warmup")));
+            assertThat(result.getOutput(), containsString("warmup"));
         });
     }
 
     @Test
-    public void warmupDoesNotWorkWithoutEnvvar() {
-        IntegUtils.run("simple-config-sources", ListUtils.of("warmup"), result -> {
-            assertThat(result.getExitCode(), equalTo(1));
-        });
+    public void canWarmupTheCli() throws IOException {
+        Path tempDir = Files.createTempDirectory("smithy-warmup-integ");
+        RunResult result = IntegUtils.run(tempDir, ListUtils.of("warmup"));
+
+        assertThat(result.getExitCode(), equalTo(0));
     }
 }

smithy-cli/src/main/java/software/amazon/smithy/cli/LoggingUtil.java

@@ -145,6 +145,8 @@ public void publish(LogRecord record) {
                 } else {
                     printer.println(formatted);
                 }
+                // We want to see log messages right away, so flush the printer.
+                printer.flush();
             }
         }
 

smithy-cli/src/main/java/software/amazon/smithy/cli/commands/SmithyCommand.java

@@ -43,7 +43,7 @@ public SmithyCommand(DependencyResolver.Factory dependencyResolverFactory) {
             new SelectCommand(getName(), dependencyResolverFactory),
             new CleanCommand(getName()),
             new Upgrade1to2Command(getName()),
-            new WarmupCommand(getName(), dependencyResolverFactory)
+            new WarmupCommand(getName())
         );
     }
 

smithy-cli/src/main/java/software/amazon/smithy/cli/commands/WarmupCommand.java

@@ -20,22 +20,53 @@
 import java.io.IOException;
 import java.nio.file.Files;
 import java.nio.file.Path;
+import java.nio.file.Paths;
+import java.util.ArrayList;
+import java.util.Collections;
 import java.util.List;
+import java.util.Locale;
+import java.util.function.Consumer;
 import java.util.logging.Logger;
 import software.amazon.smithy.build.model.MavenRepository;
 import software.amazon.smithy.build.model.SmithyBuildConfig;
+import software.amazon.smithy.cli.ArgumentReceiver;
 import software.amazon.smithy.cli.Arguments;
-import software.amazon.smithy.cli.EnvironmentVariable;
+import software.amazon.smithy.cli.CliError;
+import software.amazon.smithy.cli.CliPrinter;
 import software.amazon.smithy.cli.SmithyCli;
+import software.amazon.smithy.cli.StandardOptions;
 import software.amazon.smithy.cli.dependencies.DependencyResolver;
 import software.amazon.smithy.cli.dependencies.MavenDependencyResolver;
+import software.amazon.smithy.utils.IoUtils;
+import software.amazon.smithy.utils.MapUtils;
+import software.amazon.smithy.utils.StringUtils;
 
-final class WarmupCommand extends ClasspathCommand {
+final class WarmupCommand extends SimpleCommand {
 
     private static final Logger LOGGER = Logger.getLogger(WarmupCommand.class.getName());
 
-    WarmupCommand(String parentCommandName, DependencyResolver.Factory dependencyResolverFactory) {
-        super(parentCommandName, dependencyResolverFactory);
+    private enum Phase { WRAPPER, CLASSES, DUMP }
+
+    private static final class Config implements ArgumentReceiver {
+        private Phase phase = Phase.WRAPPER;
+
+        @Override
+        public Consumer<String> testParameter(String name) {
+            if (name.equals("--phase")) {
+                return phase -> this.phase = Phase.valueOf(phase.toUpperCase(Locale.ENGLISH));
+            } else {
+                return null;
+            }
+        }
+    }
+
+    WarmupCommand(String parentCommandName) {
+        super(parentCommandName);
+    }
+
+    @Override
+    protected List<ArgumentReceiver> createArgumentReceivers() {
+        return Collections.singletonList(new Config());
     }
 
     @Override
@@ -45,23 +76,122 @@ public String getName() {
 
     @Override
     public String getSummary() {
-        return "Creates caches for faster subsequent executions.";
+        return "Creates caches that speed up the CLI. This is typically performed during the installation.";
     }
 
     @Override
-    public boolean isHidden() {
-        return true;
+    public int run(Arguments arguments, Env env, List<String> models) {
+        boolean isDebug = arguments.getReceiver(StandardOptions.class).debug();
+        Phase phase = arguments.getReceiver(Config.class).phase;
+        LOGGER.info(() -> "Optimizing the Smithy CLI: " + phase);
+
+        switch (phase) {
+            case WRAPPER:
+                return orchestrate(isDebug, env.stderr());
+            case CLASSES:
+            case DUMP:
+            default:
+                return runCodeToOptimize(arguments, env);
+        }
     }
 
-    @Override
-    int runWithClassLoader(SmithyBuildConfig config, Arguments arguments, Env env, List<String> models) {
-        if (EnvironmentVariable.getByName("SMITHY_WARMUP_INTERNAL_ONLY") == null) {
-            throw new UnsupportedOperationException("The warmup command is for internal use only and may "
-                                                    + "be removed in the future");
+    private int orchestrate(boolean isDebug, CliPrinter printer) {
+        List<String> baseArgs = new ArrayList<>();
+        String classpath = getOrThrowIfUndefinedProperty("java.class.path");
+        Path javaHome = Paths.get(getOrThrowIfUndefinedProperty("java.home"));
+        Path lib = javaHome.resolve("lib");
+        Path bin = javaHome.resolve("bin");
+        Path windowsBinary = bin.resolve("java.exe");
+        Path posixBinary = bin.resolve("java");
+        Path jsaFile = lib.resolve("smithy.jsa");
+        Path classListFile = lib.resolve("classlist");
+
+        // Delete the archive and classlist before regenerating them.
+        classListFile.toFile().delete();
+        jsaFile.toFile().delete();
+
+        if (!Files.isDirectory(bin)) {
+            throw new CliError("$JAVA_HOME/bin directory not found: " + bin);
+        } else if (Files.exists(windowsBinary)) {
+            baseArgs.add(windowsBinary.toString());
+        } else if (Files.exists(posixBinary)) {
+            baseArgs.add(posixBinary.toString());
+        } else {
+            throw new CliError("No java binary found in " + bin);
+        }
+
+        baseArgs.add("-classpath");
+        baseArgs.add(classpath);
+
+        try {
+            // Run the command in a temp directory to avoid building whatever project the cwd might be in.
+            Path baseDir = Files.createTempDirectory("smithy-warmup");
+
+            LOGGER.info("Building class list");
+            callJava(Phase.CLASSES, isDebug, printer, baseDir, baseArgs,
+                     "-Xshare:off", "-XX:DumpLoadedClassList=" + classListFile,
+                     SmithyCli.class.getName(), "warmup");
+
+            LOGGER.info("Building archive from classlist");
+            callJava(Phase.WRAPPER, isDebug, printer, baseDir, baseArgs,
+                     "-XX:SharedClassListFile=" + classListFile, "-Xshare:dump", "-XX:SharedArchiveFile=" + jsaFile,
+                     SmithyCli.class.getName(), "warmup");
+
+            LOGGER.info("Validating that the archive was created correctly");
+            callJava(null, isDebug, printer, baseDir, baseArgs,
+                     "-Xshare:on", "-XX:SharedArchiveFile=" + jsaFile,
+                     SmithyCli.class.getName(), "--help");
+
+            classListFile.toFile().delete();
+            return 0;
+        } catch (IOException e) {
+            throw new CliError("Error running warmup command", 1, e);
+        }
+    }
+
+    private String getOrThrowIfUndefinedProperty(String property) {
+        String result = System.getProperty(property);
+        if (StringUtils.isEmpty(result)) {
+            throw new CliError(result + " system property is not defined");
+        }
+        return result;
+    }
+
+    private void callJava(
+            Phase phase,
+            boolean isDebug,
+            CliPrinter printer,
+            Path baseDir,
+            List<String> baseArgs,
+            String... args) {
+        List<String> resolved = new ArrayList<>(baseArgs);
+        Collections.addAll(resolved, args);
+
+        if (isDebug) {
+            resolved.add("--debug");
         }
 
-        LOGGER.info(() -> "Warming up Smithy CLI");
+        if (phase != null) {
+            resolved.add("--phase");
+            resolved.add(phase.toString());
+        }
+
+        LOGGER.fine(() -> "Running Java command: " + resolved);
+
+        StringBuilder builder = new StringBuilder();
+        int result = IoUtils.runCommand(resolved, baseDir, builder, MapUtils.of());
+
+        // Hide the output unless an error occurred or running in debug mode.
+        if (isDebug || result != 0) {
+            printer.println(builder.toString().trim());
+        }
+
+        if (result != 0) {
+            throw new CliError("Error warming up CLI in phase " + phase, result);
+        }
+    }
 
+    private int runCodeToOptimize(Arguments arguments, Env env) {
         try {
             Path tempDirWithPrefix = Files.createTempDirectory("smithy-warmup");
             DependencyResolver resolver = new MavenDependencyResolver(tempDirWithPrefix.toString());