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