Improve the temp location manager api

Author: jbachorik

dd-java-agent/agent-profiling/profiling-controller/src/main/java/com/datadog/profiling/controller/TempLocationManager.java

@@ -158,6 +158,8 @@ private TempLocationManager() {
     String pid = PidHelper.getPid();
 
     baseTempDir = configuredTempDir.resolve("ddprof");
+    baseTempDir.toFile().deleteOnExit();
+
     tempDir = baseTempDir.resolve("pid_" + pid);
     cleanupTask = CompletableFuture.runAsync(() -> cleanup(false));
 
@@ -180,28 +182,42 @@ private TempLocationManager() {
    * @return the temporary directory for the current process
    */
   public Path getTempDir() {
-    return getTempDir(true);
+    return getTempDir(null);
+  }
+
+  /**
+   * Get the temporary subdirectory for the current process. The directory will be removed at JVM
+   * exit.
+   *
+   * @param subPath the relative subdirectory path, may be {@literal null}
+   * @return the temporary subdirectory for the current process
+   */
+  public Path getTempDir(Path subPath) {
+    return getTempDir(subPath, true);
   }
 
   /**
-   * Get the temporary directory for the current process.
+   * Get the temporary subdirectory for the current process.
    *
+   * @param subPath the relative subdirectory path, may be {@literal null}
    * @param create if true, create the directory if it does not exist
    * @return the temporary directory for the current process
    * @throws IllegalStateException if the directory could not be created
    */
-  public Path getTempDir(boolean create) {
-    if (create && !Files.exists(tempDir)) {
+  public Path getTempDir(Path subPath, boolean create) {
+    Path rslt =
+        subPath != null && !subPath.toString().isEmpty() ? tempDir.resolve(subPath) : tempDir;
+    if (create && !Files.exists(rslt)) {
       try {
         Files.createDirectories(
-            tempDir,
+            rslt,
             PosixFilePermissions.asFileAttribute(PosixFilePermissions.fromString("rwx------")));
       } catch (Exception e) {
         log.warn("Failed to create temp directory: {}", tempDir, e);
         throw new IllegalStateException("Failed to create temp directory: " + tempDir, e);
       }
     }
-    return tempDir;
+    return rslt;
   }
 
   void cleanup(boolean cleanSelf) {

dd-java-agent/agent-profiling/profiling-controller/src/test/java/com/datadog/profiling/controller/TempLocationManagerTest.java

@@ -12,32 +12,37 @@
 import java.util.Properties;
 import java.util.UUID;
 import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.ValueSource;
 
 public class TempLocationManagerTest {
-  @Test
-  void testDefault() throws Exception {
+  @ParameterizedTest
+  @ValueSource(strings = {"", "test1"})
+  void testDefault(String subPath) throws Exception {
     TempLocationManager tempLocationManager = TempLocationManager.getInstance(true);
-    Path tempDir = tempLocationManager.getTempDir();
+    Path tempDir = tempLocationManager.getTempDir(Paths.get(subPath));
     assertNotNull(tempDir);
     assertTrue(Files.exists(tempDir));
     assertTrue(Files.isDirectory(tempDir));
     assertTrue(Files.isWritable(tempDir));
     assertTrue(Files.isReadable(tempDir));
     assertTrue(Files.isExecutable(tempDir));
-    assertTrue(tempDir.endsWith("pid_" + PidHelper.getPid()));
+    assertTrue(tempDir.toString().contains("pid_" + PidHelper.getPid()));
   }
 
-  @Test
-  void testFromConfig() throws Exception {
+  @ParameterizedTest
+  @ValueSource(strings = {"", "test1"})
+  void testFromConfig(String subPath) throws Exception {
     Path myDir =
         Files.createTempDirectory(
             "ddprof-test-",
             PosixFilePermissions.asFileAttribute(PosixFilePermissions.fromString("rwx------")));
+    myDir.toFile().deleteOnExit();
     Properties props = new Properties();
     props.put(ProfilingConfig.PROFILING_TEMP_DIR, myDir.toString());
     ConfigProvider configProvider = ConfigProvider.withPropertiesOverride(props);
     TempLocationManager tempLocationManager = new TempLocationManager(configProvider);
-    Path tempDir = tempLocationManager.getTempDir();
+    Path tempDir = tempLocationManager.getTempDir(Paths.get(subPath));
     assertNotNull(tempDir);
     assertTrue(tempDir.toString().startsWith(myDir.toString()));
   }
@@ -58,29 +63,42 @@ void testFromConfigNotWritable() throws Exception {
         Files.createTempDirectory(
             "ddprof-test-",
             PosixFilePermissions.asFileAttribute(PosixFilePermissions.fromString("r-x------")));
+    myDir.toFile().deleteOnExit();
     Properties props = new Properties();
     props.put(ProfilingConfig.PROFILING_TEMP_DIR, myDir.toString());
     ConfigProvider configProvider = ConfigProvider.withPropertiesOverride(props);
     TempLocationManager tempLocationManager = new TempLocationManager(configProvider);
     assertThrows(IllegalStateException.class, tempLocationManager::getTempDir);
   }
 
-  @Test
-  void testCleanup() throws Exception {
+  @ParameterizedTest
+  @ValueSource(strings = {"", "test1"})
+  void testCleanup(String subPath) throws Exception {
     Path myDir =
         Files.createTempDirectory(
             "ddprof-test-",
             PosixFilePermissions.asFileAttribute(PosixFilePermissions.fromString("rwx------")));
+    myDir.toFile().deleteOnExit();
     Properties props = new Properties();
     props.put(ProfilingConfig.PROFILING_TEMP_DIR, myDir.toString());
+    props.put(
+        ProfilingConfig.PROFILING_UPLOAD_PERIOD,
+        "0"); // to force immediate cleanup; must be a string value!
     ConfigProvider configProvider = ConfigProvider.withPropertiesOverride(props);
     TempLocationManager tempLocationManager = new TempLocationManager(configProvider);
-    Path tempDir = tempLocationManager.getTempDir();
+    Path tempDir = tempLocationManager.getTempDir(Paths.get(subPath));
     assertNotNull(tempDir);
 
     // fake temp location
-    Path fakeTempDir = tempDir.getParent().resolve("pid_00000");
+    Path fakeTempDir = tempDir.getParent();
+    while (fakeTempDir != null && !fakeTempDir.endsWith("ddprof")) {
+      fakeTempDir = fakeTempDir.getParent();
+    }
+    fakeTempDir = fakeTempDir.resolve("pid_0000");
     Files.createDirectories(fakeTempDir);
+    Path tmpFile = Files.createFile(fakeTempDir.resolve("test.txt"));
+    tmpFile.toFile().deleteOnExit(); // make sure this is deleted at exit
+    fakeTempDir.toFile().deleteOnExit(); // also this one
     tempLocationManager.cleanup(false);
     // fake temp location should be deleted
     // real temp location should be kept