[xray] Implements ray.wait

src/ray/object_manager/object_manager.cc

@@ -359,7 +359,19 @@ ray::Status ObjectManager::Wait(const std::vector<ObjectID> &object_ids,
                                 int64_t timeout_ms, uint64_t num_required_objects,
                                 bool wait_local, const WaitCallback &callback) {
   UniqueID wait_id = UniqueID::from_random();
+  RAY_RETURN_NOT_OK(AddWaitRequest(wait_id, object_ids, timeout_ms, num_required_objects,
+                                   wait_local, callback));
+  RAY_RETURN_NOT_OK(LookupRemainingWaitObjects(wait_id));
+  // LookupRemainingWaitObjects invokes SubscribeRemainingWaitObjects once lookup has
+  // been performed on all remaining objects.
+  return ray::Status::OK();
+}
 
+ray::Status ObjectManager::AddWaitRequest(const UniqueID &wait_id,
+                                          const std::vector<ObjectID> &object_ids,
+                                          int64_t timeout_ms,
+                                          uint64_t num_required_objects, bool wait_local,
+                                          const WaitCallback &callback) {
   if (wait_local) {
     return ray::Status::NotImplemented("Wait for local objects is not yet implemented.");
   }
@@ -385,6 +397,12 @@ ray::Status ObjectManager::Wait(const std::vector<ObjectID> &object_ids,
     }
   }
 
+  return ray::Status::OK();
+}
+
+ray::Status ObjectManager::LookupRemainingWaitObjects(const UniqueID &wait_id) {
+  auto &wait_state = active_wait_requests_.find(wait_id)->second;
+
   if (wait_state.remaining.empty()) {
     WaitComplete(wait_id);
   } else {
@@ -396,7 +414,7 @@ ray::Status ObjectManager::Wait(const std::vector<ObjectID> &object_ids,
     for (const auto &object_id : wait_state.remaining) {
       // Lookup remaining objects.
       wait_state.requested_objects.insert(object_id);
-      RAY_CHECK_OK(object_directory_->LookupLocations(
+      RAY_RETURN_NOT_OK(object_directory_->LookupLocations(
           object_id, [this, wait_id](const std::vector<ClientID> &client_ids,
                                      const ObjectID &lookup_object_id) {
             auto &wait_state = active_wait_requests_.find(wait_id)->second;
@@ -406,40 +424,47 @@ ray::Status ObjectManager::Wait(const std::vector<ObjectID> &object_ids,
             }
             wait_state.requested_objects.erase(lookup_object_id);
             if (wait_state.requested_objects.empty()) {
-              AllWaitLookupsComplete(wait_id);
+              SubscribeRemainingWaitObjects(wait_id);
             }
           }));
     }
   }
   return ray::Status::OK();
 }
 
-void ObjectManager::AllWaitLookupsComplete(const UniqueID &wait_id) {
+void ObjectManager::SubscribeRemainingWaitObjects(const UniqueID &wait_id) {
   auto &wait_state = active_wait_requests_.find(wait_id)->second;
   if (wait_state.found.size() >= wait_state.num_required_objects ||
       wait_state.timeout_ms == 0) {
     // Requirements already satisfied.
     WaitComplete(wait_id);
   } else {
-    // Subscribe to objects in order to ensure Wait-related tests are deterministic.
+    // Wait may complete during the execution of any one of the following calls to
+    // SubscribeObjectLocations, so copy the object ids that need to be iterated over.
+    // Order matters for test purposes.
+    std::vector<ObjectID> ordered_remaining_object_ids;
     for (const auto &object_id : wait_state.object_id_order) {
-      if (wait_state.remaining.count(object_id) == 0) {
-        continue;
+      if (wait_state.remaining.count(object_id) > 0) {
+        ordered_remaining_object_ids.push_back(object_id);
       }
-      // Subscribe to object notifications.
+    }
+    for (const auto &object_id : ordered_remaining_object_ids) {
       if (active_wait_requests_.find(wait_id) == active_wait_requests_.end()) {
         // This is possible if an object's location is obtained immediately,
         // within the current callstack. In this case, WaitComplete has been
         // invoked already, so we're done.
         return;
       }
       wait_state.requested_objects.insert(object_id);
+      // Subscribe to object notifications.
       RAY_CHECK_OK(object_directory_->SubscribeObjectLocations(
           wait_id, object_id, [this, wait_id](const std::vector<ClientID> &client_ids,
                                               const ObjectID &subscribe_object_id) {
             auto object_id_wait_state = active_wait_requests_.find(wait_id);
+            // We never expect to handle a subscription notification for a wait that has
+            // already completed.
             RAY_CHECK(object_id_wait_state != active_wait_requests_.end());
-            auto &wait_state = active_wait_requests_.find(wait_id)->second;
+            auto &wait_state = object_id_wait_state->second;
             RAY_CHECK(wait_state.remaining.erase(subscribe_object_id));
             wait_state.found.insert(subscribe_object_id);
             wait_state.requested_objects.erase(subscribe_object_id);

src/ray/object_manager/object_manager.h

@@ -231,9 +231,20 @@ class ObjectManager : public ObjectManagerInterface {
 
   /// A set of active wait requests.
   std::unordered_map<UniqueID, WaitState> active_wait_requests_;
+
+  /// Creates a wait request and adds it to active_wait_requests_.
+  ray::Status AddWaitRequest(const UniqueID &wait_id,
+                             const std::vector<ObjectID> &object_ids, int64_t timeout_ms,
+                             uint64_t num_required_objects, bool wait_local,
+                             const WaitCallback &callback);
+
+  /// Lookup any remaining objects that are not local. This is invoked after
+  /// the wait request is created and local objects are identified.
+  ray::Status LookupRemainingWaitObjects(const UniqueID &wait_id);
+
   /// Invoked when lookup for remaining objects has been invoked. This method subscribes
   /// to any remaining objects if wait conditions have not yet been satisfied.
-  void AllWaitLookupsComplete(const UniqueID &wait_id);
+  void SubscribeRemainingWaitObjects(const UniqueID &wait_id);
   /// Completion handler for Wait.
   void WaitComplete(const UniqueID &wait_id);
 

src/ray/object_manager/test/object_manager_test.cc

@@ -71,17 +71,16 @@ class MockServer {
   }
 
   friend class TestObjectManager;
-  friend class TestObjectManagerCommands;
 
   boost::asio::ip::tcp::acceptor object_manager_acceptor_;
   boost::asio::ip::tcp::socket object_manager_socket_;
   std::shared_ptr<gcs::AsyncGcsClient> gcs_client_;
   ObjectManager object_manager_;
 };
 
-class TestObjectManager : public ::testing::Test {
+class TestObjectManagerBase : public ::testing::Test {
  public:
-  TestObjectManager() {}
+  TestObjectManagerBase() {}
 
   std::string StartStore(const std::string &id) {
     std::string store_id = "/tmp/store";
@@ -176,10 +175,6 @@ class TestObjectManager : public ::testing::Test {
 
   void object_added_handler_2(ObjectID object_id) { v2.push_back(object_id); };
 
-  ObjectDirectoryInterface &get_object_directory(const MockServer &server) {
-    return *server.object_manager_.object_directory_;
-  }
-
  protected:
   std::thread p;
   boost::asio::io_service main_service;
@@ -199,7 +194,7 @@ class TestObjectManager : public ::testing::Test {
   uint push_timeout_ms;
 };
 
-class TestObjectManagerCommands : public TestObjectManager {
+class TestObjectManager : public TestObjectManagerBase {
  public:
   int current_wait_test = -1;
   int num_connected_clients = 0;
@@ -273,35 +268,38 @@ class TestObjectManagerCommands : public TestObjectManager {
     uint num_expected_objects1 = 1;
     uint num_expected_objects2 = 2;
     if (v1.size() == num_expected_objects1 && v2.size() == num_expected_objects2) {
-      TestWaitCallbacks1();
+      SubscribeObjectThenWait();
     }
   }
 
-  void TestWaitCallbacks1() {
+  void SubscribeObjectThenWait() {
     int data_size = 100;
     // Test to ensure Wait works properly during an active subscription to the same
     // object.
     ObjectID object_1 = WriteDataToClient(client2, data_size);
     ObjectID object_2 = WriteDataToClient(client2, data_size);
     UniqueID sub_id = ray::ObjectID::from_random();
 
-    RAY_CHECK_OK(get_object_directory(*server1).SubscribeObjectLocations(
+    RAY_CHECK_OK(server1->object_manager_.object_directory_->SubscribeObjectLocations(
         sub_id, object_1,
         [this, sub_id, object_1, object_2](const std::vector<ray::ClientID> &,
                                            const ray::ObjectID &object_id) {
-          TestWaitCallbacks2(sub_id, object_1, object_2);
+          TestWaitWhileSubscribed(sub_id, object_1, object_2);
         }));
   }
 
-  void TestWaitCallbacks2(UniqueID sub_id, ObjectID object_1, ObjectID object_2) {
+  void TestWaitWhileSubscribed(UniqueID sub_id, ObjectID object_1, ObjectID object_2) {
     int num_objects = 2;
     int required_objects = 1;
     int timeout_ms = 1000;
 
     std::vector<ObjectID> object_ids = {object_1, object_2};
     boost::posix_time::ptime start_time = boost::posix_time::second_clock::local_time();
-    RAY_CHECK_OK(server1->object_manager_.Wait(
-        object_ids, timeout_ms, required_objects, false,
+
+    UniqueID wait_id = UniqueID::from_random();
+
+    RAY_CHECK_OK(server1->object_manager_.AddWaitRequest(
+        wait_id, object_ids, timeout_ms, required_objects, false,
         [this, sub_id, object_1, object_ids, num_objects, start_time](
             const std::vector<ray::ObjectID> &found,
             const std::vector<ray::ObjectID> &remaining) {
@@ -313,10 +311,14 @@ class TestObjectManagerCommands : public TestObjectManager {
           RAY_CHECK(found.size() == 1);
           // There's nothing more to test. A check will fail if unexpected behavior is
           // triggered.
-          RAY_CHECK_OK(get_object_directory(*server1).UnsubscribeObjectLocations(
-              sub_id, object_1));
+          RAY_CHECK_OK(
+              server1->object_manager_.object_directory_->UnsubscribeObjectLocations(
+                  sub_id, object_1));
           NextWaitTest();
         }));
+
+    // Skip lookups and rely on Subscribe only to test subscribe interaction.
+    server1->object_manager_.SubscribeRemainingWaitObjects(wait_id);
   }
 
   void NextWaitTest() {
@@ -400,7 +402,8 @@ class TestObjectManagerCommands : public TestObjectManager {
 
           switch (current_wait_test) {
           case 0: {
-            // Ensure timeout_ms = 0 returns expected number of found / remaining objects.
+            // Ensure timeout_ms = 0 returns expected number of found and remaining
+            // objects.
             ASSERT_TRUE(found.size() <= required_objects);
             ASSERT_TRUE(static_cast<int>(found.size() + remaining.size()) == num_objects);
             NextWaitTest();
@@ -454,7 +457,7 @@ class TestObjectManagerCommands : public TestObjectManager {
   }
 };
 
-TEST_F(TestObjectManagerCommands, StartTestObjectManagerCommands) {
+TEST_F(TestObjectManager, StartTestObjectManager) {
   auto AsyncStartTests = main_service.wrap([this]() { WaitConnections(); });
   AsyncStartTests();
   main_service.run();