From 1499834be1324349b7f0c0ea2d0dbcb3f6120e53 Mon Sep 17 00:00:00 2001
From: atumanov <atumanov@gmail.com>
Date: Tue, 29 Nov 2016 23:04:13 -0800
Subject: [PATCH] handling partial/interrupted I/O (#69)

* read_bytes: handle EINTR + easier flow

* addressing whitespace + style comments

* handling partial i/o: write_bytes; use [read|write]_bytes in plasma send/receive
---
 src/common/io.c     | 48 ++++++++++++++++++++++++---------------------
 src/common/io.h     |  2 ++
 src/plasma/plasma.c |  8 ++------
 3 files changed, 30 insertions(+), 28 deletions(-)

diff --git a/src/common/io.c b/src/common/io.c
index fdfbc5a1d..073712e02 100644
--- a/src/common/io.c
+++ b/src/common/io.c
@@ -174,24 +174,26 @@ int accept_client(int socket_fd) {
  */
 int write_bytes(int fd, uint8_t *cursor, size_t length) {
   ssize_t nbytes = 0;
-  while (length > 0) {
+  size_t bytesleft = length;
+  size_t offset = 0;
+  while (bytesleft > 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);
+    nbytes = write(fd, cursor + offset, bytesleft);
     if (nbytes < 0) {
-      if (errno == EAGAIN || errno == EWOULDBLOCK) {
+      if (errno == EAGAIN || errno == EWOULDBLOCK || errno == EINTR) {
         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; /* Errno will be set. */
+    } else if (0 == nbytes) {
+      /* Encountered early EOF. */
       return -1;
     }
-    cursor += nbytes;
-    length -= nbytes;
+    CHECK(nbytes > 0);
+    bytesleft -= nbytes;
+    offset += nbytes;
   }
+
   return 0;
 }
 
@@ -234,28 +236,30 @@ int write_message(int fd, int64_t type, int64_t length, uint8_t *bytes) {
  * @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 int Whether there was an error while writing. 0 corresponds to
+ * @return int Whether there was an error while reading. 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;
-  while (length > 0) {
-    /* While we haven't read the whole message, read from the file descriptor,
-     * advance the cursor, and decrease the amount left to read. */
-    nbytes = read(fd, cursor, length);
+  /* Termination condition: EOF or read 'length' bytes total. */
+  size_t bytesleft = length;
+  size_t offset = 0;
+  while (bytesleft > 0) {
+    nbytes = read(fd, cursor + offset, bytesleft);
     if (nbytes < 0) {
-      if (errno == EAGAIN || errno == EWOULDBLOCK) {
+      if (errno == EAGAIN || errno == EWOULDBLOCK || errno == EINTR) {
         continue;
       }
-      /* Force an exit if there was any other type of error. */
-      CHECK(nbytes < 0);
-    }
-    if (nbytes == 0) {
+      return -1;
+    } else if (0 == nbytes) {
+      /* Encountered early EOF. */
       return -1;
     }
-    cursor += nbytes;
-    length -= nbytes;
+    CHECK(nbytes > 0);
+    bytesleft -= nbytes;
+    offset += nbytes;
   }
+
   return 0;
 }
 
diff --git a/src/common/io.h b/src/common/io.h
index efde81546..d9a7e1025 100644
--- a/src/common/io.h
+++ b/src/common/io.h
@@ -32,5 +32,7 @@ int64_t read_buffer(int fd, int64_t *type, UT_array *buffer);
 void write_log_message(int fd, char *message);
 void write_formatted_log_message(int fd, const char *format, ...);
 char *read_log_message(int fd);
+int read_bytes(int fd, uint8_t *cursor, size_t length);
+int write_bytes(int fd, uint8_t *cursor, size_t length);
 
 #endif /* IO_H */
diff --git a/src/plasma/plasma.c b/src/plasma/plasma.c
index b9eb9350d..ac9f09188 100644
--- a/src/plasma/plasma.c
+++ b/src/plasma/plasma.c
@@ -62,15 +62,11 @@ int64_t plasma_reply_size(int num_object_ids) {
 int plasma_send_reply(int sock, plasma_reply *reply) {
   DCHECK(reply);
   int64_t reply_size = plasma_reply_size(reply->num_object_ids);
-  int n = write(sock, (uint8_t *) reply, reply_size);
-  return n == reply_size ? 0 : -1;
+  return write_bytes(sock, (uint8_t *) reply, reply_size);
 }
 
 int plasma_receive_reply(int sock, int64_t reply_size, plasma_reply *reply) {
-  int r = recv(sock, reply, reply_size, 0);
-  CHECKM(r != -1, "read error");
-  CHECKM(r != 0, "connection disconnected");
-  return r == reply_size ? 0 : -1;
+  return read_bytes(sock, (uint8_t *) reply, reply_size);
 }
 
 int plasma_send_request(int sock, int64_t type, plasma_request *request) {