SERVER-85848 Treat $redact as single-doc transform for change stream stages (#18856)r7.0.6-rc0r7.0.6

The internal stages DSCSEnsureResumeTokenPresent and
DSCSHandleTopologyChange are permitted to swap with single-document
transformation stages during optimization to allow user-specified stages
to be moved closer to the front of the pipeline and to be pushed down to
the shards in a cluster deployment. The $redact stage is allowed in a
change stream pipeline, but it is not internally implemented or
recognised as a single-document transformation; it therefore does not
get moved ahead of the change stream stages mentioned above. This can in
turn decrease the performance for stages after $redact and fail the
$changeStreamSplitLargeEvents stage in a cluster deployment. In this
change we treat the $redact stage similar to single-document
transformation stages during the change stream pipeline optimisation.

GitOrigin-RevId: 66cdc1f28172cb33ff68263050d73d4ade73b9a4

8 files changed, 180 insertions, 10 deletions

diff --git a/etc/backports_required_for_multiversion_tests.yml b/etc/backports_required_for_multiversion_tests.yml
index aa6f159c08c..96ad5bb740e 100644
--- a/etc/backports_required_for_multiversion_tests.yml
+++ b/etc/backports_required_for_multiversion_tests.yml
@@ -457,6 +457,8 @@ last-continuous:
     ticket: SERVER-83119
   - test_file: jstests/sharding/clustered_coll_scan.js
     ticket: SERVER-83119
+  - test_file: jstests/change_streams/split_large_event_with_other_stages.js
+    ticket: SERVER-85848
   suites: null
 last-lts:
   all:
@@ -956,4 +958,6 @@ last-lts:
     ticket: SERVER-83119
   - test_file: jstests/sharding/clustered_coll_scan.js
     ticket: SERVER-83119
+  - test_file: jstests/change_streams/split_large_event_with_other_stages.js
+    ticket: SERVER-85848
   suites: null
diff --git a/jstests/change_streams/split_large_event_with_other_stages.js b/jstests/change_streams/split_large_event_with_other_stages.js
new file mode 100644
index 00000000000..ec0713f5be6
--- /dev/null
+++ b/jstests/change_streams/split_large_event_with_other_stages.js
@@ -0,0 +1,51 @@
+/**
+ * Tests if the $changeStreamSplitLargeEvent stage can be used together with other stages allowed in
+ * change stream pipeline.
+ */
+(function() {
+"use strict";
+
+const testDB = db.getSiblingDB("test");
+const testColl = testDB[jsTestName()];
+
+// Capture cluster time before any events.
+const startClusterTime = testDB.hello().$clusterTime.clusterTime;
+
+// Ensure there is at least one event for the test to run faster.
+testColl.insertOne({first_name: "Maya", last_name: "Ryan", homework: [10, 5, 10], pets: {cats: 2}});
+
+// Define an instance of every stage allowed in a change stream pipeline. The instances must not
+// filter out our only event when added to the change stream pipeline.
+const allowedStages = [
+    {$addFields: {totalHomework: {$sum: "$fullDocument.homework"}}},
+    {$match: {"fullDocument.first_name": "Maya"}},
+    {$project: {"fullDocument.first_name": 1, "fullDocument.last_name": 1}},
+    {
+        $replaceRoot: {
+            newRoot: {
+                full_name: {$concat: ["$fullDocument.first_name", " ", "$fullDocument.last_name"]},
+                _id: "$_id"
+            }
+        }
+    },
+    {
+        $replaceWith: {
+            _id: "$_id",
+            pets: {$mergeObjects: [{dogs: 0, cats: 0, birds: 0, fish: 0}, "$fullDocument.pets"]}
+        }
+    },
+    {$redact: {$cond: {if: {$eq: ["$level", 2]}, then: "$$PRUNE", else: "$$DESCEND"}}},
+    {$set: {totalHomework: {$sum: "$fullDocument.homework"}}},
+    {$unset: ["fullDocument.first_name", "fullDocument.last_name"]},
+];
+
+for (const stage of allowedStages) {
+    const changeStreamPipeline = [stage, {$changeStreamSplitLargeEvent: {}}];
+    const changeStreamCursor =
+        testColl.watch(changeStreamPipeline, {startAtOperationTime: startClusterTime});
+    assert.soon(
+        () => changeStreamCursor.hasNext(),
+        "Unexpected lack of events for the change stream pipeline " + tojson(changeStreamPipeline));
+    changeStreamCursor.close();
+}
+})();
diff --git a/src/mongo/db/pipeline/document_source.cpp b/src/mongo/db/pipeline/document_source.cpp
index ba245171441..93169ac861f 100644
--- a/src/mongo/db/pipeline/document_source.cpp
+++ b/src/mongo/db/pipeline/document_source.cpp
@@ -40,6 +40,7 @@
 #include "mongo/db/pipeline/document_source_internal_shard_filter.h"
 #include "mongo/db/pipeline/document_source_match.h"
 #include "mongo/db/pipeline/document_source_project.h"
+#include "mongo/db/pipeline/document_source_redact.h"
 #include "mongo/db/pipeline/document_source_replace_root.h"
 #include "mongo/db/pipeline/document_source_sample.h"
 #include "mongo/db/pipeline/document_source_sequential_document_cache.h"
@@ -212,6 +213,25 @@ bool DocumentSource::pushMatchBefore(Pipeline::SourceContainer::iterator itr,
     return false;
 }
 
+bool DocumentSource::pushRedactBefore(Pipeline::SourceContainer::iterator itr,
+                                      Pipeline::SourceContainer* container) {
+    if (constraints().canSwapWithRedact) {
+        auto nextItr = std::next(itr);
+        if (auto redactStage = dynamic_cast<DocumentSourceRedact*>(nextItr->get())) {
+            LOGV2_DEBUG(8584800,
+                        5,
+                        "Swapping a $redact stage in front of another stage: ",
+                        "redactStage"_attr = redact(redactStage->serializeToBSONForDebug()),
+                        "thisStage"_attr = redact(serializeToBSONForDebug()));
+
+            // Swap 'itr' and 'nextItr' list nodes.
+            container->splice(itr, *container, nextItr);
+            return true;
+        }
+    }
+    return false;
+}
+
 bool DocumentSource::pushSampleBefore(Pipeline::SourceContainer::iterator itr,
                                       Pipeline::SourceContainer* container) {
     auto nextSample = dynamic_cast<DocumentSourceSample*>((*std::next(itr)).get());
diff --git a/src/mongo/db/pipeline/document_source.h b/src/mongo/db/pipeline/document_source.h
index 5c2f65c1951..2b383a43e8b 100644
--- a/src/mongo/db/pipeline/document_source.h
+++ b/src/mongo/db/pipeline/document_source.h
@@ -531,6 +531,13 @@ private:
                          Pipeline::SourceContainer* container);
 
     /**
+     * Attempts to push a $redact stage directly ahead of the stage present at the 'itr' position if
+     * matches the constraints. Returns true if optimization was performed, false otherwise.
+     */
+    bool pushRedactBefore(Pipeline::SourceContainer::iterator itr,
+                          Pipeline::SourceContainer* container);
+
+    /**
      * Attempt to push a sample stage from directly ahead of the current stage given by itr to
      * before the current stage. Returns whether the optimization was performed.
      */
@@ -558,8 +565,8 @@ private:
             return false;
         }
 
-        return pushMatchBefore(itr, container) || pushSampleBefore(itr, container) ||
-            pushSingleDocumentTransformBefore(itr, container);
+        return pushMatchBefore(itr, container) || pushRedactBefore(itr, container) ||
+            pushSampleBefore(itr, container) || pushSingleDocumentTransformBefore(itr, container);
     }
 
 public:
diff --git a/src/mongo/db/pipeline/document_source_change_stream_ensure_resume_token_present.cpp b/src/mongo/db/pipeline/document_source_change_stream_ensure_resume_token_present.cpp
index 6e0ed9d9e60..383141125f7 100644
--- a/src/mongo/db/pipeline/document_source_change_stream_ensure_resume_token_present.cpp
+++ b/src/mongo/db/pipeline/document_source_change_stream_ensure_resume_token_present.cpp
@@ -72,12 +72,14 @@ StageConstraints DocumentSourceChangeStreamEnsureResumeTokenPresent::constraints
                                  UnionRequirement::kNotAllowed,
                                  ChangeStreamRequirement::kChangeStreamStage};
 
-    // The '$match' and 'DocumentSourceSingleDocumentTransformation' stages can swap with this
-    // stage, allowing filtering and reshaping to occur earlier in the pipeline. For sharded cluster
-    // pipelines, swaps can allow $match and 'DocumentSourceSingleDocumentTransformation' stages to
-    // execute on the shards, providing inter-node parallelism and potentially reducing the amount
-    // of data sent form each shard to the mongoS.
+    // The '$match', '$redact', and 'DocumentSourceSingleDocumentTransformation' stages can swap
+    // with this stage, allowing filtering and reshaping to occur earlier in the pipeline. For
+    // sharded cluster pipelines, swaps can allow $match, $redact and
+    // 'DocumentSourceSingleDocumentTransformation' stages to execute on the shards, providing
+    // inter-node parallelism and potentially reducing the amount of data sent form each shard to
+    // the mongoS.
     constraints.canSwapWithMatch = true;
+    constraints.canSwapWithRedact = true;
     constraints.canSwapWithSingleDocTransform = true;
 
     return constraints;
diff --git a/src/mongo/db/pipeline/document_source_change_stream_handle_topology_change.cpp b/src/mongo/db/pipeline/document_source_change_stream_handle_topology_change.cpp
index 28b695c89be..cdf58ba4830 100644
--- a/src/mongo/db/pipeline/document_source_change_stream_handle_topology_change.cpp
+++ b/src/mongo/db/pipeline/document_source_change_stream_handle_topology_change.cpp
@@ -129,10 +129,11 @@ StageConstraints DocumentSourceChangeStreamHandleTopologyChange::constraints(
                                  UnionRequirement::kNotAllowed,
                                  ChangeStreamRequirement::kChangeStreamStage};
 
-    // Can be swapped with the '$match' and 'DocumentSourceSingleDocumentTransformation' stages and
-    // ensures that they get pushed down to the shards, as this stage bisects the change streams
-    // pipeline.
+    // Can be swapped with the '$match', '$redact', and 'DocumentSourceSingleDocumentTransformation'
+    // stages and ensures that they get pushed down to the shards, as this stage bisects the change
+    // streams pipeline.
     constraints.canSwapWithMatch = true;
+    constraints.canSwapWithRedact = true;
     constraints.canSwapWithSingleDocTransform = true;
 
     return constraints;
diff --git a/src/mongo/db/pipeline/pipeline_test.cpp b/src/mongo/db/pipeline/pipeline_test.cpp
index a9259c05313..1eb06fee45e 100644
--- a/src/mongo/db/pipeline/pipeline_test.cpp
+++ b/src/mongo/db/pipeline/pipeline_test.cpp
@@ -42,6 +42,7 @@
 #include "mongo/db/pipeline/document_source_change_stream.h"
 #include "mongo/db/pipeline/document_source_change_stream_add_post_image.h"
 #include "mongo/db/pipeline/document_source_change_stream_add_pre_image.h"
+#include "mongo/db/pipeline/document_source_change_stream_handle_topology_change.h"
 #include "mongo/db/pipeline/document_source_facet.h"
 #include "mongo/db/pipeline/document_source_graph_lookup.h"
 #include "mongo/db/pipeline/document_source_internal_split_pipeline.h"
@@ -50,6 +51,7 @@
 #include "mongo/db/pipeline/document_source_mock.h"
 #include "mongo/db/pipeline/document_source_out.h"
 #include "mongo/db/pipeline/document_source_project.h"
+#include "mongo/db/pipeline/document_source_redact.h"
 #include "mongo/db/pipeline/document_source_sort.h"
 #include "mongo/db/pipeline/document_source_test_optimizations.h"
 #include "mongo/db/pipeline/expression_context_for_test.h"
@@ -3195,6 +3197,36 @@ TEST(PipelineOptimizationTest, FullDocumentBeforeChangeDoesNotSwapWithMatchOnPre
     ASSERT(dynamic_cast<DocumentSourceMatch*>(pipeline->getSources().back().get()));
 }
 
+TEST(PipelineOptimizationTest, ChangeStreamHandleTopologyChangeSwapsWithRedact) {
+    QueryTestServiceContext testServiceContext;
+    auto opCtx = testServiceContext.makeOperationContext();
+
+    boost::intrusive_ptr<ExpressionContext> expCtx(new ExpressionContextForTest(kTestNss));
+    expCtx->opCtx = opCtx.get();
+    expCtx->uuid = UUID::gen();
+    expCtx->inMongos = true;  // To enforce the $_internalChangeStreamHandleTopologyChange stage.
+    setMockReplicationCoordinatorOnOpCtx(expCtx->opCtx);
+
+    auto stages = DocumentSourceChangeStream::createFromBson(
+        fromjson("{$changeStream: {showExpandedEvents: true}}").firstElement(), expCtx);
+
+    // Assert that the last stage is $_internalChangeStreamHandleTopologyChange.
+    ASSERT(dynamic_cast<DocumentSourceChangeStreamHandleTopologyChange*>(stages.back().get()));
+
+    // Add $redact as the last stage.
+    stages.push_back(DocumentSourceRedact::createFromBson(
+        fromjson("{$redact: '$$PRUNE'}").firstElement(), expCtx));
+
+    auto pipeline = Pipeline::create(stages, expCtx);
+    pipeline->optimizePipeline();
+
+    // Assert that $redact swaps with $_internalChangeStreamHandleTopologyChange after optimization.
+    ASSERT(dynamic_cast<DocumentSourceRedact*>(
+        std::prev(std::prev(pipeline->getSources().end()))->get()));
+    ASSERT(dynamic_cast<DocumentSourceChangeStreamHandleTopologyChange*>(
+        pipeline->getSources().back().get()));
+}
+
 TEST(PipelineOptimizationTest, SortLimProjLimBecomesTopKSortProj) {
     std::string inputPipe =
         "[{$sort: {a: 1}}"
@@ -4785,6 +4817,55 @@ TEST_F(PipelineValidateTest, ChangeStreamIsNotValidIfNotFirstStageInFacet) {
     ASSERT_THROWS_CODE(Pipeline::parse(rawPipeline, ctx), AssertionException, 40600);
 }
 
+TEST_F(PipelineValidateTest, ChangeStreamSplitLargeEventIsValid) {
+    const std::vector<BSONObj> rawPipeline = {fromjson("{$changeStream: {}}"),
+                                              fromjson("{$changeStreamSplitLargeEvent: {}}")};
+    auto ctx = getExpCtx();
+    setMockReplicationCoordinatorOnOpCtx(ctx->opCtx);
+    ctx->ns = NamespaceString::createNamespaceString_forTest("a.collection");
+    Pipeline::parse(rawPipeline, ctx);
+}
+
+TEST_F(PipelineValidateTest, ChangeStreamSplitLargeEventIsNotValidWithoutChangeStream) {
+    const std::vector<BSONObj> rawPipeline = {fromjson("{$changeStreamSplitLargeEvent: {}}")};
+    auto ctx = getExpCtx();
+    ctx->changeStreamSpec = boost::none;
+    setMockReplicationCoordinatorOnOpCtx(ctx->opCtx);
+    ctx->ns = NamespaceString::createNamespaceString_forTest("a.collection");
+    ASSERT_THROWS_CODE(
+        Pipeline::parse(rawPipeline, ctx), DBException, ErrorCodes::IllegalOperation);
+}
+
+TEST_F(PipelineValidateTest, ChangeStreamSplitLargeEventIsNotLastStage) {
+    const std::vector<BSONObj> rawPipeline = {fromjson("{$changeStream: {}}"),
+                                              fromjson("{$changeStreamSplitLargeEvent: {}}"),
+                                              fromjson("{$match: {}}")};
+    auto ctx = getExpCtx();
+    setMockReplicationCoordinatorOnOpCtx(ctx->opCtx);
+    ctx->ns = NamespaceString::createNamespaceString_forTest("a.collection");
+    ASSERT_THROWS_CODE(Pipeline::parse(rawPipeline, ctx), DBException, 7182802);
+}
+
+TEST_F(PipelineValidateTest, ChangeStreamSplitLargeEventIsValidAfterMatch) {
+    const std::vector<BSONObj> rawPipeline = {fromjson("{$changeStream: {}}"),
+                                              fromjson("{$match: {custom: 'filter'}}"),
+                                              fromjson("{$changeStreamSplitLargeEvent: {}}")};
+    auto ctx = getExpCtx();
+    setMockReplicationCoordinatorOnOpCtx(ctx->opCtx);
+    ctx->ns = NamespaceString::createNamespaceString_forTest("a.collection");
+    Pipeline::parse(rawPipeline, ctx);
+}
+
+TEST_F(PipelineValidateTest, ChangeStreamSplitLargeEventIsValidAfterRedact) {
+    const std::vector<BSONObj> rawPipeline = {fromjson("{$changeStream: {}}"),
+                                              fromjson("{$redact: '$$PRUNE'}"),
+                                              fromjson("{$changeStreamSplitLargeEvent: {}}")};
+    auto ctx = getExpCtx();
+    setMockReplicationCoordinatorOnOpCtx(ctx->opCtx);
+    ctx->ns = NamespaceString::createNamespaceString_forTest("a.collection");
+    Pipeline::parse(rawPipeline, ctx);
+}
+
 class DocumentSourceDisallowedInTransactions : public DocumentSourceMock {
 public:
     DocumentSourceDisallowedInTransactions(const boost::intrusive_ptr<ExpressionContext>& expCtx)
diff --git a/src/mongo/db/pipeline/stage_constraints.h b/src/mongo/db/pipeline/stage_constraints.h
index 86ddac13ed0..78050420ef0 100644
--- a/src/mongo/db/pipeline/stage_constraints.h
+++ b/src/mongo/db/pipeline/stage_constraints.h
@@ -346,6 +346,9 @@ struct StageConstraints {
     // $match predicates be swapped before itself.
     bool canSwapWithMatch = false;
 
+    // True if this stage can be safely swapped with a subsequent $redact stage.
+    bool canSwapWithRedact = false;
+
     // True if this stage can be safely swapped with a stage which alters the number of documents in
     // the stream.
     //
@@ -391,6 +394,7 @@ struct StageConstraints {
             canSwapWithMatch == other.canSwapWithMatch &&
             canSwapWithSkippingOrLimitingStage == other.canSwapWithSkippingOrLimitingStage &&
             canSwapWithSingleDocTransform == other.canSwapWithSingleDocTransform &&
+            canSwapWithRedact == other.canSwapWithRedact &&
             canAppearOnlyOnceInPipeline == other.canAppearOnlyOnceInPipeline &&
             isAllowedWithinUpdatePipeline == other.isAllowedWithinUpdatePipeline &&
             unionRequirement == other.unionRequirement &&