Object reconstruction in Photon (#65)

* Object reconstruction in Photon and C test cases for Photon

* Fix hanging test case on mac

* Remove unnecessary event from photon tests

* make photon_disconnect not leak file descriptors

* fix some of the memory errors

* Fix valgrind

* lint

* Address Robert's comments and add test case for object reconstruction suppression

* Remove OWNER

.travis.yml

@@ -61,6 +61,10 @@ install:
   - make test
   - cd ../..
 
+  - cd src/photon
+  - make test
+  - cd ../..
+
   - cd numbuf
   - sudo python setup.py install
   - cd ..

src/common/test/db_tests.c

@@ -141,7 +141,7 @@ TEST task_table_test(void) {
   db_handle *db = db_connect("127.0.0.1", 6379, "local_scheduler", "", -1);
   db_attach(db, loop, false);
   node_id node = globally_unique_id();
-  task_spec *spec = example_task_spec();
+  task_spec *spec = example_task_spec(1, 1);
   task_table_test_task = alloc_task(spec, TASK_STATUS_SCHEDULED, node);
   free_task_spec(spec);
   retry_info retry = {
@@ -172,7 +172,7 @@ TEST task_table_all_test(void) {
   event_loop *loop = event_loop_create();
   db_handle *db = db_connect("127.0.0.1", 6379, "local_scheduler", "", -1);
   db_attach(db, loop, false);
-  task_spec *spec = example_task_spec();
+  task_spec *spec = example_task_spec(1, 1);
   /* Schedule two tasks on different nodes. */
   task *task1 = alloc_task(spec, TASK_STATUS_SCHEDULED, globally_unique_id());
   task *task2 = alloc_task(spec, TASK_STATUS_SCHEDULED, globally_unique_id());

src/common/test/object_table_tests.c

@@ -67,7 +67,7 @@ TEST new_object_test(void) {
   new_object_failed = 0;
   new_object_succeeded = 0;
   new_object_id = globally_unique_id();
-  new_object_task = example_task();
+  new_object_task = example_task(1, 1, TASK_STATUS_WAITING);
   new_object_task_spec = task_task_spec(new_object_task);
   new_object_task_id = task_spec_id(new_object_task_spec);
   g_loop = event_loop_create();

src/common/test/task_table_tests.c

@@ -92,7 +92,7 @@ void add_success_callback(task_id task_id, void *context) {
 }
 
 TEST add_lookup_test(void) {
-  add_lookup_task = example_task();
+  add_lookup_task = example_task(1, 1, TASK_STATUS_WAITING);
   g_loop = event_loop_create();
   db_handle *db =
       db_connect("127.0.0.1", 6379, "plasma_manager", "127.0.0.1", 1234);
@@ -180,7 +180,7 @@ TEST publish_timeout_test(void) {
   db_handle *db =
       db_connect("127.0.0.1", 6379, "plasma_manager", "127.0.0.1", 1234);
   db_attach(db, g_loop, false);
-  task *task = example_task();
+  task *task = example_task(1, 1, TASK_STATUS_WAITING);
   retry_info retry = {
       .num_retries = 5, .timeout = 100, .fail_callback = publish_fail_callback,
   };
@@ -289,7 +289,7 @@ TEST publish_retry_test(void) {
   db_handle *db =
       db_connect("127.0.0.1", 6379, "plasma_manager", "127.0.0.1", 11235);
   db_attach(db, g_loop, false);
-  task *task = example_task();
+  task *task = example_task(1, 1, TASK_STATUS_WAITING);
   retry_info retry = {
       .num_retries = 5,
       .timeout = 100,
@@ -383,7 +383,7 @@ TEST publish_late_test(void) {
   db_handle *db =
       db_connect("127.0.0.1", 6379, "plasma_manager", "127.0.0.1", 11236);
   db_attach(db, g_loop, false);
-  task *task = example_task();
+  task *task = example_task(1, 1, TASK_STATUS_WAITING);
   retry_info retry = {
       .num_retries = 0,
       .timeout = 0,

src/common/test/test_common.h

@@ -1,30 +1,99 @@
 #ifndef TEST_COMMON_H
 #define TEST_COMMON_H
 
+#include "io.h"
 #include "hiredis/hiredis.h"
+#include "utstring.h"
 
 #include "task.h"
 
-static task_spec *example_task_spec(void) {
+#ifndef _WIN32
+/* This function is actually not declared in standard POSIX, so declare it. */
+extern int usleep(useconds_t usec);
+#endif
+
+const int64_t arg_value_size = 1000;
+
+static inline task_spec *example_task_spec_with_args(int64_t num_args,
+                                                     int64_t num_returns,
+                                                     object_id arg_ids[]) {
   task_id parent_task_id = globally_unique_id();
   function_id func_id = globally_unique_id();
-  task_spec *task =
-      start_construct_task_spec(parent_task_id, 0, func_id, 2, 1, 0);
-  task_args_add_ref(task, globally_unique_id());
-  task_args_add_ref(task, globally_unique_id());
+  task_spec *task = start_construct_task_spec(
+      parent_task_id, 0, func_id, num_args, num_returns, arg_value_size);
+  for (int64_t i = 0; i < num_args; ++i) {
+    object_id arg_id;
+    if (arg_ids == NULL) {
+      arg_id = globally_unique_id();
+    } else {
+      arg_id = arg_ids[i];
+    }
+    task_args_add_ref(task, arg_id);
+  }
   finish_construct_task_spec(task);
   return task;
 }
 
-static task *example_task(void) {
-  task_spec *spec = example_task_spec();
-  task *instance = alloc_task(spec, TASK_STATUS_WAITING, NIL_ID);
+static inline task_spec *example_task_spec(int64_t num_args,
+                                           int64_t num_returns) {
+  return example_task_spec_with_args(num_args, num_returns, NULL);
+}
+
+static inline task *example_task_with_args(int64_t num_args,
+                                           int64_t num_returns,
+                                           scheduling_state task_state,
+                                           object_id arg_ids[]) {
+  task_spec *spec = example_task_spec_with_args(num_args, num_returns, arg_ids);
+  task *instance = alloc_task(spec, task_state, NIL_ID);
   free_task_spec(spec);
   return instance;
 }
 
+static inline task *example_task(int64_t num_args,
+                                 int64_t num_returns,
+                                 scheduling_state task_state) {
+  task_spec *spec = example_task_spec(num_args, num_returns);
+  task *instance = alloc_task(spec, task_state, NIL_ID);
+  free_task_spec(spec);
+  return instance;
+}
+
+/* I/O helper methods to retry binding to sockets. */
+static inline UT_string *bind_ipc_sock_retry(const char *socket_name_format,
+                                             int *fd) {
+  UT_string *socket_name = NULL;
+  for (int num_retries = 0; num_retries < 5; ++num_retries) {
+    LOG_INFO("trying to find plasma socket (attempt %d)", num_retries);
+    utstring_renew(socket_name);
+    utstring_printf(socket_name, socket_name_format, rand());
+    *fd = bind_ipc_sock(utstring_body(socket_name), true);
+    if (*fd < 0) {
+      /* Sleep for 100ms. */
+      usleep(100000);
+      continue;
+    }
+    break;
+  }
+  return socket_name;
+}
+
+static inline int bind_inet_sock_retry(int *fd) {
+  int port = -1;
+  for (int num_retries = 0; num_retries < 5; ++num_retries) {
+    port = 10000 + rand() % 40000;
+    *fd = bind_inet_sock(port, true);
+    if (*fd < 0) {
+      /* Sleep for 100ms. */
+      usleep(100000);
+      continue;
+    }
+    break;
+  }
+  return port;
+}
+
 /* Flush redis. */
-static void flushall_redis() {
+static inline void flushall_redis() {
   redisContext *context = redisConnect("127.0.0.1", 6379);
   freeReplyObject(redisCommand(context, "FLUSHALL"));
   redisFree(context);

src/photon/Makefile

@@ -1,9 +1,13 @@
 CC = gcc
 CFLAGS = -g -Wall -Wextra -Werror=implicit-function-declaration -Wno-sign-compare -Wno-unused-parameter -Wno-type-limits -Wno-missing-field-initializers --std=c99 -D_XOPEN_SOURCE=500 -D_POSIX_C_SOURCE=200809L -I.. -I../common -I../common/thirdparty -I../plasma/ -fPIC
+TEST_CFLAGS = -DPHOTON_TEST=1 -I.
 BUILD = build
 
 all: $(BUILD)/photon_scheduler $(BUILD)/photon_client.a
 
+$(BUILD)/photon_tests: test/photon_tests.c photon.h photon_scheduler.h photon_scheduler.c photon_algorithm.h photon_algorithm.c photon_client.h photon_client.c common
+	$(CC) $(CFLAGS) $(TEST_CFLAGS) -o $@ test/photon_tests.c photon_scheduler.c photon_algorithm.c photon_client.c ../common/build/libcommon.a ../common/thirdparty/hiredis/libhiredis.a -I../common/thirdparty/ -I../common/ ../plasma/build/libplasma_client.a -I../plasma/
+
 $(BUILD)/photon_client.a: photon_client.o
 	ar rcs $(BUILD)/photon_client.a photon_client.o
 
@@ -18,4 +22,14 @@ clean:
 	rm -rf $(BUILD)/*
 	rm -f *.o
 
+# Set the request timeout low and logging level at FATAL for testing purposes.
+test: CFLAGS += -DRAY_TIMEOUT=50 -DRAY_COMMON_LOG_LEVEL=4
+test: $(BUILD)/photon_tests FORCE
+	../common/thirdparty/redis/src/redis-server &
+	sleep 0.5s && ./build/photon_tests && ../common/thirdparty/redis/src/redis-cli shutdown
+
+valgrind: test
+	../common/thirdparty/redis/src/redis-server &
+	sleep 0.5s && valgrind --leak-check=full --show-leak-kinds=all --error-exitcode=1 ./build/photon_tests && ../common/thirdparty/redis/src/redis-cli shutdown
+
 FORCE:

src/photon/photon.h

@@ -2,11 +2,19 @@
 #define PHOTON_H
 
 #include "common/task.h"
+#include "common/state/table.h"
 #include "common/state/db.h"
 #include "plasma_client.h"
 #include "utarray.h"
 #include "uthash.h"
 
+/* Retry values for state table operations. For now, only try each command once
+ * and give it one second to succeed. */
+/* TODO(swang): We should set retry values in a config file somewhere. */
+static const retry_info photon_retry = {.num_retries = 0,
+                                        .timeout = 1000,
+                                        .fail_callback = NULL};
+
 enum photon_message_type {
   /** Notify the local scheduler that a task has finished. */
   TASK_DONE = 64,
@@ -15,6 +23,8 @@ enum photon_message_type {
   /** This is sent from the local scheduler to a worker to tell the worker to
    *  execute a task. */
   EXECUTE_TASK,
+  /** Reconstruct a possibly lost object. */
+  RECONSTRUCT_OBJECT,
 };
 
 // clang-format off
@@ -53,6 +63,8 @@ typedef struct {
   UT_array *workers;
   /** The handle to the database. */
   db_handle *db;
+  /** Whether there is a global scheduler. */
+  bool global_scheduler_exists;
   /** The Plasma client. */
   plasma_connection *plasma_conn;
   /** State for the scheduling algorithm. */

src/photon/photon_algorithm.c

@@ -5,12 +5,10 @@
 #include "utlist.h"
 
 #include "state/task_table.h"
+#include "state/object_table.h"
 #include "photon.h"
 #include "photon_scheduler.h"
 
-/* TODO(swang): We should set retry values in a config file somewhere. */
-const retry_info photon_retry = {0, 1000, NULL};
-
 typedef struct task_queue_entry {
   /** The task that is queued. */
   task_spec *spec;
@@ -224,6 +222,7 @@ void give_task_to_global_scheduler(local_scheduler_state *state,
                                    task_spec *spec,
                                    bool from_global_scheduler) {
   /* Pass on the task to the global scheduler. */
+  DCHECK(state->global_scheduler_exists);
   DCHECK(!from_global_scheduler);
   task *task = alloc_task(spec, TASK_STATUS_WAITING, NIL_ID);
   DCHECK(state->db != NULL);
@@ -242,11 +241,21 @@ void handle_task_submitted(local_scheduler_state *state,
   if ((utarray_len(algorithm_state->available_workers) > 0) &&
       can_run(algorithm_state, spec)) {
     run_task_immediately(state, algorithm_state, spec, false);
-  } else if (state->db == NULL) {
+  } else if (state->db == NULL || !state->global_scheduler_exists) {
     queue_task_locally(state, algorithm_state, spec, false);
   } else {
     give_task_to_global_scheduler(state, algorithm_state, spec, false);
   }
+  /* Update the result table, which holds mappings of object ID -> ID of the
+   * task that created it. */
+  if (state->db != NULL) {
+    task_id task_id = task_spec_id(spec);
+    for (int64_t i = 0; i < task_num_returns(spec); ++i) {
+      object_id return_id = task_return(spec, i);
+      result_table_add(state->db, return_id, task_id,
+                       (retry_info *) &photon_retry, NULL, NULL);
+    }
+  }
 }
 
 void handle_task_scheduled(local_scheduler_state *state,
@@ -256,6 +265,7 @@ void handle_task_scheduled(local_scheduler_state *state,
    * the global scheduler, so we can safely assert that there is a connection
    * to the database. */
   DCHECK(state->db != NULL);
+  DCHECK(state->global_scheduler_exists);
   /* Initiate fetch calls for any dependencies that are not present locally. */
   fetch_missing_dependencies(state, algorithm_state, spec);
   /* If this task's dependencies are available locally, and if there is an
@@ -326,3 +336,10 @@ void handle_object_available(local_scheduler_state *state,
     free(fetch_req);
   }
 }
+
+int num_tasks_in_queue(scheduling_algorithm_state *algorithm_state) {
+  task_queue_entry *elt;
+  int count;
+  DL_COUNT(algorithm_state->task_queue, elt, count);
+  return count;
+}

src/photon/photon_algorithm.h

@@ -82,4 +82,15 @@ void handle_worker_available(local_scheduler_state *state,
                              scheduling_algorithm_state *algorithm_state,
                              int worker_index);
 
+/** The following methods are for testing purposes only. */
+#ifdef PHOTON_TEST
+/**
+ * Get the number of tasks currently queued locally.
+ *
+ * @param algorithm_state State maintained by the scheduling algorithm.
+ * @return The number of tasks queued locally.
+ */
+int num_tasks_in_queue(scheduling_algorithm_state *algorithm_state);
+#endif
+
 #endif /* PHOTON_ALGORITHM_H */

src/photon/photon_client.c

@@ -10,6 +10,11 @@ photon_conn *photon_connect(const char *photon_socket) {
   return result;
 }
 
+void photon_disconnect(photon_conn *conn) {
+  close(conn->conn);
+  free(conn);
+}
+
 void photon_submit(photon_conn *conn, task_spec *task) {
   write_message(conn->conn, SUBMIT_TASK, task_spec_size(task),
                 (uint8_t *) task);
@@ -33,8 +38,9 @@ void photon_task_done(photon_conn *conn) {
   write_message(conn->conn, TASK_DONE, 0, NULL);
 }
 
-void photon_disconnect(photon_conn *conn) {
-  write_message(conn->conn, DISCONNECT_CLIENT, 0, NULL);
+void photon_reconstruct_object(photon_conn *conn, object_id object_id) {
+  write_message(conn->conn, RECONSTRUCT_OBJECT, sizeof(object_id),
+                (uint8_t *) &object_id);
 }
 
 void photon_log_message(photon_conn *conn) {

src/photon/photon_client.h

@@ -13,11 +13,19 @@ typedef struct {
  * Connect to the local scheduler.
  *
  * @param photon_socket The name of the socket to use to connect to the local
-          scheduler.
+ *        scheduler.
  * @return The connection information.
  */
 photon_conn *photon_connect(const char *photon_socket);
 
+/**
+ * Disconnect from the local scheduler.
+ *
+ * @param conn Photon connection information returned by photon_connect.
+ * @return Void.
+ */
+void photon_disconnect(photon_conn *conn);
+
 /**
  * Submit a task to the local scheduler.
  *
@@ -48,12 +56,13 @@ task_spec *photon_get_task(photon_conn *conn);
 void photon_task_done(photon_conn *conn);
 
 /**
- * Disconnect from the local scheduler.
+ * Tell the local scheduler to reconstruct an object.
  *
  * @param conn The connection information.
+ * @param object_id The ID of the object to reconstruct.
  * @return Void.
  */
-void photon_disconnect(photon_conn *conn);
+void photon_reconstruct_object(photon_conn *conn, object_id object_id);
 
 /**
  * Send a log message to the local scheduler.

src/photon/photon_extension.c

@@ -51,11 +51,24 @@ static PyObject *PyPhotonClient_get_task(PyObject *self) {
 }
 // clang-format on
 
+static PyObject *PyPhotonClient_reconstruct_object(PyObject *self,
+                                                   PyObject *args) {
+  object_id object_id;
+  if (!PyArg_ParseTuple(args, "O&", &PyObjectToUniqueID, &object_id)) {
+    return NULL;
+  }
+  photon_reconstruct_object(((PyPhotonClient *) self)->photon_connection,
+                            object_id);
+  Py_RETURN_NONE;
+}
+
 static PyMethodDef PyPhotonClient_methods[] = {
     {"submit", (PyCFunction) PyPhotonClient_submit, METH_VARARGS,
      "Submit a task to the local scheduler."},
     {"get_task", (PyCFunction) PyPhotonClient_get_task, METH_NOARGS,
      "Get a task from the local scheduler."},
+    {"reconstruct_object", (PyCFunction) PyPhotonClient_reconstruct_object,
+     METH_VARARGS, "Ask the local scheduler to reconstruct an object."},
     {NULL} /* Sentinel */
 };
 

src/photon/photon_scheduler.c

@@ -11,8 +11,8 @@
 #include "io.h"
 #include "object_info.h"
 #include "photon.h"
-#include "photon_algorithm.h"
 #include "photon_scheduler.h"
+#include "photon_algorithm.h"
 #include "state/db.h"
 #include "state/task_table.h"
 #include "state/object_table.h"
@@ -29,7 +29,8 @@ local_scheduler_state *init_local_scheduler(
     const char *redis_addr,
     int redis_port,
     const char *plasma_store_socket_name,
-    const char *plasma_manager_socket_name) {
+    const char *plasma_manager_socket_name,
+    bool global_scheduler_exists) {
   local_scheduler_state *state = malloc(sizeof(local_scheduler_state));
   state->loop = loop;
   state->worker_index = NULL;
@@ -52,27 +53,29 @@ local_scheduler_state *init_local_scheduler(
   /* Add the callback that processes the notification to the event loop. */
   event_loop_add_file(loop, plasma_fd, EVENT_LOOP_READ,
                       process_plasma_notification, state);
+  /* Set the flag for whether there is a global scheduler. */
+  state->global_scheduler_exists = global_scheduler_exists;
   /* Add scheduler state. */
   state->algorithm_state = make_scheduling_algorithm_state();
   utarray_new(state->input_buffer, &byte_icd);
   return state;
 };
 
-void free_local_scheduler(local_scheduler_state *s) {
-  if (s->db != NULL) {
-    db_disconnect(s->db);
+void free_local_scheduler(local_scheduler_state *state) {
+  if (state->db != NULL) {
+    db_disconnect(state->db);
   }
-  plasma_disconnect(s->plasma_conn);
+  plasma_disconnect(state->plasma_conn);
   worker_index *current_worker_index, *temp_worker_index;
-  HASH_ITER(hh, s->worker_index, current_worker_index, temp_worker_index) {
-    HASH_DEL(s->worker_index, current_worker_index);
+  HASH_ITER(hh, state->worker_index, current_worker_index, temp_worker_index) {
+    HASH_DEL(state->worker_index, current_worker_index);
     free(current_worker_index);
   }
-  utarray_free(s->workers);
-  free_scheduling_algorithm_state(s->algorithm_state);
-  utarray_free(s->input_buffer);
-  event_loop_destroy(s->loop);
-  free(s);
+  utarray_free(state->workers);
+  free_scheduling_algorithm_state(state->algorithm_state);
+  utarray_free(state->input_buffer);
+  event_loop_destroy(state->loop);
+  free(state);
 }
 
 void assign_task_to_worker(local_scheduler_state *state,
@@ -103,7 +106,7 @@ void process_plasma_notification(event_loop *loop,
                                  int client_sock,
                                  void *context,
                                  int events) {
-  local_scheduler_state *s = context;
+  local_scheduler_state *state = context;
   /* Read the notification from Plasma. */
   object_info object_info;
   int error =
@@ -117,31 +120,92 @@ void process_plasma_notification(event_loop *loop,
     close(client_sock);
     return;
   }
-  handle_object_available(s, s->algorithm_state, object_info.obj_id);
+  handle_object_available(state, state->algorithm_state, object_info.obj_id);
+}
+
+void reconstruct_object_task_lookup_callback(object_id reconstruct_object_id,
+                                             task *task,
+                                             void *user_context) {
+  /* Recursively resubmit the task and its task lineage to the scheduler. */
+  CHECKM(task != NULL,
+         "No task information found for object during reconstruction");
+  local_scheduler_state *state = user_context;
+  /* If the task's scheduling state is WAITING or SCHEDULED, assume that
+   * reconstruction is already being taken care of and cancel this
+   * reconstruction operation. NOTE: This codepath is not responsible for
+   * detecting failure of the other reconstruction, or updating the
+   * scheduling_state accordingly. */
+  /* TODO(swang): Once we add code to modify the task table properly, this
+   * should also include TASK_STATUS_RUNNING. */
+  scheduling_state task_status = task_state(task);
+  if (task_status == TASK_STATUS_WAITING ||
+      task_status == TASK_STATUS_SCHEDULED) {
+    return;
+  }
+  /* Recursively reconstruct the task's inputs, if necessary. */
+  task_spec *spec = task_task_spec(task);
+  for (int64_t i = 0; i < task_num_args(spec); ++i) {
+    object_id arg_id = task_arg_id(spec, i);
+    reconstruct_object(state, arg_id);
+  }
+  handle_task_submitted(state, state->algorithm_state, spec);
+}
+
+void reconstruct_object_object_lookup_callback(
+    object_id reconstruct_object_id,
+    int manager_count,
+    const char *manager_vector[],
+    void *user_context) {
+  /* Only continue reconstruction if we find that the object doesn't exist on
+   * any nodes. NOTE: This codepath is not responsible for checking if the
+   * object table entry is up-to-date. */
+  local_scheduler_state *state = user_context;
+  if (manager_count == 0) {
+    /* Look up the task that created the object in the result table. */
+    result_table_lookup(
+        state->db, reconstruct_object_id, (retry_info *) &photon_retry,
+        reconstruct_object_task_lookup_callback, (void *) state);
+  }
+}
+
+void reconstruct_object(local_scheduler_state *state,
+                        object_id reconstruct_object_id) {
+  LOG_DEBUG("Starting reconstruction");
+  /* TODO(swang): Track task lineage for puts. */
+  CHECK(state->db != NULL);
+  /* Determine if reconstruction is necessary by checking if the object exists
+   * on a node. */
+  object_table_lookup(
+      state->db, reconstruct_object_id, (retry_info *) &photon_retry,
+      reconstruct_object_object_lookup_callback, (void *) state);
 }
 
 void process_message(event_loop *loop,
                      int client_sock,
                      void *context,
                      int events) {
-  local_scheduler_state *s = context;
+  local_scheduler_state *state = context;
 
   int64_t type;
-  read_buffer(client_sock, &type, s->input_buffer);
+  read_buffer(client_sock, &type, state->input_buffer);
 
   LOG_DEBUG("New event of type %" PRId64, type);
 
   switch (type) {
   case SUBMIT_TASK: {
-    task_spec *spec = (task_spec *) utarray_front(s->input_buffer);
-    handle_task_submitted(s, s->algorithm_state, spec);
+    task_spec *spec = (task_spec *) utarray_front(state->input_buffer);
+    handle_task_submitted(state, state->algorithm_state, spec);
   } break;
   case TASK_DONE: {
   } break;
   case GET_TASK: {
     worker_index *wi;
-    HASH_FIND_INT(s->worker_index, &client_sock, wi);
-    handle_worker_available(s, s->algorithm_state, wi->worker_index);
+    HASH_FIND_INT(state->worker_index, &client_sock, wi);
+    handle_worker_available(state, state->algorithm_state, wi->worker_index);
+  } break;
+  case RECONSTRUCT_OBJECT: {
+    object_id *obj_id = (object_id *) utarray_front(state->input_buffer);
+    reconstruct_object(state, *obj_id);
   } break;
   case DISCONNECT_CLIENT: {
     LOG_INFO("Disconnecting client on fd %d", client_sock);
@@ -159,20 +223,21 @@ void new_client_connection(event_loop *loop,
                            int listener_sock,
                            void *context,
                            int events) {
-  local_scheduler_state *s = context;
+  local_scheduler_state *state = context;
   int new_socket = accept_client(listener_sock);
-  event_loop_add_file(loop, new_socket, EVENT_LOOP_READ, process_message, s);
+  event_loop_add_file(loop, new_socket, EVENT_LOOP_READ, process_message,
+                      state);
   LOG_DEBUG("new connection with fd %d", new_socket);
   /* Add worker to list of workers. */
   /* TODO(pcm): Where shall we free this? */
   worker_index *new_worker_index = malloc(sizeof(worker_index));
   new_worker_index->sock = new_socket;
-  new_worker_index->worker_index = utarray_len(s->workers);
-  HASH_ADD_INT(s->worker_index, sock, new_worker_index);
+  new_worker_index->worker_index = utarray_len(state->workers);
+  HASH_ADD_INT(state->worker_index, sock, new_worker_index);
   worker worker;
   memset(&worker, 0, sizeof(worker));
   worker.sock = new_socket;
-  utarray_push_back(s->workers, &worker);
+  utarray_push_back(state->workers, &worker);
 }
 
 /* We need this code so we can clean up when we get a SIGTERM signal. */
@@ -197,12 +262,13 @@ void start_server(const char *socket_name,
                   const char *redis_addr,
                   int redis_port,
                   const char *plasma_store_socket_name,
-                  const char *plasma_manager_socket_name) {
+                  const char *plasma_manager_socket_name,
+                  bool global_scheduler_exists) {
   int fd = bind_ipc_sock(socket_name, true);
   event_loop *loop = event_loop_create();
-  g_state = init_local_scheduler(loop, redis_addr, redis_port,
-                                 plasma_store_socket_name,
-                                 plasma_manager_socket_name);
+  g_state = init_local_scheduler(
+      loop, redis_addr, redis_port, plasma_store_socket_name,
+      plasma_manager_socket_name, global_scheduler_exists);
 
   /* Register a callback for registering new clients. */
   event_loop_add_file(loop, fd, EVENT_LOOP_READ, new_client_connection,
@@ -225,6 +291,9 @@ void start_server(const char *socket_name,
   event_loop_run(loop);
 }
 
+/* Only declare the main function if we are not in testing mode, since the test
+ * suite has its own declaration of main. */
+#ifndef PHOTON_TEST
 int main(int argc, char *argv[]) {
   signal(SIGTERM, signal_handler);
   /* Path of the listening socket of the local scheduler. */
@@ -236,7 +305,8 @@ int main(int argc, char *argv[]) {
   /* Socket name for the local Plasma manager. */
   char *plasma_manager_socket_name = NULL;
   int c;
-  while ((c = getopt(argc, argv, "s:r:p:m:")) != -1) {
+  bool global_scheduler_exists = true;
+  while ((c = getopt(argc, argv, "s:r:p:m:g:")) != -1) {
     switch (c) {
     case 's':
       scheduler_socket_name = optarg;
@@ -250,6 +320,9 @@ int main(int argc, char *argv[]) {
     case 'm':
       plasma_manager_socket_name = optarg;
       break;
+    case 'g':
+      global_scheduler_exists = false;
+      break;
     default:
       LOG_FATAL("unknown option %c", c);
     }
@@ -270,7 +343,7 @@ int main(int argc, char *argv[]) {
           "then a redis address must be provided with the -r switch");
     }
     start_server(scheduler_socket_name, NULL, -1, plasma_store_socket_name,
-                 NULL);
+                 NULL, global_scheduler_exists);
   } else {
     /* Parse the Redis address into an IP address and a port. */
     char redis_addr[16] = {0};
@@ -288,6 +361,8 @@ int main(int argc, char *argv[]) {
           "switch");
     }
     start_server(scheduler_socket_name, &redis_addr[0], atoi(redis_port),
-                 plasma_store_socket_name, plasma_manager_socket_name);
+                 plasma_store_socket_name, plasma_manager_socket_name,
+                 global_scheduler_exists);
   }
 }
+#endif

src/photon/photon_scheduler.h

@@ -50,4 +50,36 @@ void process_plasma_notification(event_loop *loop,
                                  void *context,
                                  int events);
 
+/**
+ * Reconstruct an object. If the object does not exist on any nodes, according
+ * to the state tables, and if the object is not already being reconstructed,
+ * this triggers a single reexecution of the task that originally created the
+ * object.
+ *
+ * @param state The local scheduler state.
+ * @param object_id The ID of the object to reconstruct.
+ * @return Void.
+ */
+void reconstruct_object(local_scheduler_state *state, object_id object_id);
+
+/** The following methods are for testing purposes only. */
+#ifdef PHOTON_TEST
+local_scheduler_state *init_local_scheduler(
+    event_loop *loop,
+    const char *redis_addr,
+    int redis_port,
+    const char *plasma_manager_socket_name,
+    const char *plasma_store_socket_name,
+    bool global_scheduler_exists);
+
+void free_local_scheduler(local_scheduler_state *state);
+
+scheduling_algorithm_state *get_algorithm_state(local_scheduler_state *state);
+
+void process_message(event_loop *loop,
+                     int client_sock,
+                     void *context,
+                     int events);
+#endif
+
 #endif /* PHOTON_SCHEDULER_H */

src/photon/test/photon_tests.c

@@ -0,0 +1,262 @@
+#include "greatest.h"
+
+#include <assert.h>
+#include <unistd.h>
+#include <poll.h>
+#include <sys/types.h>
+#include <sys/socket.h>
+#include <sys/wait.h>
+
+#include "common.h"
+#include "test/test_common.h"
+#include "event_loop.h"
+#include "io.h"
+#include "utstring.h"
+#include "task.h"
+#include "state/object_table.h"
+
+#include "photon.h"
+#include "photon_scheduler.h"
+#include "photon_algorithm.h"
+#include "photon_client.h"
+
+SUITE(photon_tests);
+
+const char *plasma_socket_name_format = "/tmp/plasma_socket_%d";
+const char *photon_socket_name_format = "/tmp/photon_socket_%d";
+
+int64_t timeout_handler(event_loop *loop, int64_t id, void *context) {
+  event_loop_stop(loop);
+  return EVENT_LOOP_TIMER_DONE;
+}
+
+typedef struct {
+  /** A socket to mock the Plasma store. */
+  int plasma_fd;
+  /** Photon's socket for IPC requests. */
+  int photon_fd;
+  /** Photon's local scheduler state. */
+  local_scheduler_state *photon_state;
+  /** Photon's event loop. */
+  event_loop *loop;
+  /** A Photon client connection. */
+  photon_conn *conn;
+} photon_mock;
+
+photon_mock *init_photon_mock() {
+  const char *redis_addr = "127.0.0.1";
+  int redis_port = 6379;
+  photon_mock *mock = malloc(sizeof(photon_mock));
+  memset(mock, 0, sizeof(photon_mock));
+  mock->loop = event_loop_create();
+  /* Bind to the Photon port and initialize the Photon scheduler. */
+  UT_string *plasma_manager_socket_name =
+      bind_ipc_sock_retry(plasma_socket_name_format, &mock->plasma_fd);
+  UT_string *plasma_store_socket_name =
+      bind_ipc_sock_retry(plasma_socket_name_format, &mock->plasma_fd);
+  UT_string *photon_socket_name =
+      bind_ipc_sock_retry(photon_socket_name_format, &mock->photon_fd);
+  CHECK(mock->plasma_fd >= 0 && mock->photon_fd >= 0);
+  mock->photon_state =
+      init_local_scheduler(mock->loop, redis_addr, redis_port,
+                           utstring_body(plasma_manager_socket_name),
+                           utstring_body(plasma_store_socket_name), false);
+  /* Connect a Photon client. */
+  mock->conn = photon_connect(utstring_body(photon_socket_name));
+  new_client_connection(mock->loop, mock->photon_fd,
+                        (void *) mock->photon_state, 0);
+  utstring_free(plasma_manager_socket_name);
+  utstring_free(plasma_store_socket_name);
+  utstring_free(photon_socket_name);
+  return mock;
+}
+
+void destroy_photon_mock(photon_mock *mock) {
+  photon_disconnect(mock->conn);
+  close(mock->photon_fd);
+  close(mock->plasma_fd);
+  /* This also frees mock->loop. */
+  free_local_scheduler(mock->photon_state);
+  free(mock);
+}
+
+/**
+ * Test that object reconstruction gets called. If a task gets submitted,
+ * assigned to a worker, and then reconstruction is triggered for its return
+ * value, the task should get assigned to a worker again.
+ */
+TEST object_reconstruction_test(void) {
+  photon_mock *photon = init_photon_mock();
+  pid_t pid = fork();
+  if (pid == 0) {
+    /* Create a task with zero dependencies and one return value. */
+    task_spec *spec = example_task_spec(0, 1);
+    /* Make sure we receive the task twice. First from the initial submission,
+     * and second from the reconstruct request. */
+    photon_submit(photon->conn, spec);
+    task_spec *task_assigned = photon_get_task(photon->conn);
+    ASSERT_EQ(memcmp(task_assigned, spec, task_spec_size(spec)), 0);
+    object_id return_id = task_return(spec, 0);
+    photon_reconstruct_object(photon->conn, return_id);
+    task_spec *reconstruct_task = photon_get_task(photon->conn);
+    ASSERT_EQ(memcmp(reconstruct_task, spec, task_spec_size(spec)), 0);
+    /* Clean up. */
+    free_task_spec(reconstruct_task);
+    free_task_spec(task_assigned);
+    free_task_spec(spec);
+    destroy_photon_mock(photon);
+    exit(0);
+  } else {
+    /* Run the event loop. NOTE: OSX appears to require the parent process to
+     * listen for events on the open file descriptors. */
+    event_loop_add_timer(photon->loop, 1000,
+                         (event_loop_timer_handler) timeout_handler, NULL);
+    event_loop_run(photon->loop);
+    /* Wait for the child process to exit and check that there are no tasks
+     * left in the local scheduler's task queue. Then, clean up. */
+    wait(NULL);
+    ASSERT_EQ(num_tasks_in_queue(photon->photon_state->algorithm_state), 0);
+    destroy_photon_mock(photon);
+    PASS();
+  }
+}
+
+/**
+ * Test that object reconstruction gets recursively called. In a chain of
+ * tasks, if all inputs are lost, then reconstruction of the final object
+ * should trigger reconstruction of all previous tasks in the lineage.
+ */
+TEST object_reconstruction_recursive_test(void) {
+  photon_mock *photon = init_photon_mock();
+  /* Create a chain of tasks, each one dependent on the one before it. Mark
+   * each object as available so that tasks will run immediately. */
+  const int NUM_TASKS = 10;
+  task_spec *specs[NUM_TASKS];
+  specs[0] = example_task_spec(0, 1);
+  for (int i = 1; i < NUM_TASKS; ++i) {
+    object_id arg_id = task_return(specs[i - 1], 0);
+    handle_object_available(photon->photon_state,
+                            photon->photon_state->algorithm_state, arg_id);
+    specs[i] = example_task_spec_with_args(1, 1, &arg_id);
+  }
+  pid_t pid = fork();
+  if (pid == 0) {
+    /* Submit the tasks, and make sure each one gets assigned to a worker. */
+    for (int i = 0; i < NUM_TASKS; ++i) {
+      photon_submit(photon->conn, specs[i]);
+    }
+    /* Make sure we receive each task from the initial submission. */
+    for (int i = 0; i < NUM_TASKS; ++i) {
+      task_spec *task_assigned = photon_get_task(photon->conn);
+      ASSERT_EQ(memcmp(task_assigned, specs[i], task_spec_size(task_assigned)),
+                0);
+      free_task_spec(task_assigned);
+    }
+    /* Request reconstruction of the last return object. */
+    object_id return_id = task_return(specs[NUM_TASKS - 1], 0);
+    photon_reconstruct_object(photon->conn, return_id);
+    /* Check that the workers receive all tasks in the final return object's
+     * lineage during reconstruction. */
+    for (int i = 0; i < NUM_TASKS; ++i) {
+      task_spec *task_assigned = photon_get_task(photon->conn);
+      bool found = false;
+      for (int j = 0; j < NUM_TASKS; ++j) {
+        if (specs[j] == NULL) {
+          continue;
+        }
+        if (memcmp(task_assigned, specs[j], task_spec_size(task_assigned)) ==
+            0) {
+          found = true;
+          free_task_spec(specs[j]);
+          specs[j] = NULL;
+        }
+      }
+      free_task_spec(task_assigned);
+      ASSERT(found);
+    }
+    destroy_photon_mock(photon);
+    exit(0);
+  } else {
+    /* Run the event loop. NOTE: OSX appears to require the parent process to
+     * listen for events on the open file descriptors. */
+    event_loop_add_timer(photon->loop, 1000,
+                         (event_loop_timer_handler) timeout_handler, NULL);
+    event_loop_run(photon->loop);
+    /* Wait for the child process to exit and check that there are no tasks
+     * left in the local scheduler's task queue. Then, clean up. */
+    wait(NULL);
+    ASSERT_EQ(num_tasks_in_queue(photon->photon_state->algorithm_state), 0);
+    for (int i = 0; i < NUM_TASKS; ++i) {
+      free_task_spec(specs[i]);
+    }
+    destroy_photon_mock(photon);
+    PASS();
+  }
+}
+
+/**
+ * Test that object reconstruction gets suppressed when there is a location
+ * listed for the object in the object table.
+ */
+task_spec *object_reconstruction_suppression_spec;
+
+void object_reconstruction_suppression_callback(object_id object_id,
+                                                void *user_context) {
+  /* Submit the task after adding the object to the object table. */
+  photon_mock *photon = user_context;
+  photon_submit(photon->conn, object_reconstruction_suppression_spec);
+}
+
+TEST object_reconstruction_suppression_test(void) {
+  photon_mock *photon = init_photon_mock();
+  object_reconstruction_suppression_spec = example_task_spec(0, 1);
+  object_id return_id = task_return(object_reconstruction_suppression_spec, 0);
+  pid_t pid = fork();
+  if (pid == 0) {
+    /* Make sure we receive the task once. This will block until the
+     * object_table_add callback completes. */
+    task_spec *task_assigned = photon_get_task(photon->conn);
+    ASSERT_EQ(memcmp(task_assigned, object_reconstruction_suppression_spec,
+                     task_spec_size(object_reconstruction_suppression_spec)),
+              0);
+    /* Trigger a reconstruction. We will check that no tasks get queued as a
+     * result of this line in the event loop process. */
+    photon_reconstruct_object(photon->conn, return_id);
+    /* Clean up. */
+    free_task_spec(task_assigned);
+    free_task_spec(object_reconstruction_suppression_spec);
+    destroy_photon_mock(photon);
+    exit(0);
+  } else {
+    object_table_add(photon->photon_state->db, return_id, 1,
+                     (unsigned char *) NIL_DIGEST, (retry_info *) &photon_retry,
+                     object_reconstruction_suppression_callback,
+                     (void *) photon);
+    /* Run the event loop. NOTE: OSX appears to require the parent process to
+     * listen for events on the open file descriptors. */
+    event_loop_add_timer(photon->loop, 1000,
+                         (event_loop_timer_handler) timeout_handler, NULL);
+    event_loop_run(photon->loop);
+    /* Wait for the child process to exit and check that there are no tasks
+     * left in the local scheduler's task queue. Then, clean up. */
+    wait(NULL);
+    ASSERT_EQ(num_tasks_in_queue(photon->photon_state->algorithm_state), 0);
+    free_task_spec(object_reconstruction_suppression_spec);
+    destroy_photon_mock(photon);
+    PASS();
+  }
+}
+
+SUITE(photon_tests) {
+  RUN_REDIS_TEST(object_reconstruction_test);
+  RUN_REDIS_TEST(object_reconstruction_recursive_test);
+  RUN_REDIS_TEST(object_reconstruction_suppression_test);
+}
+
+GREATEST_MAIN_DEFS();
+
+int main(int argc, char **argv) {
+  GREATEST_MAIN_BEGIN();
+  RUN_SUITE(photon_tests);
+  GREATEST_MAIN_END();
+}

src/plasma/lib/python/plasma.py

@@ -244,7 +244,7 @@ def start_plasma_store(plasma_store_memory=DEFAULT_PLASMA_STORE_MEMORY, use_valg
   if use_valgrind and use_profiler:
     raise Exception("Cannot use valgrind and profiler at the same time.")
   plasma_store_executable = os.path.join(os.path.abspath(os.path.dirname(__file__)), "../../build/plasma_store")
-  plasma_store_name = "/tmp/scheduler{}".format(random_name())
+  plasma_store_name = "/tmp/plasma_store{}".format(random_name())
   command = [plasma_store_executable, "-s", plasma_store_name, "-m", str(plasma_store_memory)]
   if use_valgrind:
     pid = subprocess.Popen(["valgrind", "--track-origins=yes", "--leak-check=full", "--show-leak-kinds=all", "--error-exitcode=1"] + command)
@@ -274,7 +274,7 @@ def start_plasma_manager(store_name, redis_address, num_retries=20, use_valgrind
     Exception: An exception is raised if the manager could not be started.
   """
   plasma_manager_executable = os.path.join(os.path.abspath(os.path.dirname(__file__)), "../../build/plasma_manager")
-  plasma_manager_name = "/tmp/scheduler{}".format(random_name())
+  plasma_manager_name = "/tmp/plasma_manager{}".format(random_name())
   port = None
   process = None
   counter = 0

src/plasma/test/manager_tests.c

@@ -7,6 +7,7 @@
 #include <sys/socket.h>
 
 #include "common.h"
+#include "test/test_common.h"
 #include "event_loop.h"
 #include "io.h"
 #include "utstring.h"
@@ -15,13 +16,9 @@
 #include "plasma_client.h"
 #include "plasma_manager.h"
 
-#ifndef _WIN32
-/* This function is actually not declared in standard POSIX, so declare it. */
-extern int usleep(useconds_t usec);
-#endif
-
 SUITE(plasma_manager_tests);
 
+const char *plasma_socket_name_format = "/tmp/plasma_socket_%d";
 const char *manager_addr = "127.0.0.1";
 object_id oid;
 
@@ -33,38 +30,6 @@ void wait_for_pollin(int fd) {
   CHECK(retval > 0);
 }
 
-UT_string *bind_ipc_sock_retry(int *fd) {
-  UT_string *socket_name = NULL;
-  for (int num_retries = 0; num_retries < 5; ++num_retries) {
-    LOG_INFO("trying to find plasma socket (attempt %d)", num_retries);
-    utstring_renew(socket_name);
-    utstring_printf(socket_name, "/tmp/plasma_socket_%d", rand());
-    *fd = bind_ipc_sock(utstring_body(socket_name), true);
-    if (*fd < 0) {
-      /* Sleep for 100ms. */
-      usleep(100000);
-      continue;
-    }
-    break;
-  }
-  return socket_name;
-}
-
-int bind_inet_sock_retry(int *fd) {
-  int port = -1;
-  for (int num_retries = 0; num_retries < 5; ++num_retries) {
-    port = 10000 + rand() % 40000;
-    *fd = bind_inet_sock(port, true);
-    if (*fd < 0) {
-      /* Sleep for 100ms. */
-      usleep(100000);
-      continue;
-    }
-    break;
-  }
-  return port;
-}
-
 int test_done_handler(event_loop *loop, timer_id id, void *context) {
   event_loop_stop(loop);
   return AE_NOMORE;
@@ -93,8 +58,10 @@ plasma_mock *init_plasma_mock(plasma_mock *remote_mock) {
   plasma_mock *mock = malloc(sizeof(plasma_mock));
   /* Start listening on all the ports and initiate the local plasma manager. */
   mock->port = bind_inet_sock_retry(&mock->manager_remote_fd);
-  UT_string *store_socket_name = bind_ipc_sock_retry(&mock->local_store);
-  UT_string *manager_socket_name = bind_ipc_sock_retry(&mock->manager_local_fd);
+  UT_string *store_socket_name =
+      bind_ipc_sock_retry(plasma_socket_name_format, &mock->local_store);
+  UT_string *manager_socket_name =
+      bind_ipc_sock_retry(plasma_socket_name_format, &mock->manager_local_fd);
 
   CHECK(mock->manager_local_fd >= 0 && mock->local_store >= 0);