[xray] Implements ray.wait (#2162)

Implements ray.wait for xray. Fixes #1128.

src/ray/object_manager/object_manager.h

@@ -144,23 +144,26 @@ class ObjectManager : public ObjectManagerInterface {
   ray::Status Cancel(const ObjectID &object_id);
 
   /// Callback definition for wait.
-  using WaitCallback = std::function<void(const ray::Status, uint64_t,
-                                          const std::vector<ray::ObjectID> &)>;
-  /// Wait for timeout_ms before invoking the provided callback.
-  /// If num_ready_objects is satisfied before the timeout, then
-  /// invoke the callback.
+  using WaitCallback = std::function<void(const std::vector<ray::ObjectID> &found,
+                                          const std::vector<ray::ObjectID> &remaining)>;
+  /// Wait until either num_required_objects are located or wait_ms has elapsed,
+  /// then invoke the provided callback.
   ///
   /// \param object_ids The object ids to wait on.
   /// \param timeout_ms The time in milliseconds to wait before invoking the callback.
-  /// \param num_ready_objects The minimum number of objects required before
+  /// \param num_required_objects The minimum number of objects required before
   /// invoking the callback.
+  /// \param wait_local Whether to wait until objects arrive to this node's store.
   /// \param callback Invoked when either timeout_ms is satisfied OR num_ready_objects
   /// is satisfied.
   /// \return Status of whether the wait successfully initiated.
-  ray::Status Wait(const std::vector<ObjectID> &object_ids, uint64_t timeout_ms,
-                   int num_ready_objects, const WaitCallback &callback);
+  ray::Status Wait(const std::vector<ObjectID> &object_ids, int64_t timeout_ms,
+                   uint64_t num_required_objects, bool wait_local,
+                   const WaitCallback &callback);
 
  private:
+  friend class TestObjectManager;
+
   ClientID client_id_;
   const ObjectManagerConfig config_;
   std::unique_ptr<ObjectDirectoryInterface> object_directory_;
@@ -196,12 +199,61 @@ class ObjectManager : public ObjectManagerInterface {
   /// Cache of locally available objects.
   std::unordered_map<ObjectID, ObjectInfoT> local_objects_;
 
-  /// Unfulfilled Push tasks.
-  /// The timer is for removing a push task due to unsatisfied local object.
+  /// This is used as the callback identifier in Pull for
+  /// SubscribeObjectLocations. We only need one identifier because we never need to
+  /// subscribe multiple times to the same object during Pull.
+  UniqueID object_directory_pull_callback_id_ = UniqueID::from_random();
+
+  struct WaitState {
+    WaitState(asio::io_service &service, int64_t timeout_ms, const WaitCallback &callback)
+        : timeout_ms(timeout_ms),
+          timeout_timer(std::unique_ptr<boost::asio::deadline_timer>(
+              new boost::asio::deadline_timer(
+                  service, boost::posix_time::milliseconds(timeout_ms)))),
+          callback(callback) {}
+    /// The period of time to wait before invoking the callback.
+    int64_t timeout_ms;
+    /// The timer used whenever wait_ms > 0.
+    std::unique_ptr<boost::asio::deadline_timer> timeout_timer;
+    /// The callback invoked when WaitCallback is complete.
+    WaitCallback callback;
+    /// Ordered input object_ids.
+    std::vector<ObjectID> object_id_order;
+    /// The objects that have not yet been found.
+    std::unordered_set<ObjectID> remaining;
+    /// The objects that have been found.
+    std::unordered_set<ObjectID> found;
+    /// Objects that have been requested either by Lookup or Subscribe.
+    std::unordered_set<ObjectID> requested_objects;
+    /// The number of required objects.
+    uint64_t num_required_objects;
+  };
+
+  /// 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 SubscribeRemainingWaitObjects(const UniqueID &wait_id);
+  /// Completion handler for Wait.
+  void WaitComplete(const UniqueID &wait_id);
+
+  /// Maintains a map of push requests that have not been fulfilled due to an object not
+  /// being local. Objects are removed from this map after push_timeout_ms have elapsed.
   std::unordered_map<
       ObjectID,
       std::unordered_map<ClientID, std::unique_ptr<boost::asio::deadline_timer>>>
-      unfulfilled_push_tasks_;
+      unfulfilled_push_requests_;
 
   /// Handle starting, running, and stopping asio io_service.
   void StartIOService();