folly/io/async/AsyncSocket.cpp

   1 /*
   2  * Copyright 2016 Facebook, Inc.
   3  *
   4  * Licensed under the Apache License, Version 2.0 (the "License");
   5  * you may not use this file except in compliance with the License.
   6  * You may obtain a copy of the License at
   7  *
   8  *   http://www.apache.org/licenses/LICENSE-2.0
   9  *
  10  * Unless required by applicable law or agreed to in writing, software
  11  * distributed under the License is distributed on an "AS IS" BASIS,
  12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  13  * See the License for the specific language governing permissions and
  14  * limitations under the License.
  15  */
  16
  17 #include <folly/io/async/AsyncSocket.h>
  18
  19 #include <folly/ExceptionWrapper.h>
  20 #include <folly/SocketAddress.h>
  21 #include <folly/io/IOBuf.h>
  22 #include <folly/Portability.h>
  23 #include <folly/portability/Fcntl.h>
  24 #include <folly/portability/Sockets.h>
  25 #include <folly/portability/SysUio.h>
  26 #include <folly/portability/Unistd.h>
  27
  28 #include <errno.h>
  29 #include <limits.h>
  30 #include <thread>
  31 #include <sys/types.h>
  32 #include <boost/preprocessor/control/if.hpp>
  33
  34 using std::string;
  35 using std::unique_ptr;
  36
  37 namespace fsp = folly::portability::sockets;
  38
  39 namespace folly {
  40
  41 // static members initializers
  42 const AsyncSocket::OptionMap AsyncSocket::emptyOptionMap;
  43
  44 const AsyncSocketException socketClosedLocallyEx(
  45     AsyncSocketException::END_OF_FILE, "socket closed locally");
  46 const AsyncSocketException socketShutdownForWritesEx(
  47     AsyncSocketException::END_OF_FILE, "socket shutdown for writes");
  48
  49 // TODO: It might help performance to provide a version of BytesWriteRequest that
  50 // users could derive from, so we can avoid the extra allocation for each call
  51 // to write()/writev().  We could templatize TFramedAsyncChannel just like the
  52 // protocols are currently templatized for transports.
  53 //
  54 // We would need the version for external users where they provide the iovec
  55 // storage space, and only our internal version would allocate it at the end of
  56 // the WriteRequest.
  57
  58 /* The default WriteRequest implementation, used for write(), writev() and
  59  * writeChain()
  60  *
  61  * A new BytesWriteRequest operation is allocated on the heap for all write
  62  * operations that cannot be completed immediately.
  63  */
  64 class AsyncSocket::BytesWriteRequest : public AsyncSocket::WriteRequest {
  65  public:
  66   static BytesWriteRequest* newRequest(AsyncSocket* socket,
  67                                        WriteCallback* callback,
  68                                        const iovec* ops,
  69                                        uint32_t opCount,
  70                                        uint32_t partialWritten,
  71                                        uint32_t bytesWritten,
  72                                        unique_ptr<IOBuf>&& ioBuf,
  73                                        WriteFlags flags) {
  74     assert(opCount > 0);
  75     // Since we put a variable size iovec array at the end
  76     // of each BytesWriteRequest, we have to manually allocate the memory.
  77     void* buf = malloc(sizeof(BytesWriteRequest) +
  78                        (opCount * sizeof(struct iovec)));
  79     if (buf == nullptr) {
  80       throw std::bad_alloc();
  81     }
  82
  83     return new(buf) BytesWriteRequest(socket, callback, ops, opCount,
  84                                       partialWritten, bytesWritten,
  85                                       std::move(ioBuf), flags);
  86   }
  87
  88   void destroy() override {
  89     this->~BytesWriteRequest();
  90     free(this);
  91   }
  92
  93   WriteResult performWrite() override {
  94     WriteFlags writeFlags = flags_;
  95     if (getNext() != nullptr) {
  96       writeFlags = writeFlags | WriteFlags::CORK;
  97     }
  98     auto writeResult = socket_->performWrite(
  99         getOps(), getOpCount(), writeFlags, &opsWritten_, &partialBytes_);
 100     bytesWritten_ = writeResult.writeReturn > 0 ? writeResult.writeReturn : 0;
 101     return writeResult;
 102   }
 103
 104   bool isComplete() override {
 105     return opsWritten_ == getOpCount();
 106   }
 107
 108   void consume() override {
 109     // Advance opIndex_ forward by opsWritten_
 110     opIndex_ += opsWritten_;
 111     assert(opIndex_ < opCount_);
 112
 113     // If we've finished writing any IOBufs, release them
 114     if (ioBuf_) {
 115       for (uint32_t i = opsWritten_; i != 0; --i) {
 116         assert(ioBuf_);
 117         ioBuf_ = ioBuf_->pop();
 118       }
 119     }
 120
 121     // Move partialBytes_ forward into the current iovec buffer
 122     struct iovec* currentOp = writeOps_ + opIndex_;
 123     assert((partialBytes_ < currentOp->iov_len) || (currentOp->iov_len == 0));
 124     currentOp->iov_base =
 125       reinterpret_cast<uint8_t*>(currentOp->iov_base) + partialBytes_;
 126     currentOp->iov_len -= partialBytes_;
 127
 128     // Increment the totalBytesWritten_ count by bytesWritten_;
 129     totalBytesWritten_ += bytesWritten_;
 130   }
 131
 132  private:
 133   BytesWriteRequest(AsyncSocket* socket,
 134                     WriteCallback* callback,
 135                     const struct iovec* ops,
 136                     uint32_t opCount,
 137                     uint32_t partialBytes,
 138                     uint32_t bytesWritten,
 139                     unique_ptr<IOBuf>&& ioBuf,
 140                     WriteFlags flags)
 141     : AsyncSocket::WriteRequest(socket, callback)
 142     , opCount_(opCount)
 143     , opIndex_(0)
 144     , flags_(flags)
 145     , ioBuf_(std::move(ioBuf))
 146     , opsWritten_(0)
 147     , partialBytes_(partialBytes)
 148     , bytesWritten_(bytesWritten) {
 149     memcpy(writeOps_, ops, sizeof(*ops) * opCount_);
 150   }
 151
 152   // private destructor, to ensure callers use destroy()
 153   ~BytesWriteRequest() override = default;
 154
 155   const struct iovec* getOps() const {
 156     assert(opCount_ > opIndex_);
 157     return writeOps_ + opIndex_;
 158   }
 159
 160   uint32_t getOpCount() const {
 161     assert(opCount_ > opIndex_);
 162     return opCount_ - opIndex_;
 163   }
 164
 165   uint32_t opCount_;            ///< number of entries in writeOps_
 166   uint32_t opIndex_;            ///< current index into writeOps_
 167   WriteFlags flags_;            ///< set for WriteFlags
 168   unique_ptr<IOBuf> ioBuf_;     ///< underlying IOBuf, or nullptr if N/A
 169
 170   // for consume(), how much we wrote on the last write
 171   uint32_t opsWritten_;         ///< complete ops written
 172   uint32_t partialBytes_;       ///< partial bytes of incomplete op written
 173   ssize_t bytesWritten_;        ///< bytes written altogether
 174
 175   struct iovec writeOps_[];     ///< write operation(s) list
 176 };
 177
 178 AsyncSocket::AsyncSocket()
 179     : eventBase_(nullptr),
 180       writeTimeout_(this, nullptr),
 181       ioHandler_(this, nullptr),
 182       immediateReadHandler_(this) {
 183   VLOG(5) << "new AsyncSocket()";
 184   init();
 185 }
 186
 187 AsyncSocket::AsyncSocket(EventBase* evb)
 188     : eventBase_(evb),
 189       writeTimeout_(this, evb),
 190       ioHandler_(this, evb),
 191       immediateReadHandler_(this) {
 192   VLOG(5) << "new AsyncSocket(" << this << ", evb=" << evb << ")";
 193   init();
 194 }
 195
 196 AsyncSocket::AsyncSocket(EventBase* evb,
 197                            const folly::SocketAddress& address,
 198                            uint32_t connectTimeout)
 199   : AsyncSocket(evb) {
 200   connect(nullptr, address, connectTimeout);
 201 }
 202
 203 AsyncSocket::AsyncSocket(EventBase* evb,
 204                            const std::string& ip,
 205                            uint16_t port,
 206                            uint32_t connectTimeout)
 207   : AsyncSocket(evb) {
 208   connect(nullptr, ip, port, connectTimeout);
 209 }
 210
 211 AsyncSocket::AsyncSocket(EventBase* evb, int fd)
 212     : eventBase_(evb),
 213       writeTimeout_(this, evb),
 214       ioHandler_(this, evb, fd),
 215       immediateReadHandler_(this) {
 216   VLOG(5) << "new AsyncSocket(" << this << ", evb=" << evb << ", fd="
 217           << fd << ")";
 218   init();
 219   fd_ = fd;
 220   setCloseOnExec();
 221   state_ = StateEnum::ESTABLISHED;
 222 }
 223
 224 // init() method, since constructor forwarding isn't supported in most
 225 // compilers yet.
 226 void AsyncSocket::init() {
 227   assert(eventBase_ == nullptr || eventBase_->isInEventBaseThread());
 228   shutdownFlags_ = 0;
 229   state_ = StateEnum::UNINIT;
 230   eventFlags_ = EventHandler::NONE;
 231   fd_ = -1;
 232   sendTimeout_ = 0;
 233   maxReadsPerEvent_ = 16;
 234   connectCallback_ = nullptr;
 235   readCallback_ = nullptr;
 236   writeReqHead_ = nullptr;
 237   writeReqTail_ = nullptr;
 238   shutdownSocketSet_ = nullptr;
 239   appBytesWritten_ = 0;
 240   appBytesReceived_ = 0;
 241 }
 242
 243 AsyncSocket::~AsyncSocket() {
 244   VLOG(7) << "actual destruction of AsyncSocket(this=" << this
 245           << ", evb=" << eventBase_ << ", fd=" << fd_
 246           << ", state=" << state_ << ")";
 247 }
 248
 249 void AsyncSocket::destroy() {
 250   VLOG(5) << "AsyncSocket::destroy(this=" << this << ", evb=" << eventBase_
 251           << ", fd=" << fd_ << ", state=" << state_;
 252   // When destroy is called, close the socket immediately
 253   closeNow();
 254
 255   // Then call DelayedDestruction::destroy() to take care of
 256   // whether or not we need immediate or delayed destruction
 257   DelayedDestruction::destroy();
 258 }
 259
 260 int AsyncSocket::detachFd() {
 261   VLOG(6) << "AsyncSocket::detachFd(this=" << this << ", fd=" << fd_
 262           << ", evb=" << eventBase_ << ", state=" << state_
 263           << ", events=" << std::hex << eventFlags_ << ")";
 264   // Extract the fd, and set fd_ to -1 first, so closeNow() won't
 265   // actually close the descriptor.
 266   if (shutdownSocketSet_) {
 267     shutdownSocketSet_->remove(fd_);
 268   }
 269   int fd = fd_;
 270   fd_ = -1;
 271   // Call closeNow() to invoke all pending callbacks with an error.
 272   closeNow();
 273   // Update the EventHandler to stop using this fd.
 274   // This can only be done after closeNow() unregisters the handler.
 275   ioHandler_.changeHandlerFD(-1);
 276   return fd;
 277 }
 278
 279 const folly::SocketAddress& AsyncSocket::anyAddress() {
 280   static const folly::SocketAddress anyAddress =
 281     folly::SocketAddress("0.0.0.0", 0);
 282   return anyAddress;
 283 }
 284
 285 void AsyncSocket::setShutdownSocketSet(ShutdownSocketSet* newSS) {
 286   if (shutdownSocketSet_ == newSS) {
 287     return;
 288   }
 289   if (shutdownSocketSet_ && fd_ != -1) {
 290     shutdownSocketSet_->remove(fd_);
 291   }
 292   shutdownSocketSet_ = newSS;
 293   if (shutdownSocketSet_ && fd_ != -1) {
 294     shutdownSocketSet_->add(fd_);
 295   }
 296 }
 297
 298 void AsyncSocket::setCloseOnExec() {
 299   int rv = fcntl(fd_, F_SETFD, FD_CLOEXEC);
 300   if (rv != 0) {
 301     auto errnoCopy = errno;
 302     throw AsyncSocketException(
 303         AsyncSocketException::INTERNAL_ERROR,
 304         withAddr("failed to set close-on-exec flag"),
 305         errnoCopy);
 306   }
 307 }
 308
 309 void AsyncSocket::connect(ConnectCallback* callback,
 310                            const folly::SocketAddress& address,
 311                            int timeout,
 312                            const OptionMap &options,
 313                            const folly::SocketAddress& bindAddr) noexcept {
 314   DestructorGuard dg(this);
 315   assert(eventBase_->isInEventBaseThread());
 316
 317   addr_ = address;
 318
 319   // Make sure we're in the uninitialized state
 320   if (state_ != StateEnum::UNINIT) {
 321     return invalidState(callback);
 322   }
 323
 324   connectTimeout_ = std::chrono::milliseconds(timeout);
 325   connectStartTime_ = std::chrono::steady_clock::now();
 326   // Make connect end time at least >= connectStartTime.
 327   connectEndTime_ = connectStartTime_;
 328
 329   assert(fd_ == -1);
 330   state_ = StateEnum::CONNECTING;
 331   connectCallback_ = callback;
 332
 333   sockaddr_storage addrStorage;
 334   sockaddr* saddr = reinterpret_cast<sockaddr*>(&addrStorage);
 335
 336   try {
 337     // Create the socket
 338     // Technically the first parameter should actually be a protocol family
 339     // constant (PF_xxx) rather than an address family (AF_xxx), but the
 340     // distinction is mainly just historical.  In pretty much all
 341     // implementations the PF_foo and AF_foo constants are identical.
 342     fd_ = fsp::socket(address.getFamily(), SOCK_STREAM, 0);
 343     if (fd_ < 0) {
 344       auto errnoCopy = errno;
 345       throw AsyncSocketException(
 346           AsyncSocketException::INTERNAL_ERROR,
 347           withAddr("failed to create socket"),
 348           errnoCopy);
 349     }
 350     if (shutdownSocketSet_) {
 351       shutdownSocketSet_->add(fd_);
 352     }
 353     ioHandler_.changeHandlerFD(fd_);
 354
 355     setCloseOnExec();
 356
 357     // Put the socket in non-blocking mode
 358     int flags = fcntl(fd_, F_GETFL, 0);
 359     if (flags == -1) {
 360       auto errnoCopy = errno;
 361       throw AsyncSocketException(
 362           AsyncSocketException::INTERNAL_ERROR,
 363           withAddr("failed to get socket flags"),
 364           errnoCopy);
 365     }
 366     int rv = fcntl(fd_, F_SETFL, flags | O_NONBLOCK);
 367     if (rv == -1) {
 368       auto errnoCopy = errno;
 369       throw AsyncSocketException(
 370           AsyncSocketException::INTERNAL_ERROR,
 371           withAddr("failed to put socket in non-blocking mode"),
 372           errnoCopy);
 373     }
 374
 375 #if !defined(MSG_NOSIGNAL) && defined(F_SETNOSIGPIPE)
 376     // iOS and OS X don't support MSG_NOSIGNAL; set F_SETNOSIGPIPE instead
 377     rv = fcntl(fd_, F_SETNOSIGPIPE, 1);
 378     if (rv == -1) {
 379       auto errnoCopy = errno;
 380       throw AsyncSocketException(
 381           AsyncSocketException::INTERNAL_ERROR,
 382           "failed to enable F_SETNOSIGPIPE on socket",
 383           errnoCopy);
 384     }
 385 #endif
 386
 387     // By default, turn on TCP_NODELAY
 388     // If setNoDelay() fails, we continue anyway; this isn't a fatal error.
 389     // setNoDelay() will log an error message if it fails.
 390     if (address.getFamily() != AF_UNIX) {
 391       (void)setNoDelay(true);
 392     }
 393
 394     VLOG(5) << "AsyncSocket::connect(this=" << this << ", evb=" << eventBase_
 395             << ", fd=" << fd_ << ", host=" << address.describe().c_str();
 396
 397     // bind the socket
 398     if (bindAddr != anyAddress()) {
 399       int one = 1;
 400       if (setsockopt(fd_, SOL_SOCKET, SO_REUSEADDR, &one, sizeof(one))) {
 401         auto errnoCopy = errno;
 402         doClose();
 403         throw AsyncSocketException(
 404             AsyncSocketException::NOT_OPEN,
 405             "failed to setsockopt prior to bind on " + bindAddr.describe(),
 406             errnoCopy);
 407       }
 408
 409       bindAddr.getAddress(&addrStorage);
 410
 411       if (bind(fd_, saddr, bindAddr.getActualSize()) != 0) {
 412         auto errnoCopy = errno;
 413         doClose();
 414         throw AsyncSocketException(
 415             AsyncSocketException::NOT_OPEN,
 416             "failed to bind to async socket: " + bindAddr.describe(),
 417             errnoCopy);
 418       }
 419     }
 420
 421     // Apply the additional options if any.
 422     for (const auto& opt: options) {
 423       rv = opt.first.apply(fd_, opt.second);
 424       if (rv != 0) {
 425         auto errnoCopy = errno;
 426         throw AsyncSocketException(
 427             AsyncSocketException::INTERNAL_ERROR,
 428             withAddr("failed to set socket option"),
 429             errnoCopy);
 430       }
 431     }
 432
 433     // Perform the connect()
 434     address.getAddress(&addrStorage);
 435
 436     if (tfoEnabled_) {
 437       state_ = StateEnum::FAST_OPEN;
 438       tfoAttempted_ = true;
 439     } else {
 440       if (socketConnect(saddr, addr_.getActualSize()) < 0) {
 441         return;
 442       }
 443     }
 444
 445     // If we're still here the connect() succeeded immediately.
 446     // Fall through to call the callback outside of this try...catch block
 447   } catch (const AsyncSocketException& ex) {
 448     return failConnect(__func__, ex);
 449   } catch (const std::exception& ex) {
 450     // shouldn't happen, but handle it just in case
 451     VLOG(4) << "AsyncSocket::connect(this=" << this << ", fd=" << fd_
 452                << "): unexpected " << typeid(ex).name() << " exception: "
 453                << ex.what();
 454     AsyncSocketException tex(AsyncSocketException::INTERNAL_ERROR,
 455                             withAddr(string("unexpected exception: ") +
 456                                      ex.what()));
 457     return failConnect(__func__, tex);
 458   }
 459
 460   // The connection succeeded immediately
 461   // The read callback may not have been set yet, and no writes may be pending
 462   // yet, so we don't have to register for any events at the moment.
 463   VLOG(8) << "AsyncSocket::connect succeeded immediately; this=" << this;
 464   assert(readCallback_ == nullptr);
 465   assert(writeReqHead_ == nullptr);
 466   if (state_ != StateEnum::FAST_OPEN) {
 467     state_ = StateEnum::ESTABLISHED;
 468   }
 469   invokeConnectSuccess();
 470 }
 471
 472 int AsyncSocket::socketConnect(const struct sockaddr* saddr, socklen_t len) {
 473   int rv = fsp::connect(fd_, saddr, len);
 474   if (rv < 0) {
 475     auto errnoCopy = errno;
 476     if (errnoCopy == EINPROGRESS) {
 477       scheduleConnectTimeout();
 478       registerForConnectEvents();
 479     } else {
 480       throw AsyncSocketException(
 481           AsyncSocketException::NOT_OPEN,
 482           "connect failed (immediately)",
 483           errnoCopy);
 484     }
 485   }
 486   return rv;
 487 }
 488
 489 void AsyncSocket::scheduleConnectTimeout() {
 490   // Connection in progress.
 491   int timeout = connectTimeout_.count();
 492   if (timeout > 0) {
 493     // Start a timer in case the connection takes too long.
 494     if (!writeTimeout_.scheduleTimeout(timeout)) {
 495       throw AsyncSocketException(
 496           AsyncSocketException::INTERNAL_ERROR,
 497           withAddr("failed to schedule AsyncSocket connect timeout"));
 498     }
 499   }
 500 }
 501
 502 void AsyncSocket::registerForConnectEvents() {
 503   // Register for write events, so we'll
 504   // be notified when the connection finishes/fails.
 505   // Note that we don't register for a persistent event here.
 506   assert(eventFlags_ == EventHandler::NONE);
 507   eventFlags_ = EventHandler::WRITE;
 508   if (!ioHandler_.registerHandler(eventFlags_)) {
 509     throw AsyncSocketException(
 510         AsyncSocketException::INTERNAL_ERROR,
 511         withAddr("failed to register AsyncSocket connect handler"));
 512   }
 513 }
 514
 515 void AsyncSocket::connect(ConnectCallback* callback,
 516                            const string& ip, uint16_t port,
 517                            int timeout,
 518                            const OptionMap &options) noexcept {
 519   DestructorGuard dg(this);
 520   try {
 521     connectCallback_ = callback;
 522     connect(callback, folly::SocketAddress(ip, port), timeout, options);
 523   } catch (const std::exception& ex) {
 524     AsyncSocketException tex(AsyncSocketException::INTERNAL_ERROR,
 525                             ex.what());
 526     return failConnect(__func__, tex);
 527   }
 528 }
 529
 530 void AsyncSocket::cancelConnect() {
 531   connectCallback_ = nullptr;
 532   if (state_ == StateEnum::CONNECTING || state_ == StateEnum::FAST_OPEN) {
 533     closeNow();
 534   }
 535 }
 536
 537 void AsyncSocket::setSendTimeout(uint32_t milliseconds) {
 538   sendTimeout_ = milliseconds;
 539   assert(eventBase_ == nullptr || eventBase_->isInEventBaseThread());
 540
 541   // If we are currently pending on write requests, immediately update
 542   // writeTimeout_ with the new value.
 543   if ((eventFlags_ & EventHandler::WRITE) &&
 544       (state_ != StateEnum::CONNECTING && state_ != StateEnum::FAST_OPEN)) {
 545     assert(state_ == StateEnum::ESTABLISHED);
 546     assert((shutdownFlags_ & SHUT_WRITE) == 0);
 547     if (sendTimeout_ > 0) {
 548       if (!writeTimeout_.scheduleTimeout(sendTimeout_)) {
 549         AsyncSocketException ex(AsyncSocketException::INTERNAL_ERROR,
 550             withAddr("failed to reschedule send timeout in setSendTimeout"));
 551         return failWrite(__func__, ex);
 552       }
 553     } else {
 554       writeTimeout_.cancelTimeout();
 555     }
 556   }
 557 }
 558
 559 void AsyncSocket::setReadCB(ReadCallback *callback) {
 560   VLOG(6) << "AsyncSocket::setReadCallback() this=" << this << ", fd=" << fd_
 561           << ", callback=" << callback << ", state=" << state_;
 562
 563   // Short circuit if callback is the same as the existing readCallback_.
 564   //
 565   // Note that this is needed for proper functioning during some cleanup cases.
 566   // During cleanup we allow setReadCallback(nullptr) to be called even if the
 567   // read callback is already unset and we have been detached from an event
 568   // base.  This check prevents us from asserting
 569   // eventBase_->isInEventBaseThread() when eventBase_ is nullptr.
 570   if (callback == readCallback_) {
 571     return;
 572   }
 573
 574   /* We are removing a read callback */
 575   if (callback == nullptr &&
 576       immediateReadHandler_.isLoopCallbackScheduled()) {
 577     immediateReadHandler_.cancelLoopCallback();
 578   }
 579
 580   if (shutdownFlags_ & SHUT_READ) {
 581     // Reads have already been shut down on this socket.
 582     //
 583     // Allow setReadCallback(nullptr) to be called in this case, but don't
 584     // allow a new callback to be set.
 585     //
 586     // For example, setReadCallback(nullptr) can happen after an error if we
 587     // invoke some other error callback before invoking readError().  The other
 588     // error callback that is invoked first may go ahead and clear the read
 589     // callback before we get a chance to invoke readError().
 590     if (callback != nullptr) {
 591       return invalidState(callback);
 592     }
 593     assert((eventFlags_ & EventHandler::READ) == 0);
 594     readCallback_ = nullptr;
 595     return;
 596   }
 597
 598   DestructorGuard dg(this);
 599   assert(eventBase_->isInEventBaseThread());
 600
 601   switch ((StateEnum)state_) {
 602     case StateEnum::CONNECTING:
 603     case StateEnum::FAST_OPEN:
 604       // For convenience, we allow the read callback to be set while we are
 605       // still connecting.  We just store the callback for now.  Once the
 606       // connection completes we'll register for read events.
 607       readCallback_ = callback;
 608       return;
 609     case StateEnum::ESTABLISHED:
 610     {
 611       readCallback_ = callback;
 612       uint16_t oldFlags = eventFlags_;
 613       if (readCallback_) {
 614         eventFlags_ |= EventHandler::READ;
 615       } else {
 616         eventFlags_ &= ~EventHandler::READ;
 617       }
 618
 619       // Update our registration if our flags have changed
 620       if (eventFlags_ != oldFlags) {
 621         // We intentionally ignore the return value here.
 622         // updateEventRegistration() will move us into the error state if it
 623         // fails, and we don't need to do anything else here afterwards.
 624         (void)updateEventRegistration();
 625       }
 626
 627       if (readCallback_) {
 628         checkForImmediateRead();
 629       }
 630       return;
 631     }
 632     case StateEnum::CLOSED:
 633     case StateEnum::ERROR:
 634       // We should never reach here.  SHUT_READ should always be set
 635       // if we are in STATE_CLOSED or STATE_ERROR.
 636       assert(false);
 637       return invalidState(callback);
 638     case StateEnum::UNINIT:
 639       // We do not allow setReadCallback() to be called before we start
 640       // connecting.
 641       return invalidState(callback);
 642   }
 643
 644   // We don't put a default case in the switch statement, so that the compiler
 645   // will warn us to update the switch statement if a new state is added.
 646   return invalidState(callback);
 647 }
 648
 649 AsyncSocket::ReadCallback* AsyncSocket::getReadCallback() const {
 650   return readCallback_;
 651 }
 652
 653 void AsyncSocket::write(WriteCallback* callback,
 654                          const void* buf, size_t bytes, WriteFlags flags) {
 655   iovec op;
 656   op.iov_base = const_cast<void*>(buf);
 657   op.iov_len = bytes;
 658   writeImpl(callback, &op, 1, unique_ptr<IOBuf>(), flags);
 659 }
 660
 661 void AsyncSocket::writev(WriteCallback* callback,
 662                           const iovec* vec,
 663                           size_t count,
 664                           WriteFlags flags) {
 665   writeImpl(callback, vec, count, unique_ptr<IOBuf>(), flags);
 666 }
 667
 668 void AsyncSocket::writeChain(WriteCallback* callback, unique_ptr<IOBuf>&& buf,
 669                               WriteFlags flags) {
 670   constexpr size_t kSmallSizeMax = 64;
 671   size_t count = buf->countChainElements();
 672   if (count <= kSmallSizeMax) {
 673     // suppress "warning: variable length array 'vec' is used [-Wvla]"
 674     FOLLY_PUSH_WARNING;
 675     FOLLY_GCC_DISABLE_WARNING(vla);
 676     iovec vec[BOOST_PP_IF(FOLLY_HAVE_VLA, count, kSmallSizeMax)];
 677     FOLLY_POP_WARNING;
 678
 679     writeChainImpl(callback, vec, count, std::move(buf), flags);
 680   } else {
 681     iovec* vec = new iovec[count];
 682     writeChainImpl(callback, vec, count, std::move(buf), flags);
 683     delete[] vec;
 684   }
 685 }
 686
 687 void AsyncSocket::writeChainImpl(WriteCallback* callback, iovec* vec,
 688     size_t count, unique_ptr<IOBuf>&& buf, WriteFlags flags) {
 689   size_t veclen = buf->fillIov(vec, count);
 690   writeImpl(callback, vec, veclen, std::move(buf), flags);
 691 }
 692
 693 void AsyncSocket::writeImpl(WriteCallback* callback, const iovec* vec,
 694                              size_t count, unique_ptr<IOBuf>&& buf,
 695                              WriteFlags flags) {
 696   VLOG(6) << "AsyncSocket::writev() this=" << this << ", fd=" << fd_
 697           << ", callback=" << callback << ", count=" << count
 698           << ", state=" << state_;
 699   DestructorGuard dg(this);
 700   unique_ptr<IOBuf>ioBuf(std::move(buf));
 701   assert(eventBase_->isInEventBaseThread());
 702
 703   if (shutdownFlags_ & (SHUT_WRITE | SHUT_WRITE_PENDING)) {
 704     // No new writes may be performed after the write side of the socket has
 705     // been shutdown.
 706     //
 707     // We could just call callback->writeError() here to fail just this write.
 708     // However, fail hard and use invalidState() to fail all outstanding
 709     // callbacks and move the socket into the error state.  There's most likely
 710     // a bug in the caller's code, so we abort everything rather than trying to
 711     // proceed as best we can.
 712     return invalidState(callback);
 713   }
 714
 715   uint32_t countWritten = 0;
 716   uint32_t partialWritten = 0;
 717   int bytesWritten = 0;
 718   bool mustRegister = false;
 719   if ((state_ == StateEnum::ESTABLISHED || state_ == StateEnum::FAST_OPEN) &&
 720       !connecting()) {
 721     if (writeReqHead_ == nullptr) {
 722       // If we are established and there are no other writes pending,
 723       // we can attempt to perform the write immediately.
 724       assert(writeReqTail_ == nullptr);
 725       assert((eventFlags_ & EventHandler::WRITE) == 0);
 726
 727       auto writeResult =
 728           performWrite(vec, count, flags, &countWritten, &partialWritten);
 729       bytesWritten = writeResult.writeReturn;
 730       if (bytesWritten < 0) {
 731         auto errnoCopy = errno;
 732         if (writeResult.exception) {
 733           return failWrite(__func__, callback, 0, *writeResult.exception);
 734         }
 735         AsyncSocketException ex(
 736             AsyncSocketException::INTERNAL_ERROR,
 737             withAddr("writev failed"),
 738             errnoCopy);
 739         return failWrite(__func__, callback, 0, ex);
 740       } else if (countWritten == count) {
 741         // We successfully wrote everything.
 742         // Invoke the callback and return.
 743         if (callback) {
 744           callback->writeSuccess();
 745         }
 746         return;
 747       } else { // continue writing the next writeReq
 748         if (bufferCallback_) {
 749           bufferCallback_->onEgressBuffered();
 750         }
 751       }
 752       if (!connecting()) {
 753         // Writes might put the socket back into connecting state
 754         // if TFO is enabled, and using TFO fails.
 755         // This means that write timeouts would not be active, however
 756         // connect timeouts would affect this stage.
 757         mustRegister = true;
 758       }
 759     }
 760   } else if (!connecting()) {
 761     // Invalid state for writing
 762     return invalidState(callback);
 763   }
 764
 765   // Create a new WriteRequest to add to the queue
 766   WriteRequest* req;
 767   try {
 768     req = BytesWriteRequest::newRequest(this, callback, vec + countWritten,
 769                                         count - countWritten, partialWritten,
 770                                         bytesWritten, std::move(ioBuf), flags);
 771   } catch (const std::exception& ex) {
 772     // we mainly expect to catch std::bad_alloc here
 773     AsyncSocketException tex(AsyncSocketException::INTERNAL_ERROR,
 774         withAddr(string("failed to append new WriteRequest: ") + ex.what()));
 775     return failWrite(__func__, callback, bytesWritten, tex);
 776   }
 777   req->consume();
 778   if (writeReqTail_ == nullptr) {
 779     assert(writeReqHead_ == nullptr);
 780     writeReqHead_ = writeReqTail_ = req;
 781   } else {
 782     writeReqTail_->append(req);
 783     writeReqTail_ = req;
 784   }
 785
 786   // Register for write events if are established and not currently
 787   // waiting on write events
 788   if (mustRegister) {
 789     assert(state_ == StateEnum::ESTABLISHED);
 790     assert((eventFlags_ & EventHandler::WRITE) == 0);
 791     if (!updateEventRegistration(EventHandler::WRITE, 0)) {
 792       assert(state_ == StateEnum::ERROR);
 793       return;
 794     }
 795     if (sendTimeout_ > 0) {
 796       // Schedule a timeout to fire if the write takes too long.
 797       if (!writeTimeout_.scheduleTimeout(sendTimeout_)) {
 798         AsyncSocketException ex(AsyncSocketException::INTERNAL_ERROR,
 799                                withAddr("failed to schedule send timeout"));
 800         return failWrite(__func__, ex);
 801       }
 802     }
 803   }
 804 }
 805
 806 void AsyncSocket::writeRequest(WriteRequest* req) {
 807   if (writeReqTail_ == nullptr) {
 808     assert(writeReqHead_ == nullptr);
 809     writeReqHead_ = writeReqTail_ = req;
 810     req->start();
 811   } else {
 812     writeReqTail_->append(req);
 813     writeReqTail_ = req;
 814   }
 815 }
 816
 817 void AsyncSocket::close() {
 818   VLOG(5) << "AsyncSocket::close(): this=" << this << ", fd_=" << fd_
 819           << ", state=" << state_ << ", shutdownFlags="
 820           << std::hex << (int) shutdownFlags_;
 821
 822   // close() is only different from closeNow() when there are pending writes
 823   // that need to drain before we can close.  In all other cases, just call
 824   // closeNow().
 825   //
 826   // Note that writeReqHead_ can be non-nullptr even in STATE_CLOSED or
 827   // STATE_ERROR if close() is invoked while a previous closeNow() or failure
 828   // is still running.  (e.g., If there are multiple pending writes, and we
 829   // call writeError() on the first one, it may call close().  In this case we
 830   // will already be in STATE_CLOSED or STATE_ERROR, but the remaining pending
 831   // writes will still be in the queue.)
 832   //
 833   // We only need to drain pending writes if we are still in STATE_CONNECTING
 834   // or STATE_ESTABLISHED
 835   if ((writeReqHead_ == nullptr) ||
 836       !(state_ == StateEnum::CONNECTING ||
 837       state_ == StateEnum::ESTABLISHED)) {
 838     closeNow();
 839     return;
 840   }
 841
 842   // Declare a DestructorGuard to ensure that the AsyncSocket cannot be
 843   // destroyed until close() returns.
 844   DestructorGuard dg(this);
 845   assert(eventBase_->isInEventBaseThread());
 846
 847   // Since there are write requests pending, we have to set the
 848   // SHUT_WRITE_PENDING flag, and wait to perform the real close until the
 849   // connect finishes and we finish writing these requests.
 850   //
 851   // Set SHUT_READ to indicate that reads are shut down, and set the
 852   // SHUT_WRITE_PENDING flag to mark that we want to shutdown once the
 853   // pending writes complete.
 854   shutdownFlags_ |= (SHUT_READ | SHUT_WRITE_PENDING);
 855
 856   // If a read callback is set, invoke readEOF() immediately to inform it that
 857   // the socket has been closed and no more data can be read.
 858   if (readCallback_) {
 859     // Disable reads if they are enabled
 860     if (!updateEventRegistration(0, EventHandler::READ)) {
 861       // We're now in the error state; callbacks have been cleaned up
 862       assert(state_ == StateEnum::ERROR);
 863       assert(readCallback_ == nullptr);
 864     } else {
 865       ReadCallback* callback = readCallback_;
 866       readCallback_ = nullptr;
 867       callback->readEOF();
 868     }
 869   }
 870 }
 871
 872 void AsyncSocket::closeNow() {
 873   VLOG(5) << "AsyncSocket::closeNow(): this=" << this << ", fd_=" << fd_
 874           << ", state=" << state_ << ", shutdownFlags="
 875           << std::hex << (int) shutdownFlags_;
 876   DestructorGuard dg(this);
 877   assert(eventBase_ == nullptr || eventBase_->isInEventBaseThread());
 878
 879   switch (state_) {
 880     case StateEnum::ESTABLISHED:
 881     case StateEnum::CONNECTING:
 882     case StateEnum::FAST_OPEN: {
 883       shutdownFlags_ |= (SHUT_READ | SHUT_WRITE);
 884       state_ = StateEnum::CLOSED;
 885
 886       // If the write timeout was set, cancel it.
 887       writeTimeout_.cancelTimeout();
 888
 889       // If we are registered for I/O events, unregister.
 890       if (eventFlags_ != EventHandler::NONE) {
 891         eventFlags_ = EventHandler::NONE;
 892         if (!updateEventRegistration()) {
 893           // We will have been moved into the error state.
 894           assert(state_ == StateEnum::ERROR);
 895           return;
 896         }
 897       }
 898
 899       if (immediateReadHandler_.isLoopCallbackScheduled()) {
 900         immediateReadHandler_.cancelLoopCallback();
 901       }
 902
 903       if (fd_ >= 0) {
 904         ioHandler_.changeHandlerFD(-1);
 905         doClose();
 906       }
 907
 908       invokeConnectErr(socketClosedLocallyEx);
 909
 910       failAllWrites(socketClosedLocallyEx);
 911
 912       if (readCallback_) {
 913         ReadCallback* callback = readCallback_;
 914         readCallback_ = nullptr;
 915         callback->readEOF();
 916       }
 917       return;
 918     }
 919     case StateEnum::CLOSED:
 920       // Do nothing.  It's possible that we are being called recursively
 921       // from inside a callback that we invoked inside another call to close()
 922       // that is still running.
 923       return;
 924     case StateEnum::ERROR:
 925       // Do nothing.  The error handling code has performed (or is performing)
 926       // cleanup.
 927       return;
 928     case StateEnum::UNINIT:
 929       assert(eventFlags_ == EventHandler::NONE);
 930       assert(connectCallback_ == nullptr);
 931       assert(readCallback_ == nullptr);
 932       assert(writeReqHead_ == nullptr);
 933       shutdownFlags_ |= (SHUT_READ | SHUT_WRITE);
 934       state_ = StateEnum::CLOSED;
 935       return;
 936   }
 937
 938   LOG(DFATAL) << "AsyncSocket::closeNow() (this=" << this << ", fd=" << fd_
 939               << ") called in unknown state " << state_;
 940 }
 941
 942 void AsyncSocket::closeWithReset() {
 943   // Enable SO_LINGER, with the linger timeout set to 0.
 944   // This will trigger a TCP reset when we close the socket.
 945   if (fd_ >= 0) {
 946     struct linger optLinger = {1, 0};
 947     if (setSockOpt(SOL_SOCKET, SO_LINGER, &optLinger) != 0) {
 948       VLOG(2) << "AsyncSocket::closeWithReset(): error setting SO_LINGER "
 949               << "on " << fd_ << ": errno=" << errno;
 950     }
 951   }
 952
 953   // Then let closeNow() take care of the rest
 954   closeNow();
 955 }
 956
 957 void AsyncSocket::shutdownWrite() {
 958   VLOG(5) << "AsyncSocket::shutdownWrite(): this=" << this << ", fd=" << fd_
 959           << ", state=" << state_ << ", shutdownFlags="
 960           << std::hex << (int) shutdownFlags_;
 961
 962   // If there are no pending writes, shutdownWrite() is identical to
 963   // shutdownWriteNow().
 964   if (writeReqHead_ == nullptr) {
 965     shutdownWriteNow();
 966     return;
 967   }
 968
 969   assert(eventBase_->isInEventBaseThread());
 970
 971   // There are pending writes.  Set SHUT_WRITE_PENDING so that the actual
 972   // shutdown will be performed once all writes complete.
 973   shutdownFlags_ |= SHUT_WRITE_PENDING;
 974 }
 975
 976 void AsyncSocket::shutdownWriteNow() {
 977   VLOG(5) << "AsyncSocket::shutdownWriteNow(): this=" << this
 978           << ", fd=" << fd_ << ", state=" << state_
 979           << ", shutdownFlags=" << std::hex << (int) shutdownFlags_;
 980
 981   if (shutdownFlags_ & SHUT_WRITE) {
 982     // Writes are already shutdown; nothing else to do.
 983     return;
 984   }
 985
 986   // If SHUT_READ is already set, just call closeNow() to completely
 987   // close the socket.  This can happen if close() was called with writes
 988   // pending, and then shutdownWriteNow() is called before all pending writes
 989   // complete.
 990   if (shutdownFlags_ & SHUT_READ) {
 991     closeNow();
 992     return;
 993   }
 994
 995   DestructorGuard dg(this);
 996   assert(eventBase_ == nullptr || eventBase_->isInEventBaseThread());
 997
 998   switch (static_cast<StateEnum>(state_)) {
 999     case StateEnum::ESTABLISHED:
1000     {
1001       shutdownFlags_ |= SHUT_WRITE;
1002
1003       // If the write timeout was set, cancel it.
1004       writeTimeout_.cancelTimeout();
1005
1006       // If we are registered for write events, unregister.
1007       if (!updateEventRegistration(0, EventHandler::WRITE)) {
1008         // We will have been moved into the error state.
1009         assert(state_ == StateEnum::ERROR);
1010         return;
1011       }
1012
1013       // Shutdown writes on the file descriptor
1014       shutdown(fd_, SHUT_WR);
1015
1016       // Immediately fail all write requests
1017       failAllWrites(socketShutdownForWritesEx);
1018       return;
1019     }
1020     case StateEnum::CONNECTING:
1021     {
1022       // Set the SHUT_WRITE_PENDING flag.
1023       // When the connection completes, it will check this flag,
1024       // shutdown the write half of the socket, and then set SHUT_WRITE.
1025       shutdownFlags_ |= SHUT_WRITE_PENDING;
1026
1027       // Immediately fail all write requests
1028       failAllWrites(socketShutdownForWritesEx);
1029       return;
1030     }
1031     case StateEnum::UNINIT:
1032       // Callers normally shouldn't call shutdownWriteNow() before the socket
1033       // even starts connecting.  Nonetheless, go ahead and set
1034       // SHUT_WRITE_PENDING.  Once the socket eventually connects it will
1035       // immediately shut down the write side of the socket.
1036       shutdownFlags_ |= SHUT_WRITE_PENDING;
1037       return;
1038     case StateEnum::FAST_OPEN:
1039       // In fast open state we haven't call connected yet, and if we shutdown
1040       // the writes, we will never try to call connect, so shut everything down
1041       shutdownFlags_ |= SHUT_WRITE;
1042       // Immediately fail all write requests
1043       failAllWrites(socketShutdownForWritesEx);
1044       return;
1045     case StateEnum::CLOSED:
1046     case StateEnum::ERROR:
1047       // We should never get here.  SHUT_WRITE should always be set
1048       // in STATE_CLOSED and STATE_ERROR.
1049       VLOG(4) << "AsyncSocket::shutdownWriteNow() (this=" << this
1050                  << ", fd=" << fd_ << ") in unexpected state " << state_
1051                  << " with SHUT_WRITE not set ("
1052                  << std::hex << (int) shutdownFlags_ << ")";
1053       assert(false);
1054       return;
1055   }
1056
1057   LOG(DFATAL) << "AsyncSocket::shutdownWriteNow() (this=" << this << ", fd="
1058               << fd_ << ") called in unknown state " << state_;
1059 }
1060
1061 bool AsyncSocket::readable() const {
1062   if (fd_ == -1) {
1063     return false;
1064   }
1065   struct pollfd fds[1];
1066   fds[0].fd = fd_;
1067   fds[0].events = POLLIN;
1068   fds[0].revents = 0;
1069   int rc = poll(fds, 1, 0);
1070   return rc == 1;
1071 }
1072
1073 bool AsyncSocket::isPending() const {
1074   return ioHandler_.isPending();
1075 }
1076
1077 bool AsyncSocket::hangup() const {
1078   if (fd_ == -1) {
1079     // sanity check, no one should ask for hangup if we are not connected.
1080     assert(false);
1081     return false;
1082   }
1083 #ifdef POLLRDHUP // Linux-only
1084   struct pollfd fds[1];
1085   fds[0].fd = fd_;
1086   fds[0].events = POLLRDHUP|POLLHUP;
1087   fds[0].revents = 0;
1088   poll(fds, 1, 0);
1089   return (fds[0].revents & (POLLRDHUP|POLLHUP)) != 0;
1090 #else
1091   return false;
1092 #endif
1093 }
1094
1095 bool AsyncSocket::good() const {
1096   return (
1097       (state_ == StateEnum::CONNECTING || state_ == StateEnum::FAST_OPEN ||
1098        state_ == StateEnum::ESTABLISHED) &&
1099       (shutdownFlags_ == 0) && (eventBase_ != nullptr));
1100 }
1101
1102 bool AsyncSocket::error() const {
1103   return (state_ == StateEnum::ERROR);
1104 }
1105
1106 void AsyncSocket::attachEventBase(EventBase* eventBase) {
1107   VLOG(5) << "AsyncSocket::attachEventBase(this=" << this << ", fd=" << fd_
1108           << ", old evb=" << eventBase_ << ", new evb=" << eventBase
1109           << ", state=" << state_ << ", events="
1110           << std::hex << eventFlags_ << ")";
1111   assert(eventBase_ == nullptr);
1112   assert(eventBase->isInEventBaseThread());
1113
1114   eventBase_ = eventBase;
1115   ioHandler_.attachEventBase(eventBase);
1116   writeTimeout_.attachEventBase(eventBase);
1117 }
1118
1119 void AsyncSocket::detachEventBase() {
1120   VLOG(5) << "AsyncSocket::detachEventBase(this=" << this << ", fd=" << fd_
1121           << ", old evb=" << eventBase_ << ", state=" << state_
1122           << ", events=" << std::hex << eventFlags_ << ")";
1123   assert(eventBase_ != nullptr);
1124   assert(eventBase_->isInEventBaseThread());
1125
1126   eventBase_ = nullptr;
1127   ioHandler_.detachEventBase();
1128   writeTimeout_.detachEventBase();
1129 }
1130
1131 bool AsyncSocket::isDetachable() const {
1132   DCHECK(eventBase_ != nullptr);
1133   DCHECK(eventBase_->isInEventBaseThread());
1134
1135   return !ioHandler_.isHandlerRegistered() && !writeTimeout_.isScheduled();
1136 }
1137
1138 void AsyncSocket::getLocalAddress(folly::SocketAddress* address) const {
1139   if (!localAddr_.isInitialized()) {
1140     localAddr_.setFromLocalAddress(fd_);
1141   }
1142   *address = localAddr_;
1143 }
1144
1145 void AsyncSocket::getPeerAddress(folly::SocketAddress* address) const {
1146   if (!addr_.isInitialized()) {
1147     addr_.setFromPeerAddress(fd_);
1148   }
1149   *address = addr_;
1150 }
1151
1152 bool AsyncSocket::getTFOSucceded() const {
1153   return detail::tfo_succeeded(fd_);
1154 }
1155
1156 int AsyncSocket::setNoDelay(bool noDelay) {
1157   if (fd_ < 0) {
1158     VLOG(4) << "AsyncSocket::setNoDelay() called on non-open socket "
1159                << this << "(state=" << state_ << ")";
1160     return EINVAL;
1161
1162   }
1163
1164   int value = noDelay ? 1 : 0;
1165   if (setsockopt(fd_, IPPROTO_TCP, TCP_NODELAY, &value, sizeof(value)) != 0) {
1166     int errnoCopy = errno;
1167     VLOG(2) << "failed to update TCP_NODELAY option on AsyncSocket "
1168             << this << " (fd=" << fd_ << ", state=" << state_ << "): "
1169             << strerror(errnoCopy);
1170     return errnoCopy;
1171   }
1172
1173   return 0;
1174 }
1175
1176 int AsyncSocket::setCongestionFlavor(const std::string &cname) {
1177
1178   #ifndef TCP_CONGESTION
1179   #define TCP_CONGESTION  13
1180   #endif
1181
1182   if (fd_ < 0) {
1183     VLOG(4) << "AsyncSocket::setCongestionFlavor() called on non-open "
1184                << "socket " << this << "(state=" << state_ << ")";
1185     return EINVAL;
1186
1187   }
1188
1189   if (setsockopt(fd_, IPPROTO_TCP, TCP_CONGESTION, cname.c_str(),
1190         cname.length() + 1) != 0) {
1191     int errnoCopy = errno;
1192     VLOG(2) << "failed to update TCP_CONGESTION option on AsyncSocket "
1193             << this << "(fd=" << fd_ << ", state=" << state_ << "): "
1194             << strerror(errnoCopy);
1195     return errnoCopy;
1196   }
1197
1198   return 0;
1199 }
1200
1201 int AsyncSocket::setQuickAck(bool quickack) {
1202   if (fd_ < 0) {
1203     VLOG(4) << "AsyncSocket::setQuickAck() called on non-open socket "
1204                << this << "(state=" << state_ << ")";
1205     return EINVAL;
1206
1207   }
1208
1209 #ifdef TCP_QUICKACK // Linux-only
1210   int value = quickack ? 1 : 0;
1211   if (setsockopt(fd_, IPPROTO_TCP, TCP_QUICKACK, &value, sizeof(value)) != 0) {
1212     int errnoCopy = errno;
1213     VLOG(2) << "failed to update TCP_QUICKACK option on AsyncSocket"
1214             << this << "(fd=" << fd_ << ", state=" << state_ << "): "
1215             << strerror(errnoCopy);
1216     return errnoCopy;
1217   }
1218
1219   return 0;
1220 #else
1221   return ENOSYS;
1222 #endif
1223 }
1224
1225 int AsyncSocket::setSendBufSize(size_t bufsize) {
1226   if (fd_ < 0) {
1227     VLOG(4) << "AsyncSocket::setSendBufSize() called on non-open socket "
1228                << this << "(state=" << state_ << ")";
1229     return EINVAL;
1230   }
1231
1232   if (setsockopt(fd_, SOL_SOCKET, SO_SNDBUF, &bufsize, sizeof(bufsize)) !=0) {
1233     int errnoCopy = errno;
1234     VLOG(2) << "failed to update SO_SNDBUF option on AsyncSocket"
1235             << this << "(fd=" << fd_ << ", state=" << state_ << "): "
1236             << strerror(errnoCopy);
1237     return errnoCopy;
1238   }
1239
1240   return 0;
1241 }
1242
1243 int AsyncSocket::setRecvBufSize(size_t bufsize) {
1244   if (fd_ < 0) {
1245     VLOG(4) << "AsyncSocket::setRecvBufSize() called on non-open socket "
1246                << this << "(state=" << state_ << ")";
1247     return EINVAL;
1248   }
1249
1250   if (setsockopt(fd_, SOL_SOCKET, SO_RCVBUF, &bufsize, sizeof(bufsize)) !=0) {
1251     int errnoCopy = errno;
1252     VLOG(2) << "failed to update SO_RCVBUF option on AsyncSocket"
1253             << this << "(fd=" << fd_ << ", state=" << state_ << "): "
1254             << strerror(errnoCopy);
1255     return errnoCopy;
1256   }
1257
1258   return 0;
1259 }
1260
1261 int AsyncSocket::setTCPProfile(int profd) {
1262   if (fd_ < 0) {
1263     VLOG(4) << "AsyncSocket::setTCPProfile() called on non-open socket "
1264                << this << "(state=" << state_ << ")";
1265     return EINVAL;
1266   }
1267
1268   if (setsockopt(fd_, SOL_SOCKET, SO_SET_NAMESPACE, &profd, sizeof(int)) !=0) {
1269     int errnoCopy = errno;
1270     VLOG(2) << "failed to set socket namespace option on AsyncSocket"
1271             << this << "(fd=" << fd_ << ", state=" << state_ << "): "
1272             << strerror(errnoCopy);
1273     return errnoCopy;
1274   }
1275
1276   return 0;
1277 }
1278
1279 void AsyncSocket::ioReady(uint16_t events) noexcept {
1280   VLOG(7) << "AsyncSocket::ioRead() this=" << this << ", fd" << fd_
1281           << ", events=" << std::hex << events << ", state=" << state_;
1282   DestructorGuard dg(this);
1283   assert(events & EventHandler::READ_WRITE);
1284   assert(eventBase_->isInEventBaseThread());
1285
1286   uint16_t relevantEvents = events & EventHandler::READ_WRITE;
1287   if (relevantEvents == EventHandler::READ) {
1288     handleRead();
1289   } else if (relevantEvents == EventHandler::WRITE) {
1290     handleWrite();
1291   } else if (relevantEvents == EventHandler::READ_WRITE) {
1292     EventBase* originalEventBase = eventBase_;
1293     // If both read and write events are ready, process writes first.
1294     handleWrite();
1295
1296     // Return now if handleWrite() detached us from our EventBase
1297     if (eventBase_ != originalEventBase) {
1298       return;
1299     }
1300
1301     // Only call handleRead() if a read callback is still installed.
1302     // (It's possible that the read callback was uninstalled during
1303     // handleWrite().)
1304     if (readCallback_) {
1305       handleRead();
1306     }
1307   } else {
1308     VLOG(4) << "AsyncSocket::ioRead() called with unexpected events "
1309                << std::hex << events << "(this=" << this << ")";
1310     abort();
1311   }
1312 }
1313
1314 AsyncSocket::ReadResult
1315 AsyncSocket::performRead(void** buf, size_t* buflen, size_t* /* offset */) {
1316   VLOG(5) << "AsyncSocket::performRead() this=" << this << ", buf=" << *buf
1317           << ", buflen=" << *buflen;
1318
1319   int recvFlags = 0;
1320   if (peek_) {
1321     recvFlags |= MSG_PEEK;
1322   }
1323
1324   ssize_t bytes = recv(fd_, *buf, *buflen, MSG_DONTWAIT | recvFlags);
1325   if (bytes < 0) {
1326     if (errno == EAGAIN || errno == EWOULDBLOCK) {
1327       // No more data to read right now.
1328       return ReadResult(READ_BLOCKING);
1329     } else {
1330       return ReadResult(READ_ERROR);
1331     }
1332   } else {
1333     appBytesReceived_ += bytes;
1334     return ReadResult(bytes);
1335   }
1336 }
1337
1338 void AsyncSocket::prepareReadBuffer(void** buf, size_t* buflen) {
1339   // no matter what, buffer should be preapared for non-ssl socket
1340   CHECK(readCallback_);
1341   readCallback_->getReadBuffer(buf, buflen);
1342 }
1343
1344 void AsyncSocket::handleRead() noexcept {
1345   VLOG(5) << "AsyncSocket::handleRead() this=" << this << ", fd=" << fd_
1346           << ", state=" << state_;
1347   assert(state_ == StateEnum::ESTABLISHED);
1348   assert((shutdownFlags_ & SHUT_READ) == 0);
1349   assert(readCallback_ != nullptr);
1350   assert(eventFlags_ & EventHandler::READ);
1351
1352   // Loop until:
1353   // - a read attempt would block
1354   // - readCallback_ is uninstalled
1355   // - the number of loop iterations exceeds the optional maximum
1356   // - this AsyncSocket is moved to another EventBase
1357   //
1358   // When we invoke readDataAvailable() it may uninstall the readCallback_,
1359   // which is why need to check for it here.
1360   //
1361   // The last bullet point is slightly subtle.  readDataAvailable() may also
1362   // detach this socket from this EventBase.  However, before
1363   // readDataAvailable() returns another thread may pick it up, attach it to
1364   // a different EventBase, and install another readCallback_.  We need to
1365   // exit immediately after readDataAvailable() returns if the eventBase_ has
1366   // changed.  (The caller must perform some sort of locking to transfer the
1367   // AsyncSocket between threads properly.  This will be sufficient to ensure
1368   // that this thread sees the updated eventBase_ variable after
1369   // readDataAvailable() returns.)
1370   uint16_t numReads = 0;
1371   EventBase* originalEventBase = eventBase_;
1372   while (readCallback_ && eventBase_ == originalEventBase) {
1373     // Get the buffer to read into.
1374     void* buf = nullptr;
1375     size_t buflen = 0, offset = 0;
1376     try {
1377       prepareReadBuffer(&buf, &buflen);
1378       VLOG(5) << "prepareReadBuffer() buf=" << buf << ", buflen=" << buflen;
1379     } catch (const AsyncSocketException& ex) {
1380       return failRead(__func__, ex);
1381     } catch (const std::exception& ex) {
1382       AsyncSocketException tex(AsyncSocketException::BAD_ARGS,
1383                               string("ReadCallback::getReadBuffer() "
1384                                      "threw exception: ") +
1385                               ex.what());
1386       return failRead(__func__, tex);
1387     } catch (...) {
1388       AsyncSocketException ex(AsyncSocketException::BAD_ARGS,
1389                              "ReadCallback::getReadBuffer() threw "
1390                              "non-exception type");
1391       return failRead(__func__, ex);
1392     }
1393     if (!isBufferMovable_ && (buf == nullptr || buflen == 0)) {
1394       AsyncSocketException ex(AsyncSocketException::BAD_ARGS,
1395                              "ReadCallback::getReadBuffer() returned "
1396                              "empty buffer");
1397       return failRead(__func__, ex);
1398     }
1399
1400     // Perform the read
1401     auto readResult = performRead(&buf, &buflen, &offset);
1402     auto bytesRead = readResult.readReturn;
1403     VLOG(4) << "this=" << this << ", AsyncSocket::handleRead() got "
1404             << bytesRead << " bytes";
1405     if (bytesRead > 0) {
1406       if (!isBufferMovable_) {
1407         readCallback_->readDataAvailable(bytesRead);
1408       } else {
1409         CHECK(kOpenSslModeMoveBufferOwnership);
1410         VLOG(5) << "this=" << this << ", AsyncSocket::handleRead() got "
1411                 << "buf=" << buf << ", " << bytesRead << "/" << buflen
1412                 << ", offset=" << offset;
1413         auto readBuf = folly::IOBuf::takeOwnership(buf, buflen);
1414         readBuf->trimStart(offset);
1415         readBuf->trimEnd(buflen - offset - bytesRead);
1416         readCallback_->readBufferAvailable(std::move(readBuf));
1417       }
1418
1419       // Fall through and continue around the loop if the read
1420       // completely filled the available buffer.
1421       // Note that readCallback_ may have been uninstalled or changed inside
1422       // readDataAvailable().
1423       if (size_t(bytesRead) < buflen) {
1424         return;
1425       }
1426     } else if (bytesRead == READ_BLOCKING) {
1427         // No more data to read right now.
1428         return;
1429     } else if (bytesRead == READ_ERROR) {
1430       readErr_ = READ_ERROR;
1431       if (readResult.exception) {
1432         return failRead(__func__, *readResult.exception);
1433       }
1434       auto errnoCopy = errno;
1435       AsyncSocketException ex(
1436           AsyncSocketException::INTERNAL_ERROR,
1437           withAddr("recv() failed"),
1438           errnoCopy);
1439       return failRead(__func__, ex);
1440     } else {
1441       assert(bytesRead == READ_EOF);
1442       readErr_ = READ_EOF;
1443       // EOF
1444       shutdownFlags_ |= SHUT_READ;
1445       if (!updateEventRegistration(0, EventHandler::READ)) {
1446         // we've already been moved into STATE_ERROR
1447         assert(state_ == StateEnum::ERROR);
1448         assert(readCallback_ == nullptr);
1449         return;
1450       }
1451
1452       ReadCallback* callback = readCallback_;
1453       readCallback_ = nullptr;
1454       callback->readEOF();
1455       return;
1456     }
1457     if (maxReadsPerEvent_ && (++numReads >= maxReadsPerEvent_)) {
1458       if (readCallback_ != nullptr) {
1459         // We might still have data in the socket.
1460         // (e.g. see comment in AsyncSSLSocket::checkForImmediateRead)
1461         scheduleImmediateRead();
1462       }
1463       return;
1464     }
1465   }
1466 }
1467
1468 /**
1469  * This function attempts to write as much data as possible, until no more data
1470  * can be written.
1471  *
1472  * - If it sends all available data, it unregisters for write events, and stops
1473  *   the writeTimeout_.
1474  *
1475  * - If not all of the data can be sent immediately, it reschedules
1476  *   writeTimeout_ (if a non-zero timeout is set), and ensures the handler is
1477  *   registered for write events.
1478  */
1479 void AsyncSocket::handleWrite() noexcept {
1480   VLOG(5) << "AsyncSocket::handleWrite() this=" << this << ", fd=" << fd_
1481           << ", state=" << state_;
1482   DestructorGuard dg(this);
1483
1484   if (state_ == StateEnum::CONNECTING) {
1485     handleConnect();
1486     return;
1487   }
1488
1489   // Normal write
1490   assert(state_ == StateEnum::ESTABLISHED);
1491   assert((shutdownFlags_ & SHUT_WRITE) == 0);
1492   assert(writeReqHead_ != nullptr);
1493
1494   // Loop until we run out of write requests,
1495   // or until this socket is moved to another EventBase.
1496   // (See the comment in handleRead() explaining how this can happen.)
1497   EventBase* originalEventBase = eventBase_;
1498   while (writeReqHead_ != nullptr && eventBase_ == originalEventBase) {
1499     auto writeResult = writeReqHead_->performWrite();
1500     if (writeResult.writeReturn < 0) {
1501       if (writeResult.exception) {
1502         return failWrite(__func__, *writeResult.exception);
1503       }
1504       auto errnoCopy = errno;
1505       AsyncSocketException ex(
1506           AsyncSocketException::INTERNAL_ERROR,
1507           withAddr("writev() failed"),
1508           errnoCopy);
1509       return failWrite(__func__, ex);
1510     } else if (writeReqHead_->isComplete()) {
1511       // We finished this request
1512       WriteRequest* req = writeReqHead_;
1513       writeReqHead_ = req->getNext();
1514
1515       if (writeReqHead_ == nullptr) {
1516         writeReqTail_ = nullptr;
1517         // This is the last write request.
1518         // Unregister for write events and cancel the send timer
1519         // before we invoke the callback.  We have to update the state properly
1520         // before calling the callback, since it may want to detach us from
1521         // the EventBase.
1522         if (eventFlags_ & EventHandler::WRITE) {
1523           if (!updateEventRegistration(0, EventHandler::WRITE)) {
1524             assert(state_ == StateEnum::ERROR);
1525             return;
1526           }
1527           // Stop the send timeout
1528           writeTimeout_.cancelTimeout();
1529         }
1530         assert(!writeTimeout_.isScheduled());
1531
1532         // If SHUT_WRITE_PENDING is set, we should shutdown the socket after
1533         // we finish sending the last write request.
1534         //
1535         // We have to do this before invoking writeSuccess(), since
1536         // writeSuccess() may detach us from our EventBase.
1537         if (shutdownFlags_ & SHUT_WRITE_PENDING) {
1538           assert(connectCallback_ == nullptr);
1539           shutdownFlags_ |= SHUT_WRITE;
1540
1541           if (shutdownFlags_ & SHUT_READ) {
1542             // Reads have already been shutdown.  Fully close the socket and
1543             // move to STATE_CLOSED.
1544             //
1545             // Note: This code currently moves us to STATE_CLOSED even if
1546             // close() hasn't ever been called.  This can occur if we have
1547             // received EOF from the peer and shutdownWrite() has been called
1548             // locally.  Should we bother staying in STATE_ESTABLISHED in this
1549             // case, until close() is actually called?  I can't think of a
1550             // reason why we would need to do so.  No other operations besides
1551             // calling close() or destroying the socket can be performed at
1552             // this point.
1553             assert(readCallback_ == nullptr);
1554             state_ = StateEnum::CLOSED;
1555             if (fd_ >= 0) {
1556               ioHandler_.changeHandlerFD(-1);
1557               doClose();
1558             }
1559           } else {
1560             // Reads are still enabled, so we are only doing a half-shutdown
1561             shutdown(fd_, SHUT_WR);
1562           }
1563         }
1564       }
1565
1566       // Invoke the callback
1567       WriteCallback* callback = req->getCallback();
1568       req->destroy();
1569       if (callback) {
1570         callback->writeSuccess();
1571       }
1572       // We'll continue around the loop, trying to write another request
1573     } else {
1574       // Partial write.
1575       if (bufferCallback_) {
1576         bufferCallback_->onEgressBuffered();
1577       }
1578       writeReqHead_->consume();
1579       // Stop after a partial write; it's highly likely that a subsequent write
1580       // attempt will just return EAGAIN.
1581       //
1582       // Ensure that we are registered for write events.
1583       if ((eventFlags_ & EventHandler::WRITE) == 0) {
1584         if (!updateEventRegistration(EventHandler::WRITE, 0)) {
1585           assert(state_ == StateEnum::ERROR);
1586           return;
1587         }
1588       }
1589
1590       // Reschedule the send timeout, since we have made some write progress.
1591       if (sendTimeout_ > 0) {
1592         if (!writeTimeout_.scheduleTimeout(sendTimeout_)) {
1593           AsyncSocketException ex(AsyncSocketException::INTERNAL_ERROR,
1594               withAddr("failed to reschedule write timeout"));
1595           return failWrite(__func__, ex);
1596         }
1597       }
1598       return;
1599     }
1600   }
1601   if (!writeReqHead_ && bufferCallback_) {
1602     bufferCallback_->onEgressBufferCleared();
1603   }
1604 }
1605
1606 void AsyncSocket::checkForImmediateRead() noexcept {
1607   // We currently don't attempt to perform optimistic reads in AsyncSocket.
1608   // (However, note that some subclasses do override this method.)
1609   //
1610   // Simply calling handleRead() here would be bad, as this would call
1611   // readCallback_->getReadBuffer(), forcing the callback to allocate a read
1612   // buffer even though no data may be available.  This would waste lots of
1613   // memory, since the buffer will sit around unused until the socket actually
1614   // becomes readable.
1615   //
1616   // Checking if the socket is readable now also seems like it would probably
1617   // be a pessimism.  In most cases it probably wouldn't be readable, and we
1618   // would just waste an extra system call.  Even if it is readable, waiting to
1619   // find out from libevent on the next event loop doesn't seem that bad.
1620 }
1621
1622 void AsyncSocket::handleInitialReadWrite() noexcept {
1623   // Our callers should already be holding a DestructorGuard, but grab
1624   // one here just to make sure, in case one of our calling code paths ever
1625   // changes.
1626   DestructorGuard dg(this);
1627   // If we have a readCallback_, make sure we enable read events.  We
1628   // may already be registered for reads if connectSuccess() set
1629   // the read calback.
1630   if (readCallback_ && !(eventFlags_ & EventHandler::READ)) {
1631     assert(state_ == StateEnum::ESTABLISHED);
1632     assert((shutdownFlags_ & SHUT_READ) == 0);
1633     if (!updateEventRegistration(EventHandler::READ, 0)) {
1634       assert(state_ == StateEnum::ERROR);
1635       return;
1636     }
1637     checkForImmediateRead();
1638   } else if (readCallback_ == nullptr) {
1639     // Unregister for read events.
1640     updateEventRegistration(0, EventHandler::READ);
1641   }
1642
1643   // If we have write requests pending, try to send them immediately.
1644   // Since we just finished accepting, there is a very good chance that we can
1645   // write without blocking.
1646   //
1647   // However, we only process them if EventHandler::WRITE is not already set,
1648   // which means that we're already blocked on a write attempt.  (This can
1649   // happen if connectSuccess() called write() before returning.)
1650   if (writeReqHead_ && !(eventFlags_ & EventHandler::WRITE)) {
1651     // Call handleWrite() to perform write processing.
1652     handleWrite();
1653   } else if (writeReqHead_ == nullptr) {
1654     // Unregister for write event.
1655     updateEventRegistration(0, EventHandler::WRITE);
1656   }
1657 }
1658
1659 void AsyncSocket::handleConnect() noexcept {
1660   VLOG(5) << "AsyncSocket::handleConnect() this=" << this << ", fd=" << fd_
1661           << ", state=" << state_;
1662   assert(state_ == StateEnum::CONNECTING);
1663   // SHUT_WRITE can never be set while we are still connecting;
1664   // SHUT_WRITE_PENDING may be set, be we only set SHUT_WRITE once the connect
1665   // finishes
1666   assert((shutdownFlags_ & SHUT_WRITE) == 0);
1667
1668   // In case we had a connect timeout, cancel the timeout
1669   writeTimeout_.cancelTimeout();
1670   // We don't use a persistent registration when waiting on a connect event,
1671   // so we have been automatically unregistered now.  Update eventFlags_ to
1672   // reflect reality.
1673   assert(eventFlags_ == EventHandler::WRITE);
1674   eventFlags_ = EventHandler::NONE;
1675
1676   // Call getsockopt() to check if the connect succeeded
1677   int error;
1678   socklen_t len = sizeof(error);
1679   int rv = getsockopt(fd_, SOL_SOCKET, SO_ERROR, &error, &len);
1680   if (rv != 0) {
1681     auto errnoCopy = errno;
1682     AsyncSocketException ex(
1683         AsyncSocketException::INTERNAL_ERROR,
1684         withAddr("error calling getsockopt() after connect"),
1685         errnoCopy);
1686     VLOG(4) << "AsyncSocket::handleConnect(this=" << this << ", fd="
1687                << fd_ << " host=" << addr_.describe()
1688                << ") exception:" << ex.what();
1689     return failConnect(__func__, ex);
1690   }
1691
1692   if (error != 0) {
1693     AsyncSocketException ex(AsyncSocketException::NOT_OPEN,
1694                            "connect failed", error);
1695     VLOG(1) << "AsyncSocket::handleConnect(this=" << this << ", fd="
1696             << fd_ << " host=" << addr_.describe()
1697             << ") exception: " << ex.what();
1698     return failConnect(__func__, ex);
1699   }
1700
1701   // Move into STATE_ESTABLISHED
1702   state_ = StateEnum::ESTABLISHED;
1703
1704   // If SHUT_WRITE_PENDING is set and we don't have any write requests to
1705   // perform, immediately shutdown the write half of the socket.
1706   if ((shutdownFlags_ & SHUT_WRITE_PENDING) && writeReqHead_ == nullptr) {
1707     // SHUT_READ shouldn't be set.  If close() is called on the socket while we
1708     // are still connecting we just abort the connect rather than waiting for
1709     // it to complete.
1710     assert((shutdownFlags_ & SHUT_READ) == 0);
1711     shutdown(fd_, SHUT_WR);
1712     shutdownFlags_ |= SHUT_WRITE;
1713   }
1714
1715   VLOG(7) << "AsyncSocket " << this << ": fd " << fd_
1716           << "successfully connected; state=" << state_;
1717
1718   // Remember the EventBase we are attached to, before we start invoking any
1719   // callbacks (since the callbacks may call detachEventBase()).
1720   EventBase* originalEventBase = eventBase_;
1721
1722   invokeConnectSuccess();
1723   // Note that the connect callback may have changed our state.
1724   // (set or unset the read callback, called write(), closed the socket, etc.)
1725   // The following code needs to handle these situations correctly.
1726   //
1727   // If the socket has been closed, readCallback_ and writeReqHead_ will
1728   // always be nullptr, so that will prevent us from trying to read or write.
1729   //
1730   // The main thing to check for is if eventBase_ is still originalEventBase.
1731   // If not, we have been detached from this event base, so we shouldn't
1732   // perform any more operations.
1733   if (eventBase_ != originalEventBase) {
1734     return;
1735   }
1736
1737   handleInitialReadWrite();
1738 }
1739
1740 void AsyncSocket::timeoutExpired() noexcept {
1741   VLOG(7) << "AsyncSocket " << this << ", fd " << fd_ << ": timeout expired: "
1742           << "state=" << state_ << ", events=" << std::hex << eventFlags_;
1743   DestructorGuard dg(this);
1744   assert(eventBase_->isInEventBaseThread());
1745
1746   if (state_ == StateEnum::CONNECTING) {
1747     // connect() timed out
1748     // Unregister for I/O events.
1749     if (connectCallback_) {
1750       AsyncSocketException ex(
1751           AsyncSocketException::TIMED_OUT, "connect timed out");
1752       failConnect(__func__, ex);
1753     } else {
1754       // we faced a connect error without a connect callback, which could
1755       // happen due to TFO.
1756       AsyncSocketException ex(
1757           AsyncSocketException::TIMED_OUT, "write timed out during connection");
1758       failWrite(__func__, ex);
1759     }
1760   } else {
1761     // a normal write operation timed out
1762     AsyncSocketException ex(AsyncSocketException::TIMED_OUT, "write timed out");
1763     failWrite(__func__, ex);
1764   }
1765 }
1766
1767 ssize_t AsyncSocket::tfoSendMsg(int fd, struct msghdr* msg, int msg_flags) {
1768   return detail::tfo_sendmsg(fd, msg, msg_flags);
1769 }
1770
1771 AsyncSocket::WriteResult
1772 AsyncSocket::sendSocketMessage(int fd, struct msghdr* msg, int msg_flags) {
1773   ssize_t totalWritten = 0;
1774   if (state_ == StateEnum::FAST_OPEN) {
1775     sockaddr_storage addr;
1776     auto len = addr_.getAddress(&addr);
1777     msg->msg_name = &addr;
1778     msg->msg_namelen = len;
1779     totalWritten = tfoSendMsg(fd_, msg, msg_flags);
1780     if (totalWritten >= 0) {
1781       tfoFinished_ = true;
1782       state_ = StateEnum::ESTABLISHED;
1783       // We schedule this asynchrously so that we don't end up
1784       // invoking initial read or write while a write is in progress.
1785       scheduleInitialReadWrite();
1786     } else if (errno == EINPROGRESS) {
1787       VLOG(4) << "TFO falling back to connecting";
1788       // A normal sendmsg doesn't return EINPROGRESS, however
1789       // TFO might fallback to connecting if there is no
1790       // cookie.
1791       state_ = StateEnum::CONNECTING;
1792       try {
1793         scheduleConnectTimeout();
1794         registerForConnectEvents();
1795       } catch (const AsyncSocketException& ex) {
1796         return WriteResult(
1797             WRITE_ERROR, folly::make_unique<AsyncSocketException>(ex));
1798       }
1799       // Let's fake it that no bytes were written and return an errno.
1800       errno = EAGAIN;
1801       totalWritten = -1;
1802     } else if (errno == EOPNOTSUPP) {
1803       // Try falling back to connecting.
1804       VLOG(4) << "TFO not supported";
1805       state_ = StateEnum::CONNECTING;
1806       try {
1807         int ret = socketConnect((const sockaddr*)&addr, len);
1808         if (ret == 0) {
1809           // connect succeeded immediately
1810           // Treat this like no data was written.
1811           state_ = StateEnum::ESTABLISHED;
1812           scheduleInitialReadWrite();
1813         }
1814         // If there was no exception during connections,
1815         // we would return that no bytes were written.
1816         errno = EAGAIN;
1817         totalWritten = -1;
1818       } catch (const AsyncSocketException& ex) {
1819         return WriteResult(
1820             WRITE_ERROR, folly::make_unique<AsyncSocketException>(ex));
1821       }
1822     } else if (errno == EAGAIN) {
1823       // Normally sendmsg would indicate that the write would block.
1824       // However in the fast open case, it would indicate that sendmsg
1825       // fell back to a connect. This is a return code from connect()
1826       // instead, and is an error condition indicating no fds available.
1827       return WriteResult(
1828           WRITE_ERROR,
1829           folly::make_unique<AsyncSocketException>(
1830               AsyncSocketException::UNKNOWN, "No more free local ports"));
1831     }
1832   } else {
1833     totalWritten = ::sendmsg(fd, msg, msg_flags);
1834   }
1835   return WriteResult(totalWritten);
1836 }
1837
1838 AsyncSocket::WriteResult AsyncSocket::performWrite(
1839     const iovec* vec,
1840     uint32_t count,
1841     WriteFlags flags,
1842     uint32_t* countWritten,
1843     uint32_t* partialWritten) {
1844   // We use sendmsg() instead of writev() so that we can pass in MSG_NOSIGNAL
1845   // We correctly handle EPIPE errors, so we never want to receive SIGPIPE
1846   // (since it may terminate the program if the main program doesn't explicitly
1847   // ignore it).
1848   struct msghdr msg;
1849   msg.msg_name = nullptr;
1850   msg.msg_namelen = 0;
1851   msg.msg_iov = const_cast<iovec *>(vec);
1852   msg.msg_iovlen = std::min<size_t>(count, kIovMax);
1853   msg.msg_control = nullptr;
1854   msg.msg_controllen = 0;
1855   msg.msg_flags = 0;
1856
1857   int msg_flags = MSG_DONTWAIT;
1858
1859 #ifdef MSG_NOSIGNAL // Linux-only
1860   msg_flags |= MSG_NOSIGNAL;
1861   if (isSet(flags, WriteFlags::CORK)) {
1862     // MSG_MORE tells the kernel we have more data to send, so wait for us to
1863     // give it the rest of the data rather than immediately sending a partial
1864     // frame, even when TCP_NODELAY is enabled.
1865     msg_flags |= MSG_MORE;
1866   }
1867 #endif
1868   if (isSet(flags, WriteFlags::EOR)) {
1869     // marks that this is the last byte of a record (response)
1870     msg_flags |= MSG_EOR;
1871   }
1872   auto writeResult = sendSocketMessage(fd_, &msg, msg_flags);
1873   auto totalWritten = writeResult.writeReturn;
1874   if (totalWritten < 0) {
1875     bool tryAgain = (errno == EAGAIN);
1876 #ifdef __APPLE__
1877     // Apple has a bug where doing a second write on a socket which we
1878     // have opened with TFO causes an ENOTCONN to be thrown. However the
1879     // socket is really connected, so treat ENOTCONN as a EAGAIN until
1880     // this bug is fixed.
1881     tryAgain |= (errno == ENOTCONN);
1882 #endif
1883     if (!writeResult.exception && tryAgain) {
1884       // TCP buffer is full; we can't write any more data right now.
1885       *countWritten = 0;
1886       *partialWritten = 0;
1887       return WriteResult(0);
1888     }
1889     // error
1890     *countWritten = 0;
1891     *partialWritten = 0;
1892     return writeResult;
1893   }
1894
1895   appBytesWritten_ += totalWritten;
1896
1897   uint32_t bytesWritten;
1898   uint32_t n;
1899   for (bytesWritten = totalWritten, n = 0; n < count; ++n) {
1900     const iovec* v = vec + n;
1901     if (v->iov_len > bytesWritten) {
1902       // Partial write finished in the middle of this iovec
1903       *countWritten = n;
1904       *partialWritten = bytesWritten;
1905       return WriteResult(totalWritten);
1906     }
1907
1908     bytesWritten -= v->iov_len;
1909   }
1910
1911   assert(bytesWritten == 0);
1912   *countWritten = n;
1913   *partialWritten = 0;
1914   return WriteResult(totalWritten);
1915 }
1916
1917 /**
1918  * Re-register the EventHandler after eventFlags_ has changed.
1919  *
1920  * If an error occurs, fail() is called to move the socket into the error state
1921  * and call all currently installed callbacks.  After an error, the
1922  * AsyncSocket is completely unregistered.
1923  *
1924  * @return Returns true on succcess, or false on error.
1925  */
1926 bool AsyncSocket::updateEventRegistration() {
1927   VLOG(5) << "AsyncSocket::updateEventRegistration(this=" << this
1928           << ", fd=" << fd_ << ", evb=" << eventBase_ << ", state=" << state_
1929           << ", events=" << std::hex << eventFlags_;
1930   assert(eventBase_->isInEventBaseThread());
1931   if (eventFlags_ == EventHandler::NONE) {
1932     ioHandler_.unregisterHandler();
1933     return true;
1934   }
1935
1936   // Always register for persistent events, so we don't have to re-register
1937   // after being called back.
1938   if (!ioHandler_.registerHandler(eventFlags_ | EventHandler::PERSIST)) {
1939     eventFlags_ = EventHandler::NONE; // we're not registered after error
1940     AsyncSocketException ex(AsyncSocketException::INTERNAL_ERROR,
1941         withAddr("failed to update AsyncSocket event registration"));
1942     fail("updateEventRegistration", ex);
1943     return false;
1944   }
1945
1946   return true;
1947 }
1948
1949 bool AsyncSocket::updateEventRegistration(uint16_t enable,
1950                                            uint16_t disable) {
1951   uint16_t oldFlags = eventFlags_;
1952   eventFlags_ |= enable;
1953   eventFlags_ &= ~disable;
1954   if (eventFlags_ == oldFlags) {
1955     return true;
1956   } else {
1957     return updateEventRegistration();
1958   }
1959 }
1960
1961 void AsyncSocket::startFail() {
1962   // startFail() should only be called once
1963   assert(state_ != StateEnum::ERROR);
1964   assert(getDestructorGuardCount() > 0);
1965   state_ = StateEnum::ERROR;
1966   // Ensure that SHUT_READ and SHUT_WRITE are set,
1967   // so all future attempts to read or write will be rejected
1968   shutdownFlags_ |= (SHUT_READ | SHUT_WRITE);
1969
1970   if (eventFlags_ != EventHandler::NONE) {
1971     eventFlags_ = EventHandler::NONE;
1972     ioHandler_.unregisterHandler();
1973   }
1974   writeTimeout_.cancelTimeout();
1975
1976   if (fd_ >= 0) {
1977     ioHandler_.changeHandlerFD(-1);
1978     doClose();
1979   }
1980 }
1981
1982 void AsyncSocket::invokeAllErrors(const AsyncSocketException& ex) {
1983   invokeConnectErr(ex);
1984   failAllWrites(ex);
1985
1986   if (readCallback_) {
1987     ReadCallback* callback = readCallback_;
1988     readCallback_ = nullptr;
1989     callback->readErr(ex);
1990   }
1991 }
1992
1993 void AsyncSocket::finishFail() {
1994   assert(state_ == StateEnum::ERROR);
1995   assert(getDestructorGuardCount() > 0);
1996
1997   AsyncSocketException ex(
1998       AsyncSocketException::INTERNAL_ERROR,
1999       withAddr("socket closing after error"));
2000   invokeAllErrors(ex);
2001 }
2002
2003 void AsyncSocket::finishFail(const AsyncSocketException& ex) {
2004   assert(state_ == StateEnum::ERROR);
2005   assert(getDestructorGuardCount() > 0);
2006   invokeAllErrors(ex);
2007 }
2008
2009 void AsyncSocket::fail(const char* fn, const AsyncSocketException& ex) {
2010   VLOG(4) << "AsyncSocket(this=" << this << ", fd=" << fd_ << ", state="
2011              << state_ << " host=" << addr_.describe()
2012              << "): failed in " << fn << "(): "
2013              << ex.what();
2014   startFail();
2015   finishFail();
2016 }
2017
2018 void AsyncSocket::failConnect(const char* fn, const AsyncSocketException& ex) {
2019   VLOG(5) << "AsyncSocket(this=" << this << ", fd=" << fd_ << ", state="
2020                << state_ << " host=" << addr_.describe()
2021                << "): failed while connecting in " << fn << "(): "
2022                << ex.what();
2023   startFail();
2024
2025   invokeConnectErr(ex);
2026   finishFail(ex);
2027 }
2028
2029 void AsyncSocket::failRead(const char* fn, const AsyncSocketException& ex) {
2030   VLOG(5) << "AsyncSocket(this=" << this << ", fd=" << fd_ << ", state="
2031                << state_ << " host=" << addr_.describe()
2032                << "): failed while reading in " << fn << "(): "
2033                << ex.what();
2034   startFail();
2035
2036   if (readCallback_ != nullptr) {
2037     ReadCallback* callback = readCallback_;
2038     readCallback_ = nullptr;
2039     callback->readErr(ex);
2040   }
2041
2042   finishFail();
2043 }
2044
2045 void AsyncSocket::failWrite(const char* fn, const AsyncSocketException& ex) {
2046   VLOG(5) << "AsyncSocket(this=" << this << ", fd=" << fd_ << ", state="
2047                << state_ << " host=" << addr_.describe()
2048                << "): failed while writing in " << fn << "(): "
2049                << ex.what();
2050   startFail();
2051
2052   // Only invoke the first write callback, since the error occurred while
2053   // writing this request.  Let any other pending write callbacks be invoked in
2054   // finishFail().
2055   if (writeReqHead_ != nullptr) {
2056     WriteRequest* req = writeReqHead_;
2057     writeReqHead_ = req->getNext();
2058     WriteCallback* callback = req->getCallback();
2059     uint32_t bytesWritten = req->getTotalBytesWritten();
2060     req->destroy();
2061     if (callback) {
2062       callback->writeErr(bytesWritten, ex);
2063     }
2064   }
2065
2066   finishFail();
2067 }
2068
2069 void AsyncSocket::failWrite(const char* fn, WriteCallback* callback,
2070                              size_t bytesWritten,
2071                              const AsyncSocketException& ex) {
2072   // This version of failWrite() is used when the failure occurs before
2073   // we've added the callback to writeReqHead_.
2074   VLOG(4) << "AsyncSocket(this=" << this << ", fd=" << fd_ << ", state="
2075              << state_ << " host=" << addr_.describe()
2076              <<"): failed while writing in " << fn << "(): "
2077              << ex.what();
2078   startFail();
2079
2080   if (callback != nullptr) {
2081     callback->writeErr(bytesWritten, ex);
2082   }
2083
2084   finishFail();
2085 }
2086
2087 void AsyncSocket::failAllWrites(const AsyncSocketException& ex) {
2088   // Invoke writeError() on all write callbacks.
2089   // This is used when writes are forcibly shutdown with write requests
2090   // pending, or when an error occurs with writes pending.
2091   while (writeReqHead_ != nullptr) {
2092     WriteRequest* req = writeReqHead_;
2093     writeReqHead_ = req->getNext();
2094     WriteCallback* callback = req->getCallback();
2095     if (callback) {
2096       callback->writeErr(req->getTotalBytesWritten(), ex);
2097     }
2098     req->destroy();
2099   }
2100 }
2101
2102 void AsyncSocket::invalidState(ConnectCallback* callback) {
2103   VLOG(5) << "AsyncSocket(this=" << this << ", fd=" << fd_
2104              << "): connect() called in invalid state " << state_;
2105
2106   /*
2107    * The invalidState() methods don't use the normal failure mechanisms,
2108    * since we don't know what state we are in.  We don't want to call
2109    * startFail()/finishFail() recursively if we are already in the middle of
2110    * cleaning up.
2111    */
2112
2113   AsyncSocketException ex(AsyncSocketException::ALREADY_OPEN,
2114                          "connect() called with socket in invalid state");
2115   connectEndTime_ = std::chrono::steady_clock::now();
2116   if (state_ == StateEnum::CLOSED || state_ == StateEnum::ERROR) {
2117     if (callback) {
2118       callback->connectErr(ex);
2119     }
2120   } else {
2121     // We can't use failConnect() here since connectCallback_
2122     // may already be set to another callback.  Invoke this ConnectCallback
2123     // here; any other connectCallback_ will be invoked in finishFail()
2124     startFail();
2125     if (callback) {
2126       callback->connectErr(ex);
2127     }
2128     finishFail();
2129   }
2130 }
2131
2132 void AsyncSocket::invokeConnectErr(const AsyncSocketException& ex) {
2133   connectEndTime_ = std::chrono::steady_clock::now();
2134   if (connectCallback_) {
2135     ConnectCallback* callback = connectCallback_;
2136     connectCallback_ = nullptr;
2137     callback->connectErr(ex);
2138   }
2139 }
2140
2141 void AsyncSocket::invokeConnectSuccess() {
2142   connectEndTime_ = std::chrono::steady_clock::now();
2143   if (connectCallback_) {
2144     ConnectCallback* callback = connectCallback_;
2145     connectCallback_ = nullptr;
2146     callback->connectSuccess();
2147   }
2148 }
2149
2150 void AsyncSocket::invalidState(ReadCallback* callback) {
2151   VLOG(4) << "AsyncSocket(this=" << this << ", fd=" << fd_
2152              << "): setReadCallback(" << callback
2153              << ") called in invalid state " << state_;
2154
2155   AsyncSocketException ex(AsyncSocketException::NOT_OPEN,
2156                          "setReadCallback() called with socket in "
2157                          "invalid state");
2158   if (state_ == StateEnum::CLOSED || state_ == StateEnum::ERROR) {
2159     if (callback) {
2160       callback->readErr(ex);
2161     }
2162   } else {
2163     startFail();
2164     if (callback) {
2165       callback->readErr(ex);
2166     }
2167     finishFail();
2168   }
2169 }
2170
2171 void AsyncSocket::invalidState(WriteCallback* callback) {
2172   VLOG(4) << "AsyncSocket(this=" << this << ", fd=" << fd_
2173              << "): write() called in invalid state " << state_;
2174
2175   AsyncSocketException ex(AsyncSocketException::NOT_OPEN,
2176                          withAddr("write() called with socket in invalid state"));
2177   if (state_ == StateEnum::CLOSED || state_ == StateEnum::ERROR) {
2178     if (callback) {
2179       callback->writeErr(0, ex);
2180     }
2181   } else {
2182     startFail();
2183     if (callback) {
2184       callback->writeErr(0, ex);
2185     }
2186     finishFail();
2187   }
2188 }
2189
2190 void AsyncSocket::doClose() {
2191   if (fd_ == -1) return;
2192   if (shutdownSocketSet_) {
2193     shutdownSocketSet_->close(fd_);
2194   } else {
2195     ::close(fd_);
2196   }
2197   fd_ = -1;
2198 }
2199
2200 std::ostream& operator << (std::ostream& os,
2201                            const AsyncSocket::StateEnum& state) {
2202   os << static_cast<int>(state);
2203   return os;
2204 }
2205
2206 std::string AsyncSocket::withAddr(const std::string& s) {
2207   // Don't use addr_ directly because it may not be initialized
2208   // e.g. if constructed from fd
2209   folly::SocketAddress peer, local;
2210   try {
2211     getPeerAddress(&peer);
2212     getLocalAddress(&local);
2213   } catch (const std::exception&) {
2214     // ignore
2215   } catch (...) {
2216     // ignore
2217   }
2218   return s + " (peer=" + peer.describe() + ", local=" + local.describe() + ")";
2219 }
2220
2221 void AsyncSocket::setBufferCallback(BufferCallback* cb) {
2222   bufferCallback_ = cb;
2223 }
2224
2225 } // folly