Fix muzzle mismatch causing Spark jobs on Databricks 17.3 to not be instrumented (#9872)

* relax muzzle check

* relax more

* spotlessApply

* fix muzzle block by avoiding streaming hard deps and relaxing listener bus matchers

* remove unnecessary defense on StreamExecution

* revert changes on OL listener

* improve getMicroBatchExecutionBatchIdKey

* import InvocationTargetException

* Update AbstractDatadogSparkListener.java

* no need reflection since key is known

* Update AbstractSparkInstrumentation.java

* reflection against SparkPlanInfoAdvice

* spotless

* [DJM-974] Use reflection for constructor in Scala 2.12, lookup by parameter classes (#9886)

---------

Co-authored-by: Charles Yu <charles.yu@datadoghq.com>

Author: aboitreaud

dd-java-agent/instrumentation/spark/spark_2.12/src/main/java/datadog/trace/instrumentation/spark/Spark212Instrumentation.java

@@ -6,6 +6,8 @@
 import com.google.auto.service.AutoService;
 import datadog.trace.agent.tooling.InstrumenterModule;
 import datadog.trace.api.Config;
+import de.thetaphi.forbiddenapis.SuppressForbidden;
+import java.lang.reflect.Constructor;
 import net.bytebuddy.asm.Advice;
 import org.apache.spark.SparkContext;
 import org.apache.spark.sql.execution.SparkPlan;
@@ -14,6 +16,7 @@
 import org.slf4j.LoggerFactory;
 import scala.Predef;
 import scala.collection.JavaConverters;
+import scala.collection.immutable.Map;
 
 @AutoService(InstrumenterModule.class)
 public class Spark212Instrumentation extends AbstractSparkInstrumentation {
@@ -94,21 +97,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())) {
         Spark212PlanSerializer planUtils = new Spark212PlanSerializer();
-        planInfo =
-            new SparkPlanInfo(
-                planInfo.nodeName(),
-                planInfo.simpleString(),
-                planInfo.children(),
-                JavaConverters.mapAsScalaMap(planUtils.extractFormattedProduct(plan))
-                    .toMap(Predef.$conforms()),
-                planInfo.metrics());
+        Map<String, String> meta =
+            JavaConverters.mapAsScalaMap(planUtils.extractFormattedProduct(plan))
+                .toMap(Predef.$conforms());
+        try {
+          Constructor<?> targetCtor = null;
+          for (Constructor<?> c : SparkPlanInfo.class.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/spark_2.13/src/main/java/datadog/trace/instrumentation/spark/Spark213Instrumentation.java

@@ -6,6 +6,8 @@
 import com.google.auto.service.AutoService;
 import datadog.trace.agent.tooling.InstrumenterModule;
 import datadog.trace.api.Config;
+import de.thetaphi.forbiddenapis.SuppressForbidden;
+import java.lang.reflect.Constructor;
 import net.bytebuddy.asm.Advice;
 import org.apache.spark.SparkContext;
 import org.apache.spark.sql.execution.SparkPlan;
@@ -14,6 +16,7 @@
 import org.slf4j.LoggerFactory;
 import scala.collection.JavaConverters;
 import scala.collection.immutable.HashMap;
+import scala.collection.immutable.Map;
 
 @AutoService(InstrumenterModule.class)
 public class Spark213Instrumentation extends AbstractSparkInstrumentation {
@@ -95,20 +98,36 @@ 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());
+        Map<String, String> meta =
+            HashMap.from(JavaConverters.asScala(planUtils.extractFormattedProduct(plan)));
+        try {
+          Constructor<?> targetCtor = null;
+          for (Constructor<?> c : SparkPlanInfo.class.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/AbstractDatadogSparkListener.java

@@ -43,7 +43,6 @@
 import org.apache.spark.sql.execution.SQLExecution;
 import org.apache.spark.sql.execution.SparkPlanInfo;
 import org.apache.spark.sql.execution.metric.SQLMetricInfo;
-import org.apache.spark.sql.execution.streaming.MicroBatchExecution;
 import org.apache.spark.sql.execution.streaming.StreamExecution;
 import org.apache.spark.sql.execution.ui.SparkListenerSQLExecutionEnd;
 import org.apache.spark.sql.execution.ui.SparkListenerSQLExecutionStart;
@@ -1243,7 +1242,7 @@ private static String getStreamingBatchKey(Properties properties) {
     }
 
     Object queryId = properties.get(StreamExecution.QUERY_ID_KEY());
-    Object batchId = properties.get(MicroBatchExecution.BATCH_ID_KEY());
+    Object batchId = properties.get("streaming.sql.batchId");
 
     if (queryId == null || batchId == null) {
       return null;