[GCS] Move node resource info from client table to resource table (#5050)

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

@@ -1,6 +1,7 @@
 package org.ray.runtime.gcs;
 
 import com.google.common.base.Preconditions;
+import com.google.common.collect.ImmutableMap;
 import com.google.protobuf.InvalidProtocolBufferException;
 import java.util.ArrayList;
 import java.util.HashMap;
@@ -13,9 +14,9 @@ import org.ray.api.id.BaseId;
 import org.ray.api.id.TaskId;
 import org.ray.api.id.UniqueId;
 import org.ray.api.runtimecontext.NodeInfo;
+import org.ray.runtime.generated.Gcs;
 import org.ray.runtime.generated.Gcs.ActorCheckpointIdData;
 import org.ray.runtime.generated.Gcs.ClientTableData;
-import org.ray.runtime.generated.Gcs.ClientTableData.EntryType;
 import org.ray.runtime.generated.Gcs.TablePrefix;
 import org.ray.runtime.util.IdUtil;
 import org.slf4j.Logger;
@@ -72,42 +73,47 @@ public class GcsClient {
       final UniqueId clientId = UniqueId
           .fromByteBuffer(data.getClientId().asReadOnlyByteBuffer());
 
-      if (data.getEntryType() == EntryType.INSERTION) {
+      if (data.getIsInsertion()) {
         //Code path of node insertion.
-        Map<String, Double> resources = new HashMap<>();
-        // Compute resources.
-        Preconditions.checkState(
-            data.getResourcesTotalLabelCount() == data.getResourcesTotalCapacityCount());
-        for (int i = 0; i < data.getResourcesTotalLabelCount(); i++) {
-          resources.put(data.getResourcesTotalLabel(i), data.getResourcesTotalCapacity(i));
-        }
         NodeInfo nodeInfo = new NodeInfo(
-            clientId, data.getNodeManagerAddress(), true, resources);
+            clientId, data.getNodeManagerAddress(), true, ImmutableMap.of());
         clients.put(clientId, nodeInfo);
-      } else if (data.getEntryType() == EntryType.RES_CREATEUPDATE) {
-        Preconditions.checkState(clients.containsKey(clientId));
-        NodeInfo nodeInfo = clients.get(clientId);
-        for (int i = 0; i < data.getResourcesTotalLabelCount(); i++) {
-          nodeInfo.resources.put(data.getResourcesTotalLabel(i), data.getResourcesTotalCapacity(i));
-        }
-      } else if (data.getEntryType() == EntryType.RES_DELETE) {
-        Preconditions.checkState(clients.containsKey(clientId));
-        NodeInfo nodeInfo = clients.get(clientId);
-        for (int i = 0; i < data.getResourcesTotalLabelCount(); i++) {
-          nodeInfo.resources.remove(data.getResourcesTotalLabel(i));
-        }
       } else {
         // Code path of node deletion.
-        Preconditions.checkState(data.getEntryType() == EntryType.DELETION);
         NodeInfo nodeInfo = new NodeInfo(clientId, clients.get(clientId).nodeAddress,
-            false, clients.get(clientId).resources);
+            false, ImmutableMap.of());
         clients.put(clientId, nodeInfo);
       }
     }
 
+    // Fill resources.
+    for (Map.Entry<UniqueId, NodeInfo> client : clients.entrySet()) {
+      if (client.getValue().isAlive) {
+        client.getValue().resources.putAll(getResourcesForClient(client.getKey()));
+      }
+    }
+
     return new ArrayList<>(clients.values());
   }
 
+  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());
+    }
+    return resources;
+  }
+
   /**
    * If the actor exists in GCS.
    */

java/runtime/src/main/java/org/ray/runtime/gcs/RedisClient.java

@@ -49,14 +49,18 @@ public class RedisClient {
         return jedis.hset(key, field, value);
       }
     }
-
   }
 
   public String hmset(String key, Map<String, String> hash) {
     try (Jedis jedis = jedisPool.getResource()) {
       return jedis.hmset(key, hash);
     }
+  }
 
+  public Map<byte[], byte[]> hgetAll(byte[] key) {
+    try (Jedis jedis = jedisPool.getResource()) {
+      return jedis.hgetAll(key);
+    }
   }
 
   public String get(final String key, final String field) {
@@ -67,7 +71,6 @@ public class RedisClient {
         return jedis.hget(key, field);
       }
     }
-
   }
 
   public byte[] get(byte[] key) {