SERVER-19842 support allowPartialResults query option in new mongos read path

1 files changed, 19 insertions, 3 deletions

diff --git a/src/mongo/s/query/async_results_merger.cpp b/src/mongo/s/query/async_results_merger.cpp
index ef039eadb2f..abecbba1a34 100644
--- a/src/mongo/s/query/async_results_merger.cpp
+++ b/src/mongo/s/query/async_results_merger.cpp
@@ -369,9 +369,14 @@ void AsyncResultsMerger::handleBatchResponse(
     if (!cbData.response.isOK()) {
         remote.status = cbData.response.getStatus();
 
-        // If we failed to retrieve the batch because we couldn't contact the remote, we notify that
-        // targeter that the host is unreachable. The caller can then retry on a new host.
-        if (remote.status == ErrorCodes::HostUnreachable && remote.shardId) {
+        // Errors other than HostUnreachable have no special handling.
+        if (remote.status != ErrorCodes::HostUnreachable) {
+            return;
+        }
+
+        // Notify that targeter that the host is unreachable. The caller can then retry on a new
+        // host.
+        if (remote.shardId) {
             auto shard = _params.shardRegistry->getShard(_params.txn, *remote.shardId);
             if (!shard) {
                 remote.status =
@@ -383,6 +388,17 @@ void AsyncResultsMerger::handleBatchResponse(
             }
         }
 
+        // Unreachable host errors are swallowed if the 'allowPartialResults' option is set. We
+        // remove the unreachable host entirely from consideration by marking it as exhausted.
+        if (_params.isAllowPartialResults) {
+            remote.status = Status::OK();
+
+            // Clear the results buffer and cursor id.
+            std::queue<BSONObj> emptyBuffer;
+            std::swap(remote.docBuffer, emptyBuffer);
+            remote.cursorId = 0;
+        }
+
         return;
     }