From 4884764f5c024e3d756b40681a57f80e362acae7 Mon Sep 17 00:00:00 2001
From: Hongze Zhang <hongze.zhang@intel.com>
Date: Wed, 21 Aug 2024 15:10:09 +0800
Subject: [PATCH] [GLUTEN-6930][VL] Print memory statistics during task ending
 when leak is found

---
 .../org/apache/gluten/runtime/Runtime.scala   | 22 ++++++++++---------
 1 file changed, 12 insertions(+), 10 deletions(-)

diff --git a/gluten-data/src/main/scala/org/apache/gluten/runtime/Runtime.scala b/gluten-data/src/main/scala/org/apache/gluten/runtime/Runtime.scala
index 8edceea447fc..7ba72b379ba3 100644
--- a/gluten-data/src/main/scala/org/apache/gluten/runtime/Runtime.scala
+++ b/gluten-data/src/main/scala/org/apache/gluten/runtime/Runtime.scala
@@ -97,16 +97,17 @@ object Runtime {
         throw new GlutenException(
           s"Runtime instance already released: $handle, ${resourceName()}, ${priority()}")
       }
+
+      def dump(): KnownNameAndStats = {
+        new KnownNameAndStats() {
+          override def name: String = resourceName()
+          override def stats: MemoryUsageStats = collectMemoryUsage()
+        }
+      }
+
       if (LOGGER.isDebugEnabled) {
         LOGGER.debug(
-          SparkMemoryUtil.prettyPrintStats(
-            "About to release memory manager, usage dump:",
-            new KnownNameAndStats() {
-              override def name: String = resourceName()
-
-              override def stats: MemoryUsageStats = collectMemoryUsage()
-            }
-          ))
+          SparkMemoryUtil.prettyPrintStats("About to release memory manager, usage dump:", dump()))
       }
 
       RuntimeJniWrapper.releaseRuntime(handle)
@@ -115,10 +116,11 @@ object Runtime {
         LOGGER.warn(
           String.format(
             "%s Reservation listener %s still reserved non-zero bytes, which may cause memory" +
-              " leak, size: %s. ",
+              " leak, size: %s, dump: %s ",
             name,
             rl.toString,
-            SparkMemoryUtil.bytesToString(rl.getUsedBytes)
+            SparkMemoryUtil.bytesToString(rl.getUsedBytes),
+            dump()
           ))
       }
     }