[Java] Refine tests and fix single-process mode (#4265)

java/runtime/src/main/java/org/ray/runtime/RayDevRuntime.java

@@ -28,4 +28,9 @@ public class RayDevRuntime extends AbstractRayRuntime {
   public MockObjectStore getObjectStore() {
     return store;
   }
+
+  @Override
+  public Worker getWorker() {
+    return ((MockRayletClient) rayletClient).getCurrentWorker();
+  }
 }

java/runtime/src/main/java/org/ray/runtime/RuntimeContextImpl.java

@@ -4,7 +4,6 @@ import com.google.common.base.Preconditions;
 import org.ray.api.RuntimeContext;
 import org.ray.api.id.UniqueId;
 import org.ray.runtime.config.RunMode;
-import org.ray.runtime.config.WorkerMode;
 import org.ray.runtime.task.TaskSpec;
 
 public class RuntimeContextImpl implements RuntimeContext {
@@ -22,8 +21,10 @@ public class RuntimeContextImpl implements RuntimeContext {
 
   @Override
   public UniqueId getCurrentActorId() {
-    Preconditions.checkState(runtime.rayConfig.workerMode == WorkerMode.WORKER);
-    return runtime.getWorker().getCurrentActorId();
+    Worker worker = runtime.getWorker();
+    Preconditions.checkState(worker != null && !worker.getCurrentActorId().isNil(),
+        "This method should only be called from an actor.");
+    return worker.getCurrentActorId();
   }
 
   @Override

java/runtime/src/main/java/org/ray/runtime/Worker.java

@@ -79,7 +79,6 @@ public class Worker {
    * Execute a task.
    */
   public void execute(TaskSpec spec) {
-    LOGGER.info("Executing task {}", spec.taskId);
     LOGGER.debug("Executing task {}", spec);
     UniqueId returnId = spec.returnIds[0];
     ClassLoader oldLoader = Thread.currentThread().getContextClassLoader();
@@ -123,7 +122,7 @@ public class Worker {
         maybeLoadCheckpoint(result, returnId);
         currentActor = result;
       }
-      LOGGER.info("Finished executing task {}", spec.taskId);
+      LOGGER.debug("Finished executing task {}", spec.taskId);
     } catch (Exception e) {
       LOGGER.error("Error executing task " + spec, e);
       if (!spec.isActorCreationTask()) {

java/runtime/src/main/java/org/ray/runtime/objectstore/MockObjectStore.java

@@ -95,11 +95,10 @@ public class MockObjectStore implements ObjectStoreLink {
     ArrayList<ObjectStoreData> rets = new ArrayList<>();
     for (byte[] id : objectIds) {
       try {
-        Constructor<ObjectStoreData> constructor = ObjectStoreData.class.getConstructor(
-                byte[].class, byte[].class);
+        Constructor<?> constructor = ObjectStoreData.class.getDeclaredConstructors()[0];
         constructor.setAccessible(true);
-        rets.add(constructor.newInstance(metadata.get(new UniqueId(id)),
-                data.get(new UniqueId(id))));
+        rets.add((ObjectStoreData) constructor.newInstance(data.get(new UniqueId(id)),
+                metadata.get(new UniqueId(id))));
       } catch (Exception e) {
         throw new RuntimeException(e);
       }

java/runtime/src/main/java/org/ray/runtime/raylet/MockRayletClient.java

@@ -39,6 +39,7 @@ public class MockRayletClient implements RayletClient {
   private final ExecutorService exec;
   private final Deque<Worker> idleWorkers;
   private final Map<UniqueId, Worker> actorWorkers;
+  private final ThreadLocal<Worker> currentWorker;
 
   public MockRayletClient(RayDevRuntime runtime, int numberThreads) {
     this.runtime = runtime;
@@ -48,6 +49,7 @@ public class MockRayletClient implements RayletClient {
     exec = Executors.newFixedThreadPool(numberThreads);
     idleWorkers = new LinkedList<>();
     actorWorkers = new HashMap<>();
+    currentWorker = new ThreadLocal<>();
   }
 
   public synchronized void onObjectPut(UniqueId id) {
@@ -60,22 +62,28 @@ public class MockRayletClient implements RayletClient {
     }
   }
 
+  public Worker getCurrentWorker() {
+    return currentWorker.get();
+  }
+
   /**
    * Get a worker from the worker pool to run the given task.
    */
   private Worker getWorker(TaskSpec task) {
-    if (task.isActorTask()) {
-      return actorWorkers.get(task.actorId);
-    }
     Worker worker;
-    if (idleWorkers.size() > 0) {
-      worker = idleWorkers.pop();
+    if (task.isActorTask()) {
+      worker = actorWorkers.get(task.actorId);
     } else {
-      worker = new Worker(runtime);
-    }
-    if (task.isActorCreationTask()) {
-      actorWorkers.put(task.actorCreationId, worker);
+      if (idleWorkers.size() > 0) {
+        worker = idleWorkers.pop();
+      } else {
+        worker = new Worker(runtime);
+      }
+      if (task.isActorCreationTask()) {
+        actorWorkers.put(task.actorCreationId, worker);
+      }
     }
+    currentWorker.set(worker);
     return worker;
   }
 
@@ -83,6 +91,7 @@ public class MockRayletClient implements RayletClient {
    * Return the worker to the worker pool.
    */
   private void returnWorker(Worker worker) {
+    currentWorker.remove();
     idleWorkers.push(worker);
   }
 
@@ -105,9 +114,7 @@ public class MockRayletClient implements RayletClient {
                     new byte[]{}, new byte[]{});
           }
         } finally {
-          if (!task.isActorCreationTask() && !task.isActorTask()) {
-            returnWorker(worker);
-          }
+          returnWorker(worker);
         }
       });
     } else {

java/runtime/src/main/resources/ray.default.conf

@@ -100,7 +100,7 @@ ray {
   // ----------------------------
   dev-runtime {
     // Number of threads that you process tasks
-    execution-parallelism: 5
+    execution-parallelism: 10
   }
 
 }