Rearrange files to prepare to merge into Ray.

src/common/state/db.h

@@ -0,0 +1,32 @@
+#ifndef DB_H
+#define DB_H
+
+#include "event_loop.h"
+
+typedef struct db_handle_impl db_handle;
+
+/* Connect to the global system store at address and port. Returns
+ * a handle to the database, which must be freed with db_disconnect
+ * after use. */
+db_handle *db_connect(const char *db_address,
+                      int db_port,
+                      const char *client_type,
+                      const char *client_addr,
+                      int client_port);
+
+/* Attach global system store connection to event loop. */
+void db_attach(db_handle *db, event_loop *loop);
+
+/* Disconnect from the global system store. */
+void db_disconnect(db_handle *db);
+
+/**
+ * Returns the client ID, according to the database.
+ *
+ * @param db The handle to the database.
+ * @returns int The client ID for this connection to the database. If
+ *          this client has no connection to the database, returns -1.
+ */
+int get_client_id(db_handle *db);
+
+#endif

src/common/state/object_table.h

@@ -0,0 +1,25 @@
+#include "common.h"
+#include "db.h"
+
+/* The callback that is called when the result of a lookup
+ * in the object table comes back. The callback should free
+ * the manager_vector array, but NOT the strings they are pointing to. */
+typedef void (*lookup_callback)(object_id object_id,
+                                int manager_count,
+                                const char *manager_vector[],
+                                void *context);
+
+/* Register a new object with the directory. */
+/* TODO(pcm): Retry, print for each attempt. */
+void object_table_add(db_handle *db, object_id object_id);
+
+/* Remove object from the directory. */
+void object_table_remove(db_handle *db,
+                         object_id object_id,
+                         const char *manager);
+
+/* Look up entry from the directory */
+void object_table_lookup(db_handle *db,
+                         object_id object_id,
+                         lookup_callback callback,
+                         void *context);

src/common/state/redis.c

@@ -0,0 +1,244 @@
+/* Redis implementation of the global state store */
+
+#include <assert.h>
+
+#include <stdlib.h>
+#include "hiredis/adapters/ae.h"
+#include "utstring.h"
+
+#include "common.h"
+#include "db.h"
+#include "object_table.h"
+#include "task_log.h"
+#include "event_loop.h"
+#include "redis.h"
+#include "io.h"
+
+#define LOG_REDIS_ERR(context, M, ...)                                        \
+  fprintf(stderr, "[ERROR] (%s:%d: message: %s) " M "\n", __FILE__, __LINE__, \
+          context->errstr, ##__VA_ARGS__)
+
+#define CHECK_REDIS_CONNECT(CONTEXT_TYPE, context, M, ...) \
+  do {                                                     \
+    CONTEXT_TYPE *_context = (context);                    \
+    if (!_context) {                                       \
+      LOG_ERR("could not allocate redis context");         \
+      exit(-1);                                            \
+    }                                                      \
+    if (_context->err) {                                   \
+      LOG_REDIS_ERR(_context, M, ##__VA_ARGS__);           \
+      exit(-1);                                            \
+    }                                                      \
+  } while (0);
+
+db_handle *db_connect(const char *address,
+                      int port,
+                      const char *client_type,
+                      const char *client_addr,
+                      int client_port) {
+  db_handle *db = malloc(sizeof(db_handle));
+  /* Sync connection for initial handshake */
+  redisReply *reply;
+  long long num_clients;
+  redisContext *context = redisConnect(address, port);
+  CHECK_REDIS_CONNECT(redisContext, context, "could not connect to redis %s:%d",
+                      address, port);
+  /* Add new client using optimistic locking. */
+  while (1) {
+    reply = redisCommand(context, "WATCH %s", client_type);
+    freeReplyObject(reply);
+    reply = redisCommand(context, "HLEN %s", client_type);
+    num_clients = reply->integer;
+    freeReplyObject(reply);
+    reply = redisCommand(context, "MULTI");
+    freeReplyObject(reply);
+    reply = redisCommand(context, "HSET %s %lld %s:%d", client_type,
+                         num_clients, client_addr, client_port);
+    freeReplyObject(reply);
+    reply = redisCommand(context, "EXEC");
+    CHECK(reply);
+    if (reply->type != REDIS_REPLY_NIL) {
+      freeReplyObject(reply);
+      break;
+    }
+    freeReplyObject(reply);
+  }
+
+  db->client_type = strdup(client_type);
+  db->client_id = num_clients;
+  db->service_cache = NULL;
+  db->sync_context = context;
+  utarray_new(db->callback_freelist, &ut_ptr_icd);
+
+  /* Establish async connection */
+  db->context = redisAsyncConnect(address, port);
+  CHECK_REDIS_CONNECT(redisAsyncContext, db->context,
+                      "could not connect to redis %s:%d", address, port);
+  db->context->data = (void *) db;
+  /* Establish async connection for subscription */
+  db->sub_context = redisAsyncConnect(address, port);
+  CHECK_REDIS_CONNECT(redisAsyncContext, db->sub_context,
+                      "could not connect to redis %s:%d", address, port);
+  db->sub_context->data = (void *) db;
+
+  return db;
+}
+
+void db_disconnect(db_handle *db) {
+  redisFree(db->sync_context);
+  redisAsyncFree(db->context);
+  redisAsyncFree(db->sub_context);
+  service_cache_entry *e, *tmp;
+  HASH_ITER(hh, db->service_cache, e, tmp) {
+    free(e->addr);
+    HASH_DEL(db->service_cache, e);
+    free(e);
+  }
+  free(db->client_type);
+  void **p = NULL;
+  while ((p = (void **) utarray_next(db->callback_freelist, p))) {
+    free(*p);
+  }
+  utarray_free(db->callback_freelist);
+  free(db);
+}
+
+void db_attach(db_handle *db, event_loop *loop) {
+  redisAeAttach(loop, db->context);
+  redisAeAttach(loop, db->sub_context);
+}
+
+void object_table_add(db_handle *db, unique_id object_id) {
+  redisAsyncCommand(db->context, NULL, NULL, "SADD obj:%b %d", &object_id.id[0],
+                    UNIQUE_ID_SIZE, db->client_id);
+  if (db->context->err) {
+    LOG_REDIS_ERR(db->context, "could not add object_table entry");
+  }
+}
+
+void object_table_get_entry(redisAsyncContext *c, void *r, void *privdata) {
+  db_handle *db = c->data;
+  lookup_callback_data *cb_data = privdata;
+  redisReply *reply = r;
+  if (reply == NULL)
+    return;
+  int *result = malloc(reply->elements * sizeof(int));
+  int64_t manager_count = reply->elements;
+  if (reply->type == REDIS_REPLY_ARRAY) {
+    for (int j = 0; j < reply->elements; j++) {
+      CHECK(reply->element[j]->type == REDIS_REPLY_STRING);
+      result[j] = atoi(reply->element[j]->str);
+      service_cache_entry *entry;
+      HASH_FIND_INT(db->service_cache, &result[j], entry);
+      if (!entry) {
+        redisReply *reply = redisCommand(db->sync_context, "HGET %s %lld",
+                                         db->client_type, result[j]);
+        CHECK(reply->type == REDIS_REPLY_STRING);
+        entry = malloc(sizeof(service_cache_entry));
+        entry->service_id = result[j];
+        entry->addr = strdup(reply->str);
+        HASH_ADD_INT(db->service_cache, service_id, entry);
+        freeReplyObject(reply);
+      }
+    }
+  } else {
+    LOG_ERR("expected integer or string, received type %d", reply->type);
+    exit(-1);
+  }
+  const char **manager_vector = malloc(manager_count * sizeof(char *));
+  for (int j = 0; j < manager_count; ++j) {
+    service_cache_entry *entry;
+    HASH_FIND_INT(db->service_cache, &result[j], entry);
+    manager_vector[j] = entry->addr;
+  }
+  cb_data->callback(cb_data->object_id, manager_count, manager_vector,
+                    cb_data->context);
+  free(privdata);
+  free(result);
+}
+
+void object_table_lookup(db_handle *db,
+                         object_id object_id,
+                         lookup_callback callback,
+                         void *context) {
+  lookup_callback_data *cb_data = malloc(sizeof(lookup_callback_data));
+  cb_data->callback = callback;
+  cb_data->object_id = object_id;
+  cb_data->context = context;
+  redisAsyncCommand(db->context, object_table_get_entry, cb_data,
+                    "SMEMBERS obj:%b", &object_id.id[0], UNIQUE_ID_SIZE);
+  if (db->context->err) {
+    LOG_REDIS_ERR(db->context, "error in object_table lookup");
+  }
+}
+
+void task_log_add_task(db_handle *db, task_instance *task_instance) {
+  task_iid task_iid = *task_instance_id(task_instance);
+  redisAsyncCommand(db->context, NULL, NULL, "HMSET tasklog:%b 0 %b",
+                    (char *) &task_iid.id[0], UNIQUE_ID_SIZE,
+                    (char *) task_instance, task_instance_size(task_instance));
+  if (db->context->err) {
+    LOG_REDIS_ERR(db->context, "error setting task in task_log_add_task");
+  }
+  node_id node = *task_instance_node(task_instance);
+  int32_t state = *task_instance_state(task_instance);
+  redisAsyncCommand(db->context, NULL, NULL, "PUBLISH task_log:%b:%d %b",
+                    (char *) &node.id[0], UNIQUE_ID_SIZE, state,
+                    (char *) task_instance, task_instance_size(task_instance));
+  if (db->context->err) {
+    LOG_REDIS_ERR(db->context, "error publishing task in task_log_add_task");
+  }
+}
+
+void task_log_redis_callback(redisAsyncContext *c,
+                             void *reply,
+                             void *privdata) {
+  redisReply *r = reply;
+  if (reply == NULL)
+    return;
+  CHECK(r->type == REDIS_REPLY_ARRAY);
+  /* First entry is message type, second is topic, third is payload. */
+  CHECK(r->elements > 2);
+  /* If this condition is true, we got the initial message that acknowledged the
+   * subscription. */
+  if (r->element[2]->str == NULL) {
+    return;
+  }
+  /* Otherwise, parse the task and call the callback. */
+  CHECK(privdata);
+  task_log_callback_data *callback_data = privdata;
+  task_instance *instance = malloc(r->element[2]->len);
+  memcpy(instance, r->element[2]->str, r->element[2]->len);
+  callback_data->callback(instance, callback_data->userdata);
+  task_instance_free(instance);
+}
+void task_log_register_callback(db_handle *db,
+                                task_log_callback callback,
+                                node_id node,
+                                int32_t state,
+                                void *userdata) {
+  task_log_callback_data *callback_data =
+      malloc(sizeof(task_log_callback_data));
+  utarray_push_back(db->callback_freelist, &callback_data);
+  callback_data->callback = callback;
+  callback_data->userdata = userdata;
+  if (memcmp(&node.id[0], &NIL_ID.id[0], UNIQUE_ID_SIZE) == 0) {
+    redisAsyncCommand(db->sub_context, task_log_redis_callback, callback_data,
+                      "PSUBSCRIBE task_log:*:%d", state);
+  } else {
+    redisAsyncCommand(db->sub_context, task_log_redis_callback, callback_data,
+                      "SUBSCRIBE task_log:%b:%d", (char *) &node.id[0],
+                      UNIQUE_ID_SIZE, state);
+  }
+  if (db->sub_context->err) {
+    LOG_REDIS_ERR(db->sub_context, "error in task_log_register_callback");
+  }
+}
+
+int get_client_id(db_handle *db) {
+  if (db) {
+    return db->client_id;
+  } else {
+    return -1;
+  }
+}

src/common/state/redis.h

@@ -0,0 +1,68 @@
+#ifndef REDIS_H
+#define REDIS_H
+
+#include "db.h"
+#include "object_table.h"
+#include "task_log.h"
+
+#include "hiredis/hiredis.h"
+#include "hiredis/async.h"
+#include "uthash.h"
+#include "utarray.h"
+
+typedef struct {
+  /* Unique ID for this service. */
+  int service_id;
+  /* IP address and port of this service. */
+  char *addr;
+  /* Handle for the uthash table. */
+  UT_hash_handle hh;
+} service_cache_entry;
+
+typedef struct {
+  /* The callback that will be called. */
+  task_log_callback callback;
+  /* Userdata associated with the callback. */
+  void *userdata;
+} task_log_callback_data;
+
+struct db_handle_impl {
+  /* String that identifies this client type. */
+  char *client_type;
+  /* Unique ID for this client within the type. */
+  int64_t client_id;
+  /* Redis context for this global state store connection. */
+  redisAsyncContext *context;
+  /* Redis context for "subscribe" communication.
+   * Yes, we need a separate one for that, see
+   * https://github.com/redis/hiredis/issues/55 */
+  redisAsyncContext *sub_context;
+  /* The event loop this global state store connection is part of. */
+  event_loop *loop;
+  /* Index of the database connection in the event loop */
+  int64_t db_index;
+  /* Cache for the IP addresses of services. */
+  service_cache_entry *service_cache;
+  /* Redis context for synchronous connections.
+   * Should only be used very rarely, it is not asynchronous. */
+  redisContext *sync_context;
+  /* Data structure for callbacks that needs to be freed. */
+  UT_array *callback_freelist;
+};
+
+typedef struct {
+  /* The callback that will be called. */
+  lookup_callback callback;
+  /* Object ID that is looked up. */
+  object_id object_id;
+  /* Data context for the callback. */
+  void *context;
+} lookup_callback_data;
+
+void object_table_get_entry(redisAsyncContext *c, void *r, void *privdata);
+
+void object_table_lookup_callback(redisAsyncContext *c,
+                                  void *r,
+                                  void *privdata);
+
+#endif

src/common/state/task_log.h

@@ -0,0 +1,41 @@
+#ifndef TASK_LOG_H
+#define TASK_LOG_H
+
+#include "db.h"
+#include "task.h"
+
+/* The task log is a message bus that is used for all communication between
+ * local and global schedulers (and also persisted to the state database).
+ * Here are examples of events that are recorded by the task log:
+ *
+ * 1) local scheduler writes it when submits a task to the global scheduler;
+ * 2) global scheduler reads it to get the task submitted by local schedulers;
+ * 3) global scheduler writes it when assigning the task to a local scheduler;
+ * 4) local scheduler reads it to get its tasks assigned by global scheduler;
+ * 5) local scheduler writes it when a task finishes execution;
+ * 6) global scheduler reads it to get the tasks that have finished; */
+
+/* Callback for subscribing to the task log. */
+typedef void (*task_log_callback)(task_instance *task_instance, void *userdata);
+
+/* Initially add a task instance to the task log. */
+void task_log_add_task(db_handle *db, task_instance *task_instance);
+
+/* Update task instance in the task log. */
+void task_log_update_task(db_handle *db,
+                          task_iid task_iid,
+                          int32_t state,
+                          node_id node);
+
+/* Register callback for a certain event. The node specifies the node whose
+ * events we want to listen to. If you want to listen to all events for this
+ * node, use state_filter =
+ *     TASK_WAITING | TASK_SCHEDULED | TASK_RUNNING | TASK_DONE.
+ * If you want to register to updates from all nodes, set node = NIL_ID. */
+void task_log_register_callback(db_handle *db,
+                                task_log_callback callback,
+                                node_id node,
+                                int32_t state_filter,
+                                void *userdata);
+
+#endif /* TASK_LOG_H */

src/common/state/task_table.h

@@ -0,0 +1,20 @@
+#ifndef TASK_TABLE_H
+#define TASK_TABLE_H
+
+#include "db.h"
+#include "task.h"
+
+/* Add task to the task table, handle errors here. */
+status task_table_add_task(db_handle *db, task_spec *task);
+
+/* Callback for getting an entry from the task table. Task spec will be freed
+ * by the system after the callback */
+typedef void (*task_table_callback)(task_spec *task, void *context);
+
+/* Get specific task from the task table. */
+status task_table_get_task(db_handle *db,
+                           task_id task_id,
+                           task_table_callback callback,
+                           void *context);
+
+#endif /* TASK_TABLE_H */