Add a flag to disable reconstruction for a killed actor (#7346)

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

@@ -1,21 +1,35 @@
 package org.ray.api.test;
 
 import com.google.common.collect.ImmutableList;
+import java.util.function.BiConsumer;
 import org.ray.api.Ray;
 import org.ray.api.RayActor;
 import org.ray.api.RayObject;
 import org.ray.api.TestUtils;
 import org.ray.api.exception.RayActorException;
+import org.ray.api.options.ActorCreationOptions;
 import org.testng.Assert;
+import org.testng.annotations.AfterClass;
+import org.testng.annotations.BeforeClass;
 import org.testng.annotations.Test;
 
 @Test
 public class KillActorTest extends BaseTest {
 
+  @BeforeClass
+  public void setUp() {
+    System.setProperty("ray.raylet.config.num_workers_per_process_java", "1");
+  }
+
+  @AfterClass
+  public void tearDown() {
+    System.clearProperty("ray.raylet.config.num_workers_per_process_java");
+  }
+
   public static class HangActor {
 
-    public boolean alive() {
-      return true;
+    public String ping() {
+      return "pong";
     }
 
     public boolean hang() throws InterruptedException {
@@ -25,13 +39,60 @@ public class KillActorTest extends BaseTest {
     }
   }
 
-  public void testKillActor() {
+  public static class KillerActor {
+
+    public void kill(RayActor<?> actor, boolean noReconstruction) {
+      actor.kill(noReconstruction);
+    }
+  }
+
+  private static void localKill(RayActor<?> actor, boolean noReconstruction) {
+    actor.kill(noReconstruction);
+  }
+
+  private static void remoteKill(RayActor<?> actor, boolean noReconstruction) {
+    RayActor<KillerActor> killer = Ray.createActor(KillerActor::new);
+    killer.call(KillerActor::kill, actor, noReconstruction);
+  }
+
+  private void testKillActor(BiConsumer<RayActor<?>, Boolean> kill, boolean noReconstruction) {
     TestUtils.skipTestUnderSingleProcess();
-    RayActor<HangActor> actor = Ray.createActor(HangActor::new);
-    Assert.assertTrue(actor.call(HangActor::alive).get());
+
+    ActorCreationOptions options =
+        new ActorCreationOptions.Builder().setMaxReconstructions(1).createActorCreationOptions();
+    RayActor<HangActor> actor = Ray.createActor(HangActor::new, options);
     RayObject<Boolean> result = actor.call(HangActor::hang);
+    // The actor will hang in this task.
     Assert.assertEquals(0, Ray.wait(ImmutableList.of(result), 1, 500).getReady().size());
-    Ray.killActor(actor);
+
+    // Kill the actor
+    kill.accept(actor, noReconstruction);
+    // The get operation will fail with RayActorException
     Assert.expectThrows(RayActorException.class, result::get);
+
+    try {
+      // Sleep 1s here to make sure the driver has received the actor notification
+      // (of state RECONSTRUCTING or DEAD).
+      Thread.sleep(1000);
+    } catch (InterruptedException e) {
+      throw new RuntimeException(e);
+    }
+
+    if (noReconstruction) {
+      // The actor should not be reconstructed.
+      Assert.expectThrows(RayActorException.class, () -> actor.call(HangActor::hang).get());
+    } else {
+      Assert.assertEquals(actor.call(HangActor::ping).get(), "pong");
+    }
+  }
+
+  public void testLocalKill() {
+    testKillActor(KillActorTest::localKill, false);
+    testKillActor(KillActorTest::localKill, true);
+  }
+
+  public void testRemoteKill() {
+    testKillActor(KillActorTest::remoteKill, false);
+    testKillActor(KillActorTest::remoteKill, true);
   }
 }