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 dc550a310e61ca9f8fdecaa7b69cdf72b61f0a6d..60848917f80a5b9d9bc51320fb050a4d3d82cde2 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 f8aa27f6019fa4955b6d90e4af789a42fcc08e86..d4ddf0becaf2d0254ff5dd2b88045b539cd9d2a3 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 0c46ae640e5f290b72515e0ac038a0c298ac2ef7..3d023c450d545d3cff7e2c6a3005b67b799f4ad9 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 90396f4f9d88d5658b7e72b6601cb59445f5f667..20baf7bfe3fab6292109feddadbcc00939668b62 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 fb68b2a6d47a6c9e968611465a6cbf77c7bf3366..1b48d260697660f07b26c7cf852b8e39ec2cacf9 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
  */