[xray] Track ray.get calls as task dependencies (#2362)

java/runtime-native/src/main/java/org/ray/core/impl/RayNativeRuntime.java

@@ -109,6 +109,7 @@ public class RayNativeRuntime extends RayRuntime {
           WorkerContext.currentWorkerId(),
           UniqueID.nil,
           isWorker,
+          WorkerContext.currentTask().taskId,
           0
       );
 
@@ -237,4 +238,4 @@ public class RayNativeRuntime extends RayRuntime {
       throw new TaskExecutionException(log, e);
     }
   }
-}
+}

java/runtime-native/src/main/java/org/ray/spi/impl/DefaultLocalSchedulerClient.java

@@ -26,13 +26,13 @@ public class DefaultLocalSchedulerClient implements LocalSchedulerLink {
   private long client = 0;
 
   public DefaultLocalSchedulerClient(String schedulerSockName, UniqueID clientId, UniqueID actorId,
-                                     boolean isWorker, long numGpus) {
+                                     boolean isWorker, UniqueID driverId, long numGpus) {
     client = _init(schedulerSockName, clientId.getBytes(), actorId.getBytes(), isWorker,
-        numGpus);
+        driverId.getBytes(), numGpus);
   }
 
   private static native long _init(String localSchedulerSocket, byte[] workerId, byte[] actorId,
-                                   boolean isWorker, long numGpus);
+                                   boolean isWorker, byte[] driverTaskId, long numGpus);
 
   private static native byte[] _computePutId(long client, byte[] taskId, int putIndex);