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/common/object_info.h

@@ -15,6 +15,7 @@ typedef struct {
   int64_t create_time;
   int64_t construct_duration;
   unsigned char digest[DIGEST_SIZE];
+  bool is_deletion;
 } object_info;
 
 #endif

src/common/redis_module/ray_redis_module.c

@@ -191,7 +191,8 @@ int ObjectTableLookup_RedisCommand(RedisModuleCtx *ctx,
       OpenPrefixedKey(ctx, OBJECT_LOCATION_PREFIX, argv[1], REDISMODULE_READ);
 
   int keytype = RedisModule_KeyType(key);
-  if (keytype == REDISMODULE_KEYTYPE_EMPTY) {
+  if (keytype == REDISMODULE_KEYTYPE_EMPTY ||
+      RedisModule_ValueLength(key) == 0) {
     return RedisModule_ReplyWithArray(ctx, 0);
   }
 
@@ -392,6 +393,49 @@ int ObjectTableAdd_RedisCommand(RedisModuleCtx *ctx,
     RedisModule_CloseKey(object_notification_key);
   }
 
+  RedisModule_CloseKey(table_key);
+  RedisModule_ReplyWithSimpleString(ctx, "OK");
+  return REDISMODULE_OK;
+}
+
+/**
+ * Remove a manager from a location entry in the object table.
+ *
+ * This is called from a client with the command:
+ *
+ *     RAY.OBJECT_TABLE_REMOVE <object id> <manager id>
+ *
+ * @param object_id A string representing the object ID.
+ * @param manager A string which represents the manager ID of the plasma manager
+ *        to remove.
+ * @return OK if the operation was successful or an error with string
+ *         "object not found" if the entry for the object_id doesn't exist. The
+ *         operation is counted as a success if the manager was already not in
+ *         the entry.
+ */
+int ObjectTableRemove_RedisCommand(RedisModuleCtx *ctx,
+                                   RedisModuleString **argv,
+                                   int argc) {
+  if (argc != 3) {
+    return RedisModule_WrongArity(ctx);
+  }
+
+  RedisModuleString *object_id = argv[1];
+  RedisModuleString *manager = argv[2];
+
+  /* Remove the location from the object location table. */
+  RedisModuleKey *table_key;
+  table_key = OpenPrefixedKey(ctx, OBJECT_LOCATION_PREFIX, object_id,
+                              REDISMODULE_READ | REDISMODULE_WRITE);
+  int keytype = RedisModule_KeyType(table_key);
+  if (keytype == REDISMODULE_KEYTYPE_EMPTY) {
+    RedisModule_CloseKey(table_key);
+    return RedisModule_ReplyWithError(ctx, "object not found");
+  }
+
+  RedisModule_ZsetRem(table_key, manager, NULL);
+  RedisModule_CloseKey(table_key);
+
   RedisModule_ReplyWithSimpleString(ctx, "OK");
   return REDISMODULE_OK;
 }
@@ -808,6 +852,12 @@ int RedisModule_OnLoad(RedisModuleCtx *ctx,
     return REDISMODULE_ERR;
   }
 
+  if (RedisModule_CreateCommand(ctx, "ray.object_table_remove",
+                                ObjectTableRemove_RedisCommand, "write", 0, 0,
+                                0) == REDISMODULE_ERR) {
+    return REDISMODULE_ERR;
+  }
+
   if (RedisModule_CreateCommand(ctx, "ray.object_table_request_notifications",
                                 ObjectTableRequestNotifications_RedisCommand,
                                 "write pubsub", 0, 0, 0) == REDISMODULE_ERR) {

src/common/redis_module/runtest.py

@@ -101,6 +101,40 @@ class TestGlobalStateStore(unittest.TestCase):
     with self.assertRaises(redis.ResponseError):
       self.redis.execute_command("RAY.OBJECT_TABLE_ADD", "object_id3", 1, "\x00hash2", "manager_id1")
 
+  def testObjectTableAddAndRemove(self):
+    # Try removing a manager from an object ID that has not been added yet.
+    with self.assertRaises(redis.ResponseError):
+      self.redis.execute_command("RAY.OBJECT_TABLE_REMOVE", "object_id1", "manager_id1")
+    # Try calling RAY.OBJECT_TABLE_LOOKUP with an object ID that has not been
+    # added yet.
+    response = self.redis.execute_command("RAY.OBJECT_TABLE_LOOKUP", "object_id1")
+    self.assertEqual(set(response), set([]))
+    # Add some managers and try again.
+    self.redis.execute_command("RAY.OBJECT_TABLE_ADD", "object_id1", 1, "hash1", "manager_id1")
+    self.redis.execute_command("RAY.OBJECT_TABLE_ADD", "object_id1", 1, "hash1", "manager_id2")
+    response = self.redis.execute_command("RAY.OBJECT_TABLE_LOOKUP", "object_id1")
+    self.assertEqual(set(response), {b"manager_id1", b"manager_id2"})
+    # Remove a manager that doesn't exist, and make sure we still have the same set.
+    self.redis.execute_command("RAY.OBJECT_TABLE_REMOVE", "object_id1", "manager_id3")
+    response = self.redis.execute_command("RAY.OBJECT_TABLE_LOOKUP", "object_id1")
+    self.assertEqual(set(response), {b"manager_id1", b"manager_id2"})
+    # Remove a manager that does exist. Make sure it gets removed the first
+    # time and does nothing the second time.
+    self.redis.execute_command("RAY.OBJECT_TABLE_REMOVE", "object_id1", "manager_id1")
+    response = self.redis.execute_command("RAY.OBJECT_TABLE_LOOKUP", "object_id1")
+    self.assertEqual(set(response), {b"manager_id2"})
+    self.redis.execute_command("RAY.OBJECT_TABLE_REMOVE", "object_id1", "manager_id1")
+    response = self.redis.execute_command("RAY.OBJECT_TABLE_LOOKUP", "object_id1")
+    self.assertEqual(set(response), {b"manager_id2"})
+    # Remove the last manager, and make sure we have an empty set.
+    self.redis.execute_command("RAY.OBJECT_TABLE_REMOVE", "object_id1", "manager_id2")
+    response = self.redis.execute_command("RAY.OBJECT_TABLE_LOOKUP", "object_id1")
+    self.assertEqual(set(response), set())
+    # Remove a manager from an empty set, and make sure we still have an empty set.
+    self.redis.execute_command("RAY.OBJECT_TABLE_REMOVE", "object_id1", "manager_id3")
+    response = self.redis.execute_command("RAY.OBJECT_TABLE_LOOKUP", "object_id1")
+    self.assertEqual(set(response), set())
+
   def testObjectTableSubscribeToNotifications(self):
     data_size = 0xf1f0
     p = self.redis.pubsub()

src/common/state/object_table.c

@@ -28,6 +28,23 @@ void object_table_add(db_handle *db_handle,
                       done_callback, redis_object_table_add, user_context);
 }
 
+void object_table_remove(db_handle *db_handle,
+                         object_id object_id,
+                         db_client_id *client_id,
+                         retry_info *retry,
+                         object_table_done_callback done_callback,
+                         void *user_context) {
+  CHECK(db_handle != NULL);
+  /* Copy the client ID, if one was provided. */
+  db_client_id *client_id_copy = NULL;
+  if (client_id != NULL) {
+    client_id_copy = malloc(sizeof(db_client_id));
+    *client_id_copy = *client_id;
+  }
+  init_table_callback(db_handle, object_id, __func__, client_id_copy, retry,
+                      done_callback, redis_object_table_remove, user_context);
+}
+
 void object_table_subscribe_to_notifications(
     db_handle *db_handle,
     bool subscribe_all,

src/common/state/object_table.h

@@ -86,20 +86,20 @@ typedef struct {
  *
  * @param db_handle Handle to db.
  * @param object_id Object unique identifier.
+ * @param client_id A pointer to the database client ID to remove. If this is
+ *        set to NULL, then the client ID associated with db_handle will be
+ *        removed.
  * @param retry Information about retrying the request to the database.
  * @param done_callback Callback to be called when lookup completes.
  * @param user_context User context to be passed in the callbacks.
  * @return Void.
  */
-/*
-void object_table_remove(db_handle *db,
+void object_table_remove(db_handle *db_handle,
                          object_id object_id,
-                         lookup_callback callback,
-                         void *context);
+                         db_client_id *client_id,
                          retry_info *retry,
                          object_table_done_callback done_callback,
                          void *user_context);
-*/
 
 /*
  *  ==== Subscribe to be announced when new object available ====

src/common/state/redis.c

@@ -314,7 +314,7 @@ void redis_object_table_add_callback(redisAsyncContext *c,
   CHECK(strcmp(reply->str, "OK") == 0);
   /* Call the done callback if there is one. */
   if (callback_data->done_callback != NULL) {
-    task_table_done_callback done_callback = callback_data->done_callback;
+    object_table_done_callback done_callback = callback_data->done_callback;
     done_callback(callback_data->id, callback_data->user_context);
   }
   /* Clean up the timer and callback. */
@@ -340,6 +340,49 @@ void redis_object_table_add(table_callback_data *callback_data) {
   }
 }
 
+void redis_object_table_remove_callback(redisAsyncContext *c,
+                                        void *r,
+                                        void *privdata) {
+  REDIS_CALLBACK_HEADER(db, callback_data, r);
+
+  /* Do some minimal checking. */
+  redisReply *reply = r;
+  if (strcmp(reply->str, "object not found") == 0) {
+    /* If our object entry was not in the table, it's probably a race
+     * condition with an object_table_add. */
+    return;
+  }
+  CHECK(reply->type != REDIS_REPLY_ERROR);
+  CHECK(strcmp(reply->str, "OK") == 0);
+  /* Call the done callback if there is one. */
+  if (callback_data->done_callback != NULL) {
+    object_table_done_callback done_callback = callback_data->done_callback;
+    done_callback(callback_data->id, callback_data->user_context);
+  }
+  /* Clean up the timer and callback. */
+  destroy_timer_callback(db->loop, callback_data);
+}
+
+void redis_object_table_remove(table_callback_data *callback_data) {
+  db_handle *db = callback_data->db_handle;
+
+  object_id obj_id = callback_data->id;
+  /* If the caller provided a manager ID to delete, use it. Otherwise, use our
+   * own client ID as the ID to delete. */
+  db_client_id *client_id = callback_data->data;
+  if (client_id == NULL) {
+    client_id = &db->client;
+  }
+  int status = redisAsyncCommand(
+      db->context, redis_object_table_remove_callback,
+      (void *) callback_data->timer_id, "RAY.OBJECT_TABLE_REMOVE %b %b",
+      obj_id.id, sizeof(obj_id.id), client_id->id, sizeof(client_id->id));
+
+  if ((status == REDIS_ERR) || db->context->err) {
+    LOG_REDIS_DEBUG(db->context, "error in redis_object_table_remove");
+  }
+}
+
 void redis_object_table_lookup(table_callback_data *callback_data) {
   CHECK(callback_data);
   db_handle *db = callback_data->db_handle;

src/common/state/redis.h

@@ -81,6 +81,15 @@ void redis_object_table_lookup(table_callback_data *callback_data);
  */
 void redis_object_table_add(table_callback_data *callback_data);
 
+/**
+ * Remove a location entry from the object table in redis.
+ *
+ * @param callback_data Data structure containing redis connection and timeout
+ *        information.
+ * @return Void.
+ */
+void redis_object_table_remove(table_callback_data *callback_data);
+
 /**
  * Create a client-specific channel for receiving notifications from the object
  * table.

src/common/test/object_table_tests.c

@@ -425,6 +425,63 @@ TEST add_lookup_test(void) {
   PASS();
 }
 
+/* === Test add, remove, then lookup === */
+void add_remove_lookup_done_callback(object_id object_id,
+                                     int manager_count,
+                                     const char *manager_vector[],
+                                     void *context) {
+  CHECK(context == (void *) lookup_retry_context);
+  CHECK(manager_count == 0);
+  lookup_retry_succeeded = 1;
+}
+
+void add_remove_lookup_callback(object_id object_id, void *user_context) {
+  db_handle *db = user_context;
+  retry_info retry = {
+      .num_retries = 5,
+      .timeout = 100,
+      .fail_callback = lookup_retry_fail_callback,
+  };
+  object_table_lookup(db, NIL_ID, &retry, add_remove_lookup_done_callback,
+                      (void *) lookup_retry_context);
+}
+
+void add_remove_callback(object_id object_id, void *user_context) {
+  db_handle *db = user_context;
+  retry_info retry = {
+      .num_retries = 5,
+      .timeout = 100,
+      .fail_callback = lookup_retry_fail_callback,
+  };
+  object_table_remove(db, NIL_ID, NULL, &retry, add_remove_lookup_callback,
+                      (void *) db);
+}
+
+TEST add_remove_lookup_test(void) {
+  g_loop = event_loop_create();
+  lookup_retry_succeeded = 0;
+  db_handle *db =
+      db_connect("127.0.0.1", 6379, "plasma_manager", "127.0.0.1", 11235);
+  db_attach(db, g_loop, true);
+  retry_info retry = {
+      .num_retries = 5,
+      .timeout = 100,
+      .fail_callback = lookup_retry_fail_callback,
+  };
+  object_table_add(db, NIL_ID, 0, (unsigned char *) NIL_DIGEST, &retry,
+                   add_remove_callback, (void *) db);
+  /* Install handler for terminating the event loop. */
+  event_loop_add_timer(g_loop, 750,
+                       (event_loop_timer_handler) terminate_event_loop_callback,
+                       NULL);
+  event_loop_run(g_loop);
+  db_disconnect(db);
+  destroy_outstanding_callbacks(g_loop);
+  event_loop_destroy(g_loop);
+  ASSERT(lookup_retry_succeeded);
+  PASS();
+}
+
 /* === Test subscribe retry === */
 
 const char *subscribe_retry_context = "subscribe_retry";
@@ -1023,6 +1080,7 @@ SUITE(object_table_tests) {
   RUN_REDIS_TEST(lookup_retry_test);
   RUN_REDIS_TEST(add_retry_test);
   RUN_REDIS_TEST(add_lookup_test);
+  RUN_REDIS_TEST(add_remove_lookup_test);
   RUN_REDIS_TEST(subscribe_retry_test);
   RUN_REDIS_TEST(lookup_late_test);
   RUN_REDIS_TEST(add_late_test);