reflection against SparkPlanInfoAdvice

Author: aboitreaud

dd-java-agent/instrumentation/spark/spark_2.13/src/main/java/datadog/trace/instrumentation/spark/Spark213Instrumentation.java

@@ -7,6 +7,7 @@
 import datadog.trace.agent.tooling.InstrumenterModule;
 import datadog.trace.api.Config;
 import net.bytebuddy.asm.Advice;
+import de.thetaphi.forbiddenapis.SuppressForbidden;
 import org.apache.spark.SparkContext;
 import org.apache.spark.sql.execution.SparkPlan;
 import org.apache.spark.sql.execution.SparkPlanInfo;
@@ -95,20 +96,37 @@ public static void enter(@Advice.This SparkContext sparkContext) {
 
   public static class SparkPlanInfoAdvice {
     @Advice.OnMethodExit(suppress = Throwable.class, onThrowable = Throwable.class)
+    @SuppressForbidden
     public static void exit(
         @Advice.Return(readOnly = false) SparkPlanInfo planInfo,
         @Advice.Argument(0) SparkPlan plan) {
       if (planInfo.metadata().size() == 0
           && (Config.get().isDataJobsParseSparkPlanEnabled()
               || Config.get().isDataJobsExperimentalFeaturesEnabled())) {
         Spark213PlanSerializer planUtils = new Spark213PlanSerializer();
-        planInfo =
-            new SparkPlanInfo(
-                planInfo.nodeName(),
-                planInfo.simpleString(),
-                planInfo.children(),
-                HashMap.from(JavaConverters.asScala(planUtils.extractFormattedProduct(plan))),
-                planInfo.metrics());
+        scala.collection.immutable.Map<String, String> meta =
+            HashMap.from(JavaConverters.asScala(planUtils.extractFormattedProduct(plan)));
+        try {
+          Class<?> spiClass = Class.forName("org.apache.spark.sql.execution.SparkPlanInfo");
+          java.lang.reflect.Constructor<?> targetCtor = null;
+          for (java.lang.reflect.Constructor<?> c : spiClass.getConstructors()) {
+            if (c.getParameterCount() == 5) {
+              targetCtor = c;
+              break;
+            }
+          }
+          if (targetCtor != null) {
+            Object newInst =
+                targetCtor.newInstance(
+                    planInfo.nodeName(),
+                    planInfo.simpleString(),
+                    planInfo.children(),
+                    meta,
+                    planInfo.metrics());
+            planInfo = (SparkPlanInfo) newInst;
+          }
+        } catch (Throwable ignored) {
+        }
       }
     }
   }

dd-java-agent/instrumentation/spark/src/main/java/datadog/trace/instrumentation/spark/AbstractSparkInstrumentation.java

@@ -5,7 +5,6 @@
 import static net.bytebuddy.matcher.ElementMatchers.isMethod;
 import static net.bytebuddy.matcher.ElementMatchers.nameEndsWith;
 import static net.bytebuddy.matcher.ElementMatchers.takesArgument;
-import static net.bytebuddy.matcher.ElementMatchers.takesArguments;
 
 import datadog.trace.agent.tooling.Instrumenter;
 import datadog.trace.agent.tooling.InstrumenterModule;
@@ -69,8 +68,8 @@ public void methodAdvice(MethodTransformer transformer) {
     // LiveListenerBus class is used to manage spark listeners
     transformer.applyAdvice(
         isMethod()
-            .and(named("addToSharedQueue").or(named("addToQueue")))
-            .and(takesArguments(1))
+            .and(named("addToSharedQueue"))
+            .and(takesArgument(0, named("org.apache.spark.scheduler.SparkListenerInterface")))
             .and(isDeclaredBy(named("org.apache.spark.scheduler.LiveListenerBus"))),
         AbstractSparkInstrumentation.class.getName() + "$LiveListenerBusAdvice");
   }