Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

ScanQueryFrameProcessor: Close CursorHolders as we go along. #17152

Merged
merged 4 commits into from
Sep 25, 2024
Merged

ScanQueryFrameProcessor: Close CursorHolders as we go along. #17152

Changes from 2 commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
32b7e98
ScanQueryFrameProcessor: Close CursorHolders as we go along.
gianm Sep 25, 2024
69fa29f
Fix other call sites.
gianm Sep 25, 2024
dbe19a8
Fix reference.
gianm Sep 25, 2024
cfcbb3e
Improvements.
gianm Sep 25, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
32 changes: 22 additions & 10 deletions ...stage-query/src/main/java/org/apache/druid/msq/querykit/scan/ScanQueryFrameProcessor.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down Expand Up @@ -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();
Expand Down Expand Up @@ -221,7 +223,7 @@ protected ReturnOrAwait<SegmentsInputSlice> 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();
Expand Down Expand Up @@ -256,16 +258,16 @@ protected ReturnOrAwait<Unit> 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
}
}
Expand Down Expand Up @@ -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();
Expand Down Expand Up @@ -415,9 +417,19 @@ private long flushFrameWriter() throws IOException
}
}

private long setNextCursor(final Cursor cursor, final Segment segment) throws IOException
private long setNextCursor(
@Nullable 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();
Expand Down
Loading