Use consistent approach to apply temporary system-property overrides that have a desired side-effect

Author: mcculls

dd-java-agent/agent-builder/src/main/java/datadog/trace/agent/tooling/AgentInstaller.java

@@ -33,11 +33,12 @@
 import java.util.Set;
 import java.util.concurrent.CopyOnWriteArrayList;
 import java.util.concurrent.TimeUnit;
+import java.util.function.BooleanSupplier;
 import net.bytebuddy.ByteBuddy;
 import net.bytebuddy.agent.builder.AgentBuilder;
-import net.bytebuddy.description.type.TypeDefinition;
 import net.bytebuddy.description.type.TypeDescription;
 import net.bytebuddy.dynamic.DynamicType;
+import net.bytebuddy.dynamic.NexusAccessor;
 import net.bytebuddy.dynamic.VisibilityBridgeStrategy;
 import net.bytebuddy.dynamic.scaffold.InstrumentedType;
 import net.bytebuddy.dynamic.scaffold.MethodGraph;
@@ -54,7 +55,7 @@ public class AgentInstaller {
   private static final List<Runnable> MBEAN_SERVER_BUILDER_CALLBACKS = new CopyOnWriteArrayList<>();
 
   static {
-    addByteBuddyRawSetting();
+    enableByteBuddyRawTypes();
     disableByteBuddyNexus();
     // register weak map supplier as early as possible
     WeakMaps.registerAsSupplier();
@@ -320,27 +321,39 @@ public static Set<InstrumenterModule.TargetSystem> getEnabledSystems() {
     return enabledSystems;
   }
 
-  private static void addByteBuddyRawSetting() {
-    final String savedPropertyValue = SystemProperties.get(TypeDefinition.RAW_TYPES_PROPERTY);
-    final boolean overridden = SystemProperties.set(TypeDefinition.RAW_TYPES_PROPERTY, "true");
-    final boolean rawTypes = TypeDescription.AbstractBase.RAW_TYPES;
-    if (!rawTypes && DEBUG) {
-      log.debug("Too late to enable {}", TypeDefinition.RAW_TYPES_PROPERTY);
+  private static void enableByteBuddyRawTypes() {
+    temporaryOverride("net.bytebuddy.raw", "true", AgentInstaller::rawTypesEnabled);
+  }
+
+  private static boolean rawTypesEnabled() {
+    return TypeDescription.AbstractBase.RAW_TYPES; // must avoid touching this before the override
+  }
+
+  private static void disableByteBuddyNexus() {
+    // disable byte-buddy's Nexus mechanism (we don't need it, and it triggers use of Unsafe)
+    temporaryOverride("net.bytebuddy.nexus.disabled", "true", AgentInstaller::nexusDisabled);
+  }
+
+  private static boolean nexusDisabled() {
+    return !NexusAccessor.isAlive(); // must avoid touching this before the override
+  }
+
+  /** Temporarily overrides a system property while checking it's had the intended side effect. */
+  private static void temporaryOverride(String key, String value, BooleanSupplier sideEffect) {
+    final String savedPropertyValue = SystemProperties.get(key);
+    final boolean overridden = SystemProperties.set(key, value);
+    if (!sideEffect.getAsBoolean() && DEBUG) {
+      log.debug("Too late to apply {}={}", key, value);
     }
     if (overridden) {
       if (savedPropertyValue == null) {
-        SystemProperties.clear(TypeDefinition.RAW_TYPES_PROPERTY);
+        SystemProperties.clear(key);
       } else {
-        SystemProperties.set(TypeDefinition.RAW_TYPES_PROPERTY, savedPropertyValue);
+        SystemProperties.set(key, savedPropertyValue);
       }
     }
   }
 
-  private static void disableByteBuddyNexus() {
-    // disable byte-buddy's Nexus mechanism (we don't need it, and it triggers use of Unsafe)
-    SystemProperties.set("net.bytebuddy.nexus.disabled", "true");
-  }
-
   private static AgentBuilder.RedefinitionStrategy.Listener redefinitionStrategyListener(
       final Set<InstrumenterModule.TargetSystem> enabledSystems) {
     if (enabledSystems.contains(InstrumenterModule.TargetSystem.IAST)) {