From 5cc3e1341ad1d2e59b484e19df8e64144b7b7955 Mon Sep 17 00:00:00 2001
From: Hao Chen <chenh1024@gmail.com>
Date: Wed, 11 Dec 2019 11:26:01 +0800
Subject: [PATCH] [Java] Cache result in RayObjectImpl (#6414)

---
 .../org/ray/runtime/object/RayObjectImpl.java | 23 +++++++++++++++++--
 .../main/java/org/ray/api/test/ActorTest.java | 22 +++++++++++++++++-
 2 files changed, 42 insertions(+), 3 deletions(-)

diff --git a/java/runtime/src/main/java/org/ray/runtime/object/RayObjectImpl.java b/java/runtime/src/main/java/org/ray/runtime/object/RayObjectImpl.java
index 197445831..dce123727 100644
--- a/java/runtime/src/main/java/org/ray/runtime/object/RayObjectImpl.java
+++ b/java/runtime/src/main/java/org/ray/runtime/object/RayObjectImpl.java
@@ -12,13 +12,32 @@ public final class RayObjectImpl<T> implements RayObject<T>, Serializable {
 
   private final ObjectId id;
 
+  /**
+   * Cache the result of `Ray.get()`.
+   *
+   * Note, this is necessary for direct calls, in which case, it's not allowed to call `Ray.get` on
+   * the same object twice.
+   */
+  private T object;
+
+  /**
+   * Whether the object is already gotten from the object store.
+   */
+  private boolean objectGotten;
+
   public RayObjectImpl(ObjectId id) {
     this.id = id;
+    object = null;
+    objectGotten = false;
   }
 
   @Override
-  public T get() {
-    return Ray.get(id);
+  public synchronized T get() {
+    if (!objectGotten) {
+      object = Ray.get(id);
+      objectGotten = true;
+    }
+    return object;
   }
 
   @Override
diff --git a/java/test/src/main/java/org/ray/api/test/ActorTest.java b/java/test/src/main/java/org/ray/api/test/ActorTest.java
index 500b9c31b..3bf3c79f2 100644
--- a/java/test/src/main/java/org/ray/api/test/ActorTest.java
+++ b/java/test/src/main/java/org/ray/api/test/ActorTest.java
@@ -59,6 +59,25 @@ public class ActorTest extends BaseTest {
     Assert.assertEquals(Integer.valueOf(3), Ray.call(Counter::increaseAndGet, actor, 1).get());
   }
 
+  /**
+   * Test getting a direct object (an object that is returned by a direct-call task) twice from the
+   * object store.
+   *
+   * Direct objects are stored in core worker's local memory. And it will be removed after the first
+   * get. To enable getting it twice, we cache the object in `RayObjectImpl`.
+   *
+   * NOTE(hchen): this test will run for non-direct actors as well, which doesn't have the above
+   * issue and should also succeed.
+   */
+  public void testGetDirectObjectTwice() {
+    RayActor<Counter> actor = Ray.createActor(Counter::new, 1);
+    RayObject<Integer> result = Ray.call(Counter::getValue, actor);
+    Assert.assertEquals(result.get(), Integer.valueOf(1));
+    Assert.assertEquals(result.get(), Integer.valueOf(1));
+    // TODO(hchen): The following code will still fail, and can be fixed by using ref counting.
+    // Assert.assertEquals(Ray.get(result.getId()), Integer.valueOf(1));
+  }
+
   public void testCallActorWithLargeObject() {
     RayActor<Counter> actor = Ray.createActor(Counter::new, 1);
     LargeObject largeObject = new LargeObject();
@@ -131,7 +150,8 @@ public class ActorTest extends BaseTest {
 
     try {
       // Try getting the object again, this should throw an UnreconstructableException.
-      value.get();
+      // Use `Ray.get()` to bypass the cache in `RayObjectImpl`.
+      Ray.get(value.getId());
       Assert.fail("This line should not be reachable.");
     } catch (UnreconstructableException e) {
       Assert.assertEquals(value.getId(), e.objectId);