[java] Fix the logic of generating TaskID (#2747)

## What do these changes do?
Because the logic of generating `TaskID` in java is different from python's, there are many tests fail when we change the `Ray Core` code.
In this change,  I rewrote the logic of generating `TaskID` in java which is the same as the python's.

In java, we call the native method `_generateTaskId()` to generate a `TaskID` which is also used in python. We change `computePutId()`'s logic too.

## Related issue number
[#2608](https://github.com/ray-project/ray/issues/2608)

src/local_scheduler/lib/java/org_ray_spi_impl_DefaultLocalSchedulerClient.cc

@@ -3,6 +3,7 @@
 #include "local_scheduler/lib/java/org_ray_spi_impl_DefaultLocalSchedulerClient.h"
 #include "local_scheduler_client.h"
 #include "logging.h"
+#include "ray/id.h"
 
 #ifdef __cplusplus
 extern "C" {
@@ -299,6 +300,31 @@ Java_org_ray_spi_impl_DefaultLocalSchedulerClient__1waitObject(
   return resultArray;
 }
 
+JNIEXPORT jbyteArray JNICALL
+Java_org_ray_spi_impl_DefaultLocalSchedulerClient__1generateTaskId(
+    JNIEnv *env,
+    jclass,
+    jbyteArray did,
+    jbyteArray ptid,
+    jint parent_task_counter) {
+  UniqueIdFromJByteArray o1(env, did);
+  ray::DriverID driver_id = *o1.PID;
+
+  UniqueIdFromJByteArray o2(env, ptid);
+  ray::TaskID parent_task_id = *o2.PID;
+
+  ray::TaskID task_id =
+      ray::GenerateTaskId(driver_id, parent_task_id, parent_task_counter);
+  jbyteArray result = env->NewByteArray(sizeof(ray::TaskID));
+  if (nullptr == result) {
+    return nullptr;
+  }
+  env->SetByteArrayRegion(result, 0, sizeof(TaskID),
+                          reinterpret_cast<jbyte *>(&task_id));
+
+  return result;
+}
+
 #ifdef __cplusplus
 }
 #endif

src/local_scheduler/lib/java/org_ray_spi_impl_DefaultLocalSchedulerClient.h

@@ -130,6 +130,18 @@ Java_org_ray_spi_impl_DefaultLocalSchedulerClient__1waitObject(JNIEnv *,
                                                                jint,
                                                                jboolean);
 
+/*
+ * Class:     org_ray_spi_impl_DefaultLocalSchedulerClient
+ * Method:    _generateTaskId
+ * Signature: ([B[BI)[B
+ */
+JNIEXPORT jbyteArray JNICALL
+Java_org_ray_spi_impl_DefaultLocalSchedulerClient__1generateTaskId(JNIEnv *,
+                                                                   jclass,
+                                                                   jbyteArray,
+                                                                   jbyteArray,
+                                                                   jint);
+
 #ifdef __cplusplus
 }
 #endif

src/ray/id.cc

@@ -6,6 +6,7 @@
 #include <mutex>
 #include <random>
 
+#include "common/common.h"
 #include "ray/constants.h"
 #include "ray/status.h"
 
@@ -177,6 +178,7 @@ const ObjectID ComputeReturnId(const TaskID &task_id, int64_t return_index) {
 
 const ObjectID ComputePutId(const TaskID &task_id, int64_t put_index) {
   RAY_CHECK(put_index >= 1 && put_index <= kMaxTaskPuts);
+  // We multiply put_index by -1 to distinguish from return_index.
   return ComputeObjectId(task_id, -1 * put_index);
 }
 
@@ -190,6 +192,26 @@ const TaskID ComputeTaskId(const ObjectID &object_id) {
   return task_id;
 }
 
+const TaskID GenerateTaskId(const DriverID &driver_id, const TaskID &parent_task_id,
+                            int parent_task_counter) {
+  // Compute hashes.
+  SHA256_CTX ctx;
+  sha256_init(&ctx);
+  sha256_update(&ctx, (BYTE *)&driver_id, sizeof(driver_id));
+  sha256_update(&ctx, (BYTE *)&parent_task_id, sizeof(parent_task_id));
+  sha256_update(&ctx, (BYTE *)&parent_task_counter, sizeof(parent_task_counter));
+
+  // Compute the final task ID from the hash.
+  BYTE buff[DIGEST_SIZE];
+  sha256_final(&ctx, buff);
+  TaskID task_id;
+  RAY_DCHECK(sizeof(task_id) <= DIGEST_SIZE);
+  memcpy(&task_id, buff, sizeof(task_id));
+  task_id = FinishTaskId(task_id);
+
+  return task_id;
+}
+
 int64_t ComputeObjectIndex(const ObjectID &object_id) {
   const int64_t *first_bytes = reinterpret_cast<const int64_t *>(&object_id);
   uint64_t bitmask = static_cast<uint64_t>(-1) << kObjectIdIndexSize;

src/ray/id.h

@@ -10,6 +10,10 @@
 #include "ray/constants.h"
 #include "ray/util/visibility.h"
 
+extern "C" {
+#include "sha256.h"
+}
+
 namespace ray {
 
 class RAY_EXPORT UniqueID {
@@ -81,6 +85,15 @@ const ObjectID ComputePutId(const TaskID &task_id, int64_t put_index);
 /// \return The task ID of the task that created this object.
 const TaskID ComputeTaskId(const ObjectID &object_id);
 
+/// Generate a task ID from the given info.
+///
+/// \param driver_id The driver that creates the task.
+/// \param parent_task_id The parent task of this task.
+/// \param parent_task_counter The task index of the worker.
+/// \return The task ID generated from the given info.
+const TaskID GenerateTaskId(const DriverID &driver_id, const TaskID &parent_task_id,
+                            int parent_task_counter);
+
 /// Compute the index of this object in the task that created it.
 ///
 /// \param object_id The object ID.

src/ray/raylet/node_manager.cc

@@ -1195,21 +1195,11 @@ void NodeManager::HandleTaskReconstruction(const TaskID &task_id) {
       [this](ray::gcs::AsyncGcsClient *client, const TaskID &task_id) {
         // The task was not in the GCS task table. It must therefore be in the
         // lineage cache.
-        if (!lineage_cache_.ContainsTask(task_id)) {
-          // The task was not in the lineage cache.
-          // TODO(swang): This should not ever happen, but Java TaskIDs are
-          // currently computed differently from Python TaskIDs, so
-          // reconstruction is currently broken for Java. Once the TaskID
-          // generation code matches for both frontends, we should be able to
-          // remove this warning and make it a fatal check.
-          RAY_LOG(WARNING) << "Task " << task_id << " to reconstruct was not found in "
-                                                    "the GCS or the lineage cache. This "
-                                                    "job may hang.";
-        } else {
-          // Use a copy of the cached task spec to re-execute the task.
-          const Task task = lineage_cache_.GetTask(task_id);
-          ResubmitTask(task);
-        }
+        RAY_CHECK(lineage_cache_.ContainsTask(task_id));
+        // Use a copy of the cached task spec to re-execute the task.
+        const Task task = lineage_cache_.GetTask(task_id);
+        ResubmitTask(task);
+
       }));
 }
 

src/ray/raylet/task_spec.cc

@@ -63,20 +63,8 @@ TaskSpecification::TaskSpecification(
     : spec_() {
   flatbuffers::FlatBufferBuilder fbb;
 
-  // Compute hashes.
-  SHA256_CTX ctx;
-  sha256_init(&ctx);
-  sha256_update(&ctx, (BYTE *)&driver_id, sizeof(driver_id));
-  sha256_update(&ctx, (BYTE *)&parent_task_id, sizeof(parent_task_id));
-  sha256_update(&ctx, (BYTE *)&parent_counter, sizeof(parent_counter));
+  TaskID task_id = GenerateTaskId(driver_id, parent_task_id, parent_counter);
 
-  // Compute the final task ID from the hash.
-  BYTE buff[DIGEST_SIZE];
-  sha256_final(&ctx, buff);
-  TaskID task_id;
-  RAY_DCHECK(sizeof(task_id) <= DIGEST_SIZE);
-  memcpy(&task_id, buff, sizeof(task_id));
-  task_id = FinishTaskId(task_id);
   // Add argument object IDs.
   std::vector<flatbuffers::Offset<Arg>> arguments;
   for (auto &argument : task_arguments) {