Availability after local scheduler failure (#329)

* Clean up plasma subscribers on EPIPE

First pass at a monitoring script - monitor can detect local scheduler death

Clean up task table upon local scheduler death in monitoring script

Don't schedule to dead local schedulers in global scheduler

Have global scheduler update the db clients table, monitor script cleans up state

Documentation

Monitor script should scan tables before beginning to read from subscription channel

Fix for python3

Redirect monitor output to redis logs, fix hanging in multinode tests

* Publish auxiliary addresses as part of db_client deletion notifications

* Fix test case?

* Small changes.

* Use SCAN instead of KEYS

* Address comments

* Address more comments

* Free redis module strings

src/common/redis_module/ray_redis_module.c

@@ -47,6 +47,47 @@ RedisModuleKey *OpenPrefixedKey(RedisModuleCtx *ctx,
   return key;
 }
 
+/**
+ * Publish a notification to a client's notification channel about an insertion
+ * or deletion to the db client table.
+ *
+ * @param ctx The Redis context.
+ * @param ray_client_id The ID of the database client that was inserted or
+ *        deleted.
+ * @param client_type The type of client that was inserted or deleted.
+ * @param aux_address An optional secondary address associated with the
+ *        database client.
+ * @param is_insertion A boolean that's true if the update was an insertion and
+ *        false if deletion.
+ * @return True if the publish was successful and false otherwise.
+ */
+bool PublishDBClientNotification(RedisModuleCtx *ctx,
+                                 RedisModuleString *ray_client_id,
+                                 RedisModuleString *client_type,
+                                 RedisModuleString *aux_address,
+                                 bool is_insertion) {
+  /* Construct strings to publish on the db client channel. */
+  RedisModuleString *channel_name =
+      RedisModule_CreateString(ctx, "db_clients", strlen("db_clients"));
+  RedisModuleString *client_info;
+  const char *is_insertion_string = is_insertion ? "1" : "0";
+  if (aux_address) {
+    client_info =
+        RedisString_Format(ctx, "%S:%S %S %s", ray_client_id, client_type,
+                           aux_address, is_insertion_string);
+  } else {
+    client_info = RedisString_Format(ctx, "%S:%S : %s", ray_client_id,
+                                     client_type, is_insertion_string);
+  }
+
+  /* Publish the client info on the db client channel. */
+  RedisModuleCallReply *reply;
+  reply = RedisModule_Call(ctx, "PUBLISH", "ss", channel_name, client_info);
+  RedisModule_FreeString(ctx, channel_name);
+  RedisModule_FreeString(ctx, client_info);
+  return (reply != NULL);
+}
+
 /**
  * Register a client with Redis. This is called from a client with the command:
  *
@@ -110,25 +151,65 @@ int Connect_RedisCommand(RedisModuleCtx *ctx,
   /* Clean up. */
   RedisModule_FreeString(ctx, aux_address_key);
   RedisModule_CloseKey(db_client_table_key);
-
-  /* Construct strings to publish on the db client channel. */
-  RedisModuleString *channel_name =
-      RedisModule_CreateString(ctx, "db_clients", strlen("db_clients"));
-  RedisModuleString *client_info;
-  if (aux_address) {
-    client_info = RedisString_Format(ctx, "%S:%S %S", ray_client_id,
-                                     client_type, aux_address);
-  } else {
-    client_info =
-        RedisString_Format(ctx, "%S:%S :", ray_client_id, client_type);
+  if (!PublishDBClientNotification(ctx, ray_client_id, client_type, aux_address,
+                                   true)) {
+    return RedisModule_ReplyWithError(ctx, "PUBLISH unsuccessful");
   }
 
-  /* Publish the client info on the db client channel. */
-  RedisModuleCallReply *reply;
-  reply = RedisModule_Call(ctx, "PUBLISH", "ss", channel_name, client_info);
-  RedisModule_FreeString(ctx, channel_name);
-  RedisModule_FreeString(ctx, client_info);
-  if (reply == NULL) {
+  RedisModule_ReplyWithSimpleString(ctx, "OK");
+  return REDISMODULE_OK;
+}
+
+/**
+ * Remove a client from Redis. This is called from a client with the command:
+ *
+ *     RAY.DISCONNECT <ray client id>
+ *
+ * This method also publishes a notification to all subscribers to the
+ * db_clients channel. The notification consists of a message of the form "<ray
+ * client id>:<client type>".
+ *
+ * @param ray_client_id The db client ID of the client.
+ * @return OK if the operation was successful.
+ */
+int Disconnect_RedisCommand(RedisModuleCtx *ctx,
+                            RedisModuleString **argv,
+                            int argc) {
+  if (argc != 2) {
+    return RedisModule_WrongArity(ctx);
+  }
+
+  RedisModuleString *ray_client_id = argv[1];
+
+  /* Get the client type. */
+  RedisModuleKey *db_client_table_key =
+      OpenPrefixedKey(ctx, DB_CLIENT_PREFIX, ray_client_id, REDISMODULE_WRITE);
+  if (RedisModule_KeyType(db_client_table_key) == REDISMODULE_KEYTYPE_EMPTY) {
+    /* Someone else already deleted this client. */
+    RedisModule_CloseKey(db_client_table_key);
+    RedisModule_ReplyWithSimpleString(ctx, "OK");
+    return REDISMODULE_OK;
+  }
+
+  RedisModuleString *client_type;
+  RedisModuleString *aux_address;
+  RedisModule_HashGet(db_client_table_key, REDISMODULE_HASH_CFIELDS,
+                      "client_type", &client_type, "aux_address", &aux_address,
+                      NULL);
+
+  /* Remove the client from the client table. */
+  CHECK_ERROR(RedisModule_DeleteKey(db_client_table_key),
+              "Unable to delete db client key.");
+  RedisModule_CloseKey(db_client_table_key);
+
+  /* Publish the deletion notification on the db client channel. */
+  bool published = PublishDBClientNotification(ctx, ray_client_id, client_type,
+                                               aux_address, false);
+
+  RedisModule_FreeString(ctx, aux_address);
+  RedisModule_FreeString(ctx, client_type);
+
+  if (!published) {
     return RedisModule_ReplyWithError(ctx, "PUBLISH unsuccessful");
   }
 
@@ -968,7 +1049,12 @@ int RedisModule_OnLoad(RedisModuleCtx *ctx,
   }
 
   if (RedisModule_CreateCommand(ctx, "ray.connect", Connect_RedisCommand,
-                                "write", 0, 0, 0) == REDISMODULE_ERR) {
+                                "write pubsub", 0, 0, 0) == REDISMODULE_ERR) {
+    return REDISMODULE_ERR;
+  }
+
+  if (RedisModule_CreateCommand(ctx, "ray.disconnect", Disconnect_RedisCommand,
+                                "write pubsub", 0, 0, 0) == REDISMODULE_ERR) {
     return REDISMODULE_ERR;
   }
 

src/common/state/db_client_table.cc

@@ -1,6 +1,16 @@
 #include "db_client_table.h"
 #include "redis.h"
 
+void db_client_table_remove(DBHandle *db_handle,
+                            DBClientID db_client_id,
+                            RetryInfo *retry,
+                            db_client_table_done_callback done_callback,
+                            void *user_context) {
+  init_table_callback(db_handle, db_client_id, __func__, NULL, retry,
+                      (table_done_callback) done_callback,
+                      redis_db_client_table_remove, user_context);
+}
+
 void db_client_table_subscribe(
     DBHandle *db_handle,
     db_client_table_subscribe_callback subscribe_callback,

src/common/state/db_client_table.h

@@ -7,6 +7,24 @@
 typedef void (*db_client_table_done_callback)(DBClientID db_client_id,
                                               void *user_context);
 
+/**
+ * Remove a client from the db clients table.
+ *
+ * @param db_handle Database handle.
+ * @param db_client_id The database client ID to remove.
+ * @param retry Information about retrying the request to the database.
+ * @param done_callback Function to be called when database returns result.
+ * @param user_context Data that will be passed to done_callback and
+ *        fail_callback.
+ * @return Void.
+ *
+ */
+void db_client_table_remove(DBHandle *db_handle,
+                            DBClientID db_client_id,
+                            RetryInfo *retry,
+                            db_client_table_done_callback done_callback,
+                            void *user_context);
+
 /*
  *  ==== Subscribing to the db client table ====
  */
@@ -15,6 +33,7 @@ typedef void (*db_client_table_done_callback)(DBClientID db_client_id,
 typedef void (*db_client_table_subscribe_callback)(DBClientID db_client_id,
                                                    const char *client_type,
                                                    const char *aux_address,
+                                                   bool is_insertion,
                                                    void *user_context);
 
 /**

src/common/state/redis.cc

@@ -992,6 +992,36 @@ void redis_task_table_subscribe(TableCallbackData *callback_data) {
  *  ==== db client table callbacks ====
  */
 
+void redis_db_client_table_remove_callback(redisAsyncContext *c,
+                                           void *r,
+                                           void *privdata) {
+  REDIS_CALLBACK_HEADER(db, callback_data, r);
+  redisReply *reply = (redisReply *) r;
+
+  CHECK(reply->type != REDIS_REPLY_ERROR);
+  CHECK(strcmp(reply->str, "OK") == 0);
+
+  /* Call the done callback if there is one. */
+  db_client_table_done_callback done_callback =
+      (db_client_table_done_callback) callback_data->done_callback;
+  if (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_db_client_table_remove(TableCallbackData *callback_data) {
+  DBHandle *db = callback_data->db_handle;
+  int status =
+      redisAsyncCommand(db->context, redis_db_client_table_remove_callback,
+                        (void *) callback_data->timer_id, "RAY.DISCONNECT %b",
+                        callback_data->id.id, sizeof(callback_data->id.id));
+  if ((status == REDIS_ERR) || db->context->err) {
+    LOG_REDIS_DEBUG(db->context, "error in db_client_table_remove");
+  }
+}
+
 void redis_db_client_table_subscribe_callback(redisAsyncContext *c,
                                               void *r,
                                               void *privdata) {
@@ -1024,19 +1054,30 @@ void redis_db_client_table_subscribe_callback(redisAsyncContext *c,
   /* We subtract 1 + sizeof(client.id) to compute the length of the
    * client_type string, and we add 1 to null-terminate the string. */
   int client_type_length = payload->len - 1 - sizeof(client.id) + 1;
+  CHECK(client_type_length > 0);
+
+  /* Parse the client type and auxiliary address from the response. If there is
+   * only client type, then the update was a delete. */
   char *client_type = (char *) malloc(client_type_length);
   char *aux_address = (char *) malloc(client_type_length);
+  int is_insertion;
   memset(aux_address, 0, client_type_length);
   /* Published message format: <client_id:client_type aux_addr> */
-  int rv = sscanf(&payload->str[1 + sizeof(client.id)], "%s %s", client_type,
-                  aux_address);
-  CHECKM(rv == 2,
+  int rv = sscanf(&payload->str[1 + sizeof(client.id)], "%s %s %d", client_type,
+                  aux_address, &is_insertion);
+  CHECKM(rv == 3,
          "redis_db_client_table_subscribe_callback: expected 2 parsed args, "
          "Got %d instead.",
          rv);
+  CHECKM(is_insertion == 1 || is_insertion == 0,
+         "redis_db_client_table_subscribe_callback: expected 0 or 1 for "
+         "insertion field, got %d instead.",
+         is_insertion);
+
+  /* Call the subscription callback. */
   if (data->subscribe_callback) {
     data->subscribe_callback(client, client_type, aux_address,
-                             data->subscribe_context);
+                             (bool) is_insertion, data->subscribe_context);
   }
   free(client_type);
   free(aux_address);

src/common/state/redis.h

@@ -217,6 +217,15 @@ void redis_task_table_publish_publish_callback(redisAsyncContext *c,
  */
 void redis_task_table_subscribe(TableCallbackData *callback_data);
 
+/**
+ * Remove a client from the db clients table.
+ *
+ * @param callback_data Data structure containing redis connection and timeout
+ *        information.
+ * @return Void.
+ */
+void redis_db_client_table_remove(TableCallbackData *callback_data);
+
 /**
  * Subscribe to updates from the db client table.
  *

src/common/state/table.cc

@@ -81,7 +81,7 @@ int64_t table_timeout_handler(event_loop *loop,
 
   CHECK(callback_data->retry.num_retries >= 0 ||
         callback_data->retry.num_retries == -1);
-  LOG_WARN("retrying operation, retry_count = %d",
+  LOG_WARN("retrying operation %s, retry_count = %d", callback_data->label,
            callback_data->retry.num_retries);
 
   if (callback_data->retry.num_retries == 0) {