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

diff --git a/folly/io/async/AsyncSocket.cpp b/folly/io/async/AsyncSocket.cpp
index dc550a3..6084891 100644 (file)
--- 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 f8aa27f..d4ddf0b 100644 (file)
--- 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 0c46ae6..3d023c4 100644 (file)
--- 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 90396f4..20baf7b 100644 (file)
--- 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 fb68b2a..1b48d26 100644 (file)
--- 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
  */