Add option to clean up leftover temp jars on boot (Windows-only)

DataDog · Nov 17, 2020 · 5e4a765 · 5e4a765
1 parent 447d0d9
commit 5e4a765
Show file tree

Hide file tree

Showing 3 changed files with 71 additions and 2 deletions.
diff --git a/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/HelperInjector.java b/dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/HelperInjector.java
@@ -2,10 +2,15 @@
 
 import static datadog.trace.agent.tooling.ClassLoaderMatcher.BOOTSTRAP_CLASSLOADER;
 
+import datadog.trace.api.Config;
 import java.io.File;
+import java.io.FilenameFilter;
 import java.io.IOException;
 import java.lang.ref.WeakReference;
+import java.nio.file.DirectoryStream;
 import java.nio.file.Files;
+import java.nio.file.Path;
+import java.nio.file.Paths;
 import java.security.SecureClassLoader;
 import java.util.Arrays;
 import java.util.Collection;
@@ -38,6 +43,12 @@ public String toString() {
         }
       };
 
+  static {
+    if (Config.get().isTempJarsCleanOnBoot()) {
+      cleanTempJars();
+    }
+  }
+
   private final String requestingName;
 
   private final Set<String> helperClassNames;
@@ -195,7 +206,7 @@ private void ensureModuleCanReadHelperModules(final JavaModule target) {
 
   private static File createTempDir() throws IOException {
     try {
-      return Files.createTempDirectory("datadog-temp-jars").toFile();
+      return Files.createTempDirectory(DATADOG_TEMP_JARS).toFile();
     } catch (final IOException e) {
       if (log.isErrorEnabled()) {
         log.error(
@@ -219,4 +230,50 @@ private static void deleteTempDir(final File file) {
       log.debug("file '{}' deleted", file);
     }
   }
+
+  private static final String DATADOG_TEMP_JARS = "datadog-temp-jars";
+  private static final int MAX_CLEANUP_MILLIS = 1_000;
+
+  private static void cleanTempJars() {
+    try {
+      final Path tmpDir = Paths.get(System.getProperty("java.io.tmpdir"));
+      log.debug("Cleaning temp jar directories under {}", tmpDir);
+
+      final long maxTimeMillis = System.currentTimeMillis() + MAX_CLEANUP_MILLIS;
+      try (final DirectoryStream<Path> paths =
+          Files.newDirectoryStream(tmpDir, DATADOG_TEMP_JARS + "*")) {
+        for (final Path dir : paths) {
+          if (System.currentTimeMillis() > maxTimeMillis) {
+            break; // avoid attempting too much cleanup on boot
+          }
+          cleanTempJars(dir.toFile());
+        }
+      }
+    } catch (final Throwable e) {
+      log.debug("Problem cleaning temp jar directories", e);
+    }
+  }
+
+  private static void cleanTempJars(final File tempJarDir) {
+    final File[] tempJars =
+        tempJarDir.listFiles(
+            new FilenameFilter() {
+              @Override
+              public boolean accept(final File dir, final String name) {
+                return name.startsWith("jar") && name.endsWith(".jar");
+              }
+            });
+
+    if (tempJars != null) {
+      for (final File jar : tempJars) {
+        if (jar.delete()) {
+          log.debug("file '{}' deleted", jar);
+        }
+      }
+    }
+
+    if (tempJarDir.delete()) {
+      log.debug("file '{}' deleted", tempJarDir);
+    }
+  }
 }
diff --git a/dd-trace-api/src/main/java/datadog/trace/api/config/TraceInstrumentationConfig.java b/dd-trace-api/src/main/java/datadog/trace/api/config/TraceInstrumentationConfig.java
@@ -38,5 +38,7 @@ public final class TraceInstrumentationConfig {
   public static final String SERVLET_PRINCIPAL_ENABLED = "trace.servlet.principal.enabled";
   public static final String SERVLET_ASYNC_TIMEOUT_ERROR = "trace.servlet.async-timeout.error";
 
+  public static final String TEMP_JARS_CLEAN_ON_BOOT = "temp.jars.clean.on.boot";
+
   private TraceInstrumentationConfig() {}
 }
diff --git a/internal-api/src/main/java/datadog/trace/api/Config.java b/internal-api/src/main/java/datadog/trace/api/Config.java
@@ -382,6 +382,8 @@ private String profilingProxyPasswordMasker() {
   @Getter private final boolean servletPrincipalEnabled;
   @Getter private final boolean servletAsyncTimeoutError;
 
+  @Getter private final boolean tempJarsCleanOnBoot;
+
   @Getter private final boolean traceAgentV05Enabled;
 
   @Getter private final boolean debugEnabled;
@@ -705,6 +707,10 @@ private Config(final String runtimeId, final ConfigProvider configProvider) {
     servletAsyncTimeoutError =
         configProvider.getBoolean(TraceInstrumentationConfig.SERVLET_ASYNC_TIMEOUT_ERROR, true);
 
+    tempJarsCleanOnBoot =
+        configProvider.getBoolean(TraceInstrumentationConfig.TEMP_JARS_CLEAN_ON_BOOT, false)
+            && isWindowsOS();
+
     debugEnabled = isDebugMode();
 
     internalExitOnFailure =
@@ -1044,7 +1050,7 @@ private static String getHostName() {
     String possibleHostname;
 
     // Try environment variable.  This works in almost all environments
-    if (System.getProperty("os.name").startsWith("Windows")) {
+    if (isWindowsOS()) {
       possibleHostname = System.getenv("COMPUTERNAME");
     } else {
       possibleHostname = System.getenv("HOSTNAME");
@@ -1079,6 +1085,10 @@ private static String getHostName() {
     return null;
   }
 
+  private static boolean isWindowsOS() {
+    return System.getProperty("os.name").startsWith("Windows");
+  }
+
   // This has to be placed after all other static fields to give them a chance to initialize
   private static final Config INSTANCE = new Config();