[Java] Allow actor handle to be serialized without forking (#3686)

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

@@ -67,7 +67,17 @@ public final class RayActorImpl<T> implements RayActor<T>, Externalizable {
     return taskCounter++;
   }
 
-  private UniqueId computeNextActorHandleId() {
+  public RayActorImpl<T> fork() {
+    RayActorImpl<T> ret = new RayActorImpl<>();
+    ret.id = this.id;
+    ret.taskCounter = 0;
+    ret.numForks = 0;
+    ret.taskCursor = this.taskCursor;
+    ret.handleId = this.computeNextActorHandleId();
+    return ret;
+  }
+
+  protected UniqueId computeNextActorHandleId() {
     byte[] bytes = Sha1Digestor.digest(handleId.getBytes(), ++numForks);
     return new UniqueId(bytes);
   }
@@ -75,8 +85,10 @@ public final class RayActorImpl<T> implements RayActor<T>, Externalizable {
   @Override
   public void writeExternal(ObjectOutput out) throws IOException {
     out.writeObject(this.id);
-    out.writeObject(this.computeNextActorHandleId());
+    out.writeObject(this.handleId);
     out.writeObject(this.taskCursor);
+    out.writeObject(this.taskCounter);
+    out.writeObject(this.numForks);
   }
 
   @Override
@@ -84,5 +96,7 @@ public final class RayActorImpl<T> implements RayActor<T>, Externalizable {
     this.id = (UniqueId) in.readObject();
     this.handleId = (UniqueId) in.readObject();
     this.taskCursor = (UniqueId) in.readObject();
+    this.taskCounter = (int) in.readObject();
+    this.numForks = (int) in.readObject();
   }
 }

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

@@ -0,0 +1,25 @@
+package org.ray.runtime.util;
+
+import java.io.IOException;
+import org.nustaq.serialization.FSTBasicObjectSerializer;
+import org.nustaq.serialization.FSTClazzInfo;
+import org.nustaq.serialization.FSTClazzInfo.FSTFieldInfo;
+import org.nustaq.serialization.FSTObjectInput;
+import org.nustaq.serialization.FSTObjectOutput;
+import org.ray.runtime.RayActorImpl;
+
+public class RayActorSerializer extends FSTBasicObjectSerializer {
+
+  @Override
+  public void writeObject(FSTObjectOutput out, Object toWrite, FSTClazzInfo clzInfo,
+      FSTClazzInfo.FSTFieldInfo referencedBy, int streamPosition) throws IOException {
+    ((RayActorImpl) toWrite).fork().writeExternal(out);
+  }
+
+  @Override
+  public void readObject(FSTObjectInput in, Object toRead, FSTClazzInfo clzInfo,
+                         FSTFieldInfo referencedBy) throws Exception {
+    super.readObject(in, toRead, clzInfo, referencedBy);
+    ((RayActorImpl) toRead).readExternal(in);
+  }
+}

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

@@ -1,14 +1,18 @@
 package org.ray.runtime.util;
 
 import org.nustaq.serialization.FSTConfiguration;
+import org.ray.runtime.RayActorImpl;
 
 /**
  * Java object serialization TODO: use others (e.g. Arrow) for higher performance
  */
 public class Serializer {
 
-  static final ThreadLocal<FSTConfiguration> conf = ThreadLocal.withInitial(
-      FSTConfiguration::createDefaultConfiguration);
+  private static final ThreadLocal<FSTConfiguration> conf = ThreadLocal.withInitial(() -> {
+    FSTConfiguration conf = FSTConfiguration.createDefaultConfiguration();
+    conf.registerSerializer(RayActorImpl.class, new RayActorSerializer(), true);
+    return conf;
+  });
 
   public static byte[] encode(Object obj) {
     return conf.get().asByteArray(obj);