From 14a99e2be59747df88b7695dce6fbf555056f4fd Mon Sep 17 00:00:00 2001
From: chuntaojun <liaochuntao@live.com>
Date: Tue, 15 Aug 2023 17:28:25 +0800
Subject: [PATCH] =?UTF-8?q?feat:provider=E6=94=AF=E6=8C=81=E9=85=8D?=
 =?UTF-8?q?=E7=BD=AESDKContext?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 .../provider/src/main/resources/polaris.yml   |  2 +-
 .../client/PolarisManagedChannelBuilder.java  | 21 +++++++++++++++++++
 .../grpc/loadbalance/PolarisPicker.java       |  3 ---
 .../polaris/grpc/server/GraceOffline.java     |  7 ++++++-
 .../grpc/server/PolarisGrpcServer.java        |  4 +++-
 .../grpc/server/PolarisGrpcServerBuilder.java | 19 +++++++++++++++--
 6 files changed, 48 insertions(+), 8 deletions(-)

diff --git a/grpc-java-polaris-examples/quickstart-example/provider/src/main/resources/polaris.yml b/grpc-java-polaris-examples/quickstart-example/provider/src/main/resources/polaris.yml
index 3432e5f..f360118 100644
--- a/grpc-java-polaris-examples/quickstart-example/provider/src/main/resources/polaris.yml
+++ b/grpc-java-polaris-examples/quickstart-example/provider/src/main/resources/polaris.yml
@@ -1,7 +1,7 @@
 global:
   serverConnector:
     addresses:
-      - 127.0.0.1:8091
+      - 119.91.66.223:8091
   statReporter:
     enable: true
     plugin:
diff --git a/grpc-java-polaris/src/main/java/com/tencent/polaris/grpc/client/PolarisManagedChannelBuilder.java b/grpc-java-polaris/src/main/java/com/tencent/polaris/grpc/client/PolarisManagedChannelBuilder.java
index 66bad9f..c99c10a 100644
--- a/grpc-java-polaris/src/main/java/com/tencent/polaris/grpc/client/PolarisManagedChannelBuilder.java
+++ b/grpc-java-polaris/src/main/java/com/tencent/polaris/grpc/client/PolarisManagedChannelBuilder.java
@@ -66,14 +66,35 @@ public class PolarisManagedChannelBuilder extends ManagedChannelBuilder<PolarisM
 
     private final ServiceKey sourceService;
 
+    /**
+     * follow {@link ManagedChannelBuilder#forTarget(String)}
+     *
+     * @param target 服务名
+     * @return {@link PolarisManagedChannelBuilder}
+     */
     public static PolarisManagedChannelBuilder forTarget(String target) {
         return forTarget(target, null, null);
     }
 
+    /**
+     * 增强 {@link ManagedChannelBuilder#forTarget(String)}, 在连接到目标服务时允许设置主调服务的相关信息
+     *
+     * @param target 服务名
+     * @param sourceService {@link ServiceKey} 主调服务信息以及标签
+     * @return {@link PolarisManagedChannelBuilder}
+     */
     public static PolarisManagedChannelBuilder forTarget(String target, ServiceKey sourceService) {
         return new PolarisManagedChannelBuilder(target, sourceService, null);
     }
 
+    /**
+     * 增强 {@link ManagedChannelBuilder#forTarget(String)}, 在连接到目标服务时允许设置主调服务的相关信息, 并且可以自定义北极星 SDK 的核心数据结构 {@link SDKContext}
+     *
+     * @param target 服务名
+     * @param sourceService {@link ServiceKey} 主调服务信息以及标签
+     * @param sdkContext {@link SDKContext} 可以设置北极星 SDK 的相关配置以及行为, 例如服务治理中心地址等等
+     * @return {@link PolarisManagedChannelBuilder}
+     */
     public static PolarisManagedChannelBuilder forTarget(String target, ServiceKey sourceService, SDKContext sdkContext) {
         return new PolarisManagedChannelBuilder(target, sourceService, sdkContext);
     }
diff --git a/grpc-java-polaris/src/main/java/com/tencent/polaris/grpc/loadbalance/PolarisPicker.java b/grpc-java-polaris/src/main/java/com/tencent/polaris/grpc/loadbalance/PolarisPicker.java
index f755a20..c9a0125 100644
--- a/grpc-java-polaris/src/main/java/com/tencent/polaris/grpc/loadbalance/PolarisPicker.java
+++ b/grpc-java-polaris/src/main/java/com/tencent/polaris/grpc/loadbalance/PolarisPicker.java
@@ -140,7 +140,6 @@ public PickResult pickSubchannel(PickSubchannelArgs args) {
     }
 
     Instance doLoadBalance(ServiceInstances serviceInstances) {
-
         if (serviceInstances.getInstances().size() == 1) {
             return serviceInstances.getInstances().get(0);
         }
@@ -149,12 +148,10 @@ Instance doLoadBalance(ServiceInstances serviceInstances) {
         request.setDstInstances(serviceInstances);
 
         ProcessLoadBalanceResponse response = routerAPI.processLoadBalance(request);
-
         return response.getTargetInstance();
     }
 
     ServiceInstances doRoute(ServiceInstances serviceInstances, ServiceKey target, PickSubchannelArgs args) {
-
         ProcessRoutersRequest request = new ProcessRoutersRequest();
         request.setDstInstances(serviceInstances);
 
diff --git a/grpc-java-polaris/src/main/java/com/tencent/polaris/grpc/server/GraceOffline.java b/grpc-java-polaris/src/main/java/com/tencent/polaris/grpc/server/GraceOffline.java
index f5a4c69..b4b627b 100644
--- a/grpc-java-polaris/src/main/java/com/tencent/polaris/grpc/server/GraceOffline.java
+++ b/grpc-java-polaris/src/main/java/com/tencent/polaris/grpc/server/GraceOffline.java
@@ -16,6 +16,7 @@
 
 package com.tencent.polaris.grpc.server;
 
+import com.tencent.polaris.client.api.SDKContext;
 import io.grpc.Server;
 import java.time.Duration;
 import java.util.concurrent.TimeUnit;
@@ -34,10 +35,13 @@ public final class GraceOffline {
 
     private final Duration maxWaitDuration;
 
+    private final SDKContext context;
+
     private final AtomicBoolean executed = new AtomicBoolean(false);
 
-    public GraceOffline(Server server, Duration maxWaitDuration) {
+    public GraceOffline(Server server, Duration maxWaitDuration, SDKContext context) {
         this.grpcServer = server;
+        this.context = context;
         this.maxWaitDuration = maxWaitDuration;
     }
 
@@ -60,6 +64,7 @@ public Server shutdown() {
         } catch (InterruptedException ignore) {
             Thread.currentThread().interrupt();
         }
+        context.close();
         return grpcServer;
     }
 
diff --git a/grpc-java-polaris/src/main/java/com/tencent/polaris/grpc/server/PolarisGrpcServer.java b/grpc-java-polaris/src/main/java/com/tencent/polaris/grpc/server/PolarisGrpcServer.java
index d07369b..13ec45c 100644
--- a/grpc-java-polaris/src/main/java/com/tencent/polaris/grpc/server/PolarisGrpcServer.java
+++ b/grpc-java-polaris/src/main/java/com/tencent/polaris/grpc/server/PolarisGrpcServer.java
@@ -105,11 +105,12 @@ public Server start() throws IOException {
     public Server shutdown() {
         if (shutdownOnce.compareAndSet(false, true)) {
             executorService.shutdownNow();
+            // 将自己从注册中心反注册掉
             this.deregister(targetServer.getServices());
             providerAPI.destroy();
         }
 
-        return new GraceOffline(targetServer, maxWaitDuration).shutdown();
+        return new GraceOffline(targetServer, maxWaitDuration, context).shutdown();
     }
 
     @Override
@@ -118,6 +119,7 @@ public Server shutdownNow() {
             executorService.shutdownNow();
             this.deregister(targetServer.getServices());
             providerAPI.destroy();
+            context.close();
         }
         return this.targetServer.shutdownNow();
     }
diff --git a/grpc-java-polaris/src/main/java/com/tencent/polaris/grpc/server/PolarisGrpcServerBuilder.java b/grpc-java-polaris/src/main/java/com/tencent/polaris/grpc/server/PolarisGrpcServerBuilder.java
index 0370814..de314d5 100644
--- a/grpc-java-polaris/src/main/java/com/tencent/polaris/grpc/server/PolarisGrpcServerBuilder.java
+++ b/grpc-java-polaris/src/main/java/com/tencent/polaris/grpc/server/PolarisGrpcServerBuilder.java
@@ -34,6 +34,7 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Objects;
 import java.util.concurrent.Executor;
 import javax.annotation.Nullable;
 
@@ -77,7 +78,7 @@ public final class PolarisGrpcServerBuilder extends ServerBuilder<PolarisGrpcSer
 
     private final List<ServerInterceptor> interceptors = new ArrayList<>();
 
-    private final SDKContext context = SDKContext.initContext();
+    private SDKContext context;
 
     /**
      * Static factory for creating a new PolarisGrpcServerBuilder.
@@ -98,6 +99,17 @@ public static PolarisGrpcServerBuilder forPort(int port) {
     public PolarisGrpcServerBuilder(ServerBuilder<?> builder) {
         this.builder = builder;
     }
+
+    /**
+     * Set polaris SDK Context
+     *
+     * @param context polaris sdk core object
+     * @return PolarisGrpcServerBuilder
+     */
+    public PolarisGrpcServerBuilder sdkContext(SDKContext context) {
+        this.context = context;
+        return this;
+    }
     
     /**
      * Set grpc service name.
@@ -136,7 +148,7 @@ public PolarisGrpcServerBuilder metadata(Map<String, String> metadata) {
      * set instance weight
      *
      * @param weight
-     * @return
+     * @return PolarisGrpcServerBuilder
      */
     public PolarisGrpcServerBuilder weight(int weight) {
         this.weight = weight;
@@ -290,6 +302,9 @@ public Server build() {
     }
 
     private void setDefault() {
+        if (Objects.isNull(context)) {
+            context = SDKContext.initContext();
+        }
         if (StringUtils.isBlank(namespace)) {
             this.namespace = DEFAULT_NAMESPACE;
         }