From 219fa6706ccaae5293e8aab7ca33f023d0255878 Mon Sep 17 00:00:00 2001
From: Zhiguo Wu <wuzhiguo@apache.org>
Date: Wed, 3 Jul 2024 13:58:54 +0800
Subject: [PATCH] grpc client add getFutureStub

---
 .../bigtop/manager/server/grpc/GrpcClient.java       | 12 ++++++++++++
 1 file changed, 12 insertions(+)

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
index 9d16eac6..61422456 100644
--- 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()