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