Connect local scheduler to Plasma. (#11)

* Receive notifications about sealed objects from Plasma, and schedule tasks only when the dependencies are available locally.

* Fix formatting.

* Use version of Plasma with fix.

* Fix.

* Factor out the scheduling algorithm and use worker_index instead of the client socket to identify workers

* Fixes

* clang-format

* fix remaining linter errors

.clang-format

@@ -0,0 +1,5 @@
+BasedOnStyle: Chromium
+DerivePointerAlignment: false
+IndentCaseLabels: false
+PointerAlignment: Right
+SpaceAfterCStyleCast: true

.travis.yml

@@ -45,6 +45,17 @@ matrix:
 
 install:
   - ./install-dependencies.sh
+
+  # Install Plasma side by side.
+  - cd ..
+  - git clone https://github.com/ray-project/plasma.git
+  - cd plasma
+  - git checkout f189ca746b57f22371ef10077aa535492bbd8421
+  - make
+  - source setup-env.sh
+  - cd ../photon
+
+  # Install Photon.
   - make
   - cd common/lib/python
   - python setup.py install --user

Makefile

@@ -1,5 +1,5 @@
 CC = gcc
-CFLAGS = -g -Wall --std=c99 -D_XOPEN_SOURCE=500 -D_POSIX_C_SOURCE=200809L -Icommon/thirdparty -fPIC
+CFLAGS = -g -Wall --std=c99 -D_XOPEN_SOURCE=500 -D_POSIX_C_SOURCE=200809L -Icommon -Icommon/thirdparty -fPIC
 BUILD = build
 
 all: $(BUILD)/photon_scheduler $(BUILD)/photon_client.a
@@ -7,8 +7,8 @@ all: $(BUILD)/photon_scheduler $(BUILD)/photon_client.a
 $(BUILD)/photon_client.a: photon_client.o
 	ar rcs $(BUILD)/photon_client.a photon_client.o
 
-$(BUILD)/photon_scheduler: photon.h photon_scheduler.c common
-	$(CC) $(CFLAGS) -o $@ photon_scheduler.c common/build/libcommon.a common/thirdparty/hiredis/libhiredis.a -Icommon/thirdparty -Icommon/
+$(BUILD)/photon_scheduler: photon.h photon_scheduler.c photon_algorithm.c common
+	$(CC) $(CFLAGS) -o $@ photon_scheduler.c photon_algorithm.c common/build/libcommon.a common/thirdparty/hiredis/libhiredis.a -Icommon/thirdparty/ -Icommon/ ../plasma/build/libplasma_client.a -I../plasma/src/
 
 common: FORCE
 	git submodule update --init --recursive

photon.h

@@ -1,6 +1,11 @@
 #ifndef PHOTON_H
 #define PHOTON_H
 
+#include "common/task.h"
+#include "common/state/db.h"
+#include "utarray.h"
+#include "uthash.h"
+
 enum photon_message_type {
   /** Notify the local scheduler that a task has finished. */
   TASK_DONE = 64,
@@ -11,4 +16,25 @@ enum photon_message_type {
   EXECUTE_TASK,
 };
 
-#endif
+// clang-format off
+/** Contains all information that is associated to a worker. */
+typedef struct {
+  int sock;
+} worker;
+// clang-format on
+
+/* These are needed to define the UT_arrays. */
+UT_icd task_ptr_icd;
+UT_icd worker_icd;
+
+/** Resources that are exposed to the scheduling algorithm. */
+typedef struct {
+  /** List of workers available to this node. The index into this array
+   *  is the worker_index and is used to identify workers throughout
+   *  the program. */
+  UT_array *workers;
+  /* The handle to the database. */
+  db_handle *db;
+} scheduler_info;
+
+#endif /* PHOTON_H */

photon_algorithm.c

@@ -0,0 +1,184 @@
+#include "photon_algorithm.h"
+
+#include <stdbool.h>
+#include "utarray.h"
+
+#include "state/task_log.h"
+#include "photon.h"
+#include "photon_scheduler.h"
+
+typedef struct {
+  /* Object id of this object. */
+  object_id object_id;
+  /* Handle for the uthash table. */
+  UT_hash_handle handle;
+} available_object;
+
+/** Part of the photon state that is maintained by the scheduling algorithm. */
+struct scheduler_state {
+  /** An array of pointers to tasks that are waiting to be scheduled. */
+  UT_array *task_queue;
+  /** An array of worker indices corresponding to clients that are
+   *  waiting for tasks. */
+  UT_array *available_workers;
+  /** A hash map of the objects that are available in the local Plasma store.
+   *  This information could be a little stale. */
+  available_object *local_objects;
+};
+
+scheduler_state *make_scheduler_state(void) {
+  scheduler_state *state = malloc(sizeof(scheduler_state));
+  /* Initialize an empty hash map for the cache of local available objects. */
+  state->local_objects = NULL;
+  /* Initialize the local data structures used for queuing tasks and workers. */
+  utarray_new(state->task_queue, &task_ptr_icd);
+  utarray_new(state->available_workers, &ut_int_icd);
+  return state;
+}
+
+void free_scheduler_state(scheduler_state *s) {
+  utarray_free(s->task_queue);
+  utarray_free(s->available_workers);
+  free(s);
+}
+
+/**
+ * Check if all of the remote object arguments for a task are available in the
+ * local object store.
+ *
+ * @param s The scheduler state.
+ * @param task Task specification of the task to check.
+ * @return This returns 1 if all of the remote object arguments for the task are
+ *         present in the local object store, otherwise it returns 0.
+ */
+bool can_run(scheduler_state *s, task_spec *task) {
+  int64_t num_args = task_num_args(task);
+  for (int i = 0; i < num_args; ++i) {
+    if (task_arg_type(task, i) == ARG_BY_REF) {
+      object_id obj_id = *task_arg_id(task, i);
+      available_object *entry;
+      HASH_FIND(handle, s->local_objects, &obj_id, sizeof(object_id), entry);
+      if (entry == NULL) {
+        /* The object is not present locally, so this task cannot be scheduled
+         * right now. */
+        return false;
+      }
+    }
+  }
+  return true;
+}
+
+/**
+ * If there is a task whose dependencies are available locally, assign it to the
+ * worker. This does not remove the worker from the available worker queue.
+ *
+ * @param s The scheduler state.
+ * @param worker_index The index of the worker.
+ * @return This returns 1 if it successfully assigned a task to the worker,
+ *         otherwise it returns 0.
+ */
+int find_and_schedule_task_if_possible(scheduler_info *info,
+                                       scheduler_state *state,
+                                       int worker_index) {
+  int found_task_to_schedule = 0;
+  /* Find the first task whose dependencies are available locally. */
+  task_spec *spec;
+  task_instance **task;
+  int i = 0;
+  for (; i < utarray_len(state->task_queue); ++i) {
+    task = (task_instance **) utarray_eltptr(state->task_queue, i);
+    spec = task_instance_task_spec(*task);
+    if (can_run(state, spec)) {
+      found_task_to_schedule = 1;
+      break;
+    }
+  }
+  if (found_task_to_schedule) {
+    /* This task's dependencies are available locally, so assign the task to the
+     * worker. */
+    assign_task_to_worker(info, spec, worker_index);
+    /* Update the task queue data structure and free the task. */
+    free(*task);
+    utarray_erase(state->task_queue, i, 1);
+  }
+  return found_task_to_schedule;
+}
+
+void handle_task_submitted(scheduler_info *info,
+                           scheduler_state *s,
+                           task_spec *task) {
+  /* Create a unique task instance ID. This is different from the task ID and
+   * is used to distinguish between potentially multiple executions of the
+   * task. */
+  task_iid task_iid = globally_unique_id();
+  task_instance *instance =
+      make_task_instance(task_iid, task, TASK_STATUS_WAITING, NIL_ID);
+  /* If this task's dependencies are available locally, and if there is an
+   * available worker, then assign this task to an available worker. Otherwise,
+   * add this task to the local task queue. */
+  int schedule_locally =
+      (utarray_len(s->available_workers) > 0) && can_run(s, task);
+  if (schedule_locally) {
+    /* Get the last available worker in the available worker queue. */
+    int *worker_index = (int *) utarray_back(s->available_workers);
+    /* Tell the available worker to execute the task. */
+    assign_task_to_worker(info, task, *worker_index);
+    /* Remove the available worker from the queue and free the struct. */
+    utarray_pop_back(s->available_workers);
+  } else {
+    /* Add the task to the task queue. This passes ownership of the task queue.
+     * And the task will be freed when it is assigned to a worker. */
+    utarray_push_back(s->task_queue, &instance);
+  }
+  /* Submit the task to redis. */
+  task_log_add_task(info->db, instance);
+  if (schedule_locally) {
+    /* If the task was scheduled locally, we need to free it. Otherwise,
+     * ownership of the task is passed to the task_queue, and it will be freed
+     * when it is assigned to a worker. */
+    free(instance);
+  }
+}
+
+void handle_worker_available(scheduler_info *info,
+                             scheduler_state *state,
+                             int worker_index) {
+  int scheduled_task =
+      find_and_schedule_task_if_possible(info, state, worker_index);
+  /* If we couldn't find a task to schedule, add the worker to the queue of
+   * available workers. */
+  if (!scheduled_task) {
+    for (int *p = (int *) utarray_front(state->available_workers); p != NULL;
+         p = (int *) utarray_next(state->available_workers, p)) {
+      CHECK(*p != worker_index);
+    }
+    /* Add client_sock to a list of available workers. This struct will be freed
+     * when a task is assigned to this worker. */
+    utarray_push_back(state->available_workers, &worker_index);
+    LOG_INFO("Adding worker_index %d to available workers.\n", worker_index);
+  }
+}
+
+void handle_object_available(scheduler_info *info,
+                             scheduler_state *state,
+                             object_id object_id) {
+  /* TODO(rkn): When does this get freed? */
+  available_object *entry =
+      (available_object *) malloc(sizeof(available_object));
+  entry->object_id = object_id;
+  HASH_ADD(handle, state->local_objects, object_id, sizeof(object_id), entry);
+
+  /* Check if we can schedule any tasks. */
+  int num_tasks_scheduled = 0;
+  for (int *p = (int *) utarray_front(state->available_workers); p != NULL;
+       p = (int *) utarray_next(state->available_workers, p)) {
+    /* Schedule a task on this worker if possible. */
+    int scheduled_task = find_and_schedule_task_if_possible(info, state, *p);
+    if (!scheduled_task) {
+      /* There are no tasks we can schedule, so exit the loop. */
+      break;
+    }
+    num_tasks_scheduled += 1;
+  }
+  utarray_erase(state->available_workers, 0, num_tasks_scheduled);
+}

photon_algorithm.h

@@ -0,0 +1,88 @@
+#ifndef PHOTON_ALGORITHM_H
+#define PHOTON_ALGORITHM_H
+
+#include "photon.h"
+#include "common/task.h"
+
+/* ==== The scheduling algorithm ====
+ *
+ * This file contains declaration for all functions and data structures
+ * that need to be provided if you want to implement a new algorithms
+ * for the local scheduler.
+ *
+ */
+
+/** Internal state of the scheduling algorithm. */
+typedef struct scheduler_state scheduler_state;
+
+/**
+ * Initialize the scheduler state.
+ *
+ * @return Internal state of the scheduling algorithm.
+ */
+scheduler_state *make_scheduler_state(void);
+
+/**
+ * Free the scheduler state.
+ *
+ * @param state Internal state of the scheduling algorithm.
+ * @return Void.
+ */
+void free_scheduler_state(scheduler_state *state);
+
+/**
+ * This function will be called when a new task is submitted by a worker for
+ * execution.
+ *
+ * @param info Info about resources exposed by photon to the scheduling
+ *        algorithm.
+ * @param state State of the scheduling algorithm.
+ * @param task Task that is submitted by the worker.
+ * @return Void.
+ */
+void handle_task_submitted(scheduler_info *info,
+                           scheduler_state *state,
+                           task_spec *task);
+
+/**
+ * This function will be called when a task is assigned by the global scheduler
+ * for execution on this local scheduler.
+ *
+ * @param info Info about resources exposed by photon to the scheduling
+ *        algorithm.
+ * @param state State of the scheduling algorithm.
+ * @param task Task that is assigned by the global scheduler.
+ * @return Void.
+ */
+void handle_task_assigned(scheduler_info *info,
+                          scheduler_state *state,
+                          task_spec *task);
+
+/**
+ * This function is called if a new object becomes available in the local
+ * plasma store.
+ *
+ * @param info Info about resources exposed by photon to the scheduling
+ *        algorithm.
+ * @param state State of the scheduling algorithm.
+ * @param object_id ID of the object that became available.
+ * @return Void.
+ */
+void handle_object_available(scheduler_info *info,
+                             scheduler_state *state,
+                             object_id object_id);
+
+/**
+ * This function is called when a new worker becomes available
+ *
+ * @param info Info about resources exposed by photon to the scheduling
+ *        algorithm.
+ * @param state State of the scheduling algorithm.
+ * @param worker_index The index of the worker that becomes available.
+ * @return Void.
+ */
+void handle_worker_available(scheduler_info *info,
+                             scheduler_state *state,
+                             int worker_index);
+
+#endif /* PHOTON_ALGORITHM_H */

photon_scheduler.c

@@ -10,111 +10,92 @@
 #include "event_loop.h"
 #include "io.h"
 #include "photon.h"
+#include "photon_algorithm.h"
 #include "photon_scheduler.h"
+#include "plasma_client.h"
 #include "state/db.h"
 #include "state/task_log.h"
 #include "utarray.h"
+#include "uthash.h"
 
-typedef struct {
-  /** The file descriptor used to communicate with the worker. */
-  int client_sock;
-} available_worker;
-
-/* These are needed to define the UT_arrays. */
 UT_icd task_ptr_icd = {sizeof(task_instance *), NULL, NULL, NULL};
-UT_icd worker_icd = {sizeof(available_worker), NULL, NULL, NULL};
+UT_icd worker_icd = {sizeof(worker), NULL, NULL, NULL};
+
+/** Association between the socket fd of a worker and its worker_index. */
+typedef struct {
+  /** The socket fd of a worker. */
+  int sock;
+  /** The index of the worker in scheduler_info->workers. */
+  int64_t worker_index;
+  /** Handle for the hash table. */
+  UT_hash_handle hh;
+} worker_index;
 
 struct local_scheduler_state {
   /* The local scheduler event loop. */
   event_loop *loop;
-  /* The handle to the database. */
-  db_handle *db;
-  /** This is an array of pointers to tasks that are waiting to be scheduled. */
-  UT_array *task_queue;
-  /** This is an array of file descriptors corresponding to clients that are
-   *  waiting for tasks. */
-  UT_array *available_worker_queue;
+  /* The Plasma client. */
+  plasma_store_conn *plasma_conn;
+  /* Association between client socket and worker index. */
+  worker_index *worker_index;
+  /* Info that is exposed to the scheduling algorithm. */
+  scheduler_info *scheduler_info;
+  /* State for the scheduling algorithm. */
+  scheduler_state *scheduler_state;
 };
 
-local_scheduler_state *
-init_local_scheduler(event_loop *loop, const char *redis_addr, int redis_port) {
+local_scheduler_state *init_local_scheduler(event_loop *loop,
+                                            const char *redis_addr,
+                                            int redis_port,
+                                            const char *plasma_socket_name) {
   local_scheduler_state *state = malloc(sizeof(local_scheduler_state));
   state->loop = loop;
-  state->db = db_connect(redis_addr, redis_port, "photon", "", -1);
-  db_attach(state->db, loop);
-  utarray_new(state->task_queue, &task_ptr_icd);
-  utarray_new(state->available_worker_queue, &worker_icd);
+  /* Connect to Plasma. This method will retry if Plasma hasn't started yet. */
+  state->plasma_conn = plasma_store_connect(plasma_socket_name);
+  /* Subscribe to notifications about sealed objects. */
+  int plasma_fd = plasma_subscribe(state->plasma_conn);
+  /* 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);
+  state->worker_index = NULL;
+  /* Add scheduler info. */
+  state->scheduler_info = malloc(sizeof(scheduler_info));
+  utarray_new(state->scheduler_info->workers, &worker_icd);
+  /* Connect to Redis. */
+  state->scheduler_info->db =
+      db_connect(redis_addr, redis_port, "photon", "", -1);
+  db_attach(state->scheduler_info->db, loop);
+  /* Add scheduler state. */
+  state->scheduler_state = make_scheduler_state();
   return state;
 };
 
 void free_local_scheduler(local_scheduler_state *s) {
-  db_disconnect(s->db);
-  utarray_free(s->task_queue);
-  utarray_free(s->available_worker_queue);
+  db_disconnect(s->scheduler_info->db);
+  free(s->scheduler_info);
+  free_scheduler_state(s->scheduler_state);
   event_loop_destroy(s->loop);
   free(s);
 }
 
-void handle_submit_task(local_scheduler_state *s, task_spec *task) {
-  /* Create a unique task instance ID. This is different from the task ID and
-   * is used to distinguish between potentially multiple executions of the
-   * task. */
-  task_iid task_iid = globally_unique_id();
-  task_instance *instance =
-      make_task_instance(task_iid, task, TASK_STATUS_WAITING, NIL_ID);
-  /* Assign this task to an available worker. If there are no available workers,
-   * then add this task to the local task queue. */
-  int schedule_locally = utarray_len(s->available_worker_queue) > 0;
-  if (schedule_locally) {
-    /* Get the last available worker in the available worker queue. */
-    available_worker *worker =
-        (available_worker *)utarray_back(s->available_worker_queue);
-    /* Tell the available worker to execute the task. */
-    write_message(worker->client_sock, EXECUTE_TASK, task_size(task),
-                  (uint8_t *)task);
-    /* Remove the available worker from the queue and free the struct. */
-    utarray_pop_back(s->available_worker_queue);
-  } else {
-    /* Add the task to the task queue. This passes ownership of the task queue.
-     * And the task will be freed when it is assigned to a worker. */
-    utarray_push_back(s->task_queue, &instance);
-  }
-  /* Submit the task to redis. */
-  task_log_add_task(s->db, instance);
-  if (schedule_locally) {
-    /* If the task was scheduled locally, we need to free it. Otherwise,
-     * ownership of the task is passed to the task_queue, and it will be freed
-     * when it is assigned to a worker. */
-    free(instance);
-  }
+void assign_task_to_worker(scheduler_info *info,
+                           task_spec *task,
+                           int worker_index) {
+  CHECK(worker_index < utarray_len(info->workers));
+  worker *w = (worker *) utarray_eltptr(info->workers, worker_index);
+  write_message(w->sock, EXECUTE_TASK, task_size(task), (uint8_t *) task);
 }
 
-void handle_get_task(local_scheduler_state *s, int client_sock) {
-  /* If there is an available task, assign that task to this worker. Otherwise
-   * add the worker to the queue of available workers. */
-  if (utarray_len(s->task_queue) > 0) {
-    /* Get the last task in the task queue. */
-    task_instance **back = (task_instance **)utarray_back(s->task_queue);
-    task_spec *task = task_instance_task_spec(*back);
-    /* Send a task to the worker. */
-    write_message(client_sock, EXECUTE_TASK, task_size(task), (uint8_t *)task);
-    /* Update the task queue data structure and free the task. */
-    utarray_pop_back(s->task_queue);
-    free(*back);
-  } else {
-    /* Check that client_sock is not already in the available workers. */
-    for (available_worker *p =
-             (available_worker *)utarray_front(s->available_worker_queue);
-         p != NULL;
-         p = (available_worker *)utarray_next(s->available_worker_queue, p)) {
-      CHECK(p->client_sock != client_sock);
-    }
-    /* Add client_sock to a list of available workers. This struct will be freed
-     * when a task is assigned to this worker. */
-    available_worker worker_info = {.client_sock = client_sock};
-    utarray_push_back(s->available_worker_queue, &worker_info);
-    LOG_INFO("Adding client_sock %d to available workers.\n", client_sock);
-  }
+void process_plasma_notification(event_loop *loop,
+                                 int client_sock,
+                                 void *context,
+                                 int events) {
+  local_scheduler_state *s = context;
+  /* Read the notification from Plasma. */
+  uint8_t *message = (uint8_t *) malloc(sizeof(object_id));
+  recv(client_sock, message, sizeof(object_id), 0);
+  object_id *obj_id = (object_id *) message;
+  handle_object_available(s->scheduler_info, s->scheduler_state, *obj_id);
 }
 
 void process_message(event_loop *loop, int client_sock, void *context,
@@ -126,16 +107,22 @@ void process_message(event_loop *loop, int client_sock, void *context,
   int64_t length;
   read_message(client_sock, &type, &length, &message);
 
+  LOG_DEBUG("New event of type %" PRId64, type);
+
   switch (type) {
   case SUBMIT_TASK: {
-    task_spec *task = (task_spec *)message;
-    CHECK(task_size(task) == length);
-    handle_submit_task(s, task);
+    task_spec *spec = (task_spec *) message;
+    CHECK(task_size(spec) == length);
+    handle_task_submitted(s->scheduler_info, s->scheduler_state, spec);
   } break;
   case TASK_DONE: {
   } break;
   case GET_TASK: {
-    handle_get_task(s, client_sock);
+    worker_index *wi;
+    HASH_FIND_INT(s->worker_index, &client_sock, wi);
+    printf("worker_index is %" PRId64 "\n", wi->worker_index);
+    handle_worker_available(s->scheduler_info, s->scheduler_state,
+                            wi->worker_index);
   } break;
   case DISCONNECT_CLIENT: {
     LOG_INFO("Disconnecting client on fd %d", client_sock);
@@ -156,6 +143,14 @@ void new_client_connection(event_loop *loop, int listener_sock, void *context,
   int new_socket = accept_client(listener_sock);
   event_loop_add_file(loop, new_socket, EVENT_LOOP_READ, process_message, s);
   LOG_INFO("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->scheduler_info->workers);
+  HASH_ADD_INT(s->worker_index, sock, new_worker_index);
+  worker worker = {.sock = new_socket};
+  utarray_push_back(s->scheduler_info->workers, &worker);
 }
 
 /* We need this code so we can clean up when we get a SIGTERM signal. */
@@ -171,11 +166,14 @@ void signal_handler(int signal) {
 
 /* End of the cleanup code. */
 
-void start_server(const char *socket_name, const char *redis_addr,
-                  int redis_port) {
+void start_server(const char *socket_name,
+                  const char *redis_addr,
+                  int redis_port,
+                  const char *plasma_socket_name) {
   int fd = bind_ipc_sock(socket_name);
   event_loop *loop = event_loop_create();
-  g_state = init_local_scheduler(loop, redis_addr, redis_port);
+  g_state =
+      init_local_scheduler(loop, redis_addr, redis_port, plasma_socket_name);
 
   /* Run event loop. */
   event_loop_add_file(loop, fd, EVENT_LOOP_READ, new_client_connection,
@@ -189,8 +187,10 @@ int main(int argc, char *argv[]) {
   char *scheduler_socket_name = NULL;
   /* IP address and port of redis. */
   char *redis_addr_port = NULL;
+  /* Socket name for the local Plasma store. */
+  char *plasma_socket_name = NULL;
   int c;
-  while ((c = getopt(argc, argv, "s:r:")) != -1) {
+  while ((c = getopt(argc, argv, "s:r:p:")) != -1) {
     switch (c) {
     case 's':
       scheduler_socket_name = optarg;
@@ -198,6 +198,9 @@ int main(int argc, char *argv[]) {
     case 'r':
       redis_addr_port = optarg;
       break;
+    case 'p':
+      plasma_socket_name = optarg;
+      break;
     default:
       LOG_ERR("unknown option %c", c);
       exit(-1);
@@ -207,6 +210,11 @@ int main(int argc, char *argv[]) {
     LOG_ERR("please specify socket for incoming connections with -s switch");
     exit(-1);
   }
+  if (!plasma_socket_name) {
+    LOG_ERR("please specify socket for connecting to Plasma with -p switch");
+    exit(-1);
+  }
+  /* Parse the Redis address into an IP address and a port. */
   char redis_addr[16] = {0};
   char redis_port[6] = {0};
   if (!redis_addr_port ||
@@ -215,5 +223,6 @@ int main(int argc, char *argv[]) {
     LOG_ERR("need to specify redis address like 127.0.0.1:6379 with -r switch");
     exit(-1);
   }
-  start_server(scheduler_socket_name, &redis_addr[0], atoi(redis_port));
+  start_server(scheduler_socket_name, &redis_addr[0], atoi(redis_port),
+               plasma_socket_name);
 }

photon_scheduler.h

@@ -2,6 +2,7 @@
 #define PHOTON_SCHEDULER_H
 
 #include "task.h"
+#include "event_loop.h"
 
 typedef struct local_scheduler_state local_scheduler_state;
 
@@ -15,25 +16,37 @@ typedef struct local_scheduler_state local_scheduler_state;
  * @param events Flag for events that are available on the listener socket.
  * @return Void.
  */
-void new_client_connection(event_loop *loop, int listener_sock, void *context,
+void new_client_connection(event_loop *loop,
+                           int listener_sock,
+                           void *context,
                            int events);
 
 /**
- * Assign a task to a worker.
+ * This function can be called by the scheduling algorithm to assign a task
+ * to a worker.
  *
- * @param s State of the local scheduler.
- * @param client_sock Socket by which the worker is connected.
+ * @param info
+ * @param task The task that is submitted to the worker.
+ * @param worker_index The index of the worker the task is submitted to.
  * @return Void.
  */
-void handle_get_task(local_scheduler_state *s, int client_sock);
+void assign_task_to_worker(scheduler_info *info,
+                           task_spec *task,
+                           int worker_index);
 
 /**
- * Handle incoming submit request by a worker.
+ * This is the callback that is used to process a notification from the Plasma
+ * store that an object has been sealed.
  *
- * @param s State of the local scheduler.
- * @param task Task specification of the task to be submitted.
+ * @param loop The local scheduler's event loop.
+ * @param client_sock The file descriptor to read the notification from.
+ * @param context The local scheduler state.
+ * @param events
  * @return Void.
  */
-void handle_submit_task(local_scheduler_state *s, task_spec *task);
+void process_plasma_notification(event_loop *loop,
+                                 int client_sock,
+                                 void *context,
+                                 int events);
 
 #endif /* PHOTON_SCHEDULER_H */

test/test.py

@@ -6,9 +6,11 @@ import subprocess
 import sys
 import unittest
 import random
+import threading
 import time
 
 import photon
+import plasma
 
 USE_VALGRIND = False
 
@@ -18,14 +20,19 @@ class TestPhotonClient(unittest.TestCase):
     # Start Redis.
     redis_executable = os.path.join(os.path.abspath(os.path.dirname(__file__)), "../common/thirdparty/redis-3.2.3/src/redis-server")
     self.p1 = subprocess.Popen([redis_executable, "--loglevel", "warning"])
+    # Start Plasma.
+    plasma_executable = os.path.join(os.path.abspath(os.path.dirname(__file__)), "../../plasma/build/plasma_store")
+    plasma_socket = "/tmp/plasma_store{}".format(random.randint(0, 10000))
+    self.p2 = subprocess.Popen([plasma_executable, "-s", plasma_socket])
     time.sleep(0.1)
+    self.plasma_client = plasma.PlasmaClient(plasma_socket)
     scheduler_executable = os.path.join(os.path.abspath(os.path.dirname(__file__)), "../build/photon_scheduler")
     scheduler_name = "/tmp/scheduler{}".format(random.randint(0, 10000))
-    command = [scheduler_executable, "-s", scheduler_name, "-r", "127.0.0.1:6379"]
+    command = [scheduler_executable, "-s", scheduler_name, "-r", "127.0.0.1:6379", "-p", plasma_socket]
     if USE_VALGRIND:
-      self.p2 = subprocess.Popen(["valgrind", "--track-origins=yes", "--leak-check=full", "--show-leak-kinds=all"] + command)
+      self.p3 = subprocess.Popen(["valgrind", "--track-origins=yes", "--leak-check=full", "--show-leak-kinds=all"] + command)
     else:
-      self.p2 = subprocess.Popen(command)
+      self.p3 = subprocess.Popen(command)
     if USE_VALGRIND:
       time.sleep(1.0)
     else:
@@ -36,21 +43,30 @@ class TestPhotonClient(unittest.TestCase):
   def tearDown(self):
     # Kill the Redis server.
     self.p1.kill()
+    # Kill Plasma.
+    self.p2.kill()
     # Kill the local scheduler.
     if USE_VALGRIND:
-      self.p2.send_signal(signal.SIGTERM)
-      self.p2.wait()
-      os._exit(self.p2.returncode)
+      self.p3.send_signal(signal.SIGTERM)
+      self.p3.wait()
+      os._exit(self.p3.returncode)
     else:
-      self.p2.kill()
-
+      self.p3.kill()
 
   def test_submit_and_get_task(self):
     # TODO(rkn): This should be a FunctionID.
     function_id = photon.ObjectID(20 * "a")
     object_ids = [photon.ObjectID(20 * chr(i)) for i in range(256)]
+    # Create and seal the objects in the object store so that we can schedule
+    # all of the subsequent tasks.
+    for object_id in object_ids:
+      self.plasma_client.create(object_id.id(), 0)
+      self.plasma_client.seal(object_id.id())
+    # Define some arguments to use for the tasks.
     args_list = [
       [],
+      #{},
+      #(),
       1 * [1],
       10 * [1],
       100 * [1],
@@ -104,6 +120,27 @@ class TestPhotonClient(unittest.TestCase):
       for num_return_vals in [0, 1, 2, 3, 5, 10, 100]:
         new_task = self.photon_client.get_task()
 
+  def test_scheduling_when_objects_ready(self):
+    # Create a task and submit it.
+    object_id = photon.ObjectID(20 * chr(0))
+    # TODO(rkn): This should be a FunctionID.
+    function_id = photon.ObjectID(20 * "a")
+    task = photon.Task(function_id, [object_id], 0)
+    self.photon_client.submit(task)
+    # Launch a thread to get the task.
+    def get_task():
+      self.photon_client.get_task()
+    t = threading.Thread(target=get_task)
+    t.start()
+    # Sleep to give the thread time to call get_task.
+    time.sleep(0.1)
+    # Create and seal the object ID in the object store. This should trigger a
+    # scheduling event.
+    self.plasma_client.create(object_id.id(), 0)
+    self.plasma_client.seal(object_id.id())
+    # Wait until the thread finishes so that we know the task was scheduled.
+    t.join()
+
 if __name__ == "__main__":
   if len(sys.argv) > 1:
     # pop the argument so we don't mess with unittest's own argument parser