grpc client add getFutureStub

apache · Jul 3, 2024 · 219fa67 · 219fa67
1 parent 8d4fd62
commit 219fa67
Showing 1 changed file with 12 additions and 0 deletions.
diff --git a/bigtop-manager-server/src/main/java/org/apache/bigtop/manager/server/grpc/GrpcClient.java b/bigtop-manager-server/src/main/java/org/apache/bigtop/manager/server/grpc/GrpcClient.java
@@ -29,6 +29,7 @@
 import io.grpc.ManagedChannelBuilder;
 import io.grpc.stub.AbstractAsyncStub;
 import io.grpc.stub.AbstractBlockingStub;
+import io.grpc.stub.AbstractFutureStub;
 import io.grpc.stub.AbstractStub;
 import lombok.extern.slf4j.Slf4j;
 
@@ -46,6 +47,7 @@ public class GrpcClient {
     // The key of outer map is hostname, inner map is stub class name
     private static final Map<String, Map<String, AbstractBlockingStub<?>>> BLOCKING_STUBS = new ConcurrentHashMap<>();
     private static final Map<String, Map<String, AbstractAsyncStub<?>>> ASYNC_STUBS = new ConcurrentHashMap<>();
+    private static final Map<String, Map<String, AbstractFutureStub<?>>> FUTURE_STUBS = new ConcurrentHashMap<>();
 
     public static ManagedChannel createChannel(String host) {
         int port = SpringContextHolder.getApplicationContext()
@@ -79,6 +81,16 @@ public static <T extends AbstractAsyncStub<T>> T getAsyncStub(String host, Class
         });
     }
 
+    @SuppressWarnings("unchecked")
+    public static <T extends AbstractFutureStub<T>> T getFutureStub(String host, Class<T> clazz) {
+        Map<String, AbstractFutureStub<?>> innerMap = FUTURE_STUBS.computeIfAbsent(host, k -> new HashMap<>());
+        return (T) innerMap.computeIfAbsent(clazz.getName(), k -> {
+            T instance = T.newStub(getFactory(clazz), getChannel(host));
+            log.info("Instance: {} created.", k);
+            return instance;
+        });
+    }
+
     private static ManagedChannel createChannel(String host, Integer port) {
         ManagedChannel channel = ManagedChannelBuilder.forAddress(host, port)
                 .usePlaintext()