From 4f5ce96080a200203ce20b4354cb910ac0593279 Mon Sep 17 00:00:00 2001 From: Kannan J Date: Fri, 3 Jul 2026 10:19:53 +0000 Subject: [PATCH] core: fix off-heap ByteBuf leak on RejectedExecutionException When a client or server executor is saturated and rejects task execution (e.g., throwing a RejectedExecutionException), the stream listener's messagesAvailable callback fails to schedule the task. As a result, the underlying message stream is never closed and the referenced off-heap Netty ByteBuf instances are leaked. This commit resolves the leak by: 1. Catching task execution exceptions in JumpToApplicationThreadServerStreamListener (ServerImpl), ClientStreamListenerImpl (ClientCallImpl), and Sublistener (RetriableStream). 2. Cleanly releasing the MessageProducer direct memory resources using GrpcUtil.closeQuietly(producer) upon task rejection. 3. Propagating the rejection exception (or triggering the cancel flow) to ensure the transport layers correctly clean up the stream state. Unit tests have been added to ServerImplTest, ClientCallImplTest, and RetriableStreamTest to verify that the message stream is closed and resources are released when task rejection occurs. --- .../java/io/grpc/internal/ClientCallImpl.java | 9 ++- .../io/grpc/internal/RetriableStream.java | 19 +++-- .../java/io/grpc/internal/ServerImpl.java | 8 +- .../io/grpc/internal/ClientCallImplTest.java | 42 +++++++++++ .../io/grpc/internal/RetriableStreamTest.java | 37 ++++++++++ .../java/io/grpc/internal/ServerImplTest.java | 73 +++++++++++++++++++ 6 files changed, 179 insertions(+), 9 deletions(-) diff --git a/core/src/main/java/io/grpc/internal/ClientCallImpl.java b/core/src/main/java/io/grpc/internal/ClientCallImpl.java index 4b24b1eae3d..5a0022fce98 100644 --- a/core/src/main/java/io/grpc/internal/ClientCallImpl.java +++ b/core/src/main/java/io/grpc/internal/ClientCallImpl.java @@ -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; @@ -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.")); + } } } diff --git a/core/src/main/java/io/grpc/internal/RetriableStream.java b/core/src/main/java/io/grpc/internal/RetriableStream.java index 0c37a0beaca..cb5a23bc4d1 100644 --- a/core/src/main/java/io/grpc/internal/RetriableStream.java +++ b/core/src/main/java/io/grpc/internal/RetriableStream.java @@ -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 diff --git a/core/src/main/java/io/grpc/internal/ServerImpl.java b/core/src/main/java/io/grpc/internal/ServerImpl.java index d9f64c2d473..f9cbb5817f7 100644 --- a/core/src/main/java/io/grpc/internal/ServerImpl.java +++ b/core/src/main/java/io/grpc/internal/ServerImpl.java @@ -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; @@ -857,7 +858,12 @@ public void runInContext() { } } - callExecutor.execute(new MessagesAvailable()); + try { + callExecutor.execute(new MessagesAvailable()); + } catch (RejectedExecutionException e) { + GrpcUtil.closeQuietly(producer); + throw e; + } } } diff --git a/core/src/test/java/io/grpc/internal/ClientCallImplTest.java b/core/src/test/java/io/grpc/internal/ClientCallImplTest.java index 03e613e13d9..417367c8096 100644 --- a/core/src/test/java/io/grpc/internal/ClientCallImplTest.java +++ b/core/src/test/java/io/grpc/internal/ClientCallImplTest.java @@ -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; @@ -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; @@ -1145,4 +1147,44 @@ void release() { } } } + + @Test + public void messagesAvailable_executorRejection() throws Exception { + RejectingExecutor rejectingExecutor = new RejectingExecutor(); + ClientCallImpl 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 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"); + } + } } diff --git a/core/src/test/java/io/grpc/internal/RetriableStreamTest.java b/core/src/test/java/io/grpc/internal/RetriableStreamTest.java index afbdaa395b0..ed9a95521d4 100644 --- a/core/src/test/java/io/grpc/internal/RetriableStreamTest.java +++ b/core/src/test/java/io/grpc/internal/RetriableStreamTest.java @@ -26,6 +26,7 @@ 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; @@ -33,6 +34,7 @@ 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; @@ -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; @@ -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; @@ -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 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(); + } } diff --git a/core/src/test/java/io/grpc/internal/ServerImplTest.java b/core/src/test/java/io/grpc/internal/ServerImplTest.java index 91969dd6910..6072f4b4748 100644 --- a/core/src/test/java/io/grpc/internal/ServerImplTest.java +++ b/core/src/test/java/io/grpc/internal/ServerImplTest.java @@ -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; @@ -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> callReference = new AtomicReference<>(); + mutableFallbackRegistry.addService( + ServerServiceDefinition.builder(new ServiceDescriptor("Waiter", METHOD)) + .addMethod( + METHOD, + new ServerCallHandler() { + @Override + public ServerCall.Listener startCall( + ServerCall 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); + } + } }