Add tests for concurrent cleanup

Author: jbachorik

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

@@ -7,6 +7,7 @@
 import java.nio.file.FileVisitResult;
 import java.nio.file.FileVisitor;
 import java.nio.file.Files;
+import java.nio.file.NoSuchFileException;
 import java.nio.file.Path;
 import java.nio.file.Paths;
 import java.nio.file.attribute.BasicFileAttributes;
@@ -31,6 +32,33 @@ private static final class SingletonHolder {
     private static final TempLocationManager INSTANCE = new TempLocationManager();
   }
 
+  interface CleanupHook extends FileVisitor<Path> {
+    CleanupHook EMPTY = new CleanupHook() {};
+
+    @Override
+    default FileVisitResult preVisitDirectory(Path dir, BasicFileAttributes attrs)
+        throws IOException {
+      return null;
+    }
+
+    @Override
+    default FileVisitResult visitFile(Path file, BasicFileAttributes attrs) throws IOException {
+      return null;
+    }
+
+    @Override
+    default FileVisitResult visitFileFailed(Path file, IOException exc) throws IOException {
+      return null;
+    }
+
+    @Override
+    default FileVisitResult postVisitDirectory(Path dir, IOException exc) throws IOException {
+      return null;
+    }
+
+    default void onCleanupStart() {}
+  }
+
   private class CleanupVisitor implements FileVisitor<Path> {
     private boolean shouldClean;
 
@@ -47,6 +75,8 @@ private class CleanupVisitor implements FileVisitor<Path> {
     @Override
     public FileVisitResult preVisitDirectory(Path dir, BasicFileAttributes attrs)
         throws IOException {
+      cleanupTestHook.preVisitDirectory(dir, attrs);
+
       if (dir.equals(baseTempDir)) {
         return FileVisitResult.CONTINUE;
       }
@@ -63,25 +93,40 @@ public FileVisitResult preVisitDirectory(Path dir, BasicFileAttributes attrs)
 
     @Override
     public FileVisitResult visitFile(Path file, BasicFileAttributes attrs) throws IOException {
-      if (Files.getLastModifiedTime(file).toInstant().isAfter(cutoff)) {
-        return FileVisitResult.SKIP_SUBTREE;
+      cleanupTestHook.visitFile(file, attrs);
+      try {
+        if (Files.getLastModifiedTime(file).toInstant().isAfter(cutoff)) {
+          return FileVisitResult.SKIP_SUBTREE;
+        }
+        Files.delete(file);
+      } catch (NoSuchFileException ignored) {
+        // another process has already cleaned it; ignore
       }
-      Files.delete(file);
       return FileVisitResult.CONTINUE;
     }
 
     @Override
     public FileVisitResult visitFileFailed(Path file, IOException exc) throws IOException {
-      if (log.isDebugEnabled()) {
+      cleanupTestHook.visitFileFailed(file, exc);
+      // do not log files/directories removed by another process running concurrently
+      if (!(exc instanceof NoSuchFileException) && log.isDebugEnabled()) {
         log.debug("Failed to delete file {}", file, exc);
       }
       return FileVisitResult.CONTINUE;
     }
 
     @Override
     public FileVisitResult postVisitDirectory(Path dir, IOException exc) throws IOException {
+      cleanupTestHook.postVisitDirectory(dir, exc);
+      if (exc instanceof NoSuchFileException) {
+        return FileVisitResult.CONTINUE;
+      }
       if (shouldClean) {
-        Files.delete(dir);
+        try {
+          Files.delete(dir);
+        } catch (NoSuchFileException ignored) {
+          // another process has already cleaned it; ignore
+        }
         String fileName = dir.getFileName().toString();
         // reset the flag only if we are done cleaning the top-level directory
         shouldClean = !fileName.startsWith("pid_");
@@ -96,6 +141,8 @@ public FileVisitResult postVisitDirectory(Path dir, IOException exc) throws IOEx
 
   private final CompletableFuture<Void> cleanupTask;
 
+  private final CleanupHook cleanupTestHook;
+
   /**
    * Get the singleton instance of the TempLocationManager. It will run the cleanup task in the
    * background.
@@ -125,6 +172,12 @@ private TempLocationManager() {
   }
 
   TempLocationManager(ConfigProvider configProvider) {
+    this(configProvider, CleanupHook.EMPTY);
+  }
+
+  TempLocationManager(ConfigProvider configProvider, CleanupHook testHook) {
+    cleanupTestHook = testHook;
+
     // In order to avoid racy attempts to clean up files which are currently being processed in a
     // JVM which is being shut down (the JVMs far in the shutdown routine may not be reported by
     // 'jps' but still can be eg. processing JFR chunks) we will not clean up any files not older
@@ -222,6 +275,7 @@ public Path getTempDir(Path subPath, boolean create) {
 
   void cleanup(boolean cleanSelf) {
     try {
+      cleanupTestHook.onCleanupStart();
       Files.walkFileTree(baseTempDir, new CleanupVisitor(cleanSelf));
     } catch (IOException e) {
       if (log.isDebugEnabled()) {
@@ -232,7 +286,8 @@ void cleanup(boolean cleanSelf) {
     }
   }
 
-  private void waitForCleanup() {
+  // accessible for tests
+  void waitForCleanup() {
     cleanupTask.join();
   }
 }

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

@@ -5,12 +5,16 @@
 import datadog.trace.api.config.ProfilingConfig;
 import datadog.trace.bootstrap.config.provider.ConfigProvider;
 import datadog.trace.util.PidHelper;
+import java.io.IOException;
+import java.nio.file.FileVisitResult;
 import java.nio.file.Files;
 import java.nio.file.Path;
 import java.nio.file.Paths;
+import java.nio.file.attribute.BasicFileAttributes;
 import java.nio.file.attribute.PosixFilePermissions;
 import java.util.Properties;
 import java.util.UUID;
+import java.util.concurrent.Phaser;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.ValueSource;
@@ -105,4 +109,83 @@ void testCleanup(String subPath) throws Exception {
     assertFalse(Files.exists(fakeTempDir));
     assertTrue(Files.exists(tempDir));
   }
+
+  @ParameterizedTest
+  @ValueSource(strings = {"preVisitDirectory", "visitFile", "postVisitDirectory"})
+  void testConcurrentCleanup(String section) throws Exception {
+    /* This test simulates concurrent cleanup
+      It utilizes a special hook to create synchronization points in the filetree walking routine,
+      allowing to delete the files at various points of execution.
+      The test makes sure that the cleanup is not interrupted and the file and directory being deleted
+      stays deleted.
+    */
+    Path baseDir =
+        Files.createTempDirectory(
+            "ddprof-test-",
+            PosixFilePermissions.asFileAttribute(PosixFilePermissions.fromString("rwx------")));
+
+    Path fakeTempDir = baseDir.resolve("ddprof/pid_1234/scratch");
+    Files.createDirectories(fakeTempDir);
+    Path fakeTempFile = fakeTempDir.resolve("libxxx.so");
+    Files.createFile(fakeTempFile);
+
+    fakeTempDir.toFile().deleteOnExit();
+    fakeTempFile.toFile().deleteOnExit();
+
+    Phaser phaser = new Phaser(2);
+
+    TempLocationManager.CleanupHook blocker =
+        new TempLocationManager.CleanupHook() {
+          @Override
+          public FileVisitResult preVisitDirectory(Path dir, BasicFileAttributes attrs)
+              throws IOException {
+            if (section.equals("preVisitDirectory") && dir.equals(fakeTempDir)) {
+              phaser.arriveAndAwaitAdvance();
+              phaser.arriveAndAwaitAdvance();
+            }
+            return null;
+          }
+
+          @Override
+          public FileVisitResult visitFile(Path file, BasicFileAttributes attrs)
+              throws IOException {
+            if (section.equals("visitFile") && file.equals(fakeTempFile)) {
+              phaser.arriveAndAwaitAdvance();
+              phaser.arriveAndAwaitAdvance();
+            }
+            return null;
+          }
+
+          @Override
+          public FileVisitResult postVisitDirectory(Path dir, IOException exc) throws IOException {
+            if (section.equals("postVisitDirectory") && dir.equals(fakeTempDir)) {
+              phaser.arriveAndAwaitAdvance();
+              phaser.arriveAndAwaitAdvance();
+            }
+            return null;
+          }
+        };
+
+    TempLocationManager mgr = instance(baseDir, blocker);
+
+    // wait for the cleanup start
+    phaser.arriveAndAwaitAdvance();
+    Files.deleteIfExists(fakeTempFile);
+    phaser.arriveAndAwaitAdvance();
+    mgr.waitForCleanup();
+
+    assertFalse(Files.exists(fakeTempFile));
+    assertFalse(Files.exists(fakeTempDir));
+  }
+
+  private TempLocationManager instance(Path baseDir, 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);
+  }
 }