From 34596c5d5cb8ac85598368b17d59d930197f5b1a Mon Sep 17 00:00:00 2001
From: Andy Grove <agrove@apache.org>
Date: Mon, 7 Oct 2024 14:17:05 -0600
Subject: [PATCH] fix

---
 .../main/java/org/apache/spark/CometTaskMemoryManager.java   | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/spark/src/main/java/org/apache/spark/CometTaskMemoryManager.java b/spark/src/main/java/org/apache/spark/CometTaskMemoryManager.java
index e7b4a7060..085611bde 100644
--- a/spark/src/main/java/org/apache/spark/CometTaskMemoryManager.java
+++ b/spark/src/main/java/org/apache/spark/CometTaskMemoryManager.java
@@ -63,9 +63,12 @@ public long acquireMemory(long size) {
         if (size <= available) {
           available -= size;
           return size;
+        } else {
+          long allocated = available;
+          available = 0;
+          return allocated;
         }
       }
-      return 0;
     }
   }