SERVER-71943 Convert CQF fallback to new DocumentSource visitor mechanism

10 files changed, 365 insertions, 189 deletions

diff --git a/src/mongo/db/SConscript b/src/mongo/db/SConscript
index 35064f4fd95..bcdd849f19a 100644
--- a/src/mongo/db/SConscript
+++ b/src/mongo/db/SConscript
@@ -1527,6 +1527,8 @@ env.Library(
         '$BUILD_DIR/mongo/db/concurrency/exception_util',
         '$BUILD_DIR/mongo/db/exec/sbe/query_sbe_abt',
         '$BUILD_DIR/mongo/db/internal_transactions_feature_flag',
+        '$BUILD_DIR/mongo/db/pipeline/change_stream_pipeline',
+        '$BUILD_DIR/mongo/db/pipeline/document_source_internal_apply_oplog_update',
         '$BUILD_DIR/mongo/db/query/ce/query_ce_heuristic',
         '$BUILD_DIR/mongo/db/query/ce/query_ce_histogram',
         '$BUILD_DIR/mongo/db/query/ce/query_ce_sampling',
diff --git a/src/mongo/db/pipeline/abt/abt_fallback_mechanism_pipeline_bm.cpp b/src/mongo/db/pipeline/abt/abt_fallback_mechanism_pipeline_bm.cpp
index 5ab7d94a123..6de2ee842e0 100644
--- a/src/mongo/db/pipeline/abt/abt_fallback_mechanism_pipeline_bm.cpp
+++ b/src/mongo/db/pipeline/abt/abt_fallback_mechanism_pipeline_bm.cpp
@@ -50,7 +50,7 @@ public:
         auto opCtx = testServiceContext.makeOperationContext();
 
         auto nss = NamespaceString("test.bm");
-        auto expCtx = new ExpressionContextForTest(opCtx.get(), nss);
+        auto expCtx = make_intrusive<ExpressionContextForTest>(opCtx.get(), nss);
 
         std::unique_ptr<Pipeline, PipelineDeleter> parsedPipeline =
             Pipeline::parse(pipeline, expCtx);
@@ -58,7 +58,8 @@ public:
 
         // This is where recording starts.
         for (auto keepRunning : state) {
-            benchmark::DoNotOptimize(isEligibleForBonsai_forTesting(*parsedPipeline.get()));
+            benchmark::DoNotOptimize(isEligibleForBonsai_forTesting(
+                testServiceContext.getServiceContext(), *parsedPipeline.get()));
             benchmark::ClobberMemory();
         }
     }
diff --git a/src/mongo/db/pipeline/abt/abt_translate_pipeline_bm.cpp b/src/mongo/db/pipeline/abt/abt_translate_pipeline_bm.cpp
index f706fa47144..684d676af0d 100644
--- a/src/mongo/db/pipeline/abt/abt_translate_pipeline_bm.cpp
+++ b/src/mongo/db/pipeline/abt/abt_translate_pipeline_bm.cpp
@@ -62,7 +62,8 @@ public:
                                const std::vector<BSONObj>& pipeline) override final {
         QueryTestServiceContext testServiceContext;
         auto opCtx = testServiceContext.makeOperationContext();
-        auto expCtx = new ExpressionContextForTest(opCtx.get(), NamespaceString("test.bm"));
+        auto expCtx =
+            make_intrusive<ExpressionContextForTest>(opCtx.get(), NamespaceString("test.bm"));
 
         Metadata metadata{{}};
         auto prefixId = PrefixId::createForTests();
@@ -72,7 +73,8 @@ public:
             Pipeline::parse(pipeline, expCtx);
         parsedPipeline->optimizePipeline();
 
-        if (!isEligibleForBonsai_forTesting(*parsedPipeline.get())) {
+        if (!isEligibleForBonsai_forTesting(testServiceContext.getServiceContext(),
+                                            *parsedPipeline.get())) {
             state.SkipWithError("Pipeline is not supported by CQF");
             return;
         }
diff --git a/src/mongo/db/pipeline/visitors/document_source_visitor_registry.h b/src/mongo/db/pipeline/visitors/document_source_visitor_registry.h
index 8f8dde086cb..e0eba0b7b8e 100644
--- a/src/mongo/db/pipeline/visitors/document_source_visitor_registry.h
+++ b/src/mongo/db/pipeline/visitors/document_source_visitor_registry.h
@@ -173,6 +173,8 @@ void registerVisitFuncs(DocumentSourceVisitorRegistry* reg) {
     registerVisitFuncs<VisitorCtx, Ds...>(reg);
 }
 
-// TODO SERVER-71944: Introduce global registry.
+// Declare visitor registry as a decoration on the service context.
+inline const auto getDocumentSourceVisitorRegistry =
+    ServiceContext::declareDecoration<DocumentSourceVisitorRegistry>();
 
 }  // namespace mongo
diff --git a/src/mongo/db/pipeline/visitors/document_source_visitor_registry_mongod.h b/src/mongo/db/pipeline/visitors/document_source_visitor_registry_mongod.h
new file mode 100644
index 00000000000..b586feffa5f
--- /dev/null
+++ b/src/mongo/db/pipeline/visitors/document_source_visitor_registry_mongod.h
@@ -0,0 +1,169 @@
+/**
+ *    Copyright (C) 2022-present MongoDB, Inc.
+ *
+ *    This program is free software: you can redistribute it and/or modify
+ *    it under the terms of the Server Side Public License, version 1,
+ *    as published by MongoDB, Inc.
+ *
+ *    This program is distributed in the hope that it will be useful,
+ *    but WITHOUT ANY WARRANTY; without even the implied warranty of
+ *    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ *    Server Side Public License for more details.
+ *
+ *    You should have received a copy of the Server Side Public License
+ *    along with this program. If not, see
+ *    <http://www.mongodb.com/licensing/server-side-public-license>.
+ *
+ *    As a special exception, the copyright holders give permission to link the
+ *    code of portions of this program with the OpenSSL library under certain
+ *    conditions as described in each individual source file and distribute
+ *    linked combinations including the program with the OpenSSL library. You
+ *    must comply with the Server Side Public License in all respects for
+ *    all of the code used other than as permitted herein. If you modify file(s)
+ *    with this exception, you may extend this exception to your version of the
+ *    file(s), but you are not obligated to do so. If you do not wish to do so,
+ *    delete this exception statement from your version. If you delete this
+ *    exception statement from all source files in the program, then also delete
+ *    it in the license file.
+ */
+
+#pragma once
+
+#include "mongo/db/pipeline/document_source_bucket_auto.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_check_invalidate.h"
+#include "mongo/db/pipeline/document_source_change_stream_check_resumability.h"
+#include "mongo/db/pipeline/document_source_change_stream_check_topology_change.h"
+#include "mongo/db/pipeline/document_source_change_stream_handle_topology_change.h"
+#include "mongo/db/pipeline/document_source_change_stream_transform.h"
+#include "mongo/db/pipeline/document_source_change_stream_unwind_transaction.h"
+#include "mongo/db/pipeline/document_source_coll_stats.h"
+#include "mongo/db/pipeline/document_source_current_op.h"
+#include "mongo/db/pipeline/document_source_cursor.h"
+#include "mongo/db/pipeline/document_source_densify.h"
+#include "mongo/db/pipeline/document_source_exchange.h"
+#include "mongo/db/pipeline/document_source_facet.h"
+#include "mongo/db/pipeline/document_source_find_and_modify_image_lookup.h"
+#include "mongo/db/pipeline/document_source_geo_near.h"
+#include "mongo/db/pipeline/document_source_geo_near_cursor.h"
+#include "mongo/db/pipeline/document_source_graph_lookup.h"
+#include "mongo/db/pipeline/document_source_group.h"
+#include "mongo/db/pipeline/document_source_index_stats.h"
+#include "mongo/db/pipeline/document_source_internal_all_collection_stats.h"
+#include "mongo/db/pipeline/document_source_internal_apply_oplog_update.h"
+#include "mongo/db/pipeline/document_source_internal_compute_geo_near_distance.h"
+#include "mongo/db/pipeline/document_source_internal_convert_bucket_index_stats.h"
+#include "mongo/db/pipeline/document_source_internal_inhibit_optimization.h"
+#include "mongo/db/pipeline/document_source_internal_shard_filter.h"
+#include "mongo/db/pipeline/document_source_internal_split_pipeline.h"
+#include "mongo/db/pipeline/document_source_internal_unpack_bucket.h"
+#include "mongo/db/pipeline/document_source_limit.h"
+#include "mongo/db/pipeline/document_source_list_cached_and_active_users.h"
+#include "mongo/db/pipeline/document_source_list_catalog.h"
+#include "mongo/db/pipeline/document_source_list_local_sessions.h"
+#include "mongo/db/pipeline/document_source_list_sessions.h"
+#include "mongo/db/pipeline/document_source_lookup.h"
+#include "mongo/db/pipeline/document_source_match.h"
+#include "mongo/db/pipeline/document_source_merge.h"
+#include "mongo/db/pipeline/document_source_operation_metrics.h"
+#include "mongo/db/pipeline/document_source_out.h"
+#include "mongo/db/pipeline/document_source_plan_cache_stats.h"
+#include "mongo/db/pipeline/document_source_queue.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_sample_from_random_cursor.h"
+#include "mongo/db/pipeline/document_source_sequential_document_cache.h"
+#include "mongo/db/pipeline/document_source_set_variable_from_subpipeline.h"
+#include "mongo/db/pipeline/document_source_set_window_fields.h"
+#include "mongo/db/pipeline/document_source_single_document_transformation.h"
+#include "mongo/db/pipeline/document_source_skip.h"
+#include "mongo/db/pipeline/document_source_sort.h"
+#include "mongo/db/pipeline/document_source_streaming_group.h"
+#include "mongo/db/pipeline/document_source_tee_consumer.h"
+#include "mongo/db/pipeline/document_source_telemetry.h"
+#include "mongo/db/pipeline/document_source_union_with.h"
+#include "mongo/db/pipeline/document_source_unwind.h"
+#include "mongo/db/pipeline/visitors/document_source_visitor_registry.h"
+
+namespace mongo {
+
+/**
+ * Register 'visit()' functions for all mongod DocumentSources for the visitor specified as the
+ * template parameter in the DocumentSource visitor regsitry in the given ServiceContext. Using this
+ * function helps provide compile-time safety that ensures visitor implementors have provided an
+ * implementation for all DocumentSoures. This function is intended to be used in the following
+ * manner:
+ *
+ * // Define visit functions for all DocumentSources
+ * void visit(FooVisitorCtx* ctx, const DocumentSourceMatch& match) { ... }
+ * ...
+ *
+ * const ServiceContext::ConstructorActionRegisterer fooRegisterer{
+ *   "FooRegisterer", [](ServiceContext* service) {
+ *       registerMongodVisitor<FooVisitorCtx>(service);
+ *   }};
+ */
+template <typename T>
+void registerMongodVisitor(ServiceContext* service) {
+    auto& registry = getDocumentSourceVisitorRegistry(service);
+    registerVisitFuncs<T,
+                       DocumentSourceBucketAuto,
+                       DocumentSourceChangeStreamAddPostImage,
+                       DocumentSourceChangeStreamAddPreImage,
+                       DocumentSourceChangeStreamCheckInvalidate,
+                       DocumentSourceChangeStreamCheckResumability,
+                       DocumentSourceChangeStreamCheckTopologyChange,
+                       DocumentSourceChangeStreamHandleTopologyChange,
+                       DocumentSourceChangeStreamTransform,
+                       DocumentSourceChangeStreamUnwindTransaction,
+                       DocumentSourceCollStats,
+                       DocumentSourceCurrentOp,
+                       DocumentSourceCursor,
+                       DocumentSourceExchange,
+                       DocumentSourceFacet,
+                       DocumentSourceFindAndModifyImageLookup,
+                       DocumentSourceGeoNear,
+                       DocumentSourceGeoNearCursor,
+                       DocumentSourceGraphLookUp,
+                       DocumentSourceGroup,
+                       DocumentSourceIndexStats,
+                       DocumentSourceInternalAllCollectionStats,
+                       DocumentSourceInternalApplyOplogUpdate,
+                       DocumentSourceInternalConvertBucketIndexStats,
+                       DocumentSourceInternalDensify,
+                       DocumentSourceInternalGeoNearDistance,
+                       DocumentSourceInternalInhibitOptimization,
+                       DocumentSourceInternalSetWindowFields,
+                       DocumentSourceInternalShardFilter,
+                       DocumentSourceInternalSplitPipeline,
+                       DocumentSourceInternalUnpackBucket,
+                       DocumentSourceLimit,
+                       DocumentSourceListCachedAndActiveUsers,
+                       DocumentSourceListCatalog,
+                       DocumentSourceListLocalSessions,
+                       DocumentSourceListSessions,
+                       DocumentSourceLookUp,
+                       DocumentSourceMatch,
+                       DocumentSourceMerge,
+                       DocumentSourceOperationMetrics,
+                       DocumentSourceOut,
+                       DocumentSourcePlanCacheStats,
+                       DocumentSourceQueue,
+                       DocumentSourceRedact,
+                       DocumentSourceSample,
+                       DocumentSourceSampleFromRandomCursor,
+                       DocumentSourceSequentialDocumentCache,
+                       DocumentSourceSetVariableFromSubPipeline,
+                       DocumentSourceSingleDocumentTransformation,
+                       DocumentSourceSkip,
+                       DocumentSourceSort,
+                       DocumentSourceStreamingGroup,
+                       DocumentSourceTeeConsumer,
+                       DocumentSourceTelemetry,
+                       DocumentSourceUnionWith,
+                       DocumentSourceUnwind>(&registry);
+}
+
+}  // namespace mongo
diff --git a/src/mongo/db/pipeline/visitors/document_source_visitor_registry_mongos.h b/src/mongo/db/pipeline/visitors/document_source_visitor_registry_mongos.h
new file mode 100644
index 00000000000..1993ed5fea7
--- /dev/null
+++ b/src/mongo/db/pipeline/visitors/document_source_visitor_registry_mongos.h
@@ -0,0 +1,57 @@
+/**
+ *    Copyright (C) 2022-present MongoDB, Inc.
+ *
+ *    This program is free software: you can redistribute it and/or modify
+ *    it under the terms of the Server Side Public License, version 1,
+ *    as published by MongoDB, Inc.
+ *
+ *    This program is distributed in the hope that it will be useful,
+ *    but WITHOUT ANY WARRANTY; without even the implied warranty of
+ *    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ *    Server Side Public License for more details.
+ *
+ *    You should have received a copy of the Server Side Public License
+ *    along with this program. If not, see
+ *    <http://www.mongodb.com/licensing/server-side-public-license>.
+ *
+ *    As a special exception, the copyright holders give permission to link the
+ *    code of portions of this program with the OpenSSL library under certain
+ *    conditions as described in each individual source file and distribute
+ *    linked combinations including the program with the OpenSSL library. You
+ *    must comply with the Server Side Public License in all respects for
+ *    all of the code used other than as permitted herein. If you modify file(s)
+ *    with this exception, you may extend this exception to your version of the
+ *    file(s), but you are not obligated to do so. If you do not wish to do so,
+ *    delete this exception statement from your version. If you delete this
+ *    exception statement from all source files in the program, then also delete
+ *    it in the license file.
+ */
+
+#pragma once
+
+#include "mongo/db/pipeline/visitors/document_source_visitor_registry.h"
+#include "mongo/s/query/document_source_merge_cursors.h"
+
+namespace mongo {
+
+/**
+ * Register 'visit()' functions for all mongos DocumentSources for the visitor specified as the
+ * template parameter in the DocumentSource visitor regsitry in the given ServiceContext. Using this
+ * function helps provide compile-time safety that ensures visitor implementors have provided an
+ * implementation for all DocumentSoures. This function is intended to be used in the following
+ * manner:
+ *
+ * void visit(FooVisitorCtx* ctx, const DocumentSourceMergeCursors& match) { ... }
+ *
+ * const ServiceContext::ConstructorActionRegisterer fooRegisterer{
+ *   "FooRegisterer", [](ServiceContext* service) {
+ *       registerMongosVisitor<FooVisitorCtx>(service);
+ *   }};
+ */
+template <typename T>
+void registerMongosVisitor(ServiceContext* service) {
+    auto& registry = getDocumentSourceVisitorRegistry(service);
+    registerVisitFuncs<T, DocumentSourceMergeCursors>(&registry);
+}
+
+}  // namespace mongo
diff --git a/src/mongo/db/query/cqf_command_utils.cpp b/src/mongo/db/query/cqf_command_utils.cpp
index 1ca70694a17..133cf8276cc 100644
--- a/src/mongo/db/query/cqf_command_utils.cpp
+++ b/src/mongo/db/query/cqf_command_utils.cpp
@@ -71,22 +71,37 @@
 #include "mongo/db/pipeline/abt/match_expression_visitor.h"
 #include "mongo/db/pipeline/abt/utils.h"
 #include "mongo/db/pipeline/document_source_bucket_auto.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_check_invalidate.h"
+#include "mongo/db/pipeline/document_source_change_stream_check_resumability.h"
+#include "mongo/db/pipeline/document_source_change_stream_check_topology_change.h"
+#include "mongo/db/pipeline/document_source_change_stream_handle_topology_change.h"
+#include "mongo/db/pipeline/document_source_change_stream_transform.h"
+#include "mongo/db/pipeline/document_source_change_stream_unwind_transaction.h"
 #include "mongo/db/pipeline/document_source_coll_stats.h"
 #include "mongo/db/pipeline/document_source_current_op.h"
 #include "mongo/db/pipeline/document_source_cursor.h"
+#include "mongo/db/pipeline/document_source_densify.h"
 #include "mongo/db/pipeline/document_source_exchange.h"
 #include "mongo/db/pipeline/document_source_facet.h"
+#include "mongo/db/pipeline/document_source_find_and_modify_image_lookup.h"
 #include "mongo/db/pipeline/document_source_geo_near.h"
 #include "mongo/db/pipeline/document_source_geo_near_cursor.h"
 #include "mongo/db/pipeline/document_source_graph_lookup.h"
 #include "mongo/db/pipeline/document_source_group.h"
 #include "mongo/db/pipeline/document_source_index_stats.h"
+#include "mongo/db/pipeline/document_source_internal_all_collection_stats.h"
+#include "mongo/db/pipeline/document_source_internal_apply_oplog_update.h"
+#include "mongo/db/pipeline/document_source_internal_compute_geo_near_distance.h"
+#include "mongo/db/pipeline/document_source_internal_convert_bucket_index_stats.h"
 #include "mongo/db/pipeline/document_source_internal_inhibit_optimization.h"
 #include "mongo/db/pipeline/document_source_internal_shard_filter.h"
 #include "mongo/db/pipeline/document_source_internal_split_pipeline.h"
 #include "mongo/db/pipeline/document_source_internal_unpack_bucket.h"
 #include "mongo/db/pipeline/document_source_limit.h"
 #include "mongo/db/pipeline/document_source_list_cached_and_active_users.h"
+#include "mongo/db/pipeline/document_source_list_catalog.h"
 #include "mongo/db/pipeline/document_source_list_local_sessions.h"
 #include "mongo/db/pipeline/document_source_list_sessions.h"
 #include "mongo/db/pipeline/document_source_lookup.h"
@@ -101,14 +116,18 @@
 #include "mongo/db/pipeline/document_source_sample.h"
 #include "mongo/db/pipeline/document_source_sample_from_random_cursor.h"
 #include "mongo/db/pipeline/document_source_sequential_document_cache.h"
+#include "mongo/db/pipeline/document_source_set_variable_from_subpipeline.h"
+#include "mongo/db/pipeline/document_source_set_window_fields.h"
 #include "mongo/db/pipeline/document_source_single_document_transformation.h"
 #include "mongo/db/pipeline/document_source_skip.h"
 #include "mongo/db/pipeline/document_source_sort.h"
+#include "mongo/db/pipeline/document_source_streaming_group.h"
 #include "mongo/db/pipeline/document_source_tee_consumer.h"
 #include "mongo/db/pipeline/document_source_telemetry.h"
 #include "mongo/db/pipeline/document_source_union_with.h"
 #include "mongo/db/pipeline/document_source_unwind.h"
 #include "mongo/db/pipeline/visitors/document_source_visitor.h"
+#include "mongo/db/pipeline/visitors/document_source_visitor_registry_mongod.h"
 #include "mongo/db/pipeline/visitors/document_source_walker.h"
 #include "mongo/db/pipeline/visitors/transformer_interface_walker.h"
 #include "mongo/db/query/expression_walker.h"
@@ -1012,172 +1031,6 @@ private:
     bool& _eligible;
 };
 
-/**
- * Visitor that is responsible for indicating whether a DocumentSource is eligible for Bonsai by
- * setting the 'eligible' member variable. Stages which are "test-only" and not officially supported
- * should set 'eligible' to false.
- */
-class ABTUnsupportedDocumentSourceVisitor : public DocumentSourceConstVisitor {
-public:
-    void visit(const DocumentSourceInternalUnpackBucket* source) override {
-        unsupportedStage(source);
-    }
-
-    void visit(const DocumentSourceBucketAuto* source) override {
-        unsupportedStage(source);
-    }
-
-    void visit(const DocumentSourceCollStats* source) override {
-        unsupportedStage(source);
-    }
-
-    void visit(const DocumentSourceCurrentOp* source) override {
-        unsupportedStage(source);
-    }
-
-    void visit(const DocumentSourceCursor* source) override {
-        unsupportedStage(source);
-    }
-
-    void visit(const DocumentSourceExchange* source) override {
-        unsupportedStage(source);
-    }
-
-    void visit(const DocumentSourceFacet* source) override {
-        unsupportedStage(source);
-    }
-
-    void visit(const DocumentSourceGeoNear* source) override {
-        unsupportedStage(source);
-    }
-
-    void visit(const DocumentSourceGeoNearCursor* source) override {
-        unsupportedStage(source);
-    }
-
-    void visit(const DocumentSourceGraphLookUp* source) override {
-        unsupportedStage(source);
-    }
-
-    void visit(const DocumentSourceIndexStats* source) override {
-        unsupportedStage(source);
-    }
-
-    void visit(const DocumentSourceInternalShardFilter* source) override {
-        unsupportedStage(source);
-    }
-
-    void visit(const DocumentSourceInternalSplitPipeline* source) override {
-        unsupportedStage(source);
-    }
-
-    void visit(const DocumentSourceListCachedAndActiveUsers* source) override {
-        unsupportedStage(source);
-    }
-
-    void visit(const DocumentSourceListLocalSessions* source) override {
-        unsupportedStage(source);
-    }
-
-    void visit(const DocumentSourceListSessions* source) override {
-        unsupportedStage(source);
-    }
-
-    void visit(const DocumentSourceLookUp* source) override {
-        unsupportedStage(source);
-    }
-
-    void visit(const DocumentSourceMerge* source) override {
-        unsupportedStage(source);
-    }
-
-    void visit(const DocumentSourceMergeCursors* source) override {
-        unsupportedStage(source);
-    }
-
-    void visit(const DocumentSourceOperationMetrics* source) override {
-        unsupportedStage(source);
-    }
-
-    void visit(const DocumentSourceOut* source) override {
-        unsupportedStage(source);
-    }
-
-    void visit(const DocumentSourcePlanCacheStats* source) override {
-        unsupportedStage(source);
-    }
-
-    void visit(const DocumentSourceQueue* source) override {
-        unsupportedStage(source);
-    }
-
-    void visit(const DocumentSourceRedact* source) override {
-        unsupportedStage(source);
-    }
-
-    void visit(const DocumentSourceSample* source) override {
-        unsupportedStage(source);
-    }
-
-    void visit(const DocumentSourceSampleFromRandomCursor* source) override {
-        unsupportedStage(source);
-    }
-
-    void visit(const DocumentSourceSequentialDocumentCache* source) override {
-        unsupportedStage(source);
-    }
-
-    void visit(const DocumentSourceTeeConsumer* source) override {
-        unsupportedStage(source);
-    }
-
-    void visit(const DocumentSourceTelemetry* source) override {
-        unsupportedStage(source);
-    }
-
-    void visit(const DocumentSourceGroup* source) override {
-        unsupportedStage(source);
-    }
-    void visit(const DocumentSourceLimit* source) override {
-        unsupportedStage(source);
-    }
-    void visit(const DocumentSourceSkip* source) override {
-        unsupportedStage(source);
-    }
-    void visit(const DocumentSourceSort* source) override {
-        unsupportedStage(source);
-    }
-    void visit(const DocumentSourceUnwind* source) override {
-        unsupportedStage(source);
-    }
-    void visit(const DocumentSourceUnionWith* source) override {
-        unsupportedStage(source);
-    }
-
-    void visit(const DocumentSourceInternalInhibitOptimization* source) override {
-        // Can be ignored.
-    }
-
-    void visit(const DocumentSourceMatch* source) override {
-        // Pass a reference to our local 'eligible' variable to allow the visitor to overwrite it.
-        ABTMatchExpressionVisitor visitor(eligible);
-        MatchExpressionWalker walker(nullptr /*preVisitor*/, nullptr /*inVisitor*/, &visitor);
-        tree_walker::walk<true, MatchExpression>(source->getMatchExpression(), &walker);
-    }
-
-    void visit(const DocumentSourceSingleDocumentTransformation* source) override {
-        ABTTransformerVisitor visitor(eligible);
-        TransformerInterfaceWalker walker(&visitor);
-        walker.walk(&source->getTransformer());
-    }
-
-    void unsupportedStage(const DocumentSource* source) {
-        eligible = false;
-    }
-
-    bool eligible = true;
-};
-
 template <class RequestType>
 bool isEligibleCommon(const RequestType& request,
                       OperationContext* opCtx,
@@ -1263,20 +1116,12 @@ boost::optional<bool> shouldForceEligibility() {
     MONGO_UNREACHABLE;
 }
 
-bool isEligibleForBonsai(const Pipeline& pipeline) {
-    ABTUnsupportedDocumentSourceVisitor visitor;
-    DocumentSourceWalkerLegacy walker(nullptr /*preVisitor*/, &visitor);
-
-    // The rudimentary walker may throw if it reaches a stage that it isn't aware about, so catch it
-    // here and return ineligible.
-    // TODO SERVER-62027 this should no longer be needed once all stages require a visit.
-    try {
-        walker.walk(pipeline);
-    } catch (DBException&) {
-        visitor.eligible = false;
-    }
-
-    return visitor.eligible;
+bool isEligibleForBonsai(ServiceContext* serviceCtx, const Pipeline& pipeline) {
+    ABTUnsupportedDocumentSourceVisitorContext visitorCtx;
+    auto& reg = getDocumentSourceVisitorRegistry(serviceCtx);
+    DocumentSourceWalker walker(reg, &visitorCtx);
+    walker.walk(pipeline);
+    return visitorCtx.eligible;
 }
 
 bool isEligibleForBonsai(const CanonicalQuery& cq) {
@@ -1325,7 +1170,7 @@ bool isEligibleForBonsai(const AggregateCommandRequest& request,
         return false;
     }
 
-    return isEligibleForBonsai(pipeline);
+    return isEligibleForBonsai(opCtx->getServiceContext(), pipeline);
 }
 
 bool isEligibleForBonsai(const CanonicalQuery& cq,
@@ -1362,8 +1207,34 @@ bool isEligibleForBonsai_forTesting(const CanonicalQuery& cq) {
     return isEligibleForBonsai(cq);
 }
 
-bool isEligibleForBonsai_forTesting(const Pipeline& pipeline) {
-    return isEligibleForBonsai(pipeline);
+bool isEligibleForBonsai_forTesting(ServiceContext* serviceCtx, const Pipeline& pipeline) {
+    return isEligibleForBonsai(serviceCtx, pipeline);
 }
 
 }  // namespace mongo
+
+namespace mongo::optimizer {
+// Templated visit function to mark DocumentSources as ineligible for CQF.
+template <typename T>
+void visit(ABTUnsupportedDocumentSourceVisitorContext* ctx, const T&) {
+    ctx->eligible = false;
+}
+
+void visit(ABTUnsupportedDocumentSourceVisitorContext* ctx, const DocumentSourceMatch& source) {
+    ABTMatchExpressionVisitor visitor(ctx->eligible);
+    MatchExpressionWalker walker(nullptr, nullptr, &visitor);
+    tree_walker::walk<true, MatchExpression>(source.getMatchExpression(), &walker);
+}
+
+void visit(ABTUnsupportedDocumentSourceVisitorContext* ctx,
+           const DocumentSourceSingleDocumentTransformation& source) {
+    ABTTransformerVisitor visitor(ctx->eligible);
+    TransformerInterfaceWalker walker(&visitor);
+    walker.walk(&source.getTransformer());
+}
+
+const ServiceContext::ConstructorActionRegisterer abtUnsupportedRegisterer{
+    "ABTUnsupportedRegisterer", [](ServiceContext* service) {
+        registerMongodVisitor<ABTUnsupportedDocumentSourceVisitorContext>(service);
+    }};
+}  // namespace mongo::optimizer
diff --git a/src/mongo/db/query/cqf_command_utils.h b/src/mongo/db/query/cqf_command_utils.h
index 0999fcb234d..c6948c9a18e 100644
--- a/src/mongo/db/query/cqf_command_utils.h
+++ b/src/mongo/db/query/cqf_command_utils.h
@@ -30,12 +30,26 @@
 #pragma once
 
 #include "mongo/db/catalog/collection.h"
+#include "mongo/db/pipeline/document_source_internal_inhibit_optimization.h"
+#include "mongo/db/pipeline/visitors/document_source_visitor_registry.h"
 #include "mongo/db/query/query_knobs_gen.h"
 
+
 #define MONGO_LOGV2_DEFAULT_COMPONENT ::mongo::logv2::LogComponent::kQuery
 
 namespace mongo {
 
+namespace optimizer {
+/**
+ * Visitor that is responsible for indicating whether a DocumentSource is eligible for Bonsai by
+ * setting the 'eligible' member variable. Stages which are "test-only" and not officially supported
+ * should set 'eligible' to false.
+ */
+struct ABTUnsupportedDocumentSourceVisitorContext : public DocumentSourceVisitorContextBase {
+    bool eligible{true};
+};
+}  // namespace optimizer
+
 template <typename T>
 void coutPrintAttr(const logv2::detail::NamedArg<T>& arg) {
     std::cout << arg.name << " : " << arg.value << "\n";
@@ -82,7 +96,7 @@ bool isEligibleForBonsai(const CanonicalQuery& cq,
  * options for further details.
  */
 bool isEligibleForBonsai_forTesting(const CanonicalQuery& cq);
-bool isEligibleForBonsai_forTesting(const Pipeline& pipeline);
+bool isEligibleForBonsai_forTesting(ServiceContext* serviceCtx, const Pipeline& pipeline);
 
 }  // namespace mongo
 #undef MONGO_LOGV2_DEFAULT_COMPONENT
diff --git a/src/mongo/s/query/SConscript b/src/mongo/s/query/SConscript
index 1215f24af74..1bb92050988 100644
--- a/src/mongo/s/query/SConscript
+++ b/src/mongo/s/query/SConscript
@@ -45,6 +45,7 @@ env.Library(
     LIBDEPS_PRIVATE=[
         '$BUILD_DIR/mongo/db/catalog/collection_uuid_mismatch_info',
         '$BUILD_DIR/mongo/db/timeseries/timeseries_options',
+        'cqf_utils',
     ],
 )
 
@@ -76,6 +77,17 @@ env.Library(
 )
 
 env.Library(
+    target='cqf_utils',
+    source=[
+        'cqf_utils.cpp',
+    ],
+    LIBDEPS_PRIVATE=[
+        '$BUILD_DIR/mongo/db/pipeline/pipeline',
+        'router_exec_stage',
+    ],
+)
+
+env.Library(
     target='async_results_merger',
     source=[
         'async_results_merger.cpp',
diff --git a/src/mongo/s/query/cqf_utils.cpp b/src/mongo/s/query/cqf_utils.cpp
new file mode 100644
index 00000000000..6ccd8b7273e
--- /dev/null
+++ b/src/mongo/s/query/cqf_utils.cpp
@@ -0,0 +1,46 @@
+/**
+ *    Copyright (C) 2022-present MongoDB, Inc.
+ *
+ *    This program is free software: you can redistribute it and/or modify
+ *    it under the terms of the Server Side Public License, version 1,
+ *    as published by MongoDB, Inc.
+ *
+ *    This program is distributed in the hope that it will be useful,
+ *    but WITHOUT ANY WARRANTY; without even the implied warranty of
+ *    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ *    Server Side Public License for more details.
+ *
+ *    You should have received a copy of the Server Side Public License
+ *    along with this program. If not, see
+ *    <http://www.mongodb.com/licensing/server-side-public-license>.
+ *
+ *    As a special exception, the copyright holders give permission to link the
+ *    code of portions of this program with the OpenSSL library under certain
+ *    conditions as described in each individual source file and distribute
+ *    linked combinations including the program with the OpenSSL library. You
+ *    must comply with the Server Side Public License in all respects for
+ *    all of the code used other than as permitted herein. If you modify file(s)
+ *    with this exception, you may extend this exception to your version of the
+ *    file(s), but you are not obligated to do so. If you do not wish to do so,
+ *    delete this exception statement from your version. If you delete this
+ *    exception statement from all source files in the program, then also delete
+ *    it in the license file.
+ */
+
+#include "mongo/db/pipeline/visitors/document_source_visitor_registry_mongos.h"
+#include "mongo/db/query/cqf_command_utils.h"
+#include "mongo/db/service_context.h"
+#include "mongo/s/query/document_source_merge_cursors.h"
+
+namespace mongo::optimizer {
+
+void visit(ABTUnsupportedDocumentSourceVisitorContext* ctx, const DocumentSourceMergeCursors&) {
+    ctx->eligible = false;
+}
+
+const ServiceContext::ConstructorActionRegisterer abtUnsupportedRegisterer{
+    "ABTUnsupportedRegistererMongoS", [](ServiceContext* service) {
+        registerMongosVisitor<ABTUnsupportedDocumentSourceVisitorContext>(service);
+    }};
+
+}  // namespace mongo::optimizer