diff --git a/src/common/state/object_table.c b/src/common/state/object_table.c index 35a28a012..0fc327136 100644 --- a/src/common/state/object_table.c +++ b/src/common/state/object_table.c @@ -1,5 +1,6 @@ #include "object_table.h" #include "redis.h" +#include "object_info.h" void object_table_lookup(db_handle *db_handle, object_id object_id, @@ -13,14 +14,17 @@ void object_table_lookup(db_handle *db_handle, void object_table_add(db_handle *db_handle, object_id object_id, + int64_t data_size, unsigned char digest[], retry_info *retry, object_table_done_callback done_callback, void *user_context) { CHECK(db_handle != NULL); - unsigned char *digest_copy = malloc(DIGEST_SIZE); - memcpy(digest_copy, digest, DIGEST_SIZE); - init_table_callback(db_handle, object_id, __func__, digest_copy, retry, + + object_info *info = malloc(sizeof(object_info)); + info->data_size = data_size; + memcpy(&info->digest[0], digest, DIGEST_SIZE); + init_table_callback(db_handle, object_id, __func__, info, retry, done_callback, redis_object_table_add, user_context); } @@ -43,6 +47,21 @@ void object_table_subscribe( user_context); } +void object_info_subscribe(db_handle *db_handle, + object_info_subscribe_callback subscribe_callback, + void *subscribe_context, + retry_info *retry, + object_info_done_callback done_callback, + void *user_context) { + object_info_subscribe_data *sub_data = + malloc(sizeof(object_info_subscribe_data)); + sub_data->subscribe_callback = subscribe_callback; + sub_data->subscribe_context = subscribe_context; + + init_table_callback(db_handle, NIL_OBJECT_ID, __func__, sub_data, retry, + done_callback, redis_object_info_subscribe, user_context); +} + void result_table_add(db_handle *db_handle, object_id object_id, task_id task_id_arg, diff --git a/src/common/state/object_table.h b/src/common/state/object_table.h index 941b0851e..09d10fc2e 100644 --- a/src/common/state/object_table.h +++ b/src/common/state/object_table.h @@ -49,6 +49,7 @@ typedef void (*object_table_done_callback)(object_id object_id, * * @param db_handle Handle to db. * @param object_id Object unique identifier. + * @param data_size Object data size. * @param retry Information about retrying the request to the database. * @param done_callback Callback to be called when lookup completes. * @param user_context User context to be passed in the callbacks. @@ -56,6 +57,7 @@ typedef void (*object_table_done_callback)(object_id object_id, */ void object_table_add(db_handle *db_handle, object_id object_id, + int64_t data_size, unsigned char digest[], retry_info *retry, object_table_done_callback done_callback, @@ -125,6 +127,45 @@ typedef struct { void *subscribe_context; } object_table_subscribe_data; +/* + * ==== Object info table, contains size of the object ==== + */ + +typedef void (*object_info_done_callback)(object_id object_id, + void *user_context); + +typedef void (*object_info_subscribe_callback)(object_id object_id, + int64_t object_size, + void *user_context); + +/** + * Subcribing to the object info pub/sub channel + * + * @param db_handle Handle to db. + * @param object_info_subscribe_callback callback triggered when pub/sub channel + * is notified of a new object size. + * @param subscribe_context caller context which will be passed back in the + * object_info_subscribe_callback. + * @param retry Information about retrying the request to the database. + * @param done_callback Callback to be called when subscription is installed. + * @param user_context User context to be passed into the done and fail + * callbacks. + * @return Void. + */ +void object_info_subscribe(db_handle *db_handle, + object_info_subscribe_callback subscribe_callback, + void *subscribe_context, + retry_info *retry, + object_info_done_callback done_callback, + void *user_context); + +/* Data that is needed to register new object info callbacks with the state + * database. */ +typedef struct { + object_info_subscribe_callback subscribe_callback; + void *subscribe_context; +} object_info_subscribe_data; + /* * ==== Result table ==== */ diff --git a/src/common/state/redis.c b/src/common/state/redis.c index 3ae0568d6..a3c4dbcfe 100644 --- a/src/common/state/redis.c +++ b/src/common/state/redis.c @@ -13,6 +13,7 @@ #include "db.h" #include "db_client_table.h" #include "object_table.h" +#include "object_info.h" #include "task.h" #include "task_table.h" #include "event_loop.h" @@ -260,6 +261,15 @@ task *parse_redis_task_table_entry(task_id id, * ==== object_table callbacks ==== */ +enum { + OBJECT_TABLE_ADD_CHECK_HASH_INDEX = 0, + OBJECT_TABLE_ADD_GET_HASH_INDEX, + OBJECT_TABLE_ADD_REGISTER_MANAGER_INDEX, + OBJECT_TABLE_ADD_SET_SIZE_INDEX, + OBJECT_TABLE_ADD_PUBLISH_INDEX, + OBJECT_TABLE_ADD_MAX +}; + void redis_object_table_add_callback(redisAsyncContext *c, void *r, void *privdata) { @@ -267,31 +277,27 @@ void redis_object_table_add_callback(redisAsyncContext *c, REDIS_MULTI_CALLBACK_HEADER(db, callback_data, r, requests_info); redisReply *reply = r; object_id id = callback_data->id; - unsigned char *digest = callback_data->data; -#define NUM_CHECK_AND_SET_COMMANDS 3 -#define CHECK_AND_SET_SETNX_INDEX 0 -#define CHECK_AND_SET_GET_INDEX 1 -#define CHECK_AND_SET_SADD_INDEX 2 + object_info *info = callback_data->data; /* Check that we're at a valid command index. */ int request_index = requests_info->request_index; LOG_DEBUG("Object table add request index is %d", request_index); - CHECK(request_index <= NUM_CHECK_AND_SET_COMMANDS); + CHECK(request_index <= OBJECT_TABLE_ADD_MAX); /* If we're on a valid command index, execute the current command and * register a callback that will execute the next command by incrementing the * request_index. */ int status = REDIS_OK; ++requests_info->request_index; - if (request_index == CHECK_AND_SET_SETNX_INDEX) { + if (request_index == OBJECT_TABLE_ADD_CHECK_HASH_INDEX) { /* Atomically set the object hash and get the previous value to compare to * our hash, if a previous value existed. */ requests_info->is_redis_reply = true; - status = - redisAsyncCommand(db->context, redis_object_table_add_callback, - (void *) requests_info, "SETNX objhash:%b %b", id.id, - sizeof(object_id), digest, (size_t) DIGEST_SIZE); - } else if (request_index == CHECK_AND_SET_GET_INDEX) { + status = redisAsyncCommand(db->context, redis_object_table_add_callback, + (void *) requests_info, "SETNX objhash:%b %b", + id.id, sizeof(object_id), info->digest, + (size_t) DIGEST_SIZE); + } else if (request_index == OBJECT_TABLE_ADD_GET_HASH_INDEX) { /* If there was an object hash in the table previously, check that it's * equal to ours. */ CHECKM(reply->type == REDIS_REPLY_INTEGER, @@ -308,13 +314,13 @@ void redis_object_table_add_callback(redisAsyncContext *c, (void *) requests_info, "GET objhash:%b", id.id, sizeof(object_id)); } - } else if (request_index == CHECK_AND_SET_SADD_INDEX) { + } else if (request_index == OBJECT_TABLE_ADD_REGISTER_MANAGER_INDEX) { if (requests_info->is_redis_reply) { CHECKM(reply->type == REDIS_REPLY_STRING, "Expected Redis string, received type %d %s", reply->type, reply->str); DCHECK(reply->len == DIGEST_SIZE); - if (memcmp(digest, reply->str, reply->len) != 0) { + if (memcmp(info->digest, reply->str, reply->len) != 0) { /* If our object hash doesn't match the one recorded in the table, * report the error back to the user and exit immediately. */ LOG_FATAL( @@ -329,6 +335,16 @@ void redis_object_table_add_callback(redisAsyncContext *c, (void *) requests_info, "SADD obj:%b %b", id.id, sizeof(id.id), (char *) db->client.id, sizeof(db->client.id)); + } else if (request_index == OBJECT_TABLE_ADD_SET_SIZE_INDEX) { + requests_info->is_redis_reply = true; + status = redisAsyncCommand(db->context, redis_object_table_add_callback, + (void *) requests_info, "HMSET obj:%b size %d", + (char *) id.id, sizeof(id.id), info->data_size); + } else if (request_index == OBJECT_TABLE_ADD_PUBLISH_INDEX) { + requests_info->is_redis_reply = true; + status = redisAsyncCommand(db->context, redis_object_table_add_callback, + (void *) requests_info, "PUBLISH obj:info %b:%d", + id.id, sizeof(id.id), info->data_size); } else { /* We finished executing all the Redis commands for this attempt at the * table operation. */ @@ -358,7 +374,7 @@ void redis_object_table_add(table_callback_data *callback_data) { LOG_DEBUG("Calling object table add"); redis_requests_info *requests_info = malloc(sizeof(redis_requests_info)); requests_info->timer_id = callback_data->timer_id; - requests_info->request_index = 0; + requests_info->request_index = OBJECT_TABLE_ADD_CHECK_HASH_INDEX; requests_info->is_redis_reply = false; db_handle *db = callback_data->db_handle; redis_object_table_add_callback(db->context, NULL, (void *) requests_info); @@ -587,10 +603,20 @@ void redis_object_table_subscribe(table_callback_data *callback_data) { /* subscribe to key notification associated to object id */ object_id id = callback_data->id; - int status = redisAsyncCommand( - db->sub_context, object_table_redis_subscribe_callback, - (void *) callback_data->timer_id, "SUBSCRIBE __keyspace@0__:obj:%b", - id.id, sizeof(id.id)); + int status = REDIS_OK; + + if (IS_NIL_ID(id)) { + /* Subscribe to all object events. */ + status = redisAsyncCommand( + db->sub_context, object_table_redis_subscribe_callback, + (void *) callback_data->timer_id, "PSUBSCRIBE __keyspace@0__:obj:*"); + } else { + /* Subscribe to the specified object id. */ + status = redisAsyncCommand( + db->sub_context, object_table_redis_subscribe_callback, + (void *) callback_data->timer_id, "SUBSCRIBE __keyspace@0__:obj:%b", + id.id, sizeof(id.id)); + } if ((status == REDIS_ERR) || db->sub_context->err) { LOG_REDIS_DEBUG(db->sub_context, "error in redis_object_table_subscribe_callback"); @@ -752,12 +778,12 @@ void redis_task_table_subscribe_callback(redisAsyncContext *c, redisReply *reply = r; CHECK(reply->type == REDIS_REPLY_ARRAY); - /* If this condition is true, we got the initial message that acknowledged the - * subscription. */ CHECK(reply->elements > 2); /* First entry is message type, then possibly the regex we psubscribed to, * then topic, then payload. */ redisReply *payload = reply->element[reply->elements - 1]; + /* If this condition is true, we got the initial message that acknowledged the + * subscription. */ if (payload->str == NULL) { if (callback_data->done_callback) { task_table_done_callback done_callback = callback_data->done_callback; @@ -813,12 +839,12 @@ void redis_db_client_table_subscribe_callback(redisAsyncContext *c, redisReply *reply = r; CHECK(reply->type == REDIS_REPLY_ARRAY); - /* If this condition is true, we got the initial message that acknowledged the - * subscription. */ CHECK(reply->elements > 2); /* First entry is message type, then possibly the regex we psubscribed to, * then topic, then payload. */ redisReply *payload = reply->element[reply->elements - 1]; + /* If this condition is true, we got the initial message that acknowledged the + * subscription. */ if (payload->str == NULL) { if (callback_data->done_callback) { db_client_table_done_callback done_callback = @@ -856,6 +882,54 @@ void redis_db_client_table_subscribe(table_callback_data *callback_data) { } } +void redis_object_info_subscribe_callback(redisAsyncContext *c, + void *r, + void *privdata) { + REDIS_CALLBACK_HEADER(db, callback_data, r); + redisReply *reply = r; + + CHECK(reply->type == REDIS_REPLY_ARRAY); + + CHECK(reply->elements > 2); + /* First entry is message type, then possibly the regex we psubscribed to, + * then topic, then payload. */ + redisReply *payload = reply->element[reply->elements - 1]; + /* If this condition is true, we got the initial message that acknowledged the + * subscription. */ + if (payload->str == NULL) { + if (callback_data->done_callback) { + db_client_table_done_callback done_callback = + callback_data->done_callback; + done_callback(callback_data->id, callback_data->user_context); + } + /* Note that we do not destroy the callback data yet because the + * subscription callback needs this data. */ + event_loop_remove_timer(db->loop, callback_data->timer_id); + return; + } + /* Otherwise, parse the payload and call the callback. */ + object_info_subscribe_data *data = callback_data->data; + object_id object_id; + memcpy(object_id.id, payload->str, sizeof(object_id.id)); + /* payload->str should have the format: "object_id:object_size_int" */ + LOG_DEBUG("obj:info channel received message <%s>", payload->str); + if (data->subscribe_callback) { + data->subscribe_callback( + object_id, strtol(&payload->str[1 + sizeof(object_id)], NULL, 10), + data->subscribe_context); + } +} + +void redis_object_info_subscribe(table_callback_data *callback_data) { + db_handle *db = callback_data->db_handle; + int status = redisAsyncCommand( + db->sub_context, redis_object_info_subscribe_callback, + (void *) callback_data->timer_id, "PSUBSCRIBE obj:info"); + if ((status == REDIS_ERR) || db->sub_context->err) { + LOG_REDIS_DEBUG(db->sub_context, "error in object_info_register_callback"); + } +} + db_client_id get_db_client_id(db_handle *db) { CHECK(db != NULL); return db->client; diff --git a/src/common/state/redis.h b/src/common/state/redis.h index ff8c32e1e..48f2e0570 100644 --- a/src/common/state/redis.h +++ b/src/common/state/redis.h @@ -30,7 +30,7 @@ struct db_handle { char *client_type; /** Unique ID for this client. */ db_client_id client; - /** Redis context for this global state store connection. */ + /** Redis context for all non-subscribe connections. */ redisAsyncContext *context; /** Redis context for "subscribe" communication. Yes, we need a separate one * for that, see https://github.com/redis/hiredis/issues/55. */ @@ -181,4 +181,6 @@ void redis_task_table_subscribe(table_callback_data *callback_data); */ void redis_db_client_table_subscribe(table_callback_data *callback_data); +void redis_object_info_subscribe(table_callback_data *callback_data); + #endif /* REDIS_H */ diff --git a/src/common/test/db_tests.c b/src/common/test/db_tests.c index 8c8be5791..63e2e7d19 100644 --- a/src/common/test/db_tests.c +++ b/src/common/test/db_tests.c @@ -77,9 +77,9 @@ TEST object_table_lookup_test(void) { .timeout = TIMEOUT, .fail_callback = timeout_callback, }; - object_table_add(db1, id, (unsigned char *) NIL_DIGEST, &retry, + object_table_add(db1, id, 0, (unsigned char *) NIL_DIGEST, &retry, add_done_callback, NULL); - object_table_add(db2, id, (unsigned char *) NIL_DIGEST, &retry, + object_table_add(db2, id, 0, (unsigned char *) NIL_DIGEST, &retry, add_done_callback, NULL); event_loop_add_timer(loop, 200, (event_loop_timer_handler) timeout_handler, NULL); diff --git a/src/common/test/object_table_tests.c b/src/common/test/object_table_tests.c index 619841654..b73d3905e 100644 --- a/src/common/test/object_table_tests.c +++ b/src/common/test/object_table_tests.c @@ -202,7 +202,7 @@ TEST add_timeout_test(void) { retry_info retry = { .num_retries = 5, .timeout = 100, .fail_callback = add_fail_callback, }; - object_table_add(db, NIL_ID, (unsigned char *) NIL_DIGEST, &retry, + object_table_add(db, NIL_ID, 0, (unsigned char *) NIL_DIGEST, &retry, add_done_callback, (void *) add_timeout_context); /* Disconnect the database to see if the lookup times out. */ close(db->context->c.fd); @@ -360,7 +360,7 @@ TEST add_retry_test(void) { .timeout = 100, .fail_callback = add_retry_fail_callback, }; - object_table_add(db, NIL_ID, (unsigned char *) NIL_DIGEST, &retry, + object_table_add(db, NIL_ID, 0, (unsigned char *) NIL_DIGEST, &retry, add_retry_done_callback, (void *) add_retry_context); /* Disconnect the database to let the add time out the first time. */ close(db->context->c.fd); @@ -413,7 +413,7 @@ TEST add_lookup_test(void) { .timeout = 100, .fail_callback = lookup_retry_fail_callback, }; - object_table_add(db, NIL_ID, (unsigned char *) NIL_DIGEST, &retry, + object_table_add(db, NIL_ID, 0, (unsigned char *) NIL_DIGEST, &retry, add_lookup_callback, (void *) db); /* Install handler for terminating the event loop. */ event_loop_add_timer(g_loop, 750, @@ -568,7 +568,7 @@ TEST add_late_test(void) { retry_info retry = { .num_retries = 0, .timeout = 0, .fail_callback = add_late_fail_callback, }; - object_table_add(db, NIL_ID, (unsigned char *) NIL_DIGEST, &retry, + object_table_add(db, NIL_ID, 0, (unsigned char *) NIL_DIGEST, &retry, add_late_done_callback, (void *) add_late_context); /* Install handler for terminating the event loop. */ event_loop_add_timer(g_loop, 750, @@ -653,7 +653,7 @@ void subscribe_success_done_callback(object_id object_id, retry_info retry = { .num_retries = 0, .timeout = 750, .fail_callback = NULL, }; - object_table_add((db_handle *) user_context, object_id, + object_table_add((db_handle *) user_context, object_id, 0, (unsigned char *) NIL_DIGEST, &retry, NULL, NULL); subscribe_success_done = 1; } @@ -721,7 +721,7 @@ TEST subscribe_object_present_test(void) { retry_info retry = { .num_retries = 0, .timeout = 100, .fail_callback = NULL, }; - object_table_add(db, id, (unsigned char *) NIL_DIGEST, &retry, NULL, NULL); + object_table_add(db, id, 0, (unsigned char *) NIL_DIGEST, &retry, NULL, NULL); object_table_subscribe( db, id, subscribe_object_present_object_available_callback, (void *) subscribe_object_present_context, &retry, NULL, (void *) db); @@ -802,7 +802,7 @@ int64_t add_object_callback(event_loop *loop, int64_t timer_id, void *context) { retry_info retry = { .num_retries = 0, .timeout = 100, .fail_callback = NULL, }; - object_table_add(db, NIL_ID, (unsigned char *) NIL_DIGEST, &retry, NULL, + object_table_add(db, NIL_ID, 0, (unsigned char *) NIL_DIGEST, &retry, NULL, NULL); /* Reset the timer to this large value, so it doesn't trigger again. */ return 10000; @@ -838,6 +838,76 @@ TEST subscribe_object_available_later_test(void) { PASS(); } +/* Test if object size is correctly reported by the object_info callback. */ + +typedef struct { + char *subscribe_success_msg; + int64_t data_size; + int subscribe_succeeded; + int subscribe_callback_done; +} object_info_subscribe_context; + +object_info_subscribe_context obj_info_subscribe_context = {"foo", 42, 0, 0}; + +void subscribe_object_info_done_callback(object_id object_id, + void *user_context) { + retry_info retry = { + .num_retries = 0, .timeout = 100, .fail_callback = NULL, + }; + CHECK(obj_info_subscribe_context.subscribe_succeeded == 0); + CHECK(obj_info_subscribe_context.subscribe_callback_done == 0); + + object_table_add((db_handle *) user_context, object_id, + obj_info_subscribe_context.data_size, NIL_DIGEST, &retry, + NULL, NULL); + + obj_info_subscribe_context.subscribe_callback_done = 1; +} + +void subscribe_success_object_info_available_callback(object_id object_id, + int64_t object_size, + void *user_context) { + CHECK(user_context == (void *) &obj_info_subscribe_context); + /* Check to make sure subscription done callback already fired. */ + CHECK(obj_info_subscribe_context.subscribe_callback_done == 1); + CHECK(obj_info_subscribe_context.subscribe_succeeded == 0); + CHECK(obj_info_subscribe_context.data_size == object_size); + + /* Mark success. */ + obj_info_subscribe_context.subscribe_succeeded = 1; +} + +TEST subscribe_object_info_success_test(void) { + g_loop = event_loop_create(); + db_handle *db = + db_connect("127.0.0.1", 6379, "plasma_manager", "127.0.0.1", 11236); + db_attach(db, g_loop, false); + + retry_info retry = { + .num_retries = 0, + .timeout = 100, + .fail_callback = subscribe_success_fail_callback, + }; + + object_info_subscribe(db, subscribe_success_object_info_available_callback, + (void *) &obj_info_subscribe_context, &retry, + subscribe_object_info_done_callback, (void *) db); + + /* Install handler for terminating the event loop. */ + event_loop_add_timer(g_loop, 1000, + (event_loop_timer_handler) terminate_event_loop_callback, + NULL); + + event_loop_run(g_loop); + db_disconnect(db); + destroy_outstanding_callbacks(g_loop); + event_loop_destroy(g_loop); + + ASSERT(obj_info_subscribe_context.subscribe_succeeded == 1); + ASSERT(obj_info_subscribe_context.subscribe_callback_done == 1); + PASS(); +} + SUITE(object_table_tests) { RUN_REDIS_TEST(new_object_test); RUN_REDIS_TEST(new_object_no_task_test); @@ -855,6 +925,7 @@ SUITE(object_table_tests) { RUN_REDIS_TEST(subscribe_object_present_test); RUN_REDIS_TEST(subscribe_object_not_present_test); RUN_REDIS_TEST(subscribe_object_available_later_test); + RUN_REDIS_TEST(subscribe_object_info_success_test); } GREATEST_MAIN_DEFS(); diff --git a/src/plasma/plasma_manager.c b/src/plasma/plasma_manager.c index 19fe998aa..f172e5ee0 100644 --- a/src/plasma/plasma_manager.c +++ b/src/plasma/plasma_manager.c @@ -1528,7 +1528,8 @@ void process_object_notification(event_loop *loop, if (state->db) { /* TODO(swang): Log the error if we fail to add the object, and possibly * retry later? */ - object_table_add(state->db, obj_id, object_info.digest, &retry, NULL, NULL); + object_table_add(state->db, obj_id, object_info.data_size, + object_info.digest, &retry, NULL, NULL); } /* If we were trying to fetch this object, finish up the fetch request. */ diff --git a/src/plasma/test/test.py b/src/plasma/test/test.py index 567e44ef3..8238530d7 100644 --- a/src/plasma/test/test.py +++ b/src/plasma/test/test.py @@ -342,7 +342,7 @@ class TestPlasmaClient(unittest.TestCase): metadata_sizes = [np.random.randint(1000) for _ in range(i)] data_sizes = [np.random.randint(1000) for _ in range(i)] for j in range(i): - self.plasma_client.create(object_ids[j], size=data_sizes[j], + self.plasma_client.create(object_ids[j], size=data_sizes[j], metadata=bytearray(np.random.bytes(metadata_sizes[j]))) self.plasma_client.seal(object_ids[j]) # Check that we received notifications for all of the objects. @@ -639,10 +639,17 @@ class TestPlasmaManager(unittest.TestCase): self.client2.seal(object_id) # Give the second manager some time to complete the seal, then make sure it # exited. - time.sleep(2) - self.p5.poll() + time_left = 10 + while time_left > 0: + self.p5.poll() + if self.p5.returncode != None: + self.processes_to_kill.remove(self.p5) + break + time_left -= 0.2 + time.sleep(0.2) + + print("Time waiting for plasma manager to fail = {:.2}".format(10 - time_left)) self.assertNotEqual(self.p5.returncode, None) - self.processes_to_kill.remove(self.p5) def test_illegal_functionality(self): # Create an object id string.