Switch to new wait implementation. (#113)

* Duplicate wait1 implementation and seperate out wait datastructures.

* Address Philipp's comments.

* Temporarily address test failure problem by increasing timeout and reducing load in tests.

* Update stress tests to include distributed wait.

src/plasma/plasma.h

@@ -132,7 +132,9 @@ enum plasma_message_type {
   /** Wait until an object becomes available. */
   PLASMA_WAIT,
   /** Wait until an object becomes available. */
-  PLASMA_WAIT1
+  PLASMA_WAIT1,
+  /** Wait until an object becomes available. */
+  PLASMA_WAIT2
 };
 
 typedef struct {

src/plasma/plasma_client.c

@@ -765,6 +765,60 @@ int plasma_wait_for_objects(plasma_connection *conn,
   return num_objects_ready;
 }
 
+int plasma_wait_for_objects2(plasma_connection *conn,
+                             int num_object_requests,
+                             object_request object_requests[],
+                             int num_ready_objects,
+                             uint64_t timeout_ms) {
+  CHECK(conn != NULL);
+  CHECK(conn->manager_conn >= 0);
+  CHECK(num_object_requests > 0);
+  CHECK(num_ready_objects > 0);
+  CHECK(num_ready_objects <= num_object_requests);
+
+  plasma_request *req = plasma_alloc_request(num_object_requests);
+  for (int i = 0; i < num_object_requests; ++i) {
+    CHECK(object_requests[i].type == PLASMA_QUERY_LOCAL ||
+          object_requests[i].type == PLASMA_QUERY_ANYWHERE);
+    req->object_requests[i] = object_requests[i];
+  }
+  req->num_ready_objects = num_ready_objects;
+  req->timeout = timeout_ms;
+  CHECK(plasma_send_request(conn->manager_conn, PLASMA_WAIT2, req) >= 0);
+  free(req);
+
+  plasma_reply *reply = plasma_alloc_reply(num_object_requests);
+  CHECK(plasma_receive_reply(conn->manager_conn,
+                             plasma_reply_size(num_object_requests),
+                             reply) >= 0);
+  int num_objects_ready = 0;
+  for (int i = 0; i < num_object_requests; ++i) {
+    int type = reply->object_requests[i].type;
+    int status = reply->object_requests[i].status;
+    object_requests[i].object_id = reply->object_requests[i].object_id;
+    object_requests[i].type = type;
+    object_requests[i].status = status;
+    switch (type) {
+    case PLASMA_QUERY_LOCAL:
+      if (status == PLASMA_OBJECT_LOCAL) {
+        num_objects_ready += 1;
+      }
+      break;
+    case PLASMA_QUERY_ANYWHERE:
+      if (status == PLASMA_OBJECT_LOCAL || status == PLASMA_OBJECT_REMOTE) {
+        num_objects_ready += 1;
+      } else {
+        CHECK(status == PLASMA_OBJECT_NONEXISTENT);
+      }
+      break;
+    default:
+      LOG_FATAL("This code should be unreachable.");
+    }
+  }
+  free(reply);
+  return num_objects_ready;
+}
+
 /*
  *  TODO: maybe move the plasma_client_* functions in another file.
  *

src/plasma/plasma_client.h

@@ -391,10 +391,10 @@ int plasma_info(plasma_connection *conn,
  * @param object_requests Object event array. Each element contains a request
  *        for a particular object_id. The type of request is specified in the
  *        "type" field.
- *        - A PLASMA_OBJECT_LOCAL request is satisfied when object_id becomes
+ *        - A PLASMA_QUERY_LOCAL request is satisfied when object_id becomes
  *          available in the local Plasma Store. In this case, this function
  *          sets the "status" field to PLASMA_OBJECT_LOCAL.
- *        - A PLASMA_OBJECT_ANYWHERE request is satisfied when object_id becomes
+ *        - A PLASMA_QUERY_ANYWHERE request is satisfied when object_id becomes
  *          available either at the local Plasma Store or on a remote Plasma
  *          Store. In this case, the functions sets the "status" field to
  *          PLASMA_OBJECT_LOCAL or PLASMA_OBJECT_REMOTE.
@@ -414,6 +414,39 @@ int plasma_wait_for_objects(plasma_connection *conn,
                             int num_ready_objects,
                             uint64_t timeout_ms);
 
+/**
+ * Wait for (1) a specified number of objects to be available (sealed) in the
+ * local Plasma Store or in a remote Plasma Store, or (2) for a timeout to
+ * expire. This is a blocking call.
+ *
+ * @param conn The object containing the connection state.
+ * @param num_object_requests Size of the object_requests array.
+ * @param object_requests Object event array. Each element contains a request
+ *        for a particular object_id. The type of request is specified in the
+ *        "type" field.
+ *        - A PLASMA_QUERY_LOCAL request is satisfied when object_id becomes
+ *          available in the local Plasma Store. In this case, this function
+ *          sets the "status" field to PLASMA_OBJECT_LOCAL.
+ *        - A PLASMA_QUERY_ANYWHERE request is satisfied when object_id becomes
+ *          available either at the local Plasma Store or on a remote Plasma
+ *          Store. In this case, the functions sets the "status" field to
+ *          PLASMA_OBJECT_LOCAL or PLASMA_OBJECT_REMOTE.
+ * @param num_ready_objects The number of requests in object_requests array that
+ *        must be satisfied before the function returns, unless it timeouts.
+ *        The num_ready_objects should be no larger than num_object_requests.
+ * @param timeout_ms  Timeout value in milliseconds. If this timeout expires
+ *        before min_num_ready_objects of requests are satisfied, the function
+ *        returns.
+ * @return Number of satisfied requests in the object_requests list. If the
+ *         returned number is less than min_num_ready_objects this means that
+ *         timeout expired.
+ */
+int plasma_wait_for_objects2(plasma_connection *conn,
+                             int num_object_requests,
+                             object_request object_requests[],
+                             int num_ready_objects,
+                             uint64_t timeout_ms);
+
 /**
  * TODO: maybe move the plasma_client_* functions in another file.
  *

src/plasma/plasma_extension.c

@@ -278,6 +278,81 @@ PyObject *PyPlasma_wait(PyObject *self, PyObject *args) {
   return t;
 }
 
+PyObject *PyPlasma_wait2(PyObject *self, PyObject *args) {
+  plasma_connection *conn;
+  PyObject *object_id_list;
+  long long timeout;
+  int num_returns;
+  if (!PyArg_ParseTuple(args, "O&OLi", PyObjectToPlasmaConnection, &conn,
+                        &object_id_list, &timeout, &num_returns)) {
+    return NULL;
+  }
+  Py_ssize_t n = PyList_Size(object_id_list);
+
+  if (!plasma_manager_is_connected(conn)) {
+    PyErr_SetString(PyExc_RuntimeError, "Not connected to the plasma manager");
+    return NULL;
+  }
+  if (num_returns < 0) {
+    PyErr_SetString(PyExc_RuntimeError,
+                    "The argument num_returns cannot be less than zero.");
+    return NULL;
+  }
+  if (num_returns > n) {
+    PyErr_SetString(
+        PyExc_RuntimeError,
+        "The argument num_returns cannot be greater than len(object_ids)");
+    return NULL;
+  }
+  int64_t threshold = 1 << 30;
+  if (timeout > threshold) {
+    PyErr_SetString(PyExc_RuntimeError,
+                    "The argument timeout cannot be greater than 2 ** 30.");
+    return NULL;
+  }
+
+  object_request *object_requests = malloc(sizeof(object_request) * n);
+  for (int i = 0; i < n; ++i) {
+    PyObjectToUniqueID(PyList_GetItem(object_id_list, i),
+                       &object_requests[i].object_id);
+    object_requests[i].type = PLASMA_QUERY_ANYWHERE;
+  }
+  /* Drop the global interpreter lock while we are waiting, so other threads can
+   * run. */
+  int num_return_objects;
+  Py_BEGIN_ALLOW_THREADS;
+  num_return_objects = plasma_wait_for_objects2(
+      conn, (int) n, object_requests, num_returns, (uint64_t) timeout);
+  Py_END_ALLOW_THREADS;
+
+  int num_to_return = MIN(num_return_objects, num_returns);
+  PyObject *ready_ids = PyList_New(num_to_return);
+  PyObject *waiting_ids = PySet_New(object_id_list);
+  int num_returned = 0;
+  for (int i = 0; i < n; ++i) {
+    if (num_returned == num_to_return) {
+      break;
+    }
+    if (object_requests[i].status == PLASMA_OBJECT_LOCAL ||
+        object_requests[i].status == PLASMA_OBJECT_REMOTE) {
+      PyObject *ready =
+          PyString_FromStringAndSize((char *) object_requests[i].object_id.id,
+                                     sizeof(object_requests[i].object_id));
+      PyList_SetItem(ready_ids, num_returned, ready);
+      PySet_Discard(waiting_ids, ready);
+      num_returned += 1;
+    } else {
+      CHECK(object_requests[i].status == PLASMA_OBJECT_NONEXISTENT);
+    }
+  }
+  CHECK(num_returned == num_to_return);
+  /* Return both the ready IDs and the remaining IDs. */
+  PyObject *t = PyTuple_New(2);
+  PyTuple_SetItem(t, 0, ready_ids);
+  PyTuple_SetItem(t, 1, waiting_ids);
+  return t;
+}
+
 PyObject *PyPlasma_evict(PyObject *self, PyObject *args) {
   plasma_connection *conn;
   long long num_bytes;
@@ -371,7 +446,7 @@ static PyMethodDef plasma_methods[] = {
      "Fetch the object from another plasma manager instance."},
     {"fetch2", PyPlasma_fetch2, METH_VARARGS,
      "Fetch the object from another plasma manager instance."},
-    {"wait", PyPlasma_wait, METH_VARARGS,
+    {"wait", PyPlasma_wait2, METH_VARARGS,
      "Wait until num_returns objects in object_ids are ready."},
     {"evict", PyPlasma_evict, METH_VARARGS,
      "Evict some objects until we recover some number of bytes."},

src/plasma/plasma_manager.c

@@ -162,6 +162,39 @@ typedef struct {
   UT_hash_handle hh;
 } fetch_request2;
 
+typedef struct {
+  /** The client connection that called wait. */
+  client_connection *client_conn;
+  /** The ID of the timer that will time out and cause this wait to return to
+   *  the client if it hasn't already returned. */
+  int64_t timer;
+  /** The number of objects in this wait request. */
+  int64_t num_object_requests;
+  /** The object requests for this wait request. Each object request has a
+   *  status field which is either PLASMA_QUERY_LOCAL or PLASMA_QUERY_ANYWHERE.
+   */
+  object_request *object_requests;
+  /** The minimum number of objects to wait for in this request. */
+  int64_t num_objects_to_wait_for;
+  /** The number of object requests in this wait request that are already
+   *  satisfied. */
+  int64_t num_satisfied;
+} wait_request2;
+
+/** This is used to define the utarray of wait requests in the
+ *  object_wait_requests struct. */
+UT_icd wait_request2_icd = {sizeof(wait_request2 *), NULL, NULL, NULL};
+
+typedef struct {
+  /** The ID of the object. This is used as a key in a hash table. */
+  object_id object_id;
+  /** An array of the wait requests involving this object ID. */
+  UT_array *wait_requests;
+  /** Handle for the uthash table in the manager state that keeps track of the
+   *  wait requests involving this object ID. */
+  UT_hash_handle hh;
+} object_wait_requests;
+
 struct plasma_manager_state {
   /** Event loop. */
   event_loop *loop;
@@ -183,6 +216,12 @@ struct plasma_manager_state {
   /** Hash table of outstanding fetch requests. The key is the object ID. The
    *  value is the data needed to perform the fetch. */
   fetch_request2 *fetch_requests2;
+  /** A hash table mapping object IDs to a vector of the wait requests that
+   *  are waiting for the object to arrive locally. */
+  object_wait_requests *object_wait_requests_local;
+  /** A hash table mapping object IDs to a vector of the wait requests that
+   *  are waiting for the object to be available somewhere in the system. */
+  object_wait_requests *object_wait_requests_remote;
   /** Initialize an empty hash map for the cache of local available object. */
   available_object *local_available_objects;
 };
@@ -372,6 +411,143 @@ void remove_object_request(client_connection *client_conn,
   free_client_object_request(object_req);
 }
 
+object_wait_requests **object_wait_requests_table_ptr_from_type(
+    plasma_manager_state *manager_state,
+    int type) {
+  /* We use different types of hash tables for different requests. */
+  if (type == PLASMA_QUERY_LOCAL) {
+    return &manager_state->object_wait_requests_local;
+  } else if (type == PLASMA_QUERY_ANYWHERE) {
+    return &manager_state->object_wait_requests_remote;
+  } else {
+    LOG_FATAL("This code should be unreachable.");
+  }
+}
+
+void add_wait_request_for_object(plasma_manager_state *manager_state,
+                                 object_id object_id,
+                                 int type,
+                                 wait_request2 *wait_req) {
+  object_wait_requests **object_wait_requests_table_ptr =
+      object_wait_requests_table_ptr_from_type(manager_state, type);
+  object_wait_requests *object_wait_reqs;
+  HASH_FIND(hh, *object_wait_requests_table_ptr, &object_id, sizeof(object_id),
+            object_wait_reqs);
+  /* If there are currently no wait requests involving this object ID, create a
+   * new object_wait_requests struct for this object ID and add it to the hash
+   * table. */
+  if (object_wait_reqs == NULL) {
+    object_wait_reqs = malloc(sizeof(object_wait_requests));
+    object_wait_reqs->object_id = object_id;
+    utarray_new(object_wait_reqs->wait_requests, &wait_request2_icd);
+    HASH_ADD(hh, *object_wait_requests_table_ptr, object_id,
+             sizeof(object_wait_reqs->object_id), object_wait_reqs);
+  }
+  /* Add this wait request to the vector of wait requests involving this object
+   * ID. */
+  utarray_push_back(object_wait_reqs->wait_requests, &wait_req);
+}
+
+void remove_wait_request_for_object(plasma_manager_state *manager_state,
+                                    object_id object_id,
+                                    int type,
+                                    wait_request2 *wait_req) {
+  object_wait_requests **object_wait_requests_table_ptr =
+      object_wait_requests_table_ptr_from_type(manager_state, type);
+  object_wait_requests *object_wait_reqs;
+  HASH_FIND(hh, *object_wait_requests_table_ptr, &object_id, sizeof(object_id),
+            object_wait_reqs);
+  /* If there is a vector of wait requests for this object ID, and if this
+   * vector contains the wait request, then remove the wait request from the
+   * vector. */
+  if (object_wait_reqs != NULL) {
+    for (int i = 0; i < utarray_len(object_wait_reqs->wait_requests); ++i) {
+      wait_request2 **wait_req_ptr =
+          (wait_request2 **) utarray_eltptr(object_wait_reqs->wait_requests, i);
+      if (*wait_req_ptr == wait_req) {
+        /* Remove the wait request from the array. */
+        utarray_erase(object_wait_reqs->wait_requests, i, 1);
+        break;
+      }
+    }
+    /* In principle, if there are no more wait requests involving this object
+     * ID, then we could remove the object_wait_reqs struct. */
+  }
+}
+
+void remove_wait_request2(plasma_manager_state *manager_state,
+                          wait_request2 *wait_req) {
+  if (wait_req->timer != -1) {
+    CHECK(event_loop_remove_timer(manager_state->loop, wait_req->timer) ==
+          AE_OK);
+  }
+  free(wait_req->object_requests);
+  free(wait_req);
+}
+
+void return_from_wait2(plasma_manager_state *manager_state,
+                       wait_request2 *wait_req) {
+  plasma_reply *reply = plasma_alloc_reply(wait_req->num_object_requests);
+  reply->num_object_ids = wait_req->num_object_requests;
+  for (int i = 0; i < wait_req->num_object_requests; ++i) {
+    reply->object_requests[i] = wait_req->object_requests[i];
+  }
+  /* Send the reply to the client. */
+  CHECK(plasma_send_reply(wait_req->client_conn->fd, reply) >= 0);
+  free(reply);
+  /* Remove the wait request from each of the relevant object_wait_requests hash
+   * tables if it is present there. */
+  for (int i = 0; i < wait_req->num_object_requests; ++i) {
+    remove_wait_request_for_object(manager_state,
+                                   wait_req->object_requests[i].object_id,
+                                   wait_req->object_requests[i].type, wait_req);
+  }
+  /* Remove the wait request. */
+  remove_wait_request2(manager_state, wait_req);
+}
+
+void update_object_wait_requests(plasma_manager_state *manager_state,
+                                 object_id obj_id,
+                                 int type,
+                                 int status) {
+  object_wait_requests **object_wait_requests_table_ptr =
+      object_wait_requests_table_ptr_from_type(manager_state, type);
+  /* Update the in-progress wait requests in the specified table. */
+  object_wait_requests *object_wait_reqs;
+  HASH_FIND(hh, *object_wait_requests_table_ptr, &obj_id, sizeof(obj_id),
+            object_wait_reqs);
+  if (object_wait_reqs != NULL) {
+    for (int i = 0; i < utarray_len(object_wait_reqs->wait_requests); ++i) {
+      wait_request2 **wait_req_ptr =
+          (wait_request2 **) utarray_eltptr(object_wait_reqs->wait_requests, i);
+      wait_request2 *wait_req = *wait_req_ptr;
+      wait_req->num_satisfied += 1;
+      /* Mark the object as present in the wait request. */
+      int j = 0;
+      for (; j < wait_req->num_object_requests; ++j) {
+        if (object_ids_equal(wait_req->object_requests[j].object_id, obj_id)) {
+          /* Check that this object is currently nonexistent. */
+          CHECK(wait_req->object_requests[j].status ==
+                PLASMA_OBJECT_NONEXISTENT);
+          wait_req->object_requests[j].status = status;
+          break;
+        }
+      }
+      /* Make sure that we actually marked an object as available.*/
+      CHECK(j != wait_req->num_object_requests);
+      /* If this wait request is done, reply to the client. */
+      if (wait_req->num_satisfied == wait_req->num_object_requests) {
+        return_from_wait2(manager_state, wait_req);
+      }
+    }
+    /* Remove the array of wait requests for this object, since no one should be
+     * waiting for this object anymore. */
+    HASH_DELETE(hh, *object_wait_requests_table_ptr, object_wait_reqs);
+    utarray_free(object_wait_reqs->wait_requests);
+    free(object_wait_reqs);
+  }
+}
+
 void remove_fetch_request(plasma_manager_state *manager_state,
                           fetch_request2 *fetch_req) {
   /* Remove the fetch request from the table of fetch requests. */
@@ -403,6 +579,8 @@ plasma_manager_state *init_plasma_manager_state(const char *store_socket_name,
   state->manager_connections = NULL;
   state->fetch_requests = NULL;
   state->fetch_requests2 = NULL;
+  state->object_wait_requests_local = NULL;
+  state->object_wait_requests_remote = NULL;
   if (db_addr) {
     state->db = db_connect(db_addr, db_port, "plasma_manager", manager_addr,
                            manager_port);
@@ -1038,7 +1216,7 @@ void process_fetch_requests2(client_connection *client_conn,
      * initially empty. */
     retry_info retry;
     memset(&retry, 0, sizeof(retry));
-    retry.num_retries = NUM_RETRIES;
+    retry.num_retries = 0;
     retry.timeout = MANAGER_TIMEOUT;
     retry.fail_callback = fatal_table_callback;
     object_table_subscribe(manager_state->db, obj_id, request_transfer2,
@@ -1126,6 +1304,12 @@ int wait_timeout_handler1(event_loop *loop, timer_id id, void *context) {
   return EVENT_LOOP_TIMER_DONE;
 }
 
+int wait_timeout_handler2(event_loop *loop, timer_id id, void *context) {
+  wait_request2 *wait_req = context;
+  return_from_wait2(wait_req->client_conn->manager_state, wait_req);
+  return EVENT_LOOP_TIMER_DONE;
+}
+
 void process_wait_request1(client_connection *client_conn,
                            int num_object_requests,
                            object_request object_requests[],
@@ -1223,6 +1407,97 @@ void process_wait_request1(client_connection *client_conn,
   }
 }
 
+void object_present_callback(object_id object_id,
+                             int manager_count,
+                             const char *manager_vector[],
+                             void *context) {
+  plasma_manager_state *manager_state = (plasma_manager_state *) context;
+  /* This callback is called from object_table_subscribe, which guarantees that
+   * the manager vector contains at least one element. */
+  CHECK(manager_count >= 1);
+
+  /* Update the in-progress remote wait requests. */
+  update_object_wait_requests(manager_state, object_id, PLASMA_QUERY_ANYWHERE,
+                              PLASMA_OBJECT_REMOTE);
+}
+
+void process_wait_request2(client_connection *client_conn,
+                           int num_object_requests,
+                           object_request object_requests[],
+                           uint64_t timeout_ms,
+                           int num_ready_objects) {
+  CHECK(client_conn != NULL);
+  plasma_manager_state *manager_state = client_conn->manager_state;
+
+  /* Create a wait request for this object. */
+  wait_request2 *wait_req = malloc(sizeof(wait_request2));
+  memset(wait_req, 0, sizeof(wait_request2));
+  wait_req->client_conn = client_conn;
+  wait_req->timer = -1;
+  wait_req->num_object_requests = num_object_requests;
+  wait_req->object_requests =
+      malloc(num_object_requests * sizeof(object_request));
+  for (int i = 0; i < num_object_requests; ++i) {
+    wait_req->object_requests[i].object_id = object_requests[i].object_id;
+    wait_req->object_requests[i].type = object_requests[i].type;
+    wait_req->object_requests[i].status = PLASMA_OBJECT_NONEXISTENT;
+  }
+  wait_req->num_objects_to_wait_for = num_ready_objects;
+  wait_req->num_satisfied = 0;
+
+  for (int i = 0; i < num_object_requests; ++i) {
+    object_id obj_id = object_requests[i].object_id;
+
+    /* Check if this object is already present locally. If so, mark the object
+     * as present. */
+    if (is_object_local(manager_state, obj_id)) {
+      wait_req->object_requests[i].status = PLASMA_OBJECT_LOCAL;
+      wait_req->num_satisfied += 1;
+      continue;
+    }
+
+    /* Add the wait request to the relevant data structures. */
+    add_wait_request_for_object(manager_state, obj_id,
+                                wait_req->object_requests[i].type, wait_req);
+
+    if (wait_req->object_requests[i].type == PLASMA_QUERY_LOCAL) {
+      /* TODO(rkn): If desired, we could issue a fetch command here to retrieve
+       * the object. */
+    } else if (wait_req->object_requests[i].type == PLASMA_QUERY_ANYWHERE) {
+      /* Subscribe to a notification for when the object is available somewhere
+       * in the system. */
+      retry_info retry;
+      memset(&retry, 0, sizeof(retry));
+      retry.num_retries = 0;
+      /* TODO(rkn): This timeout is excessive. However, the number of calls to
+       * object_table_subscribe here is also excessive. The issue may be the
+       * number of timers added to the manager event loop. Under heavy usage,
+       * this will trigger the fatal failure callback. The solution is probably
+       * to use Redis modules to write a special purpose command so that we only
+       * need to do a single call to Redis here (and hence create only a single
+       * timer). */
+      retry.timeout = 100000;
+      retry.fail_callback = fatal_table_callback;
+      object_table_subscribe(manager_state->db, obj_id, object_present_callback,
+                             manager_state, &retry, NULL, NULL);
+    } else {
+      /* This code should be unreachable. */
+      CHECK(0);
+    }
+  }
+
+  /* If enough of the wait requests have already been satisfied, return to the
+   * client. */
+  if (wait_req->num_satisfied >= wait_req->num_objects_to_wait_for) {
+    return_from_wait2(manager_state, wait_req);
+    return;
+  }
+
+  /* Set a timer that will cause the wait request to return to the client. */
+  wait_req->timer = event_loop_add_timer(manager_state->loop, timeout_ms,
+                                         wait_timeout_handler2, wait_req);
+}
+
 /* TODO(pcm): unify with wait_object_available_callback. */
 void wait_object_lookup_callback(object_id object_id,
                                  int manager_count,
@@ -1539,6 +1814,13 @@ void process_object_notification(event_loop *loop,
     remove_fetch_request(state, fetch_req);
     /* TODO(rkn): We also really should unsubscribe from the object table. */
   }
+
+  /* Update the in-progress local and remote wait requests. */
+  update_object_wait_requests(state, obj_id, PLASMA_QUERY_LOCAL,
+                              PLASMA_OBJECT_LOCAL);
+  update_object_wait_requests(state, obj_id, PLASMA_QUERY_ANYWHERE,
+                              PLASMA_OBJECT_LOCAL);
+
   /* Notify any clients who were waiting on a fetch to this object and tick
    * off objects we are waiting for. */
   client_object_request *object_req, *next;
@@ -1626,6 +1908,11 @@ void process_message(event_loop *loop,
     process_wait_request1(conn, req->num_object_ids, req->object_requests,
                           req->timeout, req->num_ready_objects);
     break;
+  case PLASMA_WAIT2:
+    LOG_DEBUG("Processing wait2");
+    process_wait_request2(conn, req->num_object_ids, req->object_requests,
+                          req->timeout, req->num_ready_objects);
+    break;
   case PLASMA_STATUS:
     LOG_DEBUG("Processing status");
     DCHECK(req->num_object_ids == 1);

src/plasma/test/client_tests.c

@@ -160,8 +160,8 @@ TEST plasma_wait_for_objects_tests(void) {
 
   struct timeval start, end;
   gettimeofday(&start, NULL);
-  int n = plasma_wait_for_objects(plasma_conn1, NUM_OBJ_REQUEST, obj_requests,
-                                  NUM_OBJ_REQUEST, WAIT_TIMEOUT_MS);
+  int n = plasma_wait_for_objects2(plasma_conn1, NUM_OBJ_REQUEST, obj_requests,
+                                   NUM_OBJ_REQUEST, WAIT_TIMEOUT_MS);
   ASSERT(n == 0);
   gettimeofday(&end, NULL);
   float diff_ms = (end.tv_sec - start.tv_sec);
@@ -177,34 +177,32 @@ TEST plasma_wait_for_objects_tests(void) {
   plasma_create(plasma_conn1, oid1, data_size, metadata, metadata_size, &data);
   plasma_seal(plasma_conn1, oid1);
 
-  sleep(1);
-  n = plasma_wait_for_objects(plasma_conn1, NUM_OBJ_REQUEST, obj_requests,
-                              NUM_OBJ_REQUEST, WAIT_TIMEOUT_MS);
+  n = plasma_wait_for_objects2(plasma_conn1, NUM_OBJ_REQUEST, obj_requests,
+                               NUM_OBJ_REQUEST, WAIT_TIMEOUT_MS);
   ASSERT(n == 1);
 
-  /* Create and insert an object in plasma_conn1. */
+  /* Create and insert an object in plasma_conn2. */
   plasma_create(plasma_conn2, oid2, data_size, metadata, metadata_size, &data);
   plasma_seal(plasma_conn2, oid2);
 
-  n = plasma_wait_for_objects(plasma_conn1, NUM_OBJ_REQUEST, obj_requests,
-                              NUM_OBJ_REQUEST, WAIT_TIMEOUT_MS);
+  n = plasma_wait_for_objects2(plasma_conn1, NUM_OBJ_REQUEST, obj_requests,
+                               NUM_OBJ_REQUEST, WAIT_TIMEOUT_MS);
   ASSERT(n == 2);
 
-  n = plasma_wait_for_objects(plasma_conn2, NUM_OBJ_REQUEST, obj_requests,
-                              NUM_OBJ_REQUEST, WAIT_TIMEOUT_MS);
+  n = plasma_wait_for_objects2(plasma_conn2, NUM_OBJ_REQUEST, obj_requests,
+                               NUM_OBJ_REQUEST, WAIT_TIMEOUT_MS);
   ASSERT(n == 2);
 
   obj_requests[0].type = PLASMA_QUERY_LOCAL;
   obj_requests[1].type = PLASMA_QUERY_LOCAL;
-  n = plasma_wait_for_objects(plasma_conn1, NUM_OBJ_REQUEST, obj_requests,
-                              NUM_OBJ_REQUEST, WAIT_TIMEOUT_MS);
+  n = plasma_wait_for_objects2(plasma_conn1, NUM_OBJ_REQUEST, obj_requests,
+                               NUM_OBJ_REQUEST, WAIT_TIMEOUT_MS);
   ASSERT(n == 1);
 
-  n = plasma_wait_for_objects(plasma_conn2, NUM_OBJ_REQUEST, obj_requests,
-                              NUM_OBJ_REQUEST, WAIT_TIMEOUT_MS);
+  n = plasma_wait_for_objects2(plasma_conn2, NUM_OBJ_REQUEST, obj_requests,
+                               NUM_OBJ_REQUEST, WAIT_TIMEOUT_MS);
   ASSERT(n == 1);
 
-  sleep(1);
   plasma_disconnect(plasma_conn1);
   plasma_disconnect(plasma_conn2);
 
@@ -364,7 +362,7 @@ SUITE(plasma_client_tests) {
   RUN_TEST(plasma_status_tests);
   RUN_TEST(plasma_fetch_remote_tests);
   RUN_TEST(plasma_get_local_tests);
-  // RUN_TEST(plasma_wait_for_objects_tests);
+  RUN_TEST(plasma_wait_for_objects_tests);
   RUN_TEST(plasma_get_tests);
   RUN_TEST(plasma_wait_tests);
   RUN_TEST(plasma_multiget_tests);

src/plasma/test/test.py

@@ -554,19 +554,16 @@ class TestPlasmaManager(unittest.TestCase):
     self.client1.create(obj_id1, 1000)
     self.client1.seal(obj_id1)
     ready, waiting = self.client1.wait([obj_id1], timeout=100, num_returns=1)
-    self.assertEqual(len(ready), 1)
-    self.assertEqual(ready[0], obj_id1)
-    self.assertEqual(len(waiting), 0)
+    self.assertEqual(set(ready), set([obj_id1]))
+    self.assertEqual(waiting, [])
 
     # Test wait if only one object available and only one object waited for.
     obj_id2 = random_object_id()
     self.client1.create(obj_id2, 1000)
     # Don't seal.
     ready, waiting = self.client1.wait([obj_id2, obj_id1], timeout=100, num_returns=1)
-    self.assertEqual(len(ready), 1)
-    self.assertEqual(ready[0], obj_id1)
-    self.assertEqual(len(waiting), 1)
-    self.assertEqual(waiting[0], obj_id2)
+    self.assertEqual(set(ready), set([obj_id1]))
+    self.assertEqual(set(waiting), set([obj_id2]))
 
     # Test wait if object is sealed later.
     obj_id3 = random_object_id()
@@ -574,26 +571,57 @@ class TestPlasmaManager(unittest.TestCase):
     def finish():
       self.client2.create(obj_id3, 1000)
       self.client2.seal(obj_id3)
-      self.client2.transfer("127.0.0.1", self.port1, obj_id3)
 
     t = threading.Timer(0.1, finish)
     t.start()
     ready, waiting = self.client1.wait([obj_id3, obj_id2, obj_id1], timeout=1000, num_returns=2)
-    self.assertEqual(len(ready), 2)
-    self.assertTrue((ready[0] == obj_id1 and ready[1] == obj_id3) or (ready[0] == obj_id3 and ready[1] == obj_id1))
-    self.assertEqual(len(waiting), 1)
-    self.assertTrue(waiting[0] == obj_id2)
+    self.assertEqual(set(ready), set([obj_id1, obj_id3]))
+    self.assertEqual(set(waiting), set([obj_id2]))
 
     # Test if the appropriate number of objects is shown if some objects are not ready
-    ready, wait = self.client1.wait([obj_id3, obj_id2, obj_id1], 100, 3)
-    self.assertEqual(len(ready), 2)
-    self.assertTrue((ready[0] == obj_id1 and ready[1] == obj_id3) or (ready[0] == obj_id3 and ready[1] == obj_id1))
-    self.assertEqual(len(waiting), 1)
-    self.assertTrue(waiting[0] == obj_id2)
+    ready, waiting = self.client1.wait([obj_id3, obj_id2, obj_id1], 100, 3)
+    self.assertEqual(set(ready), set([obj_id1, obj_id3]))
+    self.assertEqual(set(waiting), set([obj_id2]))
 
     # Don't forget to seal obj_id2.
     self.client1.seal(obj_id2)
 
+    # Test calling wait a bunch of times.
+    object_ids = []
+    # TODO(rkn): Increasing n to 100 (or larger) will cause failures. The
+    # problem appears to be that the number of timers added to the manager event
+    # loop slow down the manager so much that some of the asynchronous Redis
+    # commands timeout triggering fatal failure callbacks.
+    n = 40
+    for i in range(n * (n + 1) / 2):
+      if i % 2 == 0:
+        object_id, _, _ = create_object(self.client1, 200, 200)
+      else:
+        object_id, _, _ = create_object(self.client2, 200, 200)
+      object_ids.append(object_id)
+    # Try waiting for all of the object IDs on the first client.
+    waiting = object_ids
+    retrieved = []
+    for i in range(1, n + 1):
+      ready, waiting = self.client1.wait(waiting, timeout=1000, num_returns=i)
+      self.assertEqual(len(ready), i)
+      retrieved += ready
+    self.assertEqual(set(retrieved), set(object_ids))
+    ready, waiting = self.client1.wait(object_ids, timeout=1000, num_returns=len(object_ids))
+    self.assertEqual(set(ready), set(object_ids))
+    self.assertEqual(waiting, [])
+    # Try waiting for all of the object IDs on the second client.
+    waiting = object_ids
+    retrieved = []
+    for i in range(1, n + 1):
+      ready, waiting = self.client2.wait(waiting, timeout=1000, num_returns=i)
+      self.assertEqual(len(ready), i)
+      retrieved += ready
+    self.assertEqual(set(retrieved), set(object_ids))
+    ready, waiting = self.client2.wait(object_ids, timeout=1000, num_returns=len(object_ids))
+    self.assertEqual(set(ready), set(object_ids))
+    self.assertEqual(waiting, [])
+
   def test_transfer(self):
     for _ in range(100):
       # Create an object.

test/stress_tests.py

@@ -77,9 +77,7 @@ class TaskTests(unittest.TestCase):
     ray.worker.cleanup()
 
   def testWait(self):
-    # TODO(rkn): Use more local schedulers once the distributed wait
-    # implementation is in place.
-    for num_local_schedulers in [1]:
+    for num_local_schedulers in [1, 4]:
       for num_workers_per_scheduler in [4]:
         num_workers = num_local_schedulers * num_workers_per_scheduler
         ray.init(start_ray_local=True, num_workers=num_workers, num_local_schedulers=num_local_schedulers)
@@ -88,7 +86,7 @@ class TaskTests(unittest.TestCase):
         def f(x):
           return x
 
-        x_ids = [f.remote(i) for i in range(1000)]
+        x_ids = [f.remote(i) for i in range(100)]
         for i in range(len(x_ids)):
           ray.wait([x_ids[i]])
         for i in range(len(x_ids) - 1):
@@ -99,7 +97,7 @@ class TaskTests(unittest.TestCase):
           time.sleep(x)
 
         for i in range(1, 5):
-          x_ids = [g.remote(np.random.uniform(0, i)) for _ in range(4 * num_workers)]
+          x_ids = [g.remote(np.random.uniform(0, i)) for _ in range(2 * num_workers)]
           ray.wait(x_ids, num_returns=len(x_ids))
 
         self.assertTrue(ray.services.all_processes_alive())