Track external accumulators in tracer instead of using SparkInfo values

Author: charlesmyu

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

@@ -40,6 +40,7 @@
 import org.apache.spark.ExceptionFailure;
 import org.apache.spark.SparkConf;
 import org.apache.spark.TaskFailedReason;
+import org.apache.spark.executor.TaskMetrics;
 import org.apache.spark.scheduler.AccumulableInfo;
 import org.apache.spark.scheduler.JobFailed;
 import org.apache.spark.scheduler.SparkListener;
@@ -64,6 +65,7 @@
 import org.apache.spark.sql.streaming.StateOperatorProgress;
 import org.apache.spark.sql.streaming.StreamingQueryListener;
 import org.apache.spark.sql.streaming.StreamingQueryProgress;
+import org.apache.spark.util.AccumulatorV2;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import scala.Tuple2;
@@ -127,8 +129,10 @@ public abstract class AbstractDatadogSparkListener extends SparkListener {
   private final HashMap<String, SparkListenerExecutorAdded> liveExecutors = new HashMap<>();
 
   // There is no easy way to know if an accumulator is not useful anymore (meaning it is not part of
-  // an active SQL query)
-  // so capping the size of the collection storing them
+  // an active SQL query) so capping the size of the collection storing them
+  // TODO (CY): Is this potentially the reason why some Spark Plans aren't showing up consistently?
+  // If we know we don't need the accumulator values, can we drop all associated data and just map 
+  // stage ID -> accumulator ID? Put this behind some FF
   private final Map<Long, SparkSQLUtils.AccumulatorWithStage> accumulators =
       new RemoveEldestHashMap<>(MAX_ACCUMULATOR_SIZE);
 
@@ -151,6 +155,8 @@ public abstract class AbstractDatadogSparkListener extends SparkListener {
   public AbstractDatadogSparkListener(SparkConf sparkConf, String appId, String sparkVersion) {
     tracer = AgentTracer.get();
 
+    log.error("[CHARLES] HELLO WORLD");
+
     this.sparkConf = sparkConf;
     this.appId = appId;
     this.sparkVersion = sparkVersion;
@@ -229,6 +235,9 @@ public void setupOpenLineage(DDTraceId traceId) {
   /** Parent Ids of a Stage. Provide an implementation based on a specific scala version */
   protected abstract int[] getStageParentIds(StageInfo info);
 
+  /** All External Accumulators associated with a given task. Provide an implementation based on a specific scala version */
+  protected abstract List<AccumulatorV2> getExternalAccumulators(TaskMetrics metrics);
+
   @Override
   public synchronized void onApplicationStart(SparkListenerApplicationStart applicationStart) {
     this.applicationStart = applicationStart;
@@ -670,7 +679,8 @@ public synchronized void onStageCompleted(SparkListenerStageCompleted stageCompl
 
     SparkPlanInfo sqlPlan = sqlPlans.get(sqlExecutionId);
     if (sqlPlan != null) {
-      SparkSQLUtils.addSQLPlanToStageSpan(span, sqlPlan, accumulators, stageId);
+      SparkSQLUtils.addSQLPlanToStageSpan(span, sqlPlan, accumulators, stageMetric, stageId);
+      log.info("[CHARLES]", span.getTag("_dd.spark.sql_plan"));
     }
 
     span.finish(completionTimeMs * 1000);
@@ -684,7 +694,9 @@ public void onTaskEnd(SparkListenerTaskEnd taskEnd) {
 
     SparkAggregatedTaskMetrics stageMetric = stageMetrics.get(stageSpanKey);
     if (stageMetric != null) {
-      stageMetric.addTaskMetrics(taskEnd);
+      // Not happy that we have to extract external accumulators here, but needed as we're dealing with Seq
+      // which varies across Scala versions
+      stageMetric.addTaskMetrics(taskEnd, getExternalAccumulators(taskEnd.taskMetrics()));
     }
 
     if (taskEnd.taskMetrics() != null) {

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

@@ -1,13 +1,20 @@
 package datadog.trace.instrumentation.spark;
 
+import com.fasterxml.jackson.core.JsonGenerator;
 import datadog.metrics.api.Histogram;
 import datadog.trace.api.Config;
 import datadog.trace.bootstrap.instrumentation.api.AgentSpan;
+import java.io.IOException;
 import java.nio.ByteBuffer;
 import java.util.Base64;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
 import org.apache.spark.TaskFailedReason;
 import org.apache.spark.executor.TaskMetrics;
 import org.apache.spark.scheduler.SparkListenerTaskEnd;
+import org.apache.spark.sql.execution.metric.SQLMetricInfo;
+import org.apache.spark.util.AccumulatorV2;
 
 class SparkAggregatedTaskMetrics {
   private static final double HISTOGRAM_RELATIVE_ACCURACY = 1 / 32.0;
@@ -59,13 +66,16 @@ class SparkAggregatedTaskMetrics {
   private Histogram shuffleWriteBytesHistogram;
   private Histogram diskBytesSpilledHistogram;
 
+  // Used for Spark SQL Plan metrics ONLY, don't put in regular span for now
+  private Map<Long, Histogram> externalAccumulableHistograms;
+
   public SparkAggregatedTaskMetrics() {}
 
   public SparkAggregatedTaskMetrics(long availableExecutorTime) {
     this.previousAvailableExecutorTime = availableExecutorTime;
   }
 
-  public void addTaskMetrics(SparkListenerTaskEnd taskEnd) {
+  public void addTaskMetrics(SparkListenerTaskEnd taskEnd, List<AccumulatorV2> externalAccumulators) {
     taskCompletedCount += 1;
 
     if (taskEnd.taskInfo().attemptNumber() > 0) {
@@ -127,6 +137,24 @@ public void addTaskMetrics(SparkListenerTaskEnd taskEnd) {
                 shuffleWriteBytesHistogram, taskMetrics.shuffleWriteMetrics().bytesWritten());
         diskBytesSpilledHistogram =
             lazyHistogramAccept(diskBytesSpilledHistogram, taskMetrics.diskBytesSpilled());
+
+        // TODO (CY): Should we also look at TaskInfo accumulable update values as a backup? Is that
+        // only needed for SHS?
+        if (externalAccumulators != null && !externalAccumulators.isEmpty()) {
+          if (externalAccumulableHistograms == null) {
+            externalAccumulableHistograms = new HashMap<>(externalAccumulators.size());
+          }
+
+          externalAccumulators.forEach(acc -> {
+            Histogram hist = externalAccumulableHistograms.get(acc.id());
+            try {
+              // As of spark 3.5, all SQL metrics are Long, safeguard if it changes in new
+              // versions
+              externalAccumulableHistograms.put(
+                  acc.id(), lazyHistogramAccept(hist, (Long) acc.value()));
+            } catch (ClassCastException ignored) {}
+          });
+        }
       }
     }
   }
@@ -276,6 +304,19 @@ private Histogram lazyHistogramAccept(Histogram hist, double value) {
     return hist;
   }
 
+  // Used to put external accum metrics to JSON for Spark SQL plans
+  public void externalAccumToJson(JsonGenerator generator, SQLMetricInfo info) throws IOException {
+    Histogram hist = externalAccumulableHistograms.get(info.accumulatorId());
+    String name = info.name();
+
+    if (name != null && hist != null) {
+      generator.writeStartObject();
+      generator.writeStringField(name, histogramToBase64(hist));
+      generator.writeStringField("type", info.metricType());
+      generator.writeEndObject();
+    }
+  }
+
   public static long computeTaskRunTime(TaskMetrics metrics) {
     return metrics.executorDeserializeTime()
         + metrics.executorRunTime()

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

@@ -24,6 +24,7 @@ public static void addSQLPlanToStageSpan(
       AgentSpan span,
       SparkPlanInfo plan,
       Map<Long, AccumulatorWithStage> accumulators,
+      SparkAggregatedTaskMetrics stageMetric,
       int stageId) {
     Set<Integer> parentStageIds = new HashSet<>();
     SparkPlanInfoForStage planForStage =
@@ -32,7 +33,7 @@ public static void addSQLPlanToStageSpan(
     span.setTag("_dd.spark.sql_parent_stage_ids", parentStageIds.toString());
 
     if (planForStage != null) {
-      String json = planForStage.toJson(accumulators);
+      String json = planForStage.toJson(stageMetric);
       span.setTag("_dd.spark.sql_plan", json);
     }
   }
@@ -143,15 +144,15 @@ public SparkPlanInfoForStage(SparkPlanInfo plan, List<SparkPlanInfoForStage> chi
       this.children = children;
     }
 
-    public String toJson(Map<Long, AccumulatorWithStage> accumulators) {
+    public String toJson(SparkAggregatedTaskMetrics stageMetric) {
       // Using the jackson JSON lib used by spark
       // https://mvnrepository.com/artifact/org.apache.spark/spark-core_2.12/3.5.0
       ObjectMapper mapper =
           new ObjectMapper().configure(DeserializationFeature.FAIL_ON_UNKNOWN_PROPERTIES, false);
 
       ByteArrayOutputStream baos = new ByteArrayOutputStream();
       try (JsonGenerator generator = mapper.getFactory().createGenerator(baos)) {
-        this.toJson(generator, accumulators, mapper);
+        this.toJson(generator, mapper, stageMetric);
       } catch (IOException e) {
         return null;
       }
@@ -160,7 +161,7 @@ public String toJson(Map<Long, AccumulatorWithStage> accumulators) {
     }
 
     private void toJson(
-        JsonGenerator generator, Map<Long, AccumulatorWithStage> accumulators, ObjectMapper mapper)
+        JsonGenerator generator, ObjectMapper mapper, SparkAggregatedTaskMetrics stageMetric)
         throws IOException {
       generator.writeStartObject();
       generator.writeStringField("node", plan.nodeName());
@@ -199,11 +200,7 @@ private void toJson(
         generator.writeFieldName("metrics");
         generator.writeStartArray();
         for (SQLMetricInfo metric : metrics) {
-          long accumulatorId = metric.accumulatorId();
-          AccumulatorWithStage acc = accumulators.get(accumulatorId);
-          if (acc != null) {
-            acc.toJson(generator, metric);
-          }
+          stageMetric.externalAccumToJson(generator, metric);
         }
         generator.writeEndArray();
       }
@@ -213,7 +210,7 @@ private void toJson(
         generator.writeFieldName("children");
         generator.writeStartArray();
         for (SparkPlanInfoForStage child : children) {
-          child.toJson(generator, accumulators, mapper);
+          child.toJson(generator, mapper, stageMetric);
         }
         generator.writeEndArray();
       }

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

@@ -1,14 +1,22 @@
 package datadog.trace.instrumentation.spark;
 
+import java.lang.invoke.MethodHandle;
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.List;
+import java.util.Map;
+import datadog.trace.util.MethodHandles;
 import org.apache.spark.SparkConf;
+import org.apache.spark.executor.TaskMetrics;
 import org.apache.spark.scheduler.SparkListenerJobStart;
 import org.apache.spark.scheduler.StageInfo;
 import org.apache.spark.sql.execution.SparkPlanInfo;
 import org.apache.spark.sql.execution.metric.SQLMetricInfo;
+import org.apache.spark.util.AccumulatorV2;
+import scala.Function;
+import scala.Function1;
 import scala.collection.JavaConverters;
+import scala.collection.mutable.ArrayBuffer;
 
 /**
  * DatadogSparkListener compiled for Scala 2.12
@@ -17,6 +25,10 @@
  * compiled with the specific scala version
  */
 public class DatadogSpark212Listener extends AbstractDatadogSparkListener {
+  private static final MethodHandles methodLoader = new MethodHandles(ClassLoader.getSystemClassLoader());
+  private static final MethodHandle externalAccums = methodLoader.method(TaskMetrics.class, "externalAccums");
+  private static final MethodHandle withExternalAccums = methodLoader.method(TaskMetrics.class, "withExternalAccums", new Class[] {});
+
   public DatadogSpark212Listener(SparkConf sparkConf, String appId, String sparkVersion) {
     super(sparkConf, appId, sparkVersion);
   }
@@ -62,4 +74,25 @@ protected int[] getStageParentIds(StageInfo info) {
 
     return parentIds;
   }
+
+  @Override
+  protected List<AccumulatorV2> getExternalAccumulators(TaskMetrics metrics) {
+      if (metrics == null) {
+        return null;
+      }
+
+      Function1 lambda = (Function1<ArrayBuffer<AccumulatorV2>, List<AccumulatorV2>>) accumulators -> JavaConverters.seqAsJavaList(accumulators);
+      List<AccumulatorV2> res = methodLoader.invoke(withExternalAccums, metrics, lambda);
+      if (res != null) {
+        return res;
+      }
+
+      // withExternalAccums didn't work, try the legacy method
+      ArrayBuffer<AccumulatorV2> accumulators = methodLoader.invoke(externalAccums, metrics);
+      if (accumulators != null && !accumulators.isEmpty()) {
+        return JavaConverters.seqAsJavaList(accumulators);
+      }
+
+      return null;
+  }
 }

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

@@ -1,13 +1,20 @@
 package datadog.trace.instrumentation.spark;
 
+import java.lang.invoke.MethodHandle;
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.List;
+import datadog.trace.util.MethodHandles;
 import org.apache.spark.SparkConf;
+import org.apache.spark.executor.TaskMetrics;
 import org.apache.spark.scheduler.SparkListenerJobStart;
 import org.apache.spark.scheduler.StageInfo;
 import org.apache.spark.sql.execution.SparkPlanInfo;
 import org.apache.spark.sql.execution.metric.SQLMetricInfo;
+import org.apache.spark.util.AccumulatorV2;
+import scala.Function1;
+import scala.collection.JavaConverters;
+import scala.collection.mutable.ArrayBuffer;
 import scala.jdk.javaapi.CollectionConverters;
 
 /**
@@ -17,6 +24,10 @@
  * compiled with the specific scala version
  */
 public class DatadogSpark213Listener extends AbstractDatadogSparkListener {
+  private static final MethodHandles methodLoader = new MethodHandles(ClassLoader.getSystemClassLoader());
+  private static final MethodHandle externalAccums = methodLoader.method(TaskMetrics.class, "externalAccums");
+  private static final MethodHandle withExternalAccums = methodLoader.method(TaskMetrics.class, "withExternalAccums", new Class[] {});
+
   public DatadogSpark213Listener(SparkConf sparkConf, String appId, String sparkVersion) {
     super(sparkConf, appId, sparkVersion);
   }
@@ -62,4 +73,25 @@ protected int[] getStageParentIds(StageInfo info) {
 
     return parentIds;
   }
+
+  @Override
+  protected List<AccumulatorV2> getExternalAccumulators(TaskMetrics metrics) {
+    if (metrics == null) {
+      return null;
+    }
+
+    Function1 lambda = (Function1<ArrayBuffer<AccumulatorV2>, List<AccumulatorV2>>) accumulators -> CollectionConverters.asJava(accumulators);
+    List<AccumulatorV2> res = methodLoader.invoke(withExternalAccums, metrics, lambda);
+    if (res != null) {
+      return res;
+    }
+
+    // withExternalAccums didn't work, try the legacy method
+    ArrayBuffer<AccumulatorV2> accumulators = methodLoader.invoke(externalAccums, metrics);
+    if (accumulators != null && !accumulators.isEmpty()) {
+      return CollectionConverters.asJava(accumulators);
+    }
+
+    return null;
+  }
 }