Make RETRY_COVERAGE_TRACKING also able to detect nested withRetry

[thirtiseven]

Contributes to #14136

Description

This PR adds a new type, NESTED, to SPARK_RAPIDS_RETRY_COVERAGE_TRACKING, enabling the detection of nested withRetry calls, which can indicate potential performance issues.

Checklists

• This PR has added documentation for new or modified features or behaviors.
• This PR has added new tests or modified existing tests to cover new code paths.
  (Please explain in the PR description how the new code paths are tested, such as names of the new/existing tests that cover them.)
• Performance testing has been performed and its results are added in the PR description. Or, an issue has been filed with a link in the PR description.

[thirtiseven]

@greptileai full review

[greptile-apps]

Greptile Overview

Greptile Summary

This PR extends the SPARK_RAPIDS_RETRY_COVERAGE_TRACKING debug feature to detect nested withRetry calls. When allocations occur within nested retry blocks (depth > 1), they are now logged as "NESTED" instead of being ignored as properly covered allocations.

Key Changes:

• Added getRetryBlockDepth() method to RetryStateTracker to expose the current retry block nesting depth
• Modified AllocationRetryCoverageTracker.checkAllocationInternal() to detect and log nested retry scenarios
• Nested retry blocks can cause cascading retry attempts and reduced performance, making them important to track

Implementation:

• Uses existing localRetryBlockDepth ThreadLocal tracking infrastructure
• Logs allocations with depth = 0 as uncovered (HOST/DEVICE)
• Logs allocations with depth > 1 as NESTED
• Returns early for depth = 1 (properly covered, no logging needed)

The implementation correctly leverages the existing retry block depth tracking mechanism that was already in place for the coverage tracker.

Confidence Score: 5/5

• This PR is safe to merge with minimal risk
• The implementation is straightforward and leverages existing infrastructure. The new getRetryBlockDepth() method simply exposes already-tracked state, and the modified logging logic in AllocationRetryCoverageTracker correctly handles the three cases (uncovered, nested, and properly covered). This is a debug-only feature gated behind an environment variable, so it won't affect production behavior when disabled.
• No files require special attention

Important Files Changed

Filename	Overview
sql-plugin/src/main/scala/com/nvidia/spark/rapids/AllocationRetryCoverageTracker.scala	Added nested retry block detection: logs allocations with depth > 1 as NESTED
sql-plugin/src/main/scala/com/nvidia/spark/rapids/RmmRapidsRetryIterator.scala	Added getRetryBlockDepth method to expose retry block nesting depth

Sequence Diagram

sequenceDiagram
    participant User as User Code
    participant Outer as Outer withRetry
    participant Tracker as RetryStateTracker
    participant Alloc as AllocationRetryCoverageTracker
    participant Inner as Inner withRetry
    
    User->>Outer: withRetry(data) { fn }
    Outer->>Tracker: enterRetryBlock()
    Note over Tracker: depth = 1
    Outer->>User: execute fn(data)
    User->>Inner: withRetry(data2) { fn2 }
    Inner->>Tracker: enterRetryBlock()
    Note over Tracker: depth = 2 (NESTED!)
    Inner->>User: execute fn2(data2)
    User->>Alloc: checkAllocation(kind)
    Alloc->>Tracker: getRetryBlockDepth()
    Tracker-->>Alloc: depth = 2
    Note over Alloc: depth > 1, log as NESTED
    Alloc->>Alloc: writeToFile("NESTED,...")
    User-->>Inner: return result
    Inner->>Tracker: exitRetryBlock()
    Note over Tracker: depth = 1
    Inner-->>User: return result
    User-->>Outer: return result
    Outer->>Tracker: exitRetryBlock()
    Note over Tracker: depth = 0
    Outer-->>User: return result

Loading

[greptile-apps]

_{2 files reviewed, 2 comments}

_{Edit Code Review Agent Settings | Greptile}

[nvauto]

NOTE: release/26.02 has been created from main. Please retarget your PR to release/26.02 if it should be included in the release.

[thirtiseven]

@greptileai full review

[greptile-apps]

_{2 files reviewed, no comments}

_{Edit Code Review Agent Settings | Greptile}

[greptile-apps]

_{2 files reviewed, no comments}

_{Edit Code Review Agent Settings | Greptile}