[Java] Add getAsyncContext and setAsyncContext API (#6439)

* Add getAsyncContext and setAsyncContext API

* address comment

* fix bug

* Add test case

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

@@ -47,6 +47,15 @@ public class RayDevRuntime extends AbstractRayRuntime {
     LOGGER.error("Not implemented under SINGLE_PROCESS mode.");
   }
 
+  @Override
+  public Object getAsyncContext() {
+    return null;
+  }
+
+  @Override
+  public void setAsyncContext(Object asyncContext) {
+  }
+
   private JobId nextJobId() {
     return JobId.fromInt(jobCounter.getAndIncrement());
   }

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

@@ -178,20 +178,30 @@ public class RayMultiWorkerNativeRuntime implements RayRuntime {
     return getCurrentRuntime().createPyActor(moduleName, className, args, options);
   }
 
+  @Override
+  public Object getAsyncContext() {
+    return getCurrentRuntime();
+  }
+
+  @Override
+  public void setAsyncContext(Object asyncContext) {
+    currentThreadRuntime.set((RayNativeRuntime)asyncContext);
+  }
+
   @Override
   public Runnable wrapRunnable(Runnable runnable) {
-    RayNativeRuntime runtime = getCurrentRuntime();
+    Object asyncContext = getAsyncContext();
     return () -> {
-      currentThreadRuntime.set(runtime);
+      setAsyncContext(asyncContext);
       runnable.run();
     };
   }
 
   @Override
   public Callable wrapCallable(Callable callable) {
-    RayNativeRuntime runtime = getCurrentRuntime();
+    Object asyncContext = getAsyncContext();
     return () -> {
-      currentThreadRuntime.set(runtime);
+      setAsyncContext(asyncContext);
       return callable.call();
     };
   }

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

@@ -150,6 +150,15 @@ public final class RayNativeRuntime extends AbstractRayRuntime {
     nativeSetResource(nativeCoreWorkerPointer, resourceName, capacity, nodeId.getBytes());
   }
 
+  @Override
+  public Object getAsyncContext() {
+    return null;
+  }
+
+  @Override
+  public void setAsyncContext(Object asyncContext) {
+  }
+
   public void run() {
     nativeRunTaskExecutor(nativeCoreWorkerPointer, taskExecutor);
   }