SERVER-29141 Extend StageConstraints to allow kLocalOnly and kMongos host requirements, given pipeline split state

1 files changed, 6 insertions, 3 deletions

diff --git a/src/mongo/db/pipeline/document_source_check_resume_token.h b/src/mongo/db/pipeline/document_source_check_resume_token.h
index 09a45d8277b..2ff83bb2b35 100644
--- a/src/mongo/db/pipeline/document_source_check_resume_token.h
+++ b/src/mongo/db/pipeline/document_source_check_resume_token.h
@@ -60,7 +60,7 @@ public:
     GetNextResult getNext() final;
     const char* getSourceName() const final;
 
-    StageConstraints constraints() const final {
+    StageConstraints constraints(Pipeline::SplitState pipeState) const final {
         return {StreamType::kStreaming,
                 PositionRequirement::kNone,
                 HostTypeRequirement::kAnyShard,
@@ -95,10 +95,13 @@ public:
     GetNextResult getNext() final;
     const char* getSourceName() const final;
 
-    StageConstraints constraints() const final {
+    StageConstraints constraints(Pipeline::SplitState pipeState) const final {
+        // This stage should never be in the shards part of a split pipeline.
+        invariant(pipeState != Pipeline::SplitState::kSplitForShards);
         return {StreamType::kStreaming,
                 PositionRequirement::kNone,
-                HostTypeRequirement::kNone,
+                (pipeState == Pipeline::SplitState::kUnsplit ? HostTypeRequirement::kNone
+                                                             : HostTypeRequirement::kMongoS),
                 DiskUseRequirement::kNoDiskUse,
                 FacetRequirement::kNotAllowed};
     }