diff --git a/io.c b/io.c index 32f3ed4b9..1f1125bc3 100644 --- a/io.c +++ b/io.c @@ -8,14 +8,69 @@ #include #include #include +#include +#include #include #include "common.h" -/* Binds to a Unix domain streaming socket at the given - * pathname. Removes any existing file at the pathname. Returns - * a file descriptor for the socket, or -1 if an error - * occurred. */ +/** + * Binds to an Internet socket at the given port. Removes any existing file at + * the pathname. Returns a non-blocking file descriptor for the socket, or -1 + * if an error occurred. + * + * @note Since the returned file descriptor is non-blocking, it is not + * recommended to use the Linux read and write calls directly, since these + * might read or write a partial message. Instead, use the provided + * write_message and read_message methods. + * + * @param port The port to bind to. + * @return A non-blocking file descriptor for the socket, or -1 if an error + * occurs. + */ +int bind_inet_sock(const int port) { + struct sockaddr_in name; + int socket_fd = socket(PF_INET, SOCK_STREAM, 0); + if (socket_fd < 0) { + LOG_ERR("socket() failed for port %d.", port); + return -1; + } + name.sin_family = AF_INET; + name.sin_port = htons(port); + name.sin_addr.s_addr = htonl(INADDR_ANY); + int on = 1; + /* TODO(pcm): http://stackoverflow.com/q/1150635 */ + if (ioctl(socket_fd, FIONBIO, (char *) &on) < 0) { + LOG_ERR("ioctl failed"); + close(socket_fd); + return -1; + } + if (setsockopt(socket_fd, SOL_SOCKET, SO_REUSEADDR, &on, sizeof(on)) < 0) { + LOG_ERR("setsockopt failed for port %d", port); + close(socket_fd); + return -1; + } + if (bind(socket_fd, (struct sockaddr *) &name, sizeof(name)) < 0) { + LOG_ERR("Bind failed for port %d", port); + close(socket_fd); + return -1; + } + if (listen(socket_fd, 5) == -1) { + LOG_ERR("Could not listen to socket %d", port); + close(socket_fd); + return -1; + } + return socket_fd; +} + +/** + * Binds to a Unix domain streaming socket at the given + * pathname. Removes any existing file at the pathname. + * + * @param socket_pathname The pathname for the socket. + * @return A blocking file descriptor for the socket, or -1 if an error + * occurs. + */ int bind_ipc_sock(const char *socket_pathname) { struct sockaddr_un socket_address; int socket_fd = socket(AF_UNIX, SOCK_STREAM, 0); @@ -27,9 +82,9 @@ int bind_ipc_sock(const char *socket_pathname) { int on = 1; if (setsockopt(socket_fd, SOL_SOCKET, SO_REUSEADDR, (char *) &on, sizeof(on)) < 0) { - LOG_ERR("setsockopt failed"); + LOG_ERR("setsockopt failed for pathname %s", socket_pathname); close(socket_fd); - exit(-1); + return -1; } unlink(socket_pathname); @@ -37,6 +92,7 @@ int bind_ipc_sock(const char *socket_pathname) { socket_address.sun_family = AF_UNIX; if (strlen(socket_pathname) + 1 > sizeof(socket_address.sun_path)) { LOG_ERR("Socket pathname is too long."); + close(socket_fd); return -1; } strncpy(socket_address.sun_path, socket_pathname, @@ -45,16 +101,22 @@ int bind_ipc_sock(const char *socket_pathname) { if (bind(socket_fd, (struct sockaddr *) &socket_address, sizeof(struct sockaddr_un)) != 0) { LOG_ERR("Bind failed for pathname %s.", socket_pathname); + close(socket_fd); + return -1; + } + if (listen(socket_fd, 5) == -1) { + LOG_ERR("Could not listen to socket %s", socket_pathname); + close(socket_fd); return -1; } - listen(socket_fd, 5); - return socket_fd; } -/* Connects to a Unix domain streaming socket at the given +/** + * Connects to a Unix domain streaming socket at the given * pathname. Returns a file descriptor for the socket, or -1 if - * an error occurred. */ + * an error occurred. + */ int connect_ipc_sock(const char *socket_pathname) { struct sockaddr_un socket_address; int socket_fd; @@ -83,8 +145,10 @@ int connect_ipc_sock(const char *socket_pathname) { return socket_fd; } -/* Accept a new client connection on the given socket - * descriptor. Returns a descriptor for the new socket. */ +/** + * Accept a new client connection on the given socket + * descriptor. Returns a descriptor for the new socket. + */ int accept_client(int socket_fd) { int client_fd = accept(socket_fd, NULL, NULL); if (client_fd < 0) { @@ -95,55 +159,80 @@ int accept_client(int socket_fd) { } /** - * Reliably write a sequence of bytes into a file descriptor. This will block - * until one of the following happens: (1) there is an error (2) end of file, - * or (3) all length bytes have been written. + * Write a sequence of bytes into a file descriptor. This will block until one + * of the following happens: (1) there is an error (2) end of file, or (3) all + * length bytes have been written. * - * @param fd The file descriptor to write to. + * @param fd The file descriptor to write to. It can be non-blocking. * @param cursor The cursor pointing to the beginning of the bytes to send. * @param length The size of the bytes sequence to write. - * @return Void. + * @return int Whether there was an error while writing. 0 corresponds to + * success and -1 corresponds to an error (errno will be set). */ -void write_bytes(int fd, uint8_t *cursor, size_t length) { +int write_bytes(int fd, uint8_t *cursor, size_t length) { ssize_t nbytes = 0; while (length > 0) { /* While we haven't written the whole message, write to the file * descriptor, advance the cursor, and decrease the amount left to write. */ nbytes = write(fd, cursor, length); - CHECK(nbytes > 0); + if (nbytes < 0) { + if (errno == EAGAIN || errno == EWOULDBLOCK) { + continue; + } + /* TODO(swang): Return the error instead of exiting. */ + /* Force an exit if there was any other type of error. */ + CHECK(nbytes < 0); + } + if (nbytes == 0) { + return -1; + } cursor += nbytes; length -= nbytes; } + return 0; } /** * Write a sequence of bytes on a file descriptor. The bytes should then be read * by read_message. * - * @param fd The file descriptor to write to. + * @param fd The file descriptor to write to. It can be non-blocking. * @param type The type of the message to send. * @param length The size in bytes of the bytes parameter. * @param bytes The address of the message to send. - * @return Void. + * @return int Whether there was an error while writing. 0 corresponds to + * success and -1 corresponds to an error (errno will be set). */ -void write_message(int fd, int64_t type, int64_t length, uint8_t *bytes) { - write_bytes(fd, (uint8_t *) &type, sizeof(type)); - write_bytes(fd, (uint8_t *) &length, sizeof(length)); - write_bytes(fd, bytes, length * sizeof(char)); +int write_message(int fd, int64_t type, int64_t length, uint8_t *bytes) { + int closed; + closed = write_bytes(fd, (uint8_t *) &type, sizeof(type)); + if (closed) { + return closed; + } + closed = write_bytes(fd, (uint8_t *) &length, sizeof(length)); + if (closed) { + return closed; + } + closed = write_bytes(fd, bytes, length * sizeof(char)); + if (closed) { + return closed; + } + return 0; } /** - * Reliably read a sequence of bytes from a file descriptor into a buffer. This - * will block until one of the following happens: (1) there is an error (2) end - * of file, or (3) all length bytes have been written. + * Read a sequence of bytes from a file descriptor into a buffer. This will + * block until one of the following happens: (1) there is an error (2) end of + * file, or (3) all length bytes have been written. * * @note The buffer pointed to by cursor must already have length number of * bytes allocated before calling this method. * - * @param fd The file descriptor to read from. + * @param fd The file descriptor to read from. It can be non-blocking. * @param cursor The cursor pointing to the beginning of the buffer. * @param length The size of the byte sequence to read. - * @return Void. + * @return int Whether there was an error while writing. 0 corresponds to + * success and -1 corresponds to an error (errno will be set). */ int read_bytes(int fd, uint8_t *cursor, size_t length) { ssize_t nbytes = 0; @@ -173,14 +262,18 @@ int read_bytes(int fd, uint8_t *cursor, size_t length) { * * @note The caller must free the memory. * - * @param fd The file descriptor to read from. + * @param fd The file descriptor to read from. It can be non-blocking. * @param type The type of the message that is read will be written at this - address. + address. If there was an error while reading, this will be + DISCONNECT_CLIENT. * @param length The size in bytes of the message that is read will be written at this address. This size does not include the bytes used to encode - the type and length. + the type and length. If there was an error while reading, this will + be 0. * @param bytes The address at which to write the pointer to the bytes that are - read and allocated by this function. + read and allocated by this function. If there was an error while + reading, this will be NULL. + * @return Void. */ void read_message(int fd, int64_t *type, int64_t *length, uint8_t **bytes) { diff --git a/io.h b/io.h index 2299806f7..362b85fc1 100644 --- a/io.h +++ b/io.h @@ -14,6 +14,7 @@ enum common_message_type { /* Helper functions for socket communication. */ +int bind_inet_sock(const int port); int bind_ipc_sock(const char *socket_pathname); int connect_ipc_sock(const char *socket_pathname); @@ -21,7 +22,7 @@ int accept_client(int socket_fd); /* Reading and writing data */ -void write_message(int fd, int64_t type, int64_t length, uint8_t *bytes); +int write_message(int fd, int64_t type, int64_t length, uint8_t *bytes); void read_message(int fd, int64_t *type, int64_t *length, uint8_t **bytes); void write_log_message(int fd, char *message); diff --git a/state/object_table.h b/state/object_table.h index e2eb89433..bab54bc7b 100644 --- a/state/object_table.h +++ b/state/object_table.h @@ -6,7 +6,8 @@ * the manager_vector array, but NOT the strings they are pointing to. */ typedef void (*lookup_callback)(object_id object_id, int manager_count, - const char *manager_vector[]); + const char *manager_vector[], + void *context); /* Register a new object with the directory. */ /* TODO(pcm): Retry, print for each attempt. */ @@ -20,4 +21,5 @@ void object_table_remove(db_handle *db, /* Look up entry from the directory */ void object_table_lookup(db_handle *db, object_id object_id, - lookup_callback callback); + lookup_callback callback, + void *context); diff --git a/state/redis.c b/state/redis.c index db008cb5d..d53b0bf15 100644 --- a/state/redis.c +++ b/state/redis.c @@ -56,7 +56,8 @@ db_handle *db_connect(const char *address, num_clients, client_addr, client_port); freeReplyObject(reply); reply = redisCommand(context, "EXEC"); - if (reply) { + CHECK(reply); + if (reply->type != REDIS_REPLY_NIL) { freeReplyObject(reply); break; } @@ -150,17 +151,20 @@ void object_table_get_entry(redisAsyncContext *c, void *r, void *privdata) { HASH_FIND_INT(db->service_cache, &result[j], entry); manager_vector[j] = entry->addr; } - cb_data->callback(cb_data->object_id, manager_count, manager_vector); + cb_data->callback(cb_data->object_id, manager_count, manager_vector, + cb_data->context); free(privdata); free(result); } void object_table_lookup(db_handle *db, object_id object_id, - lookup_callback callback) { + lookup_callback callback, + void *context) { lookup_callback_data *cb_data = malloc(sizeof(lookup_callback_data)); cb_data->callback = callback; cb_data->object_id = object_id; + cb_data->context = context; redisAsyncCommand(db->context, object_table_get_entry, cb_data, "SMEMBERS obj:%b", &object_id.id[0], UNIQUE_ID_SIZE); if (db->context->err) { @@ -230,3 +234,11 @@ void task_log_register_callback(db_handle *db, LOG_REDIS_ERR(db->sub_context, "error in task_log_register_callback"); } } + +int get_client_id(db_handle *db) { + if (db) { + return db->client_id; + } else { + return -1; + } +} diff --git a/state/redis.h b/state/redis.h index 51479f8f0..a3e0555e7 100644 --- a/state/redis.h +++ b/state/redis.h @@ -55,6 +55,8 @@ typedef struct { lookup_callback callback; /* Object ID that is looked up. */ object_id object_id; + /* Data context for the callback. */ + void *context; } lookup_callback_data; void object_table_get_entry(redisAsyncContext *c, void *r, void *privdata); @@ -63,4 +65,13 @@ void object_table_lookup_callback(redisAsyncContext *c, void *r, void *privdata); +/** + * Returns the client ID, according to Redis. + * + * @param db The handle to the Redis database. + * @returns int The client ID for this connection to Redis. If + * this client has no connection to Redis, returns -1. + */ +int get_client_id(db_handle *db); + #endif diff --git a/test/db_tests.c b/test/db_tests.c index 95e986e33..74dd9a850 100644 --- a/test/db_tests.c +++ b/test/db_tests.c @@ -1,6 +1,8 @@ #include "greatest.h" #include +#include +#include #include "event_loop.h" #include "test/example_task.h" @@ -23,7 +25,8 @@ char received_port2[6] = {0}; /* Test if entries have been written to the database. */ void test_callback(object_id object_id, int manager_count, - const char *manager_vector[]) { + const char *manager_vector[], + void *context) { CHECK(manager_count == 2); if (!manager_vector[0] || sscanf(manager_vector[0], "%15[0-9.]:%5[0-9]", received_addr1, @@ -56,7 +59,7 @@ TEST object_table_lookup_test(void) { object_table_add(db2, id); event_loop_add_timer(loop, 100, timeout_handler, NULL); event_loop_run(loop); - object_table_lookup(db1, id, test_callback); + object_table_lookup(db1, id, test_callback, NULL); event_loop_add_timer(loop, 100, timeout_handler, NULL); event_loop_run(loop); int port1 = atoi(received_port1); @@ -130,12 +133,36 @@ TEST task_log_all_test(void) { PASS(); } +TEST unique_client_id_test(void) { + const int num_conns = 50; + + db_handle *db; + pid_t pid = fork(); + for (int i = 0; i < num_conns; ++i) { + db = db_connect("127.0.0.1", 6379, "plasma_manager", manager_addr, + manager_port1); + db_disconnect(db); + } + if (pid == 0) { + exit(0); + } else { + wait(NULL); + } + + db = db_connect("127.0.0.1", 6379, "plasma_manager", manager_addr, + manager_port1); + ASSERT_EQ(get_client_id(db), num_conns * 2); + db_disconnect(db); + PASS(); +} + SUITE(db_tests) { redisContext *context = redisConnect("127.0.0.1", 6379); freeReplyObject(redisCommand(context, "FLUSHALL")); RUN_REDIS_TEST(context, object_table_lookup_test); RUN_REDIS_TEST(context, task_log_test); RUN_REDIS_TEST(context, task_log_all_test); + RUN_REDIS_TEST(context, unique_client_id_test); redisFree(context); }