[Java] Improve JNI performance when submitting and executing tasks (#9032)

java/runtime/src/main/java/io/ray/runtime/functionmanager/PyFunctionDescriptor.java

@@ -1,5 +1,6 @@
 package io.ray.runtime.functionmanager;
 
+import com.google.common.base.Objects;
 import io.ray.runtime.generated.Common.Language;
 import java.util.Arrays;
 import java.util.List;
@@ -26,6 +27,25 @@ public class PyFunctionDescriptor implements FunctionDescriptor {
     return moduleName + "." + className + "." + functionName;
   }
 
+  @Override
+  public boolean equals(Object o) {
+    if (this == o) {
+      return true;
+    }
+    if (o == null || getClass() != o.getClass()) {
+      return false;
+    }
+    PyFunctionDescriptor that = (PyFunctionDescriptor) o;
+    return Objects.equal(moduleName, that.moduleName) &&
+      Objects.equal(className, that.className) &&
+      Objects.equal(functionName, that.functionName);
+  }
+
+  @Override
+  public int hashCode() {
+    return Objects.hashCode(moduleName, className, functionName);
+  }
+
   @Override
   public List<String> toList() {
     return Arrays.asList(moduleName, className, functionName, "" /* function hash */);

java/runtime/src/main/java/io/ray/runtime/task/NativeTaskSubmitter.java

@@ -1,6 +1,7 @@
 package io.ray.runtime.task;
 
 import com.google.common.base.Preconditions;
+import com.google.common.collect.ImmutableList;
 import io.ray.api.BaseActorHandle;
 import io.ray.api.id.ObjectId;
 import io.ray.api.options.ActorCreationOptions;
@@ -18,14 +19,19 @@ public class NativeTaskSubmitter implements TaskSubmitter {
   @Override
   public List<ObjectId> submitTask(FunctionDescriptor functionDescriptor, List<FunctionArg> args,
                                    int numReturns, CallOptions options) {
-    List<byte[]> returnIds = nativeSubmitTask(functionDescriptor, args, numReturns, options);
+    List<byte[]> returnIds = nativeSubmitTask(functionDescriptor, functionDescriptor.hashCode(),
+        args, numReturns, options);
+    if (returnIds == null) {
+      return ImmutableList.of();
+    }
     return returnIds.stream().map(ObjectId::new).collect(Collectors.toList());
   }
 
   @Override
   public BaseActorHandle createActor(FunctionDescriptor functionDescriptor, List<FunctionArg> args,
                                      ActorCreationOptions options) {
-    byte[] actorId = nativeCreateActor(functionDescriptor, args, options);
+    byte[] actorId = nativeCreateActor(functionDescriptor, functionDescriptor.hashCode(), args,
+        options);
     return NativeActorHandle.create(actorId, functionDescriptor.getLanguage());
   }
 
@@ -35,17 +41,21 @@ public class NativeTaskSubmitter implements TaskSubmitter {
       List<FunctionArg> args, int numReturns, CallOptions options) {
     Preconditions.checkState(actor instanceof NativeActorHandle);
     List<byte[]> returnIds = nativeSubmitActorTask(actor.getId().getBytes(),
-        functionDescriptor, args, numReturns, options);
+        functionDescriptor, functionDescriptor.hashCode(), args, numReturns, options);
+    if (returnIds == null) {
+      return ImmutableList.of();
+    }
     return returnIds.stream().map(ObjectId::new).collect(Collectors.toList());
   }
 
   private static native List<byte[]> nativeSubmitTask(FunctionDescriptor functionDescriptor,
-      List<FunctionArg> args, int numReturns, CallOptions callOptions);
+      int functionDescriptorHash, List<FunctionArg> args, int numReturns, CallOptions callOptions);
 
   private static native byte[] nativeCreateActor(FunctionDescriptor functionDescriptor,
-      List<FunctionArg> args, ActorCreationOptions actorCreationOptions);
+      int functionDescriptorHash, List<FunctionArg> args,
+      ActorCreationOptions actorCreationOptions);
 
   private static native List<byte[]> nativeSubmitActorTask(byte[] actorId,
-      FunctionDescriptor functionDescriptor, List<FunctionArg> args, int numReturns,
-      CallOptions callOptions);
+      FunctionDescriptor functionDescriptor, int functionDescriptorHash, List<FunctionArg> args,
+      int numReturns, CallOptions callOptions);
 }