SERVER-20689 onFinish should be set when ConnectionPool refreshes a connection

4 files changed, 15 insertions, 0 deletions

diff --git a/src/mongo/executor/connection_pool_asio.cpp b/src/mongo/executor/connection_pool_asio.cpp
index 180a85d7f69..89d418c0a43 100644
--- a/src/mongo/executor/connection_pool_asio.cpp
+++ b/src/mongo/executor/connection_pool_asio.cpp
@@ -185,6 +185,14 @@ void ASIOConnection::refresh(Milliseconds timeout, RefreshCallback cb) {
         return;
     }
 
+    // If we fail during refresh, the _onFinish function of the AsyncOp will get called. As such we
+    // need to intercept those calls so we can capture them. This will get cleared out when we fill
+    // in the real onFinish in startCommand.
+    op->setOnFinish([this, cb](StatusWith<RemoteCommandResponse> failedResponse) {
+        invariant(!failedResponse.isOK());
+        cb(this, failedResponse.getStatus());
+    });
+
     _global->_impl->_asyncRunCommand(
         op,
         [this, op](std::error_code ec, size_t bytes) {
diff --git a/src/mongo/executor/network_interface_asio.cpp b/src/mongo/executor/network_interface_asio.cpp
index 3de99bfaf1c..fe999ef8ff0 100644
--- a/src/mongo/executor/network_interface_asio.cpp
+++ b/src/mongo/executor/network_interface_asio.cpp
@@ -154,6 +154,7 @@ Date_t NetworkInterfaceASIO::now() {
 void NetworkInterfaceASIO::startCommand(const TaskExecutor::CallbackHandle& cbHandle,
                                         const RemoteCommandRequest& request,
                                         const RemoteCommandCompletionFn& onFinish) {
+    invariant(onFinish);
     {
         stdx::lock_guard<stdx::mutex> lk(_inProgressMutex);
         const auto insertResult = _inGetConnection.emplace(cbHandle);
diff --git a/src/mongo/executor/network_interface_asio.h b/src/mongo/executor/network_interface_asio.h
index 864e5402052..2c44039b48c 100644
--- a/src/mongo/executor/network_interface_asio.h
+++ b/src/mongo/executor/network_interface_asio.h
@@ -275,6 +275,8 @@ private:
 
         void reset();
 
+        void setOnFinish(RemoteCommandCompletionFn&& onFinish);
+
     private:
         NetworkInterfaceASIO* const _owner;
         // Information describing a task enqueued on the NetworkInterface
diff --git a/src/mongo/executor/network_interface_asio_operation.cpp b/src/mongo/executor/network_interface_asio_operation.cpp
index 6903d784efb..d5ca86d8cc8 100644
--- a/src/mongo/executor/network_interface_asio_operation.cpp
+++ b/src/mongo/executor/network_interface_asio_operation.cpp
@@ -242,5 +242,9 @@ void NetworkInterfaceASIO::AsyncOp::reset() {
     // _inSetup should always be false at this point.
 }
 
+void NetworkInterfaceASIO::AsyncOp::setOnFinish(RemoteCommandCompletionFn&& onFinish) {
+    _onFinish = std::move(onFinish);
+}
+
 }  // namespace executor
 }  // namespace mongo