Refactor ID Serial 1: Separate ObjectID and TaskID from UniqueID (#4776)

* Enable BaseId.

* Change TaskID and make python test pass

* Remove unnecessary functions and fix test failure and change TaskID to
16 bytes.

* Java code change draft

* Refine

* Lint

* Update java/api/src/main/java/org/ray/api/id/TaskId.java

Co-Authored-By: Hao Chen <chenh1024@gmail.com>

* Update java/api/src/main/java/org/ray/api/id/BaseId.java

Co-Authored-By: Hao Chen <chenh1024@gmail.com>

* Update java/api/src/main/java/org/ray/api/id/BaseId.java

Co-Authored-By: Hao Chen <chenh1024@gmail.com>

* Update java/api/src/main/java/org/ray/api/id/ObjectId.java

Co-Authored-By: Hao Chen <chenh1024@gmail.com>

* Address comment

* Lint

* Fix SINGLE_PROCESS

* Fix comments

* Refine code

* Refine test

* Resolve conflict

java/api/src/main/java/org/ray/api/Ray.java

@@ -1,6 +1,7 @@
 package org.ray.api;
 
 import java.util.List;
+import org.ray.api.id.ObjectId;
 import org.ray.api.id.UniqueId;
 import org.ray.api.runtime.RayRuntime;
 import org.ray.api.runtime.RayRuntimeFactory;
@@ -65,7 +66,7 @@ public final class Ray extends RayCall {
    * @param objectId The ID of the object to get.
    * @return The Java object.
    */
-  public static <T> T get(UniqueId objectId) {
+  public static <T> T get(ObjectId objectId) {
     return runtime.get(objectId);
   }
 
@@ -75,7 +76,7 @@ public final class Ray extends RayCall {
    * @param objectIds The list of object IDs.
    * @return A list of Java objects.
    */
-  public static <T> List<T> get(List<UniqueId> objectIds) {
+  public static <T> List<T> get(List<ObjectId> objectIds) {
     return runtime.get(objectIds);
   }
 

java/api/src/main/java/org/ray/api/RayObject.java

@@ -1,6 +1,6 @@
 package org.ray.api;
 
-import org.ray.api.id.UniqueId;
+import org.ray.api.id.ObjectId;
 
 /**
  * Represents an object in the object store.
@@ -17,7 +17,7 @@ public interface RayObject<T> {
   /**
    * Get the object id.
    */
-  UniqueId getId();
+  ObjectId getId();
 
 }
 

java/api/src/main/java/org/ray/api/exception/UnreconstructableException.java

@@ -1,6 +1,6 @@
 package org.ray.api.exception;
 
-import org.ray.api.id.UniqueId;
+import org.ray.api.id.ObjectId;
 
 /**
  * Indicates that an object is lost (either evicted or explicitly deleted) and cannot be
@@ -11,9 +11,9 @@ import org.ray.api.id.UniqueId;
  */
 public class UnreconstructableException extends RayException {
 
-  public final UniqueId objectId;
+  public final ObjectId objectId;
 
-  public UnreconstructableException(UniqueId objectId) {
+  public UnreconstructableException(ObjectId objectId) {
     super(String.format(
         "Object %s is lost (either evicted or explicitly deleted) and cannot be reconstructed.",
         objectId));

java/api/src/main/java/org/ray/api/id/BaseId.java

@@ -0,0 +1,99 @@
+package org.ray.api.id;
+
+import java.io.Serializable;
+import java.nio.ByteBuffer;
+import java.util.Arrays;
+import javax.xml.bind.DatatypeConverter;
+
+public abstract class BaseId implements Serializable {
+  private static final long serialVersionUID = 8588849129675565761L;
+  private final byte[] id;
+  private int hashCodeCache = 0;
+  private Boolean isNilCache = null;
+
+  /**
+   * Create a BaseId instance according to the input byte array.
+   */
+  public BaseId(byte[] id) {
+    if (id.length != size()) {
+      throw new IllegalArgumentException("Failed to construct BaseId, expect " + size()
+              + " bytes, but got " + id.length + " bytes.");
+    }
+    this.id = id;
+  }
+
+  /**
+   * Get the byte data of this id.
+   */
+  public byte[] getBytes() {
+    return id;
+  }
+
+  /**
+   * Convert the byte data to a ByteBuffer.
+   */
+  public ByteBuffer toByteBuffer() {
+    return ByteBuffer.wrap(id);
+  }
+
+  /**
+   * @return True if this id is nil.
+   */
+  public boolean isNil() {
+    if (isNilCache == null) {
+      isNilCache = true;
+      for (int i = 0; i < size(); ++i) {
+        if (id[i] != (byte) 0xff) {
+          isNilCache = false;
+          break;
+        }
+      }
+    }
+    return isNilCache;
+  }
+
+  /**
+   * Derived class should implement this function.
+   * @return The length of this id in bytes.
+   */
+  public abstract int size();
+
+  @Override
+  public int hashCode() {
+    // Lazy evaluation.
+    if (hashCodeCache == 0) {
+      hashCodeCache = Arrays.hashCode(id);
+    }
+    return hashCodeCache;
+  }
+
+  @Override
+  public boolean equals(Object obj) {
+    if (obj == null) {
+      return false;
+    }
+
+    if (!this.getClass().equals(obj.getClass())) {
+      return false;
+    }
+
+    BaseId r = (BaseId) obj;
+    return Arrays.equals(id, r.id);
+  }
+
+  @Override
+  public String toString() {
+    return DatatypeConverter.printHexBinary(id).toLowerCase();
+  }
+
+  protected static byte[] hexString2Bytes(String hex) {
+    return DatatypeConverter.parseHexBinary(hex);
+  }
+
+  protected static byte[] byteBuffer2Bytes(ByteBuffer bb) {
+    byte[] id = new byte[bb.remaining()];
+    bb.get(id);
+    return id;
+  }
+
+}

java/api/src/main/java/org/ray/api/id/ObjectId.java

@@ -0,0 +1,62 @@
+package org.ray.api.id;
+
+import java.io.Serializable;
+import java.nio.ByteBuffer;
+import java.util.Arrays;
+import java.util.Random;
+
+/**
+ * Represents the id of a Ray object.
+ */
+public class ObjectId extends BaseId implements Serializable {
+
+  public static final int LENGTH = 20;
+  public static final ObjectId NIL = genNil();
+
+  /**
+   * Create an ObjectId from a hex string.
+   */
+  public static ObjectId fromHexString(String hex) {
+    return new ObjectId(hexString2Bytes(hex));
+  }
+
+  /**
+   * Create an ObjectId from a ByteBuffer.
+   */
+  public static ObjectId fromByteBuffer(ByteBuffer bb) {
+    return new ObjectId(byteBuffer2Bytes(bb));
+  }
+
+  /**
+   * Generate a nil ObjectId.
+   */
+  private static ObjectId genNil() {
+    byte[] b = new byte[LENGTH];
+    Arrays.fill(b, (byte) 0xFF);
+    return new ObjectId(b);
+  }
+
+  /**
+   * Generate an ObjectId with random value.
+   */
+  public static ObjectId randomId() {
+    byte[] b = new byte[LENGTH];
+    new Random().nextBytes(b);
+    return new ObjectId(b);
+  }
+
+  public ObjectId(byte[] id) {
+    super(id);
+  }
+
+  @Override
+  public int size() {
+    return LENGTH;
+  }
+
+  public TaskId getTaskId() {
+    byte[] taskIdBytes = Arrays.copyOf(getBytes(), TaskId.LENGTH);
+    return new TaskId(taskIdBytes);
+  }
+
+}

java/api/src/main/java/org/ray/api/id/TaskId.java

@@ -0,0 +1,56 @@
+package org.ray.api.id;
+
+import java.io.Serializable;
+import java.nio.ByteBuffer;
+import java.util.Arrays;
+import java.util.Random;
+
+/**
+ * Represents the id of a Ray task.
+ */
+public class TaskId extends BaseId implements Serializable {
+
+  public static final int LENGTH = 16;
+  public static final TaskId NIL = genNil();
+
+  /**
+   * Create a TaskId from a hex string.
+   */
+  public static TaskId fromHexString(String hex) {
+    return new TaskId(hexString2Bytes(hex));
+  }
+
+  /**
+   * Creates a TaskId from a ByteBuffer.
+   */
+  public static TaskId fromByteBuffer(ByteBuffer bb) {
+    return new TaskId(byteBuffer2Bytes(bb));
+  }
+
+  /**
+   * Generate a nil TaskId.
+   */
+  private static TaskId genNil() {
+    byte[] b = new byte[LENGTH];
+    Arrays.fill(b, (byte) 0xFF);
+    return new TaskId(b);
+  }
+
+  /**
+   * Generate an TaskId with random value.
+   */
+  public static TaskId randomId() {
+    byte[] b = new byte[LENGTH];
+    new Random().nextBytes(b);
+    return new TaskId(b);
+  }
+
+  public TaskId(byte[] id) {
+    super(id);
+  }
+
+  @Override
+  public int size() {
+    return LENGTH;
+  }
+}

java/api/src/main/java/org/ray/api/id/UniqueId.java

@@ -4,41 +4,34 @@ import java.io.Serializable;
 import java.nio.ByteBuffer;
 import java.util.Arrays;
 import java.util.Random;
-import javax.xml.bind.DatatypeConverter;
 
 /**
  * Represents a unique id of all Ray concepts, including
- * objects, tasks, workers, actors, etc.
+ * workers, actors, checkpoints, etc.
  */
-public class UniqueId implements Serializable {
+public class UniqueId extends BaseId implements Serializable {
 
   public static final int LENGTH = 20;
   public static final UniqueId NIL = genNil();
-  private static final long serialVersionUID = 8588849129675565761L;
-  private final byte[] id;
 
   /**
    * Create a UniqueId from a hex string.
    */
   public static UniqueId fromHexString(String hex) {
-    byte[] bytes = DatatypeConverter.parseHexBinary(hex);
-    return new UniqueId(bytes);
+    return new UniqueId(hexString2Bytes(hex));
   }
 
   /**
    * Creates a UniqueId from a ByteBuffer.
    */
   public static UniqueId fromByteBuffer(ByteBuffer bb) {
-    byte[] id = new byte[bb.remaining()];
-    bb.get(id);
-
-    return new UniqueId(id);
+    return new UniqueId(byteBuffer2Bytes(bb));
   }
 
   /**
    * Generate a nil UniqueId.
    */
-  public static UniqueId genNil() {
+  private static UniqueId genNil() {
     byte[] b = new byte[LENGTH];
     Arrays.fill(b, (byte) 0xFF);
     return new UniqueId(b);
@@ -54,64 +47,11 @@ public class UniqueId implements Serializable {
   }
 
   public UniqueId(byte[] id) {
-    if (id.length != LENGTH) {
-      throw new IllegalArgumentException("Illegal argument for UniqueId, expect " + LENGTH
-          + " bytes, but got " + id.length + " bytes.");
-    }
-
-    this.id = id;
-  }
-
-  /**
-   * Get the byte data of this UniqueId.
-   */
-  public byte[] getBytes() {
-    return id;
-  }
-
-  /**
-   * Convert the byte data to a ByteBuffer.
-   */
-  public ByteBuffer toByteBuffer() {
-    return ByteBuffer.wrap(id);
-  }
-
-  /**
-   * Create a copy of this UniqueId.
-   */
-  public UniqueId copy() {
-    byte[] nid = Arrays.copyOf(id, id.length);
-    return new UniqueId(nid);
-  }
-
-  /**
-   * Returns true if this id is nil.
-   */
-  public boolean isNil() {
-    return this.equals(NIL);
+    super(id);
   }
 
   @Override
-  public int hashCode() {
-    return Arrays.hashCode(id);
-  }
-
-  @Override
-  public boolean equals(Object obj) {
-    if (obj == null) {
-      return false;
-    }
-
-    if (!(obj instanceof UniqueId)) {
-      return false;
-    }
-
-    UniqueId r = (UniqueId) obj;
-    return Arrays.equals(id, r.id);
-  }
-
-  @Override
-  public String toString() {
-    return DatatypeConverter.printHexBinary(id).toLowerCase();
+  public int size() {
+    return LENGTH;
   }
 }

java/api/src/main/java/org/ray/api/runtime/RayRuntime.java

@@ -6,6 +6,7 @@ import org.ray.api.RayObject;
 import org.ray.api.RayPyActor;
 import org.ray.api.WaitResult;
 import org.ray.api.function.RayFunc;
+import org.ray.api.id.ObjectId;
 import org.ray.api.id.UniqueId;
 import org.ray.api.options.ActorCreationOptions;
 import org.ray.api.options.CallOptions;
@@ -35,7 +36,7 @@ public interface RayRuntime {
    * @param objectId The ID of the object to get.
    * @return The Java object.
    */
-  <T> T get(UniqueId objectId);
+  <T> T get(ObjectId objectId);
 
   /**
    * Get a list of objects from the object store.
@@ -43,7 +44,7 @@ public interface RayRuntime {
    * @param objectIds The list of object IDs.
    * @return A list of Java objects.
    */
-  <T> List<T> get(List<UniqueId> objectIds);
+  <T> List<T> get(List<ObjectId> objectIds);
 
   /**
    * Wait for a list of RayObjects to be locally available, until specified number of objects are
@@ -63,7 +64,7 @@ public interface RayRuntime {
    * @param localOnly Whether only free objects for local object store or not.
    * @param deleteCreatingTasks Whether also delete objects' creating tasks from GCS.
    */
-  void free(List<UniqueId> objectIds, boolean localOnly, boolean deleteCreatingTasks);
+  void free(List<ObjectId> objectIds, boolean localOnly, boolean deleteCreatingTasks);
 
   /**
    * Set the resource for the specific node.