Introduce local scheduler heartbeats which carry load information. (#155)

* Introduce local scheduler heartbeats which carry load information.

src/common/Makefile

@@ -4,7 +4,7 @@ BUILD = build
 
 all: hiredis redis redismodule $(BUILD)/libcommon.a
 
-$(BUILD)/libcommon.a: event_loop.o common.o task.o io.o net.o state/redis.o state/table.o state/object_table.o state/task_table.o state/db_client_table.o thirdparty/ae/ae.o thirdparty/sha256.o
+$(BUILD)/libcommon.a: event_loop.o common.o task.o io.o net.o state/redis.o state/table.o state/object_table.o state/task_table.o state/db_client_table.o state/local_scheduler_table.o thirdparty/ae/ae.o thirdparty/sha256.o
 	ar rcs $@ $^
 
 $(BUILD)/common_tests: test/common_tests.c $(BUILD)/libcommon.a

src/common/state/local_scheduler_table.c

@@ -0,0 +1,27 @@
+#include "local_scheduler_table.h"
+#include "redis.h"
+
+void local_scheduler_table_subscribe(
+    db_handle *db_handle,
+    local_scheduler_table_subscribe_callback subscribe_callback,
+    void *subscribe_context,
+    retry_info *retry) {
+  local_scheduler_table_subscribe_data *sub_data =
+      malloc(sizeof(local_scheduler_table_subscribe_data));
+  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_local_scheduler_table_subscribe, NULL);
+}
+
+void local_scheduler_table_send_info(db_handle *db_handle,
+                                     local_scheduler_info *info,
+                                     retry_info *retry) {
+  local_scheduler_table_send_info_data *data =
+      malloc(sizeof(local_scheduler_table_send_info_data));
+  data->info = *info;
+
+  init_table_callback(db_handle, NIL_ID, __func__, data, retry, NULL,
+                      redis_local_scheduler_table_send_info, NULL);
+}

src/common/state/local_scheduler_table.h

@@ -0,0 +1,65 @@
+#ifndef LOCAL_SCHEDULER_TABLE_H
+#define LOCAL_SCHEDULER_TABLE_H
+
+#include "db.h"
+#include "table.h"
+
+typedef struct {
+  int task_queue_length;
+  int available_workers;
+} local_scheduler_info;
+
+/*
+ *  ==== Subscribing to the local scheduler table ====
+ */
+
+/* Callback for subscribing to the local scheduler table. */
+typedef void (*local_scheduler_table_subscribe_callback)(
+    db_client_id client_id,
+    local_scheduler_info info,
+    void *user_context);
+
+/**
+ * Register a callback for a local scheduler table event.
+ *
+ * @param db_handle Database handle.
+ * @param subscribe_callback Callback that will be called when the local
+ *        scheduler 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 local_scheduler_table_subscribe(
+    db_handle *db_handle,
+    local_scheduler_table_subscribe_callback subscribe_callback,
+    void *subscribe_context,
+    retry_info *retry);
+
+/* Data that is needed to register local scheduler table subscribe callbacks
+ * with the state database. */
+typedef struct {
+  local_scheduler_table_subscribe_callback subscribe_callback;
+  void *subscribe_context;
+} local_scheduler_table_subscribe_data;
+
+/**
+ * Send a heartbeat to all subscriers to the local scheduler table. This
+ * heartbeat contains some information about the load on the local scheduler.
+ *
+ * @param db_handle Database handle.
+ * @param info Information about the local scheduler, including the load on the
+ *        local scheduler.
+ * @param retry Information about retrying the request to the database.
+ */
+void local_scheduler_table_send_info(db_handle *db_handle,
+                                     local_scheduler_info *info,
+                                     retry_info *retry);
+
+/* Data that is needed to publish local scheduer heartbeats to the local
+ * scheduler table. */
+typedef struct {
+  local_scheduler_info info;
+} local_scheduler_table_send_info_data;
+
+#endif /* LOCAL_SCHEDULER_TABLE_H */

src/common/state/redis.c

@@ -12,6 +12,7 @@
 #include "common.h"
 #include "db.h"
 #include "db_client_table.h"
+#include "local_scheduler_table.h"
 #include "object_table.h"
 #include "object_info.h"
 #include "task.h"
@@ -1081,6 +1082,82 @@ void redis_db_client_table_subscribe(table_callback_data *callback_data) {
   }
 }
 
+void redis_local_scheduler_table_subscribe_callback(redisAsyncContext *c,
+                                                    void *r,
+                                                    void *privdata) {
+  REDIS_CALLBACK_HEADER(db, callback_data, r);
+
+  redisReply *reply = r;
+  CHECK(reply->type == REDIS_REPLY_ARRAY);
+  CHECK(reply->elements == 3);
+  redisReply *message_type = reply->element[0];
+  LOG_DEBUG("Local scheduer table subscribe callback, message %s",
+            message_type->str);
+
+  if (strcmp(message_type->str, "message") == 0) {
+    /* Handle a local scheduler heartbeat. Parse the payload and call the
+     * subscribe callback. */
+    redisReply *payload = reply->element[2];
+    local_scheduler_table_subscribe_data *data = callback_data->data;
+    db_client_id client_id;
+    local_scheduler_info info;
+    /* The payload should be the concatenation of these two structs. */
+    CHECK(sizeof(client_id) + sizeof(info) == payload->len);
+    memcpy(&client_id, payload->str, sizeof(client_id));
+    memcpy(&info, payload->str + sizeof(client_id), sizeof(info));
+    if (data->subscribe_callback) {
+      data->subscribe_callback(client_id, info, 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 local scheduler subscribe.");
+  }
+}
+
+void redis_local_scheduler_table_subscribe(table_callback_data *callback_data) {
+  db_handle *db = callback_data->db_handle;
+  int status = redisAsyncCommand(
+      db->sub_context, redis_local_scheduler_table_subscribe_callback,
+      (void *) callback_data->timer_id, "SUBSCRIBE local_schedulers");
+  if ((status == REDIS_ERR) || db->sub_context->err) {
+    LOG_REDIS_DEBUG(db->sub_context,
+                    "error in redis_local_scheduler_table_subscribe");
+  }
+}
+
+void redis_local_scheduler_table_send_info_callback(redisAsyncContext *c,
+                                                    void *r,
+                                                    void *privdata) {
+  REDIS_CALLBACK_HEADER(db, callback_data, r);
+
+  redisReply *reply = r;
+  CHECK(reply->type == REDIS_REPLY_INTEGER);
+  LOG_DEBUG("%" PRId64 " subscribers received this publish.\n", reply->integer);
+
+  CHECK(callback_data->done_callback == NULL);
+  /* Clean up the timer and callback. */
+  destroy_timer_callback(db->loop, callback_data);
+}
+
+void redis_local_scheduler_table_send_info(table_callback_data *callback_data) {
+  db_handle *db = callback_data->db_handle;
+  local_scheduler_table_send_info_data *data = callback_data->data;
+  int status = redisAsyncCommand(
+      db->context, redis_local_scheduler_table_send_info_callback,
+      (void *) callback_data->timer_id, "PUBLISH local_schedulers %b%b",
+      db->client.id, sizeof(db->client.id), &data->info, sizeof(data->info));
+  if ((status == REDIS_ERR) || db->context->err) {
+    LOG_REDIS_DEBUG(db->context,
+                    "error in redis_local_scheduler_table_send_info");
+  }
+}
+
 void redis_object_info_subscribe_callback(redisAsyncContext *c,
                                           void *r,
                                           void *privdata) {

src/common/state/redis.h

@@ -218,6 +218,24 @@ void redis_task_table_subscribe(table_callback_data *callback_data);
  */
 void redis_db_client_table_subscribe(table_callback_data *callback_data);
 
+/**
+ * Subscribe to updates from the local scheduler table.
+ *
+ * @param callback_data Data structure containing redis connection and timeout
+ *        information.
+ * @return Void.
+ */
+void redis_local_scheduler_table_subscribe(table_callback_data *callback_data);
+
+/**
+ * Publish an update to the local scheduler table.
+ *
+ * @param callback_data Data structure containing redis connection and timeout
+ *        information.
+ * @return Void.
+ */
+void redis_local_scheduler_table_send_info(table_callback_data *callback_data);
+
 void redis_object_info_subscribe(table_callback_data *callback_data);
 
 #endif /* REDIS_H */