From f84d40a04fdd8285aca048b26ba36a4df6745420 Mon Sep 17 00:00:00 2001
From: dingjiefei <1264677051@qq.com>
Date: Mon, 27 Nov 2023 22:41:36 +0800
Subject: [PATCH] add judgement for SimpleForwardingServerCall in method
 ServerStreamHelper.getServerStream (#1044)

* add judgement for SimpleForwardingServerCall in method ServerStreamHelper.getServerStream

* code format
---
 .../io/grpc/internal/ServerStreamHelper.java  | 25 +++++++++++++------
 1 file changed, 17 insertions(+), 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
index 73abb4173..390f23abd 100644
--- 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;
     }