From e920aa4f956478c7cdcb939a1689779daef04b09 Mon Sep 17 00:00:00 2001
From: Chao Sun <sunchao@apache.org>
Date: Tue, 12 Mar 2024 22:50:54 -0700
Subject: [PATCH] fix: Include active spiller when computing peak shuffle
 memory (#196)

---
 .../apache/spark/shuffle/sort/CometShuffleExternalSorter.java | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/spark/src/main/java/org/apache/spark/shuffle/sort/CometShuffleExternalSorter.java b/spark/src/main/java/org/apache/spark/shuffle/sort/CometShuffleExternalSorter.java
index aa806e2e8..9fe88ecbd 100644
--- a/spark/src/main/java/org/apache/spark/shuffle/sort/CometShuffleExternalSorter.java
+++ b/spark/src/main/java/org/apache/spark/shuffle/sort/CometShuffleExternalSorter.java
@@ -257,6 +257,9 @@ private long getMemoryUsage() {
     for (SpillSorter sorter : spillingSorters) {
       totalPageSize += sorter.getMemoryUsage();
     }
+    if (activeSpillSorter != null) {
+      totalPageSize += activeSpillSorter.getMemoryUsage();
+    }
     return totalPageSize;
   }
 
@@ -274,6 +277,7 @@ public long getPeakMemoryUsedBytes() {
   }
 
   private long freeMemory() {
+    updatePeakMemoryUsed();
     long memoryFreed = 0;
     if (isAsync) {
       for (SpillSorter sorter : spillingSorters) {