[Java worker] Refactor object store and worker context on top of core worker (#5079)

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

@@ -1,12 +1,18 @@
 package org.ray.api.test;
 
+import java.time.Duration;
+import java.time.Instant;
+import java.util.Arrays;
+import java.util.List;
 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.exception.RayException;
 import org.ray.api.exception.RayTaskException;
 import org.ray.api.exception.RayWorkerException;
+import org.ray.api.function.RayFunc0;
 import org.testng.Assert;
 import org.testng.annotations.Test;
 
@@ -23,6 +29,15 @@ public class FailureTest extends BaseTest {
     return 0;
   }
 
+  public static int slowFunc() {
+    try {
+      Thread.sleep(10000);
+    } catch (InterruptedException e) {
+      throw new RuntimeException(e);
+    }
+    return 0;
+  }
+
   public static class BadActor {
 
     public BadActor(boolean failOnCreation) {
@@ -106,5 +121,26 @@ public class FailureTest extends BaseTest {
       // RayActorException.
     }
   }
+
+  @Test
+  public void testGetThrowsQuicklyWhenFoundException() {
+    TestUtils.skipTestUnderSingleProcess();
+    List<RayFunc0<Integer>> badFunctions = Arrays.asList(FailureTest::badFunc,
+        FailureTest::badFunc2);
+    for (RayFunc0<Integer> badFunc : badFunctions) {
+      RayObject<Integer> obj1 = Ray.call(badFunc);
+      RayObject<Integer> obj2 = Ray.call(FailureTest::slowFunc);
+      Instant start = Instant.now();
+      try {
+        Ray.get(Arrays.asList(obj1.getId(), obj2.getId()));
+        Assert.fail("Should throw RayException.");
+      } catch (RayException e) {
+        Instant end = Instant.now();
+        long duration = Duration.between(start, end).toMillis();
+        Assert.assertTrue(duration < 5000, "Should fail quickly. " +
+            "Actual execution time: " + duration + " ms.");
+      }
+    }
+  }
 }
 

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

@@ -1,12 +1,10 @@
 package org.ray.api.test;
 
-import org.apache.arrow.plasma.PlasmaClient;
-import org.apache.arrow.plasma.exceptions.DuplicateObjectException;
-
 import org.ray.api.Ray;
 import org.ray.api.TestUtils;
-import org.ray.api.id.UniqueId;
+import org.ray.api.id.ObjectId;
 import org.ray.runtime.AbstractRayRuntime;
+import org.ray.runtime.objectstore.ObjectStoreProxy;
 import org.testng.Assert;
 import org.testng.annotations.Test;
 
@@ -15,15 +13,13 @@ public class PlasmaStoreTest extends BaseTest {
   @Test
   public void testPutWithDuplicateId() {
     TestUtils.skipTestUnderSingleProcess();
-    UniqueId objectId = UniqueId.randomId();
+    ObjectId objectId = ObjectId.randomId();
     AbstractRayRuntime runtime = (AbstractRayRuntime) Ray.internal();
-    PlasmaClient store = new PlasmaClient(runtime.getRayConfig().objectStoreSocketName, "", 0);
-    store.put(objectId.getBytes(), new byte[]{}, new byte[]{});
-    try {
-      store.put(objectId.getBytes(), new byte[]{}, new byte[]{});
-      Assert.fail("This line shouldn't be reached.");
-    } catch (DuplicateObjectException e) {
-      // Putting 2 objects with duplicate ID should throw DuplicateObjectException.
-    }
+    ObjectStoreProxy objectInterface = runtime.getObjectStoreProxy();
+    objectInterface.put(objectId, 1);
+    Assert.assertEquals(objectInterface.<Integer>get(objectId, -1).object, (Integer) 1);
+    objectInterface.put(objectId, 2);
+    // Putting 2 objects with duplicate ID should fail but ignored.
+    Assert.assertEquals(objectInterface.<Integer>get(objectId, -1).object, (Integer) 1);
   }
 }

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

@@ -1,7 +1,7 @@
 package org.ray.api.test;
 
 import org.ray.runtime.config.RayConfig;
-import org.ray.runtime.config.WorkerMode;
+import org.ray.runtime.generated.Common.WorkerType;
 import org.testng.Assert;
 import org.testng.annotations.Test;
 
@@ -12,7 +12,7 @@ public class RayConfigTest {
     try {
       System.setProperty("ray.job.resource-path", "path/to/ray/job/resource/path");
       RayConfig rayConfig = RayConfig.create();
-      Assert.assertEquals(WorkerMode.DRIVER, rayConfig.workerMode);
+      Assert.assertEquals(WorkerType.DRIVER, rayConfig.workerMode);
       Assert.assertEquals("path/to/ray/job/resource/path", rayConfig.jobResourcePath);
     } finally {
       // Unset system properties.