Enable EventBase switching for AsyncSocket even if it has registered events
[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
1438   updateEventRegistration();
1439
1440   writeTimeout_.attachEventBase(eventBase);
1441   if (evbChangeCb_) {
1442     evbChangeCb_->evbAttached(this);
1443   }
1444 }
1445
1446 void AsyncSocket::detachEventBase() {
1447   VLOG(5) << "AsyncSocket::detachEventBase(this=" << this << ", fd=" << fd_
1448           << ", old evb=" << eventBase_ << ", state=" << state_
1449           << ", events=" << std::hex << eventFlags_ << ")";
1450   assert(eventBase_ != nullptr);
1451   eventBase_->dcheckIsInEventBaseThread();
1452
1453   eventBase_ = nullptr;
1454
1455   ioHandler_.unregisterHandler();
1456
1457   ioHandler_.detachEventBase();
1458   writeTimeout_.detachEventBase();
1459   if (evbChangeCb_) {
1460     evbChangeCb_->evbDetached(this);
1461   }
1462 }
1463
1464 bool AsyncSocket::isDetachable() const {
1465   DCHECK(eventBase_ != nullptr);
1466   eventBase_->dcheckIsInEventBaseThread();
1467
1468   return !writeTimeout_.isScheduled();
1469 }
1470
1471 void AsyncSocket::cacheAddresses() {
1472   if (fd_ >= 0) {
1473     try {
1474       cacheLocalAddress();
1475       cachePeerAddress();
1476     } catch (const std::system_error& e) {
1477       if (e.code() != std::error_code(ENOTCONN, std::system_category())) {
1478         VLOG(1) << "Error caching addresses: " << e.code().value() << ", "
1479                 << e.code().message();
1480       }
1481     }
1482   }
1483 }
1484
1485 void AsyncSocket::cacheLocalAddress() const {
1486   if (!localAddr_.isInitialized()) {
1487     localAddr_.setFromLocalAddress(fd_);
1488   }
1489 }
1490
1491 void AsyncSocket::cachePeerAddress() const {
1492   if (!addr_.isInitialized()) {
1493     addr_.setFromPeerAddress(fd_);
1494   }
1495 }
1496
1497 bool AsyncSocket::isZeroCopyWriteInProgress() const noexcept {
1498   eventBase_->dcheckIsInEventBaseThread();
1499   return (!idZeroCopyBufPtrMap_.empty());
1500 }
1501
1502 void AsyncSocket::getLocalAddress(folly::SocketAddress* address) const {
1503   cacheLocalAddress();
1504   *address = localAddr_;
1505 }
1506
1507 void AsyncSocket::getPeerAddress(folly::SocketAddress* address) const {
1508   cachePeerAddress();
1509   *address = addr_;
1510 }
1511
1512 bool AsyncSocket::getTFOSucceded() const {
1513   return detail::tfo_succeeded(fd_);
1514 }
1515
1516 int AsyncSocket::setNoDelay(bool noDelay) {
1517   if (fd_ < 0) {
1518     VLOG(4) << "AsyncSocket::setNoDelay() called on non-open socket "
1519                << this << "(state=" << state_ << ")";
1520     return EINVAL;
1521
1522   }
1523
1524   int value = noDelay ? 1 : 0;
1525   if (setsockopt(fd_, IPPROTO_TCP, TCP_NODELAY, &value, sizeof(value)) != 0) {
1526     int errnoCopy = errno;
1527     VLOG(2) << "failed to update TCP_NODELAY option on AsyncSocket "
1528             << this << " (fd=" << fd_ << ", state=" << state_ << "): "
1529             << strerror(errnoCopy);
1530     return errnoCopy;
1531   }
1532
1533   return 0;
1534 }
1535
1536 int AsyncSocket::setCongestionFlavor(const std::string &cname) {
1537
1538   #ifndef TCP_CONGESTION
1539   #define TCP_CONGESTION  13
1540   #endif
1541
1542   if (fd_ < 0) {
1543     VLOG(4) << "AsyncSocket::setCongestionFlavor() called on non-open "
1544                << "socket " << this << "(state=" << state_ << ")";
1545     return EINVAL;
1546
1547   }
1548
1549   if (setsockopt(
1550           fd_,
1551           IPPROTO_TCP,
1552           TCP_CONGESTION,
1553           cname.c_str(),
1554           socklen_t(cname.length() + 1)) != 0) {
1555     int errnoCopy = errno;
1556     VLOG(2) << "failed to update TCP_CONGESTION option on AsyncSocket "
1557             << this << "(fd=" << fd_ << ", state=" << state_ << "): "
1558             << strerror(errnoCopy);
1559     return errnoCopy;
1560   }
1561
1562   return 0;
1563 }
1564
1565 int AsyncSocket::setQuickAck(bool quickack) {
1566   (void)quickack;
1567   if (fd_ < 0) {
1568     VLOG(4) << "AsyncSocket::setQuickAck() called on non-open socket "
1569                << this << "(state=" << state_ << ")";
1570     return EINVAL;
1571
1572   }
1573
1574 #ifdef TCP_QUICKACK // Linux-only
1575   int value = quickack ? 1 : 0;
1576   if (setsockopt(fd_, IPPROTO_TCP, TCP_QUICKACK, &value, sizeof(value)) != 0) {
1577     int errnoCopy = errno;
1578     VLOG(2) << "failed to update TCP_QUICKACK option on AsyncSocket"
1579             << this << "(fd=" << fd_ << ", state=" << state_ << "): "
1580             << strerror(errnoCopy);
1581     return errnoCopy;
1582   }
1583
1584   return 0;
1585 #else
1586   return ENOSYS;
1587 #endif
1588 }
1589
1590 int AsyncSocket::setSendBufSize(size_t bufsize) {
1591   if (fd_ < 0) {
1592     VLOG(4) << "AsyncSocket::setSendBufSize() called on non-open socket "
1593                << this << "(state=" << state_ << ")";
1594     return EINVAL;
1595   }
1596
1597   if (setsockopt(fd_, SOL_SOCKET, SO_SNDBUF, &bufsize, sizeof(bufsize)) !=0) {
1598     int errnoCopy = errno;
1599     VLOG(2) << "failed to update SO_SNDBUF option on AsyncSocket"
1600             << this << "(fd=" << fd_ << ", state=" << state_ << "): "
1601             << strerror(errnoCopy);
1602     return errnoCopy;
1603   }
1604
1605   return 0;
1606 }
1607
1608 int AsyncSocket::setRecvBufSize(size_t bufsize) {
1609   if (fd_ < 0) {
1610     VLOG(4) << "AsyncSocket::setRecvBufSize() called on non-open socket "
1611                << this << "(state=" << state_ << ")";
1612     return EINVAL;
1613   }
1614
1615   if (setsockopt(fd_, SOL_SOCKET, SO_RCVBUF, &bufsize, sizeof(bufsize)) !=0) {
1616     int errnoCopy = errno;
1617     VLOG(2) << "failed to update SO_RCVBUF option on AsyncSocket"
1618             << this << "(fd=" << fd_ << ", state=" << state_ << "): "
1619             << strerror(errnoCopy);
1620     return errnoCopy;
1621   }
1622
1623   return 0;
1624 }
1625
1626 int AsyncSocket::setTCPProfile(int profd) {
1627   if (fd_ < 0) {
1628     VLOG(4) << "AsyncSocket::setTCPProfile() called on non-open socket "
1629                << this << "(state=" << state_ << ")";
1630     return EINVAL;
1631   }
1632
1633   if (setsockopt(fd_, SOL_SOCKET, SO_SET_NAMESPACE, &profd, sizeof(int)) !=0) {
1634     int errnoCopy = errno;
1635     VLOG(2) << "failed to set socket namespace option on AsyncSocket"
1636             << this << "(fd=" << fd_ << ", state=" << state_ << "): "
1637             << strerror(errnoCopy);
1638     return errnoCopy;
1639   }
1640
1641   return 0;
1642 }
1643
1644 void AsyncSocket::ioReady(uint16_t events) noexcept {
1645   VLOG(7) << "AsyncSocket::ioRead() this=" << this << ", fd=" << fd_
1646           << ", events=" << std::hex << events << ", state=" << state_;
1647   DestructorGuard dg(this);
1648   assert(events & EventHandler::READ_WRITE);
1649   eventBase_->dcheckIsInEventBaseThread();
1650
1651   uint16_t relevantEvents = uint16_t(events & EventHandler::READ_WRITE);
1652   EventBase* originalEventBase = eventBase_;
1653   // If we got there it means that either EventHandler::READ or
1654   // EventHandler::WRITE is set. Any of these flags can
1655   // indicate that there are messages available in the socket
1656   // error message queue.
1657   handleErrMessages();
1658
1659   // Return now if handleErrMessages() detached us from our EventBase
1660   if (eventBase_ != originalEventBase) {
1661     return;
1662   }
1663
1664   if (relevantEvents == EventHandler::READ) {
1665     handleRead();
1666   } else if (relevantEvents == EventHandler::WRITE) {
1667     handleWrite();
1668   } else if (relevantEvents == EventHandler::READ_WRITE) {
1669     // If both read and write events are ready, process writes first.
1670     handleWrite();
1671
1672     // Return now if handleWrite() detached us from our EventBase
1673     if (eventBase_ != originalEventBase) {
1674       return;
1675     }
1676
1677     // Only call handleRead() if a read callback is still installed.
1678     // (It's possible that the read callback was uninstalled during
1679     // handleWrite().)
1680     if (readCallback_) {
1681       handleRead();
1682     }
1683   } else {
1684     VLOG(4) << "AsyncSocket::ioRead() called with unexpected events "
1685                << std::hex << events << "(this=" << this << ")";
1686     abort();
1687   }
1688 }
1689
1690 AsyncSocket::ReadResult
1691 AsyncSocket::performRead(void** buf, size_t* buflen, size_t* /* offset */) {
1692   VLOG(5) << "AsyncSocket::performRead() this=" << this << ", buf=" << *buf
1693           << ", buflen=" << *buflen;
1694
1695   if (preReceivedData_ && !preReceivedData_->empty()) {
1696     VLOG(5) << "AsyncSocket::performRead() this=" << this
1697             << ", reading pre-received data";
1698
1699     io::Cursor cursor(preReceivedData_.get());
1700     auto len = cursor.pullAtMost(*buf, *buflen);
1701
1702     IOBufQueue queue;
1703     queue.append(std::move(preReceivedData_));
1704     queue.trimStart(len);
1705     preReceivedData_ = queue.move();
1706
1707     appBytesReceived_ += len;
1708     return ReadResult(len);
1709   }
1710
1711   ssize_t bytes = recv(fd_, *buf, *buflen, MSG_DONTWAIT);
1712   if (bytes < 0) {
1713     if (errno == EAGAIN || errno == EWOULDBLOCK) {
1714       // No more data to read right now.
1715       return ReadResult(READ_BLOCKING);
1716     } else {
1717       return ReadResult(READ_ERROR);
1718     }
1719   } else {
1720     appBytesReceived_ += bytes;
1721     return ReadResult(bytes);
1722   }
1723 }
1724
1725 void AsyncSocket::prepareReadBuffer(void** buf, size_t* buflen) {
1726   // no matter what, buffer should be preapared for non-ssl socket
1727   CHECK(readCallback_);
1728   readCallback_->getReadBuffer(buf, buflen);
1729 }
1730
1731 void AsyncSocket::handleErrMessages() noexcept {
1732   // This method has non-empty implementation only for platforms
1733   // supporting per-socket error queues.
1734   VLOG(5) << "AsyncSocket::handleErrMessages() this=" << this << ", fd=" << fd_
1735           << ", state=" << state_;
1736   if (errMessageCallback_ == nullptr && idZeroCopyBufPtrMap_.empty()) {
1737     VLOG(7) << "AsyncSocket::handleErrMessages(): "
1738             << "no callback installed - exiting.";
1739     return;
1740   }
1741
1742 #ifdef FOLLY_HAVE_MSG_ERRQUEUE
1743   uint8_t ctrl[1024];
1744   unsigned char data;
1745   struct msghdr msg;
1746   iovec entry;
1747
1748   entry.iov_base = &data;
1749   entry.iov_len = sizeof(data);
1750   msg.msg_iov = &entry;
1751   msg.msg_iovlen = 1;
1752   msg.msg_name = nullptr;
1753   msg.msg_namelen = 0;
1754   msg.msg_control = ctrl;
1755   msg.msg_controllen = sizeof(ctrl);
1756   msg.msg_flags = 0;
1757
1758   int ret;
1759   while (true) {
1760     ret = recvmsg(fd_, &msg, MSG_ERRQUEUE);
1761     VLOG(5) << "AsyncSocket::handleErrMessages(): recvmsg returned " << ret;
1762
1763     if (ret < 0) {
1764       if (errno != EAGAIN) {
1765         auto errnoCopy = errno;
1766         LOG(ERROR) << "::recvmsg exited with code " << ret
1767                    << ", errno: " << errnoCopy;
1768         AsyncSocketException ex(
1769           AsyncSocketException::INTERNAL_ERROR,
1770           withAddr("recvmsg() failed"),
1771           errnoCopy);
1772         failErrMessageRead(__func__, ex);
1773       }
1774       return;
1775     }
1776
1777     for (struct cmsghdr* cmsg = CMSG_FIRSTHDR(&msg);
1778          cmsg != nullptr && cmsg->cmsg_len != 0;
1779          cmsg = CMSG_NXTHDR(&msg, cmsg)) {
1780       if (isZeroCopyMsg(*cmsg)) {
1781         processZeroCopyMsg(*cmsg);
1782       } else {
1783         if (errMessageCallback_) {
1784           errMessageCallback_->errMessage(*cmsg);
1785         }
1786       }
1787     }
1788   }
1789 #endif // FOLLY_HAVE_MSG_ERRQUEUE
1790 }
1791
1792 void AsyncSocket::handleRead() noexcept {
1793   VLOG(5) << "AsyncSocket::handleRead() this=" << this << ", fd=" << fd_
1794           << ", state=" << state_;
1795   assert(state_ == StateEnum::ESTABLISHED);
1796   assert((shutdownFlags_ & SHUT_READ) == 0);
1797   assert(readCallback_ != nullptr);
1798   assert(eventFlags_ & EventHandler::READ);
1799
1800   // Loop until:
1801   // - a read attempt would block
1802   // - readCallback_ is uninstalled
1803   // - the number of loop iterations exceeds the optional maximum
1804   // - this AsyncSocket is moved to another EventBase
1805   //
1806   // When we invoke readDataAvailable() it may uninstall the readCallback_,
1807   // which is why need to check for it here.
1808   //
1809   // The last bullet point is slightly subtle.  readDataAvailable() may also
1810   // detach this socket from this EventBase.  However, before
1811   // readDataAvailable() returns another thread may pick it up, attach it to
1812   // a different EventBase, and install another readCallback_.  We need to
1813   // exit immediately after readDataAvailable() returns if the eventBase_ has
1814   // changed.  (The caller must perform some sort of locking to transfer the
1815   // AsyncSocket between threads properly.  This will be sufficient to ensure
1816   // that this thread sees the updated eventBase_ variable after
1817   // readDataAvailable() returns.)
1818   uint16_t numReads = 0;
1819   EventBase* originalEventBase = eventBase_;
1820   while (readCallback_ && eventBase_ == originalEventBase) {
1821     // Get the buffer to read into.
1822     void* buf = nullptr;
1823     size_t buflen = 0, offset = 0;
1824     try {
1825       prepareReadBuffer(&buf, &buflen);
1826       VLOG(5) << "prepareReadBuffer() buf=" << buf << ", buflen=" << buflen;
1827     } catch (const AsyncSocketException& ex) {
1828       return failRead(__func__, ex);
1829     } catch (const std::exception& ex) {
1830       AsyncSocketException tex(AsyncSocketException::BAD_ARGS,
1831                               string("ReadCallback::getReadBuffer() "
1832                                      "threw exception: ") +
1833                               ex.what());
1834       return failRead(__func__, tex);
1835     } catch (...) {
1836       AsyncSocketException ex(AsyncSocketException::BAD_ARGS,
1837                              "ReadCallback::getReadBuffer() threw "
1838                              "non-exception type");
1839       return failRead(__func__, ex);
1840     }
1841     if (!isBufferMovable_ && (buf == nullptr || buflen == 0)) {
1842       AsyncSocketException ex(AsyncSocketException::BAD_ARGS,
1843                              "ReadCallback::getReadBuffer() returned "
1844                              "empty buffer");
1845       return failRead(__func__, ex);
1846     }
1847
1848     // Perform the read
1849     auto readResult = performRead(&buf, &buflen, &offset);
1850     auto bytesRead = readResult.readReturn;
1851     VLOG(4) << "this=" << this << ", AsyncSocket::handleRead() got "
1852             << bytesRead << " bytes";
1853     if (bytesRead > 0) {
1854       if (!isBufferMovable_) {
1855         readCallback_->readDataAvailable(size_t(bytesRead));
1856       } else {
1857         CHECK(kOpenSslModeMoveBufferOwnership);
1858         VLOG(5) << "this=" << this << ", AsyncSocket::handleRead() got "
1859                 << "buf=" << buf << ", " << bytesRead << "/" << buflen
1860                 << ", offset=" << offset;
1861         auto readBuf = folly::IOBuf::takeOwnership(buf, buflen);
1862         readBuf->trimStart(offset);
1863         readBuf->trimEnd(buflen - offset - bytesRead);
1864         readCallback_->readBufferAvailable(std::move(readBuf));
1865       }
1866
1867       // Fall through and continue around the loop if the read
1868       // completely filled the available buffer.
1869       // Note that readCallback_ may have been uninstalled or changed inside
1870       // readDataAvailable().
1871       if (size_t(bytesRead) < buflen) {
1872         return;
1873       }
1874     } else if (bytesRead == READ_BLOCKING) {
1875         // No more data to read right now.
1876         return;
1877     } else if (bytesRead == READ_ERROR) {
1878       readErr_ = READ_ERROR;
1879       if (readResult.exception) {
1880         return failRead(__func__, *readResult.exception);
1881       }
1882       auto errnoCopy = errno;
1883       AsyncSocketException ex(
1884           AsyncSocketException::INTERNAL_ERROR,
1885           withAddr("recv() failed"),
1886           errnoCopy);
1887       return failRead(__func__, ex);
1888     } else {
1889       assert(bytesRead == READ_EOF);
1890       readErr_ = READ_EOF;
1891       // EOF
1892       shutdownFlags_ |= SHUT_READ;
1893       if (!updateEventRegistration(0, EventHandler::READ)) {
1894         // we've already been moved into STATE_ERROR
1895         assert(state_ == StateEnum::ERROR);
1896         assert(readCallback_ == nullptr);
1897         return;
1898       }
1899
1900       ReadCallback* callback = readCallback_;
1901       readCallback_ = nullptr;
1902       callback->readEOF();
1903       return;
1904     }
1905     if (maxReadsPerEvent_ && (++numReads >= maxReadsPerEvent_)) {
1906       if (readCallback_ != nullptr) {
1907         // We might still have data in the socket.
1908         // (e.g. see comment in AsyncSSLSocket::checkForImmediateRead)
1909         scheduleImmediateRead();
1910       }
1911       return;
1912     }
1913   }
1914 }
1915
1916 /**
1917  * This function attempts to write as much data as possible, until no more data
1918  * can be written.
1919  *
1920  * - If it sends all available data, it unregisters for write events, and stops
1921  *   the writeTimeout_.
1922  *
1923  * - If not all of the data can be sent immediately, it reschedules
1924  *   writeTimeout_ (if a non-zero timeout is set), and ensures the handler is
1925  *   registered for write events.
1926  */
1927 void AsyncSocket::handleWrite() noexcept {
1928   VLOG(5) << "AsyncSocket::handleWrite() this=" << this << ", fd=" << fd_
1929           << ", state=" << state_;
1930   DestructorGuard dg(this);
1931
1932   if (state_ == StateEnum::CONNECTING) {
1933     handleConnect();
1934     return;
1935   }
1936
1937   // Normal write
1938   assert(state_ == StateEnum::ESTABLISHED);
1939   assert((shutdownFlags_ & SHUT_WRITE) == 0);
1940   assert(writeReqHead_ != nullptr);
1941
1942   // Loop until we run out of write requests,
1943   // or until this socket is moved to another EventBase.
1944   // (See the comment in handleRead() explaining how this can happen.)
1945   EventBase* originalEventBase = eventBase_;
1946   while (writeReqHead_ != nullptr && eventBase_ == originalEventBase) {
1947     auto writeResult = writeReqHead_->performWrite();
1948     if (writeResult.writeReturn < 0) {
1949       if (writeResult.exception) {
1950         return failWrite(__func__, *writeResult.exception);
1951       }
1952       auto errnoCopy = errno;
1953       AsyncSocketException ex(
1954           AsyncSocketException::INTERNAL_ERROR,
1955           withAddr("writev() failed"),
1956           errnoCopy);
1957       return failWrite(__func__, ex);
1958     } else if (writeReqHead_->isComplete()) {
1959       // We finished this request
1960       WriteRequest* req = writeReqHead_;
1961       writeReqHead_ = req->getNext();
1962
1963       if (writeReqHead_ == nullptr) {
1964         writeReqTail_ = nullptr;
1965         // This is the last write request.
1966         // Unregister for write events and cancel the send timer
1967         // before we invoke the callback.  We have to update the state properly
1968         // before calling the callback, since it may want to detach us from
1969         // the EventBase.
1970         if (eventFlags_ & EventHandler::WRITE) {
1971           if (!updateEventRegistration(0, EventHandler::WRITE)) {
1972             assert(state_ == StateEnum::ERROR);
1973             return;
1974           }
1975           // Stop the send timeout
1976           writeTimeout_.cancelTimeout();
1977         }
1978         assert(!writeTimeout_.isScheduled());
1979
1980         // If SHUT_WRITE_PENDING is set, we should shutdown the socket after
1981         // we finish sending the last write request.
1982         //
1983         // We have to do this before invoking writeSuccess(), since
1984         // writeSuccess() may detach us from our EventBase.
1985         if (shutdownFlags_ & SHUT_WRITE_PENDING) {
1986           assert(connectCallback_ == nullptr);
1987           shutdownFlags_ |= SHUT_WRITE;
1988
1989           if (shutdownFlags_ & SHUT_READ) {
1990             // Reads have already been shutdown.  Fully close the socket and
1991             // move to STATE_CLOSED.
1992             //
1993             // Note: This code currently moves us to STATE_CLOSED even if
1994             // close() hasn't ever been called.  This can occur if we have
1995             // received EOF from the peer and shutdownWrite() has been called
1996             // locally.  Should we bother staying in STATE_ESTABLISHED in this
1997             // case, until close() is actually called?  I can't think of a
1998             // reason why we would need to do so.  No other operations besides
1999             // calling close() or destroying the socket can be performed at
2000             // this point.
2001             assert(readCallback_ == nullptr);
2002             state_ = StateEnum::CLOSED;
2003             if (fd_ >= 0) {
2004               ioHandler_.changeHandlerFD(-1);
2005               doClose();
2006             }
2007           } else {
2008             // Reads are still enabled, so we are only doing a half-shutdown
2009             shutdown(fd_, SHUT_WR);
2010           }
2011         }
2012       }
2013
2014       // Invoke the callback
2015       WriteCallback* callback = req->getCallback();
2016       req->destroy();
2017       if (callback) {
2018         callback->writeSuccess();
2019       }
2020       // We'll continue around the loop, trying to write another request
2021     } else {
2022       // Partial write.
2023       if (bufferCallback_) {
2024         bufferCallback_->onEgressBuffered();
2025       }
2026       writeReqHead_->consume();
2027       // Stop after a partial write; it's highly likely that a subsequent write
2028       // attempt will just return EAGAIN.
2029       //
2030       // Ensure that we are registered for write events.
2031       if ((eventFlags_ & EventHandler::WRITE) == 0) {
2032         if (!updateEventRegistration(EventHandler::WRITE, 0)) {
2033           assert(state_ == StateEnum::ERROR);
2034           return;
2035         }
2036       }
2037
2038       // Reschedule the send timeout, since we have made some write progress.
2039       if (sendTimeout_ > 0) {
2040         if (!writeTimeout_.scheduleTimeout(sendTimeout_)) {
2041           AsyncSocketException ex(AsyncSocketException::INTERNAL_ERROR,
2042               withAddr("failed to reschedule write timeout"));
2043           return failWrite(__func__, ex);
2044         }
2045       }
2046       return;
2047     }
2048   }
2049   if (!writeReqHead_ && bufferCallback_) {
2050     bufferCallback_->onEgressBufferCleared();
2051   }
2052 }
2053
2054 void AsyncSocket::checkForImmediateRead() noexcept {
2055   // We currently don't attempt to perform optimistic reads in AsyncSocket.
2056   // (However, note that some subclasses do override this method.)
2057   //
2058   // Simply calling handleRead() here would be bad, as this would call
2059   // readCallback_->getReadBuffer(), forcing the callback to allocate a read
2060   // buffer even though no data may be available.  This would waste lots of
2061   // memory, since the buffer will sit around unused until the socket actually
2062   // becomes readable.
2063   //
2064   // Checking if the socket is readable now also seems like it would probably
2065   // be a pessimism.  In most cases it probably wouldn't be readable, and we
2066   // would just waste an extra system call.  Even if it is readable, waiting to
2067   // find out from libevent on the next event loop doesn't seem that bad.
2068   //
2069   // The exception to this is if we have pre-received data. In that case there
2070   // is definitely data available immediately.
2071   if (preReceivedData_ && !preReceivedData_->empty()) {
2072     handleRead();
2073   }
2074 }
2075
2076 void AsyncSocket::handleInitialReadWrite() noexcept {
2077   // Our callers should already be holding a DestructorGuard, but grab
2078   // one here just to make sure, in case one of our calling code paths ever
2079   // changes.
2080   DestructorGuard dg(this);
2081   // If we have a readCallback_, make sure we enable read events.  We
2082   // may already be registered for reads if connectSuccess() set
2083   // the read calback.
2084   if (readCallback_ && !(eventFlags_ & EventHandler::READ)) {
2085     assert(state_ == StateEnum::ESTABLISHED);
2086     assert((shutdownFlags_ & SHUT_READ) == 0);
2087     if (!updateEventRegistration(EventHandler::READ, 0)) {
2088       assert(state_ == StateEnum::ERROR);
2089       return;
2090     }
2091     checkForImmediateRead();
2092   } else if (readCallback_ == nullptr) {
2093     // Unregister for read events.
2094     updateEventRegistration(0, EventHandler::READ);
2095   }
2096
2097   // If we have write requests pending, try to send them immediately.
2098   // Since we just finished accepting, there is a very good chance that we can
2099   // write without blocking.
2100   //
2101   // However, we only process them if EventHandler::WRITE is not already set,
2102   // which means that we're already blocked on a write attempt.  (This can
2103   // happen if connectSuccess() called write() before returning.)
2104   if (writeReqHead_ && !(eventFlags_ & EventHandler::WRITE)) {
2105     // Call handleWrite() to perform write processing.
2106     handleWrite();
2107   } else if (writeReqHead_ == nullptr) {
2108     // Unregister for write event.
2109     updateEventRegistration(0, EventHandler::WRITE);
2110   }
2111 }
2112
2113 void AsyncSocket::handleConnect() noexcept {
2114   VLOG(5) << "AsyncSocket::handleConnect() this=" << this << ", fd=" << fd_
2115           << ", state=" << state_;
2116   assert(state_ == StateEnum::CONNECTING);
2117   // SHUT_WRITE can never be set while we are still connecting;
2118   // SHUT_WRITE_PENDING may be set, be we only set SHUT_WRITE once the connect
2119   // finishes
2120   assert((shutdownFlags_ & SHUT_WRITE) == 0);
2121
2122   // In case we had a connect timeout, cancel the timeout
2123   writeTimeout_.cancelTimeout();
2124   // We don't use a persistent registration when waiting on a connect event,
2125   // so we have been automatically unregistered now.  Update eventFlags_ to
2126   // reflect reality.
2127   assert(eventFlags_ == EventHandler::WRITE);
2128   eventFlags_ = EventHandler::NONE;
2129
2130   // Call getsockopt() to check if the connect succeeded
2131   int error;
2132   socklen_t len = sizeof(error);
2133   int rv = getsockopt(fd_, SOL_SOCKET, SO_ERROR, &error, &len);
2134   if (rv != 0) {
2135     auto errnoCopy = errno;
2136     AsyncSocketException ex(
2137         AsyncSocketException::INTERNAL_ERROR,
2138         withAddr("error calling getsockopt() after connect"),
2139         errnoCopy);
2140     VLOG(4) << "AsyncSocket::handleConnect(this=" << this << ", fd="
2141                << fd_ << " host=" << addr_.describe()
2142                << ") exception:" << ex.what();
2143     return failConnect(__func__, ex);
2144   }
2145
2146   if (error != 0) {
2147     AsyncSocketException ex(AsyncSocketException::NOT_OPEN,
2148                            "connect failed", error);
2149     VLOG(1) << "AsyncSocket::handleConnect(this=" << this << ", fd="
2150             << fd_ << " host=" << addr_.describe()
2151             << ") exception: " << ex.what();
2152     return failConnect(__func__, ex);
2153   }
2154
2155   // Move into STATE_ESTABLISHED
2156   state_ = StateEnum::ESTABLISHED;
2157
2158   // If SHUT_WRITE_PENDING is set and we don't have any write requests to
2159   // perform, immediately shutdown the write half of the socket.
2160   if ((shutdownFlags_ & SHUT_WRITE_PENDING) && writeReqHead_ == nullptr) {
2161     // SHUT_READ shouldn't be set.  If close() is called on the socket while we
2162     // are still connecting we just abort the connect rather than waiting for
2163     // it to complete.
2164     assert((shutdownFlags_ & SHUT_READ) == 0);
2165     shutdown(fd_, SHUT_WR);
2166     shutdownFlags_ |= SHUT_WRITE;
2167   }
2168
2169   VLOG(7) << "AsyncSocket " << this << ": fd " << fd_
2170           << "successfully connected; state=" << state_;
2171
2172   // Remember the EventBase we are attached to, before we start invoking any
2173   // callbacks (since the callbacks may call detachEventBase()).
2174   EventBase* originalEventBase = eventBase_;
2175
2176   invokeConnectSuccess();
2177   // Note that the connect callback may have changed our state.
2178   // (set or unset the read callback, called write(), closed the socket, etc.)
2179   // The following code needs to handle these situations correctly.
2180   //
2181   // If the socket has been closed, readCallback_ and writeReqHead_ will
2182   // always be nullptr, so that will prevent us from trying to read or write.
2183   //
2184   // The main thing to check for is if eventBase_ is still originalEventBase.
2185   // If not, we have been detached from this event base, so we shouldn't
2186   // perform any more operations.
2187   if (eventBase_ != originalEventBase) {
2188     return;
2189   }
2190
2191   handleInitialReadWrite();
2192 }
2193
2194 void AsyncSocket::timeoutExpired() noexcept {
2195   VLOG(7) << "AsyncSocket " << this << ", fd " << fd_ << ": timeout expired: "
2196           << "state=" << state_ << ", events=" << std::hex << eventFlags_;
2197   DestructorGuard dg(this);
2198   eventBase_->dcheckIsInEventBaseThread();
2199
2200   if (state_ == StateEnum::CONNECTING) {
2201     // connect() timed out
2202     // Unregister for I/O events.
2203     if (connectCallback_) {
2204       AsyncSocketException ex(
2205           AsyncSocketException::TIMED_OUT,
2206           folly::sformat(
2207               "connect timed out after {}ms", connectTimeout_.count()));
2208       failConnect(__func__, ex);
2209     } else {
2210       // we faced a connect error without a connect callback, which could
2211       // happen due to TFO.
2212       AsyncSocketException ex(
2213           AsyncSocketException::TIMED_OUT, "write timed out during connection");
2214       failWrite(__func__, ex);
2215     }
2216   } else {
2217     // a normal write operation timed out
2218     AsyncSocketException ex(
2219         AsyncSocketException::TIMED_OUT,
2220         folly::sformat("write timed out after {}ms", sendTimeout_));
2221     failWrite(__func__, ex);
2222   }
2223 }
2224
2225 ssize_t AsyncSocket::tfoSendMsg(int fd, struct msghdr* msg, int msg_flags) {
2226   return detail::tfo_sendmsg(fd, msg, msg_flags);
2227 }
2228
2229 AsyncSocket::WriteResult
2230 AsyncSocket::sendSocketMessage(int fd, struct msghdr* msg, int msg_flags) {
2231   ssize_t totalWritten = 0;
2232   if (state_ == StateEnum::FAST_OPEN) {
2233     sockaddr_storage addr;
2234     auto len = addr_.getAddress(&addr);
2235     msg->msg_name = &addr;
2236     msg->msg_namelen = len;
2237     totalWritten = tfoSendMsg(fd_, msg, msg_flags);
2238     if (totalWritten >= 0) {
2239       tfoFinished_ = true;
2240       state_ = StateEnum::ESTABLISHED;
2241       // We schedule this asynchrously so that we don't end up
2242       // invoking initial read or write while a write is in progress.
2243       scheduleInitialReadWrite();
2244     } else if (errno == EINPROGRESS) {
2245       VLOG(4) << "TFO falling back to connecting";
2246       // A normal sendmsg doesn't return EINPROGRESS, however
2247       // TFO might fallback to connecting if there is no
2248       // cookie.
2249       state_ = StateEnum::CONNECTING;
2250       try {
2251         scheduleConnectTimeout();
2252         registerForConnectEvents();
2253       } catch (const AsyncSocketException& ex) {
2254         return WriteResult(
2255             WRITE_ERROR, std::make_unique<AsyncSocketException>(ex));
2256       }
2257       // Let's fake it that no bytes were written and return an errno.
2258       errno = EAGAIN;
2259       totalWritten = -1;
2260     } else if (errno == EOPNOTSUPP) {
2261       // Try falling back to connecting.
2262       VLOG(4) << "TFO not supported";
2263       state_ = StateEnum::CONNECTING;
2264       try {
2265         int ret = socketConnect((const sockaddr*)&addr, len);
2266         if (ret == 0) {
2267           // connect succeeded immediately
2268           // Treat this like no data was written.
2269           state_ = StateEnum::ESTABLISHED;
2270           scheduleInitialReadWrite();
2271         }
2272         // If there was no exception during connections,
2273         // we would return that no bytes were written.
2274         errno = EAGAIN;
2275         totalWritten = -1;
2276       } catch (const AsyncSocketException& ex) {
2277         return WriteResult(
2278             WRITE_ERROR, std::make_unique<AsyncSocketException>(ex));
2279       }
2280     } else if (errno == EAGAIN) {
2281       // Normally sendmsg would indicate that the write would block.
2282       // However in the fast open case, it would indicate that sendmsg
2283       // fell back to a connect. This is a return code from connect()
2284       // instead, and is an error condition indicating no fds available.
2285       return WriteResult(
2286           WRITE_ERROR,
2287           std::make_unique<AsyncSocketException>(
2288               AsyncSocketException::UNKNOWN, "No more free local ports"));
2289     }
2290   } else {
2291     totalWritten = ::sendmsg(fd, msg, msg_flags);
2292   }
2293   return WriteResult(totalWritten);
2294 }
2295
2296 AsyncSocket::WriteResult AsyncSocket::performWrite(
2297     const iovec* vec,
2298     uint32_t count,
2299     WriteFlags flags,
2300     uint32_t* countWritten,
2301     uint32_t* partialWritten) {
2302   // We use sendmsg() instead of writev() so that we can pass in MSG_NOSIGNAL
2303   // We correctly handle EPIPE errors, so we never want to receive SIGPIPE
2304   // (since it may terminate the program if the main program doesn't explicitly
2305   // ignore it).
2306   struct msghdr msg;
2307   msg.msg_name = nullptr;
2308   msg.msg_namelen = 0;
2309   msg.msg_iov = const_cast<iovec *>(vec);
2310   msg.msg_iovlen = std::min<size_t>(count, kIovMax);
2311   msg.msg_flags = 0;
2312   msg.msg_controllen = sendMsgParamCallback_->getAncillaryDataSize(flags);
2313   CHECK_GE(AsyncSocket::SendMsgParamsCallback::maxAncillaryDataSize,
2314            msg.msg_controllen);
2315
2316   if (msg.msg_controllen != 0) {
2317     msg.msg_control = reinterpret_cast<char*>(alloca(msg.msg_controllen));
2318     sendMsgParamCallback_->getAncillaryData(flags, msg.msg_control);
2319   } else {
2320     msg.msg_control = nullptr;
2321   }
2322   int msg_flags = sendMsgParamCallback_->getFlags(flags, zeroCopyEnabled_);
2323
2324   auto writeResult = sendSocketMessage(fd_, &msg, msg_flags);
2325   auto totalWritten = writeResult.writeReturn;
2326   if (totalWritten < 0) {
2327     bool tryAgain = (errno == EAGAIN);
2328 #ifdef __APPLE__
2329     // Apple has a bug where doing a second write on a socket which we
2330     // have opened with TFO causes an ENOTCONN to be thrown. However the
2331     // socket is really connected, so treat ENOTCONN as a EAGAIN until
2332     // this bug is fixed.
2333     tryAgain |= (errno == ENOTCONN);
2334 #endif
2335
2336     // workaround for running with zerocopy enabled but without a proper
2337     // memlock value - see ulimit -l
2338     if (zeroCopyEnabled_ && (errno == ENOBUFS)) {
2339       tryAgain = true;
2340       zeroCopyEnabled_ = false;
2341     }
2342
2343     if (!writeResult.exception && tryAgain) {
2344       // TCP buffer is full; we can't write any more data right now.
2345       *countWritten = 0;
2346       *partialWritten = 0;
2347       return WriteResult(0);
2348     }
2349     // error
2350     *countWritten = 0;
2351     *partialWritten = 0;
2352     return writeResult;
2353   }
2354
2355   appBytesWritten_ += totalWritten;
2356
2357   uint32_t bytesWritten;
2358   uint32_t n;
2359   for (bytesWritten = uint32_t(totalWritten), n = 0; n < count; ++n) {
2360     const iovec* v = vec + n;
2361     if (v->iov_len > bytesWritten) {
2362       // Partial write finished in the middle of this iovec
2363       *countWritten = n;
2364       *partialWritten = bytesWritten;
2365       return WriteResult(totalWritten);
2366     }
2367
2368     bytesWritten -= uint32_t(v->iov_len);
2369   }
2370
2371   assert(bytesWritten == 0);
2372   *countWritten = n;
2373   *partialWritten = 0;
2374   return WriteResult(totalWritten);
2375 }
2376
2377 /**
2378  * Re-register the EventHandler after eventFlags_ has changed.
2379  *
2380  * If an error occurs, fail() is called to move the socket into the error state
2381  * and call all currently installed callbacks.  After an error, the
2382  * AsyncSocket is completely unregistered.
2383  *
2384  * @return Returns true on success, or false on error.
2385  */
2386 bool AsyncSocket::updateEventRegistration() {
2387   VLOG(5) << "AsyncSocket::updateEventRegistration(this=" << this
2388           << ", fd=" << fd_ << ", evb=" << eventBase_ << ", state=" << state_
2389           << ", events=" << std::hex << eventFlags_;
2390   eventBase_->dcheckIsInEventBaseThread();
2391   if (eventFlags_ == EventHandler::NONE) {
2392     ioHandler_.unregisterHandler();
2393     return true;
2394   }
2395
2396   // Always register for persistent events, so we don't have to re-register
2397   // after being called back.
2398   if (!ioHandler_.registerHandler(
2399           uint16_t(eventFlags_ | EventHandler::PERSIST))) {
2400     eventFlags_ = EventHandler::NONE; // we're not registered after error
2401     AsyncSocketException ex(AsyncSocketException::INTERNAL_ERROR,
2402         withAddr("failed to update AsyncSocket event registration"));
2403     fail("updateEventRegistration", ex);
2404     return false;
2405   }
2406
2407   return true;
2408 }
2409
2410 bool AsyncSocket::updateEventRegistration(uint16_t enable,
2411                                            uint16_t disable) {
2412   uint16_t oldFlags = eventFlags_;
2413   eventFlags_ |= enable;
2414   eventFlags_ &= ~disable;
2415   if (eventFlags_ == oldFlags) {
2416     return true;
2417   } else {
2418     return updateEventRegistration();
2419   }
2420 }
2421
2422 void AsyncSocket::startFail() {
2423   // startFail() should only be called once
2424   assert(state_ != StateEnum::ERROR);
2425   assert(getDestructorGuardCount() > 0);
2426   state_ = StateEnum::ERROR;
2427   // Ensure that SHUT_READ and SHUT_WRITE are set,
2428   // so all future attempts to read or write will be rejected
2429   shutdownFlags_ |= (SHUT_READ | SHUT_WRITE);
2430
2431   if (eventFlags_ != EventHandler::NONE) {
2432     eventFlags_ = EventHandler::NONE;
2433     ioHandler_.unregisterHandler();
2434   }
2435   writeTimeout_.cancelTimeout();
2436
2437   if (fd_ >= 0) {
2438     ioHandler_.changeHandlerFD(-1);
2439     doClose();
2440   }
2441 }
2442
2443 void AsyncSocket::invokeAllErrors(const AsyncSocketException& ex) {
2444   invokeConnectErr(ex);
2445   failAllWrites(ex);
2446
2447   if (readCallback_) {
2448     ReadCallback* callback = readCallback_;
2449     readCallback_ = nullptr;
2450     callback->readErr(ex);
2451   }
2452 }
2453
2454 void AsyncSocket::finishFail() {
2455   assert(state_ == StateEnum::ERROR);
2456   assert(getDestructorGuardCount() > 0);
2457
2458   AsyncSocketException ex(
2459       AsyncSocketException::INTERNAL_ERROR,
2460       withAddr("socket closing after error"));
2461   invokeAllErrors(ex);
2462 }
2463
2464 void AsyncSocket::finishFail(const AsyncSocketException& ex) {
2465   assert(state_ == StateEnum::ERROR);
2466   assert(getDestructorGuardCount() > 0);
2467   invokeAllErrors(ex);
2468 }
2469
2470 void AsyncSocket::fail(const char* fn, const AsyncSocketException& ex) {
2471   VLOG(4) << "AsyncSocket(this=" << this << ", fd=" << fd_ << ", state="
2472              << state_ << " host=" << addr_.describe()
2473              << "): failed in " << fn << "(): "
2474              << ex.what();
2475   startFail();
2476   finishFail();
2477 }
2478
2479 void AsyncSocket::failConnect(const char* fn, const AsyncSocketException& ex) {
2480   VLOG(5) << "AsyncSocket(this=" << this << ", fd=" << fd_ << ", state="
2481                << state_ << " host=" << addr_.describe()
2482                << "): failed while connecting in " << fn << "(): "
2483                << ex.what();
2484   startFail();
2485
2486   invokeConnectErr(ex);
2487   finishFail(ex);
2488 }
2489
2490 void AsyncSocket::failRead(const char* fn, const AsyncSocketException& ex) {
2491   VLOG(5) << "AsyncSocket(this=" << this << ", fd=" << fd_ << ", state="
2492                << state_ << " host=" << addr_.describe()
2493                << "): failed while reading in " << fn << "(): "
2494                << ex.what();
2495   startFail();
2496
2497   if (readCallback_ != nullptr) {
2498     ReadCallback* callback = readCallback_;
2499     readCallback_ = nullptr;
2500     callback->readErr(ex);
2501   }
2502
2503   finishFail();
2504 }
2505
2506 void AsyncSocket::failErrMessageRead(const char* fn,
2507                                      const AsyncSocketException& ex) {
2508   VLOG(5) << "AsyncSocket(this=" << this << ", fd=" << fd_ << ", state="
2509                << state_ << " host=" << addr_.describe()
2510                << "): failed while reading message in " << fn << "(): "
2511                << ex.what();
2512   startFail();
2513
2514   if (errMessageCallback_ != nullptr) {
2515     ErrMessageCallback* callback = errMessageCallback_;
2516     errMessageCallback_ = nullptr;
2517     callback->errMessageError(ex);
2518   }
2519
2520   finishFail();
2521 }
2522
2523 void AsyncSocket::failWrite(const char* fn, const AsyncSocketException& ex) {
2524   VLOG(5) << "AsyncSocket(this=" << this << ", fd=" << fd_ << ", state="
2525                << state_ << " host=" << addr_.describe()
2526                << "): failed while writing in " << fn << "(): "
2527                << ex.what();
2528   startFail();
2529
2530   // Only invoke the first write callback, since the error occurred while
2531   // writing this request.  Let any other pending write callbacks be invoked in
2532   // finishFail().
2533   if (writeReqHead_ != nullptr) {
2534     WriteRequest* req = writeReqHead_;
2535     writeReqHead_ = req->getNext();
2536     WriteCallback* callback = req->getCallback();
2537     uint32_t bytesWritten = req->getTotalBytesWritten();
2538     req->destroy();
2539     if (callback) {
2540       callback->writeErr(bytesWritten, ex);
2541     }
2542   }
2543
2544   finishFail();
2545 }
2546
2547 void AsyncSocket::failWrite(const char* fn, WriteCallback* callback,
2548                              size_t bytesWritten,
2549                              const AsyncSocketException& ex) {
2550   // This version of failWrite() is used when the failure occurs before
2551   // we've added the callback to writeReqHead_.
2552   VLOG(4) << "AsyncSocket(this=" << this << ", fd=" << fd_ << ", state="
2553              << state_ << " host=" << addr_.describe()
2554              <<"): failed while writing in " << fn << "(): "
2555              << ex.what();
2556   startFail();
2557
2558   if (callback != nullptr) {
2559     callback->writeErr(bytesWritten, ex);
2560   }
2561
2562   finishFail();
2563 }
2564
2565 void AsyncSocket::failAllWrites(const AsyncSocketException& ex) {
2566   // Invoke writeError() on all write callbacks.
2567   // This is used when writes are forcibly shutdown with write requests
2568   // pending, or when an error occurs with writes pending.
2569   while (writeReqHead_ != nullptr) {
2570     WriteRequest* req = writeReqHead_;
2571     writeReqHead_ = req->getNext();
2572     WriteCallback* callback = req->getCallback();
2573     if (callback) {
2574       callback->writeErr(req->getTotalBytesWritten(), ex);
2575     }
2576     req->destroy();
2577   }
2578 }
2579
2580 void AsyncSocket::invalidState(ConnectCallback* callback) {
2581   VLOG(5) << "AsyncSocket(this=" << this << ", fd=" << fd_
2582           << "): connect() called in invalid state " << state_;
2583
2584   /*
2585    * The invalidState() methods don't use the normal failure mechanisms,
2586    * since we don't know what state we are in.  We don't want to call
2587    * startFail()/finishFail() recursively if we are already in the middle of
2588    * cleaning up.
2589    */
2590
2591   AsyncSocketException ex(AsyncSocketException::ALREADY_OPEN,
2592                          "connect() called with socket in invalid state");
2593   connectEndTime_ = std::chrono::steady_clock::now();
2594   if (state_ == StateEnum::CLOSED || state_ == StateEnum::ERROR) {
2595     if (callback) {
2596       callback->connectErr(ex);
2597     }
2598   } else {
2599     // We can't use failConnect() here since connectCallback_
2600     // may already be set to another callback.  Invoke this ConnectCallback
2601     // here; any other connectCallback_ will be invoked in finishFail()
2602     startFail();
2603     if (callback) {
2604       callback->connectErr(ex);
2605     }
2606     finishFail();
2607   }
2608 }
2609
2610 void AsyncSocket::invalidState(ErrMessageCallback* callback) {
2611   VLOG(4) << "AsyncSocket(this=" << this << ", fd=" << fd_
2612           << "): setErrMessageCB(" << callback
2613           << ") called in invalid state " << state_;
2614
2615   AsyncSocketException ex(
2616       AsyncSocketException::NOT_OPEN,
2617       msgErrQueueSupported
2618       ? "setErrMessageCB() called with socket in invalid state"
2619       : "This platform does not support socket error message notifications");
2620   if (state_ == StateEnum::CLOSED || state_ == StateEnum::ERROR) {
2621     if (callback) {
2622       callback->errMessageError(ex);
2623     }
2624   } else {
2625     startFail();
2626     if (callback) {
2627       callback->errMessageError(ex);
2628     }
2629     finishFail();
2630   }
2631 }
2632
2633 void AsyncSocket::invokeConnectErr(const AsyncSocketException& ex) {
2634   connectEndTime_ = std::chrono::steady_clock::now();
2635   if (connectCallback_) {
2636     ConnectCallback* callback = connectCallback_;
2637     connectCallback_ = nullptr;
2638     callback->connectErr(ex);
2639   }
2640 }
2641
2642 void AsyncSocket::invokeConnectSuccess() {
2643   connectEndTime_ = std::chrono::steady_clock::now();
2644   if (connectCallback_) {
2645     ConnectCallback* callback = connectCallback_;
2646     connectCallback_ = nullptr;
2647     callback->connectSuccess();
2648   }
2649 }
2650
2651 void AsyncSocket::invalidState(ReadCallback* callback) {
2652   VLOG(4) << "AsyncSocket(this=" << this << ", fd=" << fd_
2653              << "): setReadCallback(" << callback
2654              << ") called in invalid state " << state_;
2655
2656   AsyncSocketException ex(AsyncSocketException::NOT_OPEN,
2657                          "setReadCallback() called with socket in "
2658                          "invalid state");
2659   if (state_ == StateEnum::CLOSED || state_ == StateEnum::ERROR) {
2660     if (callback) {
2661       callback->readErr(ex);
2662     }
2663   } else {
2664     startFail();
2665     if (callback) {
2666       callback->readErr(ex);
2667     }
2668     finishFail();
2669   }
2670 }
2671
2672 void AsyncSocket::invalidState(WriteCallback* callback) {
2673   VLOG(4) << "AsyncSocket(this=" << this << ", fd=" << fd_
2674              << "): write() called in invalid state " << state_;
2675
2676   AsyncSocketException ex(AsyncSocketException::NOT_OPEN,
2677                          withAddr("write() called with socket in invalid state"));
2678   if (state_ == StateEnum::CLOSED || state_ == StateEnum::ERROR) {
2679     if (callback) {
2680       callback->writeErr(0, ex);
2681     }
2682   } else {
2683     startFail();
2684     if (callback) {
2685       callback->writeErr(0, ex);
2686     }
2687     finishFail();
2688   }
2689 }
2690
2691 void AsyncSocket::doClose() {
2692   if (fd_ == -1) {
2693     return;
2694   }
2695   if (const auto shutdownSocketSet = wShutdownSocketSet_.lock()) {
2696     shutdownSocketSet->close(fd_);
2697   } else {
2698     ::close(fd_);
2699   }
2700   fd_ = -1;
2701 }
2702
2703 std::ostream& operator << (std::ostream& os,
2704                            const AsyncSocket::StateEnum& state) {
2705   os << static_cast<int>(state);
2706   return os;
2707 }
2708
2709 std::string AsyncSocket::withAddr(const std::string& s) {
2710   // Don't use addr_ directly because it may not be initialized
2711   // e.g. if constructed from fd
2712   folly::SocketAddress peer, local;
2713   try {
2714     getPeerAddress(&peer);
2715     getLocalAddress(&local);
2716   } catch (const std::exception&) {
2717     // ignore
2718   } catch (...) {
2719     // ignore
2720   }
2721   return s + " (peer=" + peer.describe() + ", local=" + local.describe() + ")";
2722 }
2723
2724 void AsyncSocket::setBufferCallback(BufferCallback* cb) {
2725   bufferCallback_ = cb;
2726 }
2727
2728 } // namespace folly