[BUGFIX] Fix resource leak when shuffle read

[izchen]

What changes were proposed in this pull request?

Use org.apache.spark.TaskContext#addTaskCompletionListener to clean up resources used by RssShuffleDataIterator. This avoids possible resource leaks.

Why are the changes needed?

Before this PR, RssShuffleDataIterator would only clean up used resources after all records read.

When the Spark Task fails or cancels, or runs some special logic such as LocalLimit, the resource will not be cleaned up. This creates potential resource leaks.

Does this PR introduce any user-facing change?

No

How was this patch tested?

Added a UT case

[jerqi]

If we use AQE, we could have many iterators, if we don't release the resource of them after we use them, we may occur OOM. I means that we use one iterator, and then release the iterator. We shouldn't release all the iterators at the end of task.

[izchen]

org.apache.spark.util.CompletionIterator will clean up as soon as its wrapped RssShuffleDataIterator.hasNext returns false :

// org.apache.spark.util.CompletionIterator
  def hasNext: Boolean = {
    val r = iter.hasNext  // iter => RssShuffleDataIterator
    if (!r && !completed) {
      completed = true
      // reassign to release resources of highly resource consuming iterators early
      iter = Iterator.empty.asInstanceOf[I]
      completion()
    }
    r
  }

  def completion(): Unit  // completion() => RssShuffleDataIterator.cleanup

After this PR, if there is no special case for the Spark Task, the timing of resource cleanup is still when the RssShuffleDataIterator ends, not when the Spark Task ends.

This is the same behavior as before the PR.

[jerqi]

org.apache.spark.util.CompletionIterator will clean up as soon as its wrapped RssShuffleDataIterator.hasNext returns false :

// org.apache.spark.util.CompletionIterator
  def hasNext: Boolean = {
    val r = iter.hasNext  // iter => RssShuffleDataIterator
    if (!r && !completed) {
      completed = true
      // reassign to release resources of highly resource consuming iterators early
      iter = Iterator.empty.asInstanceOf[I]
      completion()
    }
    r
  }

  def completion(): Unit  // completion() => RssShuffleDataIterator.cleanup

After this PR, if there is no special case for the Spark Task, the timing of resource cleanup is still when the RssShuffleDataIterator ends, not when the Spark Task ends.

This is the same behavior as before the PR.

OK, I got it. Good catch.

[codecov-commenter]

Codecov Report

Merging #174 (d492aea) into master (74cd40b) will decrease coverage by 1.37%.
The diff coverage is n/a.

@@             Coverage Diff              @@
##             master     #174      +/-   ##
============================================
- Coverage     58.29%   56.92%   -1.38%     
+ Complexity     1262     1183      -79     
============================================
  Files           158      149       -9     
  Lines          8397     7902     -495     
  Branches        779      749      -30     
============================================
- Hits           4895     4498     -397     
+ Misses         3251     3162      -89     
+ Partials        251      242       -9     

Impacted Files	Coverage Δ
...storage/handler/impl/DataSkippableReadHandler.java	81.25% <0.00%> (-3.13%)	⬇️
...org/apache/uniffle/server/ShuffleFlushManager.java	77.65% <0.00%> (-1.68%)	⬇️
...e/spark/shuffle/reader/RssShuffleDataIterator.java
...ava/org/apache/spark/shuffle/RssShuffleHandle.java
...org/apache/spark/shuffle/RssSparkShuffleUtils.java
...che/spark/shuffle/writer/BufferManagerOptions.java
.../java/org/apache/spark/shuffle/RssSparkConfig.java
...org/apache/spark/shuffle/writer/AddBlockEvent.java
...k/shuffle/writer/WrappedByteArrayOutputStream.java
...pache/spark/shuffle/writer/WriteBufferManager.java
... and 1 more

📣 We’re building smart automated test selection to slash your CI/CD build times. Learn more

[jerqi]

LGTM, thanks @izchen

[izchen]

Thanks for your review, @jerqi ！