[Java] Support exchange ObjectRef between processes (#10729)

java/runtime/src/main/java/io/ray/runtime/object/LocalModeObjectStore.java

@@ -105,4 +105,14 @@ public class LocalModeObjectStore extends ObjectStore {
   @Override
   public void removeLocalReference(UniqueId workerId, ObjectId objectId) {
   }
+
+  @Override
+  public byte[] promoteAndGetOwnershipInfo(ObjectId objectId) {
+    return new byte[0];
+  }
+
+  @Override
+  public void registerOwnershipInfoAndResolveFuture(
+      ObjectId objectId, ObjectId outerObjectId, byte[] ownerAddress) {
+  }
 }

java/runtime/src/main/java/io/ray/runtime/object/NativeObjectStore.java

@@ -68,6 +68,21 @@ public class NativeObjectStore extends ObjectStore {
     }
   }
 
+  @Override
+  public byte[] promoteAndGetOwnershipInfo(ObjectId objectId) {
+    return nativePromoteAndGetOwnershipInfo(objectId.getBytes());
+  }
+
+  @Override
+  public void registerOwnershipInfoAndResolveFuture(ObjectId objectId, ObjectId outerObjectId,
+      byte[] ownerAddress) {
+    byte[] outer = null;
+    if (outerObjectId != null) {
+      outer = outerObjectId.getBytes();
+    }
+    nativeRegisterOwnershipInfoAndResolveFuture(objectId.getBytes(), outer, ownerAddress);
+  }
+
   public Map<ObjectId, long[]> getAllReferenceCounts() {
     Map<ObjectId, long[]> referenceCounts = new HashMap<>();
     for (Map.Entry<byte[], long[]> entry :
@@ -98,4 +113,9 @@ public class NativeObjectStore extends ObjectStore {
   private static native void nativeRemoveLocalReference(byte[] workerId, byte[] objectId);
 
   private static native Map<byte[], long[]> nativeGetAllReferenceCounts();
+
+  private static native byte[] nativePromoteAndGetOwnershipInfo(byte[] objectId);
+
+  private static native void nativeRegisterOwnershipInfoAndResolveFuture(byte[] objectId,
+      byte[] outerObjectId, byte[] ownerAddress);
 }

java/runtime/src/main/java/io/ray/runtime/object/ObjectRefImpl.java

@@ -65,6 +65,10 @@ public final class ObjectRefImpl<T> implements ObjectRef<T>, Externalizable {
   public void writeExternal(ObjectOutput out) throws IOException {
     out.writeObject(this.getId());
     out.writeObject(this.getType());
+    RayRuntimeInternal runtime = (RayRuntimeInternal) Ray.internal();
+    byte[] ownerAddress = runtime.getObjectStore().promoteAndGetOwnershipInfo(this.getId());
+    out.writeInt(ownerAddress.length);
+    out.write(ownerAddress);
     ObjectSerializer.addContainedObjectId(this.getId());
   }
 
@@ -72,7 +76,13 @@ public final class ObjectRefImpl<T> implements ObjectRef<T>, Externalizable {
   public void readExternal(ObjectInput in) throws IOException, ClassNotFoundException {
     this.id = (ObjectId) in.readObject();
     this.type = (Class<T>) in.readObject();
+    int len = in.readInt();
+    byte[] ownerAddress = new byte[len];
+    in.readFully(ownerAddress);
     addLocalReference();
+    RayRuntimeInternal runtime = (RayRuntimeInternal) Ray.internal();
+    runtime.getObjectStore().registerOwnershipInfoAndResolveFuture(
+        this.id, ObjectSerializer.getOuterObjectId(), ownerAddress);
   }
 
   private void addLocalReference() {

java/runtime/src/main/java/io/ray/runtime/object/ObjectSerializer.java

@@ -47,6 +47,8 @@ public class ObjectSerializer {
   // field will contain all the nested object IDs.
   static ThreadLocal<Set<ObjectId>> containedObjectIds = ThreadLocal.withInitial(HashSet::new);
 
+  static ThreadLocal<ObjectId> outerObjectId = ThreadLocal.withInitial(() -> null);
+
   /**
    * Deserialize an object from an {@link NativeRayObject} instance.
    *
@@ -170,4 +172,16 @@ public class ObjectSerializer {
     containedObjectIds.get().clear();
     return ids;
   }
+
+  static void setOuterObjectId(ObjectId objectId) {
+    outerObjectId.set(objectId);
+  }
+
+  static ObjectId getOuterObjectId() {
+    return outerObjectId.get();
+  }
+
+  static void resetOuterObjectId() {
+    outerObjectId.set(null);
+  }
 }

java/runtime/src/main/java/io/ray/runtime/object/ObjectStore.java

@@ -96,8 +96,13 @@ public abstract class ObjectStore {
       NativeRayObject dataAndMeta = dataAndMetaList.get(i);
       Object object = null;
       if (dataAndMeta != null) {
-        object = ObjectSerializer
+        try {
+          ObjectSerializer.setOuterObjectId(ids.get(i));
+          object = ObjectSerializer
             .deserialize(dataAndMeta, ids.get(i), elementType);
+        } finally {
+          ObjectSerializer.resetOuterObjectId();
+        }
       }
       if (object instanceof RayException) {
         // If the object is a `RayException`, it means that an error occurred during task
@@ -181,4 +186,27 @@ public abstract class ObjectStore {
    * @param objectId The object ID to decrease the reference count for.
    */
   public abstract void removeLocalReference(UniqueId workerId, ObjectId objectId);
+
+  /**
+   * Promote the given object to the underlying object store, and get the ownership info.
+   *
+   * @param objectId The ID of the object to promote
+   * @return the serialized ownership address
+   */
+  public abstract byte[] promoteAndGetOwnershipInfo(ObjectId objectId);
+
+  /**
+   * Add a reference to an ObjectID that will deserialized. This will also start the process to
+   * resolve the future. Specifically, we will periodically contact the owner, until we learn that
+   * the object has been created or the owner is no longer reachable. This will then unblock any
+   * Gets or submissions of tasks dependent on the object.
+   *
+   * @param objectId The object ID to deserialize.
+   * @param outerObjectId The object ID that contained objectId, if any. This may be nil if the
+   *                      object ID was inlined directly in a task spec or if it was passed
+   *                      out-of-band by the application (deserialized from a byte string).
+   * @param ownerAddress The address of the object's owner.
+   */
+  public abstract void registerOwnershipInfoAndResolveFuture(ObjectId objectId,
+      ObjectId outerObjectId, byte[] ownerAddress);
 }