Object table remove (#139)

* Object table remove redis module

* Test case for object table remove redis module

* Client code for object_table_remove

* Delete object notifications in plasma

* Test for object deletion notifications

* Fix subscribe deletion test

* Address Robert's comments

* free hash table entry

src/plasma/plasma_extension.c

@@ -328,8 +328,10 @@ PyObject *PyPlasma_receive_notification(PyObject *self, PyObject *args) {
   if (!PyArg_ParseTuple(args, "i", &plasma_sock)) {
     return NULL;
   }
-  /* Receive object notification from the plasma connection socket,
-   * return a tuple of its fields: object_id, data_size, metadata_size. */
+  /* Receive object notification from the plasma connection socket. If the
+   * object was added, return a tuple of its fields: object_id, data_size,
+   * metadata_size. If the object was deleted, data_size and metadata_size will
+   * be set to -1. */
   int nbytes =
       read_bytes(plasma_sock, (uint8_t *) &object_info, sizeof(object_info));
 
@@ -342,8 +344,13 @@ PyObject *PyPlasma_receive_notification(PyObject *self, PyObject *args) {
   PyObject *t = PyTuple_New(3);
   PyTuple_SetItem(t, 0, PyBytes_FromStringAndSize(
                             (char *) object_info.obj_id.id, UNIQUE_ID_SIZE));
-  PyTuple_SetItem(t, 1, PyLong_FromLong(object_info.data_size));
-  PyTuple_SetItem(t, 2, PyLong_FromLong(object_info.metadata_size));
+  if (object_info.is_deletion) {
+    PyTuple_SetItem(t, 1, PyLong_FromLong(-1));
+    PyTuple_SetItem(t, 2, PyLong_FromLong(-1));
+  } else {
+    PyTuple_SetItem(t, 1, PyLong_FromLong(object_info.data_size));
+    PyTuple_SetItem(t, 2, PyLong_FromLong(object_info.metadata_size));
+  }
 
   return t;
 }

src/plasma/plasma_manager.c

@@ -1206,33 +1206,35 @@ void process_status_request(client_connection *client_conn,
                       request_status_done, client_conn);
 }
 
-void process_object_notification(event_loop *loop,
-                                 int client_sock,
-                                 void *context,
-                                 int events) {
-  plasma_manager_state *state = context;
-  object_id obj_id;
-  object_info object_info;
-  retry_info retry = {
-      .num_retries = NUM_RETRIES,
-      .timeout = MANAGER_TIMEOUT,
-      .fail_callback = NULL,
-  };
-  /* Read the notification from Plasma. */
-  int error =
-      read_bytes(client_sock, (uint8_t *) &object_info, sizeof(object_info));
-  if (error < 0) {
-    /* The store has closed the socket. */
-    LOG_DEBUG(
-        "The plasma store has closed the object notification socket, or some "
-        "other error has occurred.");
-    event_loop_remove_file(loop, client_sock);
-    close(client_sock);
-    return;
+void process_delete_object_notification(plasma_manager_state *state,
+                                        object_info object_info) {
+  object_id obj_id = object_info.obj_id;
+  available_object *entry;
+  HASH_FIND(hh, state->local_available_objects, &obj_id, sizeof(obj_id), entry);
+  if (entry != NULL) {
+    HASH_DELETE(hh, state->local_available_objects, entry);
+    free(entry);
   }
-  obj_id = object_info.obj_id;
-  /* Add object to locally available object. */
-  /* TODO(pcm): Where is this deallocated? */
+
+  /* Remove this object from the (redis) object table. */
+  if (state->db) {
+    retry_info retry = {
+        .num_retries = NUM_RETRIES,
+        .timeout = MANAGER_TIMEOUT,
+        .fail_callback = NULL,
+    };
+    object_table_remove(state->db, obj_id, NULL, &retry, NULL, NULL);
+  }
+
+  /* NOTE: There could be pending wait requests for this object that will now
+   * return when the object is not actually available. For simplicity, we allow
+   * this scenario rather than try to keep the wait request statuses exactly
+   * up-to-date. */
+}
+
+void process_add_object_notification(plasma_manager_state *state,
+                                     object_info object_info) {
+  object_id obj_id = object_info.obj_id;
   available_object *entry =
       (available_object *) malloc(sizeof(available_object));
   entry->object_id = obj_id;
@@ -1243,6 +1245,11 @@ void process_object_notification(event_loop *loop,
   if (state->db) {
     /* TODO(swang): Log the error if we fail to add the object, and possibly
      * retry later? */
+    retry_info retry = {
+        .num_retries = NUM_RETRIES,
+        .timeout = MANAGER_TIMEOUT,
+        .fail_callback = NULL,
+    };
     object_table_add(state->db, obj_id,
                      object_info.data_size + object_info.metadata_size,
                      object_info.digest, &retry, NULL, NULL);
@@ -1263,6 +1270,32 @@ void process_object_notification(event_loop *loop,
                               PLASMA_OBJECT_LOCAL);
 }
 
+void process_object_notification(event_loop *loop,
+                                 int client_sock,
+                                 void *context,
+                                 int events) {
+  plasma_manager_state *state = context;
+  object_info object_info;
+  /* Read the notification from Plasma. */
+  int error =
+      read_bytes(client_sock, (uint8_t *) &object_info, sizeof(object_info));
+  if (error < 0) {
+    /* The store has closed the socket. */
+    LOG_DEBUG(
+        "The plasma store has closed the object notification socket, or some "
+        "other error has occurred.");
+    event_loop_remove_file(loop, client_sock);
+    close(client_sock);
+    return;
+  }
+  /* Add object to locally available object. */
+  if (object_info.is_deletion) {
+    process_delete_object_notification(state, object_info);
+  } else {
+    process_add_object_notification(state, object_info);
+  }
+}
+
 void process_message(event_loop *loop,
                      int client_sock,
                      void *context,

src/plasma/plasma_store.c

@@ -103,6 +103,8 @@ plasma_store_state *init_plasma_store(event_loop *loop, int64_t system_memory) {
   return state;
 }
 
+void push_notification(plasma_store_state *state, object_id object_id);
+
 /* If this client is not already using the object, add the client to the
  * object's list of clients, otherwise do nothing. */
 void add_client_to_object_clients(object_table_entry *entry,
@@ -321,12 +323,7 @@ void seal_object(client *client_context,
   /* Set the object digest. */
   memcpy(entry->info.digest, digest, DIGEST_SIZE);
   /* Inform all subscribers that a new object has been sealed. */
-  notification_queue *queue, *temp_queue;
-  HASH_ITER(hh, plasma_state->pending_notifications, queue, temp_queue) {
-    utarray_push_back(queue->object_ids, &object_id);
-    send_notifications(plasma_state->loop, queue->subscriber_fd, plasma_state,
-                       0);
-  }
+  push_notification(plasma_state, object_id);
 
   /* Inform processes getting this object that the object is ready now. */
   object_notify_entry *notify_entry;
@@ -375,6 +372,8 @@ void delete_object(plasma_store_state *plasma_state, object_id object_id) {
   dlfree(pointer);
   utarray_free(entry->clients);
   free(entry);
+  /* Inform all subscribers that the object has been deleted. */
+  push_notification(plasma_state, object_id);
 }
 
 void remove_objects(plasma_store_state *plasma_state,
@@ -390,6 +389,15 @@ void remove_objects(plasma_store_state *plasma_state,
   }
 }
 
+void push_notification(plasma_store_state *plasma_state, object_id object_id) {
+  notification_queue *queue, *temp_queue;
+  HASH_ITER(hh, plasma_state->pending_notifications, queue, temp_queue) {
+    utarray_push_back(queue->object_ids, &object_id);
+    send_notifications(plasma_state->loop, queue->subscriber_fd, plasma_state,
+                       0);
+  }
+}
+
 /* Send more notifications to a subscriber. */
 void send_notifications(event_loop *loop,
                         int client_sock,
@@ -409,9 +417,16 @@ void send_notifications(event_loop *loop,
     /* This object should already exist in plasma store state. */
     HASH_FIND(handle, plasma_state->plasma_store_info->objects, obj_id,
               sizeof(object_id), entry);
-    CHECK(entry != NULL);
 
-    object_info object_info = entry->info;
+    object_info object_info;
+    if (entry == NULL) {
+      memset(&object_info, 0, sizeof(object_info));
+      object_info.obj_id = *obj_id;
+      object_info.is_deletion = true;
+    } else {
+      object_info = entry->info;
+      object_info.is_deletion = false;
+    }
 
     /* Attempt to send a notification about this object ID. */
     int nbytes =

src/plasma/test/test.py

@@ -351,6 +351,63 @@ class TestPlasmaClient(unittest.TestCase):
         self.assertEqual(data_sizes[j], recv_dsize)
         self.assertEqual(metadata_sizes[j], recv_msize)
 
+  def test_subscribe_deletions(self):
+    # Subscribe to notifications from the Plasma Store. We use plasma_client2
+    # to make sure that all used objects will get evicted properly.
+    sock = self.plasma_client2.subscribe()
+    for i in [1, 10, 100, 1000, 10000, 100000]:
+      object_ids = [random_object_id() for _ in range(i)]
+      # Add 1 to the sizes to make sure we have nonzero object sizes.
+      metadata_sizes = [np.random.randint(1000) + 1 for _ in range(i)]
+      data_sizes = [np.random.randint(1000) + 1 for _ in range(i)]
+      for j in range(i):
+        x = self.plasma_client2.create(object_ids[j], size=data_sizes[j],
+                                  metadata=bytearray(np.random.bytes(metadata_sizes[j])))
+        self.plasma_client2.seal(object_ids[j])
+      del x
+      # Check that we received notifications for creating all of the objects.
+      for j in range(i):
+        recv_objid, recv_dsize, recv_msize = self.plasma_client2.get_next_notification()
+        self.assertEqual(object_ids[j], recv_objid)
+        self.assertEqual(data_sizes[j], recv_dsize)
+        self.assertEqual(metadata_sizes[j], recv_msize)
+
+      # Check that we receive notifications for deleting all objects, as we
+      # evict them.
+      for j in range(i):
+        self.assertEqual(self.plasma_client2.evict(1), data_sizes[j] + metadata_sizes[j])
+        recv_objid, recv_dsize, recv_msize = self.plasma_client2.get_next_notification()
+        self.assertEqual(object_ids[j], recv_objid)
+        self.assertEqual(-1, recv_dsize)
+        self.assertEqual(-1, recv_msize)
+
+    # Test multiple deletion notifications. The first 9 object IDs have size 0,
+    # and the last has a nonzero size. When Plasma evicts 1 byte, it will evict
+    # all objects, so we should receive deletion notifications for each.
+    num_object_ids = 10
+    object_ids = [random_object_id() for _ in range(num_object_ids)]
+    metadata_sizes = [0] * (num_object_ids - 1)
+    data_sizes = [0] * (num_object_ids - 1)
+    metadata_sizes.append(np.random.randint(1000))
+    data_sizes.append(np.random.randint(1000))
+    for i in range(num_object_ids):
+      x = self.plasma_client2.create(object_ids[i], size=data_sizes[i],
+                                metadata=bytearray(np.random.bytes(metadata_sizes[i])))
+      self.plasma_client2.seal(object_ids[i])
+    del x
+    for i in range(num_object_ids):
+      recv_objid, recv_dsize, recv_msize = self.plasma_client2.get_next_notification()
+      self.assertEqual(object_ids[i], recv_objid)
+      self.assertEqual(data_sizes[i], recv_dsize)
+      self.assertEqual(metadata_sizes[i], recv_msize)
+    self.assertEqual(self.plasma_client2.evict(1), data_sizes[-1] + metadata_sizes[-1])
+    for i in range(num_object_ids):
+      recv_objid, recv_dsize, recv_msize = self.plasma_client2.get_next_notification()
+      self.assertEqual(object_ids[i], recv_objid)
+      self.assertEqual(-1, recv_dsize)
+      self.assertEqual(-1, recv_msize)
+
+
 class TestPlasmaManager(unittest.TestCase):
 
   def setUp(self):