Merge pull request #15212 from Andro999b/master

GRPC client onClose and onMessage should be hanled on same thread
quarkusio · Feb 21, 2021 · aa802d8 · aa802d8
2 parents d034fff + 664c16f
commit aa802d8
Showing 1 changed file with 11 additions and 2 deletions.
diff --git a/...rpc/runtime/src/main/java/io/quarkus/grpc/runtime/supports/IOThreadClientInterceptor.java b/...rpc/runtime/src/main/java/io/quarkus/grpc/runtime/supports/IOThreadClientInterceptor.java
@@ -23,10 +23,19 @@ public void start(Listener<RespT> responseListener, Metadata headers) {
                 super.start(new ForwardingClientCallListener.SimpleForwardingClientCallListener<RespT>(responseListener) {
                     @Override
                     public void onMessage(RespT message) {
+                        runInContextIfNeed(() -> super.onMessage(message));
+                    }
+
+                    @Override
+                    public void onClose(Status status, Metadata trailers) {
+                        runInContextIfNeed(() -> super.onClose(status, trailers));
+                    }
+
+                    private void runInContextIfNeed(Runnable fun) {
                         if (context != null) {
-                            context.runOnContext(unused -> super.onMessage(message));
+                            context.runOnContext(unused -> fun.run());
                         } else {
-                            super.onMessage(message);
+                            fun.run();
                         }
                     }
                 }, headers);