From 32b7e985b50c8abae38d059697bcad3173ab4778 Mon Sep 17 00:00:00 2001 From: Gian Merlino Date: Tue, 24 Sep 2024 20:49:01 -0700 Subject: [PATCH 1/4] 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 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 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(); From 69fa29f885a9d1b18cdeb9e4816207804debe6e2 Mon Sep 17 00:00:00 2001 From: Gian Merlino Date: Tue, 24 Sep 2024 20:55:02 -0700 Subject: [PATCH 2/4] Fix other call sites. --- .../msq/querykit/scan/ScanQueryFrameProcessor.java | 12 ++++++------ 1 file changed, 6 insertions(+), 6 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 6ef6c569d035..c0211e22521e 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 @@ -223,7 +223,7 @@ protected ReturnOrAwait runWithDataServerQuery(final DataSer cursorYielder.close(); return ReturnOrAwait.returnObject(handedOffSegments); } else { - final long rowsFlushed = setNextCursor(cursorYielder.get(), null); + final long rowsFlushed = setNextCursor(null, cursorYielder.get(), null); closer.register(cursorYielder); if (rowsFlushed > 0) { return ReturnOrAwait.runAgain(); @@ -258,16 +258,16 @@ protected ReturnOrAwait runWithSegment(final SegmentWithDescriptor segment ); } - 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 cursors! + nextCursorHolder.close(); return ReturnOrAwait.returnObject(Unit.instance()); } else { - final long rowsFlushed = setNextCursor(nextCursor, segmentHolder.get().getSegment()); + final long rowsFlushed = setNextCursor(nextCursorHolder, nextCursor, segmentHolder.get().getSegment()); assert rowsFlushed == 0; // There's only ever one cursor when running with a segment } } @@ -418,7 +418,7 @@ private long flushFrameWriter() throws IOException } private long setNextCursor( - final CursorHolder cursorHolder, + @Nullable final CursorHolder cursorHolder, final Cursor cursor, final Segment segment ) throws IOException From dbe19a8daff04796d3ee8ea41a5cbc888786ab66 Mon Sep 17 00:00:00 2001 From: Gian Merlino Date: Tue, 24 Sep 2024 22:05:11 -0700 Subject: [PATCH 3/4] Fix reference. --- .../apache/druid/msq/querykit/scan/ScanQueryFrameProcessor.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 c0211e22521e..013cab795b13 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 @@ -260,7 +260,7 @@ protected ReturnOrAwait runWithSegment(final SegmentWithDescriptor segment final CursorHolder nextCursorHolder = cursorFactory.makeCursorHolder(ScanQueryEngine.makeCursorBuildSpec(query, null)); - final Cursor nextCursor = cursorHolder.asCursor(); + final Cursor nextCursor = nextCursorHolder.asCursor(); if (nextCursor == null) { // No cursors! From cfcbb3ecebe4b86945d32b112ee9c55774da7c22 Mon Sep 17 00:00:00 2001 From: Gian Merlino Date: Wed, 25 Sep 2024 08:10:11 -0700 Subject: [PATCH 4/4] Improvements. --- .../scan/ScanQueryFrameProcessor.java | 20 +++++++++---------- 1 file changed, 10 insertions(+), 10 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 013cab795b13..dbcd271cccfb 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,8 +101,8 @@ public class ScanQueryFrameProcessor extends BaseLeafFrameProcessor private final VirtualColumns frameWriterVirtualColumns; private final Closer closer = Closer.create(); - private CursorHolder cursorHolder; private Cursor cursor; + private Closeable cursorCloser; private Segment segment; private final SimpleSettableOffset cursorOffset = new SimpleAscendingOffset(Integer.MAX_VALUE); private FrameWriter frameWriter; @@ -157,7 +157,7 @@ public ReturnOrAwait runIncrementally(final IntSet readableInputs) throw @Override public void cleanup() throws IOException { - closer.register(cursorHolder); + closer.register(cursorCloser); closer.register(frameWriter); closer.register(super::cleanup); closer.close(); @@ -223,7 +223,7 @@ protected ReturnOrAwait runWithDataServerQuery(final DataSer cursorYielder.close(); return ReturnOrAwait.returnObject(handedOffSegments); } else { - final long rowsFlushed = setNextCursor(null, cursorYielder.get(), null); + final long rowsFlushed = setNextCursor(cursorYielder.get(), null, null); closer.register(cursorYielder); if (rowsFlushed > 0) { return ReturnOrAwait.runAgain(); @@ -267,7 +267,7 @@ protected ReturnOrAwait runWithSegment(final SegmentWithDescriptor segment nextCursorHolder.close(); return ReturnOrAwait.returnObject(Unit.instance()); } else { - final long rowsFlushed = setNextCursor(nextCursorHolder, nextCursor, segmentHolder.get().getSegment()); + final long rowsFlushed = setNextCursor(nextCursor, nextCursorHolder, segmentHolder.get().getSegment()); assert rowsFlushed == 0; // There's only ever one cursor when running with a segment } } @@ -306,14 +306,14 @@ protected ReturnOrAwait runWithInputChannel( final CursorHolder nextCursorHolder = cursorFactory.makeCursorHolder(ScanQueryEngine.makeCursorBuildSpec(query, null)); - final Cursor nextCursor = cursorHolder.asCursor(); + final Cursor nextCursor = nextCursorHolder.asCursor(); if (nextCursor == null) { // no cursor nextCursorHolder.close(); return ReturnOrAwait.returnObject(Unit.instance()); } - final long rowsFlushed = setNextCursor(nextCursorHolder, nextCursor, frameSegment); + final long rowsFlushed = setNextCursor(nextCursor, nextCursorHolder, frameSegment); if (rowsFlushed > 0) { return ReturnOrAwait.runAgain(); @@ -418,19 +418,19 @@ private long flushFrameWriter() throws IOException } private long setNextCursor( - @Nullable final CursorHolder cursorHolder, final Cursor cursor, + @Nullable final Closeable cursorCloser, final Segment segment ) throws IOException { final long rowsFlushed = flushFrameWriter(); - if (this.cursorHolder != null) { + if (this.cursorCloser != 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.cursorCloser.close(); } - this.cursorHolder = cursorHolder; this.cursor = cursor; + this.cursorCloser = cursorCloser; this.segment = segment; this.cursorOffset.reset(); return rowsFlushed;