Skip to content

core: fix off-heap ByteBuf leak on RejectedExecutionException #12891

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

Closed
kannanjgithub wants to merge 1 commit into from
Closed

core: fix off-heap ByteBuf leak on RejectedExecutionException #12891

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
9 changes: 8 additions & 1 deletion core/src/main/java/io/grpc/internal/ClientCallImpl.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -58,6 +58,7 @@
import java.util.Locale;
import java.util.concurrent.CancellationException;
import java.util.concurrent.Executor;
import java.util.concurrent.RejectedExecutionException;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.ScheduledFuture;
import java.util.concurrent.TimeUnit;
Expand Down Expand Up @@ -674,7 +675,13 @@ private void runInternal() {
}
}

callExecutor.execute(new MessagesAvailable());
try {
callExecutor.execute(new MessagesAvailable());
} catch (RejectedExecutionException e) {
GrpcUtil.closeQuietly(producer);
exceptionThrown(
Status.CANCELLED.withCause(e).withDescription("Failed to read message."));
Comment on lines +682 to +683

@AgraVator AgraVator Jul 6, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

Prior to this PR, exceptionThrown(Status status) was only ever called inside runInternal() of HeadersRead, MessagesAvailable, and StreamOnReady. All three executed inside callExecutor.execute(...), meaning exceptionThrown() ran strictly on the application callExecutor thread. In fact, its javadoc explicitly relies on this:

/**
 * Cancels call and schedules onClose() notification. May only be called from the application
 * thread.
 */
private void exceptionThrown(Status status) {
  exceptionStatus = status;
  stream.cancel(status);
}

When callExecutor.execute(...) throws RejectedExecutionException, messagesAvailable() catches it on the transport thread and directly invokes exceptionThrown(status) on the transport thread.

Because private Status exceptionStatus; is not volatile, writing exceptionStatus = status from the transport thread introduces a potential data race when if (exceptionStatus != null) is subsequently checked on the application thread (e.g., inside StreamClosed.runInternal()).

Consider making exceptionStatus volatile to guarantee visibility across threads:

- private Status exceptionStatus;
+ private volatile Status exceptionStatus;

And updating or clarifying the corresponding javadoc on exceptionThrown() (May only be called from the application thread.) to reflect that it can now be invoked from the transport thread upon task rejection.

}
}
}

Expand Down
19 changes: 12 additions & 7 deletions core/src/main/java/io/grpc/internal/RetriableStream.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1135,13 +1135,18 @@ public void messagesAvailable(final MessageProducer producer) {
GrpcUtil.closeQuietly(producer);
return;
}
listenerSerializeExecutor.execute(
new Runnable() {
@Override
public void run() {
masterListener.messagesAvailable(producer);
}
});
try {
listenerSerializeExecutor.execute(
new Runnable() {
@Override
public void run() {
masterListener.messagesAvailable(producer);
}
});
} catch (Throwable e) {
GrpcUtil.closeQuietly(producer);
throw e;
}
}

@Override
Expand Down
8 changes: 7 additions & 1 deletion core/src/main/java/io/grpc/internal/ServerImpl.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -73,6 +73,7 @@
import java.util.concurrent.Executor;
import java.util.concurrent.Future;
import java.util.concurrent.FutureTask;
import java.util.concurrent.RejectedExecutionException;
import java.util.concurrent.TimeUnit;
import java.util.logging.Level;
import java.util.logging.Logger;
Expand Down Expand Up @@ -857,7 +858,12 @@ public void runInContext() {
}
}

callExecutor.execute(new MessagesAvailable());
try {
callExecutor.execute(new MessagesAvailable());
} catch (RejectedExecutionException e) {
GrpcUtil.closeQuietly(producer);
throw e;
}
}
}

Expand Down
42 changes: 42 additions & 0 deletions core/src/test/java/io/grpc/internal/ClientCallImplTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -33,6 +33,7 @@
import static org.mockito.ArgumentMatchers.same;
import static org.mockito.Mockito.doAnswer;
import static org.mockito.Mockito.doThrow;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.never;
import static org.mockito.Mockito.timeout;
import static org.mockito.Mockito.times;
Expand Down Expand Up @@ -75,6 +76,7 @@
import java.util.concurrent.Executor;
import java.util.concurrent.Executors;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.RejectedExecutionException;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
Expand Down Expand Up @@ -1145,4 +1147,44 @@ void release() {
}
}
}

@Test
public void messagesAvailable_executorRejection() throws Exception {
RejectingExecutor rejectingExecutor = new RejectingExecutor();
ClientCallImpl<Void, Void> call = new ClientCallImpl<>(
method,
rejectingExecutor,
baseCallOptions,
clientStreamProvider,
deadlineCancellationExecutor,
channelCallTracer, configSelector);
call.start(callListener, new Metadata());
verify(stream).start(listenerArgumentCaptor.capture());
final ClientStreamListener streamListener = listenerArgumentCaptor.getValue();

StreamListener.MessageProducer producer = mock(StreamListener.MessageProducer.class);
InputStream message = mock(InputStream.class);
when(producer.next()).thenReturn(message).thenReturn(null);

// Call messagesAvailable. The executor will reject the runnable,
// which should catch the RejectedExecutionException, close the message,
// and cancel the stream.
streamListener.messagesAvailable(producer);

// Verify that the producer's message was closed to release resources
verify(message).close();

// Verify that the stream was cancelled with a CANCELLED status containing the rejection cause
ArgumentCaptor<Status> statusCaptor = ArgumentCaptor.forClass(Status.class);
verify(stream).cancel(statusCaptor.capture());
assertEquals(Status.Code.CANCELLED, statusCaptor.getValue().getCode());
assertTrue(statusCaptor.getValue().getCause() instanceof RejectedExecutionException);
}

private static final class RejectingExecutor implements Executor {
@Override
public void execute(Runnable command) {
throw new RejectedExecutionException("rejected");
}
}
}
37 changes: 37 additions & 0 deletions core/src/test/java/io/grpc/internal/RetriableStreamTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -26,13 +26,15 @@
import static org.junit.Assert.assertFalse;
import static org.junit.Assert.assertNotSame;
import static org.junit.Assert.assertNull;
import static org.junit.Assert.assertThrows;
import static org.junit.Assert.assertTrue;
import static org.mockito.AdditionalAnswers.delegatesTo;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.anyInt;
import static org.mockito.ArgumentMatchers.same;
import static org.mockito.Mockito.doAnswer;
import static org.mockito.Mockito.doReturn;
import static org.mockito.Mockito.doThrow;
import static org.mockito.Mockito.inOrder;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.never;
Expand All @@ -54,6 +56,7 @@
import io.grpc.MethodDescriptor.MethodType;
import io.grpc.Status;
import io.grpc.Status.Code;
import io.grpc.StatusRuntimeException;
import io.grpc.StringMarshaller;
import io.grpc.internal.ClientStreamListener.RpcProgress;
import io.grpc.internal.RetriableStream.ChannelBufferMeter;
Expand All @@ -67,6 +70,7 @@
import java.util.Random;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.Executor;
import java.util.concurrent.RejectedExecutionException;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicReference;
Expand Down Expand Up @@ -2992,4 +2996,37 @@ public InputStream next() {
}
}
}

@Test
public void messagesAvailable_executorRejection() throws Exception {
ClientStream mockStream = mock(ClientStream.class);
doReturn(mockStream).when(retriableStreamRecorder).newSubstream(0);
retriableStream.start(masterListener);

ArgumentCaptor<ClientStreamListener> sublistenerCaptor =
ArgumentCaptor.forClass(ClientStreamListener.class);
verify(mockStream).start(sublistenerCaptor.capture());
ClientStreamListener sublistener = sublistenerCaptor.getValue();

// Transition state to committed/winningSubstream
sublistener.headersRead(new Metadata());
verify(retriableStreamRecorder).postCommit();

// Mock masterListener to throw RejectedExecutionException when messagesAvailable is called
doThrow(new RejectedExecutionException("rejected"))
.when(masterListener)
.messagesAvailable(any(StreamListener.MessageProducer.class));

StreamListener.MessageProducer producer = mock(StreamListener.MessageProducer.class);
InputStream message = mock(InputStream.class);
when(producer.next()).thenReturn(message).thenReturn(null);

// Call messagesAvailable and assert that it throws StatusRuntimeException
// (wrapped by SynchronizationContext)
assertThrows(
StatusRuntimeException.class, () -> sublistener.messagesAvailable(producer));

// Verify that the message was closed (releasing ByteBufs)
verify(message).close();
}
}
73 changes: 73 additions & 0 deletions core/src/test/java/io/grpc/internal/ServerImplTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -96,6 +96,7 @@
import java.util.List;
import java.util.concurrent.CyclicBarrier;
import java.util.concurrent.Executor;
import java.util.concurrent.RejectedExecutionException;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
Expand Down Expand Up @@ -1895,4 +1896,76 @@ public void drain() {
}
}
}

@Test
public void messagesAvailable_executorRejection() throws Exception {
final RejectingExecutor rejectingExecutor =
new RejectingExecutor(executor.getScheduledExecutorService());
when(executorPool.getObject()).thenReturn(rejectingExecutor);

final AtomicReference<ServerCall<String, Integer>> callReference = new AtomicReference<>();
mutableFallbackRegistry.addService(
ServerServiceDefinition.builder(new ServiceDescriptor("Waiter", METHOD))
.addMethod(
METHOD,
new ServerCallHandler<String, Integer>() {
@Override
public ServerCall.Listener<String> startCall(
ServerCall<String, Integer> call, Metadata headers) {
callReference.set(call);
return callListener;
}
})
.build());

createAndStartServer();
ServerTransportListener transportListener =
transportServer.registerNewServerTransport(new SimpleServerTransport());
transportListener.transportReady(Attributes.EMPTY);

Metadata requestHeaders = new Metadata();
StatsTraceContext statsTraceCtx =
StatsTraceContext.newServerContext(streamTracerFactories, "Waiter/serve", requestHeaders);
when(stream.statsTraceContext()).thenReturn(statsTraceCtx);

transportListener.streamCreated(stream, "Waiter/serve", requestHeaders);
verify(stream).setListener(streamListenerCaptor.capture());
ServerStreamListener streamListener = streamListenerCaptor.getValue();
assertNotNull(streamListener);

// Run setup tasks (MethodLookup, HandleServerCall)
assertEquals(1, executor.runDueTasks());
assertNotNull(callReference.get());

// Enable rejection
rejectingExecutor.reject = true;

StreamListener.MessageProducer producer = mock(StreamListener.MessageProducer.class);
InputStream message = mock(InputStream.class);
when(producer.next()).thenReturn(message).thenReturn(null);

// Call messagesAvailable and assert that it throws RejectedExecutionException
assertThrows(
RejectedExecutionException.class, () -> streamListener.messagesAvailable(producer));

// Verify that the message was closed (releasing ByteBufs)
verify(message).close();
}

private static final class RejectingExecutor implements Executor {
private final Executor delegate;
boolean reject = false;

RejectingExecutor(Executor delegate) {
this.delegate = delegate;
}

@Override
public void execute(Runnable command) {
if (reject) {
throw new RejectedExecutionException("rejected");
}
delegate.execute(command);
}
}
}
Loading