From 32b7e985b50c8abae38d059697bcad3173ab4778 Mon Sep 17 00:00:00 2001
From: Gian Merlino <gianmerlino@gmail.com>
Date: Tue, 24 Sep 2024 20:49:01 -0700
Subject: [PATCH] ScanQueryFrameProcessor: Close CursorHolders as we go along.

The change in #16533 added CursorHolders to the processor-level Closer.
This is problematic when running on an input channel: it means we started
keeping around all CursorHolders for all frames we process and closing them
when the channel is complete, rather than closing them as we go along.

This patch restores the prior behavior, where resources are closed as we go.
---
 .../scan/ScanQueryFrameProcessor.java         | 22 ++++++++++++++-----
 1 file changed, 17 insertions(+), 5 deletions(-)

diff --git a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/scan/ScanQueryFrameProcessor.java b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/scan/ScanQueryFrameProcessor.java
index e5fa0a03d621..6ef6c569d035 100644
--- a/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/scan/ScanQueryFrameProcessor.java
+++ b/extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/scan/ScanQueryFrameProcessor.java
@@ -101,6 +101,7 @@ public class ScanQueryFrameProcessor extends BaseLeafFrameProcessor
   private final VirtualColumns frameWriterVirtualColumns;
   private final Closer closer = Closer.create();
 
+  private CursorHolder cursorHolder;
   private Cursor cursor;
   private Segment segment;
   private final SimpleSettableOffset cursorOffset = new SimpleAscendingOffset(Integer.MAX_VALUE);
@@ -156,6 +157,7 @@ public ReturnOrAwait<Object> runIncrementally(final IntSet readableInputs) throw
   @Override
   public void cleanup() throws IOException
   {
+    closer.register(cursorHolder);
     closer.register(frameWriter);
     closer.register(super::cleanup);
     closer.close();
@@ -302,16 +304,16 @@ protected ReturnOrAwait<Unit> runWithInputChannel(
           );
         }
 
-        final CursorHolder cursorHolder = closer.register(
-            cursorFactory.makeCursorHolder(ScanQueryEngine.makeCursorBuildSpec(query, null))
-        );
+        final CursorHolder nextCursorHolder =
+            cursorFactory.makeCursorHolder(ScanQueryEngine.makeCursorBuildSpec(query, null));
         final Cursor nextCursor = cursorHolder.asCursor();
 
         if (nextCursor == null) {
           // no cursor
+          nextCursorHolder.close();
           return ReturnOrAwait.returnObject(Unit.instance());
         }
-        final long rowsFlushed = setNextCursor(nextCursor, frameSegment);
+        final long rowsFlushed = setNextCursor(nextCursorHolder, nextCursor, frameSegment);
 
         if (rowsFlushed > 0) {
           return ReturnOrAwait.runAgain();
@@ -415,9 +417,19 @@ private long flushFrameWriter() throws IOException
     }
   }
 
-  private long setNextCursor(final Cursor cursor, final Segment segment) throws IOException
+  private long setNextCursor(
+      final CursorHolder cursorHolder,
+      final Cursor cursor,
+      final Segment segment
+  ) throws IOException
   {
     final long rowsFlushed = flushFrameWriter();
+    if (this.cursorHolder != null) {
+      // Close here, don't add to the processor-level Closer, to avoid leaking CursorHolders. We may generate many
+      // CursorHolders per instance of this processor, and we need to close them as we go, not all at the end.
+      this.cursorHolder.close();
+    }
+    this.cursorHolder = cursorHolder;
     this.cursor = cursor;
     this.segment = segment;
     this.cursorOffset.reset();