diff --git a/src/plasma/plasma.c b/src/plasma/plasma.c index 550df0f8f..b9eb9350d 100644 --- a/src/plasma/plasma.c +++ b/src/plasma/plasma.c @@ -88,3 +88,14 @@ int plasma_receive_request(int sock, int64_t *type, plasma_request **request) { } return length == plasma_request_size((*request)->num_object_ids) ? 0 : -1; } + +bool plasma_object_ids_distinct(int num_object_ids, object_id object_ids[]) { + for (int i = 0; i < num_object_ids; ++i) { + for (int j = 0; j < i; ++j) { + if (object_ids_equal(object_ids[i], object_ids[j])) { + return false; + } + } + } + return true; +} diff --git a/src/plasma/plasma.h b/src/plasma/plasma.h index 037a8b75b..58336fd12 100644 --- a/src/plasma/plasma.h +++ b/src/plasma/plasma.h @@ -255,4 +255,13 @@ int plasma_send_request(int sock, int64_t type, plasma_request *request); */ int plasma_receive_request(int sock, int64_t *type, plasma_request **request); +/** + * Check if a collection of object IDs contains any duplicates. + * + * @param num_object_ids The number of object IDs. + * @param object_ids[] The list of object IDs to check. + * @return True if the object IDs are all distinct and false otherwise. + */ +bool plasma_object_ids_distinct(int num_object_ids, object_id object_ids[]); + #endif /* PLASMA_H */ diff --git a/src/plasma/plasma_client.c b/src/plasma/plasma_client.c index dd971984f..a441351c0 100644 --- a/src/plasma/plasma_client.c +++ b/src/plasma/plasma_client.c @@ -458,6 +458,9 @@ void plasma_fetch(plasma_connection *conn, object_id object_ids[], int is_fetched[]) { CHECK(conn->manager_conn >= 0); + /* Make sure that there are no duplicated object IDs. TODO(rkn): we should + * allow this case in the future. */ + CHECK(plasma_object_ids_distinct(num_object_ids, object_ids)); plasma_request *req = plasma_alloc_request(num_object_ids, object_ids); LOG_DEBUG("Requesting fetch"); CHECK(plasma_send_request(conn->manager_conn, PLASMA_FETCH, req) >= 0); @@ -481,15 +484,15 @@ void plasma_fetch(plasma_connection *conn, /* Update the correct index in is_fetched. */ int i = 0; for (; i < num_object_ids; i++) { - if (object_ids_equal(object_ids[i], reply.object_ids[0])) { - /* Check that this isn't a duplicate response. */ - CHECK(!is_fetched[i]); + if (object_ids_equal(object_ids[i], reply.object_ids[0]) && + !is_fetched[i]) { is_fetched[i] = success; break; } } CHECKM(i != num_object_ids, - "Received unexpected object ID from manager during fetch."); + "Received an unexpected object ID from manager during fetch or the " + "object ID was received multiple times."); } } diff --git a/src/plasma/plasma_extension.c b/src/plasma/plasma_extension.c index aa6152c0b..441c7a146 100644 --- a/src/plasma/plasma_extension.c +++ b/src/plasma/plasma_extension.c @@ -135,18 +135,25 @@ PyObject *PyPlasma_fetch(PyObject *self, PyObject *args) { &object_id_list)) { return NULL; } - if (!plasma_manager_is_connected(conn)) { PyErr_SetString(PyExc_RuntimeError, "Not connected to the plasma manager"); return NULL; } - Py_ssize_t n = PyList_Size(object_id_list); object_id *object_ids = malloc(sizeof(object_id) * n); for (int i = 0; i < n; ++i) { PyObjectToUniqueID(PyList_GetItem(object_id_list, i), &object_ids[i]); } + /* Check that there are no duplicate object IDs. TODO(rkn): we should allow + * this in the future. */ + if (!plasma_object_ids_distinct(n, object_ids)) { + PyErr_SetString(PyExc_RuntimeError, + "The same object ID is used multiple times in this call to " + "fetch."); + return NULL; + } int *success_array = malloc(sizeof(int) * n); + memset(success_array, 0, sizeof(int) * n); plasma_fetch(conn, (int) n, object_ids, success_array); PyObject *success_list = PyList_New(n); for (int i = 0; i < n; ++i) { diff --git a/src/plasma/plasma_manager.c b/src/plasma/plasma_manager.c index 30a8b7df7..4e2c5a7f7 100644 --- a/src/plasma/plasma_manager.c +++ b/src/plasma/plasma_manager.c @@ -190,6 +190,11 @@ client_object_connection *add_object_connection(client_connection *client_conn, object_conn->manager_vector = NULL; object_conn->next_manager = 0; /* Register the object context with the client context. */ + client_object_connection *temp_object_conn = NULL; + HASH_FIND(active_hh, client_conn->active_objects, &object_id, + sizeof(object_id), temp_object_conn); + CHECKM(temp_object_conn == NULL, + "The hash table already has an object connection for this object ID."); HASH_ADD(active_hh, client_conn->active_objects, object_id, sizeof(object_id), object_conn); /* Register the object context with the manager state. */ @@ -211,6 +216,7 @@ client_object_connection *add_object_connection(client_connection *client_conn, void remove_object_connection(client_connection *client_conn, client_object_connection *object_conn) { /* Deregister the object context with the client context. */ + /* TODO(rkn): Check that object_conn is actually in the hash table. */ HASH_DELETE(active_hh, client_conn->active_objects, object_conn); /* Deregister the object context with the manager state. */ client_object_connection *object_conns; @@ -611,7 +617,6 @@ void request_transfer(object_id object_id, * register a Redis callback for changes to this object table entry. */ free(manager_vector); send_client_failure_reply(object_id, client_conn); - remove_object_connection(client_conn, object_conn); return; } /* Register the new outstanding fetch with the current client connection. */ diff --git a/src/plasma/test/test.py b/src/plasma/test/test.py index af4577185..5112fc5e3 100644 --- a/src/plasma/test/test.py +++ b/src/plasma/test/test.py @@ -291,60 +291,64 @@ class TestPlasmaManager(unittest.TestCase): self.p5.kill() self.redis_process.kill() - # def test_fetch(self): - # if self.redis_process is None: - # print("Cannot test fetch without a running redis instance.") - # self.assertTrue(False) - # for _ in range(100): - # # Create an object. - # object_id1, memory_buffer1, metadata1 = create_object(self.client1, 2000, 2000) - # # Fetch the object from the other plasma store. - # # TODO(swang): This line is a hack! It makes sure that the entry will be - # # in the object table once we call the fetch operation. Remove once - # # retries are implemented by Ray common. - # time.sleep(0.1) - # successes = self.client2.fetch([object_id1]) - # self.assertEqual(successes, [True]) - # # Compare the two buffers. - # assert_get_object_equal(self, self.client1, self.client2, object_id1, - # memory_buffer=memory_buffer1, metadata=metadata1) - # # Fetch in the other direction. These should return quickly because - # # client1 already has the object. - # successes = self.client1.fetch([object_id1]) - # self.assertEqual(successes, [True]) - # assert_get_object_equal(self, self.client2, self.client1, object_id1, - # memory_buffer=memory_buffer1, metadata=metadata1) + def test_fetch(self): + if self.redis_process is None: + print("Cannot test fetch without a running redis instance.") + self.assertTrue(False) + for _ in range(100): + # Create an object. + object_id1, memory_buffer1, metadata1 = create_object(self.client1, 2000, 2000) + # Fetch the object from the other plasma store. + # TODO(swang): This line is a hack! It makes sure that the entry will be + # in the object table once we call the fetch operation. Remove once + # retries are implemented by Ray common. + time.sleep(0.1) + successes = self.client2.fetch([object_id1]) + self.assertEqual(successes, [True]) + # Compare the two buffers. + assert_get_object_equal(self, self.client1, self.client2, object_id1, + memory_buffer=memory_buffer1, metadata=metadata1) + # Fetch in the other direction. These should return quickly because + # client1 already has the object. + successes = self.client1.fetch([object_id1]) + self.assertEqual(successes, [True]) + assert_get_object_equal(self, self.client2, self.client1, object_id1, + memory_buffer=memory_buffer1, metadata=metadata1) - # def test_fetch_multiple(self): - # if self.redis_process is None: - # print("Cannot test fetch without a running redis instance.") - # self.assertTrue(False) - # for _ in range(20): - # # Create two objects and a third fake one that doesn't exist. - # object_id1, memory_buffer1, metadata1 = create_object(self.client1, 2000, 2000) - # missing_object_id = random_object_id() - # object_id2, memory_buffer2, metadata2 = create_object(self.client1, 2000, 2000) - # object_ids = [object_id1, missing_object_id, object_id2] - # # Fetch the objects from the other plasma store. The second object ID - # # should timeout since it does not exist. - # # TODO(swang): This line is a hack! It makes sure that the entry will be - # # in the object table once we call the fetch operation. Remove once - # # retries are implemented by Ray common. - # time.sleep(0.1) - # successes = self.client2.fetch(object_ids) - # self.assertEqual(successes, [True, False, True]) - # # Compare the buffers of the objects that do exist. - # assert_get_object_equal(self, self.client1, self.client2, object_id1, - # memory_buffer=memory_buffer1, metadata=metadata1) - # assert_get_object_equal(self, self.client1, self.client2, object_id2, - # memory_buffer=memory_buffer2, metadata=metadata2) - # # Fetch in the other direction. The fake object still does not exist. - # successes = self.client1.fetch(object_ids) - # self.assertEqual(successes, [True, False, True]) - # assert_get_object_equal(self, self.client2, self.client1, object_id1, - # memory_buffer=memory_buffer1, metadata=metadata1) - # assert_get_object_equal(self, self.client2, self.client1, object_id2, - # memory_buffer=memory_buffer2, metadata=metadata2) + def test_fetch_multiple(self): + if self.redis_process is None: + print("Cannot test fetch without a running redis instance.") + self.assertTrue(False) + for _ in range(20): + # Create two objects and a third fake one that doesn't exist. + object_id1, memory_buffer1, metadata1 = create_object(self.client1, 2000, 2000) + missing_object_id = random_object_id() + object_id2, memory_buffer2, metadata2 = create_object(self.client1, 2000, 2000) + object_ids = [object_id1, missing_object_id, object_id2] + # Fetch the objects from the other plasma store. The second object ID + # should timeout since it does not exist. + # TODO(swang): This line is a hack! It makes sure that the entry will be + # in the object table once we call the fetch operation. Remove once + # retries are implemented by Ray common. + time.sleep(0.1) + successes = self.client2.fetch(object_ids) + self.assertEqual(successes, [True, False, True]) + # Compare the buffers of the objects that do exist. + assert_get_object_equal(self, self.client1, self.client2, object_id1, + memory_buffer=memory_buffer1, metadata=metadata1) + assert_get_object_equal(self, self.client1, self.client2, object_id2, + memory_buffer=memory_buffer2, metadata=metadata2) + # Fetch in the other direction. The fake object still does not exist. + successes = self.client1.fetch(object_ids) + self.assertEqual(successes, [True, False, True]) + assert_get_object_equal(self, self.client2, self.client1, object_id1, + memory_buffer=memory_buffer1, metadata=metadata1) + assert_get_object_equal(self, self.client2, self.client1, object_id2, + memory_buffer=memory_buffer2, metadata=metadata2) + + # Check that calling fetch with the same object ID fails. + object_id = random_object_id() + self.assertRaises(Exception, lambda : self.client1.fetch([object_id, object_id])) def test_wait(self): # Test timeout.