Cleanup spark tmp dir on startup

3 files changed, 122 insertions, 20 deletions

diff --git a/spark-common/src/main/java/me/lucko/spark/common/SparkPlatform.java b/spark-common/src/main/java/me/lucko/spark/common/SparkPlatform.java
index f5886b5..11c419c 100644
--- a/spark-common/src/main/java/me/lucko/spark/common/SparkPlatform.java
+++ b/spark-common/src/main/java/me/lucko/spark/common/SparkPlatform.java
@@ -45,6 +45,7 @@ import me.lucko.spark.common.monitor.ping.PingStatistics;
 import me.lucko.spark.common.monitor.ping.PlayerPingProvider;
 import me.lucko.spark.common.monitor.tick.SparkTickStatistics;
 import me.lucko.spark.common.monitor.tick.TickStatistics;
+import me.lucko.spark.common.platform.PlatformInfo;
 import me.lucko.spark.common.platform.PlatformStatisticsProvider;
 import me.lucko.spark.common.sampler.BackgroundSamplerManager;
 import me.lucko.spark.common.sampler.SamplerContainer;
@@ -123,7 +124,10 @@ public class SparkPlatform {
         this.plugin = plugin;
         SparkStaticLogger.setLogger(plugin::log);
 
-        this.temporaryFiles = new TemporaryFiles(this.plugin.getPluginDirectory().resolve("tmp"));
+        this.temporaryFiles = new TemporaryFiles(this.plugin.getPlatformInfo().getType() == PlatformInfo.Type.CLIENT
+                ? this.plugin.getPluginDirectory().resolve("tmp")
+                : this.plugin.getPluginDirectory().resolve("tmp-client")
+        );
         this.configuration = Configuration.combining(
                 RuntimeConfiguration.SYSTEM_PROPERTIES,
                 RuntimeConfiguration.ENVIRONMENT_VARIABLES,
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 91a474c..01dfccf 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
@@ -25,8 +25,11 @@ import com.google.common.collect.ImmutableList;
 import java.io.IOException;
 import java.nio.charset.StandardCharsets;
 import java.nio.file.FileSystems;
+import java.nio.file.FileVisitResult;
 import java.nio.file.Files;
 import java.nio.file.Path;
+import java.nio.file.SimpleFileVisitor;
+import java.nio.file.attribute.BasicFileAttributes;
 import java.nio.file.attribute.FileAttribute;
 import java.nio.file.attribute.PosixFilePermission;
 import java.nio.file.attribute.PosixFilePermissions;
@@ -58,16 +61,21 @@ public final class TemporaryFiles {
     private final Set<Path> files = Collections.synchronizedSet(new HashSet<>());
 
     public TemporaryFiles(Path tmpDirectory) {
-        this.tmpDirectory = tmpDirectory;
+        boolean useOsTmpDir = Boolean.parseBoolean(System.getProperty("spark.useOsTmpDir", "false"));
+        if (useOsTmpDir) {
+            this.tmpDirectory = null;
+        } else {
+            this.tmpDirectory = init(tmpDirectory);
+        }
     }
 
     public Path create(String prefix, String suffix) throws IOException {
         Path file;
-        if (ensureDirectoryIsReady()) {
+        if (this.tmpDirectory == null) {
+            file = Files.createTempFile(prefix, suffix);
+        } else {
             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);
     }
@@ -92,19 +100,33 @@ public final class TemporaryFiles {
         }
     }
 
-    private boolean ensureDirectoryIsReady() {
-        if (Boolean.parseBoolean(System.getProperty("spark.useOsTmpDir", "false"))) {
-            return false;
-        }
-
-        if (Files.isDirectory(this.tmpDirectory)) {
-            return true;
-        }
-
+    private static Path init(Path tmpDirectory) {
         try {
-            Files.createDirectories(this.tmpDirectory);
-
-            Files.write(this.tmpDirectory.resolve("about.txt"), ImmutableList.of(
+            Files.createDirectories(tmpDirectory);
+            Path readmePath = tmpDirectory.resolve("about.txt");
+
+            Files.walkFileTree(
+                    tmpDirectory,
+                    new SimpleFileVisitor<Path>() {
+                        @Override
+                        public FileVisitResult postVisitDirectory(Path dir, IOException exc) throws IOException {
+                            if (!dir.equals(tmpDirectory)) {
+                                Files.delete(dir);
+                            }
+                            return FileVisitResult.CONTINUE;
+                        }
+
+                        @Override
+                        public FileVisitResult visitFile(Path file, BasicFileAttributes attrs) throws IOException {
+                            if (!file.equals(readmePath)) {
+                                Files.delete(file);
+                            }
+                            return FileVisitResult.CONTINUE;
+                        }
+                    }
+            );
+
+            Files.write(readmePath, ImmutableList.of(
                     "# What is this directory?",
                     "",
                     "* In order to perform certain functions, spark sometimes needs to write temporary data to the disk. ",
@@ -116,11 +138,10 @@ public final class TemporaryFiles {
                     "",
                     "tl;dr: spark uses this folder to store some temporary data."
             ), StandardCharsets.UTF_8);
-
-            return true;
         } catch (IOException e) {
-            return false;
+            // ignore
         }
+        return tmpDirectory;
     }
 
 }
diff --git a/spark-common/src/test/java/me/lucko/spark/common/util/TemporaryFilesTest.java b/spark-common/src/test/java/me/lucko/spark/common/util/TemporaryFilesTest.java
new file mode 100644
index 0000000..9e07381
--- /dev/null
+++ b/spark-common/src/test/java/me/lucko/spark/common/util/TemporaryFilesTest.java
@@ -0,0 +1,77 @@
+/*
+ * This file is part of spark.
+ *
+ *  Copyright (c) lucko (Luck) <luck@lucko.me>
+ *  Copyright (c) contributors
+ *
+ *  This program is free software: you can redistribute it and/or modify
+ *  it under the terms of the GNU General Public License as published by
+ *  the Free Software Foundation, either version 3 of the License, or
+ *  (at your option) any later version.
+ *
+ *  This program is distributed in the hope that it will be useful,
+ *  but WITHOUT ANY WARRANTY; without even the implied warranty of
+ *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ *  GNU General Public License for more details.
+ *
+ *  You should have received a copy of the GNU General Public License
+ *  along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ */
+
+package me.lucko.spark.common.util;
+
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.io.TempDir;
+
+import java.io.IOException;
+import java.nio.file.Files;
+import java.nio.file.Path;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+public class TemporaryFilesTest {
+
+    @Test
+    public void testDelete(@TempDir Path tempDir) throws IOException {
+        Path dir = tempDir.resolve("test");
+        TemporaryFiles temporaryFiles = new TemporaryFiles(dir);
+
+        assertTrue(Files.exists(dir) && Files.isDirectory(dir));
+        assertTrue(Files.exists(dir.resolve("about.txt")));
+        assertEquals("# What is this directory?", Files.readAllLines(dir.resolve("about.txt")).get(0));
+
+        Path temporaryFile = temporaryFiles.create("test", ".txt");
+        Files.write(temporaryFile, "Hello, world!".getBytes());
+
+        assertTrue(Files.exists(temporaryFile));
+        temporaryFiles.deleteTemporaryFiles();
+        assertFalse(Files.exists(temporaryFile));
+    }
+
+    @Test
+    public void testCleanupOnInit(@TempDir Path tempDir) throws IOException {
+        Path dir = tempDir.resolve("test");
+
+        Path nestedDirectory = dir.resolve("hello").resolve("world");
+        Files.createDirectories(nestedDirectory);
+
+        Path testFile = nestedDirectory.resolve("file.txt");
+        Files.write(testFile, "Hello, world!".getBytes());
+        assertTrue(Files.exists(testFile));
+
+        TemporaryFiles temporaryFiles = new TemporaryFiles(dir);
+
+        assertFalse(Files.exists(testFile));
+    }
+
+    @Test
+    public void testSecondInit(@TempDir Path tempDir) throws IOException {
+        Path dir = tempDir.resolve("test");
+
+        TemporaryFiles temporaryFiles = new TemporaryFiles(dir);
+        TemporaryFiles temporaryFiles2 = new TemporaryFiles(dir);
+    }
+
+}