[GCS] Global state accessor from node resource table (#8658)

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

@@ -10,7 +10,6 @@ import io.ray.api.id.TaskId;
 import io.ray.api.id.UniqueId;
 import io.ray.api.runtimecontext.NodeInfo;
 import io.ray.runtime.config.RayConfig;
-import io.ray.runtime.gcs.GlobalStateAccessor;
 import io.ray.runtime.generated.Gcs;
 import io.ray.runtime.generated.Gcs.ActorCheckpointIdData;
 import io.ray.runtime.generated.Gcs.GcsNodeInfo;
@@ -96,19 +95,16 @@ public class GcsClient {
   }
 
   private Map<String, Double> getResourcesForClient(UniqueId clientId) {
-    final String prefix = TablePrefix.NODE_RESOURCE.toString();
-    final byte[] key = ArrayUtils.addAll(prefix.getBytes(), clientId.getBytes());
-    Map<byte[], byte[]> results = primary.hgetAll(key);
-    Map<String, Double> resources = new HashMap<>();
-    for (Map.Entry<byte[], byte[]> entry : results.entrySet()) {
-      String resourceName = new String(entry.getKey());
-      Gcs.ResourceTableData resourceTableData;
-      try {
-        resourceTableData = Gcs.ResourceTableData.parseFrom(entry.getValue());
-      } catch (InvalidProtocolBufferException e) {
-        throw new RuntimeException("Received invalid protobuf data from GCS.");
-      }
-      resources.put(resourceName, resourceTableData.getResourceCapacity());
+    byte[] resourceMapBytes = globalStateAccessor.getNodeResourceInfo(clientId);
+    Gcs.ResourceMap resourceMap;
+    try {
+      resourceMap = Gcs.ResourceMap.parseFrom(resourceMapBytes);
+    } catch (InvalidProtocolBufferException e) {
+      throw new RuntimeException("Received invalid protobuf data from GCS.");
+    }
+    HashMap<String, Double> resources = new HashMap<>();
+    for (Map.Entry<String, Gcs.ResourceTableData> entry : resourceMap.getItemsMap().entrySet()) {
+      resources.put(entry.getKey(), entry.getValue().getResourceCapacity());
     }
     return resources;
   }

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

@@ -2,6 +2,7 @@ package io.ray.runtime.gcs;
 
 import com.google.common.base.Preconditions;
 import io.ray.api.id.ActorId;
+import io.ray.api.id.UniqueId;
 import java.util.List;
 
 /**
@@ -65,6 +66,18 @@ public class GlobalStateAccessor {
     }
   }
 
+  /**
+   * @param nodeId node unique id.
+   * @return A map of node resource info in protobuf schema.
+   */
+  public byte[] getNodeResourceInfo(UniqueId nodeId) {
+    synchronized (GlobalStateAccessor.class) {
+      Preconditions.checkState(globalStateAccessorNativePointer != 0,
+          "Get resource info by node id when global state accessor have been destroyed.");
+      return nativeGetNodeResourceInfo(globalStateAccessorNativePointer, nodeId.getBytes());
+    }
+  }
+
   /**
    * @return A list of actor info with ActorInfo protobuf schema.
    */
@@ -120,6 +133,8 @@ public class GlobalStateAccessor {
 
   private native List<byte[]> nativeGetAllNodeInfo(long nativePtr);
 
+  private native byte[] nativeGetNodeResourceInfo(long nativePtr, byte[] nodeId);
+
   private native List<byte[]> nativeGetAllActorInfo(long nativePtr);
 
   private native byte[] nativeGetActorInfo(long nativePtr, byte[] actorId);