[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/runtime/src/main/java/org/ray/runtime/AbstractRayRuntime.java

@@ -19,6 +19,8 @@ import org.ray.api.RayPyActor;
 import org.ray.api.WaitResult;
 import org.ray.api.exception.RayException;
 import org.ray.api.function.RayFunc;
+import org.ray.api.id.ActorId;
+import org.ray.api.id.JobId;
 import org.ray.api.id.ObjectId;
 import org.ray.api.id.TaskId;
 import org.ray.api.id.UniqueId;
@@ -34,10 +36,10 @@ import org.ray.runtime.functionmanager.PyFunctionDescriptor;
 import org.ray.runtime.gcs.GcsClient;
 import org.ray.runtime.objectstore.ObjectStoreProxy;
 import org.ray.runtime.raylet.RayletClient;
+import org.ray.runtime.raylet.RayletClientImpl;
 import org.ray.runtime.task.ArgumentsBuilder;
 import org.ray.runtime.task.TaskLanguage;
 import org.ray.runtime.task.TaskSpec;
-import org.ray.runtime.util.IdUtil;
 import org.ray.runtime.util.StringUtil;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -123,9 +125,8 @@ public abstract class AbstractRayRuntime implements RayRuntime {
 
   @Override
   public <T> RayObject<T> put(T obj) {
-    ObjectId objectId = IdUtil.computePutId(
-        workerContext.getCurrentTaskId(), workerContext.nextPutIndex());
-
+    ObjectId objectId = ObjectId.forPut(workerContext.getCurrentTaskId(),
+        workerContext.nextPutIndex());
     put(objectId, obj);
     return new RayObjectImpl<>(objectId);
   }
@@ -144,8 +145,8 @@ public abstract class AbstractRayRuntime implements RayRuntime {
    * @return A RayObject instance that represents the in-store object.
    */
   public RayObject<Object> putSerialized(byte[] obj) {
-    ObjectId objectId = IdUtil.computePutId(
-        workerContext.getCurrentTaskId(), workerContext.nextPutIndex());
+    ObjectId objectId = ObjectId.forPut(workerContext.getCurrentTaskId(),
+        workerContext.nextPutIndex());
     TaskId taskId = workerContext.getCurrentTaskId();
     LOGGER.debug("Putting serialized object {}, for task {} ", objectId, taskId);
     objectStoreProxy.putSerialized(objectId, obj);
@@ -212,7 +213,7 @@ public abstract class AbstractRayRuntime implements RayRuntime {
       Object[] args, ActorCreationOptions options) {
     TaskSpec spec = createTaskSpec(actorFactoryFunc, null, RayActorImpl.NIL,
         args, true, false, options);
-    RayActorImpl<?> actor = new RayActorImpl(new UniqueId(spec.returnIds[0].getBytes()));
+    RayActorImpl<?> actor = new RayActorImpl(spec.taskId.getActorId());
     actor.increaseTaskCounter();
     actor.setTaskCursor(spec.returnIds[0]);
     rayletClient.submitTask(spec);
@@ -272,7 +273,7 @@ public abstract class AbstractRayRuntime implements RayRuntime {
    *
    * @param func The target remote function.
    * @param pyFunctionDescriptor Descriptor of the target Python function, if the task is a Python
-   * task.
+   *                             task.
    * @param actor The actor handle. If the task is not an actor task, actor id must be NIL.
    * @param args The arguments for the remote function.
    * @param isActorCreationTask Whether this task is an actor creation task.
@@ -284,16 +285,22 @@ public abstract class AbstractRayRuntime implements RayRuntime {
       boolean isActorCreationTask, boolean isActorTask, BaseTaskOptions taskOptions) {
     Preconditions.checkArgument((func == null) != (pyFunctionDescriptor == null));
 
-    TaskId taskId = rayletClient.generateTaskId(workerContext.getCurrentJobId(),
-        workerContext.getCurrentTaskId(), workerContext.nextTaskIndex());
-    int numReturns = actor.getId().isNil() ? 1 : 2;
-    ObjectId[] returnIds = IdUtil.genReturnIds(taskId, numReturns);
-
-    UniqueId actorCreationId = UniqueId.NIL;
+    ActorId actorCreationId = ActorId.NIL;
+    TaskId taskId = null;
+    final JobId currentJobId = workerContext.getCurrentJobId();
+    final TaskId currentTaskId = workerContext.getCurrentTaskId();
+    final int taskIndex = workerContext.nextTaskIndex();
     if (isActorCreationTask) {
-      actorCreationId = new UniqueId(returnIds[0].getBytes());
+      taskId = RayletClientImpl.generateActorCreationTaskId(currentJobId, currentTaskId, taskIndex);
+      actorCreationId = taskId.getActorId();
+    } else if (isActorTask) {
+      taskId = RayletClientImpl.generateActorTaskId(currentJobId, currentTaskId, taskIndex, actor.getId());
+    } else {
+      taskId = RayletClientImpl.generateNormalTaskId(currentJobId, currentTaskId, taskIndex);
     }
 
+    int numReturns = actor.getId().isNil() ? 1 : 2;
+
     Map<String, Double> resources;
     if (null == taskOptions) {
       resources = new HashMap<>();
@@ -337,7 +344,7 @@ public abstract class AbstractRayRuntime implements RayRuntime {
         actor.getId(),
         actor.getHandleId(),
         actor.increaseTaskCounter(),
-	previousActorTaskDummyObjectId,
+        previousActorTaskDummyObjectId,
         actor.getNewActorHandles().toArray(new UniqueId[0]),
         ArgumentsBuilder.wrap(args, language == TaskLanguage.PYTHON),
         numReturns,

java/runtime/src/main/java/org/ray/runtime/RayActorImpl.java

@@ -7,6 +7,7 @@ import java.io.ObjectOutput;
 import java.util.ArrayList;
 import java.util.List;
 import org.ray.api.RayActor;
+import org.ray.api.id.ActorId;
 import org.ray.api.id.ObjectId;
 import org.ray.api.id.UniqueId;
 import org.ray.runtime.util.Sha1Digestor;
@@ -18,7 +19,7 @@ public class RayActorImpl<T> implements RayActor<T>, Externalizable {
   /**
    * Id of this actor.
    */
-  protected UniqueId id;
+  protected ActorId id;
   /**
    * Handle id of this actor.
    */
@@ -47,14 +48,14 @@ public class RayActorImpl<T> implements RayActor<T>, Externalizable {
   protected List<UniqueId> newActorHandles;
 
   public RayActorImpl() {
-    this(UniqueId.NIL, UniqueId.NIL);
+    this(ActorId.NIL, UniqueId.NIL);
   }
 
-  public RayActorImpl(UniqueId id) {
+  public RayActorImpl(ActorId id) {
     this(id, UniqueId.NIL);
   }
 
-  public RayActorImpl(UniqueId id, UniqueId handleId) {
+  public RayActorImpl(ActorId id, UniqueId handleId) {
     this.id = id;
     this.handleId = handleId;
     this.taskCounter = 0;
@@ -64,7 +65,7 @@ public class RayActorImpl<T> implements RayActor<T>, Externalizable {
   }
 
   @Override
-  public UniqueId getId() {
+  public ActorId getId() {
     return id;
   }
 
@@ -120,7 +121,7 @@ public class RayActorImpl<T> implements RayActor<T>, Externalizable {
 
   @Override
   public void readExternal(ObjectInput in) throws IOException, ClassNotFoundException {
-    this.id = (UniqueId) in.readObject();
+    this.id = (ActorId) in.readObject();
     this.handleId = (UniqueId) in.readObject();
     this.taskCursor = (ObjectId) in.readObject();
     this.taskCounter = (int) in.readObject();

java/runtime/src/main/java/org/ray/runtime/RayPyActorImpl.java

@@ -4,11 +4,11 @@ import java.io.IOException;
 import java.io.ObjectInput;
 import java.io.ObjectOutput;
 import org.ray.api.RayPyActor;
-import org.ray.api.id.UniqueId;
+import org.ray.api.id.ActorId;
 
 public class RayPyActorImpl extends RayActorImpl implements RayPyActor {
 
-  public static final RayPyActorImpl NIL = new RayPyActorImpl(UniqueId.NIL, null, null);
+  public static final RayPyActorImpl NIL = new RayPyActorImpl(ActorId.NIL, null, null);
 
   /**
    * Module name of the Python actor class.
@@ -24,7 +24,7 @@ public class RayPyActorImpl extends RayActorImpl implements RayPyActor {
   // since it'll be needed when deserializing.
   public RayPyActorImpl() {}
 
-  public RayPyActorImpl(UniqueId id, String moduleName, String className) {
+  public RayPyActorImpl(ActorId id, String moduleName, String className) {
     super(id);
     this.moduleName = moduleName;
     this.className = className;

java/runtime/src/main/java/org/ray/runtime/RuntimeContextImpl.java

@@ -2,8 +2,9 @@ package org.ray.runtime;
 
 import com.google.common.base.Preconditions;
 import java.util.List;
+
+import org.ray.api.id.ActorId;
 import org.ray.api.id.JobId;
-import org.ray.api.id.UniqueId;
 import org.ray.api.runtimecontext.NodeInfo;
 import org.ray.api.runtimecontext.RuntimeContext;
 import org.ray.runtime.config.RunMode;
@@ -23,7 +24,7 @@ public class RuntimeContextImpl implements RuntimeContext {
   }
 
   @Override
-  public UniqueId getCurrentActorId() {
+  public ActorId getCurrentActorId() {
     Worker worker = runtime.getWorker();
     Preconditions.checkState(worker != null && !worker.getCurrentActorId().isNil(),
         "This method should only be called from an actor.");

java/runtime/src/main/java/org/ray/runtime/Worker.java

@@ -7,12 +7,14 @@ import org.ray.api.Checkpointable;
 import org.ray.api.Checkpointable.Checkpoint;
 import org.ray.api.Checkpointable.CheckpointContext;
 import org.ray.api.exception.RayTaskException;
+import org.ray.api.id.ActorId;
 import org.ray.api.id.ObjectId;
 import org.ray.api.id.UniqueId;
 import org.ray.runtime.config.RunMode;
 import org.ray.runtime.functionmanager.RayFunction;
 import org.ray.runtime.task.ArgumentsBuilder;
 import org.ray.runtime.task.TaskSpec;
+import org.ray.runtime.util.IdUtil;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -37,7 +39,7 @@ public class Worker {
   /**
    * Id of the current actor object, if the worker is an actor, otherwise NIL.
    */
-  private UniqueId currentActorId = UniqueId.NIL;
+  private ActorId currentActorId = ActorId.NIL;
 
   /**
    * The exception that failed the actor creation task, if any.
@@ -64,7 +66,7 @@ public class Worker {
     this.runtime = runtime;
   }
 
-  public UniqueId getCurrentActorId() {
+  public ActorId getCurrentActorId() {
     return currentActorId;
   }
 
@@ -92,7 +94,7 @@ public class Worker {
       Thread.currentThread().setContextClassLoader(rayFunction.classLoader);
 
       if (spec.isActorCreationTask()) {
-        currentActorId = new UniqueId(returnId.getBytes());
+        currentActorId = spec.taskId.getActorId();
       }
 
       // Get local actor object and arguments.
@@ -118,9 +120,10 @@ public class Worker {
         if (spec.isActorTask()) {
           maybeSaveCheckpoint(actor, spec.actorId);
         }
+
         runtime.put(returnId, result);
       } else {
-        maybeLoadCheckpoint(result, new UniqueId(returnId.getBytes()));
+        maybeLoadCheckpoint(result, spec.taskId.getActorId());
         currentActor = result;
       }
       LOGGER.debug("Finished executing task {}", spec.taskId);
@@ -136,7 +139,7 @@ public class Worker {
     }
   }
 
-  private void maybeSaveCheckpoint(Object actor, UniqueId actorId) {
+  private void maybeSaveCheckpoint(Object actor, ActorId actorId) {
     if (!(actor instanceof Checkpointable)) {
       return;
     }
@@ -161,7 +164,7 @@ public class Worker {
     checkpointable.saveCheckpoint(actorId, checkpointId);
   }
 
-  private void maybeLoadCheckpoint(Object actor, UniqueId actorId) {
+  private void maybeLoadCheckpoint(Object actor, ActorId actorId) {
     if (!(actor instanceof Checkpointable)) {
       return;
     }

java/runtime/src/main/java/org/ray/runtime/WorkerContext.java

@@ -48,7 +48,7 @@ public class WorkerContext {
    * for other threads, this method returns a random ID.
    */
   public TaskId getCurrentTaskId() {
-    return new TaskId(nativeGetCurrentTaskId(nativeWorkerContextPointer));
+    return TaskId.fromBytes(nativeGetCurrentTaskId(nativeWorkerContextPointer));
   }
 
   /**

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

@@ -9,6 +9,7 @@ import java.util.Map;
 import java.util.stream.Collectors;
 import org.apache.commons.lang3.ArrayUtils;
 import org.ray.api.Checkpointable.Checkpoint;
+import org.ray.api.id.ActorId;
 import org.ray.api.id.BaseId;
 import org.ray.api.id.JobId;
 import org.ray.api.id.TaskId;
@@ -117,7 +118,7 @@ public class GcsClient {
   /**
    * If the actor exists in GCS.
    */
-  public boolean actorExists(UniqueId actorId) {
+  public boolean actorExists(ActorId actorId) {
     byte[] key = ArrayUtils.addAll(
         TablePrefix.ACTOR.toString().getBytes(), actorId.getBytes());
     return primary.exists(key);
@@ -136,7 +137,7 @@ public class GcsClient {
   /**
    * Get the available checkpoints for the given actor ID.
    */
-  public List<Checkpoint> getCheckpointsForActor(UniqueId actorId) {
+  public List<Checkpoint> getCheckpointsForActor(ActorId actorId) {
     List<Checkpoint> checkpoints = new ArrayList<>();
     final String prefix = TablePrefix.ACTOR_CHECKPOINT_ID.toString();
     final byte[] key = ArrayUtils.addAll(prefix.getBytes(), actorId.getBytes());

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

@@ -9,7 +9,6 @@ import java.util.function.Consumer;
 import java.util.stream.Collectors;
 import org.ray.api.id.ObjectId;
 import org.ray.runtime.WorkerContext;
-import org.ray.runtime.util.IdUtil;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -37,7 +36,7 @@ public class MockObjectInterface implements ObjectInterface {
 
   @Override
   public ObjectId put(NativeRayObject obj) {
-    ObjectId objectId = IdUtil.computePutId(workerContext.getCurrentTaskId(),
+    ObjectId objectId = ObjectId.forPut(workerContext.getCurrentTaskId(),
         workerContext.nextPutIndex());
     put(obj, objectId);
     return objectId;

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

@@ -57,7 +57,8 @@ public class ObjectInterfaceImpl implements ObjectInterface {
 
   @Override
   public void delete(List<ObjectId> objectIds, boolean localOnly, boolean deleteCreatingTasks) {
-    nativeDelete(nativeObjectInterfacePointer, toBinaryList(objectIds), localOnly, deleteCreatingTasks);
+    nativeDelete(nativeObjectInterfacePointer,
+        toBinaryList(objectIds), localOnly, deleteCreatingTasks);
   }
 
   public void destroy() {

java/runtime/src/main/java/org/ray/runtime/raylet/MockRayletClient.java

@@ -18,7 +18,7 @@ import java.util.stream.Collectors;
 import org.apache.commons.lang3.NotImplementedException;
 import org.ray.api.RayObject;
 import org.ray.api.WaitResult;
-import org.ray.api.id.JobId;
+import org.ray.api.id.ActorId;
 import org.ray.api.id.ObjectId;
 import org.ray.api.id.TaskId;
 import org.ray.api.id.UniqueId;
@@ -43,7 +43,7 @@ public class MockRayletClient implements RayletClient {
   private final RayDevRuntime runtime;
   private final ExecutorService exec;
   private final Deque<Worker> idleWorkers;
-  private final Map<UniqueId, Worker> actorWorkers;
+  private final Map<ActorId, Worker> actorWorkers;
   private final ThreadLocal<Worker> currentWorker;
 
   public MockRayletClient(RayDevRuntime runtime, int numberThreads) {
@@ -154,11 +154,6 @@ public class MockRayletClient implements RayletClient {
     throw new RuntimeException("invalid execution flow here");
   }
 
-  @Override
-  public TaskId generateTaskId(JobId jobId, TaskId parentTaskId, int taskIndex) {
-    return TaskId.randomId();
-  }
-
   @Override
   public <T> WaitResult<T> wait(List<RayObject<T>> waitFor, int numReturns, int
       timeoutMs, TaskId currentTaskId) {
@@ -188,12 +183,12 @@ public class MockRayletClient implements RayletClient {
 
 
   @Override
-  public UniqueId prepareCheckpoint(UniqueId actorId) {
+  public UniqueId prepareCheckpoint(ActorId actorId) {
     throw new NotImplementedException("Not implemented.");
   }
 
   @Override
-  public void notifyActorResumedFromCheckpoint(UniqueId actorId, UniqueId checkpointId) {
+  public void notifyActorResumedFromCheckpoint(ActorId actorId, UniqueId checkpointId) {
     throw new NotImplementedException("Not implemented.");
   }
 

java/runtime/src/main/java/org/ray/runtime/raylet/RayletClient.java

@@ -3,7 +3,7 @@ package org.ray.runtime.raylet;
 import java.util.List;
 import org.ray.api.RayObject;
 import org.ray.api.WaitResult;
-import org.ray.api.id.JobId;
+import org.ray.api.id.ActorId;
 import org.ray.api.id.ObjectId;
 import org.ray.api.id.TaskId;
 import org.ray.api.id.UniqueId;
@@ -18,16 +18,14 @@ public interface RayletClient {
 
   TaskSpec getTask();
 
-  TaskId generateTaskId(JobId jobId, TaskId parentTaskId, int taskIndex);
-
   <T> WaitResult<T> wait(List<RayObject<T>> waitFor, int numReturns, int
       timeoutMs, TaskId currentTaskId);
 
   void freePlasmaObjects(List<ObjectId> objectIds, boolean localOnly, boolean deleteCreatingTasks);
 
-  UniqueId prepareCheckpoint(UniqueId actorId);
+  UniqueId prepareCheckpoint(ActorId actorId);
 
-  void notifyActorResumedFromCheckpoint(UniqueId actorId, UniqueId checkpointId);
+  void notifyActorResumedFromCheckpoint(ActorId actorId, UniqueId checkpointId);
 
   void setResource(String resourceName, double capacity, UniqueId nodeId);
 

java/runtime/src/main/java/org/ray/runtime/raylet/RayletClientImpl.java

@@ -4,6 +4,7 @@ import com.google.common.base.Preconditions;
 import com.google.common.collect.ImmutableList;
 import com.google.protobuf.ByteString;
 import com.google.protobuf.InvalidProtocolBufferException;
+import java.nio.ByteBuffer;
 import java.nio.charset.Charset;
 import java.util.ArrayList;
 import java.util.Arrays;
@@ -13,10 +14,11 @@ import java.util.stream.Collectors;
 import org.ray.api.RayObject;
 import org.ray.api.WaitResult;
 import org.ray.api.exception.RayException;
-import org.ray.api.id.JobId;
-import org.ray.api.id.ObjectId;
-import org.ray.api.id.TaskId;
+import org.ray.api.id.ActorId;
 import org.ray.api.id.UniqueId;
+import org.ray.api.id.JobId;
+import org.ray.api.id.TaskId;
+import org.ray.api.id.ObjectId;
 import org.ray.runtime.functionmanager.JavaFunctionDescriptor;
 import org.ray.runtime.generated.Common;
 import org.ray.runtime.generated.Common.TaskType;
@@ -93,12 +95,6 @@ public class RayletClientImpl implements RayletClient {
     return parseTaskSpecFromProtobuf(bytes);
   }
 
-  @Override
-  public TaskId generateTaskId(JobId jobId, TaskId parentTaskId, int taskIndex) {
-    byte[] bytes = nativeGenerateTaskId(jobId.getBytes(), parentTaskId.getBytes(), taskIndex);
-    return new TaskId(bytes);
-  }
-
   @Override
   public void freePlasmaObjects(List<ObjectId> objectIds, boolean localOnly,
       boolean deleteCreatingTasks) {
@@ -107,15 +103,30 @@ public class RayletClientImpl implements RayletClient {
   }
 
   @Override
-  public UniqueId prepareCheckpoint(UniqueId actorId) {
+  public UniqueId prepareCheckpoint(ActorId actorId) {
     return new UniqueId(nativePrepareCheckpoint(client, actorId.getBytes()));
   }
 
   @Override
-  public void notifyActorResumedFromCheckpoint(UniqueId actorId, UniqueId checkpointId) {
+  public void notifyActorResumedFromCheckpoint(ActorId actorId, UniqueId checkpointId) {
     nativeNotifyActorResumedFromCheckpoint(client, actorId.getBytes(), checkpointId.getBytes());
   }
 
+  public static TaskId generateActorCreationTaskId(JobId jobId, TaskId parentTaskId, int taskIndex) {
+    byte[] bytes = nativeGenerateActorCreationTaskId(jobId.getBytes(), parentTaskId.getBytes(), taskIndex);
+    return TaskId.fromBytes(bytes);
+  }
+
+  public static TaskId generateActorTaskId(JobId jobId, TaskId parentTaskId, int taskIndex, ActorId actorId) {
+    byte[] bytes = nativeGenerateActorTaskId(jobId.getBytes(), parentTaskId.getBytes(), taskIndex, actorId.getBytes());
+    return TaskId.fromBytes(bytes);
+  }
+
+  public static TaskId generateNormalTaskId(JobId jobId, TaskId parentTaskId, int taskIndex) {
+    byte[] bytes = nativeGenerateNormalTaskId(jobId.getBytes(), parentTaskId.getBytes(), taskIndex);
+    return TaskId.fromBytes(bytes);
+  }
+
   /**
    * Parse `TaskSpec` protobuf bytes.
    */
@@ -160,13 +171,13 @@ public class RayletClientImpl implements RayletClient {
     );
 
     // Parse ActorCreationTaskSpec.
-    UniqueId actorCreationId = UniqueId.NIL;
+    ActorId actorCreationId = ActorId.NIL;
     int maxActorReconstructions = 0;
     UniqueId[] newActorHandles = new UniqueId[0];
     List<String> dynamicWorkerOptions = new ArrayList<>();
     if (taskSpec.getType() == Common.TaskType.ACTOR_CREATION_TASK) {
       Common.ActorCreationTaskSpec actorCreationTaskSpec = taskSpec.getActorCreationTaskSpec();
-      actorCreationId = UniqueId
+      actorCreationId = ActorId
           .fromByteBuffer(actorCreationTaskSpec.getActorId().asReadOnlyByteBuffer());
       maxActorReconstructions = (int) actorCreationTaskSpec.getMaxActorReconstructions();
       dynamicWorkerOptions = ImmutableList
@@ -174,18 +185,18 @@ public class RayletClientImpl implements RayletClient {
     }
 
     // Parse ActorTaskSpec.
-    UniqueId actorId = UniqueId.NIL;
+    ActorId actorId = ActorId.NIL;
     UniqueId actorHandleId = UniqueId.NIL;
     ObjectId previousActorTaskDummyObjectId = ObjectId.NIL;
     int actorCounter = 0;
     if (taskSpec.getType() == Common.TaskType.ACTOR_TASK) {
       Common.ActorTaskSpec actorTaskSpec = taskSpec.getActorTaskSpec();
-      actorId = UniqueId.fromByteBuffer(actorTaskSpec.getActorId().asReadOnlyByteBuffer());
+      actorId = ActorId.fromByteBuffer(actorTaskSpec.getActorId().asReadOnlyByteBuffer());
       actorHandleId = UniqueId
           .fromByteBuffer(actorTaskSpec.getActorHandleId().asReadOnlyByteBuffer());
       actorCounter = (int) actorTaskSpec.getActorCounter();
       previousActorTaskDummyObjectId = ObjectId.fromByteBuffer(
-		      actorTaskSpec.getPreviousActorTaskDummyObjectId().asReadOnlyByteBuffer());
+          actorTaskSpec.getPreviousActorTaskDummyObjectId().asReadOnlyByteBuffer());
       newActorHandles = actorTaskSpec.getNewActorHandlesList().stream()
           .map(byteString -> UniqueId.fromByteBuffer(byteString.asReadOnlyByteBuffer()))
           .toArray(UniqueId[]::new);
@@ -193,8 +204,8 @@ public class RayletClientImpl implements RayletClient {
 
     return new TaskSpec(jobId, taskId, parentTaskId, parentCounter, actorCreationId,
         maxActorReconstructions, actorId, actorHandleId, actorCounter,
-	previousActorTaskDummyObjectId, newActorHandles, args, numReturns, resources,
-	TaskLanguage.JAVA, functionDescriptor, dynamicWorkerOptions);
+        previousActorTaskDummyObjectId, newActorHandles, args, numReturns, resources,
+        TaskLanguage.JAVA, functionDescriptor, dynamicWorkerOptions);
   }
 
   /**
@@ -255,13 +266,16 @@ public class RayletClientImpl implements RayletClient {
       builder.setType(TaskType.ACTOR_TASK);
       List<ByteString> newHandles = Arrays.stream(task.newActorHandles)
           .map(id -> ByteString.copyFrom(id.getBytes())).collect(Collectors.toList());
+      final ObjectId actorCreationDummyObjectId = IdUtil.computeActorCreationDummyObjectId(
+          ActorId.fromByteBuffer(ByteBuffer.wrap(task.actorId.getBytes())));
       builder.setActorTaskSpec(
           Common.ActorTaskSpec.newBuilder()
               .setActorId(ByteString.copyFrom(task.actorId.getBytes()))
               .setActorHandleId(ByteString.copyFrom(task.actorHandleId.getBytes()))
-              .setActorCreationDummyObjectId(ByteString.copyFrom(task.actorId.getBytes()))
+              .setActorCreationDummyObjectId(
+                  ByteString.copyFrom(actorCreationDummyObjectId.getBytes()))
               .setPreviousActorTaskDummyObjectId(
-		      ByteString.copyFrom(task.previousActorTaskDummyObjectId.getBytes()))
+                  ByteString.copyFrom(task.previousActorTaskDummyObjectId.getBytes()))
               .setActorCounter(task.actorCounter)
               .addAllNewActorHandles(newHandles)
       );
@@ -307,9 +321,6 @@ public class RayletClientImpl implements RayletClient {
   private static native boolean[] nativeWaitObject(long conn, byte[][] objectIds,
       int numReturns, int timeout, boolean waitLocal, byte[] currentTaskId) throws RayException;
 
-  private static native byte[] nativeGenerateTaskId(byte[] jobId, byte[] parentTaskId,
-      int taskIndex);
-
   private static native void nativeFreePlasmaObjects(long conn, byte[][] objectIds,
       boolean localOnly, boolean deleteCreatingTasks) throws RayException;
 
@@ -320,4 +331,13 @@ public class RayletClientImpl implements RayletClient {
 
   private static native void nativeSetResource(long conn, String resourceName, double capacity,
       byte[] nodeId) throws RayException;
+
+  private static native byte[] nativeGenerateActorCreationTaskId(byte[] jobId, byte[] parentTaskId,
+                                                                 int taskIndex);
+
+  private static native byte[] nativeGenerateActorTaskId(byte[] jobId, byte[] parentTaskId,
+                                                          int taskIndex, byte[] actorId);
+
+  private static native byte[] nativeGenerateNormalTaskId(byte[] jobId, byte[] parentTaskId,
+                                                          int taskIndex);
 }

java/runtime/src/main/java/org/ray/runtime/task/TaskSpec.java

@@ -1,18 +1,17 @@
 package org.ray.runtime.task;
 
 import com.google.common.base.Preconditions;
-import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.List;
 import java.util.Map;
+import org.ray.api.id.ActorId;
 import org.ray.api.id.JobId;
-import org.ray.api.id.ObjectId;
 import org.ray.api.id.TaskId;
+import org.ray.api.id.ObjectId;
 import org.ray.api.id.UniqueId;
 import org.ray.runtime.functionmanager.FunctionDescriptor;
 import org.ray.runtime.functionmanager.JavaFunctionDescriptor;
 import org.ray.runtime.functionmanager.PyFunctionDescriptor;
-import org.ray.runtime.util.IdUtil;
 
 /**
  * Represents necessary information of a task for scheduling and executing.
@@ -32,13 +31,13 @@ public class TaskSpec {
   public final int parentCounter;
 
   // Id for createActor a target actor
-  public final UniqueId actorCreationId;
+  public final ActorId actorCreationId;
 
   public final int maxActorReconstructions;
 
   // Actor ID of the task. This is the actor that this task is executed on
   // or NIL_ACTOR_ID if the task is just a normal task.
-  public final UniqueId actorId;
+  public final ActorId actorId;
 
   // ID per actor client for session consistency
   public final UniqueId actorHandleId;
@@ -87,9 +86,9 @@ public class TaskSpec {
       TaskId taskId,
       TaskId parentTaskId,
       int parentCounter,
-      UniqueId actorCreationId,
+      ActorId actorCreationId,
       int maxActorReconstructions,
-      UniqueId actorId,
+      ActorId actorId,
       UniqueId actorHandleId,
       int actorCounter,
       ObjectId previousActorTaskDummyObjectId,
@@ -117,7 +116,7 @@ public class TaskSpec {
 
     returnIds = new ObjectId[numReturns];
     for (int i = 0; i < numReturns; ++i) {
-      returnIds[i] = IdUtil.computeReturnId(taskId, i + 1);
+      returnIds[i] = ObjectId.forReturn(taskId, i + 1);
     }
     this.resources = resources;
     this.language = language;

java/runtime/src/main/java/org/ray/runtime/util/IdUtil.java

@@ -1,15 +1,11 @@
 package org.ray.runtime.util;
 
-import com.google.common.base.Preconditions;
 import java.nio.ByteBuffer;
-import java.nio.ByteOrder;
 import java.util.Arrays;
 import java.util.List;
 import org.ray.api.id.BaseId;
-import org.ray.api.id.JobId;
 import org.ray.api.id.ObjectId;
-import org.ray.api.id.TaskId;
-import org.ray.api.id.UniqueId;
+import org.ray.api.id.ActorId;
 
 /**
  * Helper method for different Ids.
@@ -17,60 +13,6 @@ import org.ray.api.id.UniqueId;
  * in src/ray/common/id.h
  */
 public class IdUtil {
-  public static final int OBJECT_INDEX_POS = 16;
-
-  /**
-   * Compute the object ID of an object returned by the task.
-   *
-   * @param taskId The task ID of the task that created the object.
-   * @param returnIndex What number return value this object is in the task.
-   * @return The computed object ID.
-   */
-  public static ObjectId computeReturnId(TaskId taskId, int returnIndex) {
-    return computeObjectId(taskId, returnIndex);
-  }
-
-  /**
-   * Compute the object ID from the task ID and the index.
-   * @param taskId The task ID of the task that created the object.
-   * @param index The index which can distinguish different objects in one task.
-   * @return The computed object ID.
-   */
-  private static ObjectId computeObjectId(TaskId taskId, int index) {
-    byte[] bytes = new byte[ObjectId.LENGTH];
-    System.arraycopy(taskId.getBytes(), 0, bytes, 0, taskId.size());
-    ByteBuffer wbb = ByteBuffer.wrap(bytes);
-    wbb.order(ByteOrder.LITTLE_ENDIAN);
-    wbb.putInt(OBJECT_INDEX_POS, index);
-    return new ObjectId(bytes);
-  }
-
-  /**
-   * Compute the object ID of an object put by the task.
-   *
-   * @param taskId The task ID of the task that created the object.
-   * @param putIndex What number put this object was created by in the task.
-   * @return The computed object ID.
-   */
-  public static ObjectId computePutId(TaskId taskId, int putIndex) {
-    // We multiply putIndex by -1 to distinguish from returnIndex.
-    return computeObjectId(taskId, -1 * putIndex);
-  }
-
-  /**
-   * Generate the return ids of a task.
-   *
-   * @param taskId The ID of the task that generates returnsIds.
-   * @param numReturns The number of returnIds.
-   * @return The Return Ids of this task.
-   */
-  public static ObjectId[] genReturnIds(TaskId taskId, int numReturns) {
-    ObjectId[] ret = new ObjectId[numReturns];
-    for (int i = 0; i < numReturns; i++) {
-      ret[i] = IdUtil.computeReturnId(taskId, i + 1);
-    }
-    return ret;
-  }
 
   public static <T extends BaseId> byte[][] getIdBytes(List<T> objectIds) {
     int size = objectIds.size();
@@ -81,79 +23,6 @@ public class IdUtil {
     return ids;
   }
 
-  public static byte[][] getByteListFromByteBuffer(ByteBuffer byteBufferOfIds, int length) {
-    Preconditions.checkArgument(byteBufferOfIds != null);
-
-    byte[] bytesOfIds = new byte[byteBufferOfIds.remaining()];
-    byteBufferOfIds.get(bytesOfIds, 0, byteBufferOfIds.remaining());
-
-    int count = bytesOfIds.length / length;
-    byte[][] idBytes = new byte[count][];
-
-    for (int i = 0; i < count; ++i) {
-      byte[] id = new byte[length];
-      System.arraycopy(bytesOfIds, i * length, id, 0, length);
-      idBytes[i] = id;
-    }
-
-    return idBytes;
-  }
-
-  /**
-   * Get unique IDs from concatenated ByteBuffer.
-   *
-   * @param byteBufferOfIds The ByteBuffer concatenated from IDs.
-   * @return The array of unique IDs.
-   */
-  public static UniqueId[] getUniqueIdsFromByteBuffer(ByteBuffer byteBufferOfIds) {
-    byte[][]idBytes = getByteListFromByteBuffer(byteBufferOfIds, UniqueId.LENGTH);
-    UniqueId[] uniqueIds = new UniqueId[idBytes.length];
-
-    for (int i = 0; i < idBytes.length; ++i) {
-      uniqueIds[i] = UniqueId.fromByteBuffer(ByteBuffer.wrap(idBytes[i]));
-    }
-
-    return uniqueIds;
-  }
-
-  /**
-   * Get object IDs from concatenated ByteBuffer.
-   *
-   * @param byteBufferOfIds The ByteBuffer concatenated from IDs.
-   * @return The array of object IDs.
-   */
-  public static ObjectId[] getObjectIdsFromByteBuffer(ByteBuffer byteBufferOfIds) {
-    byte[][]idBytes = getByteListFromByteBuffer(byteBufferOfIds, UniqueId.LENGTH);
-    ObjectId[] objectIds = new ObjectId[idBytes.length];
-
-    for (int i = 0; i < idBytes.length; ++i) {
-      objectIds[i] = ObjectId.fromByteBuffer(ByteBuffer.wrap(idBytes[i]));
-    }
-
-    return objectIds;
-  }
-
-  /**
-   * Concatenate IDs to a ByteBuffer.
-   *
-   * @param ids The array of IDs that will be concatenated.
-   * @return A ByteBuffer that contains bytes of concatenated IDs.
-   */
-  public static <T extends BaseId> ByteBuffer concatIds(T[] ids) {
-    int length = 0;
-    if (ids != null && ids.length != 0) {
-      length = ids[0].size() * ids.length;
-    }
-    byte[] bytesOfIds = new byte[length];
-    for (int i = 0; i < ids.length; ++i) {
-      System.arraycopy(ids[i].getBytes(), 0, bytesOfIds,
-          i * ids[i].size(), ids[i].size());
-    }
-
-    return ByteBuffer.wrap(bytesOfIds);
-  }
-
-
   /**
    * Compute the murmur hash code of this ID.
    */
@@ -221,4 +90,16 @@ public class IdUtil {
 
     return h;
   }
+
+  /*
+   * A helper function to compute actor creation dummy object id according
+   * the given actor id.
+   */
+  public static ObjectId computeActorCreationDummyObjectId(ActorId actorId) {
+    byte[] bytes = new byte[ObjectId.LENGTH];
+    System.arraycopy(actorId.getBytes(), 0, bytes, 0, ActorId.LENGTH);
+    Arrays.fill(bytes, ActorId.LENGTH, bytes.length, (byte) 0xFF);
+    return ObjectId.fromByteBuffer(ByteBuffer.wrap(bytes));
+  }
+
 }