Fix issues with temporary files going missing (#225)

1 files changed, 72 insertions, 9 deletions

diff --git a/spark-common/src/main/java/me/lucko/spark/common/util/TemporaryFiles.java b/spark-common/src/main/java/me/lucko/spark/common/util/TemporaryFiles.java
index 8a4a621..91a474c 100644
--- a/spark-common/src/main/java/me/lucko/spark/common/util/TemporaryFiles.java
+++ b/spark-common/src/main/java/me/lucko/spark/common/util/TemporaryFiles.java
@@ -20,10 +20,18 @@
 
 package me.lucko.spark.common.util;
 
+import com.google.common.collect.ImmutableList;
+
 import java.io.IOException;
+import java.nio.charset.StandardCharsets;
+import java.nio.file.FileSystems;
 import java.nio.file.Files;
 import java.nio.file.Path;
+import java.nio.file.attribute.FileAttribute;
+import java.nio.file.attribute.PosixFilePermission;
+import java.nio.file.attribute.PosixFilePermissions;
 import java.util.Collections;
+import java.util.EnumSet;
 import java.util.HashSet;
 import java.util.Iterator;
 import java.util.Set;
@@ -32,23 +40,47 @@ import java.util.Set;
  * Utility for handling temporary files.
  */
 public final class TemporaryFiles {
-    private TemporaryFiles() {}
 
-    private static final Set<Path> DELETE_SET = Collections.synchronizedSet(new HashSet<>());
+    public static final FileAttribute<?>[] OWNER_ONLY_FILE_PERMISSIONS;
+
+    static {
+        boolean isPosix = FileSystems.getDefault().supportedFileAttributeViews().contains("posix");
+        if (isPosix) {
+            OWNER_ONLY_FILE_PERMISSIONS = new FileAttribute[]{PosixFilePermissions.asFileAttribute(EnumSet.of(
+                    PosixFilePermission.OWNER_READ, PosixFilePermission.OWNER_WRITE
+            ))};
+        } else {
+            OWNER_ONLY_FILE_PERMISSIONS = new FileAttribute[0];
+        }
+    }
+
+    private final Path tmpDirectory;
+    private final Set<Path> files = Collections.synchronizedSet(new HashSet<>());
 
-    public static Path create(String prefix, String suffix) throws IOException {
-        return register(Files.createTempFile(prefix, suffix));
+    public TemporaryFiles(Path tmpDirectory) {
+        this.tmpDirectory = tmpDirectory;
     }
 
-    public static Path register(Path path) {
+    public Path create(String prefix, String suffix) throws IOException {
+        Path file;
+        if (ensureDirectoryIsReady()) {
+            String name = prefix + Long.toHexString(System.nanoTime()) + suffix;
+            file = Files.createFile(this.tmpDirectory.resolve(name), OWNER_ONLY_FILE_PERMISSIONS);
+        } else {
+            file = Files.createTempFile(prefix, suffix);
+        }
+        return register(file);
+    }
+
+    public Path register(Path path) {
         path.toFile().deleteOnExit();
-        DELETE_SET.add(path);
+        this.files.add(path);
         return path;
     }
 
-    public static void deleteTemporaryFiles() {
-        synchronized (DELETE_SET) {
-            for (Iterator<Path> iterator = DELETE_SET.iterator(); iterator.hasNext(); ) {
+    public void deleteTemporaryFiles() {
+        synchronized (this.files) {
+            for (Iterator<Path> iterator = this.files.iterator(); iterator.hasNext(); ) {
                 Path path = iterator.next();
                 try {
                     Files.deleteIfExists(path);
@@ -60,4 +92,35 @@ public final class TemporaryFiles {
         }
     }
 
+    private boolean ensureDirectoryIsReady() {
+        if (Boolean.parseBoolean(System.getProperty("spark.useOsTmpDir", "false"))) {
+            return false;
+        }
+
+        if (Files.isDirectory(this.tmpDirectory)) {
+            return true;
+        }
+
+        try {
+            Files.createDirectories(this.tmpDirectory);
+
+            Files.write(this.tmpDirectory.resolve("about.txt"), ImmutableList.of(
+                    "# What is this directory?",
+                    "",
+                    "* In order to perform certain functions, spark sometimes needs to write temporary data to the disk. ",
+                    "* Previously, a temporary directory provided by the operating system was used for this purpose. ",
+                    "* However, this proved to be unreliable in some circumstances, so spark now stores temporary data here instead!",
+                    "",
+                    "spark will automatically cleanup the contents of this directory. " ,
+                    "(but if for some reason it doesn't, if the server is stopped, you can freely delete any files ending in .tmp)",
+                    "",
+                    "tl;dr: spark uses this folder to store some temporary data."
+            ), StandardCharsets.UTF_8);
+
+            return true;
+        } catch (IOException e) {
+            return false;
+        }
+    }
+
 }