From f3bbb0b1faa8f430a41dc2f49b10d56af9851e89 Mon Sep 17 00:00:00 2001
From: Kevin Chen <kevinch@fb.com>
Date: Wed, 10 May 2017 19:42:00 -0700
Subject: [PATCH] Add writable() to AsyncTransport

Summary:
This is useful for checking if it's possible to still write to a transport,
even if its read side is closed (for transports that support half shutdown,
like AsyncSocket). Default implementation just returns true for now (up
to implementers to override).

Reviewed By: yfeldblum

Differential Revision: D4982649

fbshipit-source-id: 0a9a2e2b745ea3db57e9f151f3a8634e1bda2465
---
 folly/io/async/AsyncSocket.cpp                | 12 ++++++
 folly/io/async/AsyncSocket.h                  |  1 +
 folly/io/async/AsyncTransport.h               | 10 +++++
 .../io/async/DecoratedAsyncTransportWrapper.h |  4 ++
 folly/io/async/test/AsyncSocketTest2.cpp      | 38 +++++++++++++++++++
 5 files changed, 65 insertions(+)

diff --git a/folly/io/async/AsyncSocket.cpp b/folly/io/async/AsyncSocket.cpp
index dc550a31..60848917 100644
--- a/folly/io/async/AsyncSocket.cpp
+++ b/folly/io/async/AsyncSocket.cpp
@@ -1189,6 +1189,18 @@ bool AsyncSocket::readable() const {
   return rc == 1;
 }
 
+bool AsyncSocket::writable() const {
+  if (fd_ == -1) {
+    return false;
+  }
+  struct pollfd fds[1];
+  fds[0].fd = fd_;
+  fds[0].events = POLLOUT;
+  fds[0].revents = 0;
+  int rc = poll(fds, 1, 0);
+  return rc == 1;
+}
+
 bool AsyncSocket::isPending() const {
   return ioHandler_.isPending();
 }
diff --git a/folly/io/async/AsyncSocket.h b/folly/io/async/AsyncSocket.h
index f8aa27f6..d4ddf0be 100644
--- a/folly/io/async/AsyncSocket.h
+++ b/folly/io/async/AsyncSocket.h
@@ -522,6 +522,7 @@ class AsyncSocket : virtual public AsyncTransportWrapper {
   void shutdownWriteNow() override;
 
   bool readable() const override;
+  bool writable() const override;
   bool isPending() const override;
   virtual bool hangup() const;
   bool good() const override;
diff --git a/folly/io/async/AsyncTransport.h b/folly/io/async/AsyncTransport.h
index 0c46ae64..3d023c45 100644
--- a/folly/io/async/AsyncTransport.h
+++ b/folly/io/async/AsyncTransport.h
@@ -238,6 +238,16 @@ class AsyncTransport : public DelayedDestruction, public AsyncSocketBase {
    */
   virtual bool readable() const = 0;
 
+  /**
+   * Determine if the transport is writable or not.
+   *
+   * @return  true iff the transport is writable, false otherwise.
+   */
+  virtual bool writable() const {
+    // By default return good() - leave it to implementers to override.
+    return good();
+  }
+
   /**
    * Determine if the there is pending data on the transport.
    *
diff --git a/folly/io/async/DecoratedAsyncTransportWrapper.h b/folly/io/async/DecoratedAsyncTransportWrapper.h
index 90396f4f..20baf7bf 100644
--- a/folly/io/async/DecoratedAsyncTransportWrapper.h
+++ b/folly/io/async/DecoratedAsyncTransportWrapper.h
@@ -152,6 +152,10 @@ class DecoratedAsyncTransportWrapper : public folly::AsyncTransportWrapper {
     return transport_->readable();
   }
 
+  virtual bool writable() const override {
+    return transport_->writable();
+  }
+
   virtual void setEorTracking(bool track) override {
     return transport_->setEorTracking(track);
   }
diff --git a/folly/io/async/test/AsyncSocketTest2.cpp b/folly/io/async/test/AsyncSocketTest2.cpp
index fb68b2a6..1b48d260 100644
--- a/folly/io/async/test/AsyncSocketTest2.cpp
+++ b/folly/io/async/test/AsyncSocketTest2.cpp
@@ -1096,6 +1096,44 @@ TEST(AsyncSocketTest, WritePipeError) {
   ASSERT_FALSE(socket->isClosedByPeer());
 }
 
+/**
+ * Test writing to a socket that has its read side closed
+ */
+TEST(AsyncSocketTest, WriteAfterReadEOF) {
+  TestServer server;
+
+  // connect()
+  EventBase evb;
+  std::shared_ptr<AsyncSocket> socket =
+      AsyncSocket::newSocket(&evb, server.getAddress(), 30);
+  evb.loop(); // loop until the socket is connected
+
+  // Accept the connection
+  std::shared_ptr<AsyncSocket> acceptedSocket = server.acceptAsync(&evb);
+  ReadCallback rcb;
+  acceptedSocket->setReadCB(&rcb);
+
+  // Shutdown the write side of client socket (read side of server socket)
+  socket->shutdownWrite();
+  evb.loop();
+
+  // Check that accepted socket is still writable
+  ASSERT_FALSE(acceptedSocket->good());
+  ASSERT_TRUE(acceptedSocket->writable());
+
+  // Write data to accepted socket
+  constexpr size_t simpleBufLength = 5;
+  char simpleBuf[simpleBufLength];
+  memset(simpleBuf, 'a', simpleBufLength);
+  WriteCallback wcb;
+  acceptedSocket->write(&wcb, simpleBuf, simpleBufLength);
+  evb.loop();
+
+  // Make sure we were able to write even after getting a read EOF
+  ASSERT_EQ(rcb.state, STATE_SUCCEEDED); // this indicates EOF
+  ASSERT_EQ(wcb.state, STATE_SUCCEEDED);
+}
+
 /**
  * Test that bytes written is correctly computed in case of write failure
  */
-- 
2.34.1