SERVER-34725: Group and count plan stages do not set the WorkingSetID output on PlanStage::DEAD state

1 files changed, 14 insertions, 3 deletions

diff --git a/src/mongo/db/exec/queued_data_stage.cpp b/src/mongo/db/exec/queued_data_stage.cpp
index 6f3f3b0ed2f..69ff26330f5 100644
--- a/src/mongo/db/exec/queued_data_stage.cpp
+++ b/src/mongo/db/exec/queued_data_stage.cpp
@@ -51,9 +51,20 @@ PlanStage::StageState QueuedDataStage::doWork(WorkingSetID* out) {
     StageState state = _results.front();
     _results.pop();
 
-    if (PlanStage::ADVANCED == state) {
-        *out = _members.front();
-        _members.pop();
+    switch (state) {
+        case PlanStage::ADVANCED:
+            *out = _members.front();
+            _members.pop();
+            break;
+        case PlanStage::DEAD:
+        case PlanStage::FAILURE:
+            // On DEAD or FAILURE, this stage is reponsible for allocating the WorkingSetMember with
+            // the error details.
+            *out = WorkingSetCommon::allocateStatusMember(
+                _ws, Status(ErrorCodes::InternalError, "Queued data stage failure"));
+            break;
+        default:
+            break;
     }
 
     return state;