Allow reading/writing generic message types, not just tasks. (#24)

* Allow reading/writing generic message types, not just tasks.

* Allow messages of length 0 to be read/written, and handle closed sockets.

* Address comments.

* Simplify accept_client.

* Allow ports to be reused in bind_ipc_sock.

io.c

@@ -18,13 +18,19 @@
  * occurred. */
 int bind_ipc_sock(const char *socket_pathname) {
   struct sockaddr_un socket_address;
-  int socket_fd;
-
-  socket_fd = socket(AF_UNIX, SOCK_STREAM, 0);
+  int socket_fd = socket(AF_UNIX, SOCK_STREAM, 0);
   if (socket_fd < 0) {
     LOG_ERR("socket() failed for pathname %s.", socket_pathname);
     return -1;
   }
+  /* Tell the system to allow the port to be reused. */
+  int on = 1;
+  if (setsockopt(socket_fd, SOL_SOCKET, SO_REUSEADDR, (char *) &on,
+                 sizeof(on)) < 0) {
+    LOG_ERR("setsockopt failed");
+    close(socket_fd);
+    exit(-1);
+  }
 
   unlink(socket_pathname);
   memset(&socket_address, 0, sizeof(struct sockaddr_un));
@@ -80,11 +86,7 @@ int connect_ipc_sock(const char *socket_pathname) {
 /* Accept a new client connection on the given socket
  * descriptor. Returns a descriptor for the new socket. */
 int accept_client(int socket_fd) {
-  struct sockaddr_un client_addr;
-  int client_fd, client_len;
-  client_len = sizeof(client_addr);
-  client_fd = accept(socket_fd, (struct sockaddr *) &client_addr,
-                     (socklen_t *) &client_len);
+  int client_fd = accept(socket_fd, NULL, NULL);
   if (client_fd < 0) {
     LOG_ERR("Error reading from socket.");
     return -1;
@@ -92,57 +94,77 @@ int accept_client(int socket_fd) {
   return client_fd;
 }
 
-/* Write a sequence of bytes on a file descriptor. */
-void write_bytes(int fd, uint8_t *bytes, int64_t length) {
-  ssize_t nbytes = write(fd, (char *) &length, sizeof(length));
-  if (nbytes == -1) {
-    LOG_ERR("Error sending to socket.\n");
-    return;
-  }
+/**
+ * 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 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.
+ */
+void write_message(int fd, int64_t type, int64_t length, uint8_t *bytes) {
+  ssize_t nbytes = write(fd, (char *) &type, sizeof(type));
+  CHECK(nbytes == sizeof(int64_t));
+  nbytes = write(fd, (char *) &length, sizeof(length));
+  CHECK(nbytes == sizeof(int64_t));
   nbytes = write(fd, (char *) bytes, length * sizeof(char));
-  if (nbytes == -1) {
-    LOG_ERR("Error sending to socket.\n");
-    return;
-  }
+  CHECK(nbytes >= 0);
 }
 
-/* Read a sequence of bytes written by write_bytes from a file descriptor.
- * Allocates and returns a pointer to the bytes.
- * NOTE: Caller must free the memory! */
-void read_bytes(int fd, uint8_t **bytes, int64_t *length) {
-  ssize_t nbytes = read(fd, length, sizeof(int64_t));
-  if (nbytes < 0) {
-    LOG_ERR("Error reading length of message from socket.");
+/**
+ * Read a sequence of bytes written by write_bytes from a file descriptor. This
+ * allocates space for the message.
+ *
+ * @note The caller must free the memory.
+ *
+ * @param fd The file descriptor to read from.
+ * @param type The type of the message that is read will be written at this
+          address.
+ * @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.
+ * @param bytes The address at which to write the pointer to the bytes that are
+          read and allocated by this function.
+ * @return Void.
+ */
+void read_message(int fd, int64_t *type, int64_t *length, uint8_t **bytes) {
+  ssize_t nbytes = read(fd, type, sizeof(int64_t));
+  CHECK(nbytes >= 0);
+  /* Handle the case in which the socket is closed. */
+  if (nbytes == 0) {
+    *type = DISCONNECT_CLIENT;
+    *length = 0;
     *bytes = NULL;
     return;
   }
-
+  nbytes = read(fd, length, sizeof(int64_t));
+  CHECK(nbytes == sizeof(int64_t));
   *bytes = malloc(*length * sizeof(uint8_t));
   nbytes = read(fd, *bytes, *length);
-  if (nbytes < 0) {
-    LOG_ERR("Error reading message from socket.");
-    free(*bytes);
-    *bytes = NULL;
-  }
+  CHECK(nbytes >= 0);
 }
 
 /* Write a null-terminated string to a file descriptor. */
-void write_string(int fd, char *message) {
+void write_log_message(int fd, char *message) {
   /* Account for the \0 at the end of the string. */
-  write_bytes(fd, (uint8_t *) message, strlen(message) + 1);
+  write_message(fd, LOG_MESSAGE, strlen(message) + 1, (uint8_t *) message);
 }
 
 /* Reads a null-terminated string from the file descriptor that has been
- * written by write_string. Allocates and returns a pointer to the string.
+ * written by write_log_message. Allocates and returns a pointer to the string.
  * NOTE: Caller must free the memory! */
-char *read_string(int fd) {
+char *read_log_message(int fd) {
   uint8_t *bytes;
+  int64_t type;
   int64_t length;
-  read_bytes(fd, &bytes, &length);
+  read_message(fd, &type, &length, &bytes);
+  CHECK(type == LOG_MESSAGE);
   return (char *) bytes;
 }
 
-void write_formatted_string(int socket_fd, const char *format, ...) {
+void write_formatted_log_message(int socket_fd, const char *format, ...) {
   UT_string *cmd;
   va_list ap;
 
@@ -151,6 +173,6 @@ void write_formatted_string(int socket_fd, const char *format, ...) {
   utstring_printf_va(cmd, format, ap);
   va_end(ap);
 
-  write_string(socket_fd, utstring_body(cmd));
+  write_log_message(socket_fd, utstring_body(cmd));
   utstring_free(cmd);
 }

io.h

@@ -3,6 +3,15 @@
 
 #include <stdint.h>
 
+enum common_message_type {
+  /** Disconnect a client. */
+  DISCONNECT_CLIENT,
+  /** Log a message from a client. */
+  LOG_MESSAGE,
+  /** Submit a task to the local scheduler. */
+  SUBMIT_TASK,
+};
+
 /* Helper functions for socket communication. */
 
 int bind_ipc_sock(const char *socket_pathname);
@@ -12,11 +21,11 @@ int accept_client(int socket_fd);
 
 /* Reading and writing data */
 
-void write_bytes(int fd, uint8_t *bytes, int64_t length);
-void read_bytes(int fd, uint8_t **bytes, int64_t *length);
+void 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_string(int fd, char *message);
-void write_formatted_string(int fd, const char *format, ...);
-char *read_string(int fd);
+void write_log_message(int fd, char *message);
+void write_formatted_log_message(int fd, const char *format, ...);
+char *read_log_message(int fd);
 
 #endif

logging.c

@@ -68,10 +68,10 @@ void ray_log(ray_logger *logger,
      * ID to be filled in by someone else. */
     utstring_printf(origin_id, "%s:%s", "%ld", "%ld");
     int *socket_fd = (int *) logger->conn;
-    write_formatted_string(*socket_fd, log_fmt, utstring_body(timestamp),
-                           logger->client_type, utstring_body(origin_id),
-                           log_levels[log_level], event_type, message,
-                           utstring_body(timestamp));
+    write_formatted_log_message(*socket_fd, log_fmt, utstring_body(timestamp),
+                                logger->client_type, utstring_body(origin_id),
+                                log_levels[log_level], event_type, message,
+                                utstring_body(timestamp));
   }
   utstring_free(origin_id);
   utstring_free(timestamp);

task.c

@@ -148,19 +148,6 @@ void free_task_spec(task_spec *spec) {
   free(spec);
 }
 
-void write_task(int fd, task_spec *spec) {
-  write_bytes(fd, (uint8_t *) spec, task_size(spec));
-}
-
-task_spec *read_task(int fd) {
-  uint8_t *bytes;
-  int64_t length;
-  read_bytes(fd, &bytes, &length);
-  task_spec *spec = (task_spec *) bytes;
-  CHECK(task_size(spec) == length);
-  return 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'). */

test/io_tests.c

@@ -20,20 +20,23 @@ TEST ipc_socket_test(void) {
     close(socket_fd);
     socket_fd = connect_ipc_sock(socket_pathname);
     ASSERT(socket_fd >= 0);
-    write_string(socket_fd, test_string);
-    write_bytes(socket_fd, (uint8_t *) test_bytes, strlen(test_bytes));
+    write_log_message(socket_fd, test_string);
+    write_message(socket_fd, LOG_MESSAGE, strlen(test_bytes),
+                  (uint8_t *) test_bytes);
     close(socket_fd);
     exit(0);
   } else {
     int client_fd = accept_client(socket_fd);
     ASSERT(client_fd >= 0);
-    char *message = read_string(client_fd);
+    char *message = read_log_message(client_fd);
     ASSERT(message != NULL);
     ASSERT_STR_EQ(test_string, message);
     free(message);
+    int64_t type;
     int64_t len;
     uint8_t *bytes;
-    read_bytes(client_fd, &bytes, &len);
+    read_message(client_fd, &type, &len, &bytes);
+    ASSERT(type == LOG_MESSAGE);
     ASSERT(memcmp(test_bytes, bytes, len) == 0);
     free(bytes);
     close(client_fd);

test/redis_tests.c

@@ -45,10 +45,10 @@ TEST redis_socket_test(void) {
 
   int client_fd = connect_ipc_sock(socket_pathname);
   ASSERT(client_fd >= 0);
-  write_formatted_string(client_fd, test_set_format, test_key, test_value);
+  write_formatted_log_message(client_fd, test_set_format, test_key, test_value);
 
   int server_fd = accept_client(socket_fd);
-  char *cmd = read_string(server_fd);
+  char *cmd = read_log_message(server_fd);
   close(client_fd);
   close(server_fd);
   close(socket_fd);
@@ -69,7 +69,7 @@ TEST redis_socket_test(void) {
 
 void redis_read_callback(event_loop *loop, int fd, void *context, int events) {
   db_handle *db = context;
-  char *cmd = read_string(fd);
+  char *cmd = read_log_message(fd);
   redisAsyncCommand(db->context, async_redis_socket_test_callback, NULL, cmd,
                     db->client_id, 0);
   free(cmd);
@@ -109,7 +109,7 @@ TEST async_redis_socket_test(void) {
   int client_fd = connect_ipc_sock(socket_pathname);
   ASSERT(client_fd >= 0);
   utarray_push_back(connections, &client_fd);
-  write_formatted_string(client_fd, test_set_format, test_key, test_value);
+  write_formatted_log_message(client_fd, test_set_format, test_key, test_value);
 
   event_loop_add_file(loop, client_fd, EVENT_LOOP_READ, redis_read_callback,
                       db);
@@ -148,7 +148,7 @@ void logging_read_callback(event_loop *loop,
                            void *context,
                            int events) {
   db_handle *conn = context;
-  char *cmd = read_string(fd);
+  char *cmd = read_log_message(fd);
   redisAsyncCommand(conn->context, logging_test_callback, NULL, cmd,
                     conn->client_id, 0);
   free(cmd);

test/task_tests.c

@@ -7,6 +7,7 @@
 #include "common.h"
 #include "test/example_task.h"
 #include "task.h"
+#include "io.h"
 
 SUITE(task_tests);
 
@@ -48,8 +49,13 @@ TEST send_task(void) {
   *task_return(task, 1) = globally_unique_id();
   int fd[2];
   socketpair(AF_UNIX, SOCK_STREAM, 0, fd);
-  write_task(fd[0], task);
-  task_spec *result = read_task(fd[1]);
+  write_message(fd[0], SUBMIT_TASK, task_size(task), task);
+  int64_t type;
+  int64_t length;
+  uint8_t *message;
+  read_message(fd[1], &type, &length, &message);
+  task_spec *result = (task_spec *) message;
+  ASSERT(type == SUBMIT_TASK);
   ASSERT(memcmp(task, result, task_size(task)) == 0);
   ASSERT(memcmp(task, result, task_size(result)) == 0);
   free(task);