Add coroutine-aware Read and Write function to FileDescriptor

Author: dallison

toolbelt/fd.cc

@@ -8,11 +8,96 @@ void CloseAllFds(std::function<bool(int)> predicate) {
   int e = getrlimit(RLIMIT_NOFILE, &lim);
   if (e == 0) {
     for (int fd = 0; fd < lim.rlim_cur; fd++) {
-      if (fcntl(fd, F_GETFD) == 0 && predicate(fd) ) {
+      if (fcntl(fd, F_GETFD) == 0 && predicate(fd)) {
         (void)close(fd);
       }
     }
   }
 }
 
-}
+absl::StatusOr<ssize_t> FileDescriptor::Read(void *buffer, size_t length,
+                                             co::Coroutine *c) {
+  char *buf = reinterpret_cast<char *>(buffer);
+  size_t total = 0;
+  while (total < length) {
+    if (c != nullptr) {
+      // For a coroutine we need to get it to wait for the fd to be ready.
+      // This is a coroutine yield point.
+      int fd = c->Wait(data_->fd, POLLIN);
+      if (fd != Fd()) {
+        return absl::InternalError(
+            absl::StrFormat("Unexpected file descriptor from Wait: %d", fd));
+      }
+    }
+    ssize_t n = ::read(Fd(), buf + total, length - total);
+    if (n == 0) {
+      break;
+    }
+    if (n == -1) {
+      if (errno == EAGAIN || errno == EWOULDBLOCK) {
+        if (c == nullptr) {
+          return absl::InternalError("Operation would block");
+        }
+        // If we are nonblocking yield the coroutine now.  When we
+        // are resumed we can write to the socket again.
+        if (!data_->non_blocking_) {
+          int fd = c->Wait(Fd(), POLLIN);
+          if (fd != Fd()) {
+            return absl::InternalError(absl::StrFormat(
+                "Unexpected file descriptor from Wait: %d", fd));
+          }
+        }
+        continue;
+      }
+      return absl::InternalError(
+          absl::StrFormat("Read failed: %s", strerror(errno)));
+    }
+    total += n;
+  }
+  return total;
+}
+
+absl::StatusOr<ssize_t> FileDescriptor::Write(const void *buffer, size_t length,
+                                              co::Coroutine *c) {
+  const char *buf = reinterpret_cast<const char *>(buffer);
+
+  size_t total = 0;
+  while (total < length) {
+    if (c != nullptr) {
+      // For a coroutine we need to get it to wait for the fd to be ready.
+      // This is a coroutine yield point.
+      int fd = c->Wait(data_->fd, POLLOUT);
+      if (fd != Fd()) {
+        return absl::InternalError(
+            absl::StrFormat("Unexpected file descriptor from Wait: %d", fd));
+      }
+    }
+    ssize_t n = ::write(Fd(), buf + total, length - total);
+    if (n == 0) {
+      break;
+    }
+    if (n == -1) {
+      if (errno == EAGAIN || errno == EWOULDBLOCK) {
+        if (c == nullptr) {
+          return absl::InternalError("Operation would block");
+        }
+        // If we are nonblocking yield the coroutine now.  When we
+        // are resumed we can write to the socket again.
+        if (!data_->non_blocking_) {
+          int fd = c->Wait(Fd(), POLLOUT);
+          if (fd != Fd()) {
+            return absl::InternalError(absl::StrFormat(
+                "Unexpected file descriptor from Wait: %d", fd));
+          }
+        }
+        continue;
+      }
+      return absl::InternalError(
+          absl::StrFormat("Write failed: %s", strerror(errno)));
+    }
+    total += n;
+  }
+  return total;
+}
+
+} // namespace toolbelt

toolbelt/fd.h

@@ -6,6 +6,7 @@
 #define __TOOLBELT_FD_H
 
 #include "absl/status/status.h"
+#include "absl/status/statusor.h"
 #include "absl/strings/str_format.h"
 #include <cassert>
 #include <cerrno>
@@ -17,6 +18,7 @@
 #include <sys/resource.h>
 #include <sys/stat.h>
 #include <unistd.h>
+#include "coroutine.h"
 
 namespace toolbelt {
 
@@ -160,6 +162,7 @@ class FileDescriptor {
       return absl::InternalError(absl::StrFormat(
           "Failed to set nonblocking mode on fd: %s", strerror(errno)));
     }
+    data_->non_blocking_ = true;
     return absl::OkStatus();
   }
 
@@ -180,6 +183,9 @@ class FileDescriptor {
     return absl::OkStatus();
   }
 
+  absl::StatusOr<ssize_t> Read(void* buffer, size_t length, co::Coroutine* c = nullptr);
+  absl::StatusOr<ssize_t> Write(const void* buffer, size_t length,
+                                      co::Coroutine* c = nullptr);
 private:
   // Reference counted OS fd, shared among all FileDescriptors with the
   // same OS fd, provided you don't create two FileDescriptors with the
@@ -193,6 +199,7 @@ class FileDescriptor {
       }
     }
     int fd = -1; // OS file descriptor.
+    bool non_blocking_ = false;
   };
 
   // The actual shared data.  If nullptr the FileDescriptor is invalid.