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.

Author: kannanjgithub

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."));
+        }
       }
     }
 

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

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;
+        }
       }
     }
 

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<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");
+    }
+  }
 }

core/src/test/java/io/grpc/internal/RetriableStreamTest.java

@@ -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;
@@ -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<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();
+  }
 }

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<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);
+    }
+  }
 }