Add delete_creating_tasks option for internal.free() (#4588)

* add delete creating task objects.

* format code style

* Fix lint

* add tests add address comments.

* Refine test

* Refine java test

* Fix CI

* Refine

* Fix lint

* Fix CI

java/test/src/main/java/org/ray/api/TestUtils.java

@@ -1,15 +1,45 @@
 package org.ray.api;
 
+import java.util.function.Supplier;
 import org.ray.runtime.AbstractRayRuntime;
 import org.ray.runtime.config.RunMode;
 import org.testng.SkipException;
 
 public class TestUtils {
 
+  private static final int WAIT_INTERVAL_MS = 5;
+
   public static void skipTestUnderSingleProcess() {
     AbstractRayRuntime runtime = (AbstractRayRuntime)Ray.internal();
     if (runtime.getRayConfig().runMode == RunMode.SINGLE_PROCESS) {
       throw new SkipException("This test doesn't work under single-process mode.");
     }
   }
+
+  /**
+   * Wait until the given condition is met.
+   *
+   * @param condition A function that predicts the condition.
+   * @param timeoutMs Timeout in milliseconds.
+   * @return True if the condition is met within the timeout, false otherwise.
+   */
+  public static boolean waitForCondition(Supplier<Boolean> condition, int timeoutMs) {
+    int waitTime = 0;
+    while (true) {
+      if (condition.get()) {
+        return true;
+      }
+
+      try {
+        java.util.concurrent.TimeUnit.MILLISECONDS.sleep(WAIT_INTERVAL_MS);
+      } catch (InterruptedException e) {
+        throw new RuntimeException(e);
+      }
+      waitTime += WAIT_INTERVAL_MS;
+      if (waitTime > timeoutMs) {
+        break;
+      }
+    }
+    return false;
+  }
 }

java/test/src/main/java/org/ray/api/test/ActorTest.java

@@ -97,7 +97,7 @@ public class ActorTest extends BaseTest {
     RayObject value = Ray.call(Counter::getValue, counter);
     Assert.assertEquals(100, value.get());
     // Delete the object from the object store.
-    Ray.internal().free(ImmutableList.of(value.getId()), false);
+    Ray.internal().free(ImmutableList.of(value.getId()), false, false);
     // Wait until the object is deleted, because the above free operation is async.
     while (true) {
       GetResult<Integer> result = ((AbstractRayRuntime)

java/test/src/main/java/org/ray/api/test/PlasmaFreeTest.java

@@ -1,17 +1,16 @@
 package org.ray.api.test;
 
 import com.google.common.collect.ImmutableList;
-import java.util.ArrayList;
-import java.util.List;
 import org.ray.api.Ray;
 import org.ray.api.RayObject;
-import org.ray.api.WaitResult;
+import org.ray.api.TestUtils;
 import org.ray.api.annotation.RayRemote;
-import org.ray.api.id.UniqueId;
+import org.ray.runtime.AbstractRayRuntime;
+import org.ray.runtime.RayNativeRuntime;
+import org.ray.runtime.util.UniqueIdUtil;
 import org.testng.Assert;
 import org.testng.annotations.Test;
 
-
 public class PlasmaFreeTest extends BaseTest {
 
   @RayRemote
@@ -20,31 +19,27 @@ public class PlasmaFreeTest extends BaseTest {
   }
 
   @Test
-  public void test() {
+  public void testDeleteObjects() {
     RayObject<String> helloId = Ray.call(PlasmaFreeTest::hello);
     String helloString = helloId.get();
     Assert.assertEquals("hello", helloString);
-    List<RayObject<String>> waitFor = ImmutableList.of(helloId);
-    WaitResult<String> waitResult = Ray.wait(waitFor, 1, 2 * 1000);
-    List<RayObject<String>> readyOnes = waitResult.getReady();
-    List<RayObject<String>> unreadyOnes = waitResult.getUnready();
-    Assert.assertEquals(1, readyOnes.size());
-    Assert.assertEquals(0, unreadyOnes.size());
+    Ray.internal().free(ImmutableList.of(helloId.getId()), true, false);
 
-    List<UniqueId> freeList = new ArrayList<>();
-    freeList.add(helloId.getId());
-    Ray.internal().free(freeList, true);
-    // Flush: trigger the release function because Plasma Client has cache.
-    for (int i = 0; i < 128; i++) {
-      Ray.call(PlasmaFreeTest::hello).get();
-    }
-
-    // Check if the object has been evicted. Don't give ray.wait enough
-    // time to reconstruct the object.
-    waitResult = Ray.wait(waitFor, 1, 0);
-    readyOnes = waitResult.getReady();
-    unreadyOnes = waitResult.getUnready();
-    Assert.assertEquals(0, readyOnes.size());
-    Assert.assertEquals(1, unreadyOnes.size());
+    final boolean result = TestUtils.waitForCondition(() -> !((AbstractRayRuntime) Ray.internal())
+        .getObjectStoreProxy().get(helloId.getId(), 0).exists, 50);
+    Assert.assertTrue(result);
   }
+
+  @Test
+  public void testDeleteCreatingTasks() {
+    TestUtils.skipTestUnderSingleProcess();
+    RayObject<String> helloId = Ray.call(PlasmaFreeTest::hello);
+    Assert.assertEquals("hello", helloId.get());
+    Ray.internal().free(ImmutableList.of(helloId.getId()), true, true);
+
+    final boolean result = TestUtils.waitForCondition(() -> !((RayNativeRuntime) Ray.internal())
+          .rayletTaskExistsInGcs(UniqueIdUtil.computeTaskId(helloId.getId())), 50);
+    Assert.assertTrue(result);
+  }
+
 }