[ID Refactor] Refactor ActorID, TaskID and ObjectID (#5286)

* Refactor ActorID, TaskID on the Java side.

Left a TODO comment

WIP for ObjectID

ADD test

Fix

Add java part

Fix Java test

Fix

Refine test.

Enable test in CI

* Extra a helper function.

* Resolve TODOs

* Fix Python CI

* Fix Java lint

* Update .travis.yml

Co-Authored-By: Stephanie Wang <swang@cs.berkeley.edu>

* Address some comments.

Address some comments.

Add id_specification.rst

Reanme id_specification.rst to id_specification.md

typo

Address zhijun's comments.

Fix test

Address comments.

Fix lint

Address comments

* Fix test

* Address comments.

* Fix build error

* Update src/ray/design_docs/id_specification.md

Co-Authored-By: Stephanie Wang <swang@cs.berkeley.edu>

* Update src/ray/design_docs/id_specification.md

Co-Authored-By: Stephanie Wang <swang@cs.berkeley.edu>

* Update src/ray/design_docs/id_specification.md

Co-Authored-By: Stephanie Wang <swang@cs.berkeley.edu>

* Update src/ray/design_docs/id_specification.md

Co-Authored-By: Stephanie Wang <swang@cs.berkeley.edu>

* Update src/ray/design_docs/id_specification.md

Co-Authored-By: Stephanie Wang <swang@cs.berkeley.edu>

* Address comments

* Update src/ray/common/id.h

Co-Authored-By: Stephanie Wang <swang@cs.berkeley.edu>

* Update src/ray/common/id.h

Co-Authored-By: Stephanie Wang <swang@cs.berkeley.edu>

* Update src/ray/common/id.h

Co-Authored-By: Stephanie Wang <swang@cs.berkeley.edu>

* Update src/ray/design_docs/id_specification.md

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

* Update src/ray/design_docs/id_specification.md

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

* Address comments.

* Address comments.

* Address comments.

* Update C++ part to make sure task id is generated determantic

* WIP

* Fix core worker

* Fix Java part

* Fix comments.

* Add Python side

* Fix python

* Address comments

* Fix linting

* Fix

* Fix C++ linting

* Add JobId() method to TaskID

* Fix linting

* Update src/ray/common/id.h

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

* 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/TaskId.java

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

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

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

* Address comments

* Add DriverTaskId embeding job id

* Fix tests

* Add python dor_fake_driver_id

* Address comments and fix linting

* Fix CI

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

@@ -1,6 +1,8 @@
 package org.ray.api;
 
 import java.util.List;
+
+import org.ray.api.id.ActorId;
 import org.ray.api.id.UniqueId;
 
 public interface Checkpointable {
@@ -10,7 +12,7 @@ public interface Checkpointable {
     /**
      * Actor's ID.
      */
-    public final UniqueId actorId;
+    public final ActorId actorId;
     /**
      * Number of tasks executed since last checkpoint.
      */
@@ -20,8 +22,8 @@ public interface Checkpointable {
      */
     public final long timeElapsedMsSinceLastCheckpoint;
 
-    public CheckpointContext(UniqueId actorId, int numTasksSinceLastCheckpoint,
-        long timeElapsedMsSinceLastCheckpoint) {
+    public CheckpointContext(ActorId actorId, int numTasksSinceLastCheckpoint,
+                             long timeElapsedMsSinceLastCheckpoint) {
       this.actorId = actorId;
       this.numTasksSinceLastCheckpoint = numTasksSinceLastCheckpoint;
       this.timeElapsedMsSinceLastCheckpoint = timeElapsedMsSinceLastCheckpoint;
@@ -67,7 +69,7 @@ public interface Checkpointable {
    * @param checkpointId An ID that represents this actor's current state in GCS. You should
    *     save this checkpoint ID together with actor's checkpoint data.
    */
-  void saveCheckpoint(UniqueId actorId, UniqueId checkpointId);
+  void saveCheckpoint(ActorId actorId, UniqueId checkpointId);
 
   /**
    * Load actor's previous checkpoint, and restore actor's state.
@@ -83,7 +85,7 @@ public interface Checkpointable {
    * @return The ID of the checkpoint from which the actor was resumed, or null if the actor should
    *     restart from the beginning.
    */
-  UniqueId loadCheckpoint(UniqueId actorId, List<Checkpoint> availableCheckpoints);
+  UniqueId loadCheckpoint(ActorId actorId, List<Checkpoint> availableCheckpoints);
 
   /**
    * Delete an expired checkpoint;
@@ -95,5 +97,5 @@ public interface Checkpointable {
    * @param actorId ID of the actor.
    * @param checkpointId ID of the checkpoint that has expired.
    */
-  void checkpointExpired(UniqueId actorId, UniqueId checkpointId);
+  void checkpointExpired(ActorId actorId, UniqueId checkpointId);
 }

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

@@ -0,0 +1,6 @@
+package org.ray.api;
+
+public enum ObjectType {
+  PUT_OBJECT,
+  RETURN_OBJECT,
+}

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

@@ -1,5 +1,6 @@
 package org.ray.api;
 
+import org.ray.api.id.ActorId;
 import org.ray.api.id.UniqueId;
 
 /**
@@ -12,7 +13,7 @@ public interface RayActor<T> {
   /**
    * @return The id of this actor.
    */
-  UniqueId getId();
+  ActorId getId();
 
   /**
    * @return The id of this actor handle.

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

@@ -0,0 +1,54 @@
+package org.ray.api.id;
+
+import java.io.Serializable;
+import java.nio.ByteBuffer;
+import java.nio.ByteOrder;
+import java.util.Arrays;
+import java.util.Random;
+
+public class ActorId extends BaseId implements Serializable {
+  private static final int UNIQUE_BYTES_LENGTH = 4;
+
+  public static final int LENGTH = UNIQUE_BYTES_LENGTH + JobId.LENGTH;
+
+  public static final ActorId NIL = nil();
+
+  private ActorId(byte[] id) {
+    super(id);
+  }
+
+  public static ActorId fromByteBuffer(ByteBuffer bb) {
+    return new ActorId(byteBuffer2Bytes(bb));
+  }
+
+  public static ActorId fromBytes(byte[] bytes) {
+    return new ActorId(bytes);
+  }
+
+  public static ActorId generateActorId(JobId jobId) {
+    byte[] uniqueBytes = new byte[ActorId.UNIQUE_BYTES_LENGTH];
+    new Random().nextBytes(uniqueBytes);
+
+    byte[] bytes = new byte[ActorId.LENGTH];
+    ByteBuffer wbb = ByteBuffer.wrap(bytes);
+    wbb.order(ByteOrder.LITTLE_ENDIAN);
+
+    System.arraycopy(uniqueBytes, 0, bytes, 0, ActorId.UNIQUE_BYTES_LENGTH);
+    System.arraycopy(jobId.getBytes(), 0, bytes, ActorId.UNIQUE_BYTES_LENGTH, JobId.LENGTH);
+    return new ActorId(bytes);
+  }
+
+  /**
+   * Generate a nil ActorId.
+   */
+  private static ActorId nil() {
+    byte[] b = new byte[LENGTH];
+    Arrays.fill(b, (byte) 0xFF);
+    return new ActorId(b);
+  }
+
+  @Override
+  public int size() {
+    return LENGTH;
+  }
+}

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

@@ -14,7 +14,7 @@ public abstract class BaseId implements Serializable {
   /**
    * Create a BaseId instance according to the input byte array.
    */
-  public BaseId(byte[] id) {
+  protected BaseId(byte[] id) {
     if (id.length != size()) {
       throw new IllegalArgumentException("Failed to construct BaseId, expect " + size()
               + " bytes, but got " + id.length + " bytes.");

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

@@ -2,8 +2,10 @@ package org.ray.api.id;
 
 import java.io.Serializable;
 import java.nio.ByteBuffer;
+import java.nio.ByteOrder;
 import java.util.Arrays;
 import java.util.Random;
+import org.ray.api.ObjectType;
 
 /**
  * Represents the id of a Ray object.
@@ -11,14 +13,22 @@ import java.util.Random;
 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));
-  }
+  private static int CREATED_BY_TASK_FLAG_BITS_OFFSET = 15;
+
+  private static int OBJECT_TYPE_FLAG_BITS_OFFSET = 14;
+
+  private static int TRANSPORT_TYPE_FLAG_BITS_OFFSET = 11;
+
+  private static int FLAGS_BYTES_POS = TaskId.LENGTH;
+
+  private static int FLAGS_BYTES_LENGTH = 2;
+
+  private static int INDEX_BYTES_POS = FLAGS_BYTES_POS + FLAGS_BYTES_LENGTH;
+
+  private static int INDEX_BYTES_LENGTH = 4;
 
   /**
    * Create an ObjectId from a ByteBuffer.
@@ -39,12 +49,54 @@ public class ObjectId extends BaseId implements Serializable {
   /**
    * Generate an ObjectId with random value.
    */
-  public static ObjectId randomId() {
+  public static ObjectId fromRandom() {
     byte[] b = new byte[LENGTH];
     new Random().nextBytes(b);
     return new ObjectId(b);
   }
 
+  /**
+   * Compute the object ID of an object put by the task.
+   */
+  public static ObjectId forPut(TaskId taskId, int putIndex) {
+    short flags = 0;
+    flags = setCreatedByTaskFlag(flags, true);
+    // Set a default transport type with value 0.
+    flags = (short) (flags | (0x0 << TRANSPORT_TYPE_FLAG_BITS_OFFSET));
+    flags = setObjectTypeFlag(flags, ObjectType.PUT_OBJECT);
+
+    byte[] bytes = new byte[ObjectId.LENGTH];
+    System.arraycopy(taskId.getBytes(), 0, bytes, 0, TaskId.LENGTH);
+
+    ByteBuffer wbb = ByteBuffer.wrap(bytes);
+    wbb.order(ByteOrder.LITTLE_ENDIAN);
+    wbb.putShort(FLAGS_BYTES_POS, flags);
+
+    wbb.putInt(INDEX_BYTES_POS, putIndex);
+    return new ObjectId(bytes);
+  }
+
+  /**
+   * Compute the object ID of an object return by the task.
+   */
+  public static ObjectId forReturn(TaskId taskId, int returnIndex) {
+    short flags = 0;
+    flags = setCreatedByTaskFlag(flags, true);
+    // Set a default transport type with value 0.
+    flags = (short) (flags | (0x0 << TRANSPORT_TYPE_FLAG_BITS_OFFSET));
+    flags = setObjectTypeFlag(flags, ObjectType.RETURN_OBJECT);
+
+    byte[] bytes = new byte[ObjectId.LENGTH];
+    System.arraycopy(taskId.getBytes(), 0, bytes, 0, TaskId.LENGTH);
+
+    ByteBuffer wbb = ByteBuffer.wrap(bytes);
+    wbb.order(ByteOrder.LITTLE_ENDIAN);
+    wbb.putShort(FLAGS_BYTES_POS, flags);
+
+    wbb.putInt(INDEX_BYTES_POS, returnIndex);
+    return new ObjectId(bytes);
+  }
+
   public ObjectId(byte[] id) {
     super(id);
   }
@@ -56,7 +108,23 @@ public class ObjectId extends BaseId implements Serializable {
 
   public TaskId getTaskId() {
     byte[] taskIdBytes = Arrays.copyOf(getBytes(), TaskId.LENGTH);
-    return new TaskId(taskIdBytes);
+    return TaskId.fromBytes(taskIdBytes);
+  }
+
+  private static short setCreatedByTaskFlag(short flags, boolean createdByTask) {
+    if (createdByTask) {
+      return (short) (flags | (0x1 << CREATED_BY_TASK_FLAG_BITS_OFFSET));
+    } else {
+      return (short) (flags | (0x0 << CREATED_BY_TASK_FLAG_BITS_OFFSET));
+    }
+  }
+
+  private static short setObjectTypeFlag(short flags, ObjectType objectType) {
+    if (objectType == ObjectType.RETURN_OBJECT) {
+      return (short)(flags | (0x1 << OBJECT_TYPE_FLAG_BITS_OFFSET));
+    } else {
+      return (short)(flags | (0x0 << OBJECT_TYPE_FLAG_BITS_OFFSET));
+    }
   }
 
 }

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

@@ -2,6 +2,7 @@ package org.ray.api.id;
 
 import java.io.Serializable;
 import java.nio.ByteBuffer;
+import java.nio.ByteOrder;
 import java.util.Arrays;
 import java.util.Random;
 
@@ -10,7 +11,10 @@ import java.util.Random;
  */
 public class TaskId extends BaseId implements Serializable {
 
-  public static final int LENGTH = 16;
+  private static final int UNIQUE_BYTES_LENGTH = 6;
+
+  public static final int LENGTH = UNIQUE_BYTES_LENGTH + ActorId.LENGTH;
+
   public static final TaskId NIL = genNil();
 
   /**
@@ -27,6 +31,22 @@ public class TaskId extends BaseId implements Serializable {
     return new TaskId(byteBuffer2Bytes(bb));
   }
 
+  /**
+   * Creates a TaskId from given bytes.
+   */
+  public static TaskId fromBytes(byte[] bytes) {
+    return new TaskId(bytes);
+  }
+
+  /**
+   * Get the id of the actor to which this task belongs
+   */
+  public ActorId getActorId() {
+    byte[] actorIdBytes = new byte[ActorId.LENGTH];
+    System.arraycopy(getBytes(), UNIQUE_BYTES_LENGTH, actorIdBytes, 0, ActorId.LENGTH);
+    return ActorId.fromByteBuffer(ByteBuffer.wrap(actorIdBytes));
+  }
+
   /**
    * Generate a nil TaskId.
    */
@@ -36,16 +56,7 @@ public class TaskId extends BaseId implements Serializable {
     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) {
+  private TaskId(byte[] id) {
     super(id);
   }
 

java/api/src/main/java/org/ray/api/runtimecontext/RuntimeContext.java

@@ -1,8 +1,8 @@
 package org.ray.api.runtimecontext;
 
 import java.util.List;
+import org.ray.api.id.ActorId;
 import org.ray.api.id.JobId;
-import org.ray.api.id.UniqueId;
 
 /**
  * A class used for getting information of Ray runtime.
@@ -19,7 +19,7 @@ public interface RuntimeContext {
    *
    * Note, this can only be called in actors.
    */
-  UniqueId getCurrentActorId();
+  ActorId getCurrentActorId();
 
   /**
    * Returns true if the current actor was reconstructed, false if it's created for the first time.