[gRPC] Migrate gcs data structures to protobuf (#5024)

java/BUILD.bazel

@@ -1,4 +1,5 @@
 load("//bazel:ray.bzl", "flatbuffer_java_library", "define_java_module")
+load("@build_stack_rules_proto//java:java_proto_compile.bzl", "java_proto_compile")
 
 exports_files([
     "testng.xml",
@@ -50,6 +51,7 @@ define_java_module(
     name = "runtime",
     additional_srcs = [
         ":generate_java_gcs_fbs",
+        ":gcs_java_proto",
     ],
     additional_resources = [
         ":java_native_deps",
@@ -68,6 +70,7 @@ define_java_module(
         "@plasma//:org_apache_arrow_arrow_plasma",
         "@maven//:com_github_davidmoten_flatbuffers_java",
         "@maven//:com_google_guava_guava",
+        "@maven//:com_google_protobuf_protobuf_java",
         "@maven//:com_typesafe_config",
         "@maven//:commons_io_commons_io",
         "@maven//:de_ruedigermoeller_fst",
@@ -148,38 +151,16 @@ java_binary(
     ],
 )
 
+java_proto_compile(
+    name = "gcs_java_proto",
+    deps = ["@//:gcs_proto"],
+)
+
 flatbuffers_generated_files = [
-    "ActorCheckpointData.java",
-    "ActorCheckpointIdData.java",
-    "ActorState.java",
-    "ActorTableData.java",
     "Arg.java",
-    "ClassTableData.java",
-    "ClientTableData.java",
-    "ConfigTableData.java",
-    "CustomSerializerData.java",
-    "DriverTableData.java",
-    "EntryType.java",
-    "ErrorTableData.java",
-    "ErrorType.java",
-    "FunctionTableData.java",
-    "GcsEntry.java",
-    "HeartbeatBatchTableData.java",
-    "HeartbeatTableData.java",
     "Language.java",
-    "ObjectTableData.java",
-    "ProfileEvent.java",
-    "ProfileTableData.java",
-    "RayResource.java",
-    "ResourcePair.java",
-    "SchedulingState.java",
-    "TablePrefix.java",
-    "TablePubsub.java",
     "TaskInfo.java",
-    "TaskLeaseData.java",
-    "TaskReconstructionData.java",
-    "TaskTableData.java",
-    "TaskTableTestAndUpdate.java",
+    "ResourcePair.java",
 ]
 
 flatbuffer_java_library(
@@ -198,7 +179,7 @@ genrule(
     cmd = """
         for f in $(locations //java:java_gcs_fbs); do
             chmod +w $$f
-            cp -f $$f $(@D)/runtime/src/main/java/org/ray/runtime/generated
+            mv -f $$f $(@D)/runtime/src/main/java/org/ray/runtime/generated
         done
         python $$(pwd)/java/modify_generated_java_flatbuffers_files.py $(@D)/..
     """,
@@ -221,8 +202,10 @@ filegroup(
 genrule(
     name = "gen_maven_deps",
     srcs = [
-        ":java_native_deps",
+        ":gcs_java_proto",
         ":generate_java_gcs_fbs",
+        ":java_native_deps",
+        ":copy_pom_file",
         "@plasma//:org_apache_arrow_arrow_plasma",
     ],
     outs = ["gen_maven_deps.out"],
@@ -237,10 +220,15 @@ genrule(
             chmod +w $$f
             cp $$f $$NATIVE_DEPS_DIR
         done
-        # Copy flatbuffers-generated files
+        # Copy protobuf-generated files.
         GENERATED_DIR=$$WORK_DIR/java/runtime/src/main/java/org/ray/runtime/generated
         rm -rf $$GENERATED_DIR
         mkdir -p $$GENERATED_DIR
+        for f in $(locations //java:gcs_java_proto); do
+            unzip $$f
+            mv org/ray/runtime/generated/* $$GENERATED_DIR
+        done
+        # Copy flatbuffers-generated files
         for f in $(locations //java:generate_java_gcs_fbs); do
             cp $$f $$GENERATED_DIR
         done
@@ -250,6 +238,7 @@ genrule(
         echo $$(date) > $@
     """,
     local = 1,
+    tags = ["no-cache"],
 )
 
 genrule(

java/dependencies.bzl

@@ -6,6 +6,7 @@ def gen_java_deps():
             "com.beust:jcommander:1.72",
             "com.github.davidmoten:flatbuffers-java:1.9.0.1",
             "com.google.guava:guava:27.0.1-jre",
+            "com.google.protobuf:protobuf-java:3.8.0",
             "com.puppycrawl.tools:checkstyle:8.15",
             "com.sun.xml.bind:jaxb-core:2.3.0",
             "com.sun.xml.bind:jaxb-impl:2.3.0",

java/modify_generated_java_flatbuffers_files.py

@@ -4,7 +4,6 @@ from __future__ import print_function
 
 import os
 import sys
-
 """
 This script is used for modifying the generated java flatbuffer
 files for the reason: The package declaration in Java is different
@@ -21,19 +20,18 @@ RAY_HOME: The root directory of Ray project.
 PACKAGE_DECLARATION = "package org.ray.runtime.generated;"
 
 
-def add_new_line(file, line_num, text):
+def add_package(file):
     with open(file, "r") as file_handler:
         lines = file_handler.readlines()
-        if (line_num <= 0) or (line_num > len(lines) + 1):
-            return False
 
-    lines.insert(line_num - 1, text + os.linesep)
+    if "FlatBuffers" not in lines[0]:
+        return
+
+    lines.insert(1, PACKAGE_DECLARATION + os.linesep)
     with open(file, "w") as file_handler:
         for line in lines:
             file_handler.write(line)
 
-    return True
-
 
 def add_package_declarations(generated_root_path):
     file_names = os.listdir(generated_root_path)
@@ -41,15 +39,11 @@ def add_package_declarations(generated_root_path):
         if not file_name.endswith(".java"):
             continue
         full_name = os.path.join(generated_root_path, file_name)
-        success = add_new_line(full_name, 2, PACKAGE_DECLARATION)
-        if not success:
-            raise RuntimeError("Failed to add package declarations, "
-                               "file name is %s" % full_name)
+        add_package(full_name)
 
 
 if __name__ == "__main__":
     ray_home = sys.argv[1]
     root_path = os.path.join(
-        ray_home,
-        "java/runtime/src/main/java/org/ray/runtime/generated")
+        ray_home, "java/runtime/src/main/java/org/ray/runtime/generated")
     add_package_declarations(root_path)

java/runtime/pom.xml

@@ -41,6 +41,11 @@
   <artifactId>guava</artifactId>
   <version>27.0.1-jre</version>
 </dependency>
+<dependency>
+  <groupId>com.google.protobuf</groupId>
+  <artifactId>protobuf-java</artifactId>
+  <version>3.8.0</version>
+</dependency>
 <dependency>
   <groupId>com.typesafe</groupId>
   <artifactId>config</artifactId>

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

@@ -1,7 +1,7 @@
 package org.ray.runtime.gcs;
 
 import com.google.common.base.Preconditions;
-import java.nio.ByteBuffer;
+import com.google.protobuf.InvalidProtocolBufferException;
 import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.List;
@@ -13,10 +13,10 @@ 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.ActorCheckpointIdData;
-import org.ray.runtime.generated.ClientTableData;
-import org.ray.runtime.generated.EntryType;
-import org.ray.runtime.generated.TablePrefix;
+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;
 import org.slf4j.LoggerFactory;
@@ -51,7 +51,7 @@ public class GcsClient {
   }
 
   public List<NodeInfo> getAllNodeInfo() {
-    final String prefix = TablePrefix.name(TablePrefix.CLIENT);
+    final String prefix = TablePrefix.CLIENT.toString();
     final byte[] key = ArrayUtils.addAll(prefix.getBytes(), UniqueId.NIL.getBytes());
     List<byte[]> results = primary.lrange(key, 0, -1);
 
@@ -63,36 +63,42 @@ public class GcsClient {
     Map<UniqueId, NodeInfo> clients = new HashMap<>();
     for (byte[] result : results) {
       Preconditions.checkNotNull(result);
-      ClientTableData data = ClientTableData.getRootAsClientTableData(ByteBuffer.wrap(result));
-      final UniqueId clientId = UniqueId.fromByteBuffer(data.clientIdAsByteBuffer());
+      ClientTableData data = null;
+      try {
+        data = ClientTableData.parseFrom(result);
+      } catch (InvalidProtocolBufferException e) {
+        throw new RuntimeException("Received invalid protobuf data from GCS.");
+      }
+      final UniqueId clientId = UniqueId
+          .fromByteBuffer(data.getClientId().asReadOnlyByteBuffer());
 
-      if (data.entryType() == EntryType.INSERTION) {
+      if (data.getEntryType() == EntryType.INSERTION) {
         //Code path of node insertion.
         Map<String, Double> resources = new HashMap<>();
         // Compute resources.
         Preconditions.checkState(
-            data.resourcesTotalLabelLength() == data.resourcesTotalCapacityLength());
-        for (int i = 0; i < data.resourcesTotalLabelLength(); i++) {
-          resources.put(data.resourcesTotalLabel(i), data.resourcesTotalCapacity(i));
+            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.nodeManagerAddress(), true, resources);
+            clientId, data.getNodeManagerAddress(), true, resources);
         clients.put(clientId, nodeInfo);
-      } else if (data.entryType() == EntryType.RES_CREATEUPDATE) {
+      } else if (data.getEntryType() == EntryType.RES_CREATEUPDATE) {
         Preconditions.checkState(clients.containsKey(clientId));
         NodeInfo nodeInfo = clients.get(clientId);
-        for (int i = 0; i < data.resourcesTotalLabelLength(); i++) {
-          nodeInfo.resources.put(data.resourcesTotalLabel(i), data.resourcesTotalCapacity(i));
+        for (int i = 0; i < data.getResourcesTotalLabelCount(); i++) {
+          nodeInfo.resources.put(data.getResourcesTotalLabel(i), data.getResourcesTotalCapacity(i));
         }
-      } else if (data.entryType() == EntryType.RES_DELETE) {
+      } else if (data.getEntryType() == EntryType.RES_DELETE) {
         Preconditions.checkState(clients.containsKey(clientId));
         NodeInfo nodeInfo = clients.get(clientId);
-        for (int i = 0; i < data.resourcesTotalLabelLength(); i++) {
-          nodeInfo.resources.remove(data.resourcesTotalLabel(i));
+        for (int i = 0; i < data.getResourcesTotalLabelCount(); i++) {
+          nodeInfo.resources.remove(data.getResourcesTotalLabel(i));
         }
       } else {
         // Code path of node deletion.
-        Preconditions.checkState(data.entryType() == EntryType.DELETION);
+        Preconditions.checkState(data.getEntryType() == EntryType.DELETION);
         NodeInfo nodeInfo = new NodeInfo(clientId, clients.get(clientId).nodeAddress,
             false, clients.get(clientId).resources);
         clients.put(clientId, nodeInfo);
@@ -107,7 +113,7 @@ public class GcsClient {
    */
   public boolean actorExists(UniqueId actorId) {
     byte[] key = ArrayUtils.addAll(
-        TablePrefix.name(TablePrefix.ACTOR).getBytes(), actorId.getBytes());
+        TablePrefix.ACTOR.toString().getBytes(), actorId.getBytes());
     return primary.exists(key);
   }
 
@@ -115,7 +121,7 @@ public class GcsClient {
    * Query whether the raylet task exists in Gcs.
    */
   public boolean rayletTaskExistsInGcs(TaskId taskId) {
-    byte[] key = ArrayUtils.addAll(TablePrefix.name(TablePrefix.RAYLET_TASK).getBytes(),
+    byte[] key = ArrayUtils.addAll(TablePrefix.RAYLET_TASK.toString().getBytes(),
         taskId.getBytes());
     RedisClient client = getShardClient(taskId);
     return client.exists(key);
@@ -126,19 +132,26 @@ public class GcsClient {
    */
   public List<Checkpoint> getCheckpointsForActor(UniqueId actorId) {
     List<Checkpoint> checkpoints = new ArrayList<>();
-    final String prefix = TablePrefix.name(TablePrefix.ACTOR_CHECKPOINT_ID);
+    final String prefix = TablePrefix.ACTOR_CHECKPOINT_ID.toString();
     final byte[] key = ArrayUtils.addAll(prefix.getBytes(), actorId.getBytes());
     RedisClient client = getShardClient(actorId);
 
     byte[] result = client.get(key);
     if (result != null) {
-      ActorCheckpointIdData data =
-          ActorCheckpointIdData.getRootAsActorCheckpointIdData(ByteBuffer.wrap(result));
-      UniqueId[] checkpointIds = IdUtil.getUniqueIdsFromByteBuffer(
-          data.checkpointIdsAsByteBuffer());
+      ActorCheckpointIdData data = null;
+      try {
+        data = ActorCheckpointIdData.parseFrom(result);
+      } catch (InvalidProtocolBufferException e) {
+        throw new RuntimeException("Received invalid protobuf data from GCS.");
+      }
+      UniqueId[] checkpointIds = new UniqueId[data.getCheckpointIdsCount()];
+      for (int i = 0; i < checkpointIds.length; i++) {
+        checkpointIds[i] = UniqueId
+            .fromByteBuffer(data.getCheckpointIds(i).asReadOnlyByteBuffer());
+      }
 
       for (int i = 0; i < checkpointIds.length; i++) {
-        checkpoints.add(new Checkpoint(checkpointIds[i], data.timestamps(i)));
+        checkpoints.add(new Checkpoint(checkpointIds[i], data.getTimestamps(i)));
       }
     }
     checkpoints.sort((x, y) -> Long.compare(y.timestamp, x.timestamp));

java/runtime/src/main/java/org/ray/runtime/objectstore/ObjectStoreProxy.java

@@ -16,7 +16,7 @@ import org.ray.api.id.ObjectId;
 import org.ray.runtime.AbstractRayRuntime;
 import org.ray.runtime.RayDevRuntime;
 import org.ray.runtime.config.RunMode;
-import org.ray.runtime.generated.ErrorType;
+import org.ray.runtime.generated.Gcs.ErrorType;
 import org.ray.runtime.util.IdUtil;
 import org.ray.runtime.util.Serializer;
 import org.slf4j.Logger;
@@ -29,12 +29,12 @@ public class ObjectStoreProxy {
 
   private static final Logger LOGGER = LoggerFactory.getLogger(ObjectStoreProxy.class);
 
-  private static final byte[] WORKER_EXCEPTION_META = String.valueOf(ErrorType.WORKER_DIED)
-      .getBytes();
-  private static final byte[] ACTOR_EXCEPTION_META = String.valueOf(ErrorType.ACTOR_DIED)
-      .getBytes();
+  private static final byte[] WORKER_EXCEPTION_META = String
+      .valueOf(ErrorType.WORKER_DIED.getNumber()).getBytes();
+  private static final byte[] ACTOR_EXCEPTION_META = String
+      .valueOf(ErrorType.ACTOR_DIED.getNumber()).getBytes();
   private static final byte[] UNRECONSTRUCTABLE_EXCEPTION_META = String
-      .valueOf(ErrorType.OBJECT_UNRECONSTRUCTABLE).getBytes();
+      .valueOf(ErrorType.OBJECT_UNRECONSTRUCTABLE.getNumber()).getBytes();
 
   private static final byte[] RAW_TYPE_META = "RAW".getBytes();