From e09d151fff0a15b6e3804883b5d8cd7733a4592d Mon Sep 17 00:00:00 2001
From: Andrii Rosa <andriyrosa@gmail.com>
Date: Wed, 17 Aug 2022 12:30:50 -0400
Subject: [PATCH] Ensure memory is freed in SpoolingExchangeOutputBuffer

Close memory context to prevent race condition when memory is reserved
in enqueuePage after output buffer destroy following the pattern used in
other output buffers
---
 .../buffer/SpoolingExchangeOutputBuffer.java         | 12 ++++++++++--
 1 file changed, 10 insertions(+), 2 deletions(-)

diff --git a/core/trino-main/src/main/java/io/trino/execution/buffer/SpoolingExchangeOutputBuffer.java b/core/trino-main/src/main/java/io/trino/execution/buffer/SpoolingExchangeOutputBuffer.java
index ca1a4d23d2ea..2d5f0bb7146e 100644
--- a/core/trino-main/src/main/java/io/trino/execution/buffer/SpoolingExchangeOutputBuffer.java
+++ b/core/trino-main/src/main/java/io/trino/execution/buffer/SpoolingExchangeOutputBuffer.java
@@ -203,7 +203,7 @@ public void setNoMorePages()
                 stateMachine.finish();
             }
             exchangeSink = null;
-            updateMemoryUsage(0);
+            forceFreeMemory();
         });
     }
 
@@ -234,7 +234,7 @@ public void abort()
                 log.warn(failure, "Error aborting exchange sink");
             }
             exchangeSink = null;
-            updateMemoryUsage(0);
+            forceFreeMemory();
         });
     }
 
@@ -272,6 +272,14 @@ private void updatePeakMemoryUsage(long bytes)
         }
     }
 
+    private void forceFreeMemory()
+    {
+        LocalMemoryContext context = getSystemMemoryContextOrNull();
+        if (context != null) {
+            context.close();
+        }
+    }
+
     private LocalMemoryContext getSystemMemoryContextOrNull()
     {
         try {