Convert actor dummy objects to task execution edges. (#1281)

* Define execution dependencies flatbuffer and add to Redis commands

* Convert TaskSpec to TaskExecutionSpec

* Add execution dependencies to Python bindings

* Submitting actor tasks uses execution dependency API instead of dummy argument

* Fix dependency getters and some cleanup for fetching missing dependencies

* C++ convention

* Make TaskExecutionSpec a C++ class

* Convert local scheduler to use TaskExecutionSpec class

* Convert some pointers to references

* Finish conversion to TaskExecutionSpec class

* fix

* Fix

* Fix memory errors?

* Cast flatbuffers GetSize to size_t

* Fixes

* add more retries in global scheduler unit test

* fix linting and cast fbb.GetSize to size_t

* Style and doc

* Fix linting and simplify from_flatbuf.

src/common/common_protocol.cc

@@ -13,6 +13,16 @@ ObjectID from_flatbuf(const flatbuffers::String &string) {
   return object_id;
 }
 
+const std::vector<ObjectID> from_flatbuf(
+    const flatbuffers::Vector<flatbuffers::Offset<flatbuffers::String>>
+        &vector) {
+  std::vector<ObjectID> object_ids;
+  for (int64_t i = 0; i < vector.Length(); i++) {
+    object_ids.push_back(from_flatbuf(*vector.Get(i)));
+  }
+  return object_ids;
+}
+
 flatbuffers::Offset<
     flatbuffers::Vector<flatbuffers::Offset<flatbuffers::String>>>
 to_flatbuf(flatbuffers::FlatBufferBuilder &fbb,
@@ -25,6 +35,17 @@ to_flatbuf(flatbuffers::FlatBufferBuilder &fbb,
   return fbb.CreateVector(results);
 }
 
+flatbuffers::Offset<
+    flatbuffers::Vector<flatbuffers::Offset<flatbuffers::String>>>
+to_flatbuf(flatbuffers::FlatBufferBuilder &fbb,
+           const std::vector<ObjectID> &object_ids) {
+  std::vector<flatbuffers::Offset<flatbuffers::String>> results;
+  for (auto object_id : object_ids) {
+    results.push_back(to_flatbuf(fbb, object_id));
+  }
+  return fbb.CreateVector(results);
+}
+
 std::string string_from_flatbuf(const flatbuffers::String &string) {
   return std::string(string.data(), string.size());
 }

src/common/common_protocol.h

@@ -24,6 +24,14 @@ flatbuffers::Offset<flatbuffers::String> to_flatbuf(
 /// @return The object ID.
 ObjectID from_flatbuf(const flatbuffers::String &string);
 
+/// Convert a flatbuffer vector of strings to a vector of object IDs.
+///
+/// @param vector The flatbuffer vector.
+/// @return The vector of object IDs.
+const std::vector<ObjectID> from_flatbuf(
+    const flatbuffers::Vector<flatbuffers::Offset<flatbuffers::String>>
+        &vector);
+
 /// Convert an array of object IDs to a flatbuffer vector of strings.
 ///
 /// @param fbb Reference to the flatbuffer builder.
@@ -36,6 +44,16 @@ to_flatbuf(flatbuffers::FlatBufferBuilder &fbb,
            ObjectID object_ids[],
            int64_t num_objects);
 
+/// Convert a vector of object IDs to a flatbuffer vector of strings.
+///
+/// @param fbb Reference to the flatbuffer builder.
+/// @param object_ids Vector of object IDs.
+/// @return Flatbuffer vector of strings.
+flatbuffers::Offset<
+    flatbuffers::Vector<flatbuffers::Offset<flatbuffers::String>>>
+to_flatbuf(flatbuffers::FlatBufferBuilder &fbb,
+           const std::vector<ObjectID> &object_ids);
+
 /// Convert a flatbuffer string to a std::string.
 ///
 /// @param fbb Reference to the flatbuffer builder.

src/common/format/common.fbs

@@ -70,6 +70,14 @@ table ObjectInfo {
 
 root_type TaskInfo;
 
+table TaskExecutionDependencies {
+  // A list of object IDs representing this task's dependencies at execution
+  // time.
+  execution_dependencies: [string];
+}
+
+root_type TaskExecutionDependencies;
+
 table SubscribeToNotificationsReply {
   // The object ID of the object that the notification is about.
   object_id: string;
@@ -89,6 +97,8 @@ table TaskReply {
   state: long;
   // A local scheduler ID.
   local_scheduler_id: string;
+  // A string of bytes representing the task's TaskExecutionDependencies.
+  execution_dependencies: string;
   // A string of bytes representing the task specification.
   task_spec: string;
   // A boolean representing whether the update was successful. This field

src/common/lib/python/common_extension.cc

@@ -4,6 +4,7 @@
 
 #include "common.h"
 #include "common_extension.h"
+#include "common_protocol.h"
 #include "task.h"
 
 #include <string>
@@ -104,6 +105,8 @@ PyObject *PyTask_from_string(PyObject *self, PyObject *args) {
   result = (PyTask *) PyObject_Init((PyObject *) result, &PyTaskType);
   result->size = size;
   result->spec = TaskSpec_copy((TaskSpec *) data, size);
+  /* The created task does not include any execution dependencies. */
+  result->execution_dependencies = new std::vector<ObjectID>();
   /* TODO(pcm): Use flatbuffers validation here. */
   return (PyObject *) result;
 }
@@ -288,14 +291,18 @@ static int PyTask_init(PyTask *self, PyObject *args, PyObject *kwds) {
   TaskID parent_task_id;
   /* The number of tasks that the parent task has called prior to this one. */
   int parent_counter;
+  /* Arguments of the task that are execution-dependent. These must be
+   * PyObjectIDs). */
+  PyObject *execution_arguments = NULL;
   /* Dictionary of resource requirements for this task. */
   PyObject *resource_map = NULL;
-  if (!PyArg_ParseTuple(
-          args, "O&O&OiO&i|O&O&iOO", &PyObjectToUniqueID, &driver_id,
-          &PyObjectToUniqueID, &function_id, &arguments, &num_returns,
-          &PyObjectToUniqueID, &parent_task_id, &parent_counter,
-          &PyObjectToUniqueID, &actor_id, &PyObjectToUniqueID, &actor_handle_id,
-          &actor_counter, &is_actor_checkpoint_method_object, &resource_map)) {
+  if (!PyArg_ParseTuple(args, "O&O&OiO&i|O&O&iOOO", &PyObjectToUniqueID,
+                        &driver_id, &PyObjectToUniqueID, &function_id,
+                        &arguments, &num_returns, &PyObjectToUniqueID,
+                        &parent_task_id, &parent_counter, &PyObjectToUniqueID,
+                        &actor_id, &PyObjectToUniqueID, &actor_handle_id,
+                        &actor_counter, &is_actor_checkpoint_method_object,
+                        &execution_arguments, &resource_map)) {
     return -1;
   }
 
@@ -371,6 +378,23 @@ static int PyTask_init(PyTask *self, PyObject *args, PyObject *kwds) {
 
   /* Compute the task ID and the return object IDs. */
   self->spec = TaskSpec_finish_construct(g_task_builder, &self->size);
+
+  /* Set the task's execution dependencies. */
+  self->execution_dependencies = new std::vector<ObjectID>();
+  if (execution_arguments != NULL) {
+    size = PyList_Size(execution_arguments);
+    for (Py_ssize_t i = 0; i < size; ++i) {
+      PyObject *execution_arg = PyList_GetItem(execution_arguments, i);
+      if (!PyObject_IsInstance(execution_arg, (PyObject *) &PyObjectIDType)) {
+        PyErr_SetString(PyExc_TypeError,
+                        "Execution arguments must be an ObjectID.");
+        return -1;
+      }
+      self->execution_dependencies->push_back(
+          ((PyObjectID *) execution_arg)->object_id);
+    }
+  }
+
   return 0;
 }
 
@@ -378,6 +402,7 @@ static void PyTask_dealloc(PyTask *self) {
   if (self->spec != NULL) {
     TaskSpec_free(self->spec);
   }
+  delete self->execution_dependencies;
   Py_TYPE(self)->tp_free((PyObject *) self);
 }
 
@@ -471,6 +496,15 @@ static PyObject *PyTask_returns(PyObject *self) {
   return return_id_list;
 }
 
+static PyObject *PyTask_execution_dependencies_string(PyTask *self) {
+  flatbuffers::FlatBufferBuilder fbb;
+  auto execution_dependencies = CreateTaskExecutionDependencies(
+      fbb, to_flatbuf(fbb, *self->execution_dependencies));
+  fbb.Finish(execution_dependencies);
+  return PyBytes_FromStringAndSize((char *) fbb.GetBufferPointer(),
+                                   fbb.GetSize());
+}
+
 static PyMethodDef PyTask_methods[] = {
     {"function_id", (PyCFunction) PyTask_function_id, METH_NOARGS,
      "Return the function ID for this task."},
@@ -492,6 +526,9 @@ static PyMethodDef PyTask_methods[] = {
      "Return the resource vector of the task."},
     {"returns", (PyCFunction) PyTask_returns, METH_NOARGS,
      "Return the object IDs for the return values of the task."},
+    {"execution_dependencies_string",
+     (PyCFunction) PyTask_execution_dependencies_string, METH_NOARGS,
+     "Return the execution dependencies for the task as a string."},
     {NULL} /* Sentinel */
 };
 
@@ -543,6 +580,8 @@ PyObject *PyTask_make(TaskSpec *task_spec, int64_t task_size) {
   result = (PyTask *) PyObject_Init((PyObject *) result, &PyTaskType);
   result->spec = task_spec;
   result->size = task_size;
+  /* The created task does not include any execution dependencies. */
+  result->execution_dependencies = new std::vector<ObjectID>();
   return (PyObject *) result;
 }
 

src/common/lib/python/common_extension.h

@@ -1,6 +1,8 @@
 #ifndef COMMON_EXTENSION_H
 #define COMMON_EXTENSION_H
 
+#include <vector>
+
 #include <Python.h>
 #include "marshal.h"
 #include "structmember.h"
@@ -22,6 +24,7 @@ typedef struct {
   PyObject_HEAD
   int64_t size;
   TaskSpec *spec;
+  std::vector<ObjectID> *execution_dependencies;
 } PyTask;
 // clang-format on
 

src/common/redis_module/ray_redis_module.cc

@@ -760,11 +760,14 @@ int ReplyWithTask(RedisModuleCtx *ctx,
     /* If the key exists, look up the fields and return them in an array. */
     RedisModuleString *state = NULL;
     RedisModuleString *local_scheduler_id = NULL;
+    RedisModuleString *execution_dependencies = NULL;
     RedisModuleString *task_spec = NULL;
     RedisModule_HashGet(key, REDISMODULE_HASH_CFIELDS, "state", &state,
-                        "local_scheduler_id", &local_scheduler_id, "TaskSpec",
-                        &task_spec, NULL);
-    if (state == NULL || local_scheduler_id == NULL || task_spec == NULL) {
+                        "local_scheduler_id", &local_scheduler_id,
+                        "execution_dependencies", &execution_dependencies,
+                        "TaskSpec", &task_spec, NULL);
+    if (state == NULL || local_scheduler_id == NULL ||
+        execution_dependencies == NULL || task_spec == NULL) {
       /* We must have either all fields or no fields. */
       RedisModule_CloseKey(key);
       return RedisModule_ReplyWithError(
@@ -777,6 +780,7 @@ int ReplyWithTask(RedisModuleCtx *ctx,
       RedisModule_CloseKey(key);
       RedisModule_FreeString(ctx, state);
       RedisModule_FreeString(ctx, local_scheduler_id);
+      RedisModule_FreeString(ctx, execution_dependencies);
       RedisModule_FreeString(ctx, task_spec);
       return RedisModule_ReplyWithError(ctx, "Found invalid scheduling state.");
     }
@@ -785,6 +789,7 @@ int ReplyWithTask(RedisModuleCtx *ctx,
     auto message =
         CreateTaskReply(fbb, RedisStringToFlatbuf(fbb, task_id), state_integer,
                         RedisStringToFlatbuf(fbb, local_scheduler_id),
+                        RedisStringToFlatbuf(fbb, execution_dependencies),
                         RedisStringToFlatbuf(fbb, task_spec), updated);
     fbb.Finish(message);
 
@@ -794,6 +799,7 @@ int ReplyWithTask(RedisModuleCtx *ctx,
 
     RedisModule_FreeString(ctx, state);
     RedisModule_FreeString(ctx, local_scheduler_id);
+    RedisModule_FreeString(ctx, execution_dependencies);
     RedisModule_FreeString(ctx, task_spec);
   } else {
     /* If the key does not exist, return nil. */
@@ -904,6 +910,7 @@ int TaskTableWrite(RedisModuleCtx *ctx,
                    RedisModuleString *task_id,
                    RedisModuleString *state,
                    RedisModuleString *local_scheduler_id,
+                   RedisModuleString *execution_dependencies,
                    RedisModuleString *task_spec) {
   /* Extract the scheduling state. */
   long long state_value;
@@ -917,7 +924,8 @@ int TaskTableWrite(RedisModuleCtx *ctx,
       OpenPrefixedKey(ctx, TASK_PREFIX, task_id, REDISMODULE_WRITE);
   if (task_spec == NULL) {
     RedisModule_HashSet(key, REDISMODULE_HASH_CFIELDS, "state", state,
-                        "local_scheduler_id", local_scheduler_id, NULL);
+                        "local_scheduler_id", local_scheduler_id,
+                        "execution_dependencies", execution_dependencies, NULL);
     RedisModule_HashGet(key, REDISMODULE_HASH_CFIELDS, "TaskSpec",
                         &existing_task_spec, NULL);
     if (existing_task_spec == NULL) {
@@ -927,8 +935,9 @@ int TaskTableWrite(RedisModuleCtx *ctx,
     }
   } else {
     RedisModule_HashSet(key, REDISMODULE_HASH_CFIELDS, "state", state,
-                        "local_scheduler_id", local_scheduler_id, "TaskSpec",
-                        task_spec, NULL);
+                        "local_scheduler_id", local_scheduler_id,
+                        "execution_dependencies", execution_dependencies,
+                        "TaskSpec", task_spec, NULL);
   }
   RedisModule_CloseKey(key);
 
@@ -953,6 +962,7 @@ int TaskTableWrite(RedisModuleCtx *ctx,
     auto message =
         CreateTaskReply(fbb, RedisStringToFlatbuf(fbb, task_id), state_value,
                         RedisStringToFlatbuf(fbb, local_scheduler_id),
+                        RedisStringToFlatbuf(fbb, execution_dependencies),
                         RedisStringToFlatbuf(fbb, task_spec_to_use));
     fbb.Finish(message);
 
@@ -996,13 +1006,16 @@ int TaskTableWrite(RedisModuleCtx *ctx,
  *
  * This is called from a client with the command:
  *
- *     RAY.TASK_TABLE_ADD <task ID> <state> <local scheduler ID> <task spec>
+ *     RAY.TASK_TABLE_ADD <task ID> <state> <local scheduler ID>
+ *         <execution dependencies> <task spec>
  *
  * @param task_id A string that is the ID of the task.
  * @param state A string that is the current scheduling state (a
  *        scheduling_state enum instance).
  * @param local_scheduler_id A string that is the ray client ID of the
  *        associated local scheduler, if any.
+ * @param execution_dependencies A string that is the list of execution
+ *        dependencies.
  * @param task_spec A string that is the specification of the task, which can
  *        be cast to a `task_spec`.
  * @return OK if the operation was successful.
@@ -1010,11 +1023,11 @@ int TaskTableWrite(RedisModuleCtx *ctx,
 int TaskTableAddTask_RedisCommand(RedisModuleCtx *ctx,
                                   RedisModuleString **argv,
                                   int argc) {
-  if (argc != 5) {
+  if (argc != 6) {
     return RedisModule_WrongArity(ctx);
   }
 
-  return TaskTableWrite(ctx, argv[1], argv[2], argv[3], argv[4]);
+  return TaskTableWrite(ctx, argv[1], argv[2], argv[3], argv[4], argv[5]);
 }
 
 /**
@@ -1024,22 +1037,25 @@ int TaskTableAddTask_RedisCommand(RedisModuleCtx *ctx,
  * This is called from a client with the command:
  *
  *     RAY.TASK_TABLE_UPDATE <task ID> <state> <local scheduler ID>
+ *         <execution dependencies>
  *
  * @param task_id A string that is the ID of the task.
  * @param state A string that is the current scheduling state (a
  *        scheduling_state enum instance).
  * @param ray_client_id A string that is the ray client ID of the associated
  *        local scheduler, if any.
+ * @param execution_dependencies A string that is the list of execution
+ *        dependencies.
  * @return OK if the operation was successful.
  */
 int TaskTableUpdate_RedisCommand(RedisModuleCtx *ctx,
                                  RedisModuleString **argv,
                                  int argc) {
-  if (argc != 4) {
+  if (argc != 5) {
     return RedisModule_WrongArity(ctx);
   }
 
-  return TaskTableWrite(ctx, argv[1], argv[2], argv[3], NULL);
+  return TaskTableWrite(ctx, argv[1], argv[2], argv[3], argv[4], NULL);
 }
 
 /**

src/common/state/redis.cc

@@ -531,8 +531,13 @@ Task *parse_and_construct_task_from_redis_reply(redisReply *reply) {
     auto message = flatbuffers::GetRoot<TaskReply>(reply->str);
     TaskSpec *spec = (TaskSpec *) message->task_spec()->data();
     int64_t task_spec_size = message->task_spec()->size();
-    task = Task_alloc(spec, task_spec_size, message->state(),
-                      from_flatbuf(*message->local_scheduler_id()));
+    auto execution_dependencies =
+        flatbuffers::GetRoot<TaskExecutionDependencies>(
+            message->execution_dependencies()->data());
+    task = Task_alloc(
+        spec, task_spec_size, message->state(),
+        from_flatbuf(*message->local_scheduler_id()),
+        from_flatbuf(*execution_dependencies->execution_dependencies()));
   } else {
     LOG_FATAL("Unexpected reply type %d", reply->type);
   }
@@ -859,7 +864,9 @@ void redis_task_table_get_task_callback(redisAsyncContext *c,
     done_callback(task, callback_data->user_context);
   }
   /* Free the task if it is not NULL. */
-  Task_free(task);
+  if (task != NULL) {
+    Task_free(task);
+  }
 
   /* Clean up the timer and callback. */
   destroy_timer_callback(db->loop, callback_data);
@@ -917,18 +924,27 @@ void redis_task_table_add_task_callback(redisAsyncContext *c,
 void redis_task_table_add_task(TableCallbackData *callback_data) {
   DBHandle *db = callback_data->db_handle;
   Task *task = (Task *) callback_data->data->Get();
+  CHECKM(task != NULL, "NULL task passed to redis_task_table_add_task.");
+
   TaskID task_id = Task_task_id(task);
   DBClientID local_scheduler_id = Task_local_scheduler(task);
   redisAsyncContext *context = get_redis_context(db, task_id);
   int state = Task_state(task);
-  TaskSpec *spec = Task_task_spec(task);
 
-  CHECKM(task != NULL, "NULL task passed to redis_task_table_add_task.");
+  TaskExecutionSpec *execution_spec = Task_task_execution_spec(task);
+  TaskSpec *spec = execution_spec->Spec();
+
+  flatbuffers::FlatBufferBuilder fbb;
+  auto execution_dependencies = CreateTaskExecutionDependencies(
+      fbb, to_flatbuf(fbb, execution_spec->ExecutionDependencies()));
+  fbb.Finish(execution_dependencies);
+
   int status = redisAsyncCommand(
       context, redis_task_table_add_task_callback,
-      (void *) callback_data->timer_id, "RAY.TASK_TABLE_ADD %b %d %b %b",
+      (void *) callback_data->timer_id, "RAY.TASK_TABLE_ADD %b %d %b %b %b",
       task_id.id, sizeof(task_id.id), state, local_scheduler_id.id,
-      sizeof(local_scheduler_id.id), spec, Task_task_spec_size(task));
+      sizeof(local_scheduler_id.id), fbb.GetBufferPointer(),
+      (size_t) fbb.GetSize(), spec, execution_spec->SpecSize());
   if ((status == REDIS_ERR) || context->err) {
     LOG_REDIS_DEBUG(context, "error in redis_task_table_add_task");
   }
@@ -972,17 +988,25 @@ void redis_task_table_update_callback(redisAsyncContext *c,
 void redis_task_table_update(TableCallbackData *callback_data) {
   DBHandle *db = callback_data->db_handle;
   Task *task = (Task *) callback_data->data->Get();
+  CHECKM(task != NULL, "NULL task passed to redis_task_table_update.");
+
   TaskID task_id = Task_task_id(task);
   redisAsyncContext *context = get_redis_context(db, task_id);
   DBClientID local_scheduler_id = Task_local_scheduler(task);
   int state = Task_state(task);
 
-  CHECKM(task != NULL, "NULL task passed to redis_task_table_update.");
+  TaskExecutionSpec *execution_spec = Task_task_execution_spec(task);
+  flatbuffers::FlatBufferBuilder fbb;
+  auto execution_dependencies = CreateTaskExecutionDependencies(
+      fbb, to_flatbuf(fbb, execution_spec->ExecutionDependencies()));
+  fbb.Finish(execution_dependencies);
+
   int status = redisAsyncCommand(
       context, redis_task_table_update_callback,
-      (void *) callback_data->timer_id, "RAY.TASK_TABLE_UPDATE %b %d %b",
+      (void *) callback_data->timer_id, "RAY.TASK_TABLE_UPDATE %b %d %b %b",
       task_id.id, sizeof(task_id.id), state, local_scheduler_id.id,
-      sizeof(local_scheduler_id.id));
+      sizeof(local_scheduler_id.id), fbb.GetBufferPointer(),
+      (size_t) fbb.GetSize());
   if ((status == REDIS_ERR) || context->err) {
     LOG_REDIS_DEBUG(context, "error in redis_task_table_update");
   }
@@ -1081,11 +1105,17 @@ void redis_task_table_subscribe_callback(redisAsyncContext *c,
     /* Extract the local scheduler ID. */
     DBClientID local_scheduler_id =
         from_flatbuf(*message->local_scheduler_id());
+    /* Extract the execution dependencies. */
+    auto execution_dependencies =
+        flatbuffers::GetRoot<TaskExecutionDependencies>(
+            message->execution_dependencies()->data());
     /* Extract the task spec. */
     TaskSpec *spec = (TaskSpec *) message->task_spec()->data();
     int64_t task_spec_size = message->task_spec()->size();
     /* Create a task. */
-    Task *task = Task_alloc(spec, task_spec_size, state, local_scheduler_id);
+    Task *task = Task_alloc(
+        spec, task_spec_size, state, local_scheduler_id,
+        from_flatbuf(*execution_dependencies->execution_dependencies()));
 
     /* Call the subscribe callback if there is one. */
     TaskTableSubscribeData *data =
@@ -1382,7 +1412,7 @@ void redis_local_scheduler_table_disconnect(DBHandle *db) {
 
   redisReply *reply = (redisReply *) redisCommand(
       db->sync_context, "PUBLISH local_schedulers %b", fbb.GetBufferPointer(),
-      fbb.GetSize());
+      (size_t) fbb.GetSize());
   CHECK(reply->type != REDIS_REPLY_ERROR);
   CHECK(reply->type == REDIS_REPLY_INTEGER);
   LOG_DEBUG("%" PRId64 " subscribers received this publish.\n", reply->integer);
@@ -1467,7 +1497,7 @@ void redis_driver_table_send_driver_death(TableCallbackData *callback_data) {
   int status = redisAsyncCommand(
       db->context, redis_driver_table_send_driver_death_callback,
       (void *) callback_data->timer_id, "PUBLISH driver_deaths %b",
-      fbb.GetBufferPointer(), fbb.GetSize());
+      fbb.GetBufferPointer(), (size_t) fbb.GetSize());
   if ((status == REDIS_ERR) || db->context->err) {
     LOG_REDIS_DEBUG(db->context,
                     "error in redis_driver_table_send_driver_death");

src/common/task.cc

@@ -282,6 +282,17 @@ int64_t TaskSpec_num_args(TaskSpec *spec) {
   return message->args()->size();
 }
 
+int64_t TaskSpec_num_args_by_ref(TaskSpec *spec) {
+  int64_t num_args = TaskSpec_num_args(spec);
+  int64_t num_args_by_ref = 0;
+  for (int64_t i = 0; i < num_args; i++) {
+    if (TaskSpec_arg_by_ref(spec, i)) {
+      num_args_by_ref++;
+    }
+  }
+  return num_args_by_ref;
+}
+
 int TaskSpec_arg_id_count(TaskSpec *spec, int64_t arg_index) {
   CHECK(spec);
   auto message = flatbuffers::GetRoot<TaskInfo>(spec);
@@ -348,20 +359,6 @@ const std::unordered_map<std::string, double> TaskSpec_get_required_resources(
   return map_from_flatbuf(*message->required_resources());
 }
 
-bool TaskSpec_is_dependent_on(TaskSpec *spec, ObjectID object_id) {
-  int64_t num_args = TaskSpec_num_args(spec);
-  for (int i = 0; i < num_args; ++i) {
-    int count = TaskSpec_arg_id_count(spec, i);
-    for (int j = 0; j < count; j++) {
-      ObjectID arg_id = TaskSpec_arg_id(spec, i, j);
-      if (ObjectID_equal(arg_id, object_id)) {
-        return true;
-      }
-    }
-  }
-  return false;
-}
-
 TaskSpec *TaskSpec_copy(TaskSpec *spec, int64_t task_spec_size) {
   TaskSpec *copy = (TaskSpec *) malloc(task_spec_size);
   memcpy(copy, spec, task_spec_size);
@@ -372,32 +369,147 @@ void TaskSpec_free(TaskSpec *spec) {
   free(spec);
 }
 
+TaskExecutionSpec::TaskExecutionSpec(
+    const std::vector<ObjectID> &execution_dependencies,
+    TaskSpec *spec,
+    int64_t task_spec_size) {
+  execution_dependencies_ = execution_dependencies;
+  task_spec_size_ = task_spec_size;
+  TaskSpec *spec_copy = new TaskSpec[task_spec_size_];
+  memcpy(spec_copy, spec, task_spec_size);
+  spec_ = std::unique_ptr<TaskSpec[]>(spec_copy);
+}
+
+TaskExecutionSpec::TaskExecutionSpec(TaskExecutionSpec *other) {
+  execution_dependencies_ = other->execution_dependencies_;
+  task_spec_size_ = other->task_spec_size_;
+  TaskSpec *spec_copy = new TaskSpec[task_spec_size_];
+  memcpy(spec_copy, other->spec_.get(), task_spec_size_);
+  spec_ = std::unique_ptr<TaskSpec[]>(spec_copy);
+}
+
+std::vector<ObjectID> TaskExecutionSpec::ExecutionDependencies() {
+  return execution_dependencies_;
+}
+
+int64_t TaskExecutionSpec::SpecSize() {
+  return task_spec_size_;
+}
+
+TaskSpec *TaskExecutionSpec::Spec() {
+  return spec_.get();
+}
+
+int64_t TaskExecutionSpec::NumDependencies() {
+  TaskSpec *spec = Spec();
+  int64_t num_dependencies = TaskSpec_num_args(spec);
+  num_dependencies += execution_dependencies_.size();
+  return num_dependencies;
+}
+
+int TaskExecutionSpec::DependencyIdCount(int64_t dependency_index) {
+  TaskSpec *spec = Spec();
+  /* The first dependencies are the arguments of the task itself, followed by
+   * the execution dependencies. Find the total number of task arguments so
+   * that we can index into the correct list. */
+  int64_t num_args = TaskSpec_num_args(spec);
+  if (dependency_index < num_args) {
+    /* Index into the task arguments. */
+    return TaskSpec_arg_id_count(spec, dependency_index);
+  } else {
+    /* Index into the execution dependencies. */
+    dependency_index -= num_args;
+    CHECK((size_t) dependency_index < execution_dependencies_.size());
+    /* All elements in the execution dependency list have exactly one ID. */
+    return 1;
+  }
+}
+
+ObjectID TaskExecutionSpec::DependencyId(int64_t dependency_index,
+                                         int64_t id_index) {
+  TaskSpec *spec = Spec();
+  /* The first dependencies are the arguments of the task itself, followed by
+   * the execution dependencies. Find the total number of task arguments so
+   * that we can index into the correct list. */
+  int64_t num_args = TaskSpec_num_args(spec);
+  if (dependency_index < num_args) {
+    /* Index into the task arguments. */
+    return TaskSpec_arg_id(spec, dependency_index, id_index);
+  } else {
+    /* Index into the execution dependencies. */
+    dependency_index -= num_args;
+    CHECK((size_t) dependency_index < execution_dependencies_.size());
+    return execution_dependencies_[dependency_index];
+  }
+}
+
+bool TaskExecutionSpec::DependsOn(ObjectID object_id) {
+  // Iterate through the task arguments to see if it contains object_id.
+  TaskSpec *spec = Spec();
+  int64_t num_args = TaskSpec_num_args(spec);
+  for (int i = 0; i < num_args; ++i) {
+    int count = TaskSpec_arg_id_count(spec, i);
+    for (int j = 0; j < count; j++) {
+      ObjectID arg_id = TaskSpec_arg_id(spec, i, j);
+      if (ObjectID_equal(arg_id, object_id)) {
+        return true;
+      }
+    }
+  }
+  // Iterate through the execution dependencies to see if it contains object_id.
+  for (auto dependency_id : execution_dependencies_) {
+    if (ObjectID_equal(dependency_id, object_id)) {
+      return true;
+    }
+  }
+  // The requested object ID was not a task argument or an execution dependency.
+  // This task is not dependent on it.
+  return false;
+}
+
+bool TaskExecutionSpec::IsStaticDependency(int64_t dependency_index) {
+  TaskSpec *spec = Spec();
+  /* The first dependencies are the arguments of the task itself, followed by
+   * the execution dependencies. If the requested dependency index is a task
+   * argument, then it is a task dependency. */
+  int64_t num_args = TaskSpec_num_args(spec);
+  return (dependency_index < num_args);
+}
+
 /* TASK INSTANCES */
 
 Task *Task_alloc(TaskSpec *spec,
                  int64_t task_spec_size,
                  int state,
+                 DBClientID local_scheduler_id,
+                 const std::vector<ObjectID> &execution_dependencies) {
+  Task *result = new Task();
+  auto execution_spec =
+      new TaskExecutionSpec(execution_dependencies, spec, task_spec_size);
+  result->execution_spec = std::unique_ptr<TaskExecutionSpec>(execution_spec);
+  result->state = state;
+  result->local_scheduler_id = local_scheduler_id;
+  return result;
+}
+
+Task *Task_alloc(TaskExecutionSpec &execution_spec,
+                 int state,
                  DBClientID local_scheduler_id) {
-  int64_t size = sizeof(Task) - sizeof(TaskSpec) + task_spec_size;
-  Task *result = (Task *) malloc(size);
-  memset(result, 0, size);
+  Task *result = new Task();
+  result->execution_spec = std::unique_ptr<TaskExecutionSpec>(
+      new TaskExecutionSpec(&execution_spec));
   result->state = state;
   result->local_scheduler_id = local_scheduler_id;
-  result->task_spec_size = task_spec_size;
-  memcpy(&result->spec, spec, task_spec_size);
   return result;
 }
 
 Task *Task_copy(Task *other) {
-  int64_t size = Task_size(other);
-  Task *copy = (Task *) malloc(size);
-  CHECK(copy != NULL);
-  memcpy(copy, other, size);
-  return copy;
+  return Task_alloc(*Task_task_execution_spec(other), other->state,
+                    other->local_scheduler_id);
 }
 
 int64_t Task_size(Task *task_arg) {
-  return sizeof(Task) - sizeof(TaskSpec) + task_arg->task_spec_size;
+  return sizeof(Task) - sizeof(TaskSpec) + task_arg->execution_spec->SpecSize();
 }
 
 int Task_state(Task *task) {
@@ -416,19 +528,16 @@ void Task_set_local_scheduler(Task *task, DBClientID local_scheduler_id) {
   task->local_scheduler_id = local_scheduler_id;
 }
 
-TaskSpec *Task_task_spec(Task *task) {
-  return &task->spec;
-}
-
-int64_t Task_task_spec_size(Task *task) {
-  return task->task_spec_size;
+TaskExecutionSpec *Task_task_execution_spec(Task *task) {
+  return task->execution_spec.get();
 }
 
 TaskID Task_task_id(Task *task) {
-  TaskSpec *spec = Task_task_spec(task);
+  TaskExecutionSpec *execution_spec = Task_task_execution_spec(task);
+  TaskSpec *spec = execution_spec->Spec();
   return TaskSpec_task_id(spec);
 }
 
 void Task_free(Task *task) {
-  free(task);
+  delete task;
 }

src/common/task.h

@@ -13,6 +13,73 @@
 
 typedef uint8_t TaskSpec;
 
+class TaskExecutionSpec {
+ public:
+  TaskExecutionSpec(const std::vector<ObjectID> &execution_dependencies,
+                    TaskSpec *spec,
+                    int64_t task_spec_size);
+  TaskExecutionSpec(TaskExecutionSpec *execution_spec);
+
+  /// Get the task's execution dependencies.
+  ///
+  /// @return A vector of object IDs representing this task's execution
+  ///         dependencies.
+  std::vector<ObjectID> ExecutionDependencies();
+
+  /// Get the task spec size.
+  ///
+  /// @return The size of the immutable task spec.
+  int64_t SpecSize();
+
+  /// Get the task spec.
+  ///
+  /// @return A pointer to the immutable task spec.
+  TaskSpec *Spec();
+
+  /// Get the number of dependencies. This comprises the immutable task
+  /// arguments and the mutable execution dependencies.
+  ///
+  /// @return The number of dependencies.
+  int64_t NumDependencies();
+
+  /// Get the number of object IDs at the given dependency index.
+  ///
+  /// @param dependency_index The dependency index whose object IDs to count.
+  /// @return The number of object IDs at the given dependency_index.
+  int DependencyIdCount(int64_t dependency_index);
+
+  /// Get the object ID of a given dependency index.
+  ///
+  /// @param dependency_index The index at which we should look up the object
+  ///        ID.
+  /// @param id_index The index of the object ID.
+  ObjectID DependencyId(int64_t dependency_index, int64_t id_index);
+
+  /// Compute whether the task is dependent on an object ID.
+  ///
+  /// @param object_id The object ID that the task may be dependent on.
+  /// @return bool This returns true if the task is dependent on the given
+  ///         object ID and false otherwise.
+  bool DependsOn(ObjectID object_id);
+
+  /// Returns whether the given dependency index is a static dependency (an
+  /// argument of the immutable task).
+  ///
+  /// @param dependency_index The requested dependency index.
+  /// @return bool This returns true if the requested dependency index is
+  ///         immutable (an argument of the task).
+  bool IsStaticDependency(int64_t dependency_index);
+
+ private:
+  /** A list of object IDs representing this task's dependencies at execution
+   *  time. */
+  std::vector<ObjectID> execution_dependencies_;
+  /** The size of the task specification for this task. */
+  int64_t task_spec_size_;
+  /** The task specification for this task. */
+  std::unique_ptr<TaskSpec[]> spec_;
+};
+
 class TaskBuilder;
 
 #define NIL_TASK_ID NIL_ID
@@ -346,16 +413,6 @@ double TaskSpec_get_required_resource(const TaskSpec *spec,
 const std::unordered_map<std::string, double> TaskSpec_get_required_resources(
     const TaskSpec *spec);
 
-/**
- * Compute whether the task is dependent on an object ID.
- *
- * @param spec Task specification.
- * @param object_id The object ID that the task may be dependent on.
- * @return bool This returns true if the task is dependent on the given object
- *         ID and false otherwise.
- */
-bool TaskSpec_is_dependent_on(TaskSpec *spec, ObjectID object_id);
-
 /**
  * Compute the object id associated to a put call.
  *
@@ -426,10 +483,8 @@ struct Task {
   int state;
   /** The ID of the local scheduler involved. */
   DBClientID local_scheduler_id;
-  /** The size of the task specification for this task. */
-  int64_t task_spec_size;
-  /** The task specification for this task. */
-  TaskSpec spec;
+  /** The execution specification for this task. */
+  std::unique_ptr<TaskExecutionSpec> execution_spec;
 };
 
 /**
@@ -442,6 +497,11 @@ struct Task {
  */
 Task *Task_alloc(TaskSpec *spec,
                  int64_t task_spec_size,
+                 int state,
+                 DBClientID local_scheduler_id,
+                 const std::vector<ObjectID> &execution_dependencies);
+
+Task *Task_alloc(TaskExecutionSpec &execution_spec,
                  int state,
                  DBClientID local_scheduler_id);
 
@@ -468,10 +528,7 @@ DBClientID Task_local_scheduler(Task *task);
 /** Set the local scheduler ID for this task. */
 void Task_set_local_scheduler(Task *task, DBClientID local_scheduler_id);
 
-/** Task specification of this task. */
-TaskSpec *Task_task_spec(Task *task);
-
-int64_t Task_task_spec_size(Task *task);
+TaskExecutionSpec *Task_task_execution_spec(Task *task);
 
 /** Task ID of this task. */
 TaskID Task_task_id(Task *task);

src/common/test/db_tests.cc

@@ -138,8 +138,7 @@ void task_table_test_callback(Task *callback_task, void *user_data) {
   task_table_test_callback_called = 1;
   CHECK(Task_state(callback_task) == TASK_STATUS_SCHEDULED);
   CHECK(Task_size(callback_task) == Task_size(task_table_test_task));
-  CHECK(memcmp(callback_task, task_table_test_task, Task_size(callback_task)) ==
-        0);
+  CHECK(Task_equals(callback_task, task_table_test_task));
   event_loop *loop = (event_loop *) user_data;
   event_loop_stop(loop);
 }
@@ -151,11 +150,9 @@ TEST task_table_test(void) {
                             "127.0.0.1", std::vector<std::string>());
   db_attach(db, loop, false);
   DBClientID local_scheduler_id = globally_unique_id();
-  int64_t task_spec_size;
-  TaskSpec *spec = example_task_spec(1, 1, &task_spec_size);
-  task_table_test_task = Task_alloc(spec, task_spec_size, TASK_STATUS_SCHEDULED,
-                                    local_scheduler_id);
-  TaskSpec_free(spec);
+  TaskExecutionSpec spec = example_task_execution_spec(1, 1);
+  task_table_test_task =
+      Task_alloc(spec, TASK_STATUS_SCHEDULED, local_scheduler_id);
   RetryInfo retry = {
       .num_retries = NUM_RETRIES,
       .timeout = TIMEOUT,
@@ -186,13 +183,10 @@ TEST task_table_all_test(void) {
   DBHandle *db = db_connect(std::string("127.0.0.1"), 6379, "local_scheduler",
                             "127.0.0.1", std::vector<std::string>());
   db_attach(db, loop, false);
-  int64_t task_spec_size;
-  TaskSpec *spec = example_task_spec(1, 1, &task_spec_size);
+  TaskExecutionSpec spec = example_task_execution_spec(1, 1);
   /* Schedule two tasks on different local local schedulers. */
-  Task *task1 = Task_alloc(spec, task_spec_size, TASK_STATUS_SCHEDULED,
-                           globally_unique_id());
-  Task *task2 = Task_alloc(spec, task_spec_size, TASK_STATUS_SCHEDULED,
-                           globally_unique_id());
+  Task *task1 = Task_alloc(spec, TASK_STATUS_SCHEDULED, globally_unique_id());
+  Task *task2 = Task_alloc(spec, TASK_STATUS_SCHEDULED, globally_unique_id());
   RetryInfo retry = {
       .num_retries = NUM_RETRIES, .timeout = TIMEOUT, .fail_callback = NULL,
   };
@@ -207,7 +201,6 @@ TEST task_table_all_test(void) {
   event_loop_add_timer(loop, 200, (event_loop_timer_handler) timeout_handler,
                        NULL);
   event_loop_run(loop);
-  TaskSpec_free(spec);
   db_disconnect(db);
   destroy_outstanding_callbacks(loop);
   event_loop_destroy(loop);

src/common/test/example_task.h

@@ -7,10 +7,10 @@ extern TaskBuilder *g_task_builder;
 
 const int64_t arg_value_size = 1000;
 
-static inline TaskSpec *example_task_spec_with_args(int64_t num_args,
-                                                    int64_t num_returns,
-                                                    ObjectID arg_ids[],
-                                                    int64_t *task_spec_size) {
+static inline TaskExecutionSpec example_task_execution_spec_with_args(
+    int64_t num_args,
+    int64_t num_returns,
+    ObjectID arg_ids[]) {
   TaskID parent_task_id = globally_unique_id();
   FunctionID func_id = globally_unique_id();
   TaskSpec_start_construct(g_task_builder, NIL_ID, parent_task_id, 0,
@@ -25,36 +25,53 @@ static inline TaskSpec *example_task_spec_with_args(int64_t num_args,
     }
     TaskSpec_args_add_ref(g_task_builder, &arg_id, 1);
   }
-  return TaskSpec_finish_construct(g_task_builder, task_spec_size);
+  int64_t task_spec_size;
+  TaskSpec *spec = TaskSpec_finish_construct(g_task_builder, &task_spec_size);
+  std::vector<ObjectID> execution_dependencies;
+  auto execution_spec =
+      TaskExecutionSpec(execution_dependencies, spec, task_spec_size);
+  TaskSpec_free(spec);
+  return execution_spec;
 }
 
-static inline TaskSpec *example_task_spec(int64_t num_args,
-                                          int64_t num_returns,
-                                          int64_t *task_spec_size) {
-  return example_task_spec_with_args(num_args, num_returns, NULL,
-                                     task_spec_size);
+static inline TaskExecutionSpec example_task_execution_spec(
+    int64_t num_args,
+    int64_t num_returns) {
+  return example_task_execution_spec_with_args(num_args, num_returns, NULL);
 }
 
 static inline Task *example_task_with_args(int64_t num_args,
                                            int64_t num_returns,
                                            int task_state,
                                            ObjectID arg_ids[]) {
-  int64_t task_spec_size;
-  TaskSpec *spec = example_task_spec_with_args(num_args, num_returns, arg_ids,
-                                               &task_spec_size);
-  Task *instance = Task_alloc(spec, task_spec_size, task_state, NIL_ID);
-  TaskSpec_free(spec);
+  TaskExecutionSpec spec =
+      example_task_execution_spec_with_args(num_args, num_returns, arg_ids);
+  Task *instance = Task_alloc(spec, task_state, NIL_ID);
   return instance;
 }
 
 static inline Task *example_task(int64_t num_args,
                                  int64_t num_returns,
                                  int task_state) {
-  int64_t task_spec_size;
-  TaskSpec *spec = example_task_spec(num_args, num_returns, &task_spec_size);
-  Task *instance = Task_alloc(spec, task_spec_size, task_state, NIL_ID);
-  TaskSpec_free(spec);
+  TaskExecutionSpec spec = example_task_execution_spec(num_args, num_returns);
+  Task *instance = Task_alloc(spec, task_state, NIL_ID);
   return instance;
 }
 
+static inline bool Task_equals(Task *task1, Task *task2) {
+  if (task1->state != task2->state) {
+    return false;
+  }
+  if (!DBClientID_equal(task1->local_scheduler_id, task2->local_scheduler_id)) {
+    return false;
+  }
+  auto execution_spec1 = Task_task_execution_spec(task1);
+  auto execution_spec2 = Task_task_execution_spec(task2);
+  if (execution_spec1->SpecSize() != execution_spec2->SpecSize()) {
+    return false;
+  }
+  return memcmp(execution_spec1->Spec(), execution_spec2->Spec(),
+                execution_spec1->SpecSize()) == 0;
+}
+
 #endif /* EXAMPLE_TASK_H */

src/common/test/object_table_tests.cc

@@ -80,7 +80,7 @@ TEST new_object_test(void) {
   new_object_succeeded = 0;
   new_object_id = globally_unique_id();
   new_object_task = example_task(1, 1, TASK_STATUS_WAITING);
-  new_object_task_spec = Task_task_spec(new_object_task);
+  new_object_task_spec = Task_task_execution_spec(new_object_task)->Spec();
   new_object_task_id = TaskSpec_task_id(new_object_task_spec);
   g_loop = event_loop_create();
   DBHandle *db = db_connect(std::string("127.0.0.1"), 6379, "plasma_manager",

src/common/test/task_table_tests.cc

@@ -75,7 +75,7 @@ void add_lookup_fail_callback(UniqueID id,
 
 void lookup_success_callback(Task *task, void *context) {
   lookup_success = 1;
-  CHECK(memcmp(task, add_lookup_task, Task_size(task)) == 0);
+  CHECK(Task_equals(task, add_lookup_task));
   event_loop_stop(g_loop);
 }