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

Author: jbachorik

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

@@ -20,6 +20,7 @@
 import java.util.Set;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.TimeUnit;
+import java.util.stream.Stream;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -31,6 +32,8 @@
  */
 public final class TempLocationManager {
   private static final Logger log = LoggerFactory.getLogger(TempLocationManager.class);
+  // accessible to tests
+  static final String TEMPDIR_PREFIX = "pid_";
 
   private static final class SingletonHolder {
     private static final TempLocationManager INSTANCE = new TempLocationManager();
@@ -106,7 +109,7 @@ public FileVisitResult preVisitDirectory(Path dir, BasicFileAttributes attrs)
       }
       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) {
@@ -168,7 +171,7 @@ 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;
     }
@@ -236,10 +239,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
@@ -278,8 +282,11 @@ private TempLocationManager() {
     baseTempDir = configuredTempDir.resolve("ddprof");
     baseTempDir.toFile().deleteOnExit();
 
-    tempDir = baseTempDir.resolve("pid_" + pid);
-    AgentTaskScheduler.INSTANCE.execute(() -> 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(
@@ -364,6 +371,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);
@@ -394,4 +414,9 @@ boolean waitForCleanup(long timeout, TimeUnit unit) {
     }
     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);
   }
 }