TempLocationManager Fixes and Improvements (#8191)

* Remove CompletableFuture from TempLocationManager

* Do not run the cleanup if there is nothing to clean up

* Exclude JFR repository from temp files self-cleanup

Author: jbachorik

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

@@ -4,6 +4,7 @@
 
 import datadog.trace.api.config.ProfilingConfig;
 import datadog.trace.bootstrap.config.provider.ConfigProvider;
+import datadog.trace.util.AgentTaskScheduler;
 import datadog.trace.util.PidHelper;
 import java.io.IOException;
 import java.nio.file.FileVisitResult;
@@ -17,10 +18,10 @@
 import java.time.Instant;
 import java.time.temporal.ChronoUnit;
 import java.util.Set;
-import java.util.concurrent.CompletableFuture;
-import java.util.concurrent.ExecutionException;
+import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.TimeUnit;
-import java.util.concurrent.TimeoutException;
+import java.util.regex.Pattern;
+import java.util.stream.Stream;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -32,6 +33,9 @@
  */
 public final class TempLocationManager {
   private static final Logger log = LoggerFactory.getLogger(TempLocationManager.class);
+  private static final Pattern JFR_DIR_PATTERN =
+      Pattern.compile("\\d{4}_\\d{2}_\\d{2}_\\d{2}_\\d{2}_\\d{2}_\\d{6}");
+  private static final String TEMPDIR_PREFIX = "pid_";
 
   private static final class SingletonHolder {
     private static final TempLocationManager INSTANCE = new TempLocationManager();
@@ -64,7 +68,7 @@ default FileVisitResult postVisitDirectory(Path dir, IOException exc) throws IOE
     default void onCleanupStart(boolean selfCleanup, long timeout, TimeUnit unit) {}
   }
 
-  private class CleanupVisitor implements FileVisitor<Path> {
+  private final class CleanupVisitor implements FileVisitor<Path> {
     private boolean shouldClean;
 
     private final Set<String> pidSet = PidHelper.getJavaPids();
@@ -100,14 +104,19 @@ public FileVisitResult preVisitDirectory(Path dir, BasicFileAttributes attrs)
         terminated = true;
         return FileVisitResult.TERMINATE;
       }
+      if (cleanSelf && JFR_DIR_PATTERN.matcher(dir.getFileName().toString()).matches()) {
+        // do not delete JFR repository on 'self-cleanup' - it conflicts with the JFR's own cleanup
+        return FileVisitResult.SKIP_SUBTREE;
+      }
+
       cleanupTestHook.preVisitDirectory(dir, attrs);
 
       if (dir.equals(baseTempDir)) {
         return FileVisitResult.CONTINUE;
       }
       String fileName = dir.getFileName().toString();
       // the JFR repository directories are under <basedir>/pid_<pid>
-      String pid = fileName.startsWith("pid_") ? fileName.substring(4) : null;
+      String pid = fileName.startsWith(TEMPDIR_PREFIX) ? fileName.substring(4) : null;
       boolean isSelfPid = pid != null && pid.equals(PidHelper.getPid());
       shouldClean |= cleanSelf ? isSelfPid : !isSelfPid && !pidSet.contains(pid);
       if (shouldClean) {
@@ -169,18 +178,43 @@ public FileVisitResult postVisitDirectory(Path dir, IOException exc) throws IOEx
         }
         String fileName = dir.getFileName().toString();
         // reset the flag only if we are done cleaning the top-level directory
-        shouldClean = !fileName.startsWith("pid_");
+        shouldClean = !fileName.startsWith(TEMPDIR_PREFIX);
       }
       return FileVisitResult.CONTINUE;
     }
   }
 
+  private final class CleanupTask implements Runnable {
+    private final CountDownLatch latch = new CountDownLatch(1);
+    private volatile Throwable throwable = null;
+
+    @Override
+    public void run() {
+      try {
+        cleanup(false);
+      } catch (OutOfMemoryError oom) {
+        throw oom;
+      } catch (Throwable t) {
+        throwable = t;
+      } finally {
+        latch.countDown();
+      }
+    }
+
+    boolean await(long timeout, TimeUnit unit) throws Throwable {
+      boolean ret = latch.await(timeout, unit);
+      if (throwable != null) {
+        throw throwable;
+      }
+      return ret;
+    }
+  }
+
   private final Path baseTempDir;
   private final Path tempDir;
   private final long cutoffSeconds;
 
-  private final CompletableFuture<Void> cleanupTask;
-
+  private final CleanupTask cleanupTask = new CleanupTask();
   private final CleanupHook cleanupTestHook;
 
   /**
@@ -202,11 +236,7 @@ public static TempLocationManager getInstance() {
   static TempLocationManager getInstance(boolean waitForCleanup) {
     TempLocationManager instance = SingletonHolder.INSTANCE;
     if (waitForCleanup) {
-      try {
-        instance.waitForCleanup(5, TimeUnit.SECONDS);
-      } catch (TimeoutException ignored) {
-
-      }
+      instance.waitForCleanup(5, TimeUnit.SECONDS);
     }
     return instance;
   }
@@ -216,10 +246,11 @@ private TempLocationManager() {
   }
 
   TempLocationManager(ConfigProvider configProvider) {
-    this(configProvider, CleanupHook.EMPTY);
+    this(configProvider, true, CleanupHook.EMPTY);
   }
 
-  TempLocationManager(ConfigProvider configProvider, CleanupHook testHook) {
+  TempLocationManager(
+      ConfigProvider configProvider, boolean runStartupCleanup, CleanupHook testHook) {
     cleanupTestHook = testHook;
 
     // In order to avoid racy attempts to clean up files which are currently being processed in a
@@ -258,21 +289,21 @@ private TempLocationManager() {
     baseTempDir = configuredTempDir.resolve("ddprof");
     baseTempDir.toFile().deleteOnExit();
 
-    tempDir = baseTempDir.resolve("pid_" + pid);
-    cleanupTask = CompletableFuture.runAsync(() -> cleanup(false));
+    tempDir = baseTempDir.resolve(TEMPDIR_PREFIX + pid);
+    if (runStartupCleanup) {
+      // do not execute the background cleanup task when running in tests
+      AgentTaskScheduler.INSTANCE.execute(() -> cleanup(false));
+    }
 
     Thread selfCleanup =
         new Thread(
             () -> {
-              try {
-                waitForCleanup(1, TimeUnit.SECONDS);
-              } catch (TimeoutException e) {
+              if (!waitForCleanup(1, TimeUnit.SECONDS)) {
                 log.info(
                     "Cleanup task timed out. {} temp directory might not have been cleaned up properly",
                     tempDir);
-              } finally {
-                cleanup(true);
               }
+              cleanup(true);
             },
             "Temp Location Manager Cleanup");
     Runtime.getRuntime().addShutdownHook(selfCleanup);
@@ -347,6 +378,19 @@ boolean cleanup(boolean cleanSelf) {
    */
   boolean cleanup(boolean cleanSelf, long timeout, TimeUnit unit) {
     try {
+      if (!Files.exists(baseTempDir)) {
+        // not event the main temp location exists; nothing to clean up
+        return true;
+      }
+      try (Stream<Path> paths = Files.walk(baseTempDir)) {
+        if (paths.noneMatch(
+            path ->
+                Files.isDirectory(path)
+                    && path.getFileName().toString().startsWith(TEMPDIR_PREFIX))) {
+          // nothing to clean up; bail out early
+          return true;
+        }
+      }
       cleanupTestHook.onCleanupStart(cleanSelf, timeout, unit);
       CleanupVisitor visitor = new CleanupVisitor(cleanSelf, timeout, unit);
       Files.walkFileTree(baseTempDir, visitor);
@@ -362,21 +406,24 @@ boolean cleanup(boolean cleanSelf, long timeout, TimeUnit unit) {
   }
 
   // accessible for tests
-  void waitForCleanup(long timeout, TimeUnit unit) throws TimeoutException {
+  boolean waitForCleanup(long timeout, TimeUnit unit) {
     try {
-      cleanupTask.get(timeout, unit);
+      return cleanupTask.await(timeout, unit);
     } catch (InterruptedException e) {
-      cleanupTask.cancel(true);
+      log.debug("Temp directory cleanup was interrupted");
       Thread.currentThread().interrupt();
-    } catch (TimeoutException e) {
-      cleanupTask.cancel(true);
-      throw e;
-    } catch (ExecutionException e) {
+    } catch (Throwable t) {
       if (log.isDebugEnabled()) {
-        log.debug("Failed to cleanup temp directory: {}", tempDir, e);
+        log.debug("Failed to cleanup temp directory: {}", tempDir, t);
       } else {
         log.debug("Failed to cleanup temp directory: {}", tempDir);
       }
     }
+    return false;
+  }
+
+  // accessible for tests
+  void createDirStructure() throws IOException {
+    Files.createDirectories(baseTempDir);
   }
 }

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

@@ -18,6 +18,7 @@
 import java.util.UUID;
 import java.util.concurrent.Phaser;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.locks.LockSupport;
 import java.util.stream.Stream;
 import org.junit.jupiter.api.Test;
@@ -173,7 +174,7 @@ public FileVisitResult postVisitDirectory(Path dir, IOException exc) throws IOEx
           }
         };
 
-    TempLocationManager mgr = instance(baseDir, blocker);
+    TempLocationManager mgr = instance(baseDir, true, blocker);
 
     // wait for the cleanup start
     phaser.arriveAndAwaitAdvance();
@@ -187,8 +188,9 @@ public FileVisitResult postVisitDirectory(Path dir, IOException exc) throws IOEx
 
   @ParameterizedTest
   @MethodSource("timeoutTestArguments")
-  void testCleanupWithTimeout(boolean selfCleanup, String section) throws Exception {
-    long timeoutMs = 500;
+  void testCleanupWithTimeout(boolean selfCleanup, boolean shouldSucceed, String section)
+      throws Exception {
+    long timeoutMs = 10;
     TempLocationManager.CleanupHook delayer =
         new TempLocationManager.CleanupHook() {
           @Override
@@ -229,30 +231,63 @@ public FileVisitResult visitFile(Path file, BasicFileAttributes attrs)
         Files.createTempDirectory(
             "ddprof-test-",
             PosixFilePermissions.asFileAttribute(PosixFilePermissions.fromString("rwx------")));
-    TempLocationManager instance = instance(baseDir, delayer);
-    boolean rslt = instance.cleanup(selfCleanup, timeoutMs, TimeUnit.MILLISECONDS);
-    assertTrue(rslt);
+    TempLocationManager instance = instance(baseDir, false, delayer);
+    Path mytempdir = instance.getTempDir();
+    Path otherTempdir = mytempdir.getParent().resolve("pid_0000");
+    Files.createDirectories(otherTempdir);
+    Files.createFile(mytempdir.resolve("dummy"));
+    Files.createFile(otherTempdir.resolve("dummy"));
+    boolean rslt =
+        instance.cleanup(
+            selfCleanup, (long) (timeoutMs * (shouldSucceed ? 10 : 0.5d)), TimeUnit.MILLISECONDS);
+    assertEquals(shouldSucceed, rslt);
+  }
+
+  @Test
+  void testShortCircuit() throws Exception {
+    Path baseDir =
+        Files.createTempDirectory(
+            "ddprof-test-",
+            PosixFilePermissions.asFileAttribute(PosixFilePermissions.fromString("rwx------")));
+    AtomicBoolean executed = new AtomicBoolean();
+    TempLocationManager.CleanupHook hook =
+        new TempLocationManager.CleanupHook() {
+          @Override
+          public void onCleanupStart(boolean selfCleanup, long timeout, TimeUnit unit) {
+            executed.set(true);
+          }
+        };
+    TempLocationManager instance = instance(baseDir, false, hook);
+
+    instance.createDirStructure();
+
+    boolean ret = instance.cleanup(false);
+    assertTrue(ret);
+    assertFalse(executed.get());
   }
 
   private static Stream<Arguments> timeoutTestArguments() {
     List<Arguments> argumentsList = new ArrayList<>();
     for (boolean selfCleanup : new boolean[] {true, false}) {
       for (String intercepted :
           new String[] {"preVisitDirectory", "visitFile", "postVisitDirectory"}) {
-        argumentsList.add(Arguments.of(selfCleanup, intercepted));
+        argumentsList.add(Arguments.of(selfCleanup, true, intercepted));
+        argumentsList.add(Arguments.of(selfCleanup, false, intercepted));
       }
     }
     return argumentsList.stream();
   }
 
-  private TempLocationManager instance(Path baseDir, TempLocationManager.CleanupHook cleanupHook)
+  private TempLocationManager instance(
+      Path baseDir, boolean withStartupCleanup, TempLocationManager.CleanupHook cleanupHook)
       throws IOException {
     Properties props = new Properties();
     props.put(ProfilingConfig.PROFILING_TEMP_DIR, baseDir.toString());
     props.put(
         ProfilingConfig.PROFILING_UPLOAD_PERIOD,
         "0"); // to force immediate cleanup; must be a string value!
 
-    return new TempLocationManager(ConfigProvider.withPropertiesOverride(props), cleanupHook);
+    return new TempLocationManager(
+        ConfigProvider.withPropertiesOverride(props), withStartupCleanup, cleanupHook);
   }
 }