diff --git a/src/plasma/plasma.h b/src/plasma/plasma.h index 0c815be47..f17bfbbf8 100644 --- a/src/plasma/plasma.h +++ b/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 { diff --git a/src/plasma/plasma_client.c b/src/plasma/plasma_client.c index 93eec0ded..56c58557a 100644 --- a/src/plasma/plasma_client.c +++ b/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. * diff --git a/src/plasma/plasma_client.h b/src/plasma/plasma_client.h index bb7feec30..c13464feb 100644 --- a/src/plasma/plasma_client.h +++ b/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. * diff --git a/src/plasma/plasma_extension.c b/src/plasma/plasma_extension.c index 8c9427a73..e9084ff3f 100644 --- a/src/plasma/plasma_extension.c +++ b/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."}, diff --git a/src/plasma/plasma_manager.c b/src/plasma/plasma_manager.c index f172e5ee0..f3d1e7d9b 100644 --- a/src/plasma/plasma_manager.c +++ b/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); diff --git a/src/plasma/test/client_tests.c b/src/plasma/test/client_tests.c index e3406e7d9..2199a8c24 100644 --- a/src/plasma/test/client_tests.c +++ b/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); diff --git a/src/plasma/test/test.py b/src/plasma/test/test.py index 948830b83..391d13284 100644 --- a/src/plasma/test/test.py +++ b/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. diff --git a/test/stress_tests.py b/test/stress_tests.py index 385a3aceb..1cbb83fe8 100644 --- a/test/stress_tests.py +++ b/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())