Reconstruction for evicted objects (#181)

* First pass at reconstruction in the worker

Modify reconstruction stress testing to start Plasma service before rest of Ray cluster

TODO about reconstructing ray.puts

Fix ray.put error for double creates

Distinguish between empty entry and no entry in object table

Fix test case

Fix Python test

Fix tests

* Only call reconstruct on objects we have not yet received

* Address review comments

* Fix reconstruction for Python3

* remove unused code

* Address Robert's comments, stress tests are crashing

* Test and update the task's scheduling state to suppress duplicate
reconstruction requests.

* Split result table into two lookups, one for task ID and the other as a
test-and-set for the task state

* Fix object table tests

* Fix redis module result_table_lookup test case

* Multinode reconstruction tests

* Fix python3 test case

* rename

* Use new start_redis

* Remove unused code

* lint

* indent

* Address Robert's comments

* Use start_redis from ray.services in state table tests

* Remove unnecessary memset

src/common/lib/python/common_extension.c

@@ -35,6 +35,16 @@ void init_pickle_module(void) {
 
 /* Define the PyObjectID class. */
 
+int PyStringToUniqueID(PyObject *object, object_id *object_id) {
+  if (PyBytes_Check(object)) {
+    memcpy(&object_id->id[0], PyBytes_AsString(object), UNIQUE_ID_SIZE);
+    return 1;
+  } else {
+    PyErr_SetString(PyExc_TypeError, "must be a 20 character string");
+    return 0;
+  }
+}
+
 int PyObjectToUniqueID(PyObject *object, object_id *objectid) {
   if (PyObject_IsInstance(object, (PyObject *) &PyObjectIDType)) {
     *objectid = ((PyObjectID *) object)->object_id;

src/common/lib/python/common_extension.h

@@ -33,6 +33,8 @@ extern PyObject *pickle_loads;
 
 void init_pickle_module(void);
 
+int PyStringToUniqueID(PyObject *object, object_id *object_id);
+
 int PyObjectToUniqueID(PyObject *object, object_id *objectid);
 
 PyObject *PyObjectID_make(object_id object_id);

src/common/redis_module/ray_redis_module.c

@@ -188,8 +188,9 @@ int GetClientAddress_RedisCommand(RedisModuleCtx *ctx,
  *     RAY.OBJECT_TABLE_LOOKUP <object id>
  *
  * @param object_id A string representing the object ID.
- * @return A list of plasma manager IDs that are listed in the object table as
- *         having the object.
+ * @return A list, possibly empty, of plasma manager IDs that are listed in the
+ *         object table as having the object. If there was no entry found in
+ *         the object table, returns nil.
  */
 int ObjectTableLookup_RedisCommand(RedisModuleCtx *ctx,
                                    RedisModuleString **argv,
@@ -201,8 +202,12 @@ int ObjectTableLookup_RedisCommand(RedisModuleCtx *ctx,
   RedisModuleKey *key =
       OpenPrefixedKey(ctx, OBJECT_LOCATION_PREFIX, argv[1], REDISMODULE_READ);
 
-  if (RedisModule_KeyType(key) == REDISMODULE_KEYTYPE_EMPTY ||
-      RedisModule_ValueLength(key) == 0) {
+  if (RedisModule_KeyType(key) == REDISMODULE_KEYTYPE_EMPTY) {
+    /* Return nil if no entry was found. */
+    return RedisModule_ReplyWithNull(ctx);
+  }
+  if (RedisModule_ValueLength(key) == 0) {
+    /* Return empty list if there are no managers. */
     return RedisModule_ReplyWithArray(ctx, 0);
   }
 
@@ -581,6 +586,35 @@ int ResultTableAdd_RedisCommand(RedisModuleCtx *ctx,
   return REDISMODULE_OK;
 }
 
+int ParseTaskState(RedisModuleString *state) {
+  size_t state_length;
+  const char *state_string = RedisModule_StringPtrLen(state, &state_length);
+  int state_integer;
+  int scanned = sscanf(state_string, "%2d", &state_integer);
+  if (scanned != 1 || state_length != 2) {
+    return -1;
+  }
+  return state_integer;
+}
+
+RedisModuleString *NormalizeTaskState(RedisModuleCtx *ctx,
+                                      RedisModuleString *state) {
+  /* Pad the state integer to a fixed-width integer, and make sure it has width
+   * less than or equal to 2. */
+  long long state_integer;
+  int status = RedisModule_StringToLongLong(state, &state_integer);
+  if (status != REDISMODULE_OK) {
+    return NULL;
+  }
+  state = RedisModule_CreateStringPrintf(ctx, "%2d", state_integer);
+  size_t length;
+  RedisModule_StringPtrLen(state, &length);
+  if (length != 2) {
+    return NULL;
+  }
+  return state;
+}
+
 /**
  * Reply with information about a task ID. This is used by
  * RAY.RESULT_TABLE_LOOKUP and RAY.TASK_TABLE_GET.
@@ -609,11 +643,8 @@ int ReplyWithTask(RedisModuleCtx *ctx, RedisModuleString *task_id) {
           ctx, "Missing fields in the task table entry");
     }
 
-    size_t state_length;
-    const char *state_string = RedisModule_StringPtrLen(state, &state_length);
-    int state_integer;
-    int scanned = sscanf(state_string, "%2d", &state_integer);
-    if (scanned != 1 || state_length != 2) {
+    int state_integer = ParseTaskState(state);
+    if (state_integer < 0) {
       RedisModule_CloseKey(key);
       RedisModule_FreeString(ctx, state);
       RedisModule_FreeString(ctx, local_scheduler_id);
@@ -668,23 +699,21 @@ int ResultTableLookup_RedisCommand(RedisModuleCtx *ctx,
   key = OpenPrefixedKey(ctx, OBJECT_INFO_PREFIX, object_id, REDISMODULE_READ);
 
   if (RedisModule_KeyType(key) == REDISMODULE_KEYTYPE_EMPTY) {
+    RedisModule_CloseKey(key);
     return RedisModule_ReplyWithNull(ctx);
   }
 
   RedisModuleString *task_id;
   RedisModule_HashGet(key, REDISMODULE_HASH_CFIELDS, "task", &task_id, NULL);
+  RedisModule_CloseKey(key);
   if (task_id == NULL) {
     return RedisModule_ReplyWithNull(ctx);
   }
 
-  /* Construct a reply by getting the task from the task ID. */
-  int status = ReplyWithTask(ctx, task_id);
-
-  /* Clean up. */
+  RedisModule_ReplyWithString(ctx, task_id);
   RedisModule_FreeString(ctx, task_id);
-  RedisModule_CloseKey(key);
 
-  return status;
+  return REDISMODULE_OK;
 }
 
 int TaskTableWrite(RedisModuleCtx *ctx,
@@ -694,18 +723,11 @@ int TaskTableWrite(RedisModuleCtx *ctx,
                    RedisModuleString *task_spec) {
   /* Pad the state integer to a fixed-width integer, and make sure it has width
    * less than or equal to 2. */
-  long long state_integer;
-  int status = RedisModule_StringToLongLong(state, &state_integer);
-  if (status != REDISMODULE_OK) {
+  state = NormalizeTaskState(ctx, state);
+  if (state == NULL) {
     return RedisModule_ReplyWithError(
-        ctx, "Invalid scheduling state (must be an integer)");
-  }
-  state = RedisModule_CreateStringPrintf(ctx, "%2d", state_integer);
-  size_t length;
-  RedisModule_StringPtrLen(state, &length);
-  if (length != 2) {
-    return RedisModule_ReplyWithError(
-        ctx, "Invalid scheduling state width (must have width 2)");
+        ctx,
+        "Invalid scheduling state (must be an integer of width at most 2)");
   }
 
   /* Add the task to the task table. If no spec was provided, get the existing
@@ -720,6 +742,7 @@ int TaskTableWrite(RedisModuleCtx *ctx,
                         &existing_task_spec, NULL);
     if (existing_task_spec == NULL) {
       RedisModule_CloseKey(key);
+      RedisModule_FreeString(ctx, state);
       return RedisModule_ReplyWithError(
           ctx, "Cannot update a task that doesn't exist yet");
     }
@@ -743,6 +766,7 @@ int TaskTableWrite(RedisModuleCtx *ctx,
     publish_message = RedisString_Format(ctx, "%S %S %S %S", task_id, state,
                                          node_id, existing_task_spec);
   }
+  RedisModule_FreeString(ctx, state);
 
   RedisModuleCallReply *reply =
       RedisModule_Call(ctx, "PUBLISH", "ss", publish_topic, publish_message);
@@ -820,6 +844,93 @@ int TaskTableUpdate_RedisCommand(RedisModuleCtx *ctx,
   return TaskTableWrite(ctx, argv[1], argv[2], argv[3], NULL);
 }
 
+/**
+ * Test and update an entry in the task table if the current value matches the
+ * test value. This does not update the task specification in the table.
+ *
+ * This is called from a client with the command:
+ *
+ *     RAY.TASK_TABLE_TEST_AND_UPDATE <task ID> <test state> <state>
+ *         <local scheduler ID>
+ *
+ * @param task_id A string that is the ID of the task.
+ * @param test_state A string that is the test value for the scheduling state.
+ *        The update happens if and only if the current scheduling state
+ *        matches this value.
+ * @param state A string that is the scheduling state (a scheduling_state enum
+ *        instance) to update the task entry with. The string's value must be a
+ *        nonnegative integer less than 100, so that it has width at most 2. If
+ *        less than 2, the value will be left-padded with spaces to a width of
+ *        2.
+ * @param ray_client_id A string that is the ray client ID of the associated
+ *        local scheduler, if any, to update the task entry with.
+ * @return If the current scheduling state does not match the test value,
+ *         returns nil. Else, returns the same as RAY.TASK_TABLE_GET: an array
+ *         of strings representing the updated task fields in the following
+ *         order: 1) (integer) scheduling state 2) (string) associated node ID,
+ *         if any 3) (string) the task specification, which can be casted to a
+ *         task_spec.
+ */
+int TaskTableTestAndUpdate_RedisCommand(RedisModuleCtx *ctx,
+                                        RedisModuleString **argv,
+                                        int argc) {
+  if (argc != 5) {
+    return RedisModule_WrongArity(ctx);
+  }
+
+  RedisModuleString *state = NormalizeTaskState(ctx, argv[3]);
+  if (state == NULL) {
+    return RedisModule_ReplyWithError(
+        ctx,
+        "Invalid scheduling state (must be an integer of width at most 2)");
+  }
+
+  RedisModuleKey *key = OpenPrefixedKey(ctx, TASK_PREFIX, argv[1],
+                                        REDISMODULE_READ | REDISMODULE_WRITE);
+  if (RedisModule_KeyType(key) == REDISMODULE_KEYTYPE_EMPTY) {
+    RedisModule_CloseKey(key);
+    RedisModule_FreeString(ctx, state);
+    return RedisModule_ReplyWithNull(ctx);
+  }
+
+  /* If the key exists, look up the fields and return them in an array. */
+  RedisModuleString *current_state = NULL;
+  RedisModule_HashGet(key, REDISMODULE_HASH_CFIELDS, "state", &current_state,
+                      NULL);
+  int current_state_integer = ParseTaskState(current_state);
+  if (current_state_integer < 0) {
+    RedisModule_CloseKey(key);
+    RedisModule_FreeString(ctx, state);
+    return RedisModule_ReplyWithError(ctx,
+                                      "Found invalid scheduling state (must "
+                                      "be an integer of width 2");
+  }
+  long long test_state_integer;
+  int status = RedisModule_StringToLongLong(argv[2], &test_state_integer);
+  if (status != REDISMODULE_OK) {
+    RedisModule_CloseKey(key);
+    RedisModule_FreeString(ctx, state);
+    return RedisModule_ReplyWithError(
+        ctx, "Invalid test value for scheduling state");
+  }
+  if (current_state_integer != test_state_integer) {
+    /* The current value does not match the test value, so do not perform the
+     * update. */
+    RedisModule_CloseKey(key);
+    RedisModule_FreeString(ctx, state);
+    return RedisModule_ReplyWithNull(ctx);
+  }
+
+  /* The test passed, so perform the update. */
+  RedisModule_HashSet(key, REDISMODULE_HASH_CFIELDS, "state", state, "node",
+                      argv[4], NULL);
+  /* Clean up. */
+  RedisModule_CloseKey(key);
+  RedisModule_FreeString(ctx, state);
+  /* Construct a reply by getting the task from the task ID. */
+  return ReplyWithTask(ctx, argv[1]);
+}
+
 /**
  * Get an entry from the task table.
  *
@@ -922,6 +1033,12 @@ int RedisModule_OnLoad(RedisModuleCtx *ctx,
     return REDISMODULE_ERR;
   }
 
+  if (RedisModule_CreateCommand(ctx, "ray.task_table_test_and_update",
+                                TaskTableTestAndUpdate_RedisCommand,
+                                "write pubsub", 0, 0, 0) == REDISMODULE_ERR) {
+    return REDISMODULE_ERR;
+  }
+
   if (RedisModule_CreateCommand(ctx, "ray.task_table_get",
                                 TaskTableGet_RedisCommand, "readonly", 0, 0,
                                 0) == REDISMODULE_ERR) {

src/common/state/object_table.h

@@ -11,7 +11,9 @@
  */
 
 /* Callback called when the lookup completes. The callback should free
- * the manager_vector array, but NOT the strings they are pointing to.
+ * the manager_vector array, but NOT the strings they are pointing to. If there
+ * was no entry at all for the object (the object had never been created
+ * before), then manager_count will be -1.
  */
 typedef void (*object_table_lookup_done_callback)(
     object_id object_id,
@@ -234,11 +236,12 @@ void result_table_add(db_handle *db_handle,
 
 /** Callback called when the result table lookup completes. */
 typedef void (*result_table_lookup_callback)(object_id object_id,
-                                             task *task,
+                                             task_id task_id,
                                              void *user_context);
 
 /**
- * Lookup the task that created an object in the result table.
+ * Lookup the task that created an object in the result table. The return value
+ * is the task ID.
  *
  * @param db_handle Handle to object_table database.
  * @param object_id ID of the object to lookup.

src/common/state/redis.c

@@ -401,16 +401,20 @@ void redis_result_table_lookup_callback(redisAsyncContext *c,
                                         void *privdata) {
   REDIS_CALLBACK_HEADER(db, callback_data, r);
   redisReply *reply = r;
+  CHECKM(reply->type == REDIS_REPLY_NIL || reply->type == REDIS_REPLY_STRING,
+         "Unexpected reply type %d in redis_result_table_lookup_callback",
+         reply->type);
   /* Parse the task from the reply. */
-  task *task = parse_and_construct_task_from_redis_reply(reply);
+  task_id result_id = NIL_TASK_ID;
+  if (reply->type == REDIS_REPLY_STRING) {
+    CHECK(reply->len == sizeof(result_id));
+    memcpy(&result_id, reply->str, reply->len);
+  }
+
   /* Call the done callback if there is one. */
   result_table_lookup_callback done_callback = callback_data->done_callback;
   if (done_callback != NULL) {
-    done_callback(callback_data->id, task, callback_data->user_context);
-  }
-  /* Free the task if it is not NULL. */
-  if (task != NULL) {
-    free_task(task);
+    done_callback(callback_data->id, result_id, callback_data->user_context);
   }
   /* Clean up timer and callback. */
   destroy_timer_callback(db->loop, callback_data);
@@ -465,24 +469,33 @@ void redis_object_table_lookup_callback(redisAsyncContext *c,
                                         void *privdata) {
   REDIS_CALLBACK_HEADER(db, callback_data, r);
   redisReply *reply = r;
+  LOG_DEBUG("Object table lookup callback");
+  CHECK(reply->type == REDIS_REPLY_NIL || reply->type == REDIS_REPLY_ARRAY);
 
   object_id obj_id = callback_data->id;
-
-  LOG_DEBUG("Object table lookup callback");
-  CHECK(reply->type == REDIS_REPLY_ARRAY);
-
-  int64_t manager_count = reply->elements;
+  int64_t manager_count = 0;
   db_client_id *managers = NULL;
   const char **manager_vector = NULL;
-  if (manager_count > 0) {
-    managers = malloc(reply->elements * sizeof(db_client_id));
-    manager_vector = malloc(manager_count * sizeof(char *));
-  }
-  for (int j = 0; j < reply->elements; ++j) {
-    CHECK(reply->element[j]->type == REDIS_REPLY_STRING);
-    memcpy(managers[j].id, reply->element[j]->str, sizeof(managers[j].id));
-    redis_get_cached_db_client(db, managers[j], manager_vector + j);
+
+  /* Parse the Redis reply. */
+  if (reply->type == REDIS_REPLY_NIL) {
+    /* The object entry did not exist. */
+    manager_count = -1;
+  } else if (reply->type == REDIS_REPLY_ARRAY) {
+    manager_count = reply->elements;
+    if (manager_count > 0) {
+      managers = malloc(reply->elements * sizeof(db_client_id));
+      manager_vector = malloc(manager_count * sizeof(char *));
+    }
+    for (int j = 0; j < reply->elements; ++j) {
+      CHECK(reply->element[j]->type == REDIS_REPLY_STRING);
+      memcpy(managers[j].id, reply->element[j]->str, sizeof(managers[j].id));
+      redis_get_cached_db_client(db, managers[j], manager_vector + j);
+    }
+  } else {
+    LOG_FATAL("Unexpected reply type from object table lookup.");
   }
+
   object_table_lookup_done_callback done_callback =
       callback_data->done_callback;
   if (done_callback) {
@@ -821,6 +834,43 @@ void redis_task_table_update(table_callback_data *callback_data) {
   }
 }
 
+void redis_task_table_test_and_update_callback(redisAsyncContext *c,
+                                               void *r,
+                                               void *privdata) {
+  REDIS_CALLBACK_HEADER(db, callback_data, r);
+  redisReply *reply = r;
+  /* Parse the task from the reply. */
+  task *task = parse_and_construct_task_from_redis_reply(reply);
+  /* Call the done callback if there is one. */
+  task_table_get_callback done_callback = callback_data->done_callback;
+  if (done_callback != NULL) {
+    done_callback(task, callback_data->user_context);
+  }
+  /* Free the task if it is not NULL. */
+  if (task != NULL) {
+    free_task(task);
+  }
+  /* Clean up timer and callback. */
+  destroy_timer_callback(db->loop, callback_data);
+}
+
+void redis_task_table_test_and_update(table_callback_data *callback_data) {
+  db_handle *db = callback_data->db_handle;
+  task_id task_id = callback_data->id;
+  task_table_test_and_update_data *update_data = callback_data->data;
+
+  int status = redisAsyncCommand(
+      db->context, redis_task_table_test_and_update_callback,
+      (void *) callback_data->timer_id,
+      "RAY.TASK_TABLE_TEST_AND_UPDATE %b %d %d %b", task_id.id,
+      sizeof(task_id.id), update_data->test_state, update_data->update_state,
+      update_data->local_scheduler_id.id,
+      sizeof(update_data->local_scheduler_id.id));
+  if ((status == REDIS_ERR) || db->context->err) {
+    LOG_REDIS_DEBUG(db->context, "error in redis_task_table_test_and_update");
+  }
+}
+
 /* The format of the payload is described in ray_redis_module.c and is
  * "<task ID> <state> <local scheduler ID> <task specification>". TODO(rkn):
  * Make this code nicer. */

src/common/state/redis.h

@@ -121,8 +121,7 @@ void redis_object_table_request_notifications(
 void redis_result_table_add(table_callback_data *callback_data);
 
 /**
- * Lookup the object in the object table in redis. The entry in
- * the object table contains metadata about the object.
+ * Lookup the task that created the object in redis. The result is the task ID.
  *
  * @param callback_data Data structure containing redis connection and timeout
  *        information.
@@ -176,6 +175,16 @@ void redis_task_table_add_task(table_callback_data *callback_data);
  */
 void redis_task_table_update(table_callback_data *callback_data);
 
+/**
+ * Update a task table entry with the task's scheduling information, if the
+ * task's current scheduling information matches the test value.
+ *
+ * @param callback_data Data structure containing redis connection and timeout
+ *        information.
+ * @return Void.
+ */
+void redis_task_table_test_and_update(table_callback_data *callback_data);
+
 /**
  * Callback invoked when the reply from the task push command is received.
  *

src/common/state/task_table.c

@@ -30,6 +30,24 @@ void task_table_update(db_handle *db_handle,
                       done_callback, redis_task_table_update, user_context);
 }
 
+void task_table_test_and_update(db_handle *db_handle,
+                                task_id task_id,
+                                scheduling_state test_state,
+                                scheduling_state update_state,
+                                retry_info *retry,
+                                task_table_get_callback done_callback,
+                                void *user_context) {
+  task_table_test_and_update_data *update_data =
+      malloc(sizeof(task_table_test_and_update_data));
+  update_data->test_state = test_state;
+  update_data->update_state = update_state;
+  /* Update the task entry's local scheduler with this client's ID. */
+  update_data->local_scheduler_id = db_handle->client;
+  init_table_callback(db_handle, task_id, __func__, update_data, retry,
+                      done_callback, redis_task_table_test_and_update,
+                      user_context);
+}
+
 /* TODO(swang): A corresponding task_table_unsubscribe. */
 void task_table_subscribe(db_handle *db_handle,
                           db_client_id local_scheduler_id,

src/common/state/task_table.h

@@ -87,6 +87,40 @@ void task_table_update(db_handle *db_handle,
                        task_table_done_callback done_callback,
                        void *user_context);
 
+/**
+ * Update a task's scheduling information in the task table, if the current
+ * value matches the given test value. If the update succeeds, it also updates
+ * the task entry's local scheduler ID with the ID of the client who called
+ * this function. This assumes that the task spec already exists in the task
+ * table entry.
+ *
+ * @param db_handle Database handle.
+ * @param task_id The task ID of the task entry to update.
+ * @param test_state The value to test the current task entry's scheduling
+ *        state against.
+ * @param update_state The value to update the task entry's scheduling state
+ *        with, if the current state matches test_state.
+ * @param retry Information about retrying the request to the database.
+ * @param done_callback Function to be called when database returns result.
+ * @param user_context Data that will be passed to done_callback and
+ *        fail_callback.
+ * @return Void.
+ */
+void task_table_test_and_update(db_handle *db_handle,
+                                task_id task_id,
+                                scheduling_state test_state,
+                                scheduling_state update_state,
+                                retry_info *retry,
+                                task_table_get_callback done_callback,
+                                void *user_context);
+
+/* Data that is needed to test and set the task's scheduling state. */
+typedef struct {
+  scheduling_state test_state;
+  scheduling_state update_state;
+  db_client_id local_scheduler_id;
+} task_table_test_and_update_data;
+
 /*
  *  ==== Subscribing to the task table ====
  */

src/common/task.c

@@ -173,14 +173,6 @@ void finish_construct_task_spec(task_spec *spec) {
   }
 }
 
-task_spec *alloc_nil_task_spec(task_id task_id) {
-  task_spec *spec =
-      start_construct_task_spec(NIL_ID, NIL_ID, 0, NIL_FUNCTION_ID, 0, 0, 0);
-  finish_construct_task_spec(spec);
-  spec->task_id = task_id;
-  return spec;
-}
-
 int64_t task_spec_size(task_spec *spec) {
   return TASK_SPEC_SIZE(spec->num_args, spec->num_returns,
                         spec->args_value_size);
@@ -332,13 +324,6 @@ task *copy_task(task *other) {
   return copy;
 }
 
-task *alloc_nil_task(task_id task_id) {
-  task_spec *nil_spec = alloc_nil_task_spec(task_id);
-  task *nil_task = alloc_task(nil_spec, 0, NIL_ID);
-  free_task_spec(nil_spec);
-  return nil_task;
-}
-
 int64_t task_size(task *task_arg) {
   return sizeof(task) - sizeof(task_spec) + task_spec_size(&task_arg->spec);
 }

src/common/task.h

@@ -273,7 +273,9 @@ typedef enum {
   /** The task is running on a worker. */
   TASK_STATUS_RUNNING = 8,
   /** The task is done executing. */
-  TASK_STATUS_DONE = 16
+  TASK_STATUS_DONE = 16,
+  /** The task will be submitted for reexecution. */
+  TASK_STATUS_RECONSTRUCTING = 32
 } scheduling_state;
 
 /** A task is an execution of a task specification.  It has a state of execution
@@ -325,14 +327,4 @@ task_id task_task_id(task *task);
 /** Free this task datastructure. */
 void free_task(task *task);
 
-/**
- * ==== Task update ====
- * Contains the information necessary to update a task in the task log.
- */
-
-typedef struct {
-  scheduling_state state;
-  db_client_id local_scheduler_id;
-} task_update;
-
 #endif

src/common/test/object_table_tests.c

@@ -31,12 +31,11 @@ void new_object_fail_callback(unique_id id,
 /* === Test adding an object with an associated task === */
 
 void new_object_done_callback(object_id object_id,
-                              task *task,
+                              task_id task_id,
                               void *user_context) {
   new_object_succeeded = 1;
   CHECK(object_ids_equal(object_id, new_object_id));
-  CHECK(task);
-  CHECK(memcmp(task, new_object_task, task_size(task)) == 0);
+  CHECK(task_ids_equal(task_id, new_object_task_id));
   event_loop_stop(g_loop);
 }
 
@@ -92,26 +91,14 @@ TEST new_object_test(void) {
 
 /* === Test adding an object without an associated task === */
 
-void new_object_no_task_lookup_callback(object_id object_id,
-                                        task *task,
-                                        void *user_context) {
+void new_object_no_task_callback(object_id object_id,
+                                 task_id task_id,
+                                 void *user_context) {
   new_object_succeeded = 1;
-  CHECK(task == NULL);
+  CHECK(IS_NIL_ID(task_id));
   event_loop_stop(g_loop);
 }
 
-void new_object_no_task_callback(object_id object_id, void *user_context) {
-  CHECK(object_ids_equal(object_id, new_object_id));
-  retry_info retry = {
-      .num_retries = 5,
-      .timeout = 100,
-      .fail_callback = new_object_fail_callback,
-  };
-  db_handle *db = user_context;
-  result_table_lookup(db, object_id, &retry, new_object_no_task_lookup_callback,
-                      NULL);
-}
-
 TEST new_object_no_task_test(void) {
   new_object_failed = 0;
   new_object_succeeded = 0;
@@ -126,8 +113,8 @@ TEST new_object_no_task_test(void) {
       .timeout = 100,
       .fail_callback = new_object_fail_callback,
   };
-  result_table_add(db, new_object_id, new_object_task_id, &retry,
-                   new_object_no_task_callback, db);
+  result_table_lookup(db, new_object_id, &retry, new_object_no_task_callback,
+                      NULL);
   event_loop_run(g_loop);
   db_disconnect(db);
   destroy_outstanding_callbacks(g_loop);