[GCS] global state query node info table from GCS. (#8498)

java/runtime/src/main/java/io/ray/runtime/RayNativeRuntime.java

@@ -108,8 +108,12 @@ public final class RayNativeRuntime extends AbstractRayRuntime {
         manager.cleanup();
         manager = null;
       }
-      RayConfig.reset();
     }
+    if (null != gcsClient) {
+      gcsClient.destroy();
+      gcsClient = null;
+    }
+    RayConfig.reset();
     LOGGER.info("RayNativeRuntime shutdown");
   }
 

java/runtime/src/main/java/io/ray/runtime/gcs/GcsClient.java

@@ -32,6 +32,7 @@ public class GcsClient {
   private RedisClient primary;
 
   private List<RedisClient> shards;
+  private GlobalStateAccessor globalStateAccessor;
 
   public GcsClient(String redisAddress, String redisPassword) {
     primary = new RedisClient(redisAddress, redisPassword);
@@ -49,16 +50,11 @@ public class GcsClient {
     shards = shardAddresses.stream().map((byte[] address) -> {
       return new RedisClient(new String(address), redisPassword);
     }).collect(Collectors.toList());
+    globalStateAccessor = GlobalStateAccessor.getInstance(redisAddress, redisPassword);
   }
 
   public List<NodeInfo> getAllNodeInfo() {
-    final String prefix = TablePrefix.CLIENT.toString();
-    final byte[] key = ArrayUtils.addAll(prefix.getBytes(), UniqueId.NIL.getBytes());
-    List<byte[]> results = primary.lrange(key, 0, -1);
-
-    if (results == null) {
-      return new ArrayList<>();
-    }
+    List<byte[]> results = globalStateAccessor.getAllNodeInfo();
 
     // This map is used for deduplication of node entries.
     Map<UniqueId, NodeInfo> nodes = new HashMap<>();
@@ -191,6 +187,15 @@ public class GcsClient {
     return JobId.fromInt(jobCounter);
   }
 
+  /**
+   * Destroy global state accessor when ray native runtime will be shutdown.
+   */
+  public void destroy() {
+    // Only ray shutdown should call gcs client destroy.
+    LOGGER.debug("Destroying global state accessor.");
+    GlobalStateAccessor.destroyInstance();
+  }
+
   private RedisClient getShardClient(BaseId key) {
     return shards.get((int) Long.remainderUnsigned(IdUtil.murmurHashCode(key),
         shards.size()));

java/runtime/src/main/java/io/ray/runtime/gcs/GlobalStateAccessor.java

@@ -24,6 +24,7 @@ public class GlobalStateAccessor {
   public static synchronized void destroyInstance() {
     if (null != globalStateAccessor) {
       globalStateAccessor.destroyGlobalStateAccessor();
+      globalStateAccessor = null;
     }
   }
 
@@ -45,7 +46,8 @@ public class GlobalStateAccessor {
   public List<byte[]> getAllJobInfo() {
     // Fetch a job list with protobuf bytes format from GCS.
     synchronized (GlobalStateAccessor.class) {
-      Preconditions.checkState(globalStateAccessorNativePointer != 0);
+      Preconditions.checkState(globalStateAccessorNativePointer != 0,
+          "Get all job info when global state accessor have been destroyed.");
       return this.nativeGetAllJobInfo(globalStateAccessorNativePointer);
     }
   }
@@ -56,7 +58,8 @@ public class GlobalStateAccessor {
   public List<byte[]> getAllNodeInfo() {
     // Fetch a node list with protobuf bytes format from GCS.
     synchronized (GlobalStateAccessor.class) {
-      Preconditions.checkState(globalStateAccessorNativePointer != 0);
+      Preconditions.checkState(globalStateAccessorNativePointer != 0,
+          "Get all node info when global state accessor have been destroyed.");
       return this.nativeGetAllNodeInfo(globalStateAccessorNativePointer);
     }
   }