Skip to content

[CDAP-21235] : Introduce Spark Task attempt based record count metrics, buffer and dedup while run and Emit only on task success #16182

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
sahusanket wants to merge 1 commit into
base: develop
Choose a base branch
from
Open

[CDAP-21235] : Introduce Spark Task attempt based record count metrics, buffer and dedup while run and Emit only on task success #16182

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 10 additions & 0 deletions cdap-common/src/main/java/io/cdap/cdap/common/conf/Constants.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -982,6 +982,14 @@ public static final class Metrics {
public static final String MESSAGING_TOPIC_NUM = "metrics.messaging.topic.num";

public static final String TWILL_INSTANCE_ID = "metrics.twill.instance.id";

public static final String RAW_SUFFIX = ".raw";
// Suffixes for pipeline stage metrics. Defined here instead of etl-core Constants
// because cdap-spark-core needs to match them but cannot depend on etl-core (circular dependency).
public static final String RECORDS_IN_SUFFIX = ".records.in";
public static final String RECORDS_OUT_SUFFIX = ".records.out";
public static final String RECORDS_ERROR_SUFFIX = ".records.error";
public static final String RECORDS_ALERT_SUFFIX = ".records.alert";

public static final String METRICS_WRITER_EXTENSIONS_DIR = "metrics.writer.extensions.dir";
public static final String METRICS_WRITER_PREFIX = "metrics.writer.";
Expand Down Expand Up @@ -1087,6 +1095,8 @@ public static final class Tag {

public static final String PROVISIONER = "prv";
public static final String SPARK = "sp";
public static final String SPARK_PARTITION = "spark_part";
public static final String SPARK_ATTEMPT = "spark_att";
public static final String STATUS = "st";
public static final String CLUSTER_STATUS = "clst";

Expand Down
96 changes: 94 additions & 2 deletions cdap-spark-core-base/src/main/java/io/cdap/cdap/app/runtime/spark/SparkRuntimeContext.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -45,12 +45,18 @@
import io.cdap.cdap.security.spi.authentication.AuthenticationContext;
import io.cdap.cdap.security.spi.authorization.AccessEnforcer;
import org.apache.hadoop.conf.Configuration;
import org.apache.spark.TaskContext;
import org.apache.spark.util.TaskFailureListener;
import org.apache.tephra.TransactionSystemClient;
import org.apache.twill.api.ServiceAnnouncer;
import org.apache.twill.filesystem.LocationFactory;

import java.io.Closeable;
import java.util.Arrays;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Set;
import javax.annotation.Nullable;

/**
Expand Down Expand Up @@ -112,35 +118,121 @@ true, metricsCollectionService, createMetricsTags(workflowProgramInfo),
this.closeable = closeable;
}

private static final List<String> TARGET_METRIC_SUFFIXES = Arrays.asList(
Constants.Metrics.RECORDS_IN_SUFFIX,
Constants.Metrics.RECORDS_OUT_SUFFIX,
Constants.Metrics.RECORDS_ERROR_SUFFIX,
Constants.Metrics.RECORDS_ALERT_SUFFIX
);

private final ThreadLocal<Map<String, Long>> bufferedCounts = ThreadLocal.withInitial(HashMap::new);

@vsethi09 vsethi09 Aug 13, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why is ThreadLocal needed here?

What is the scope of SparkRuntimeContext object?

@sahusanket sahusanket Aug 15, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

  • On a Spark Executor, SparkRuntimeContext is a process-level singleton shared across the entire JVM.

  • A executor can have 4 threads and each of them will handle 1 partition each . And we are buffering records per stage (example : user.<StageName>.records.in = 5 ) , so there will be collisions accross threads if we don't handle at Thread ( task ) level.

private final ThreadLocal<Boolean> listenerRegistered = ThreadLocal.withInitial(() -> false);
private final ThreadLocal<Boolean> taskFailed = ThreadLocal.withInitial(() -> false);

@Override
public void close() {
super.close();
Closeables.closeQuietly(closeable);
}

private Metrics getTaskMetrics() {
TaskContext tc = TaskContext.get();
if (tc == null) {
return getMetrics();
}
Map<String, String> taskTags = new HashMap<>();
taskTags.put(Constants.Metrics.Tag.SPARK_PARTITION, String.valueOf(tc.partitionId()));
taskTags.put(Constants.Metrics.Tag.SPARK_ATTEMPT, String.valueOf(tc.attemptNumber()));
return getMetrics().child(taskTags);
}

@Override
public void count(String metricName, int delta) {
getMetrics().count(metricName, delta);
bufferMetric(metricName, delta, false);
}

@Override
public void countLong(String metricName, long delta) {
getMetrics().countLong(metricName, delta);
bufferMetric(metricName, delta, false);
}

@Override
public void gauge(String metricName, long value) {
getMetrics().gauge(metricName, value);
bufferMetric(metricName, value, true);
}

private void bufferMetric(String metricName, long value, boolean isGauge) {
TaskContext tc = TaskContext.get();
if (tc == null) {
return;
}
boolean isTarget = false;
for (String suffix : TARGET_METRIC_SUFFIXES) {
if (metricName.endsWith(suffix)) {
isTarget = true;
break;
}
}
if (!isTarget) {
return;
}
registerCompletionListenerIfNeeded(tc);
Map<String, Long> buffer = bufferedCounts.get();
if (isGauge) {
buffer.put(metricName, value);
} else {
buffer.put(metricName, buffer.getOrDefault(metricName, 0L) + value);
}
}

private void registerCompletionListenerIfNeeded(TaskContext tc) {
if (listenerRegistered.get()) {
return;
}
tc.addTaskFailureListener((context, error) -> {
taskFailed.set(true);
});
tc.addTaskCompletionListener(context -> {
try {
if (!taskFailed.get() && !context.isInterrupted()) {
flushBufferedMetrics();
}
} finally {
clearBufferedMetrics();
}
});
listenerRegistered.set(true);
}
Comment thread
sahusanket marked this conversation as resolved.

private void flushBufferedMetrics() {
Map<String, Long> buffer = bufferedCounts.get();
if (buffer.isEmpty()) {
return;
}
Metrics taskMetrics = getTaskMetrics();
for (Map.Entry<String, Long> entry : buffer.entrySet()) {
String rawMetricName = entry.getKey() + Constants.Metrics.RAW_SUFFIX;
taskMetrics.countLong(rawMetricName, entry.getValue());
}
}
Comment thread
sahusanket marked this conversation as resolved.

private void clearBufferedMetrics() {
bufferedCounts.remove();
listenerRegistered.remove();
taskFailed.remove();
}

@Override
public Metrics child(Map<String, String> tags) {
return getMetrics().child(tags);
return getTaskMetrics().child(tags);
}

@Override
public Map<String, String> getTags() {
return getMetrics().getTags();
return getTaskMetrics().getTags();
}

/**
Expand Down
Loading
Loading