SERVER-24711 properly timeout asio connections

ASIO Connections that have already timed out before setup now result
in timeouts instead of invariants.

(cherry picked from commit 587afac9ca7ff084aeb8bf74e98ecb9688b3afdb)

3 files changed, 55 insertions, 1 deletions

diff --git a/src/mongo/executor/connection_pool_asio.cpp b/src/mongo/executor/connection_pool_asio.cpp
index 92ac06e14b9..a69ba48c52f 100644
--- a/src/mongo/executor/connection_pool_asio.cpp
+++ b/src/mongo/executor/connection_pool_asio.cpp
@@ -188,8 +188,36 @@ void ASIOConnection::setup(Milliseconds timeout, SetupCallback cb) {
             cb(ptr, status);
         };
 
+        std::size_t generation;
+        {
+            stdx::lock_guard<stdx::mutex> lk(_impl->_access->mutex);
+            generation = _impl->_access->id;
+        }
+
         // Actually timeout setup
-        setTimeout(timeout, [this] { _impl->connection().stream().cancel(); });
+        setTimeout(timeout,
+                   [this, generation] {
+                       // For operations that time out immediately, we can't simply cancel the
+                       // connection, because it may not have been initialized.
+                       stdx::lock_guard<stdx::mutex> lk(_impl->_access->mutex);
+
+                       auto access = _impl->_access;
+                       if (generation != access->id) {
+                           // The operation has been cleaned up, do not access.
+                           return;
+                       }
+
+                       // Time out the op.
+                       _impl->_strand.post([this, access, generation] {
+                           stdx::lock_guard<stdx::mutex> lk(access->mutex);
+                           if (generation == access->id) {
+                               _impl->_timedOut = true;
+                               if (_impl->_connection) {
+                                   _impl->_connection->cancel();
+                               }
+                           }
+                       });
+                   });
 
         _global->_impl->_connect(_impl.get());
     });
diff --git a/src/mongo/executor/connection_pool_asio_integration_test.cpp b/src/mongo/executor/connection_pool_asio_integration_test.cpp
index 3a8460e1f99..a1a347ac762 100644
--- a/src/mongo/executor/connection_pool_asio_integration_test.cpp
+++ b/src/mongo/executor/connection_pool_asio_integration_test.cpp
@@ -149,6 +149,31 @@ TEST(ConnectionPoolASIO, TestHostTimeoutRace) {
     }
 }
 
+
+/**
+ * Verify that a connections that timeout immediately don't invariant.
+ */
+TEST(ConnectionPoolASIO, ConnSetupTimeout) {
+    auto fixture = unittest::getFixtureConnectionString();
+
+    NetworkInterfaceASIO::Options options;
+    options.streamFactory = stdx::make_unique<AsyncStreamFactory>();
+    options.timerFactory = stdx::make_unique<AsyncTimerFactoryASIO>();
+    options.connectionPoolOptions.refreshTimeout = Milliseconds(-2);
+    NetworkInterfaceASIO net{std::move(options)};
+
+    net.startup();
+    auto guard = MakeGuard([&] { net.shutdown(); });
+
+    Deferred<StatusWith<RemoteCommandResponse>> deferred;
+    net.startCommand(
+        makeCallbackHandle(),
+        RemoteCommandRequest{fixture.getServers()[0], "admin", BSON("ping" << 1), BSONObj()},
+        [&](StatusWith<RemoteCommandResponse> resp) { deferred.emplace(std::move(resp)); });
+
+    ASSERT_EQ(deferred.get().getStatus().code(), ErrorCodes::ExceededTimeLimit);
+}
+
 }  // namespace
 }  // namespace executor
 }  // namespace mongo
diff --git a/src/mongo/executor/network_interface_asio.h b/src/mongo/executor/network_interface_asio.h
index f4eb6eac110..9632a37bef9 100644
--- a/src/mongo/executor/network_interface_asio.h
+++ b/src/mongo/executor/network_interface_asio.h
@@ -226,6 +226,7 @@ private:
      */
     class AsyncOp {
         friend class NetworkInterfaceASIO;
+        friend class connection_pool_asio::ASIOConnection;
 
     public:
         AsyncOp(NetworkInterfaceASIO* net,