From 3845c294c3767601a2e18970bfe285e0042081cf Mon Sep 17 00:00:00 2001
From: Wang Qing <kingchin1218@126.com>
Date: Mon, 6 Aug 2018 14:49:54 +0800
Subject: [PATCH] [java] Fix java raylet wait (#2553)

---
 .../spi/impl/DefaultLocalSchedulerClient.java |  4 +-
 .../main/java/org/ray/api/test/WaitTest.java  | 52 +++++++++++++++++++
 2 files changed, 54 insertions(+), 2 deletions(-)
 create mode 100644 java/test/src/main/java/org/ray/api/test/WaitTest.java

diff --git a/java/runtime-native/src/main/java/org/ray/spi/impl/DefaultLocalSchedulerClient.java b/java/runtime-native/src/main/java/org/ray/spi/impl/DefaultLocalSchedulerClient.java
index f732efb37..78deb5413 100644
--- a/java/runtime-native/src/main/java/org/ray/spi/impl/DefaultLocalSchedulerClient.java
+++ b/java/runtime-native/src/main/java/org/ray/spi/impl/DefaultLocalSchedulerClient.java
@@ -50,6 +50,7 @@ public class DefaultLocalSchedulerClient implements LocalSchedulerLink {
     assert (useRaylet == true);
 
     boolean[] readys = _waitObject(client, objectIds, numReturns, timeoutMs, false);
+    assert (readys.length == objectIds.length);
 
     List<byte[]> ret = new ArrayList<>();
     for (int i = 0; i < readys.length; i++) {
@@ -58,8 +59,7 @@ public class DefaultLocalSchedulerClient implements LocalSchedulerLink {
       }
     }
 
-    assert (ret.size() == readys.length);
-    return ret; 
+    return ret;
   }
 
   @Override
diff --git a/java/test/src/main/java/org/ray/api/test/WaitTest.java b/java/test/src/main/java/org/ray/api/test/WaitTest.java
new file mode 100644
index 000000000..440cf2373
--- /dev/null
+++ b/java/test/src/main/java/org/ray/api/test/WaitTest.java
@@ -0,0 +1,52 @@
+package org.ray.api.test;
+
+import org.junit.Assert;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.ray.api.Ray;
+import org.ray.api.RayList;
+import org.ray.api.RayObject;
+import org.ray.api.RayRemote;
+import org.ray.api.WaitResult;
+
+@RunWith(MyRunner.class)
+public class WaitTest {
+
+  @RayRemote
+  public static String hi() {
+    return "hi";
+  }
+
+  @RayRemote
+  public static String delayHi() {
+    try {
+      Thread.sleep(100 * 1000);
+    } catch (Exception e) {
+      e.printStackTrace();
+    }
+
+    return "hi";
+  }
+
+  @Test
+  public void test() {
+    RayObject<String> obj1 = Ray.call(WaitTest::hi);
+    RayObject<String> obj2 = Ray.call(WaitTest::delayHi);
+
+    RayList<String> waitfor = new RayList<>();
+    waitfor.add(obj1);
+    waitfor.add(obj2);
+    WaitResult<String> waitResult = Ray.wait(waitfor, 2, 2 * 1000);
+    RayList<String> readys = waitResult.getReadyOnes();
+
+    if (!readys.isEmpty()) {
+      Assert.assertEquals(1, waitResult.getReadyOnes().size());
+      Assert.assertEquals(1, waitResult.getRemainOnes().size());
+      Assert.assertEquals("hi", readys.get(0));
+    } else {
+      Assert.assertEquals(0, waitResult.getReadyOnes().size());
+      Assert.assertEquals(2, waitResult.getRemainOnes().size());
+    }
+  }
+
+}