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

[sahusanket]

Deduplicate Spark Retry Metrics by Introducing Executor-Side Buffering

1. Context

During CDAP Spark pipeline executions, stages report processing metrics (e.g. records.in and records.out) to the CDAP Master. To facilitate diagnostic monitoring, these are enriched with partition and attempt tags under a .raw suffix (e.g. user.stage_name.records.out.raw).

Currently, when Spark tasks fail and retry, both the failed and successful attempts write their counts immediately to the metrics database. Consequently, database queries using standard SUM aggregates double-count records processed during failed attempts, causing audit mismatches.

---

2. Approach

We introduce executor-side metrics buffering to separate real-time execution feedback from the final database records:

1. Immediate Progress: Legacy flat progress metrics (e.g. records.out) bypass the buffer and are sent immediately to preserve real-time UI tracking.
2. Deduplicated raw metrics: Suffix-matched counts (inputs, outputs, errors, alerts) are intercepted on worker task threads and accumulated in memory inside isolated ThreadLocal maps.
3. Task Hooks: We register Spark TaskFailureListener and TaskCompletionListener hooks on task startup:
   • On Success: The listener flushes the buffered metrics exactly once with the successful attempt's tags.
   • On Failure/Interruption: The listener catches task exceptions/kills and discards the buffered metrics from memory. No database entries are written for failed attempts.
4. Resource Safety:
   • Checks for TaskContext.get() == null first to prevent thread-local memory leaks on driver and helper threads.
   • Uses a simple stack-allocated loop instead of stream mappings to avoid GC allocations on hot paths.
   • Caches the Metrics reference (getTaskMetrics()) outside the task flush loop to avoid redundant tag maps and child wrapper instantiations.

---

3. Major Changes

cdap-spark-core-base

• SparkRuntimeContext.java:
  • Declared ThreadLocal Map and Boolean registers to isolate memory buffers across concurrent worker threads.
  • Updated metric counting hooks to intercept and buffer target suffix-matched counts.
  • Registered Spark completion and failure listeners to flush .raw metrics conditionally and clean up ThreadLocal allocations at thread reuse.

---

4. Verification Example

For a partition processing 1,000 records that fails once (processing 5 records) and succeeds on the retry:

Metrics Query payload (POST /v3/metrics/query)

{
  "qid": {
    "tags": {
      "namespace": "default",
      "app": "fail_retry_records",
      "run": "d9faaa02-75e5-11f1-b343-d65e96f69bd1"
    },
    "metrics": [
      "user.JavaScript.records.out",
      "user.JavaScript.records.out.raw"
    ],
    "aggregate": true,
    "timeRange": {"startTime": 0, "endTime": "now"}
  }
}

Metrics Query Response

The legacy metric double-counts the failure, while the new .raw metric is successfully deduplicated:

{
  "qid": {
    "series": [
      {
        "metricName": "user.JavaScript.records.out",
        "data": [{ "time": 0, "value": 1005 }]
      },
      {
        "metricName": "user.JavaScript.records.out.raw",
        "data": [{ "time": 0, "value": 1000 }]
      }
    ]
  }
}

Diagnostic Breakdown payload (by partition & attempt)

{
  "qid": {
    "tags": {
      "namespace": "default",
      "app": "fail_retry_records",
      "run": "d9faaa02-75e5-11f1-b343-d65e96f69bd1"
    },
    "metrics": [
      "user.JavaScript.records.out.raw"
    ],
    "groupBy": [
      "spark_part",
      "spark_att"
    ],
    "timeRange": {"startTime": 0, "endTime": "now"}
  }
}

Diagnostic Response

Only successful attempt metrics exist in the series list:

{
  "qid": {
    "series": [
      {
        "metricName": "user.JavaScript.records.out.raw",
        "grouping": {
          "spark_part": "0",
          "spark_att": "1"
        },
        "data": [{ "time": 0, "value": 1000 }]
      }
    ]
  }
}

[gemini-code-assist]

Code Review

This pull request introduces buffering for pipeline stage metrics in Spark tasks, flushing raw metrics upon successful task completion and adding Spark partition and attempt tags to metric aggregations. The reviewer identified a potential memory leak in SparkRuntimeContext when TaskContext.get() is null (as the thread-local map is populated but never cleared), as well as performance overheads from stream allocations in bufferMetric and redundant getTaskMetrics() calls inside the loop in flushBufferedMetrics.

[sonarqubecloud]

Quality Gate passed

Issues
3 New issues
0 Accepted issues

Measures
0 Security Hotspots
88.6% Coverage on New Code
0.0% Duplication on New Code

See analysis details on SonarQube Cloud

[vsethi09]

Why is ThreadLocal needed here?

What is the scope of SparkRuntimeContext object?

[sahusanket]

• 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.

[vsethi09]

Return early if tc == null

TaskContext tc = TaskContext.get();
if (tc == null) {
    return getMetrics();
}
...
...

[vsethi09]

return early to avoid long nested blocks.

if (listenerRegistered.get()) {
    return;
}

...
...

[vsethi09]

Please add unit tests wherever applicable.

[sonarqubecloud]

Quality Gate passed

Issues
47 New issues
0 Accepted issues

Measures
0 Security Hotspots
92.5% Coverage on New Code
0.0% Duplication on New Code

See analysis details on SonarQube Cloud