zhztheplayer commented on code in PR #5447:
URL: https://github.com/apache/incubator-gluten/pull/5447#discussion_r1570100475
##########
gluten-core/src/main/scala/org/apache/spark/memory/SparkMemoryUtil.scala:
##########
@@ -209,9 +209,50 @@ object SparkMemoryUtil {
sb.toString()
}
+ def addLeakSafeTaskCompletionListener[U](f: TaskContext => U): TaskContext =
{
+ TaskContext.get().addTaskCompletionListener(f)
+ }
+
private case class PrintableMemoryUsageStats(
name: String,
used: Option[Long],
peak: Option[Long],
children: Iterable[PrintableMemoryUsageStats])
+
+ class GenericRetainer[T <: AutoCloseable] {
+ private var retained: Option[T] = None
+
+ def retain(batch: T): Unit = {
+ if (retained.isDefined) {
+ throw new IllegalStateException
+ }
+ retained = Some(batch)
+ }
+
+ def release(): Unit = {
+ retained.foreach(b => b.close())
+ retained = None
+ }
+ }
+
+ class UnsafeItr[T <: AutoCloseable](delegate: Iterator[T]) extends
Iterator[T] {
+ val holder = new GenericRetainer[T]()
+
+ addLeakSafeTaskCompletionListener[Unit](
+ (_: TaskContext) => {
+ holder.release()
+ })
+
+ override def hasNext: Boolean = {
+ holder.release()
+ val hasNext = delegate.hasNext
+ hasNext
+ }
+
+ override def next(): T = {
+ val b = delegate.next()
+ holder.retain(b)
+ b
+ }
+ }
Review Comment:
Had some new APIs to replace this sort of wrappers
https://github.com/apache/incubator-gluten/blob/main/gluten-core/src/main/scala/org/apache/gluten/utils/Iterators.scala
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
