Clean up when a driver disconnects. (#462)

* Clean up state when drivers exit.

* Remove unnecessary field in ActorMapEntry struct.

* Have monitor release GPU resources in Redis when driver exits.

* Enable multiple drivers in multi-node tests and test driver cleanup.

* Make redis GPU allocation a redis transaction and small cleanups.

* Fix multi-node test.

* Small cleanups.

* Make global scheduler take node_ip_address so it appears in the right place in the client table.

* Cleanups.

* Fix linting and cleanups in local scheduler.

* Fix removed_driver_test.

* Fix bug related to vector -> list.

* Fix linting.

* Cleanup.

* Fix multi node tests.

* Fix jenkins tests.

* Add another multi node test with many drivers.

* Fix linting.

* Make the actor creation notification a flatbuffer message.

* Revert "Make the actor creation notification a flatbuffer message."

This reverts commit af99099c8084dbf9177fb4e34c0c9b1a12c78f39.

* Add comment explaining flatbuffer problems.

src/common/CMakeLists.txt

@@ -61,6 +61,7 @@ add_library(common STATIC
   state/object_table.cc
   state/task_table.cc
   state/db_client_table.cc
+  state/driver_table.cc
   state/actor_notification_table.cc
   state/local_scheduler_table.cc
   state/error_table.cc

src/common/common.cc

@@ -47,6 +47,10 @@ bool DBClientID_equal(DBClientID first_id, DBClientID second_id) {
   return UNIQUE_ID_EQ(first_id, second_id);
 }
 
+bool WorkerID_equal(WorkerID first_id, WorkerID second_id) {
+  return UNIQUE_ID_EQ(first_id, second_id);
+}
+
 char *ObjectID_to_string(ObjectID obj_id, char *id_string, int id_length) {
   CHECK(id_length >= ID_STRING_SIZE);
   static const char hex[] = "0123456789abcdef";

src/common/common.h

@@ -153,7 +153,9 @@ extern const UniqueID NIL_ID;
 UniqueID globally_unique_id(void);
 
 #define NIL_OBJECT_ID NIL_ID
+#define NIL_WORKER_ID NIL_ID
 
+/** The object ID is the type used to identify objects. */
 typedef UniqueID ObjectID;
 
 #ifdef __cplusplus
@@ -202,6 +204,18 @@ bool ObjectID_equal(ObjectID first_id, ObjectID second_id);
  */
 bool ObjectID_is_nil(ObjectID id);
 
+/** The worker ID is the ID of a worker or driver. */
+typedef UniqueID WorkerID;
+
+/**
+ * Compare two worker IDs.
+ *
+ * @param first_id The first worker ID to compare.
+ * @param second_id The first worker ID to compare.
+ * @return True if the worker IDs are the same and false otherwise.
+ */
+bool WorkerID_equal(WorkerID first_id, WorkerID second_id);
+
 typedef UniqueID DBClientID;
 
 /**

src/common/format/common.fbs

@@ -139,3 +139,8 @@ table ResultTableReply {
 }
 
 root_type ResultTableReply;
+
+table DriverTableMessage {
+  // The driver ID of the driver that died.
+  driver_id: string;
+}

src/common/state/actor_notification_table.h

@@ -5,20 +5,16 @@
 #include "db.h"
 #include "table.h"
 
-typedef struct {
-  /** The ID of the actor. */
-  ActorID actor_id;
-  /** The ID of the local scheduler that is responsible for the actor. */
-  DBClientID local_scheduler_id;
-} ActorInfo;
-
 /*
  *  ==== Subscribing to the actor notification table ====
  */
 
 /* Callback for subscribing to the local scheduler table. */
-typedef void (*actor_notification_table_subscribe_callback)(ActorInfo info,
-                                                            void *user_context);
+typedef void (*actor_notification_table_subscribe_callback)(
+    ActorID actor_id,
+    WorkerID driver_id,
+    DBClientID local_scheduler_id,
+    void *user_context);
 
 /**
  * Register a callback to process actor notification events.

src/common/state/driver_table.cc

@@ -0,0 +1,21 @@
+#include "driver_table.h"
+#include "redis.h"
+
+void driver_table_subscribe(DBHandle *db_handle,
+                            driver_table_subscribe_callback subscribe_callback,
+                            void *subscribe_context,
+                            RetryInfo *retry) {
+  DriverTableSubscribeData *sub_data =
+      (DriverTableSubscribeData *) malloc(sizeof(DriverTableSubscribeData));
+  sub_data->subscribe_callback = subscribe_callback;
+  sub_data->subscribe_context = subscribe_context;
+  init_table_callback(db_handle, NIL_ID, __func__, sub_data, retry, NULL,
+                      redis_driver_table_subscribe, NULL);
+}
+
+void driver_table_send_driver_death(DBHandle *db_handle,
+                                    WorkerID driver_id,
+                                    RetryInfo *retry) {
+  init_table_callback(db_handle, driver_id, __func__, NULL, retry, NULL,
+                      redis_driver_table_send_driver_death, NULL);
+}

src/common/state/driver_table.h

@@ -0,0 +1,50 @@
+#ifndef DRIVER_TABLE_H
+#define DRIVER_TABLE_H
+
+#include "db.h"
+#include "table.h"
+#include "task.h"
+
+/*
+ *  ==== Subscribing to the driver table ====
+ */
+
+/* Callback for subscribing to the driver table. */
+typedef void (*driver_table_subscribe_callback)(WorkerID driver_id,
+                                                void *user_context);
+
+/**
+ * Register a callback for a driver table event.
+ *
+ * @param db_handle Database handle.
+ * @param subscribe_callback Callback that will be called when the driver event
+ *        happens.
+ * @param subscribe_context Context that will be passed into the
+ *        subscribe_callback.
+ * @param retry Information about retrying the request to the database.
+ * @return Void.
+ */
+void driver_table_subscribe(DBHandle *db_handle,
+                            driver_table_subscribe_callback subscribe_callback,
+                            void *subscribe_context,
+                            RetryInfo *retry);
+
+/* Data that is needed to register driver table subscribe callbacks with the
+ * state database. */
+typedef struct {
+  driver_table_subscribe_callback subscribe_callback;
+  void *subscribe_context;
+} DriverTableSubscribeData;
+
+/**
+ * Send driver death update to all subscribers.
+ *
+ * @param db_handle Database handle.
+ * @param driver_id The ID of the driver that died.
+ * @param retry Information about retrying the request to the database.
+ */
+void driver_table_send_driver_death(DBHandle *db_handle,
+                                    WorkerID driver_id,
+                                    RetryInfo *retry);
+
+#endif /* DRIVER_TABLE_H */

src/common/state/redis.cc

@@ -17,6 +17,7 @@ extern "C" {
 #include "db.h"
 #include "db_client_table.h"
 #include "actor_notification_table.h"
+#include "driver_table.h"
 #include "local_scheduler_table.h"
 #include "object_table.h"
 #include "task.h"
@@ -1089,6 +1090,90 @@ void redis_local_scheduler_table_send_info(TableCallbackData *callback_data) {
   }
 }
 
+void redis_driver_table_subscribe_callback(redisAsyncContext *c,
+                                           void *r,
+                                           void *privdata) {
+  REDIS_CALLBACK_HEADER(db, callback_data, r);
+
+  redisReply *reply = (redisReply *) r;
+  CHECK(reply->type == REDIS_REPLY_ARRAY);
+  CHECK(reply->elements == 3);
+  redisReply *message_type = reply->element[0];
+  LOG_DEBUG("Driver table subscribe callback, message %s", message_type->str);
+
+  if (strcmp(message_type->str, "message") == 0) {
+    /* Handle a driver heartbeat. Parse the payload and call the subscribe
+     * callback. */
+    auto message =
+        flatbuffers::GetRoot<DriverTableMessage>(reply->element[2]->str);
+    /* Extract the client ID. */
+    WorkerID driver_id = from_flatbuf(message->driver_id());
+
+    /* Call the subscribe callback. */
+    DriverTableSubscribeData *data =
+        (DriverTableSubscribeData *) callback_data->data;
+    if (data->subscribe_callback) {
+      data->subscribe_callback(driver_id, data->subscribe_context);
+    }
+  } else if (strcmp(message_type->str, "subscribe") == 0) {
+    /* The reply for the initial SUBSCRIBE command. */
+    CHECK(callback_data->done_callback == NULL);
+    /* If the initial SUBSCRIBE was successful, clean up the timer, but don't
+     * destroy the callback data. */
+    event_loop_remove_timer(db->loop, callback_data->timer_id);
+
+  } else {
+    LOG_FATAL("Unexpected reply type from driver subscribe.");
+  }
+}
+
+void redis_driver_table_subscribe(TableCallbackData *callback_data) {
+  DBHandle *db = callback_data->db_handle;
+  int status = redisAsyncCommand(
+      db->sub_context, redis_driver_table_subscribe_callback,
+      (void *) callback_data->timer_id, "SUBSCRIBE driver_deaths");
+  if ((status == REDIS_ERR) || db->sub_context->err) {
+    LOG_REDIS_DEBUG(db->sub_context, "error in redis_driver_table_subscribe");
+  }
+}
+
+void redis_driver_table_send_driver_death_callback(redisAsyncContext *c,
+                                                   void *r,
+                                                   void *privdata) {
+  REDIS_CALLBACK_HEADER(db, callback_data, r);
+
+  redisReply *reply = (redisReply *) r;
+  CHECK(reply->type == REDIS_REPLY_INTEGER);
+  LOG_DEBUG("%" PRId64 " subscribers received this publish.\n", reply->integer);
+  /* At the very least, the local scheduler that publishes this message should
+   * also receive it. */
+  CHECK(reply->integer >= 1);
+
+  CHECK(callback_data->done_callback == NULL);
+  /* Clean up the timer and callback. */
+  destroy_timer_callback(db->loop, callback_data);
+}
+
+void redis_driver_table_send_driver_death(TableCallbackData *callback_data) {
+  DBHandle *db = callback_data->db_handle;
+  WorkerID driver_id = callback_data->id;
+
+  /* Create a flatbuffer object to serialize and publish. */
+  flatbuffers::FlatBufferBuilder fbb;
+  /* Create the flatbuffers message. */
+  auto message = CreateDriverTableMessage(fbb, to_flatbuf(fbb, driver_id));
+  fbb.Finish(message);
+
+  int status = redisAsyncCommand(
+      db->context, redis_driver_table_send_driver_death_callback,
+      (void *) callback_data->timer_id, "PUBLISH driver_deaths %b",
+      fbb.GetBufferPointer(), fbb.GetSize());
+  if ((status == REDIS_ERR) || db->context->err) {
+    LOG_REDIS_DEBUG(db->context,
+                    "error in redis_driver_table_send_driver_death");
+  }
+}
+
 void redis_plasma_manager_send_heartbeat(TableCallbackData *callback_data) {
   DBHandle *db = callback_data->db_handle;
   /* NOTE(swang): We purposefully do not provide a callback, leaving the table
@@ -1124,15 +1209,20 @@ void redis_actor_notification_table_subscribe_callback(redisAsyncContext *c,
     redisReply *payload = reply->element[2];
     ActorNotificationTableSubscribeData *data =
         (ActorNotificationTableSubscribeData *) callback_data->data;
-    ActorInfo info;
-    /* The payload should be the concatenation of these two structs. */
-    CHECK(sizeof(info.actor_id) + sizeof(info.local_scheduler_id) ==
+    /* The payload should be the concatenation of three IDs. */
+    ActorID actor_id;
+    WorkerID driver_id;
+    DBClientID local_scheduler_id;
+    CHECK(sizeof(actor_id) + sizeof(driver_id) + sizeof(local_scheduler_id) ==
           payload->len);
-    memcpy(&info.actor_id, payload->str, sizeof(info.actor_id));
-    memcpy(&info.local_scheduler_id, payload->str + sizeof(info.actor_id),
-           sizeof(info.local_scheduler_id));
+    memcpy(&actor_id, payload->str, sizeof(actor_id));
+    memcpy(&driver_id, payload->str + sizeof(actor_id), sizeof(driver_id));
+    memcpy(&local_scheduler_id,
+           payload->str + sizeof(actor_id) + sizeof(driver_id),
+           sizeof(local_scheduler_id));
     if (data->subscribe_callback) {
-      data->subscribe_callback(info, data->subscribe_context);
+      data->subscribe_callback(actor_id, driver_id, local_scheduler_id,
+                               data->subscribe_context);
     }
   } else if (strcmp(message_type->str, "subscribe") == 0) {
     /* The reply for the initial SUBSCRIBE command. */

src/common/state/redis.h

@@ -253,6 +253,24 @@ void redis_local_scheduler_table_subscribe(TableCallbackData *callback_data);
  */
 void redis_local_scheduler_table_send_info(TableCallbackData *callback_data);
 
+/**
+ * Subscribe to updates from the driver table.
+ *
+ * @param callback_data Data structure containing redis connection and timeout
+ *        information.
+ * @return Void.
+ */
+void redis_driver_table_subscribe(TableCallbackData *callback_data);
+
+/**
+ * Publish an update to the driver table.
+ *
+ * @param callback_data Data structure containing redis connection and timeout
+ *        information.
+ * @return Void.
+ */
+void redis_driver_table_send_driver_death(TableCallbackData *callback_data);
+
 void redis_plasma_manager_send_heartbeat(TableCallbackData *callback_data);
 
 /**

src/common/task.h

@@ -16,7 +16,6 @@ struct TaskBuilder;
 #define NIL_TASK_ID NIL_ID
 #define NIL_ACTOR_ID NIL_ID
 #define NIL_FUNCTION_ID NIL_ID
-#define NIL_WORKER_ID NIL_ID
 
 typedef UniqueID FunctionID;