[ISSUE apache#7833] Fix invokeImpl() in RemotingAbstract

Jianing-Xu · Feb 19, 2024 · 7dc0e5a · 7dc0e5a
1 parent 1b46d18
commit 7dc0e5a
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 7 deletions.
diff --git a/remoting/src/main/java/org/apache/rocketmq/remoting/netty/NettyRemotingAbstract.java b/remoting/src/main/java/org/apache/rocketmq/remoting/netty/NettyRemotingAbstract.java
@@ -496,13 +496,7 @@ public RemotingCommand invokeSyncImpl(final Channel channel, final RemotingComma
 
     public CompletableFuture<ResponseFuture> invokeImpl(final Channel channel, final RemotingCommand request,
         final long timeoutMillis) {
-        String channelRemoteAddr = RemotingHelper.parseChannelRemoteAddr(channel);
-        doBeforeRpcHooks(channelRemoteAddr, request);
-        return invoke0(channel, request, timeoutMillis).whenComplete((v, t) -> {
-            if (t == null) {
-                doAfterRpcHooks(channelRemoteAddr, request, v.getResponseCommand());
-            }
-        });
+        return invoke0(channel, request, timeoutMillis);
     }
 
     protected CompletableFuture<ResponseFuture> invoke0(final Channel channel, final RemotingCommand request,

diff --git a/remoting/src/main/java/org/apache/rocketmq/remoting/netty/NettyRemotingClient.java b/remoting/src/main/java/org/apache/rocketmq/remoting/netty/NettyRemotingClient.java
@@ -804,6 +804,9 @@ public CompletableFuture<RemotingCommand> invoke(String addr, RemotingCommand re
     public CompletableFuture<ResponseFuture> invokeImpl(final Channel channel, final RemotingCommand request,
         final long timeoutMillis) {
         Stopwatch stopwatch = Stopwatch.createStarted();
+        String channelRemoteAddr = RemotingHelper.parseChannelRemoteAddr(channel);
+        doBeforeRpcHooks(channelRemoteAddr, request);
+
         return super.invokeImpl(channel, request, timeoutMillis).thenCompose(responseFuture -> {
             RemotingCommand response = responseFuture.getResponseCommand();
             if (response.getCode() == ResponseCode.GO_AWAY) {
@@ -839,6 +842,10 @@ public CompletableFuture<ResponseFuture> invokeImpl(final Channel channel, final
                 }
             }
             return CompletableFuture.completedFuture(responseFuture);
+        }).whenComplete((v, t) -> {
+            if (t == null) {
+                doAfterRpcHooks(channelRemoteAddr, request, v.getResponseCommand());
+            }
         });
     }