Configure Circuit breaker in OpenLineage instrumentation (#9554)

pawel-big-lebowski · web-flow · commit 40ebe1b180c5 · 2025-10-01T10:17:57.000+02:00
* configure circuit breaker in OpenLineage instrumentation
diff --git a/dd-java-agent/instrumentation/spark/src/main/java/datadog/trace/instrumentation/spark/AbstractDatadogSparkListener.java b/dd-java-agent/instrumentation/spark/src/main/java/datadog/trace/instrumentation/spark/AbstractDatadogSparkListener.java
@@ -76,6 +76,7 @@ public abstract class AbstractDatadogSparkListener extends SparkListener {
   private final int MAX_ACCUMULATOR_SIZE = 50000;
   private final String RUNTIME_TAGS_PREFIX = "spark.datadog.tags.";
   private static final String AGENT_OL_ENDPOINT = "openlineage/api/v1/lineage";
+  private static final int OL_CIRCUIT_BREAKER_TIMEOUT_IN_SECONDS = 60;
 
   public volatile SparkListenerInterface openLineageSparkListener = null;
   public volatile SparkConf openLineageSparkConf = null;
@@ -186,6 +187,7 @@ public void setupOpenLineage(DDTraceId traceId) {
               + ProcessTags.getTagsForSerialization()
               + ";_dd.ol_app_id:"
               + appId);
+      setupOpenLineageCircuitBreaker();
       return;
     }
     log.debug(
@@ -1324,6 +1326,32 @@ private static String getServiceForOpenLineage(SparkConf conf, boolean isRunning
     return sparkAppName;
   }
 
+  private void setupOpenLineageCircuitBreaker() {
+    if (!Config.get().isDataJobsOpenLineageTimeoutEnabled()) {
+      log.debug("Data Jobs OpenLineage timeout is not enabled");
+      return;
+    }
+    if (!classIsLoadable("io.openlineage.client.circuitBreaker.TimeoutCircuitBreaker")) {
+      log.debug(
+          "OpenLineage version without timeout circuit breaker. Probably OL version < 1.35.0");
+      return;
+    }
+    if (openLineageSparkConf.contains("spark.openlineage.circuitBreaker.type")) {
+      log.debug(
+          "Other OpenLineage circuit breaker already configured: {}",
+          openLineageSparkConf.get("spark.openlineage.circuitBreaker.type"));
+      return;
+    }
+
+    openLineageSparkConf.set("spark.openlineage.circuitBreaker.type", "timeout");
+    openLineageSparkConf.setIfMissing(
+        "spark.openlineage.circuitBreaker.timeoutInSeconds",
+        String.valueOf(OL_CIRCUIT_BREAKER_TIMEOUT_IN_SECONDS));
+    log.debug(
+        "Setting OpenLineage circuit breaker with timeout {} seconds",
+        openLineageSparkConf.get("spark.openlineage.circuitBreaker.timeoutInSeconds"));
+  }
+
   private static void reportKafkaOffsets(
       final String appName, final AgentSpan span, final SourceProgress progress) {
     if (!traceConfig().isDataStreamsEnabled()
diff --git a/dd-java-agent/instrumentation/spark/src/testFixtures/groovy/datadog/trace/instrumentation/spark/AbstractSparkListenerTest.groovy b/dd-java-agent/instrumentation/spark/src/testFixtures/groovy/datadog/trace/instrumentation/spark/AbstractSparkListenerTest.groovy
@@ -568,6 +568,50 @@ abstract class AbstractSparkListenerTest extends InstrumentationSpecification {
     .contains("_dd.ol_intake.process_tags:" + ProcessTags.getTagsForSerialization())
   }
 
+  def "test setupOpenLineage fills circuit breaker config"(
+    Boolean configEnabled,
+    String sparkConfCircuitBreakerType,
+    String expectedCircuitBreakerType
+  ) {
+    setup:
+    injectSysConfig("data.jobs.openlineage.timeout.enabled", configEnabled.toString())
+    def listener = getTestDatadogSparkListener()
+    listener.openLineageSparkListener = Mock(SparkListenerInterface)
+    listener.openLineageSparkConf = new SparkConf()
+    if (sparkConfCircuitBreakerType != null) {
+      listener.openLineageSparkConf.set("spark.openlineage.circuitBreaker.type", sparkConfCircuitBreakerType)
+    }
+    listener.setupOpenLineage(Mock(DDTraceId))
+
+    expect:
+    assert listener
+    .openLineageSparkConf
+    .getOption("spark.openlineage.circuitBreaker.type") == Option.apply(expectedCircuitBreakerType)
+    assert listener
+    .openLineageSparkConf
+    .getOption("spark.openlineage.circuitBreaker.timeoutInSeconds") == ((expectedCircuitBreakerType == "timeout") ? Option.apply("60") : Option.apply(null))
+
+    where:
+    configEnabled | sparkConfCircuitBreakerType  | expectedCircuitBreakerType
+    true          | null                         | "timeout"
+    true          | "other"                      | "other"
+    false         | null                         | null
+  }
+
+  def "test OpenLineage's circuit breaker timeout can be overriden"() {
+    setup:
+    def listener = getTestDatadogSparkListener()
+    listener.openLineageSparkListener = Mock(SparkListenerInterface)
+    listener.openLineageSparkConf = new SparkConf()
+    listener.openLineageSparkConf.set("spark.openlineage.circuitBreaker.timeoutInSeconds", "120")
+    listener.setupOpenLineage(Mock(DDTraceId))
+
+    expect:
+    assert listener
+    .openLineageSparkConf
+    .getOption("spark.openlineage.circuitBreaker.timeoutInSeconds") == Option.apply("120")
+  }
+
   protected validateRelativeError(double value, double expected, double relativeAccuracy) {
     double relativeError = Math.abs(value - expected) / expected
     assert relativeError < relativeAccuracy
diff --git a/dd-java-agent/instrumentation/spark/src/testFixtures/java/io/openlineage/client/circuitBreaker/TimeoutCircuitBreaker.java b/dd-java-agent/instrumentation/spark/src/testFixtures/java/io/openlineage/client/circuitBreaker/TimeoutCircuitBreaker.java
@@ -0,0 +1,7 @@
+package io.openlineage.client.circuitBreaker;
+
+/**
+ * Existence of the class with this name on a classpath is required to test OpenLineage circuit
+ * breaker setup
+ */
+public class TimeoutCircuitBreaker {}
diff --git a/dd-trace-api/src/main/java/datadog/trace/api/ConfigDefaults.java b/dd-trace-api/src/main/java/datadog/trace/api/ConfigDefaults.java
@@ -228,6 +228,7 @@ public final class ConfigDefaults {
 
   static final boolean DEFAULT_DATA_JOBS_ENABLED = false;
   static final boolean DEFAULT_DATA_JOBS_OPENLINEAGE_ENABLED = false;
+  static final boolean DEFAULT_DATA_JOBS_OPENLINEAGE_TIMEOUT_ENABLED = true;
 
   static final boolean DEFAULT_DATA_STREAMS_ENABLED = false;
   static final int DEFAULT_DATA_STREAMS_BUCKET_DURATION = 10; // seconds
diff --git a/dd-trace-api/src/main/java/datadog/trace/api/config/GeneralConfig.java b/dd-trace-api/src/main/java/datadog/trace/api/config/GeneralConfig.java
@@ -77,6 +77,8 @@ public final class GeneralConfig {
   public static final String DATA_JOBS_ENABLED = "data.jobs.enabled";
   public static final String DATA_JOBS_COMMAND_PATTERN = "data.jobs.command.pattern";
   public static final String DATA_JOBS_OPENLINEAGE_ENABLED = "data.jobs.openlineage.enabled";
+  public static final String DATA_JOBS_OPENLINEAGE_TIMEOUT_ENABLED =
+      "data.jobs.openlineage.timeout.enabled";
 
   public static final String DATA_STREAMS_ENABLED = "data.streams.enabled";
   public static final String DATA_STREAMS_BUCKET_DURATION_SECONDS =
diff --git a/internal-api/src/main/java/datadog/trace/api/Config.java b/internal-api/src/main/java/datadog/trace/api/Config.java
@@ -47,6 +47,7 @@
 import static datadog.trace.api.ConfigDefaults.DEFAULT_CWS_TLS_REFRESH;
 import static datadog.trace.api.ConfigDefaults.DEFAULT_DATA_JOBS_ENABLED;
 import static datadog.trace.api.ConfigDefaults.DEFAULT_DATA_JOBS_OPENLINEAGE_ENABLED;
+import static datadog.trace.api.ConfigDefaults.DEFAULT_DATA_JOBS_OPENLINEAGE_TIMEOUT_ENABLED;
 import static datadog.trace.api.ConfigDefaults.DEFAULT_DATA_STREAMS_BUCKET_DURATION;
 import static datadog.trace.api.ConfigDefaults.DEFAULT_DATA_STREAMS_ENABLED;
 import static datadog.trace.api.ConfigDefaults.DEFAULT_DB_CLIENT_HOST_SPLIT_BY_HOST;
@@ -331,6 +332,7 @@
 import static datadog.trace.api.config.GeneralConfig.AZURE_APP_SERVICES;
 import static datadog.trace.api.config.GeneralConfig.DATA_JOBS_ENABLED;
 import static datadog.trace.api.config.GeneralConfig.DATA_JOBS_OPENLINEAGE_ENABLED;
+import static datadog.trace.api.config.GeneralConfig.DATA_JOBS_OPENLINEAGE_TIMEOUT_ENABLED;
 import static datadog.trace.api.config.GeneralConfig.DATA_STREAMS_BUCKET_DURATION_SECONDS;
 import static datadog.trace.api.config.GeneralConfig.DATA_STREAMS_ENABLED;
 import static datadog.trace.api.config.GeneralConfig.DOGSTATSD_ARGS;
@@ -1172,6 +1174,7 @@ public static String getHostName() {
 
   private final boolean dataJobsEnabled;
   private final boolean dataJobsOpenLineageEnabled;
+  private final boolean dataJobsOpenLineageTimeoutEnabled;
 
   private final boolean dataStreamsEnabled;
   private final float dataStreamsBucketDurationSeconds;
@@ -2595,6 +2598,9 @@ PROFILING_DATADOG_PROFILER_ENABLED, isDatadogProfilerSafeInCurrentEnvironment())
     dataJobsOpenLineageEnabled =
         configProvider.getBoolean(
             DATA_JOBS_OPENLINEAGE_ENABLED, DEFAULT_DATA_JOBS_OPENLINEAGE_ENABLED);
+    dataJobsOpenLineageTimeoutEnabled =
+        configProvider.getBoolean(
+            DATA_JOBS_OPENLINEAGE_TIMEOUT_ENABLED, DEFAULT_DATA_JOBS_OPENLINEAGE_TIMEOUT_ENABLED);
 
     dataStreamsEnabled =
         configProvider.getBoolean(DATA_STREAMS_ENABLED, DEFAULT_DATA_STREAMS_ENABLED);
@@ -4499,6 +4505,10 @@ public boolean isDataJobsOpenLineageEnabled() {
     return dataJobsOpenLineageEnabled;
   }
 
+  public boolean isDataJobsOpenLineageTimeoutEnabled() {
+    return dataJobsOpenLineageTimeoutEnabled;
+  }
+
   public boolean isApmTracingEnabled() {
     return apmTracingEnabled;
   }
@@ -5827,6 +5837,10 @@ public String toString() {
         + appSecRaspEnabled
         + ", dataJobsEnabled="
         + dataJobsEnabled
+        + ", dataJobsOpenLineageEnabled="
+        + dataJobsOpenLineageEnabled
+        + ", dataJobsOpenLineageTimeoutEnabled="
+        + dataJobsOpenLineageTimeoutEnabled
         + ", apmTracingEnabled="
         + apmTracingEnabled
         + ", jdkSocketEnabled="
