Use ControllableTimeSource for deterministic TempLocationManagerTest (#10538)

Author: jbachorik

internal-api/src/main/java/datadog/trace/util/TempLocationManager.java

@@ -3,6 +3,8 @@
 import datadog.environment.SystemProperties;
 import datadog.trace.api.config.ProfilingConfig;
 import datadog.trace.api.profiling.ProfilerFlareLogger;
+import datadog.trace.api.time.SystemTimeSource;
+import datadog.trace.api.time.TimeSource;
 import datadog.trace.bootstrap.config.provider.ConfigProvider;
 import datadog.trace.config.inversion.ConfigHelper;
 import java.io.IOException;
@@ -16,8 +18,6 @@
 import java.nio.file.attribute.BasicFileAttributes;
 import java.nio.file.attribute.PosixFilePermission;
 import java.nio.file.attribute.PosixFilePermissions;
-import java.time.Instant;
-import java.time.temporal.ChronoUnit;
 import java.util.Set;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.TimeUnit;
@@ -71,25 +71,23 @@ private final class CleanupVisitor implements FileVisitor<Path> {
 
     private Set<String> pidSet;
 
-    private final Instant cutoff;
-    private final Instant timeoutTarget;
+    private final long cutoffMillis;
+    private final long timeoutTargetMillis;
 
     private boolean terminated = false;
 
     CleanupVisitor(long timeout, TimeUnit unit) {
-      this.cutoff = Instant.now().minus(cutoffSeconds, ChronoUnit.SECONDS);
-      this.timeoutTarget =
-          timeout > -1
-              ? Instant.now().plus(TimeUnit.MILLISECONDS.convert(timeout, unit), ChronoUnit.MILLIS)
-              : null;
+      long now = timeSource.getCurrentTimeMillis();
+      this.cutoffMillis = now - TimeUnit.SECONDS.toMillis(cutoffSeconds);
+      this.timeoutTargetMillis = timeout > -1 ? now + unit.toMillis(timeout) : Long.MAX_VALUE;
     }
 
     boolean isTerminated() {
       return terminated;
     }
 
     private boolean isTimedOut() {
-      return timeoutTarget != null && Instant.now().isAfter(timeoutTarget);
+      return timeSource.getCurrentTimeMillis() > timeoutTargetMillis;
     }
 
     @Override
@@ -133,7 +131,7 @@ public FileVisitResult visitFile(Path file, BasicFileAttributes attrs) throws IO
       }
       cleanupTestHook.visitFile(file, attrs, false);
       try {
-        if (Files.getLastModifiedTime(file).toInstant().isAfter(cutoff)) {
+        if (Files.getLastModifiedTime(file).toMillis() > cutoffMillis) {
           return FileVisitResult.SKIP_SUBTREE;
         }
         Files.delete(file);
@@ -218,6 +216,7 @@ boolean await(long timeout, TimeUnit unit) throws Throwable {
 
   private final CleanupTask cleanupTask = new CleanupTask();
   private final CleanupHook cleanupTestHook;
+  private final TimeSource timeSource;
 
   /**
    * Get the singleton instance of the TempLocationManager. It will run the cleanup task in the
@@ -248,12 +247,21 @@ private TempLocationManager() {
   }
 
   TempLocationManager(ConfigProvider configProvider) {
-    this(configProvider, true, CleanupHook.EMPTY);
+    this(configProvider, true, CleanupHook.EMPTY, SystemTimeSource.INSTANCE);
   }
 
   TempLocationManager(
       ConfigProvider configProvider, boolean runStartupCleanup, CleanupHook testHook) {
+    this(configProvider, runStartupCleanup, testHook, SystemTimeSource.INSTANCE);
+  }
+
+  TempLocationManager(
+      ConfigProvider configProvider,
+      boolean runStartupCleanup,
+      CleanupHook testHook,
+      TimeSource timeSource) {
     cleanupTestHook = testHook;
+    this.timeSource = timeSource;
 
     Set<String> supportedViews = FileSystems.getDefault().supportedFileAttributeViews();
     isPosixFs = supportedViews.contains("posix");

internal-api/src/test/java/datadog/trace/util/TempLocationManagerTest.java

@@ -8,6 +8,9 @@
 import static org.junit.jupiter.api.Assertions.assertTrue;
 
 import datadog.trace.api.config.ProfilingConfig;
+import datadog.trace.api.time.ControllableTimeSource;
+import datadog.trace.api.time.SystemTimeSource;
+import datadog.trace.api.time.TimeSource;
 import datadog.trace.bootstrap.config.provider.ConfigProvider;
 import java.io.IOException;
 import java.nio.file.FileVisitResult;
@@ -16,7 +19,6 @@
 import java.nio.file.Paths;
 import java.nio.file.attribute.BasicFileAttributes;
 import java.nio.file.attribute.PosixFilePermissions;
-import java.time.Duration;
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Properties;
@@ -25,7 +27,6 @@
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicReference;
-import java.util.concurrent.locks.LockSupport;
 import java.util.stream.Stream;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.params.ParameterizedTest;
@@ -232,23 +233,33 @@ public FileVisitResult postVisitDirectory(Path dir, IOException exc, boolean tim
   void testCleanupWithTimeout(boolean shouldSucceed, String section) throws Exception {
     /*
      * Test that cleanup correctly handles timeout conditions.
-     * Uses a hook that introduces delays at specific points in the file tree walk.
+     * Uses a ControllableTimeSource for fully deterministic behavior:
+     * - Success case: time is never advanced, so timeout never occurs
+     * - Failure case: time is advanced past timeout on first matching hook call
      *
-     * Timing strategy (more generous for JDK 8 compatibility):
-     * - Base delay: 100ms (was 10ms - too tight for JDK 8 timer resolution)
-     * - Success case: 500ms timeout with 100ms delays → plenty of margin
-     * - Failure case: 20ms timeout with 100ms delays → guaranteed timeout
+     * This approach is independent of how many times hooks are invoked.
      */
-    long delayMs = 100;
+    long timeoutMs = 100;
+    ControllableTimeSource timeSource = new ControllableTimeSource();
+
     AtomicBoolean withTimeout = new AtomicBoolean(false);
+    AtomicBoolean timeAdvanced = new AtomicBoolean(false);
     TempLocationManager.CleanupHook delayer =
         new TempLocationManager.CleanupHook() {
+          private void maybeAdvanceTime() {
+            // Only advance time once, and only for failure case
+            if (!shouldSucceed && timeAdvanced.compareAndSet(false, true)) {
+              // Advance well past the timeout
+              timeSource.advance(TimeUnit.MILLISECONDS.toNanos(timeoutMs * 10));
+            }
+          }
+
           @Override
           public FileVisitResult preVisitDirectory(
               Path dir, BasicFileAttributes attrs, boolean timeout) throws IOException {
             withTimeout.compareAndSet(false, timeout);
             if (section.equals("preVisitDirectory")) {
-              waitFor(Duration.ofMillis(delayMs));
+              maybeAdvanceTime();
             }
             return TempLocationManager.CleanupHook.super.preVisitDirectory(dir, attrs, timeout);
           }
@@ -258,7 +269,7 @@ public FileVisitResult visitFileFailed(Path file, IOException exc, boolean timeo
               throws IOException {
             withTimeout.compareAndSet(false, timeout);
             if (section.equals("visitFileFailed")) {
-              waitFor(Duration.ofMillis(delayMs));
+              maybeAdvanceTime();
             }
             return TempLocationManager.CleanupHook.super.visitFileFailed(file, exc, timeout);
           }
@@ -268,7 +279,7 @@ public FileVisitResult postVisitDirectory(Path dir, IOException exc, boolean tim
               throws IOException {
             withTimeout.compareAndSet(false, timeout);
             if (section.equals("postVisitDirectory")) {
-              waitFor(Duration.ofMillis(delayMs));
+              maybeAdvanceTime();
             }
             return TempLocationManager.CleanupHook.super.postVisitDirectory(dir, exc, timeout);
           }
@@ -278,7 +289,7 @@ public FileVisitResult visitFile(Path file, BasicFileAttributes attrs, boolean t
               throws IOException {
             withTimeout.compareAndSet(false, timeout);
             if (section.equals("visitFile")) {
-              waitFor(Duration.ofMillis(delayMs));
+              maybeAdvanceTime();
             }
             return TempLocationManager.CleanupHook.super.visitFile(file, attrs, timeout);
           }
@@ -287,16 +298,15 @@ public FileVisitResult visitFile(Path file, BasicFileAttributes attrs, boolean t
         Files.createTempDirectory(
             "ddprof-test-",
             PosixFilePermissions.asFileAttribute(PosixFilePermissions.fromString("rwx------")));
-    TempLocationManager instance = instance(baseDir, false, delayer);
+    TempLocationManager instance = instance(baseDir, false, delayer, timeSource);
     Path mytempdir = instance.getTempDir();
     Path otherTempdir = mytempdir.getParent().resolve("pid_0000");
     Files.createDirectories(otherTempdir);
     Files.createFile(mytempdir.resolve("dummy"));
     Files.createFile(otherTempdir.resolve("dummy"));
-    // Success: 500ms timeout (5 * 100ms) with 100ms delays → should complete
-    // Failure: 20ms timeout (0.2 * 100ms) with 100ms delays → should timeout
-    boolean rslt =
-        instance.cleanup((long) (delayMs * (shouldSucceed ? 5 : 0.2d)), TimeUnit.MILLISECONDS);
+    // Set controllable time AFTER files are created so their modification times are in the past
+    timeSource.set(TimeUnit.MILLISECONDS.toNanos(System.currentTimeMillis()));
+    boolean rslt = instance.cleanup(timeoutMs, TimeUnit.MILLISECONDS);
     assertEquals(shouldSucceed, rslt);
     assertNotEquals(shouldSucceed, withTimeout.get()); // timeout = !shouldSucceed
   }
@@ -313,23 +323,21 @@ private static Stream<Arguments> timeoutTestArguments() {
 
   private TempLocationManager instance(
       Path baseDir, boolean withStartupCleanup, TempLocationManager.CleanupHook cleanupHook) {
+    return instance(baseDir, withStartupCleanup, cleanupHook, SystemTimeSource.INSTANCE);
+  }
+
+  private TempLocationManager instance(
+      Path baseDir,
+      boolean withStartupCleanup,
+      TempLocationManager.CleanupHook cleanupHook,
+      TimeSource timeSource) {
     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), withStartupCleanup, cleanupHook);
-  }
-
-  private void waitFor(Duration timeout) {
-    long target = System.nanoTime() + timeout.toNanos();
-    while (System.nanoTime() < target) {
-      long toSleep = target - System.nanoTime();
-      if (toSleep > 0) {
-        LockSupport.parkNanos(toSleep);
-      }
-    }
+        ConfigProvider.withPropertiesOverride(props), withStartupCleanup, cleanupHook, timeSource);
   }
 }