Allow error message callback cancellation regardless of socket state.

[folly.git] / folly / io / async / AsyncSocket.cpp

diff --git a/folly/io/async/AsyncSocket.cpp b/folly/io/async/AsyncSocket.cpp
index 013813f92d0eef971e1b84fcb76796beb813c93a..17040bf9737df7ff04bd1e4908c064834001aa16 100644 (file)
--- a/folly/io/async/AsyncSocket.cpp
+++ b/folly/io/async/AsyncSocket.cpp
@@ -612,7 +612,7 @@ void AsyncSocket::setErrMessageCB(ErrMessageCallback* callback) {
           << ", fd=" << fd_ << ", callback=" << callback
           << ", state=" << state_;
 
-  // Short circuit if callback is the same as the existing timestampCallback_.
+  // Short circuit if callback is the same as the existing errMessageCallback_.
   if (callback == errMessageCallback_) {
     return;
   }
@@ -625,6 +625,14 @@ void AsyncSocket::setErrMessageCB(ErrMessageCallback* callback) {
   DestructorGuard dg(this);
   assert(eventBase_->isInEventBaseThread());
 
+  if (callback == nullptr) {
+    // We should be able to reset the callback regardless of the
+    // socket state. It's important to have a reliable callback
+    // cancellation mechanism.
+    errMessageCallback_ = callback;
+    return;
+  }
+
   switch ((StateEnum)state_) {
     case StateEnum::CONNECTING:
     case StateEnum::FAST_OPEN: