fix: release late response ByteBuf when callback executor is shut down (#16415)

dubbo-rpc/dubbo-rpc-triple/src/main/java/org/apache/dubbo/rpc/protocol/tri/stream/AbstractStream.java

@@ -18,6 +18,7 @@
 
 import org.apache.dubbo.common.threadpool.serial.SerializingExecutor;
 import org.apache.dubbo.common.utils.ClassUtils;
+import org.apache.dubbo.common.utils.ExecutorUtil;
 import org.apache.dubbo.rpc.model.FrameworkModel;
 
 import java.util.concurrent.Executor;
@@ -30,17 +31,38 @@ public abstract class AbstractStream implements Stream {
     protected Executor executor;
     protected final FrameworkModel frameworkModel;
 
+    /**
+     * The executor passed by the caller, kept for shutdown-state checks. It is wrapped by
+     * {@link SerializingExecutor} exposed via {@link #executor}.
+     */
+    private Executor callbackExecutor;
+
     private static final boolean HAS_PROTOBUF = ClassUtils.hasProtobuf();
 
     public AbstractStream(Executor executor, FrameworkModel frameworkModel) {
+        this.callbackExecutor = executor;
         this.executor = new SerializingExecutor(executor);
         this.frameworkModel = frameworkModel;
     }
 
     public void setExecutor(Executor executor) {
+        this.callbackExecutor = executor;
         this.executor = new SerializingExecutor(executor);
     }
 
+    /**
+     * Whether the callback executor has been shut down (e.g. the {@code ThreadlessExecutor}
+     * of a sync call is shut down after a request timeout). {@link SerializingExecutor}
+     * silently drops tasks submitted to a shut-down executor, so callers must release any
+     * task-owned resources (e.g. a ByteBuf) when this returns {@code true}.
+     *
+     * @return true if the callback executor is an {@link java.util.concurrent.ExecutorService}
+     *         that has been shut down
+     */
+    protected boolean isCallbackExecutorShutdown() {
+        return callbackExecutor != null && ExecutorUtil.isShutdown(callbackExecutor);
+    }
+
     public static boolean getGrpcStatusDetailEnabled() {
         return HAS_PROTOBUF;
     }

dubbo-rpc/dubbo-rpc-triple/src/main/java/org/apache/dubbo/rpc/protocol/tri/stream/AbstractTripleClientStream.java

@@ -565,14 +565,24 @@ public void onData(ByteBuf data, boolean endStream) {
             if (LOGGER.isDebugEnabled()) {
                 LOGGER.debug("endStream: {} DATA: {}", endStream, data.toString(StandardCharsets.UTF_8));
             }
+            if (isCallbackExecutorShutdown()) {
+                // The callback executor (e.g. ThreadlessExecutor) has been shut down, e.g.
+                // after the request timed out in {@link AsyncRpcResult}. SerializingExecutor would
+                // silently drop the submitted task so doOnData would never run; release the
+                // ByteBuf now to avoid an off-heap memory leak.
+                ReferenceCountUtil.release(data);
+                LOGGER.warn(
+                        PROTOCOL_FAILED_RESPONSE, "", "", "Drop late response data: callback executor is shut down");
+                return;
+            }
             try {
                 executor.execute(() -> doOnData(data, endStream));
             } catch (Throwable t) {
-                // Tasks will be rejected when the thread pool is closed or full,
-                // ByteBuf needs to be released to avoid out of heap memory leakage.
-                // For example, ThreadLessExecutor will be shutdown when request timeout {@link AsyncRpcResult}
+                // The task can be rejected when the thread pool is closed or full (e.g. the
+                // ThreadlessExecutor is shut down after a request timeout {@link AsyncRpcResult}).
+                // Release the ByteBuf to avoid a memory leak.
                 ReferenceCountUtil.release(data);
-                LOGGER.error(PROTOCOL_FAILED_RESPONSE, "", "", "submit onData task failed", t);
+                LOGGER.warn(PROTOCOL_FAILED_RESPONSE, "", "", "Drop response data, executor rejected the task", t);
             }
         }
 

dubbo-rpc/dubbo-rpc-triple/src/test/java/org/apache/dubbo/rpc/protocol/tri/stream/TripleClientStreamTest.java

@@ -17,6 +17,7 @@
 package org.apache.dubbo.rpc.protocol.tri.stream;
 
 import org.apache.dubbo.common.URL;
+import org.apache.dubbo.common.threadpool.ThreadlessExecutor;
 import org.apache.dubbo.remoting.http12.HttpHeaderNames;
 import org.apache.dubbo.remoting.http12.message.MediaType;
 import org.apache.dubbo.rpc.TriRpcStatus;
@@ -45,7 +46,6 @@
 import io.netty.buffer.ByteBuf;
 import io.netty.buffer.Unpooled;
 import io.netty.channel.embedded.EmbeddedChannel;
-import io.netty.channel.nio.NioEventLoopGroup;
 import io.netty.handler.codec.http.HttpResponseStatus;
 import io.netty.handler.codec.http.HttpScheme;
 import io.netty.handler.codec.http2.DefaultHttp2Headers;
@@ -79,7 +79,7 @@ void progress() {
         Http2StreamChannel http2StreamChannel = mock(Http2StreamChannel.class);
         when(http2StreamChannel.isActive()).thenReturn(true);
         when(http2StreamChannel.newSucceededFuture()).thenReturn(channel.newSucceededFuture());
-        when(http2StreamChannel.eventLoop()).thenReturn(new NioEventLoopGroup().next());
+        when(http2StreamChannel.eventLoop()).thenReturn(channel.eventLoop());
         when(http2StreamChannel.newPromise()).thenReturn(channel.newPromise());
         when(http2StreamChannel.parent()).thenReturn(channel);
         AbstractTripleClientStream stream = new Http2TripleClientStream(
@@ -131,4 +131,43 @@ void progress() {
         transportListener.onData(buf, false);
         Assertions.assertEquals(1, listener.message.length);
     }
+
+    @Test
+    void testOnDataReleaseByteBufAfterCallbackExecutorShutdown() {
+        final URL url = URL.valueOf("tri://127.0.0.1:8080/foo.bar.service");
+        final ModuleServiceRepository repo =
+                ApplicationModel.defaultModel().getDefaultModule().getServiceRepository();
+        repo.registerService(IGreeter.class);
+        final ServiceDescriptor serviceDescriptor = repo.getService(IGreeter.class.getName());
+        final MethodDescriptor methodDescriptor = serviceDescriptor.getMethod("echo", new Class<?>[] {String.class});
+
+        MockClientStreamListener listener = new MockClientStreamListener();
+        TripleWriteQueue writeQueue = mock(TripleWriteQueue.class);
+        final EmbeddedChannel channel = new EmbeddedChannel();
+        when(writeQueue.enqueueFuture(any(QueuedCommand.class), any(Executor.class)))
+                .thenReturn(channel.newPromise());
+        Http2StreamChannel http2StreamChannel = mock(Http2StreamChannel.class);
+        when(http2StreamChannel.isActive()).thenReturn(true);
+        when(http2StreamChannel.newSucceededFuture()).thenReturn(channel.newSucceededFuture());
+        when(http2StreamChannel.eventLoop()).thenReturn(channel.eventLoop());
+        when(http2StreamChannel.newPromise()).thenReturn(channel.newPromise());
+        when(http2StreamChannel.parent()).thenReturn(channel);
+
+        // Mirror the sync call path (TripleInvoker#doInvoke): a per-call ThreadlessExecutor
+        // wrapped by SerializingExecutor inside AbstractStream.
+        ThreadlessExecutor callbackExecutor = new ThreadlessExecutor();
+        AbstractTripleClientStream stream = new Http2TripleClientStream(
+                url.getOrDefaultFrameworkModel(), callbackExecutor, writeQueue, listener, http2StreamChannel);
+
+        // Simulate request timeout: AsyncRpcResult#get(timeout) shuts down the executor in finally.
+        callbackExecutor.shutdown();
+
+        H2TransportListener transportListener = stream.createTransportListener();
+        final ByteBuf buf = Unpooled.buffer(16);
+        buf.writeByte(1);
+        transportListener.onData(buf, false);
+        // A late DATA frame must not leak the ByteBuf just because the callback executor
+        // was shut down by the timeout.
+        Assertions.assertEquals(0, buf.refCnt());
+    }
 }