Fixes for implementing Plasma fetch (#39)

* Add ability to pass callback context to object table lookup

* Propagate errors during socket writes up to caller.

* Use recv and MSG_WAITALL flag instead of looping read

* Error checking in write_bytes

* Method to listen on a network port

* Revert "Use recv and MSG_WAITALL flag instead of looping read"

This reverts commit 32d9333bc6a185729aadb4b41b70b3d7f150a9c2.

* Some documentation

* Clearer documentation

* Fix bug where database clients were getting assigned the same ID

* Regression test for unique client IDs

io.c

@@ -8,14 +8,69 @@
 #include <stdio.h>
 #include <inttypes.h>
 #include <stdarg.h>
+#include <sys/ioctl.h>
+#include <netinet/in.h>
 #include <utstring.h>
 
 #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) {

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);

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);

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;
+  }
+}

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

test/db_tests.c

@@ -1,6 +1,8 @@
 #include "greatest.h"
 
 #include <assert.h>
+#include <unistd.h>
+#include <sys/wait.h>
 
 #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);
 }