add judgement for SimpleForwardingServerCall in method ServerStreamHe…

…lper.getServerStream (#1044) * add judgement for SimpleForwardingServerCall in method ServerStreamHelper.getServerStream * code format
sofastack · Nov 27, 2023 · f84d40a · f84d40a
1 parent 5cae8a5
commit f84d40a
Showing 1 changed file with 17 additions and 8 deletions.
diff --git a/jraft-extension/rpc-grpc-impl/src/main/java/io/grpc/internal/ServerStreamHelper.java b/jraft-extension/rpc-grpc-impl/src/main/java/io/grpc/internal/ServerStreamHelper.java
@@ -16,10 +16,10 @@
  */
 package io.grpc.internal;
 
-import io.grpc.ServerCall;
-
 import com.alipay.sofa.jraft.util.internal.ReferenceFieldUpdater;
 import com.alipay.sofa.jraft.util.internal.Updaters;
+import io.grpc.ForwardingServerCall;
+import io.grpc.ServerCall;
 
 /**
  * Get grpc's server stream.
@@ -28,14 +28,23 @@
  */
 public class ServerStreamHelper {
 
-    private static final ReferenceFieldUpdater<ServerCallImpl<?, ?>, ServerStream> STREAM_GETTER = Updaters
-                                                                                                     .newReferenceFieldUpdater(
-                                                                                                         ServerCallImpl.class,
-                                                                                                         "stream");
+    private static final ReferenceFieldUpdater<ServerCallImpl<?, ?>, ServerStream>                                      STREAM_GETTER      = Updaters
+                                                                                                                                               .newReferenceFieldUpdater(
+                                                                                                                                                   ServerCallImpl.class,
+                                                                                                                                                   "stream");
+
+    private static final ReferenceFieldUpdater<ForwardingServerCall.SimpleForwardingServerCall<?, ?>, ServerCall<?, ?>> SERVER_CALL_GETTER = Updaters
+                                                                                                                                               .newReferenceFieldUpdater(
+                                                                                                                                                   ForwardingServerCall.SimpleForwardingServerCall.class,
+                                                                                                                                                   "delegate");
 
     public static ServerStream getServerStream(final ServerCall<?, ?> call) {
-        if (call instanceof ServerCallImpl) {
-            return STREAM_GETTER.get((ServerCallImpl<?, ?>) call);
+        ServerCall<?, ?> lastServerCall = call;
+        if (call instanceof ForwardingServerCall.SimpleForwardingServerCall) {
+            lastServerCall = SERVER_CALL_GETTER.get((ForwardingServerCall.SimpleForwardingServerCall<?, ?>) call);
+        }
+        if (lastServerCall instanceof ServerCallImpl) {
+            return STREAM_GETTER.get((ServerCallImpl<?, ?>) lastServerCall);
         }
         return null;
     }