8240349: jlink should not leave partial image output directory on failure

src/jdk.jlink/share/classes/jdk/tools/jlink/internal/ImageFileCreator.java

@@ -158,8 +158,10 @@ private void writeImage(Set<Archive> archives,
         BasicImageWriter writer = new BasicImageWriter(byteOrder);
         ResourcePoolManager allContent = createPoolManager(archives,
                 entriesForModule, byteOrder, writer);
-        ResourcePool result = generateJImage(allContent,
-             writer, plugins, plugins.getJImageFileOutputStream());
+        ResourcePool result;
+        try (DataOutputStream out = plugins.getJImageFileOutputStream()) {
+            result = generateJImage(allContent, writer, plugins, out);
+        }
 
         //Handle files.
         try {

src/jdk.jlink/share/classes/jdk/tools/jlink/internal/JlinkTask.java

@@ -39,8 +39,11 @@
 import java.net.URI;
 import java.nio.ByteOrder;
 import java.nio.file.Files;
+import java.nio.file.FileVisitResult;
 import java.nio.file.Path;
 import java.nio.file.Paths;
+import java.nio.file.SimpleFileVisitor;
+import java.nio.file.attribute.BasicFileAttributes;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Comparator;
@@ -225,6 +228,7 @@ int run(String[] args) {
             setLog(new PrintWriter(System.out, true),
                    new PrintWriter(System.err, true));
         }
+        Path outputPath = null;
         try {
             List<String> remaining = optionsHelper.handleOptions(this, args);
             if (remaining.size() > 0 && !options.suggestProviders) {
@@ -265,6 +269,7 @@ int run(String[] args) {
             }
 
             JlinkConfiguration config = initJlinkConfig();
+            outputPath = config.getOutput();
             if (options.suggestProviders) {
                 suggestProviders(config, remaining);
             } else {
@@ -279,8 +284,14 @@ int run(String[] args) {
             log.println(taskHelper.getMessage("error.prefix") + " " + e.getMessage());
             e.printStackTrace(log);
             return EXIT_ERROR;
-        } catch (PluginException | IllegalArgumentException |
-                 UncheckedIOException |IOException | ResolutionException e) {
+        } catch (PluginException | UncheckedIOException | IOException e) {
+            log.println(taskHelper.getMessage("error.prefix") + " " + e.getMessage());
+            if (DEBUG) {
+                e.printStackTrace(log);
+            }
+            cleanupOutput(outputPath);
+            return EXIT_ERROR;
+        } catch (IllegalArgumentException | ResolutionException e) {
             log.println(taskHelper.getMessage("error.prefix") + " " + e.getMessage());
             if (DEBUG) {
                 e.printStackTrace(log);
@@ -298,12 +309,26 @@ int run(String[] args) {
         } catch (Throwable x) {
             log.println(taskHelper.getMessage("error.prefix") + " " + x.getMessage());
             x.printStackTrace(log);
+            cleanupOutput(outputPath);
             return EXIT_ABNORMAL;
         } finally {
             log.flush();
         }
     }
 
+    private void cleanupOutput(Path dir) {
+        try {
+            if (dir != null && Files.isDirectory(dir)) {
+                deleteDirectory(dir);
+            }
+        } catch (IOException io) {
+            log.println(taskHelper.getMessage("error.prefix") + " " + io.getMessage());
+            if (DEBUG) {
+                io.printStackTrace(log);
+            }
+        }
+    }
+
     /*
      * Jlink API entry point.
      */
@@ -468,6 +493,28 @@ public static ModuleFinder newModuleFinder(List<Path> paths,
         return finder;
     }
 
+    private static void deleteDirectory(Path dir) throws IOException {
+        Files.walkFileTree(dir, new SimpleFileVisitor<Path>() {
+            @Override
+            public FileVisitResult visitFile(Path file, BasicFileAttributes attrs)
+                    throws IOException {
+                Files.delete(file);
+                return FileVisitResult.CONTINUE;
+            }
+            @Override
+            public FileVisitResult postVisitDirectory(Path dir, IOException e)
+                    throws IOException {
+                if (e == null) {
+                    Files.delete(dir);
+                    return FileVisitResult.CONTINUE;
+                } else {
+                    // directory iteration failed.
+                    throw e;
+                }
+            }
+        });
+    }
+
     private static Path toPathLocation(ResolvedModule m) {
         Optional<URI> ouri = m.reference().location();
         if (!ouri.isPresent())

src/jdk.jlink/share/classes/jdk/tools/jlink/internal/TaskHelper.java

@@ -420,7 +420,7 @@ private PluginsConfiguration getPluginsConfig(Path output, Map<String, String> l
                     ) throws IOException, BadArgs {
             if (output != null) {
                 if (Files.exists(output)) {
-                    throw new PluginException(PluginsResourceBundle.
+                    throw new IllegalArgumentException(PluginsResourceBundle.
                             getMessage("err.dir.already.exits", output));
                 }
             }

test/jdk/tools/jlink/JLinkNegativeTest.java

@@ -130,6 +130,9 @@ public void testOutputIsFile() throws IOException {
                 .output(image)
                 .addMods("leaf1")
                 .call().assertFailure("Error: directory already exists: .*failure4.image(\n|\r|.)*");
+        if (Files.notExists(image)) {
+            throw new RuntimeException("output directory should not have been deleted");
+        }
     }
 
     public void testModuleNotFound() {

test/jdk/tools/jlink/JLinkTest.java

@@ -47,6 +47,7 @@
  * @bug 8189777
  * @bug 8194922
  * @bug 8206962
+ * @bug 8240349
  * @author Jean-Francois Denise
  * @requires (vm.compMode != "Xcomp" & os.maxMemory >= 2g)
  * @library ../lib
@@ -366,6 +367,19 @@ public static void main(String[] args) throws Exception {
                     .option("--release-info=del")
                     .call().assertFailure("Error: No key specified for delete");
         }
+        {
+            String imageDir = "bug8240349";
+            Path imagePath = helper.createNewImageDir(imageDir);
+            JImageGenerator.getJLinkTask()
+                    .modulePath(helper.defaultModulePath())
+                    .output(imagePath)
+                    .addMods("java.base")
+                    .option("--vm=client")
+                    .call().assertFailure("Error: Selected VM client doesn't exist");
+            if (!Files.notExists(imagePath)) {
+                throw new RuntimeException("bug8240349 directory not deleted");
+            }
+        }
     }
 
     private static void testCompress(Helper helper, String moduleName, String... userOptions) throws IOException {