Make client process all buffered reads on write failure

Summary: When connection is reset by server it's possible for a write to fail before the read callback is called. Make sure we still try to process all the data that was received before closing the socket.

Reviewed By: rhodo

Differential Revision: D19851139

fbshipit-source-id: af50b939ccd7f25384e48bcf29b4892a23c2ef83

Author: Andrii Grynenko

folly/io/async/AsyncSocket.cpp

@@ -2747,13 +2747,17 @@ void AsyncSocket::failWrite(
   VLOG(4) << "AsyncSocket(this=" << this << ", fd=" << fd_
           << ", state=" << state_ << " host=" << addr_.describe()
           << "): failed while writing in " << fn << "(): " << ex.what();
-  startFail();
+  if (closeOnFailedWrite_) {
+    startFail();
+  }
 
   if (callback != nullptr) {
     callback->writeErr(bytesWritten, ex);
   }
 
-  finishFail();
+  if (closeOnFailedWrite_) {
+    finishFail();
+  }
 }
 
 void AsyncSocket::failAllWrites(const AsyncSocketException& ex) {

folly/io/async/AsyncSocket.h

@@ -905,6 +905,13 @@ class AsyncSocket : virtual public AsyncTransportWrapper {
     return selfCertData_.get();
   }
 
+  /**
+   * Whether socket should be closed on write failure (true by default).
+   */
+  void setCloseOnFailedWrite(bool closeOnFailedWrite) {
+    closeOnFailedWrite_ = closeOnFailedWrite;
+  }
+
   /**
    * writeReturn is the total number of bytes written, or WRITE_ERROR on error.
    * If no data has been written, 0 is returned.
@@ -1349,6 +1356,8 @@ class AsyncSocket : virtual public AsyncTransportWrapper {
       nullptr};
   mutable std::unique_ptr<const AsyncTransportCertificate> selfCertData_{
       nullptr};
+
+  bool closeOnFailedWrite_{true};
 };
 #ifdef _MSC_VER
 #pragma vtordisp(pop)