Skip to content

fix: Race between run and cancellation in FrameProcessorExecutor. #19006

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

Open
gianm wants to merge 2 commits into
base: master
Choose a base branch
from
Open

fix: Race between run and cancellation in FrameProcessorExecutor. #19006

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
9623956
fix: Race between run and cancellation in FrameProcessorExecutor.
gianm Apr 2, 2026
8b7397e
Fix another race.
gianm Apr 2, 2026
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
18 changes: 12 additions & 6 deletions ...query/src/test/java/org/apache/druid/msq/shuffle/output/ChannelStageOutputReaderTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,7 @@
import com.google.common.math.IntMath;
import com.google.common.util.concurrent.ListenableFuture;
import org.apache.druid.common.guava.FutureUtils;
import org.apache.druid.error.DruidException;
import org.apache.druid.frame.Frame;
import org.apache.druid.frame.FrameType;
import org.apache.druid.frame.channel.BlockingQueueFrameChannel;
Expand Down Expand Up @@ -169,8 +170,8 @@ public void test_remote_threeFrames_readOneByteAtATime() throws Exception
channel.writable().write(frame);

// See that we can't write another frame.
final IllegalStateException e = Assert.assertThrows(
IllegalStateException.class,
final DruidException e = Assert.assertThrows(
DruidException.class,
() -> channel.writable().write(frame)
);

Expand All @@ -195,8 +196,8 @@ public void test_remote_threeFrames_readOneByteAtATime() throws Exception
channel.writable().write(frame);

// See that we can't write a fourth frame.
final IllegalStateException e2 = Assert.assertThrows(
IllegalStateException.class,
final DruidException e2 = Assert.assertThrows(
DruidException.class,
() -> channel.writable().write(frame)
);

Expand Down Expand Up @@ -280,7 +281,7 @@ public static class WithMaximalBuffering extends InitializedNullHandlingTest
private static final int EXPECTED_NUM_ROWS = 1209;

private final BlockingQueueFrameChannel channel = new BlockingQueueFrameChannel(MAX_FRAMES);
private final ChannelStageOutputReader reader = new ChannelStageOutputReader(channel.readable(), FrameTestUtil.WT_CONTEXT_LEGACY);
private ChannelStageOutputReader reader = new ChannelStageOutputReader(channel.readable(), FrameTestUtil.WT_CONTEXT_LEGACY);

@Rule
public TemporaryFolder temporaryFolder = new TemporaryFolder();
Expand Down Expand Up @@ -310,7 +311,9 @@ public void setUp()
@After
public void tearDown()
{
reader.close();
if (reader != null) {
reader.close();
}
}

@Test
Expand All @@ -337,6 +340,9 @@ public void test_readLocally_closePriorToRead() throws IOException
IllegalStateException.class,
reader::readLocally
);

// Prevent close() in tearDown, which would fail.
reader = null;
}

@Test
Expand Down
87 changes: 62 additions & 25 deletions processing/src/main/java/org/apache/druid/frame/channel/BlockingQueueFrameChannel.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,9 +24,9 @@
import com.google.common.util.concurrent.SettableFuture;
import com.google.errorprone.annotations.concurrent.GuardedBy;
import it.unimi.dsi.fastutil.objects.ObjectIntPair;
import org.apache.druid.error.DruidException;
import org.apache.druid.java.util.common.Either;
import org.apache.druid.java.util.common.IAE;
import org.apache.druid.java.util.common.ISE;
import org.apache.druid.query.rowsandcols.RowsAndColumns;

import javax.annotation.Nullable;
Expand All @@ -52,14 +52,31 @@ public class BlockingQueueFrameChannel
private final Writable writable;
private final Readable readable;

/**
* Queue of items from the writer. Ends with {@link #END_MARKER} if the writable channel is closed.
* Only ever updated by the writer.
*/
@GuardedBy("lock")
private final ArrayDeque<Optional<Either<Throwable, ObjectIntPair<RowsAndColumns>>>> queue;

/**
* Whether {@link Readable#close()} has been called.
*/
@GuardedBy("lock")
private SettableFuture<?> readyForWritingFuture = null;
private boolean readerClosed;

/**
* Future that is set to null by {@link #notifyWriter()} when the reader has read from {@link #queue}
* or been closed.
*/
@GuardedBy("lock")
private SettableFuture<?> readyForReadingFuture = null;
private SettableFuture<?> readyForWritingFuture;

/**
* Future that is set to null by {@link #notifyReader()} when the writer has written to {@link #queue}.
*/
@GuardedBy("lock")
private SettableFuture<?> readyForReadingFuture;

/**
* Create a channel with a particular buffer size (expressed in number of frames).
Expand Down Expand Up @@ -100,13 +117,6 @@ public static BlockingQueueFrameChannel minimal()
return new BlockingQueueFrameChannel(1);
}

private boolean isFinished()
{
synchronized (lock) {
return END_MARKER.equals(queue.peek());
}
}

@GuardedBy("lock")
private void notifyWriter()
{
Expand All @@ -133,16 +143,14 @@ private class Writable implements WritableFrameChannel
public void write(RowsAndColumns rac, int partitionNumber)
{
synchronized (lock) {
if (isFinished()) {
throw new ISE("Channel cannot accept new frames");
if (isClosed()) {
throw DruidException.defensive("Channel cannot accept new frames");
} else if (queue.size() >= maxQueuedFrames) {
// Caller should have checked if this channel was ready for writing.
throw new ISE("Channel has no capacity");
} else {
if (!queue.offer(Optional.of(Either.value(ObjectIntPair.of(rac, partitionNumber))))) {
// If this happens, it's a bug in this class's capacity-counting.
throw new ISE("Channel had capacity, but could not add frame");
}
throw DruidException.defensive("Channel has no capacity");
} else if (!queue.offer(Optional.of(Either.value(ObjectIntPair.of(rac, partitionNumber))))) {
// If this happens, it's a bug in this class's capacity-counting.
throw DruidException.defensive("Channel had capacity, but could not add frame");
}

notifyReader();
Expand All @@ -153,7 +161,9 @@ public void write(RowsAndColumns rac, int partitionNumber)
public ListenableFuture<?> writabilityFuture()
{
synchronized (lock) {
if (queue.size() < maxQueuedFrames) {
if (isClosed()) {
throw DruidException.defensive("Closed, cannot call writabilityFuture()");
} else if (queue.size() < maxQueuedFrames) {
return Futures.immediateFuture(null);
} else if (readyForWritingFuture != null) {
return readyForWritingFuture;
Expand All @@ -167,11 +177,15 @@ public ListenableFuture<?> writabilityFuture()
public void fail(@Nullable Throwable cause)
{
synchronized (lock) {
if (isClosed()) {
throw DruidException.defensive("Closed, cannot call fail()");
}

queue.clear();

if (!queue.offer(Optional.of(Either.error(cause != null ? cause : new ISE("Failed"))))) {
if (!queue.offer(Optional.of(Either.error(cause != null ? cause : new RuntimeException("Failed"))))) {
// If this happens, it's a bug, potentially due to incorrectly using this class with multiple writers.
throw new ISE("Could not write error to channel");
throw DruidException.defensive("Could not write error to channel");
}

notifyReader();
Expand All @@ -183,12 +197,12 @@ public void close()
{
synchronized (lock) {
if (isClosed()) {
throw new ISE("Already closed");
throw DruidException.defensive("Closed, cannot call close() again");
}

if (!queue.offer(END_MARKER)) {
// If this happens, it's a bug, potentially due to incorrectly using this class with multiple writers.
throw new ISE("Channel had capacity, but could not add end marker");
throw DruidException.defensive("Channel had capacity, but could not add end marker");
}

notifyReader();
Expand All @@ -210,13 +224,23 @@ private class Readable implements ReadableFrameChannel
@Override
public boolean isFinished()
{
return BlockingQueueFrameChannel.this.isFinished();
synchronized (lock) {
if (readerClosed) {
throw DruidException.defensive("Closed, cannot call isFinished()");
}

return END_MARKER.equals(queue.peek());
}
}

@Override
public boolean canRead()
{
synchronized (lock) {
if (readerClosed) {
throw DruidException.defensive("Closed, cannot call canRead()");
}

return !queue.isEmpty() && !isFinished();
}
}
Expand All @@ -227,6 +251,10 @@ public RowsAndColumns read()
final Optional<Either<Throwable, ObjectIntPair<RowsAndColumns>>> next;

synchronized (lock) {
if (readerClosed) {
throw DruidException.defensive("Closed, cannot call read()");
}

if (isFinished()) {
throw new NoSuchElementException();
}
Expand All @@ -247,6 +275,10 @@ public RowsAndColumns read()
public ListenableFuture<?> readabilityFuture()
{
synchronized (lock) {
if (readerClosed) {
throw DruidException.defensive("Closed, cannot call readabilityFuture()");
}

if (!queue.isEmpty()) {
return Futures.immediateFuture(null);
} else if (readyForReadingFuture != null) {
Expand All @@ -261,7 +293,12 @@ public ListenableFuture<?> readabilityFuture()
public void close()
{
synchronized (lock) {
queue.clear();
if (readerClosed) {
// close() should not be called twice.
throw DruidException.defensive("Closed, cannot call close() again");
}

readerClosed = true;
Copy link
Copy Markdown
Member

@FrankChen021 FrankChen021 Apr 25, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Readable close can leave writers permanently blocked

Readable.close() now only sets readerClosed and notifies the current writer future, but it leaves the queue contents in place and Writable.writabilityFuture() does not consider readerClosed. If the queue is full when the reader is closed, the notified writer can re-check writability, see the queue still full, create a new readyForWritingFuture, and wait forever because no reader will drain the queue or notify again. This can hang upstream processors or result producers when a consumer cancels or closes early. Consider clearing the queue as before, or making readerClosed close/fail the writable side so future waiters resume deterministically.

Copy link
Copy Markdown
Contributor Author

@gianm gianm Apr 26, 2026 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It's true although the old code is the same way. The old code cleared the queue here, but typically the queue is of length just 1 or 2, and it's easy for it to fill back up again. I believe isn't a problem in practice, because in production early-close scenarios (mainly LIMIT) processors are directly canceled by Worker#postCleanupStage. Generally we aren't relying on the channels themselves to propagate such information.

notifyWriter();
}
}
Expand Down
15 changes: 7 additions & 8 deletions processing/src/main/java/org/apache/druid/frame/channel/WritableFrameChannel.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -92,14 +92,19 @@ default void write(Frame frame) throws IOException
*/
void fail(@Nullable Throwable cause) throws IOException;

/**
* Returns a future that resolves when {@link #write} is able to receive a new batch of data without blocking or
* throwing an exception. The future never resolves to an exception.
*/
ListenableFuture<?> writabilityFuture();

/**
* Finish writing to this channel.
*
* When this method is called without {@link #fail(Throwable)} having previously been called, the writer is
* understood to have completed successfully.
*
* After calling this method, no additional calls to {@link #write}, {@link #fail(Throwable)}, or this method
* are permitted.
* After calling this method, no additional calls to are permitted to any methods other than {@link #isClosed()}.
*/
@Override
void close() throws IOException;
Expand All @@ -108,10 +113,4 @@ default void write(Frame frame) throws IOException
* Whether {@link #close()} has been called on this channel.
*/
boolean isClosed();

/**
* Returns a future that resolves when {@link #write} is able to receive a new batch of data without blocking or
* throwing an exception. The future never resolves to an exception.
*/
ListenableFuture<?> writabilityFuture();
}
Loading
Loading