[Java] Local and distributed ref counting in Java (#9371)

java/test/src/main/java/io/ray/test/ActorTest.java

@@ -42,6 +42,10 @@ public class ActorTest extends BaseTest {
       value += largeObject.data.length;
       return value;
     }
+
+    public TestUtils.LargeObject createLargeObject() {
+      return new TestUtils.LargeObject();
+    }
   }
 
   public void testCreateAndCallActor() {
@@ -68,8 +72,7 @@ public class ActorTest extends BaseTest {
     ObjectRef<Integer> result = actor.task(Counter::getValue).remote();
     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));
+    Assert.assertEquals(Ray.get(result), Integer.valueOf(1));
   }
 
   public void testCallActorWithLargeObject() {
@@ -117,33 +120,28 @@ public class ActorTest extends BaseTest {
             Collections.singletonList(actor), 100).remote().get());
   }
 
-  // TODO(qwang): Will re-enable this test case once ref counting is supported in Java.
-  @Test(enabled = false, groups = {"cluster"})
+  // This test case follows `test_internal_free` in `python/ray/tests/test_advanced.py`.
+  @Test(groups = {"cluster"})
   public void testUnreconstructableActorObject() throws InterruptedException {
-    // The UnreconstructableException is created by raylet.
     ActorHandle<Counter> counter = Ray.actor(Counter::new, 100).remote();
     // Call an actor method.
     ObjectRef value = counter.task(Counter::getValue).remote();
     Assert.assertEquals(100, value.get());
     // Delete the object from the object store.
-    Ray.internal().free(ImmutableList.of(value.getId()), false, false);
-    // Wait until the object is deleted, because the above free operation is async.
-    while (true) {
-      Boolean result = TestUtils.getRuntime().getObjectStore()
-          .wait(ImmutableList.of(value.getId()), 1, 0).get(0);
-      if (!result) {
-        break;
-      }
-      TimeUnit.MILLISECONDS.sleep(100);
-    }
+    Ray.internal().free(ImmutableList.of(value), false, false);
+    // Wait for delete RPC to propagate
+    TimeUnit.SECONDS.sleep(1);
+    // Free deletes from in-memory store.
+    Assert.expectThrows(UnreconstructableException.class, () -> value.get());
 
-    try {
-      // Try getting the object again, this should throw an UnreconstructableException.
-      // Use `Ray.get()` to bypass the cache in `RayObjectImpl`.
-      Ray.get(value.getId(), value.getType());
-      Assert.fail("This line should not be reachable.");
-    } catch (UnreconstructableException e) {
-      Assert.assertEquals(value.getId(), e.objectId);
-    }
+    // Call an actor method.
+    ObjectRef<TestUtils.LargeObject> largeValue = counter.task(Counter::createLargeObject).remote();
+    Assert.assertTrue(largeValue.get() instanceof TestUtils.LargeObject);
+    // Delete the object from the object store.
+    Ray.internal().free(ImmutableList.of(largeValue), false, false);
+    // Wait for delete RPC to propagate
+    TimeUnit.SECONDS.sleep(1);
+    // Free deletes big objects from plasma store.
+    Assert.expectThrows(UnreconstructableException.class, () -> largeValue.get());
   }
 }

java/test/src/main/java/io/ray/test/BaseMultiLanguageTest.java

@@ -51,7 +51,7 @@ public abstract class BaseMultiLanguageTest {
     }
   }
 
-  @BeforeClass(alwaysRun = true)
+  @BeforeClass(alwaysRun = true, inheritGroups = false)
   public void setUp() {
     // Delete existing socket files.
     for (String socket : ImmutableList.of(RAYLET_SOCKET_NAME, PLASMA_STORE_SOCKET_NAME)) {
@@ -106,7 +106,7 @@ public abstract class BaseMultiLanguageTest {
     return ImmutableMap.of();
   }
 
-  @AfterClass(alwaysRun = true)
+  @AfterClass(alwaysRun = true, inheritGroups = false)
   public void tearDown() {
     // Disconnect to the cluster.
     Ray.shutdown();

java/test/src/main/java/io/ray/test/GcsClientTest.java

@@ -33,7 +33,7 @@ public class GcsClientTest extends BaseTest {
     Assert.assertEquals(allNodeInfo.size(), 1);
     Assert.assertEquals(allNodeInfo.get(0).nodeAddress, config.nodeIp);
     Assert.assertTrue(allNodeInfo.get(0).isAlive);
-    Assert.assertEquals(allNodeInfo.get(0).resources.get("A"), 8.0);
+    Assert.assertEquals((double) allNodeInfo.get(0).resources.get("A"), 8.0);
   }
 
   @Test

java/test/src/main/java/io/ray/test/GlobalGcTest.java

@@ -0,0 +1,96 @@
+package io.ray.test;
+
+import io.ray.api.ActorHandle;
+import io.ray.api.ObjectRef;
+import io.ray.api.Ray;
+import java.lang.ref.WeakReference;
+import java.util.List;
+import java.util.stream.Collectors;
+import java.util.stream.IntStream;
+import org.testng.Assert;
+import org.testng.annotations.AfterClass;
+import org.testng.annotations.BeforeClass;
+import org.testng.annotations.Test;
+
+@Test(groups = {"cluster"})
+public class GlobalGcTest extends BaseTest {
+
+  @BeforeClass
+  public void setUp() {
+    System.setProperty("ray.object-store.size", "140 MB");
+  }
+
+  @AfterClass
+  public void tearDown() {
+    System.clearProperty("ray.object-store.size");
+  }
+
+  public static class LargeObjectWithCyclicRef {
+
+    private final LargeObjectWithCyclicRef loop;
+
+    private final ObjectRef<TestUtils.LargeObject> largeObject;
+
+    public LargeObjectWithCyclicRef() {
+      this.loop = this;
+      this.largeObject = Ray.put(new TestUtils.LargeObject(40 * 1024 * 1024));
+    }
+  }
+
+  public static class GarbageHolder {
+
+    private WeakReference<LargeObjectWithCyclicRef> garbage;
+
+    public GarbageHolder() {
+      LargeObjectWithCyclicRef x = new LargeObjectWithCyclicRef();
+      garbage = new WeakReference<>(x);
+    }
+
+    public boolean hasGarbage() {
+      return garbage.get() != null;
+    }
+
+    public TestUtils.LargeObject returnLargeObject() {
+      return new TestUtils.LargeObject(80 * 1024 * 1024);
+    }
+  }
+
+  private void testGlobalGcWhenFull(boolean withPut) {
+    // Local driver.
+    WeakReference<LargeObjectWithCyclicRef> localRef = new WeakReference<>(
+        new LargeObjectWithCyclicRef());
+
+    // Remote workers.
+    List<ActorHandle<GarbageHolder>> actors = IntStream
+        .range(0, 2).mapToObj(i -> Ray.actor(GarbageHolder::new).remote())
+        .collect(Collectors.toList());
+
+    Assert.assertNotNull(localRef.get());
+    for (ActorHandle<GarbageHolder> actor : actors) {
+      Assert.assertTrue(actor.task(GarbageHolder::hasGarbage).remote().get());
+    }
+
+    if (withPut) {
+      // GC should be triggered for all workers, including the local driver,
+      // when the driver tries to Ray.put a value that doesn't fit in the
+      // object store. This should cause the captured ObjectRefs to be evicted.
+      Ray.put(new TestUtils.LargeObject(80 * 1024 * 1024));
+    } else {
+      // GC should be triggered for all workers, including the local driver,
+      // when a remote task tries to put a return value that doesn't fit in
+      // the object store. This should cause the captured ObjectRefs' to be evicted.
+      actors.get(0).task(GarbageHolder::returnLargeObject).remote().get();
+    }
+
+    TestUtils.waitForCondition(() -> localRef.get() == null && actors.stream().noneMatch(
+        a -> a.task(GarbageHolder::hasGarbage).remote().get()), 10 * 1000);
+  }
+
+  public void testGlobalGcWhenFullWithPut() {
+    testGlobalGcWhenFull(true);
+  }
+
+  public void testGlobalGcWhenFullWithReturn() {
+    testGlobalGcWhenFull(false);
+  }
+}

java/test/src/main/java/io/ray/test/MultiThreadingTest.java

@@ -143,7 +143,7 @@ public class MultiThreadingTest extends BaseTest {
     final ActorHandle<Echo> fooActor = Ray.actor(Echo::new).remote();
     return new Runnable[]{
         () -> Ray.put(1),
-        () -> Ray.get(fooObject.getId(), fooObject.getType()),
+        () -> Ray.get(ImmutableList.of(fooObject)),
         fooObject::get,
         () -> Ray.wait(ImmutableList.of(fooObject)),
         Ray::getRuntimeContext,

java/test/src/main/java/io/ray/test/ObjectStoreTest.java

@@ -3,7 +3,6 @@ package io.ray.test;
 import com.google.common.collect.ImmutableList;
 import io.ray.api.ObjectRef;
 import io.ray.api.Ray;
-import io.ray.api.id.ObjectId;
 import java.util.List;
 import java.util.stream.Collectors;
 import org.testng.Assert;
@@ -37,8 +36,8 @@ public class ObjectStoreTest extends BaseTest {
   @Test
   public void testGetMultipleObjects() {
     List<Integer> ints = ImmutableList.of(1, 2, 3, 4, 5);
-    List<ObjectId> ids = ints.stream().map(obj -> Ray.put(obj).getId())
+    List<ObjectRef<Integer>> refs = ints.stream().map(Ray::put)
         .collect(Collectors.toList());
-    Assert.assertEquals(ints, Ray.get(ids, Integer.class));
+    Assert.assertEquals(ints, Ray.get(refs));
   }
 }

java/test/src/main/java/io/ray/test/PlasmaFreeTest.java

@@ -4,6 +4,7 @@ import com.google.common.collect.ImmutableList;
 import io.ray.api.ObjectRef;
 import io.ray.api.Ray;
 import io.ray.api.id.TaskId;
+import io.ray.runtime.object.ObjectRefImpl;
 import java.util.Arrays;
 import org.testng.Assert;
 import org.testng.annotations.Test;
@@ -19,11 +20,11 @@ public class PlasmaFreeTest extends BaseTest {
     ObjectRef<String> helloId = Ray.task(PlasmaFreeTest::hello).remote();
     String helloString = helloId.get();
     Assert.assertEquals("hello", helloString);
-    Ray.internal().free(ImmutableList.of(helloId.getId()), true, false);
+    Ray.internal().free(ImmutableList.of(helloId), true, false);
 
     final boolean result = TestUtils.waitForCondition(() ->
         !TestUtils.getRuntime().getObjectStore()
-          .wait(ImmutableList.of(helloId.getId()), 1, 0).get(0), 50);
+          .wait(ImmutableList.of(((ObjectRefImpl<String>) helloId).getId()), 1, 0).get(0), 50);
     if (TestUtils.isSingleProcessMode()) {
       Assert.assertTrue(result);
     } else {
@@ -36,9 +37,10 @@ public class PlasmaFreeTest extends BaseTest {
   public void testDeleteCreatingTasks() {
     ObjectRef<String> helloId = Ray.task(PlasmaFreeTest::hello).remote();
     Assert.assertEquals("hello", helloId.get());
-    Ray.internal().free(ImmutableList.of(helloId.getId()), true, true);
+    Ray.internal().free(ImmutableList.of(helloId), true, true);
 
-    TaskId taskId = TaskId.fromBytes(Arrays.copyOf(helloId.getId().getBytes(), TaskId.LENGTH));
+    TaskId taskId = TaskId.fromBytes(
+        Arrays.copyOf(((ObjectRefImpl<String>) helloId).getId().getBytes(), TaskId.LENGTH));
     final boolean result = TestUtils.waitForCondition(
         () -> !TestUtils.getRuntime().getGcsClient()
             .rayletTaskExistsInGcs(taskId), 50);

java/test/src/main/java/io/ray/test/PlasmaStoreTest.java

@@ -1,21 +0,0 @@
-package io.ray.test;
-
-import io.ray.api.Ray;
-import io.ray.api.id.ObjectId;
-import io.ray.runtime.object.ObjectStore;
-import org.testng.Assert;
-import org.testng.annotations.Test;
-
-public class PlasmaStoreTest extends BaseTest {
-
-  @Test(groups = {"cluster"})
-  public void testPutWithDuplicateId() {
-    ObjectId objectId = ObjectId.fromRandom();
-    ObjectStore objectStore = TestUtils.getRuntime().getObjectStore();
-    objectStore.put("1", objectId);
-    Assert.assertEquals(Ray.get(objectId, String.class), "1");
-    objectStore.put("2", objectId);
-    // Putting the second object with duplicate ID should fail but ignored.
-    Assert.assertEquals(Ray.get(objectId, String.class), "1");
-  }
-}

java/test/src/main/java/io/ray/test/ReferenceCountingTest.java

@@ -0,0 +1,362 @@
+package io.ray.test;
+
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
+import com.google.gson.Gson;
+import io.ray.api.ActorHandle;
+import io.ray.api.ObjectRef;
+import io.ray.api.Ray;
+import io.ray.api.id.ObjectId;
+import io.ray.runtime.object.NativeObjectStore;
+import io.ray.runtime.object.ObjectRefImpl;
+import java.lang.ref.Reference;
+import java.lang.reflect.Field;
+import java.lang.reflect.Method;
+import java.time.Duration;
+import java.time.Instant;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import org.testng.Assert;
+import org.testng.annotations.AfterClass;
+import org.testng.annotations.BeforeClass;
+import org.testng.annotations.BeforeMethod;
+import org.testng.annotations.Test;
+
+@Test(groups = {"cluster"})
+public class ReferenceCountingTest extends BaseTest {
+  @BeforeClass
+  public void setUp() {
+    System.setProperty("ray.object-store.size", "100 MB");
+  }
+
+  @AfterClass
+  public void tearDown() {
+    System.clearProperty("ray.object-store.size");
+  }
+
+  /**
+   * Because we can't explicitly GC an Java object. We use this helper method to manually remove an
+   * local reference.
+   */
+  private static void del(ObjectRef<?> obj) {
+    try {
+      Field referencesField = ObjectRefImpl.class.getDeclaredField("REFERENCES");
+      referencesField.setAccessible(true);
+      Set<?> references = (Set<?>) referencesField.get(null);
+      Class<?> referenceClass = Class
+          .forName("io.ray.runtime.object.ObjectRefImpl$ObjectRefImplReference");
+      Method finalizeReferentMethod = referenceClass.getDeclaredMethod("finalizeReferent");
+      finalizeReferentMethod.setAccessible(true);
+      for (Object reference : references) {
+        if (obj.equals(((Reference<?>)reference).get())) {
+          finalizeReferentMethod.invoke(reference);
+          break;
+        }
+      }
+    } catch (Exception e) {
+      throw new RuntimeException(e);
+    }
+  }
+
+  private void checkRefCounts(Map<ObjectId, long[]> expected, Duration timeout) {
+    Instant start = Instant.now();
+    while (true) {
+      Map<ObjectId, long[]> actual =
+          ((NativeObjectStore) TestUtils.getRuntime().getObjectStore()).getAllReferenceCounts();
+      try {
+        Assert.assertEqualsDeep(actual, expected);
+        return;
+      } catch (AssertionError e) {
+        if (Duration.between(start, Instant.now()).compareTo(timeout) >= 0) {
+          System.out.println("Actual: " + new Gson().toJson(actual));
+          System.out.println("Expected: " + new Gson().toJson(expected));
+          throw e;
+        } else {
+          try {
+            Thread.sleep(100);
+          } catch (InterruptedException ex) {
+            throw new RuntimeException(e);
+          }
+        }
+      }
+    }
+  }
+
+  private void checkRefCounts(Map<ObjectId, long[]> expected) {
+    checkRefCounts(expected, Duration.ofSeconds(10));
+  }
+
+  private void checkRefCounts(ObjectId objectId, long localRefCount, long submittedTaskRefCount) {
+    checkRefCounts(ImmutableMap.of(objectId, new long[] {localRefCount, submittedTaskRefCount}));
+  }
+
+  private void checkRefCounts(ObjectId objectId1, long localRefCount1, long submittedTaskRefCount1,
+      ObjectId objectId2, long localRefCount2, long submittedTaskRefCount2) {
+    checkRefCounts(ImmutableMap.of(objectId1, new long[] {localRefCount1, submittedTaskRefCount1},
+        objectId2, new long[] {localRefCount2, submittedTaskRefCount2}));
+  }
+
+  private static void fillObjectStoreAndGet(ObjectId objectId, boolean succeed) {
+    fillObjectStoreAndGet(objectId, succeed, 40 * 1024 * 1024, 5);
+  }
+
+  private static void fillObjectStoreAndGet(ObjectId objectId, boolean succeed, int objectSize,
+      int numObjects) {
+    for (int i = 0; i < numObjects; i++) {
+      Ray.put(new TestUtils.LargeObject(objectSize));
+    }
+    if (succeed) {
+      TestUtils.getRuntime().getObjectStore().getRaw(ImmutableList.of(objectId), Long.MAX_VALUE);
+    } else {
+      List<Boolean> result =
+          TestUtils.getRuntime().getObjectStore().wait(ImmutableList.of(objectId), 1, 100);
+      Assert.assertFalse(result.get(0));
+    }
+  }
+
+  /**
+   * Based on Python test case `test_local_refcounts`.
+   */
+  public void testLocalRefCounts() {
+    ObjectRefImpl<Object> obj1 = (ObjectRefImpl<Object>) Ray.put(null);
+    checkRefCounts(obj1.getId(), 1, 0);
+    ObjectRef<Object> obj1Copy = new ObjectRefImpl<>(obj1.getId(), obj1.getType());
+    checkRefCounts(obj1.getId(), 2, 0);
+
+    del(obj1);
+    checkRefCounts(obj1.getId(), 1, 0);
+    del(obj1Copy);
+    checkRefCounts(ImmutableMap.of());
+  }
+
+  private static int oneDep(Object obj) {
+    return oneDep(obj, null);
+  }
+
+  private static int oneDep(Object obj, ActorHandle<SignalActor> singal) {
+    return oneDep(obj, singal, false);
+  }
+
+  private static int oneDep(Object obj, ActorHandle<SignalActor> singal, boolean fail) {
+    if (singal != null) {
+      singal.task(SignalActor::waitSignal).remote().get();
+    }
+    if (fail) {
+      throw new RuntimeException("failed on purpose");
+    }
+    return 0;
+  }
+
+  private static TestUtils.LargeObject oneDepLarge(Object obj, ActorHandle<SignalActor> singal) {
+    if (singal != null) {
+      singal.task(SignalActor::waitSignal).remote().get();
+    }
+    // This will be spilled to plasma.
+    return new TestUtils.LargeObject(10 * 1024 * 1024);
+  }
+
+  private static void sendSignal(ActorHandle<SignalActor> signal) {
+    ObjectRef<Integer> result = signal.task(SignalActor::sendSignal).remote();
+    result.get();
+    // Remove the reference immediately, otherwise it will affect subsequent tests.
+    del(result);
+  }
+
+  /**
+   * Based on Python test case `test_dependency_refcounts`.
+   */
+  public void testDependencyRefCounts() {
+    {
+      // Test that regular plasma dependency refcounts are decremented once the
+      // task finishes.
+      ActorHandle<SignalActor> signal = SignalActor.create();
+      ObjectRefImpl<TestUtils.LargeObject> largeDep = (ObjectRefImpl<TestUtils.LargeObject>) Ray
+          .put(new TestUtils.LargeObject());
+      ObjectRefImpl<Object> result = (ObjectRefImpl<Object>)
+          Ray.<TestUtils.LargeObject, ActorHandle<SignalActor>, Object>task(
+              ReferenceCountingTest::oneDep, largeDep, signal).remote();
+      checkRefCounts(largeDep.getId(), 1, 1, result.getId(), 1, 0);
+      sendSignal(signal);
+      // Reference count should be removed once the task finishes.
+      checkRefCounts(largeDep.getId(), 1, 0, result.getId(), 1, 0);
+      del(largeDep);
+      del(result);
+      checkRefCounts(ImmutableMap.of());
+    }
+
+    {
+      // Test that inlined dependency refcounts are decremented once they are
+      // inlined.
+      ActorHandle<SignalActor> signal = SignalActor.create();
+      ObjectRefImpl<Integer> dep = (ObjectRefImpl<Integer>)
+          Ray.<Integer, ActorHandle<SignalActor>, Integer>task(ReferenceCountingTest::oneDep,
+              Integer.valueOf(1), signal).remote();
+      checkRefCounts(dep.getId(), 1, 0);
+      ObjectRefImpl<Object> result = (ObjectRefImpl<Object>)
+          Ray.<Integer, Object>task(ReferenceCountingTest::oneDep, dep).remote();
+      checkRefCounts(dep.getId(), 1, 1, result.getId(), 1, 0);
+      sendSignal(signal);
+      // Reference count should be removed as soon as the dependency is inlined.
+      checkRefCounts(dep.getId(), 1, 0, result.getId(), 1, 0);
+      del(dep);
+      del(result);
+      checkRefCounts(ImmutableMap.of());
+    }
+
+    {
+      // Test that spilled plasma dependency refcounts are decremented once
+      // the task finishes.
+      ActorHandle<SignalActor> signal1 = SignalActor.create();
+      ActorHandle<SignalActor> signal2 = SignalActor.create();
+      ObjectRefImpl<TestUtils.LargeObject> dep = (ObjectRefImpl<TestUtils.LargeObject>)
+          Ray.<TestUtils.LargeObject, ActorHandle<SignalActor>, TestUtils.LargeObject>task(
+              ReferenceCountingTest::oneDepLarge, (TestUtils.LargeObject) null, signal1).remote();
+      checkRefCounts(dep.getId(), 1, 0);
+      ObjectRefImpl<Integer> result = (ObjectRefImpl<Integer>)
+          Ray.<TestUtils.LargeObject, ActorHandle<SignalActor>, Integer>task(
+              ReferenceCountingTest::oneDep, dep, signal2).remote();
+      checkRefCounts(dep.getId(), 1, 1, result.getId(), 1, 0);
+      sendSignal(signal1);
+      dep.get(); // TODO(kfstorm): timeout=10
+      // Reference count should remain because the dependency is in plasma.
+      checkRefCounts(dep.getId(), 1, 1, result.getId(), 1, 0);
+      sendSignal(signal2);
+      // Reference count should be removed because the task finished.
+      checkRefCounts(dep.getId(), 1, 0, result.getId(), 1, 0);
+      del(dep);
+      del(result);
+      checkRefCounts(ImmutableMap.of());
+    }
+
+    {
+      // Test that regular plasma dependency refcounts are decremented if a task
+      // fails.
+      ActorHandle<SignalActor> signal = SignalActor.create();
+      ObjectRefImpl<TestUtils.LargeObject> largeDep = (ObjectRefImpl<TestUtils.LargeObject>)
+          Ray.put(new TestUtils.LargeObject(10 * 1024 * 1024));
+      ObjectRefImpl<Integer> result = (ObjectRefImpl<Integer>)
+          Ray.<TestUtils.LargeObject, ActorHandle<SignalActor>, Boolean, Integer>task(
+              ReferenceCountingTest::oneDep, largeDep, signal, /* fail= */true).remote();
+      checkRefCounts(largeDep.getId(), 1, 1, result.getId(), 1, 0);
+      sendSignal(signal);
+      // Reference count should be removed once the task finishes.
+      checkRefCounts(largeDep.getId(), 1, 0, result.getId(), 1, 0);
+      del(largeDep);
+      del(result);
+      checkRefCounts(ImmutableMap.of());
+    }
+
+    {
+      // Test that spilled plasma dependency refcounts are decremented if a task
+      // fails.
+      ActorHandle<SignalActor> signal1 = SignalActor.create();
+      ActorHandle<SignalActor> signal2 = SignalActor.create();
+      ObjectRefImpl<TestUtils.LargeObject> dep = (ObjectRefImpl<TestUtils.LargeObject>)
+          Ray.<Integer, ActorHandle<SignalActor>, TestUtils.LargeObject>task(
+              ReferenceCountingTest::oneDepLarge, (Integer) null, signal1).remote();
+      checkRefCounts(dep.getId(), 1, 0);
+      ObjectRefImpl<Integer> result = (ObjectRefImpl<Integer>)
+          Ray.<TestUtils.LargeObject, ActorHandle<SignalActor>, Boolean, Integer>task(
+              ReferenceCountingTest::oneDep, dep, signal2, /* fail= */true).remote();
+      checkRefCounts(dep.getId(), 1, 1, result.getId(), 1, 0);
+      sendSignal(signal1);
+      dep.get(); // TODO(kfstorm): timeout=10
+      // Reference count should remain because the dependency is in plasma.
+      checkRefCounts(dep.getId(), 1, 1, result.getId(), 1, 0);
+      sendSignal(signal2);
+      // Reference count should be removed because the task finished.
+      checkRefCounts(dep.getId(), 1, 0, result.getId(), 1, 0);
+      del(dep);
+      del(result);
+      checkRefCounts(ImmutableMap.of());
+    }
+  }
+
+  private static int fooBasicPinning(Object arg) {
+    return 0;
+  }
+
+  public static class ActorBasicPinning {
+    private ObjectRef<TestUtils.LargeObject> largeObject;
+
+    public ActorBasicPinning() {
+      // Hold a long-lived reference to a ray.put object's ID. The object
+      // should not be garbage collected while the actor is alive because
+      // the object is pinned by the raylet.
+      largeObject = Ray.put(new TestUtils.LargeObject(25 * 1024 * 1024));
+    }
+
+    public TestUtils.LargeObject getLargeObject() {
+      return largeObject.get();
+    }
+  }
+
+  /**
+   * Based on Python test case `test_basic_pinning`.
+   */
+  public void testBasicPinning() {
+    ActorHandle<ActorBasicPinning> actor = Ray.actor(ActorBasicPinning::new).remote();
+
+    // Fill up the object store with short-lived objects. These should be
+    // evicted before the long-lived object whose reference is held by
+    // the actor.
+    for (int i = 0; i < 10; i++) {
+      ObjectRef<Integer> intermediateResult = Ray
+          .task(ReferenceCountingTest::fooBasicPinning, new TestUtils.LargeObject(10 * 1024 * 1024))
+          .remote();
+      intermediateResult.get();
+    }
+    // The ray.get below would fail with only LRU eviction, as the object
+    // that was ray.put by the actor would have been evicted.
+    actor.task(ActorBasicPinning::getLargeObject).remote().get();
+  }
+
+  private static Object pending(TestUtils.LargeObject input1, Integer input2) {
+    return null;
+  }
+
+  /**
+   * Based on Python test case `test_pending_task_dependency_pinning`.
+   */
+  public void testPendingTaskDependencyPinning() {
+    // The object that is ray.put here will go out of scope immediately, so if
+    // pending task dependencies aren't considered, it will be evicted before
+    // the ray.get below due to the subsequent ray.puts that fill up the object
+    // store.
+    TestUtils.LargeObject input1 = new TestUtils.LargeObject(40 * 1024 * 1024);
+    ActorHandle<SignalActor> signal = SignalActor.create();
+    ObjectRef<Object> result = Ray
+        .task(ReferenceCountingTest::pending, input1, signal.task(SignalActor::waitSignal).remote())
+        .remote();
+
+    for (int i = 0; i < 2; i++) {
+      Ray.put(new TestUtils.LargeObject(40 * 1024 * 1024));
+    }
+
+    sendSignal(signal);
+    result.get();
+  }
+
+  /**
+   * Test that an object containing object IDs within it pins the inner IDs. Based on Python test
+   * case `test_basic_nested_ids`.
+   */
+  public void testBasicNestedIds() {
+    ObjectRefImpl<byte[]> inner = (ObjectRefImpl<byte[]>) Ray.put(new byte[40 * 1024 * 1024]);
+    ObjectRef<List<ObjectRef<byte[]>>> outer = Ray.put(Collections.singletonList(inner));
+
+    // Remove the local reference to the inner object.
+    del(inner);
+
+    // Check that the outer reference pins the inner object.
+    fillObjectStoreAndGet(inner.getId(), true);
+
+    // Remove the outer reference and check that the inner object gets evicted.
+    del(outer);
+    fillObjectStoreAndGet(inner.getId(), false);
+  }
+
+  // TODO(kfstorm): Add more test cases
+}

java/test/src/main/java/io/ray/test/SignalActor.java

@@ -0,0 +1,28 @@
+package io.ray.test;
+
+import io.ray.api.ActorHandle;
+import io.ray.api.Ray;
+import java.util.concurrent.Semaphore;
+
+public class SignalActor {
+
+  private Semaphore semaphore;
+
+  public SignalActor() {
+    this.semaphore = new Semaphore(0);
+  }
+
+  public int sendSignal() {
+    this.semaphore.release();
+    return 0;
+  }
+
+  public int waitSignal() throws InterruptedException {
+    this.semaphore.acquire();
+    return 0;
+  }
+
+  public static ActorHandle<SignalActor> create() {
+    return Ray.actor(SignalActor::new).setMaxConcurrency(2).remote();
+  }
+}

java/test/src/main/java/io/ray/test/StressTest.java

@@ -4,7 +4,6 @@ import com.google.common.collect.ImmutableList;
 import io.ray.api.ActorHandle;
 import io.ray.api.ObjectRef;
 import io.ray.api.Ray;
-import io.ray.api.id.ObjectId;
 import java.util.ArrayList;
 import java.util.List;
 import org.testng.Assert;
@@ -21,12 +20,12 @@ public class StressTest extends BaseTest {
     for (int numIterations : ImmutableList.of(1, 10, 100, 1000)) {
       int numTasks = 1000 / numIterations;
       for (int i = 0; i < numIterations; i++) {
-        List<ObjectId> resultIds = new ArrayList<>();
+        List<ObjectRef<Integer>> results = new ArrayList<>();
         for (int j = 0; j < numTasks; j++) {
-          resultIds.add(Ray.task(StressTest::echo, 1).remote().getId());
+          results.add(Ray.task(StressTest::echo, 1).remote());
         }
 
-        for (Integer result : Ray.<Integer>get(resultIds, Integer.class)) {
+        for (Integer result : Ray.get(results)) {
           Assert.assertEquals(result, Integer.valueOf(1));
         }
       }
@@ -58,12 +57,12 @@ public class StressTest extends BaseTest {
     }
 
     public int ping(int n) {
-      List<ObjectId> objectIds = new ArrayList<>();
+      List<ObjectRef<Integer>> objectRefs = new ArrayList<>();
       for (int i = 0; i < n; i++) {
-        objectIds.add(actor.task(Actor::ping).remote().getId());
+        objectRefs.add(actor.task(Actor::ping).remote());
       }
       int sum = 0;
-      for (Integer result : Ray.<Integer>get(objectIds, Integer.class)) {
+      for (Integer result : Ray.get(objectRefs)) {
         sum += result;
       }
       return sum;
@@ -72,13 +71,13 @@ public class StressTest extends BaseTest {
 
   public void testSubmittingManyTasksToOneActor() throws Exception {
     ActorHandle<Actor> actor = Ray.actor(Actor::new).remote();
-    List<ObjectId> objectIds = new ArrayList<>();
+    List<ObjectRef<Integer>> objectRefs = new ArrayList<>();
     for (int i = 0; i < 10; i++) {
       ActorHandle<Worker> worker = Ray.actor(Worker::new, actor).remote();
-      objectIds.add(worker.task(Worker::ping, 100).remote().getId());
+      objectRefs.add(worker.task(Worker::ping, 100).remote());
     }
 
-    for (Integer result : Ray.<Integer>get(objectIds, Integer.class)) {
+    for (Integer result : Ray.get(objectRefs)) {
       Assert.assertEquals(result, Integer.valueOf(100));
     }
   }

java/test/src/main/java/io/ray/test/TestUtils.java

@@ -1,10 +1,12 @@
 package io.ray.test;
 
+import com.google.common.base.Preconditions;
 import io.ray.api.ObjectRef;
 import io.ray.api.Ray;
 import io.ray.runtime.RayRuntimeInternal;
 import io.ray.runtime.RayRuntimeProxy;
 import io.ray.runtime.config.RunMode;
+import io.ray.runtime.task.ArgumentsBuilder;
 import java.io.Serializable;
 import java.util.function.Supplier;
 import org.testng.Assert;
@@ -13,7 +15,16 @@ public class TestUtils {
 
   public static class LargeObject implements Serializable {
 
-    public byte[] data = new byte[1024 * 1024];
+    public byte[] data;
+
+    public LargeObject() {
+      this(1024 * 1024);
+    }
+
+    public LargeObject(int size) {
+      Preconditions.checkState(size > ArgumentsBuilder.LARGEST_SIZE_PASS_BY_VALUE);
+      data = new byte[size];
+    }
   }
 
   private static final int WAIT_INTERVAL_MS = 5;
@@ -55,11 +66,12 @@ public class TestUtils {
 
   /**
    * Warm up the cluster to make sure there's at least one idle worker.
-   * <p>
-   * This is needed before calling `wait`. Because, in Travis CI, starting a new worker
-   * process could be slower than the wait timeout.
-   * TODO(hchen): We should consider supporting always reversing a certain number of
-   * idle workers in Raylet's worker pool.
+   * <p/>
+   * This is needed before calling `wait`. Because, in Travis CI, starting a new worker process
+   * could be slower than the wait timeout.
+   * <p/>
+   * TODO(hchen): We should consider supporting always reversing a certain number of idle workers in
+   * Raylet's worker pool.
    */
   public static void warmUpCluster() {
     ObjectRef<String> obj = Ray.task(TestUtils::hi).remote();