[3.0-Triple] Add health service cancel handler (#9004)

* add health service cancel handler * change log * change indent style
apache · Oct 12, 2021 · 23e4456 · 23e4456
1 parent 1dc1822
commit 23e4456
Showing 1 changed file with 20 additions and 7 deletions.
diff --git a/...bbo-rpc-triple/src/main/java/org/apache/dubbo/rpc/protocol/tri/service/TriHealthImpl.java b/...bbo-rpc-triple/src/main/java/org/apache/dubbo/rpc/protocol/tri/service/TriHealthImpl.java
@@ -17,7 +17,10 @@
 
 package org.apache.dubbo.rpc.protocol.tri.service;
 
+import org.apache.dubbo.common.logger.Logger;
+import org.apache.dubbo.common.logger.LoggerFactory;
 import org.apache.dubbo.common.stream.StreamObserver;
+import org.apache.dubbo.rpc.RpcContext;
 import org.apache.dubbo.rpc.RpcException;
 
 import grpc.health.v1.Health;
@@ -28,14 +31,12 @@
 import java.util.IdentityHashMap;
 import java.util.Map;
 import java.util.concurrent.ConcurrentHashMap;
-import java.util.logging.Level;
-import java.util.logging.Logger;
 
 import static org.apache.dubbo.rpc.RpcException.METHOD_NOT_FOUND;
 
 public class TriHealthImpl implements Health {
 
-    private static final Logger logger = Logger.getLogger(TriHealthImpl.class.getName());
+    private static final Logger logger = LoggerFactory.getLogger(TriHealthImpl.class);
 
     // Due to the latency of rpc calls, synchronization of the map does not help with consistency.
     // However, need use ConcurrentHashMap to allow concurrent reading by check().
@@ -82,13 +83,25 @@ public void watch(HealthCheckRequest request, StreamObserver<HealthCheckResponse
             }
             serviceWatchers.put(responseObserver, Boolean.TRUE);
         }
-        // todo add client cancel listener
+        RpcContext.getCancellationContext()
+            .addListener(context -> {
+                synchronized (watchLock) {
+                    IdentityHashMap<StreamObserver<HealthCheckResponse>, Boolean> serviceWatchers =
+                        watchers.get(service);
+                    if (serviceWatchers != null) {
+                        serviceWatchers.remove(responseObserver);
+                        if (serviceWatchers.isEmpty()) {
+                            watchers.remove(service);
+                        }
+                    }
+                }
+            });
     }
 
     void setStatus(String service, HealthCheckResponse.ServingStatus status) {
         synchronized (watchLock) {
             if (terminal) {
-                logger.log(Level.FINE, "Ignoring status {} for {}", new Object[]{status, service});
+                logger.info("Ignoring status " + status + " for " + service);
                 return;
             }
             setStatusInternal(service, status);
@@ -105,7 +118,7 @@ private void setStatusInternal(String service, HealthCheckResponse.ServingStatus
     void clearStatus(String service) {
         synchronized (watchLock) {
             if (terminal) {
-                logger.log(Level.FINE, "Ignoring status clearing for {}", new Object[]{service});
+                logger.info("Ignoring status clearing for " + service);
                 return;
             }
             HealthCheckResponse.ServingStatus prevStatus = statusMap.remove(service);
@@ -118,7 +131,7 @@ void clearStatus(String service) {
     void enterTerminalState() {
         synchronized (watchLock) {
             if (terminal) {
-                logger.log(Level.WARNING, "Already terminating", new RuntimeException());
+                logger.warn("Already terminating", new RuntimeException());
                 return;
             }
             terminal = true;