Re-enable fetch tests. (#63)

* Bring back fetch tests.

* Zero initialize success array in PyPlasma_fetch.

* Fix bug in fetch in case where the object ID doesn't have any managers in the object table.

* Temporarily disallow calling fetch with multiple copies of the same object ID.

* Fix.

* Factor out code for checking if list of object IDs are all distinct.

* Remove commented out code.

* Fix.

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;
+}

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 */

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.");
   }
 }
 

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) {

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. */

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.