Convert task_spec to flatbuffers (#255)

* convert Ray to C++

* convert task_spec to flatbuffers

* fix

* it compiles

* latest

* tests are passing

* task2 -> task

* fix

* fix

* fix

* fix

* fix

* linting

* fix valgrind

* upgrade flatbuffers

* use debug mode for valgrind

* fix naming and comments

* downgrade flatbuffers

* fix linting

* reintroduce TaskSpec_free

* rename TaskSpec -> TaskInfo

* refactoring

* linting

src/common/CMakeLists.txt

@@ -10,6 +10,22 @@ set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -fPIC")
 
 include_directories(thirdparty/ae)
 
+# Compile flatbuffers
+
+set(COMMON_FBS_SRC "${CMAKE_CURRENT_LIST_DIR}/format/common.fbs")
+set(OUTPUT_DIR ${CMAKE_CURRENT_LIST_DIR}/format/)
+
+add_custom_target(gen_common_fbs ALL)
+
+add_custom_command(
+  TARGET gen_common_fbs
+  COMMAND ${FLATBUFFERS_COMPILER} -c -o ${OUTPUT_DIR} ${COMMON_FBS_SRC}
+  DEPENDS ${FBS_DEPENDS}
+  COMMENT "Running flatc compiler on ${COMMON_FBS_SRC}"
+  VERBATIM)
+
+add_dependencies(gen_common_fbs flatbuffers_ep)
+
 add_custom_target(
    hiredis
    COMMAND make
@@ -18,6 +34,7 @@ add_custom_target(
 add_library(common STATIC
   event_loop.cc
   common.cc
+  common_protocol.cc
   task.cc
   io.cc
   net.cc
@@ -32,6 +49,8 @@ add_library(common STATIC
   thirdparty/ae/ae.c
   thirdparty/sha256.c)
 
+add_dependencies(common gen_common_fbs)
+
 target_link_libraries(common "${CMAKE_CURRENT_LIST_DIR}/thirdparty/hiredis/libhiredis.a")
 
 function(define_test test_name library)

src/common/common.h

@@ -12,7 +12,13 @@
 #endif
 
 #include "utarray.h"
+#ifdef __cplusplus
+extern "C" {
+#endif
 #include "sha256.h"
+#ifdef __cplusplus
+}
+#endif
 
 /** Definitions for Ray logging levels. */
 #define RAY_COMMON_DEBUG 0

src/common/common_protocol.cc

@@ -0,0 +1,26 @@
+#include "common_protocol.h"
+
+flatbuffers::Offset<flatbuffers::String> to_flatbuf(
+    flatbuffers::FlatBufferBuilder &fbb,
+    ObjectID object_id) {
+  return fbb.CreateString((char *) &object_id.id[0], sizeof(object_id.id));
+}
+
+ObjectID from_flatbuf(const flatbuffers::String *string) {
+  ObjectID object_id;
+  CHECK(string->size() == sizeof(object_id.id));
+  memcpy(&object_id.id[0], string->data(), sizeof(object_id.id));
+  return object_id;
+}
+
+flatbuffers::Offset<
+    flatbuffers::Vector<flatbuffers::Offset<flatbuffers::String>>>
+to_flatbuf(flatbuffers::FlatBufferBuilder &fbb,
+           ObjectID object_ids[],
+           int64_t num_objects) {
+  std::vector<flatbuffers::Offset<flatbuffers::String>> results;
+  for (size_t i = 0; i < num_objects; i++) {
+    results.push_back(to_flatbuf(fbb, object_ids[i]));
+  }
+  return fbb.CreateVector(results);
+}

src/common/common_protocol.h

@@ -0,0 +1,41 @@
+#ifndef COMMON_PROTOCOL_H
+#define COMMON_PROTOCOL_H
+
+#include "format/common_generated.h"
+
+#include "common.h"
+
+/**
+ * Convert an object ID to a flatbuffer string.
+ *
+ * @param fbb Reference to the flatbuffer builder.
+ * @param object_id The object ID to be converted.
+ * @return The flatbuffer string contining the object ID.
+ */
+flatbuffers::Offset<flatbuffers::String> to_flatbuf(
+    flatbuffers::FlatBufferBuilder &fbb,
+    ObjectID object_id);
+
+/**
+ * Convert a flatbuffer string to an object ID.
+ *
+ * @param string The flatbuffer string.
+ * @return The object ID.
+ */
+ObjectID from_flatbuf(const flatbuffers::String *string);
+
+/**
+ * Convert an array of object IDs to a flatbuffer vector of strings.
+ *
+ * @param fbb Reference to the flatbuffer builder.
+ * @param object_ids Array of object IDs.
+ * @param num_objects Number of elements in the array.
+ * @return Flatbuffer vector of strings.
+ */
+flatbuffers::Offset<
+    flatbuffers::Vector<flatbuffers::Offset<flatbuffers::String>>>
+to_flatbuf(flatbuffers::FlatBufferBuilder &fbb,
+           ObjectID object_ids[],
+           int64_t num_objects);
+
+#endif

src/common/format/common.fbs

@@ -0,0 +1,52 @@
+
+// Indices into resource vectors.
+// A resource vector maps a resource index to the number
+// of units of that resource required.
+
+// The total length of the resource vector is ResourceIndex_MAX.
+enum ResourceIndex:int {
+  // A central processing unit.
+  CPU = 0,
+  // A graphics processing unit.
+  GPU = 1,
+  // A dummy entry to make ResourceIndex_MAX equal to the length of
+  // a resource vector.
+  DUMMY = 2
+}
+
+table Arg {
+  // Object ID for pass-by-reference arguments.
+  object_id: string;
+  // Data for pass-by-value arguments.
+  data: string;
+}
+
+table TaskInfo {
+  // ID of the driver that created this task.
+  driver_id: string;
+  // Task ID of the task.
+  task_id: string;
+  // Task ID of the parent task.
+  parent_task_id: string;
+  // A count of the number of tasks submitted by the parent task before this one.
+  parent_counter: int;
+  // Actor ID of the task. This is the actor that this task is executed on
+  // or NIL_ACTOR_ID if the task is just a normal task.
+  actor_id: string;
+  // Number of tasks that have been submitted to this actor so far.
+  actor_counter: int;
+  // Function ID of the task.
+  function_id: string;
+  // Task arguments.
+  args: [Arg];
+  // Object IDs of return values.
+  returns: [string];
+  // The required_resources vector indicates the quantities of the different
+  // resources required by this task. The index in this vector corresponds to
+  // the resource type defined in the ResourceIndex enum. For example,
+  // required_resources[0] is the number of CPUs required, and
+  // required_resources[1] is the number of GPUs required.
+  required_resources: [double];
+}
+
+root_type TaskInfo;

src/common/lib/python/common_extension.cc

@@ -33,6 +33,8 @@ void init_pickle_module(void) {
   CHECK(pickle_protocol != NULL);
 }
 
+TaskBuilder *g_task_builder = NULL;
+
 /* Define the PyObjectID class. */
 
 int PyStringToUniqueID(PyObject *object, ObjectID *object_id) {
@@ -96,14 +98,10 @@ PyObject *PyTask_from_string(PyObject *self, PyObject *args) {
   }
   PyTask *result = PyObject_New(PyTask, &PyTaskType);
   result = (PyTask *) PyObject_Init((PyObject *) result, &PyTaskType);
-  result->spec = (task_spec *) malloc(size);
+  result->size = size;
+  result->spec = (TaskSpec *) malloc(size);
   memcpy(result->spec, data, size);
-  /* TODO(pcm): Better error checking once we use flatbuffers. */
-  if (size != task_spec_size(result->spec)) {
-    PyErr_SetString(CommonError,
-                    "task_from_string: task specification string malformed");
-    return NULL;
-  }
+  /* TODO(pcm): Use flatbuffers validation here. */
   return (PyObject *) result;
 }
 
@@ -123,8 +121,7 @@ PyObject *PyTask_to_string(PyObject *self, PyObject *args) {
     return NULL;
   }
   PyTask *task = (PyTask *) arg;
-  return PyBytes_FromStringAndSize((char *) task->spec,
-                                   task_spec_size(task->spec));
+  return PyBytes_FromStringAndSize((char *) task->spec, task->size);
 }
 
 static PyObject *PyObjectID_id(PyObject *self) {
@@ -271,9 +268,7 @@ static int PyTask_init(PyTask *self, PyObject *args, PyObject *kwds) {
   FunctionID function_id;
   /* Arguments of the task (can be PyObjectIDs or Python values). */
   PyObject *arguments;
-  /* Array of pointers to string representations of pass-by-value args. */
-  UT_array *val_repr_ptrs;
-  utarray_new(val_repr_ptrs, &ut_ptr_icd);
+  /* Number of return values of this task. */
   int num_returns;
   /* The ID of the task that called this task. */
   TaskID parent_task_id;
@@ -289,101 +284,84 @@ static int PyTask_init(PyTask *self, PyObject *args, PyObject *kwds) {
     return -1;
   }
   Py_ssize_t size = PyList_Size(arguments);
-  /* Determine the size of pass by value data in bytes. */
-  Py_ssize_t value_data_bytes = 0;
-  for (Py_ssize_t i = 0; i < size; ++i) {
-    PyObject *arg = PyList_GetItem(arguments, i);
-    if (!PyObject_IsInstance(arg, (PyObject *) &PyObjectIDType)) {
-      CHECK(pickle_module != NULL);
-      CHECK(pickle_dumps != NULL);
-      PyObject *data = PyObject_CallMethodObjArgs(pickle_module, pickle_dumps,
-                                                  arg, pickle_protocol, NULL);
-      value_data_bytes += PyBytes_Size(data);
-      utarray_push_back(val_repr_ptrs, &data);
-    }
-  }
   /* Construct the task specification. */
-  int val_repr_index = 0;
-  self->spec = start_construct_task_spec(
-      driver_id, parent_task_id, parent_counter, actor_id, actor_counter,
-      function_id, size, num_returns, value_data_bytes);
+  TaskSpec_start_construct(g_task_builder, driver_id, parent_task_id,
+                           parent_counter, actor_id, actor_counter, function_id,
+                           num_returns);
   /* Add the task arguments. */
   for (Py_ssize_t i = 0; i < size; ++i) {
     PyObject *arg = PyList_GetItem(arguments, i);
     if (PyObject_IsInstance(arg, (PyObject *) &PyObjectIDType)) {
-      task_args_add_ref(self->spec, ((PyObjectID *) arg)->object_id);
+      TaskSpec_args_add_ref(g_task_builder, ((PyObjectID *) arg)->object_id);
     } else {
       /* We do this check because we cast a signed int to an unsigned int. */
-      CHECK(val_repr_index >= 0);
-      PyObject *data = *((PyObject **) utarray_eltptr(
-          val_repr_ptrs, (uint64_t) val_repr_index));
-      task_args_add_val(self->spec, (uint8_t *) PyBytes_AS_STRING(data),
-                        PyBytes_GET_SIZE(data));
+      PyObject *data = PyObject_CallMethodObjArgs(pickle_module, pickle_dumps,
+                                                  arg, pickle_protocol, NULL);
+      TaskSpec_args_add_val(g_task_builder, (uint8_t *) PyBytes_AS_STRING(data),
+                            PyBytes_GET_SIZE(data));
       Py_DECREF(data);
-      val_repr_index += 1;
     }
   }
-  utarray_free(val_repr_ptrs);
   /* Set the resource vector of the task. */
   if (resource_vector != NULL) {
-    CHECK(PyList_Size(resource_vector) == MAX_RESOURCE_INDEX);
-    for (int i = 0; i < MAX_RESOURCE_INDEX; ++i) {
+    CHECK(PyList_Size(resource_vector) == ResourceIndex_MAX);
+    for (int i = 0; i < ResourceIndex_MAX; ++i) {
       PyObject *resource_entry = PyList_GetItem(resource_vector, i);
-      task_spec_set_required_resource(self->spec, i,
-                                      PyFloat_AsDouble(resource_entry));
+      TaskSpec_set_required_resource(g_task_builder, i,
+                                     PyFloat_AsDouble(resource_entry));
     }
   } else {
-    for (int i = 0; i < MAX_RESOURCE_INDEX; ++i) {
-      task_spec_set_required_resource(self->spec, i,
-                                      i == CPU_RESOURCE_INDEX ? 1.0 : 0.0);
+    for (int i = 0; i < ResourceIndex_MAX; ++i) {
+      TaskSpec_set_required_resource(g_task_builder, i,
+                                     i == ResourceIndex_CPU ? 1.0 : 0.0);
     }
   }
   /* Compute the task ID and the return object IDs. */
-  finish_construct_task_spec(self->spec);
+  self->spec = TaskSpec_finish_construct(g_task_builder, &self->size);
   return 0;
 }
 
 static void PyTask_dealloc(PyTask *self) {
   if (self->spec != NULL) {
-    free_task_spec(self->spec);
+    TaskSpec_free(self->spec);
   }
   Py_TYPE(self)->tp_free((PyObject *) self);
 }
 
 static PyObject *PyTask_function_id(PyObject *self) {
-  FunctionID function_id = task_function(((PyTask *) self)->spec);
+  FunctionID function_id = TaskSpec_function(((PyTask *) self)->spec);
   return PyObjectID_make(function_id);
 }
 
 static PyObject *PyTask_actor_id(PyObject *self) {
-  ActorID actor_id = task_spec_actor_id(((PyTask *) self)->spec);
+  ActorID actor_id = TaskSpec_actor_id(((PyTask *) self)->spec);
   return PyObjectID_make(actor_id);
 }
 
 static PyObject *PyTask_driver_id(PyObject *self) {
-  UniqueID driver_id = task_spec_driver_id(((PyTask *) self)->spec);
+  UniqueID driver_id = TaskSpec_driver_id(((PyTask *) self)->spec);
   return PyObjectID_make(driver_id);
 }
 
 static PyObject *PyTask_task_id(PyObject *self) {
-  TaskID task_id = task_spec_id(((PyTask *) self)->spec);
+  TaskID task_id = TaskSpec_task_id(((PyTask *) self)->spec);
   return PyObjectID_make(task_id);
 }
 
 static PyObject *PyTask_arguments(PyObject *self) {
-  task_spec *task = ((PyTask *) self)->spec;
-  int64_t num_args = task_num_args(task);
+  TaskSpec *task = ((PyTask *) self)->spec;
+  int64_t num_args = TaskSpec_num_args(task);
   PyObject *arg_list = PyList_New((Py_ssize_t) num_args);
   for (int i = 0; i < num_args; ++i) {
-    if (task_arg_type(task, i) == ARG_BY_REF) {
-      ObjectID object_id = task_arg_id(task, i);
+    if (TaskSpec_arg_by_ref(task, i)) {
+      ObjectID object_id = TaskSpec_arg_id(task, i);
       PyList_SetItem(arg_list, i, PyObjectID_make(object_id));
     } else {
       CHECK(pickle_module != NULL);
       CHECK(pickle_loads != NULL);
       PyObject *str =
-          PyBytes_FromStringAndSize((char *) task_arg_val(task, i),
-                                    (Py_ssize_t) task_arg_length(task, i));
+          PyBytes_FromStringAndSize((char *) TaskSpec_arg_val(task, i),
+                                    (Py_ssize_t) TaskSpec_arg_length(task, i));
       PyObject *val =
           PyObject_CallMethodObjArgs(pickle_module, pickle_loads, str, NULL);
       Py_XDECREF(str);
@@ -394,21 +372,21 @@ static PyObject *PyTask_arguments(PyObject *self) {
 }
 
 static PyObject *PyTask_required_resources(PyObject *self) {
-  task_spec *task = ((PyTask *) self)->spec;
-  PyObject *required_resources = PyList_New((Py_ssize_t) MAX_RESOURCE_INDEX);
-  for (int i = 0; i < MAX_RESOURCE_INDEX; ++i) {
-    double r = task_spec_get_required_resource(task, i);
+  TaskSpec *task = ((PyTask *) self)->spec;
+  PyObject *required_resources = PyList_New((Py_ssize_t) ResourceIndex_MAX);
+  for (int i = 0; i < ResourceIndex_MAX; ++i) {
+    double r = TaskSpec_get_required_resource(task, i);
     PyList_SetItem(required_resources, i, PyFloat_FromDouble(r));
   }
   return required_resources;
 }
 
 static PyObject *PyTask_returns(PyObject *self) {
-  task_spec *task = ((PyTask *) self)->spec;
-  int64_t num_returns = task_num_returns(task);
+  TaskSpec *task = ((PyTask *) self)->spec;
+  int64_t num_returns = TaskSpec_num_returns(task);
   PyObject *return_id_list = PyList_New((Py_ssize_t) num_returns);
   for (int i = 0; i < num_returns; ++i) {
-    ObjectID object_id = task_return(task, i);
+    ObjectID object_id = TaskSpec_return(task, i);
     PyList_SetItem(return_id_list, i, PyObjectID_make(object_id));
   }
   return return_id_list;
@@ -474,11 +452,12 @@ PyTypeObject PyTaskType = {
 };
 
 /* Create a PyTask from a C struct. The resulting PyTask takes ownership of the
- * task_spec and will deallocate the task_spec in the PyTask destructor. */
-PyObject *PyTask_make(task_spec *task_spec) {
+ * TaskSpec and will deallocate the TaskSpec in the PyTask destructor. */
+PyObject *PyTask_make(TaskSpec *task_spec, int64_t task_size) {
   PyTask *result = PyObject_New(PyTask, &PyTaskType);
   result = (PyTask *) PyObject_Init((PyObject *) result, &PyTaskType);
   result->spec = task_spec;
+  result->size = task_size;
   return (PyObject *) result;
 }
 

src/common/lib/python/common_extension.h

@@ -6,7 +6,9 @@
 #include "structmember.h"
 
 #include "common.h"
-#include "task.h"
+
+typedef uint8_t TaskSpec;
+struct TaskBuilder;
 
 extern PyObject *CommonError;
 
@@ -18,7 +20,8 @@ typedef struct {
 
 typedef struct {
   PyObject_HEAD
-  task_spec *spec;
+  int64_t size;
+  TaskSpec *spec;
 } PyTask;
 // clang-format on
 
@@ -33,9 +36,11 @@ extern PyObject *pickle_loads;
 
 void init_pickle_module(void);
 
+extern TaskBuilder *g_task_builder;
+
 int PyStringToUniqueID(PyObject *object, ObjectID *object_id);
 
-int PyObjectToUniqueID(PyObject *object, ObjectID *objectid);
+int PyObjectToUniqueID(PyObject *object, ObjectID *object_id);
 
 PyObject *PyObjectID_make(ObjectID object_id);
 
@@ -46,6 +51,6 @@ PyObject *PyTask_from_string(PyObject *, PyObject *args);
 
 PyObject *compute_put_id(PyObject *self, PyObject *args);
 
-PyObject *PyTask_make(task_spec *task_spec);
+PyObject *PyTask_make(TaskSpec *task_spec, int64_t task_size);
 
 #endif /* COMMON_EXTENSION_H */

src/common/redis_module/ray_redis_module.c

@@ -714,7 +714,7 @@ int ReplyWithTask(RedisModuleCtx *ctx, RedisModuleString *task_id) {
     RedisModuleString *local_scheduler_id = NULL;
     RedisModuleString *task_spec = NULL;
     RedisModule_HashGet(key, REDISMODULE_HASH_CFIELDS, "state", &state, "node",
-                        &local_scheduler_id, "task_spec", &task_spec, NULL);
+                        &local_scheduler_id, "TaskSpec", &task_spec, NULL);
     if (state == NULL || local_scheduler_id == NULL || task_spec == NULL) {
       /* We must have either all fields or no fields. */
       RedisModule_CloseKey(key);
@@ -817,7 +817,7 @@ int TaskTableWrite(RedisModuleCtx *ctx,
   if (task_spec == NULL) {
     RedisModule_HashSet(key, REDISMODULE_HASH_CFIELDS, "state", state, "node",
                         node_id, NULL);
-    RedisModule_HashGet(key, REDISMODULE_HASH_CFIELDS, "task_spec",
+    RedisModule_HashGet(key, REDISMODULE_HASH_CFIELDS, "TaskSpec",
                         &existing_task_spec, NULL);
     if (existing_task_spec == NULL) {
       RedisModule_CloseKey(key);
@@ -827,7 +827,7 @@ int TaskTableWrite(RedisModuleCtx *ctx,
     }
   } else {
     RedisModule_HashSet(key, REDISMODULE_HASH_CFIELDS, "state", state, "node",
-                        node_id, "task_spec", task_spec, NULL);
+                        node_id, "TaskSpec", task_spec, NULL);
   }
   RedisModule_CloseKey(key);
 

src/common/state/local_scheduler_table.h

@@ -17,10 +17,10 @@ typedef struct {
   int available_workers;
   /** The resource vector of resources generally available to this local
    *  scheduler. */
-  double static_resources[MAX_RESOURCE_INDEX];
+  double static_resources[ResourceIndex_MAX];
   /** The resource vector of resources currently available to this local
    *  scheduler. */
-  double dynamic_resources[MAX_RESOURCE_INDEX];
+  double dynamic_resources[ResourceIndex_MAX];
 } LocalSchedulerInfo;
 
 /*

src/common/state/redis.cc

@@ -358,7 +358,7 @@ Task *parse_and_construct_task_from_redis_reply(redisReply *reply) {
   } else if (reply->type == REDIS_REPLY_ARRAY) {
     /* Check that the reply is as expected. The 0th element is the scheduling
      * state. The 1st element is the db_client_id of the associated local
-     * scheduler, and the 2nd element is the task_spec. */
+     * scheduler, and the 2nd element is the TaskSpec. */
     CHECK(reply->elements == 3);
     CHECK(reply->element[0]->type == REDIS_REPLY_INTEGER);
     CHECK(reply->element[1]->type == REDIS_REPLY_STRING);
@@ -371,12 +371,11 @@ Task *parse_and_construct_task_from_redis_reply(redisReply *reply) {
     memcpy(local_scheduler_id.id, reply->element[1]->str,
            reply->element[1]->len);
     /* Parse the task spec. */
-    task_spec *spec = (task_spec *) malloc(reply->element[2]->len);
+    TaskSpec *spec = (TaskSpec *) malloc(reply->element[2]->len);
     memcpy(spec, reply->element[2]->str, reply->element[2]->len);
-    CHECK(task_spec_size(spec) == reply->element[2]->len);
-    task = Task_alloc(spec, state, local_scheduler_id);
+    task = Task_alloc(spec, reply->element[2]->len, state, local_scheduler_id);
     /* Free the task spec. */
-    free_task_spec(spec);
+    TaskSpec_free(spec);
   } else {
     LOG_FATAL("Unexpected reply type %d", reply->type);
   }
@@ -777,16 +776,16 @@ void redis_task_table_add_task(TableCallbackData *callback_data) {
   DBHandle *db = callback_data->db_handle;
   Task *task = (Task *) callback_data->data;
   TaskID task_id = Task_task_id(task);
-  DBClientID local_scheduler_id = Task_local_scheduler_id(task);
+  DBClientID local_scheduler_id = Task_local_scheduler(task);
   int state = Task_state(task);
-  task_spec *spec = Task_task_spec(task);
+  TaskSpec *spec = Task_task_spec(task);
 
   CHECKM(task != NULL, "NULL task passed to redis_task_table_add_task.");
   int status = redisAsyncCommand(
       db->context, redis_task_table_add_task_callback,
       (void *) callback_data->timer_id, "RAY.TASK_TABLE_ADD %b %d %b %b",
       task_id.id, sizeof(task_id.id), state, local_scheduler_id.id,
-      sizeof(local_scheduler_id.id), spec, task_spec_size(spec));
+      sizeof(local_scheduler_id.id), spec, Task_task_spec_size(task));
   if ((status == REDIS_ERR) || db->context->err) {
     LOG_REDIS_DEBUG(db->context, "error in redis_task_table_add_task");
   }
@@ -814,7 +813,7 @@ void redis_task_table_update(TableCallbackData *callback_data) {
   DBHandle *db = callback_data->db_handle;
   Task *task = (Task *) callback_data->data;
   TaskID task_id = Task_task_id(task);
-  DBClientID local_scheduler_id = Task_local_scheduler_id(task);
+  DBClientID local_scheduler_id = Task_local_scheduler(task);
   int state = Task_state(task);
 
   CHECKM(task != NULL, "NULL task passed to redis_task_table_update.");
@@ -875,7 +874,8 @@ void parse_task_table_subscribe_callback(char *payload,
                                          TaskID *task_id,
                                          int *state,
                                          DBClientID *local_scheduler_id,
-                                         task_spec **spec) {
+                                         TaskSpec **spec,
+                                         int64_t *task_spec_size) {
   /* Note that the state is padded with spaces to consist of precisely two
    * characters. */
   int task_spec_payload_size =
@@ -902,9 +902,9 @@ void parse_task_table_subscribe_callback(char *payload,
   CHECK(memcmp(space_str, &payload[offset], strlen(space_str)) == 0);
   offset += strlen(space_str);
   /* Read in the task spec. */
-  *spec = (task_spec *) malloc(task_spec_payload_size);
+  *spec = (TaskSpec *) malloc(task_spec_payload_size);
   memcpy(*spec, &payload[offset], task_spec_payload_size);
-  CHECK(task_spec_size(*spec) == task_spec_payload_size);
+  *task_spec_size = task_spec_payload_size;
 }
 
 void redis_task_table_subscribe_callback(redisAsyncContext *c,
@@ -932,11 +932,13 @@ void redis_task_table_subscribe_callback(redisAsyncContext *c,
     TaskID task_id;
     int state;
     DBClientID local_scheduler_id;
-    task_spec *spec;
+    TaskSpec *spec;
+    int64_t task_spec_size;
     parse_task_table_subscribe_callback(payload->str, payload->len, &task_id,
-                                        &state, &local_scheduler_id, &spec);
-    Task *task = Task_alloc(spec, state, local_scheduler_id);
-    free(spec);
+                                        &state, &local_scheduler_id, &spec,
+                                        &task_spec_size);
+    Task *task = Task_alloc(spec, task_spec_size, state, local_scheduler_id);
+    TaskSpec_free(spec);
     /* Call the subscribe callback if there is one. */
     if (data->subscribe_callback != NULL) {
       data->subscribe_callback(task, data->subscribe_context);

src/common/task.cc

@@ -1,135 +1,12 @@
-#include <stdlib.h>
-#include <stdio.h>
-#include <string.h>
+#include <limits.h>
+
+#include "common_protocol.h"
+
+#include "task.h"
 
 extern "C" {
 #include "sha256.h"
 }
-#include "utarray.h"
-
-#include "task.h"
-#include "common.h"
-#include "io.h"
-
-/* TASK SPECIFICATIONS */
-
-/* Tasks are stored in a consecutive chunk of memory, the first
- * sizeof(task_spec) bytes are arranged according to the struct
- * task_spec. Then there is an array of task_args of length
- * (num_args + num_returns), and then follows the data of
- * pass-by-value arguments of size args_value_size. The offsets in the
- * task_arg.val are with respect to the end of the augmented structure,
- * i.e. with respect to the address &task_spec.args_and_returns[0] +
- * (task_spec->num_args + task_spec->num_returns) * sizeof(task_arg). */
-
-typedef struct {
-  /* Either ARG_BY_REF or ARG_BY_VAL. */
-  int8_t type;
-  union {
-    ObjectID obj_id;
-    struct {
-      /* Offset where the data associated to this arg is located relative
-       * to &task_spec.args_and_returns[0]. */
-      ptrdiff_t offset;
-      int64_t length;
-    } value;
-  };
-} task_arg;
-
-struct task_spec_impl {
-  /** ID of the driver that created this task. */
-  UniqueID driver_id;
-  /** Task ID of the task. */
-  TaskID task_id;
-  /** Task ID of the parent task. */
-  TaskID parent_task_id;
-  /** A count of the number of tasks submitted by the parent task before this
-   *  one. */
-  int64_t parent_counter;
-  /** Actor ID of the task. This is the actor that this task is executed on
-   *  or NIL_ACTOR_ID if the task is just a normal task. */
-  ActorID actor_id;
-  /** Number of tasks that have been submitted to this actor so far. */
-  int64_t actor_counter;
-  /** Function ID of the task. */
-  FunctionID function_id;
-  /** Total number of arguments. */
-  int64_t num_args;
-  /** Index of the last argument that has been constructed. */
-  int64_t arg_index;
-  /** Number of return values. */
-  int64_t num_returns;
-  /** Number of bytes the pass-by-value arguments are occupying. */
-  int64_t args_value_size;
-  /** The offset of the number of bytes of pass-by-value data that
-   *  has been written so far, relative to &task_spec->args_and_returns[0] +
-   *  (task_spec->num_args + task_spec->num_returns) * sizeof(task_arg) */
-  int64_t args_value_offset;
-  /** Resource vector for this task. A resource vector maps a resource index
-   *  (like "cpu" or "gpu") to the number of units of that resource required.
-   *  Note that this will allow us to support arbitrary attributes:
-   *  For example, we can have a coloring of nodes and "red" can correspond
-   *  to 0.0, "green" to 1.0 and "yellow" to 2.0. */
-  double required_resources[MAX_RESOURCE_INDEX];
-  /** Argument and return IDs as well as offsets for pass-by-value args. */
-  task_arg args_and_returns[0];
-};
-
-/* The size of a task specification is given by the following expression. */
-#define TASK_SPEC_SIZE(NUM_ARGS, NUM_RETURNS, ARGS_VALUE_SIZE)           \
-  (sizeof(task_spec) + ((NUM_ARGS) + (NUM_RETURNS)) * sizeof(task_arg) + \
-   (ARGS_VALUE_SIZE))
-
-bool TaskID_equal(TaskID first_id, TaskID second_id) {
-  return UNIQUE_ID_EQ(first_id, second_id);
-}
-
-bool TaskID_is_nil(TaskID id) {
-  return TaskID_equal(id, NIL_TASK_ID);
-}
-
-bool ActorID_equal(ActorID first_id, ActorID second_id) {
-  return UNIQUE_ID_EQ(first_id, second_id);
-}
-
-bool FunctionID_equal(FunctionID first_id, FunctionID second_id) {
-  return UNIQUE_ID_EQ(first_id, second_id);
-}
-
-bool FunctionID_is_nil(FunctionID id) {
-  return FunctionID_equal(id, NIL_FUNCTION_ID);
-}
-
-TaskID *task_return_ptr(task_spec *spec, int64_t return_index) {
-  DCHECK(0 <= return_index && return_index < spec->num_returns);
-  task_arg *ret = &spec->args_and_returns[spec->num_args + return_index];
-  DCHECK(ret->type == ARG_BY_REF);
-  return &ret->obj_id;
-}
-
-/* Compute the task ID. This assumes that all of the other fields have been set
- * and that the return IDs have not been set. It assumes the task_spec was
- * zero-initialized so that uninitialized fields will not make the task ID
- * nondeterministic. */
-TaskID compute_task_id(task_spec *spec) {
-  /* Check that the task ID and return ID fields of the task_spec are
-   * uninitialized. */
-  DCHECK(TaskID_equal(spec->task_id, NIL_TASK_ID));
-  for (int i = 0; i < spec->num_returns; ++i) {
-    DCHECK(ObjectID_equal(*task_return_ptr(spec, i), NIL_ID));
-  }
-  /* Compute a SHA256 hash of the task_spec. */
-  SHA256_CTX ctx;
-  BYTE buff[DIGEST_SIZE];
-  sha256_init(&ctx);
-  sha256_update(&ctx, (BYTE *) spec, task_spec_size(spec));
-  sha256_final(&ctx, buff);
-  /* Create a task ID out of the hash. This will truncate the hash. */
-  TaskID task_id;
-  CHECK(sizeof(task_id) <= DIGEST_SIZE);
-  memcpy(&task_id.id, buff, sizeof(task_id.id));
-  return task_id;
-}
 
 ObjectID task_compute_return_id(TaskID task_id, int64_t return_index) {
   /* Here, return_indices need to be >= 0, so we can use negative
@@ -155,209 +32,266 @@ ObjectID task_compute_put_id(TaskID task_id, int64_t put_index) {
   return put_id;
 }
 
-task_spec *start_construct_task_spec(UniqueID driver_id,
-                                     TaskID parent_task_id,
-                                     int64_t parent_counter,
-                                     ActorID actor_id,
-                                     int64_t actor_counter,
-                                     FunctionID function_id,
-                                     int64_t num_args,
-                                     int64_t num_returns,
-                                     int64_t args_value_size) {
-  int64_t size = TASK_SPEC_SIZE(num_args, num_returns, args_value_size);
-  task_spec *task = (task_spec *) malloc(size);
-  memset(task, 0, size);
-  task->driver_id = driver_id;
-  task->task_id = NIL_TASK_ID;
-  task->parent_task_id = parent_task_id;
-  task->parent_counter = parent_counter;
-  task->actor_id = actor_id;
-  task->actor_counter = actor_counter;
-  task->function_id = function_id;
-  task->num_args = num_args;
-  task->arg_index = 0;
-  task->num_returns = num_returns;
-  task->args_value_size = args_value_size;
-  for (int i = 0; i < num_returns; ++i) {
-    *task_return_ptr(task, i) = NIL_ID;
+class TaskBuilder {
+ public:
+  void Start(UniqueID driver_id,
+             TaskID parent_task_id,
+             int64_t parent_counter,
+             ActorID actor_id,
+             int64_t actor_counter,
+             FunctionID function_id,
+             int64_t num_returns) {
+    driver_id_ = driver_id;
+    parent_task_id_ = parent_task_id;
+    parent_counter_ = parent_counter;
+    actor_id_ = actor_id;
+    actor_counter_ = actor_counter;
+    function_id_ = function_id;
+    num_returns_ = num_returns;
+
+    /* Compute hashes. */
+    sha256_init(&ctx);
+    sha256_update(&ctx, (BYTE *) &driver_id, sizeof(driver_id));
+    sha256_update(&ctx, (BYTE *) &parent_task_id, sizeof(parent_task_id));
+    sha256_update(&ctx, (BYTE *) &parent_counter, sizeof(parent_counter));
+    sha256_update(&ctx, (BYTE *) &actor_id, sizeof(actor_id));
+    sha256_update(&ctx, (BYTE *) &actor_counter, sizeof(actor_counter));
+    sha256_update(&ctx, (BYTE *) &function_id, sizeof(function_id));
   }
-  return task;
-}
 
-void finish_construct_task_spec(task_spec *spec) {
-  /* Check that all of the arguments were added to the task. */
-  DCHECK(spec->arg_index == spec->num_args);
-  spec->task_id = compute_task_id(spec);
-  /* Set the object IDs for the return values. */
-  for (int64_t i = 0; i < spec->num_returns; ++i) {
-    *task_return_ptr(spec, i) = task_compute_return_id(spec->task_id, i);
+  void NextReferenceArgument(ObjectID object_id) {
+    args.push_back(CreateArg(fbb, to_flatbuf(fbb, object_id)));
+    sha256_update(&ctx, (BYTE *) &object_id, sizeof(object_id));
   }
+
+  void NextValueArgument(uint8_t *value, int64_t length) {
+    auto arg = fbb.CreateString((const char *) value, length);
+    auto empty_id = fbb.CreateString("", 0);
+    args.push_back(CreateArg(fbb, empty_id, arg));
+    sha256_update(&ctx, (BYTE *) &value, length);
+  }
+
+  void SetRequiredResource(int64_t resource_index, double value) {
+    if (resource_index >= resource_vector_.size()) {
+      /* Make sure the resource vector is constructed entry by entry,
+       * in order. */
+      CHECK(resource_index == resource_vector_.size());
+      resource_vector_.resize(resource_index + 1);
+    }
+    resource_vector_[resource_index] = value;
+  }
+
+  uint8_t *Finish(int64_t *size) {
+    /* Add arguments. */
+    auto arguments = fbb.CreateVector(args);
+    /* Update hash. */
+    BYTE buff[DIGEST_SIZE];
+    sha256_final(&ctx, buff);
+    TaskID task_id;
+    CHECK(sizeof(task_id) <= DIGEST_SIZE);
+    memcpy(&task_id, buff, sizeof(task_id));
+    /* Add return object IDs. */
+    std::vector<flatbuffers::Offset<flatbuffers::String>> returns;
+    for (int64_t i = 0; i < num_returns_; i++) {
+      ObjectID return_id = task_compute_return_id(task_id, i);
+      returns.push_back(to_flatbuf(fbb, return_id));
+    }
+    /* Create TaskInfo. */
+    for (int64_t i = resource_vector_.size(); i < ResourceIndex_MAX; ++i) {
+      resource_vector_.push_back(0.0);
+    }
+    auto message = CreateTaskInfo(
+        fbb, to_flatbuf(fbb, driver_id_), to_flatbuf(fbb, task_id),
+        to_flatbuf(fbb, parent_task_id_), parent_counter_,
+        to_flatbuf(fbb, actor_id_), actor_counter_,
+        to_flatbuf(fbb, function_id_), arguments, fbb.CreateVector(returns),
+        fbb.CreateVector(resource_vector_));
+    /* Finish the TaskInfo. */
+    fbb.Finish(message);
+    *size = fbb.GetSize();
+    uint8_t *result = (uint8_t *) malloc(*size);
+    memcpy(result, fbb.GetBufferPointer(), *size);
+    fbb.Clear();
+    args.clear();
+    return result;
+  }
+
+ private:
+  flatbuffers::FlatBufferBuilder fbb;
+  std::vector<flatbuffers::Offset<Arg>> args;
+  SHA256_CTX ctx;
+
+  /* Data for the builder. */
+  UniqueID driver_id_;
+  TaskID parent_task_id_;
+  int64_t parent_counter_;
+  ActorID actor_id_;
+  int64_t actor_counter_;
+  FunctionID function_id_;
+  int64_t num_returns_;
+  std::vector<double> resource_vector_;
+};
+
+TaskBuilder *make_task_builder(void) {
+  return new TaskBuilder();
 }
 
-int64_t task_spec_size(task_spec *spec) {
-  return TASK_SPEC_SIZE(spec->num_args, spec->num_returns,
-                        spec->args_value_size);
+void free_task_builder(TaskBuilder *builder) {
+  delete builder;
 }
 
-FunctionID task_function(task_spec *spec) {
-  /* Check that the task has been constructed. */
-  DCHECK(!TaskID_equal(spec->task_id, NIL_TASK_ID));
-  return spec->function_id;
+bool TaskID_equal(TaskID first_id, TaskID second_id) {
+  return UNIQUE_ID_EQ(first_id, second_id);
 }
 
-ActorID task_spec_actor_id(task_spec *spec) {
-  /* Check that the task has been constructed. */
-  DCHECK(!TaskID_equal(spec->task_id, NIL_TASK_ID));
-  return spec->actor_id;
+bool TaskID_is_nil(TaskID id) {
+  return TaskID_equal(id, NIL_TASK_ID);
 }
 
-int64_t task_spec_actor_counter(task_spec *spec) {
-  /* Check that the task has been constructed. */
-  DCHECK(!TaskID_equal(spec->task_id, NIL_TASK_ID));
-  return spec->actor_counter;
+bool ActorID_equal(ActorID first_id, ActorID second_id) {
+  return UNIQUE_ID_EQ(first_id, second_id);
 }
 
-UniqueID task_spec_driver_id(task_spec *spec) {
-  /* Check that the task has been constructed. */
-  DCHECK(!TaskID_equal(spec->task_id, NIL_TASK_ID));
-  return spec->driver_id;
+bool FunctionID_equal(FunctionID first_id, FunctionID second_id) {
+  return UNIQUE_ID_EQ(first_id, second_id);
 }
 
-TaskID task_spec_id(task_spec *spec) {
-  /* Check that the task has been constructed. */
-  DCHECK(!TaskID_equal(spec->task_id, NIL_TASK_ID));
-  return spec->task_id;
+bool FunctionID_is_nil(FunctionID id) {
+  return FunctionID_equal(id, NIL_FUNCTION_ID);
 }
 
-int64_t task_num_args(task_spec *spec) {
-  return spec->num_args;
+/* Functions for building tasks. */
+
+void TaskSpec_start_construct(TaskBuilder *builder,
+                              UniqueID driver_id,
+                              TaskID parent_task_id,
+                              int64_t parent_counter,
+                              ActorID actor_id,
+                              int64_t actor_counter,
+                              FunctionID function_id,
+                              int64_t num_returns) {
+  builder->Start(driver_id, parent_task_id, parent_counter, actor_id,
+                 actor_counter, function_id, num_returns);
 }
 
-int64_t task_num_returns(task_spec *spec) {
-  return spec->num_returns;
+uint8_t *TaskSpec_finish_construct(TaskBuilder *builder, int64_t *size) {
+  return builder->Finish(size);
 }
 
-int8_t task_arg_type(task_spec *spec, int64_t arg_index) {
-  DCHECK(0 <= arg_index && arg_index < spec->num_args);
-  return spec->args_and_returns[arg_index].type;
+void TaskSpec_args_add_ref(TaskBuilder *builder, ObjectID object_id) {
+  builder->NextReferenceArgument(object_id);
 }
 
-ObjectID task_arg_id(task_spec *spec, int64_t arg_index) {
-  /* Check that the task has been constructed. */
-  DCHECK(!TaskID_equal(spec->task_id, NIL_TASK_ID));
-  DCHECK(0 <= arg_index && arg_index < spec->num_args);
-  task_arg *arg = &spec->args_and_returns[arg_index];
-  DCHECK(arg->type == ARG_BY_REF)
-  return arg->obj_id;
+void TaskSpec_args_add_val(TaskBuilder *builder,
+                           uint8_t *value,
+                           int64_t length) {
+  builder->NextValueArgument(value, length);
 }
 
-uint8_t *task_arg_val(task_spec *spec, int64_t arg_index) {
-  DCHECK(0 <= arg_index && arg_index < spec->num_args);
-  task_arg *arg = &spec->args_and_returns[arg_index];
-  DCHECK(arg->type == ARG_BY_VAL);
-  uint8_t *data = (uint8_t *) &spec->args_and_returns[0];
-  data += (spec->num_args + spec->num_returns) * sizeof(task_arg);
-  return data + arg->value.offset;
+void TaskSpec_set_required_resource(TaskBuilder *builder,
+                                    int64_t resource_index,
+                                    double value) {
+  builder->SetRequiredResource(resource_index, value);
 }
 
-int64_t task_arg_length(task_spec *spec, int64_t arg_index) {
-  DCHECK(0 <= arg_index && arg_index < spec->num_args);
-  task_arg *arg = &spec->args_and_returns[arg_index];
-  DCHECK(arg->type == ARG_BY_VAL);
-  return arg->value.length;
+/* Functions for reading tasks. */
+
+TaskID TaskSpec_task_id(TaskSpec *spec) {
+  CHECK(spec);
+  auto message = flatbuffers::GetRoot<TaskInfo>(spec);
+  return from_flatbuf(message->task_id());
 }
 
-int64_t task_args_add_ref(task_spec *spec, ObjectID obj_id) {
-  /* Check that the task is still under construction. */
-  DCHECK(TaskID_equal(spec->task_id, NIL_TASK_ID));
-  task_arg *arg = &spec->args_and_returns[spec->arg_index];
-  arg->type = ARG_BY_REF;
-  arg->obj_id = obj_id;
-  return spec->arg_index++;
+FunctionID TaskSpec_function(TaskSpec *spec) {
+  CHECK(spec);
+  auto message = flatbuffers::GetRoot<TaskInfo>(spec);
+  return from_flatbuf(message->function_id());
 }
 
-int64_t task_args_add_val(task_spec *spec, uint8_t *data, int64_t length) {
-  /* Check that the task is still under construction. */
-  DCHECK(TaskID_equal(spec->task_id, NIL_TASK_ID));
-  task_arg *arg = &spec->args_and_returns[spec->arg_index];
-  arg->type = ARG_BY_VAL;
-  arg->value.offset = spec->args_value_offset;
-  arg->value.length = length;
-  uint8_t *addr = task_arg_val(spec, spec->arg_index);
-  DCHECK(spec->args_value_offset + length <= spec->args_value_size);
-  DCHECK(spec->arg_index != spec->num_args - 1 ||
-         spec->args_value_offset + length == spec->args_value_size);
-  memcpy(addr, data, length);
-  spec->args_value_offset += length;
-  return spec->arg_index++;
+ActorID TaskSpec_actor_id(TaskSpec *spec) {
+  CHECK(spec);
+  auto message = flatbuffers::GetRoot<TaskInfo>(spec);
+  return from_flatbuf(message->actor_id());
 }
 
-void task_spec_set_required_resource(task_spec *spec,
-                                     int64_t resource_index,
-                                     double value) {
-  spec->required_resources[resource_index] = value;
+int64_t TaskSpec_actor_counter(TaskSpec *spec) {
+  CHECK(spec);
+  auto message = flatbuffers::GetRoot<TaskInfo>(spec);
+  return message->actor_counter();
 }
 
-ObjectID task_return(task_spec *spec, int64_t return_index) {
-  /* Check that the task has been constructed. */
-  DCHECK(!TaskID_equal(spec->task_id, NIL_TASK_ID));
-  DCHECK(0 <= return_index && return_index < spec->num_returns);
-  task_arg *ret = &spec->args_and_returns[spec->num_args + return_index];
-  DCHECK(ret->type == ARG_BY_REF);
-  return ret->obj_id;
+UniqueID TaskSpec_driver_id(TaskSpec *spec) {
+  CHECK(spec);
+  auto message = flatbuffers::GetRoot<TaskInfo>(spec);
+  return from_flatbuf(message->driver_id());
 }
 
-double task_spec_get_required_resource(const task_spec *spec,
-                                       int64_t resource_index) {
-  return spec->required_resources[resource_index];
+int64_t TaskSpec_num_args(TaskSpec *spec) {
+  CHECK(spec);
+  auto message = flatbuffers::GetRoot<TaskInfo>(spec);
+  return message->args()->size();
 }
 
-void free_task_spec(task_spec *spec) {
-  /* Check that the task has been constructed. */
-  DCHECK(!TaskID_equal(spec->task_id, NIL_TASK_ID));
-  DCHECK(spec->arg_index == spec->num_args);
+ObjectID TaskSpec_arg_id(TaskSpec *spec, int64_t arg_index) {
+  CHECK(spec);
+  auto message = flatbuffers::GetRoot<TaskInfo>(spec);
+  return from_flatbuf(message->args()->Get(arg_index)->object_id());
+}
+
+const uint8_t *TaskSpec_arg_val(TaskSpec *spec, int64_t arg_index) {
+  CHECK(spec);
+  auto message = flatbuffers::GetRoot<TaskInfo>(spec);
+  return (uint8_t *) message->args()->Get(arg_index)->data()->c_str();
+}
+
+int64_t TaskSpec_arg_length(TaskSpec *spec, int64_t arg_index) {
+  CHECK(spec);
+  auto message = flatbuffers::GetRoot<TaskInfo>(spec);
+  return message->args()->Get(arg_index)->data()->size();
+}
+
+int64_t TaskSpec_num_returns(TaskSpec *spec) {
+  CHECK(spec);
+  auto message = flatbuffers::GetRoot<TaskInfo>(spec);
+  return message->returns()->size();
+}
+
+bool TaskSpec_arg_by_ref(TaskSpec *spec, int64_t arg_index) {
+  CHECK(spec);
+  auto message = flatbuffers::GetRoot<TaskInfo>(spec);
+  return message->args()->Get(arg_index)->object_id()->size() != 0;
+}
+
+ObjectID TaskSpec_return(TaskSpec *spec, int64_t return_index) {
+  CHECK(spec);
+  auto message = flatbuffers::GetRoot<TaskInfo>(spec);
+  return from_flatbuf(message->returns()->Get(return_index));
+}
+
+double TaskSpec_get_required_resource(const TaskSpec *spec,
+                                      int64_t resource_index) {
+  CHECK(spec);
+  auto message = flatbuffers::GetRoot<TaskInfo>(spec);
+  return message->required_resources()->Get(resource_index);
+}
+
+void TaskSpec_free(TaskSpec *spec) {
   free(spec);
 }
 
-void print_task(task_spec *spec, UT_string *output) {
-  /* For converting an id to hex, which has double the number
-   * of bytes compared to the id (+ 1 byte for '\0'). */
-  static char hex[ID_STRING_SIZE];
-  /* Print function id. */
-  ObjectID_to_string((ObjectID) task_function(spec), &hex[0], ID_STRING_SIZE);
-  utstring_printf(output, "fun %s ", &hex[0]);
-  /* Print arguments. */
-  for (int i = 0; i < task_num_args(spec); ++i) {
-    ObjectID_to_string((ObjectID) task_arg_id(spec, i), &hex[0],
-                       ID_STRING_SIZE);
-    utstring_printf(output, " id:%d %s", i, &hex[0]);
-  }
-  /* Print return ids. */
-  for (int i = 0; i < task_num_returns(spec); ++i) {
-    ObjectID obj_id = task_return(spec, i);
-    ObjectID_to_string(obj_id, &hex[0], ID_STRING_SIZE);
-    utstring_printf(output, " ret:%d %s", i, &hex[0]);
-  }
-}
-
 /* TASK INSTANCES */
 
-struct TaskImpl {
-  /** The scheduling state of the task. */
-  int state;
-  /** The ID of the local scheduler involved. */
-  DBClientID local_scheduler_id;
-  /** The task specification for this task. */
-  task_spec spec;
-};
-
-Task *Task_alloc(task_spec *spec, int state, DBClientID local_scheduler_id) {
-  int64_t size = sizeof(Task) - sizeof(task_spec) + task_spec_size(spec);
+Task *Task_alloc(TaskSpec *spec,
+                 int64_t task_spec_size,
+                 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);
   result->state = state;
   result->local_scheduler_id = local_scheduler_id;
-  memcpy(&result->spec, spec, task_spec_size(spec));
+  result->task_spec_size = task_spec_size;
+  memcpy(&result->spec, spec, task_spec_size);
   return result;
 }
 
@@ -370,7 +304,7 @@ Task *Task_copy(Task *other) {
 }
 
 int64_t Task_size(Task *task_arg) {
-  return sizeof(Task) - sizeof(task_spec) + task_spec_size(&task_arg->spec);
+  return sizeof(Task) - sizeof(TaskSpec) + task_arg->task_spec_size;
 }
 
 int Task_state(Task *task) {
@@ -381,21 +315,25 @@ void Task_set_state(Task *task, int state) {
   task->state = state;
 }
 
-DBClientID Task_local_scheduler_id(Task *task) {
+DBClientID Task_local_scheduler(Task *task) {
   return task->local_scheduler_id;
 }
 
-void Task_set_local_scheduler_id(Task *task, DBClientID local_scheduler_id) {
+void Task_set_local_scheduler(Task *task, DBClientID local_scheduler_id) {
   task->local_scheduler_id = local_scheduler_id;
 }
 
-task_spec *Task_task_spec(Task *task) {
+TaskSpec *Task_task_spec(Task *task) {
   return &task->spec;
 }
 
+int64_t Task_task_spec_size(Task *task) {
+  return task->task_spec_size;
+}
+
 TaskID Task_task_id(Task *task) {
-  task_spec *spec = Task_task_spec(task);
-  return task_spec_id(spec);
+  TaskSpec *spec = Task_task_spec(task);
+  return TaskSpec_task_id(spec);
 }
 
 void Task_free(Task *task) {

src/common/task.h

@@ -1,19 +1,18 @@
 #ifndef TASK_H
 #define TASK_H
 
-/**
- * This API specifies the task data structures. It is in C so we can
- * easily construct tasks from other languages like Python. The data structures
- * are also defined in such a way that memory is contiguous and all pointers
- * are relative, so that we can memcpy the datastructure and ship it over the
- * network without serialization and deserialization. */
-
 #include <stdbool.h>
 #include <stddef.h>
 #include <stdint.h>
 #include "common.h"
 #include "utstring.h"
 
+#include "format/common_generated.h"
+
+typedef uint8_t TaskSpec;
+
+struct TaskBuilder;
+
 #define NIL_TASK_ID NIL_ID
 #define NIL_ACTOR_ID NIL_ID
 #define NIL_FUNCTION_ID NIL_ID
@@ -28,17 +27,6 @@ typedef UniqueID TaskID;
  *  not run on an actor, then NIL_ACTOR_ID should be used. */
 typedef UniqueID ActorID;
 
-/**
- * ==== Task specifications ====
- * Contain all the information neccessary to execute the
- * task (function id, arguments, return object ids).
- */
-
-typedef struct task_spec_impl task_spec;
-
-/** If argument is passed by value or reference. */
-enum arg_type { ARG_BY_REF, ARG_BY_VAL };
-
 /**
  * Compare two task IDs.
  *
@@ -84,6 +72,10 @@ bool FunctionID_is_nil(FunctionID id);
 
 /* Construct and modify task specifications. */
 
+TaskBuilder *make_task_builder(void);
+
+void free_task_builder(TaskBuilder *builder);
+
 /**
  * Begin constructing a task_spec. After this is called, the arguments must be
  * added to the task_spec and then finish_construct_task_spec must be called.
@@ -93,8 +85,6 @@ bool FunctionID_is_nil(FunctionID id);
  * @param parent_task_id The task ID of the task that submitted this task.
  * @param parent_counter A counter indicating how many tasks were submitted by
  *        the parent task prior to this one.
- * @param actor_id The ID of the actor this task belongs to.
- * @param actor_counter Number of tasks that have been executed on this actor.
  * @param function_id The function ID of the function to execute in this task.
  * @param num_args The number of arguments that this task has.
  * @param num_returns The number of return values that this task has.
@@ -102,15 +92,14 @@ bool FunctionID_is_nil(FunctionID id);
           ignoring object ID arguments.
  * @return The partially constructed task_spec.
  */
-task_spec *start_construct_task_spec(UniqueID driver_id,
-                                     TaskID parent_task_id,
-                                     int64_t parent_counter,
-                                     UniqueID actor_id,
-                                     int64_t actor_counter,
-                                     FunctionID function_id,
-                                     int64_t num_args,
-                                     int64_t num_returns,
-                                     int64_t args_value_size);
+void TaskSpec_start_construct(TaskBuilder *B,
+                              UniqueID driver_id,
+                              TaskID parent_task_id,
+                              int64_t parent_counter,
+                              UniqueID actor_id,
+                              int64_t actor_counter,
+                              FunctionID function_id,
+                              int64_t num_returns);
 
 /**
  * Finish constructing a task_spec. This computes the task ID and the object IDs
@@ -120,15 +109,7 @@ task_spec *start_construct_task_spec(UniqueID driver_id,
  * @param spec The task spec whose ID and return object IDs should be computed.
  * @return Void.
  */
-void finish_construct_task_spec(task_spec *spec);
-
-/**
- * The size of the task in bytes.
- *
- * @param spec The task_spec in question.
- * @return The size of the task_spec in bytes.
- */
-int64_t task_spec_size(task_spec *spec);
+uint8_t *TaskSpec_finish_construct(TaskBuilder *builder, int64_t *size);
 
 /**
  * Return the function ID of the task.
@@ -136,7 +117,7 @@ int64_t task_spec_size(task_spec *spec);
  * @param spec The task_spec in question.
  * @return The function ID of the function to execute in this task.
  */
-FunctionID task_function(task_spec *spec);
+FunctionID TaskSpec_function(TaskSpec *spec);
 
 /**
  * Return the actor ID of the task.
@@ -144,7 +125,7 @@ FunctionID task_function(task_spec *spec);
  * @param spec The task_spec in question.
  * @return The actor ID of the actor the task is part of.
  */
-UniqueID task_spec_actor_id(task_spec *spec);
+UniqueID TaskSpec_actor_id(TaskSpec *spec);
 
 /**
  * Return the actor counter of the task. This starts at 0 and increments by 1
@@ -153,7 +134,7 @@ UniqueID task_spec_actor_id(task_spec *spec);
  * @param spec The task_spec in question.
  * @return The actor counter of the task.
  */
-int64_t task_spec_actor_counter(task_spec *spec);
+int64_t TaskSpec_actor_counter(TaskSpec *spec);
 
 /**
  * Return the driver ID of the task.
@@ -161,7 +142,7 @@ int64_t task_spec_actor_counter(task_spec *spec);
  * @param spec The task_spec in question.
  * @return The driver ID of the task.
  */
-UniqueID task_spec_driver_id(task_spec *spec);
+UniqueID TaskSpec_driver_id(TaskSpec *spec);
 
 /**
  * Return the task ID of the task.
@@ -169,7 +150,7 @@ UniqueID task_spec_driver_id(task_spec *spec);
  * @param spec The task_spec in question.
  * @return The task ID of the task.
  */
-TaskID task_spec_id(task_spec *spec);
+TaskID TaskSpec_task_id(TaskSpec *spec);
 
 /**
  * Get the number of arguments to this task.
@@ -177,7 +158,7 @@ TaskID task_spec_id(task_spec *spec);
  * @param spec The task_spec in question.
  * @return The number of arguments to this task.
  */
-int64_t task_num_args(task_spec *spec);
+int64_t TaskSpec_num_args(TaskSpec *spec);
 
 /**
  * Get the number of return values expected from this task.
@@ -185,17 +166,16 @@ int64_t task_num_args(task_spec *spec);
  * @param spec The task_spec in question.
  * @return The number of return values expected from this task.
  */
-int64_t task_num_returns(task_spec *spec);
+int64_t TaskSpec_num_returns(TaskSpec *spec);
 
 /**
- * Get the type of an argument to this task. It should be either ARG_BY_REF or
- * ARG_BY_VAL.
+ * Return true if this argument is passed by reference.
  *
  * @param spec The task_spec in question.
  * @param arg_index The index of the argument in question.
- * @return The type of the argument.
+ * @return True if this argument is passed by reference.
  */
-int8_t task_arg_type(task_spec *spec, int64_t arg_index);
+bool TaskSpec_arg_by_ref(TaskSpec *spec, int64_t arg_index);
 
 /**
  * Get a particular argument to this task. This assumes the argument is an
@@ -205,7 +185,7 @@ int8_t task_arg_type(task_spec *spec, int64_t arg_index);
  * @param arg_index The index of the argument in question.
  * @return The argument at that index.
  */
-ObjectID task_arg_id(task_spec *spec, int64_t arg_index);
+ObjectID TaskSpec_arg_id(TaskSpec *spec, int64_t arg_index);
 
 /**
  * Get a particular argument to this task. This assumes the argument is a value.
@@ -214,7 +194,7 @@ ObjectID task_arg_id(task_spec *spec, int64_t arg_index);
  * @param arg_index The index of the argument in question.
  * @return The argument at that index.
  */
-uint8_t *task_arg_val(task_spec *spec, int64_t arg_index);
+const uint8_t *TaskSpec_arg_val(TaskSpec *spec, int64_t arg_index);
 
 /**
  * Get the number of bytes in a particular argument to this task. This assumes
@@ -224,7 +204,7 @@ uint8_t *task_arg_val(task_spec *spec, int64_t arg_index);
  * @param arg_index The index of the argument in question.
  * @return The number of bytes in the argument.
  */
-int64_t task_arg_length(task_spec *spec, int64_t arg_index);
+int64_t TaskSpec_arg_length(TaskSpec *spec, int64_t arg_index);
 
 /**
  * Set the next task argument. Note that this API only allows you to set the
@@ -235,7 +215,7 @@ int64_t task_arg_length(task_spec *spec, int64_t arg_index);
  * @return The number of task arguments that have been set before this one. This
  *         is only used for testing.
  */
-int64_t task_args_add_ref(task_spec *spec, ObjectID obj_id);
+void TaskSpec_args_add_ref(TaskBuilder *spec, ObjectID obj_id);
 
 /**
  * Set the next task argument. Note that this API only allows you to set the
@@ -247,32 +227,9 @@ int64_t task_args_add_ref(task_spec *spec, ObjectID obj_id);
  * @return The number of task arguments that have been set before this one. This
  *         is only used for testing.
  */
-int64_t task_args_add_val(task_spec *spec, uint8_t *data, int64_t length);
-
-/**
- * Get a particular return object ID of a task.
- *
- * @param spec The task_spec in question.
- * @param return_index The index of the return object ID in question.
- * @return The relevant return object ID.
- */
-ObjectID task_return(task_spec *spec, int64_t return_index);
-
-/**
- * Indices into resource vectors.
- * A resource vector maps a resource index to the number
- * of units of that resource required.
- *
- * The total length of the resource vector is NUM_RESOURCE_INDICES.
- */
-typedef enum {
-  /** Index for number of cpus the task requires. */
-  CPU_RESOURCE_INDEX = 0,
-  /** Index for number of gpus the task requires. */
-  GPU_RESOURCE_INDEX,
-  /** Total number of different resources in the system. */
-  MAX_RESOURCE_INDEX
-} resource_vector_index;
+void TaskSpec_args_add_val(TaskBuilder *builder,
+                           uint8_t *value,
+                           int64_t length);
 
 /**
  * Set the value associated to a resource index.
@@ -283,9 +240,18 @@ typedef enum {
  *        this task needs or a value for an attribute this task requires.
  * @return Void.
  */
-void task_spec_set_required_resource(task_spec *spec,
-                                     int64_t resource_index,
-                                     double value);
+void TaskSpec_set_required_resource(TaskBuilder *builder,
+                                    int64_t resource_index,
+                                    double value);
+
+/**
+ * Get a particular return object ID of a task.
+ *
+ * @param spec The task_spec in question.
+ * @param return_index The index of the return object ID in question.
+ * @return The relevant return object ID.
+ */
+ObjectID TaskSpec_return(TaskSpec *data, int64_t return_index);
 
 /**
  * Get the value associated to a resource index.
@@ -294,8 +260,8 @@ void task_spec_set_required_resource(task_spec *spec,
  * @param resource_index Index of the resource.
  * @return How many of this resource the task needs to execute.
  */
-double task_spec_get_required_resource(const task_spec *spec,
-                                       int64_t resource_index);
+double TaskSpec_get_required_resource(const TaskSpec *spec,
+                                      int64_t resource_index);
 
 /**
  * Compute the object id associated to a put call.
@@ -306,14 +272,6 @@ double task_spec_get_required_resource(const task_spec *spec,
  */
 ObjectID task_compute_put_id(TaskID task_id, int64_t put_index);
 
-/**
- * Free a task_spec.
- *
- * @param The task_spec in question.
- * @return Void.
- */
-void free_task_spec(task_spec *spec);
-
 /**
  * Print the task as a humanly readable string.
  *
@@ -321,7 +279,15 @@ void free_task_spec(task_spec *spec);
  * @param output The buffer to write the string to.
  * @return Void.
  */
-void print_task(task_spec *spec, UT_string *output);
+void TaskSpec_print(TaskSpec *spec, UT_string *output);
+
+/**
+ * Free a task_spec.
+ *
+ * @param The task_spec in question.
+ * @return Void.
+ */
+void TaskSpec_free(TaskSpec *spec);
 
 /**
  * ==== Task ====
@@ -353,7 +319,17 @@ typedef enum {
 /** A task is an execution of a task specification.  It has a state of execution
  * (see scheduling_state) and the ID of the local scheduler it is scheduled on
  * or running on. */
-typedef struct TaskImpl Task;
+
+struct Task {
+  /** The scheduling state of the 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;
+};
 
 /**
  * Allocate a new task. Must be freed with free_task after use.
@@ -363,7 +339,10 @@ typedef struct TaskImpl Task;
  * @param local_scheduler_id The ID of the local scheduler that the task is
  *        scheduled on, if any.
  */
-Task *Task_alloc(task_spec *spec, int state, DBClientID local_scheduler_id);
+Task *Task_alloc(TaskSpec *spec,
+                 int64_t task_spec_size,
+                 int state,
+                 DBClientID local_scheduler_id);
 
 /**
  * Create a copy of the task. Must be freed with free_task after use.
@@ -383,13 +362,15 @@ int Task_state(Task *task);
 void Task_set_state(Task *task, int state);
 
 /** Local scheduler this task has been assigned to or is running on. */
-DBClientID Task_local_scheduler_id(Task *task);
+DBClientID Task_local_scheduler(Task *task);
 
 /** Set the local scheduler ID for this task. */
-void Task_set_local_scheduler_id(Task *task, DBClientID local_scheduler_id);
+void Task_set_local_scheduler(Task *task, DBClientID local_scheduler_id);
 
 /** Task specification of this task. */
-task_spec *Task_task_spec(Task *task);
+TaskSpec *Task_task_spec(Task *task);
+
+int64_t Task_task_spec_size(Task *task);
 
 /** Task ID of this task. */
 TaskID Task_task_id(Task *task);
@@ -397,4 +378,4 @@ TaskID Task_task_id(Task *task);
 /** Free this task datastructure. */
 void Task_free(Task *task);
 
-#endif
+#endif /* TASK_H */

src/common/test/db_tests.cc

@@ -6,6 +6,7 @@
 
 #include "event_loop.h"
 #include "test_common.h"
+#include "example_task.h"
 #include "state/db.h"
 #include "state/object_table.h"
 #include "state/task_table.h"
@@ -16,6 +17,8 @@
 
 SUITE(db_tests);
 
+TaskBuilder *g_task_builder = NULL;
+
 /* Retry 10 times with an 100ms timeout. */
 const int NUM_RETRIES = 10;
 const uint64_t TIMEOUT = 50;
@@ -149,10 +152,11 @@ TEST task_table_test(void) {
       db_connect("127.0.0.1", 6379, "local_scheduler", "127.0.0.1", 0, NULL);
   db_attach(db, loop, false);
   DBClientID local_scheduler_id = globally_unique_id();
-  task_spec *spec = example_task_spec(1, 1);
-  task_table_test_task =
-      Task_alloc(spec, TASK_STATUS_SCHEDULED, local_scheduler_id);
-  free_task_spec(spec);
+  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);
   RetryInfo retry = {
       .num_retries = NUM_RETRIES,
       .timeout = TIMEOUT,
@@ -183,10 +187,13 @@ TEST task_table_all_test(void) {
   DBHandle *db =
       db_connect("127.0.0.1", 6379, "local_scheduler", "127.0.0.1", 0, NULL);
   db_attach(db, loop, false);
-  task_spec *spec = example_task_spec(1, 1);
+  int64_t task_spec_size;
+  TaskSpec *spec = example_task_spec(1, 1, &task_spec_size);
   /* Schedule two tasks on different local local schedulers. */
-  Task *task1 = Task_alloc(spec, TASK_STATUS_SCHEDULED, globally_unique_id());
-  Task *task2 = Task_alloc(spec, TASK_STATUS_SCHEDULED, globally_unique_id());
+  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());
   RetryInfo retry = {
       .num_retries = NUM_RETRIES, .timeout = TIMEOUT, .fail_callback = NULL,
   };
@@ -201,7 +208,7 @@ TEST task_table_all_test(void) {
   event_loop_add_timer(loop, 200, (event_loop_timer_handler) timeout_handler,
                        NULL);
   event_loop_run(loop);
-  free_task_spec(spec);
+  TaskSpec_free(spec);
   db_disconnect(db);
   destroy_outstanding_callbacks(loop);
   event_loop_destroy(loop);
@@ -237,6 +244,7 @@ SUITE(db_tests) {
 GREATEST_MAIN_DEFS();
 
 int main(int argc, char **argv) {
+  g_task_builder = make_task_builder();
   GREATEST_MAIN_BEGIN();
   RUN_SUITE(db_tests);
   GREATEST_MAIN_END();

src/common/test/example_task.h

@@ -0,0 +1,59 @@
+#ifndef EXAMPLE_TASK_H
+#define EXAMPLE_TASK_H
+
+#include "task.h"
+
+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) {
+  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,
+                           NIL_ACTOR_ID, 0, func_id, num_returns);
+  for (int64_t i = 0; i < num_args; ++i) {
+    ObjectID arg_id;
+    if (arg_ids == NULL) {
+      arg_id = globally_unique_id();
+    } else {
+      arg_id = arg_ids[i];
+    }
+    TaskSpec_args_add_ref(g_task_builder, arg_id);
+  }
+  return TaskSpec_finish_construct(g_task_builder, task_spec_size);
+}
+
+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 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);
+  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);
+  return instance;
+}
+
+#endif /* EXAMPLE_TASK_H */

src/common/test/object_table_tests.cc

@@ -1,6 +1,7 @@
 #include "greatest.h"
 
 #include "event_loop.h"
+#include "example_task.h"
 #include "test_common.h"
 #include "common.h"
 #include "state/object_table.h"
@@ -11,6 +12,7 @@
 SUITE(object_table_tests);
 
 static event_loop *g_loop;
+TaskBuilder *g_task_builder = NULL;
 
 /* ==== Test adding and looking up metadata ==== */
 
@@ -18,7 +20,7 @@ int new_object_failed = 0;
 int new_object_succeeded = 0;
 ObjectID new_object_id;
 Task *new_object_task;
-task_spec *new_object_task_spec;
+TaskSpec *new_object_task_spec;
 TaskID new_object_task_id;
 
 void new_object_fail_callback(UniqueID id,
@@ -68,7 +70,7 @@ TEST new_object_test(void) {
   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_id = task_spec_id(new_object_task_spec);
+  new_object_task_id = TaskSpec_task_id(new_object_task_spec);
   g_loop = event_loop_create();
   DBHandle *db =
       db_connect("127.0.0.1", 6379, "plasma_manager", "127.0.0.1", 0, NULL);
@@ -902,6 +904,7 @@ SUITE(object_table_tests) {
 GREATEST_MAIN_DEFS();
 
 int main(int argc, char **argv) {
+  g_task_builder = make_task_builder();
   GREATEST_MAIN_BEGIN();
   RUN_SUITE(object_table_tests);
   GREATEST_MAIN_END();

src/common/test/task_table_tests.cc

@@ -1,6 +1,7 @@
 #include "greatest.h"
 
 #include "event_loop.h"
+#include "example_task.h"
 #include "test_common.h"
 #include "common.h"
 #include "state/object_table.h"
@@ -12,6 +13,7 @@
 SUITE(task_table_tests);
 
 event_loop *g_loop;
+TaskBuilder *g_task_builder = NULL;
 
 /* ==== Test operations in non-failure scenario ==== */
 
@@ -418,6 +420,7 @@ SUITE(task_table_tests) {
 GREATEST_MAIN_DEFS();
 
 int main(int argc, char **argv) {
+  g_task_builder = make_task_builder();
   GREATEST_MAIN_BEGIN();
   RUN_SUITE(task_table_tests);
   GREATEST_MAIN_END();

src/common/test/task_tests.cc

@@ -14,37 +14,39 @@ SUITE(task_tests);
 TEST task_test(void) {
   TaskID parent_task_id = globally_unique_id();
   FunctionID func_id = globally_unique_id();
-  task_spec *spec = start_construct_task_spec(
-      NIL_ID, parent_task_id, 0, NIL_ACTOR_ID, 0, func_id, 4, 2, 10);
-  ASSERT(task_num_args(spec) == 4);
-  ASSERT(task_num_returns(spec) == 2);
+  TaskBuilder *builder = make_task_builder();
+  TaskSpec_start_construct(builder, NIL_ID, parent_task_id, 0, NIL_ACTOR_ID, 0,
+                           func_id, 2);
 
   UniqueID arg1 = globally_unique_id();
-  ASSERT(task_args_add_ref(spec, arg1) == 0);
-  ASSERT(task_args_add_val(spec, (uint8_t *) "hello", 5) == 1);
+  TaskSpec_args_add_ref(builder, arg1);
+  TaskSpec_args_add_val(builder, (uint8_t *) "hello", 5);
   UniqueID arg2 = globally_unique_id();
-  ASSERT(task_args_add_ref(spec, arg2) == 2);
-  ASSERT(task_args_add_val(spec, (uint8_t *) "world", 5) == 3);
+  TaskSpec_args_add_ref(builder, arg2);
+  TaskSpec_args_add_val(builder, (uint8_t *) "world", 5);
   /* Finish constructing the spec. This constructs the task ID and the
    * return IDs. */
-  finish_construct_task_spec(spec);
+  int64_t size;
+  uint8_t *spec = TaskSpec_finish_construct(builder, &size);
 
   /* Check that the spec was constructed as expected. */
-  ASSERT(task_num_args(spec) == 4);
-  ASSERT(task_num_returns(spec) == 2);
-  ASSERT(FunctionID_equal(task_function(spec), func_id));
-  ASSERT(ObjectID_equal(task_arg_id(spec, 0), arg1));
-  ASSERT(memcmp(task_arg_val(spec, 1), (uint8_t *) "hello",
-                task_arg_length(spec, 1)) == 0);
-  ASSERT(ObjectID_equal(task_arg_id(spec, 2), arg2));
-  ASSERT(memcmp(task_arg_val(spec, 3), (uint8_t *) "world",
-                task_arg_length(spec, 3)) == 0);
+  ASSERT(TaskSpec_num_args(spec) == 4);
+  ASSERT(TaskSpec_num_returns(spec) == 2);
+  ASSERT(FunctionID_equal(TaskSpec_function(spec), func_id));
+  ASSERT(ObjectID_equal(TaskSpec_arg_id(spec, 0), arg1));
+  ASSERT(memcmp(TaskSpec_arg_val(spec, 1), (uint8_t *) "hello",
+                TaskSpec_arg_length(spec, 1)) == 0);
+  ASSERT(ObjectID_equal(TaskSpec_arg_id(spec, 2), arg2));
+  ASSERT(memcmp(TaskSpec_arg_val(spec, 3), (uint8_t *) "world",
+                TaskSpec_arg_length(spec, 3)) == 0);
 
-  free_task_spec(spec);
+  TaskSpec_free(spec);
+  free_task_builder(builder);
   PASS();
 }
 
 TEST deterministic_ids_test(void) {
+  TaskBuilder *builder = make_task_builder();
   /* Define the inputs to the task construction. */
   TaskID parent_task_id = globally_unique_id();
   FunctionID func_id = globally_unique_id();
@@ -52,123 +54,132 @@ TEST deterministic_ids_test(void) {
   uint8_t *arg2 = (uint8_t *) "hello world";
 
   /* Construct a first task. */
-  task_spec *spec1 = start_construct_task_spec(
-      NIL_ID, parent_task_id, 0, NIL_ACTOR_ID, 0, func_id, 2, 3, 11);
-  task_args_add_ref(spec1, arg1);
-  task_args_add_val(spec1, arg2, 11);
-  finish_construct_task_spec(spec1);
+  TaskSpec_start_construct(builder, NIL_ID, parent_task_id, 0, NIL_ACTOR_ID, 0,
+                           func_id, 3);
+  TaskSpec_args_add_ref(builder, arg1);
+  TaskSpec_args_add_val(builder, arg2, 11);
+  int64_t size1;
+  TaskSpec *spec1 = TaskSpec_finish_construct(builder, &size1);
 
   /* Construct a second identical task. */
-  task_spec *spec2 = start_construct_task_spec(
-      NIL_ID, parent_task_id, 0, NIL_ACTOR_ID, 0, func_id, 2, 3, 11);
-  task_args_add_ref(spec2, arg1);
-  task_args_add_val(spec2, arg2, 11);
-  finish_construct_task_spec(spec2);
+  TaskSpec_start_construct(builder, NIL_ID, parent_task_id, 0, NIL_ACTOR_ID, 0,
+                           func_id, 3);
+  TaskSpec_args_add_ref(builder, arg1);
+  TaskSpec_args_add_val(builder, arg2, 11);
+  int64_t size2;
+  TaskSpec *spec2 = TaskSpec_finish_construct(builder, &size2);
 
-  /* Check that these tasks have the same task IDs and the same return IDs.*/
-  ASSERT(TaskID_equal(task_spec_id(spec1), task_spec_id(spec2)));
-  ASSERT(ObjectID_equal(task_return(spec1, 0), task_return(spec2, 0)));
-  ASSERT(ObjectID_equal(task_return(spec1, 1), task_return(spec2, 1)));
-  ASSERT(ObjectID_equal(task_return(spec1, 2), task_return(spec2, 2)));
+  /* Check that these tasks have the same task IDs and the same return IDs. */
+  ASSERT(TaskID_equal(TaskSpec_task_id(spec1), TaskSpec_task_id(spec2)));
+  ASSERT(ObjectID_equal(TaskSpec_return(spec1, 0), TaskSpec_return(spec2, 0)));
+  ASSERT(ObjectID_equal(TaskSpec_return(spec1, 1), TaskSpec_return(spec2, 1)));
+  ASSERT(ObjectID_equal(TaskSpec_return(spec1, 2), TaskSpec_return(spec2, 2)));
   /* Check that the return IDs are all distinct. */
-  ASSERT(!ObjectID_equal(task_return(spec1, 0), task_return(spec2, 1)));
-  ASSERT(!ObjectID_equal(task_return(spec1, 0), task_return(spec2, 2)));
-  ASSERT(!ObjectID_equal(task_return(spec1, 1), task_return(spec2, 2)));
+  ASSERT(!ObjectID_equal(TaskSpec_return(spec1, 0), TaskSpec_return(spec2, 1)));
+  ASSERT(!ObjectID_equal(TaskSpec_return(spec1, 0), TaskSpec_return(spec2, 2)));
+  ASSERT(!ObjectID_equal(TaskSpec_return(spec1, 1), TaskSpec_return(spec2, 2)));
 
   /* Create more tasks that are only mildly different. */
 
   /* Construct a task with a different parent task ID. */
-  task_spec *spec3 = start_construct_task_spec(
-      NIL_ID, globally_unique_id(), 0, NIL_ACTOR_ID, 0, func_id, 2, 3, 11);
-  task_args_add_ref(spec3, arg1);
-  task_args_add_val(spec3, arg2, 11);
-  finish_construct_task_spec(spec3);
+  TaskSpec_start_construct(builder, NIL_ID, globally_unique_id(), 0,
+                           NIL_ACTOR_ID, 0, func_id, 3);
+  TaskSpec_args_add_ref(builder, arg1);
+  TaskSpec_args_add_val(builder, arg2, 11);
+  int64_t size3;
+  TaskSpec *spec3 = TaskSpec_finish_construct(builder, &size3);
 
   /* Construct a task with a different parent counter. */
-  task_spec *spec4 = start_construct_task_spec(
-      NIL_ID, parent_task_id, 1, NIL_ACTOR_ID, 0, func_id, 2, 3, 11);
-  task_args_add_ref(spec4, arg1);
-  task_args_add_val(spec4, arg2, 11);
-  finish_construct_task_spec(spec4);
+  TaskSpec_start_construct(builder, NIL_ID, parent_task_id, 1, NIL_ACTOR_ID, 0,
+                           func_id, 3);
+  TaskSpec_args_add_ref(builder, arg1);
+  TaskSpec_args_add_val(builder, arg2, 11);
+  int64_t size4;
+  TaskSpec *spec4 = TaskSpec_finish_construct(builder, &size4);
 
   /* Construct a task with a different function ID. */
-  task_spec *spec5 =
-      start_construct_task_spec(NIL_ID, parent_task_id, 0, NIL_ACTOR_ID, 0,
-                                globally_unique_id(), 2, 3, 11);
-  task_args_add_ref(spec5, arg1);
-  task_args_add_val(spec5, arg2, 11);
-  finish_construct_task_spec(spec5);
+  TaskSpec_start_construct(builder, NIL_ID, parent_task_id, 0, NIL_ACTOR_ID, 0,
+                           globally_unique_id(), 3);
+  TaskSpec_args_add_ref(builder, arg1);
+  TaskSpec_args_add_val(builder, arg2, 11);
+  int64_t size5;
+  TaskSpec *spec5 = TaskSpec_finish_construct(builder, &size5);
 
   /* Construct a task with a different object ID argument. */
-  task_spec *spec6 = start_construct_task_spec(
-      NIL_ID, parent_task_id, 0, NIL_ACTOR_ID, 0, func_id, 2, 3, 11);
-  task_args_add_ref(spec6, globally_unique_id());
-  task_args_add_val(spec6, arg2, 11);
-  finish_construct_task_spec(spec6);
+  TaskSpec_start_construct(builder, NIL_ID, parent_task_id, 0, NIL_ACTOR_ID, 0,
+                           func_id, 3);
+  TaskSpec_args_add_ref(builder, globally_unique_id());
+  TaskSpec_args_add_val(builder, arg2, 11);
+  int64_t size6;
+  TaskSpec *spec6 = TaskSpec_finish_construct(builder, &size6);
 
   /* Construct a task with a different value argument. */
-  task_spec *spec7 = start_construct_task_spec(
-      NIL_ID, parent_task_id, 0, NIL_ACTOR_ID, 0, func_id, 2, 3, 11);
-  task_args_add_ref(spec7, arg1);
-  task_args_add_val(spec7, (uint8_t *) "hello_world", 11);
-  finish_construct_task_spec(spec7);
+  TaskSpec_start_construct(builder, NIL_ID, parent_task_id, 0, NIL_ACTOR_ID, 0,
+                           func_id, 3);
+  TaskSpec_args_add_ref(builder, arg1);
+  TaskSpec_args_add_val(builder, (uint8_t *) "hello_world", 11);
+  int64_t size7;
+  TaskSpec *spec7 = TaskSpec_finish_construct(builder, &size7);
 
   /* Check that the task IDs are all distinct from the original. */
-  ASSERT(!TaskID_equal(task_spec_id(spec1), task_spec_id(spec3)));
-  ASSERT(!TaskID_equal(task_spec_id(spec1), task_spec_id(spec4)));
-  ASSERT(!TaskID_equal(task_spec_id(spec1), task_spec_id(spec5)));
-  ASSERT(!TaskID_equal(task_spec_id(spec1), task_spec_id(spec6)));
-  ASSERT(!TaskID_equal(task_spec_id(spec1), task_spec_id(spec7)));
+  ASSERT(!TaskID_equal(TaskSpec_task_id(spec1), TaskSpec_task_id(spec3)));
+  ASSERT(!TaskID_equal(TaskSpec_task_id(spec1), TaskSpec_task_id(spec4)));
+  ASSERT(!TaskID_equal(TaskSpec_task_id(spec1), TaskSpec_task_id(spec5)));
+  ASSERT(!TaskID_equal(TaskSpec_task_id(spec1), TaskSpec_task_id(spec6)));
+  ASSERT(!TaskID_equal(TaskSpec_task_id(spec1), TaskSpec_task_id(spec7)));
 
   /* Check that the return object IDs are distinct from the originals. */
-  task_spec *specs[6] = {spec1, spec3, spec4, spec5, spec6, spec7};
+  TaskSpec *specs[6] = {spec1, spec3, spec4, spec5, spec6, spec7};
   for (int task_index1 = 0; task_index1 < 6; ++task_index1) {
     for (int return_index1 = 0; return_index1 < 3; ++return_index1) {
       for (int task_index2 = 0; task_index2 < 6; ++task_index2) {
         for (int return_index2 = 0; return_index2 < 3; ++return_index2) {
           if (task_index1 != task_index2 && return_index1 != return_index2) {
             ASSERT(!ObjectID_equal(
-                task_return(specs[task_index1], return_index1),
-                task_return(specs[task_index2], return_index2)));
+                TaskSpec_return(specs[task_index1], return_index1),
+                TaskSpec_return(specs[task_index2], return_index2)));
           }
         }
       }
     }
   }
 
-  free_task_spec(spec1);
-  free_task_spec(spec2);
-  free_task_spec(spec3);
-  free_task_spec(spec4);
-  free_task_spec(spec5);
-  free_task_spec(spec6);
-  free_task_spec(spec7);
+  TaskSpec_free(spec1);
+  TaskSpec_free(spec2);
+  TaskSpec_free(spec3);
+  TaskSpec_free(spec4);
+  TaskSpec_free(spec5);
+  TaskSpec_free(spec6);
+  TaskSpec_free(spec7);
+  free_task_builder(builder);
   PASS();
 }
 
 TEST send_task(void) {
+  TaskBuilder *builder = make_task_builder();
   TaskID parent_task_id = globally_unique_id();
   FunctionID func_id = globally_unique_id();
-  task_spec *spec = start_construct_task_spec(
-      NIL_ID, parent_task_id, 0, NIL_ACTOR_ID, 0, func_id, 4, 2, 10);
-  task_args_add_ref(spec, globally_unique_id());
-  task_args_add_val(spec, (uint8_t *) "Hello", 5);
-  task_args_add_val(spec, (uint8_t *) "World", 5);
-  task_args_add_ref(spec, globally_unique_id());
-  finish_construct_task_spec(spec);
+  TaskSpec_start_construct(builder, NIL_ID, parent_task_id, 0, NIL_ACTOR_ID, 0,
+                           func_id, 2);
+  TaskSpec_args_add_ref(builder, globally_unique_id());
+  TaskSpec_args_add_val(builder, (uint8_t *) "Hello", 5);
+  TaskSpec_args_add_val(builder, (uint8_t *) "World", 5);
+  TaskSpec_args_add_ref(builder, globally_unique_id());
+  int64_t size;
+  TaskSpec *spec = TaskSpec_finish_construct(builder, &size);
   int fd[2];
   socketpair(AF_UNIX, SOCK_STREAM, 0, fd);
-  write_message(fd[0], SUBMIT_TASK, task_spec_size(spec), (uint8_t *) spec);
+  write_message(fd[0], SUBMIT_TASK, size, (uint8_t *) spec);
   int64_t type;
   int64_t length;
   uint8_t *message;
   read_message(fd[1], &type, &length, &message);
-  task_spec *result = (task_spec *) message;
+  TaskSpec *result = (TaskSpec *) message;
   ASSERT(type == SUBMIT_TASK);
-  ASSERT(memcmp(spec, result, task_spec_size(spec)) == 0);
-  ASSERT(memcmp(spec, result, task_spec_size(result)) == 0);
-  free(spec);
+  ASSERT(memcmp(spec, result, size) == 0);
+  TaskSpec_free(spec);
   free(result);
+  free_task_builder(builder);
   PASS();
 }
 

src/common/test/test_common.h

@@ -3,64 +3,16 @@
 
 #include <unistd.h>
 
+#include "common.h"
 #include "io.h"
 #include "hiredis/hiredis.h"
 #include "utstring.h"
 
-#include "task.h"
-
 #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,
-                                                     ObjectID arg_ids[]) {
-  TaskID parent_task_id = globally_unique_id();
-  FunctionID func_id = globally_unique_id();
-  task_spec *task =
-      start_construct_task_spec(NIL_ID, parent_task_id, 0, NIL_ACTOR_ID, 0,
-                                func_id, num_args, num_returns, arg_value_size);
-  for (int64_t i = 0; i < num_args; ++i) {
-    ObjectID 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 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,
-                                           int Task_state,
-                                           ObjectID arg_ids[]) {
-  task_spec *spec = example_task_spec_with_args(num_args, num_returns, arg_ids);
-  Task *instance = Task_alloc(spec, Task_state, NIL_ID);
-  free_task_spec(spec);
-  return instance;
-}
-
-static inline Task *example_task(int64_t num_args,
-                                 int64_t num_returns,
-                                 int Task_state) {
-  task_spec *spec = example_task_spec(num_args, num_returns);
-  Task *instance = Task_alloc(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) {