SERVER-27253 Bump index usage stats during $lookup.

8 files changed, 86 insertions, 18 deletions

diff --git a/jstests/core/index_stats.js b/jstests/core/index_stats.js
index 16d5a16d8d2..60b37fd571e 100644
--- a/jstests/core/index_stats.js
+++ b/jstests/core/index_stats.js
@@ -7,8 +7,9 @@
     var col = db[colName];
     col.drop();
 
-    var getUsageCount = function(indexName) {
-        var cursor = col.aggregate([{$indexStats: {}}]);
+    var getUsageCount = function(indexName, collection) {
+        collection = collection || col;
+        var cursor = collection.aggregate([{$indexStats: {}}]);
         while (cursor.hasNext()) {
             var doc = cursor.next();
 
@@ -213,4 +214,64 @@
     assert.throws(function() {
         col.aggregate([{$match: {}}, {$indexStats: {}}]);
     });
+
+    //
+    // Confirm index use is recorded for $lookup.
+    //
+    const foreignCollection = db[colName + "_foreign"];
+    foreignCollection.drop();
+    assert.writeOK(foreignCollection.insert([{_id: 0}, {_id: 1}, {_id: 2}]));
+    col.drop();
+    assert.writeOK(col.insert([{_id: 0, foreignId: 1}, {_id: 1, foreignId: 2}]));
+    assert.eq(0, getUsageCount("_id_"));
+    assert.eq(2,
+              col.aggregate([
+                     {$match: {_id: {$in: [0, 1]}}},
+                     {
+                       $lookup: {
+                           from: foreignCollection.getName(),
+                           localField: 'foreignId',
+                           foreignField: '_id',
+                           as: 'results'
+                       }
+                     }
+                 ])
+                  .itcount());
+    assert.eq(1, getUsageCount("_id_", col), "Expected aggregation to use _id index");
+    assert.eq(2,
+              getUsageCount("_id_", foreignCollection),
+              "Expected each lookup to be tracked as an index use");
+
+    //
+    // Confirm index use is recorded for $graphLookup.
+    //
+    foreignCollection.drop();
+    assert.writeOK(foreignCollection.insert([
+        {_id: 0, connectedTo: 1},
+        {_id: 1, connectedTo: "X"},
+        {_id: 2, connectedTo: 3},
+        {_id: 3, connectedTo: "Y"},  // Be sure to use a different value here to make sure
+                                     // $graphLookup doesn't cache the query.
+    ]));
+    col.drop();
+    assert.writeOK(col.insert([{_id: 0, foreignId: 0}, {_id: 1, foreignId: 2}]));
+    assert.eq(0, getUsageCount("_id_"));
+    assert.eq(2,
+              col.aggregate([
+                     {$match: {_id: {$in: [0, 1]}}},
+                     {
+                       $graphLookup: {
+                           from: foreignCollection.getName(),
+                           startWith: '$foreignId',
+                           connectToField: '_id',
+                           connectFromField: 'connectedTo',
+                           as: 'results'
+                       }
+                     }
+                 ])
+                  .itcount());
+    assert.eq(1, getUsageCount("_id_", col), "Expected aggregation to use _id index");
+    assert.eq(2 * 3,
+              getUsageCount("_id_", foreignCollection),
+              "Expected each of two graph searches to issue 3 queries, each using the _id index");
 })();
diff --git a/src/mongo/db/commands/pipeline_command.cpp b/src/mongo/db/commands/pipeline_command.cpp
index bb426950b44..4887f61d49b 100644
--- a/src/mongo/db/commands/pipeline_command.cpp
+++ b/src/mongo/db/commands/pipeline_command.cpp
@@ -506,12 +506,6 @@ public:
             }
 
             if (collection) {
-                PlanSummaryStats stats;
-                Explain::getSummaryStats(*exec, &stats);
-                collection->infoCache()->notifyOfQuery(txn, stats.indexesUsed);
-            }
-
-            if (collection) {
                 const bool isAggCursor = true;  // enable special locking behavior
                 pin.emplace(collection->getCursorManager()->registerCursor(
                     {exec.release(),
diff --git a/src/mongo/db/pipeline/document_source_cursor.cpp b/src/mongo/db/pipeline/document_source_cursor.cpp
index 2def98d25bd..61fd276c2ec 100644
--- a/src/mongo/db/pipeline/document_source_cursor.cpp
+++ b/src/mongo/db/pipeline/document_source_cursor.cpp
@@ -30,7 +30,7 @@
 
 #include "mongo/db/pipeline/document_source_cursor.h"
 
-#include "mongo/db/catalog/database_holder.h"
+#include "mongo/db/catalog/collection.h"
 #include "mongo/db/db_raii.h"
 #include "mongo/db/exec/working_set_common.h"
 #include "mongo/db/pipeline/document.h"
@@ -233,7 +233,8 @@ void DocumentSourceCursor::reattachToOperationContext(OperationContext* opCtx) {
     }
 }
 
-DocumentSourceCursor::DocumentSourceCursor(const string& ns,
+DocumentSourceCursor::DocumentSourceCursor(Collection* collection,
+                                           const string& ns,
                                            std::unique_ptr<PlanExecutor> exec,
                                            const intrusive_ptr<ExpressionContext>& pCtx)
     : DocumentSource(pCtx),
@@ -245,14 +246,18 @@ DocumentSourceCursor::DocumentSourceCursor(const string& ns,
 
     // We record execution metrics here to allow for capture of indexes used prior to execution.
     recordPlanSummaryStats();
+    if (collection) {
+        collection->infoCache()->notifyOfQuery(pCtx->opCtx, _planSummaryStats.indexesUsed);
+    }
 }
 
 intrusive_ptr<DocumentSourceCursor> DocumentSourceCursor::create(
+    Collection* collection,
     const string& ns,
     std::unique_ptr<PlanExecutor> exec,
     const intrusive_ptr<ExpressionContext>& pExpCtx) {
     intrusive_ptr<DocumentSourceCursor> source(
-        new DocumentSourceCursor(ns, std::move(exec), pExpCtx));
+        new DocumentSourceCursor(collection, ns, std::move(exec), pExpCtx));
     return source;
 }
 
diff --git a/src/mongo/db/pipeline/document_source_cursor.h b/src/mongo/db/pipeline/document_source_cursor.h
index 799834f1487..360b855ad39 100644
--- a/src/mongo/db/pipeline/document_source_cursor.h
+++ b/src/mongo/db/pipeline/document_source_cursor.h
@@ -75,6 +75,7 @@ public:
      * in order to fetch data from the database.
      */
     static boost::intrusive_ptr<DocumentSourceCursor> create(
+        Collection* collection,
         const std::string& ns,
         std::unique_ptr<PlanExecutor> exec,
         const boost::intrusive_ptr<ExpressionContext>& pExpCtx);
@@ -133,7 +134,8 @@ public:
     const PlanSummaryStats& getPlanSummaryStats() const;
 
 private:
-    DocumentSourceCursor(const std::string& ns,
+    DocumentSourceCursor(Collection* collection,
+                         const std::string& ns,
                          std::unique_ptr<PlanExecutor> exec,
                          const boost::intrusive_ptr<ExpressionContext>& pExpCtx);
 
diff --git a/src/mongo/db/pipeline/pipeline_d.cpp b/src/mongo/db/pipeline/pipeline_d.cpp
index b2a7bee2fa8..57ede6d6955 100644
--- a/src/mongo/db/pipeline/pipeline_d.cpp
+++ b/src/mongo/db/pipeline/pipeline_d.cpp
@@ -360,6 +360,7 @@ void PipelineD::prepareCursorSource(Collection* collection,
                     expCtx, sampleSize, idString, numRecords));
 
                 addCursorSource(
+                    collection,
                     pipeline,
                     expCtx,
                     std::move(exec),
@@ -420,7 +421,8 @@ void PipelineD::prepareCursorSource(Collection* collection,
                                                 &sortObj,
                                                 &projForQuery));
 
-    addCursorSource(pipeline, expCtx, std::move(exec), deps, queryObj, sortObj, projForQuery);
+    addCursorSource(
+        collection, pipeline, expCtx, std::move(exec), deps, queryObj, sortObj, projForQuery);
 }
 
 StatusWith<std::unique_ptr<PlanExecutor>> PipelineD::prepareExecutor(
@@ -542,7 +544,8 @@ StatusWith<std::unique_ptr<PlanExecutor>> PipelineD::prepareExecutor(
         txn, collection, expCtx, queryObj, *projectionObj, *sortObj, plannerOpts);
 }
 
-void PipelineD::addCursorSource(const intrusive_ptr<Pipeline>& pipeline,
+void PipelineD::addCursorSource(Collection* collection,
+                                const intrusive_ptr<Pipeline>& pipeline,
                                 const intrusive_ptr<ExpressionContext>& expCtx,
                                 unique_ptr<PlanExecutor> exec,
                                 DepsTracker deps,
@@ -557,7 +560,7 @@ void PipelineD::addCursorSource(const intrusive_ptr<Pipeline>& pipeline,
 
     // Put the PlanExecutor into a DocumentSourceCursor and add it to the front of the pipeline.
     intrusive_ptr<DocumentSourceCursor> pSource =
-        DocumentSourceCursor::create(fullName, std::move(exec), expCtx);
+        DocumentSourceCursor::create(collection, fullName, std::move(exec), expCtx);
 
     // Note the query, sort, and projection for explain.
     pSource->setQuery(queryObj);
diff --git a/src/mongo/db/pipeline/pipeline_d.h b/src/mongo/db/pipeline/pipeline_d.h
index 5be4893cf3b..20b6cf16965 100644
--- a/src/mongo/db/pipeline/pipeline_d.h
+++ b/src/mongo/db/pipeline/pipeline_d.h
@@ -108,7 +108,8 @@ private:
      * Creates a DocumentSourceCursor from the given PlanExecutor and adds it to the front of the
      * Pipeline.
      */
-    static void addCursorSource(const boost::intrusive_ptr<Pipeline>& pipeline,
+    static void addCursorSource(Collection* collection,
+                                const boost::intrusive_ptr<Pipeline>& pipeline,
                                 const boost::intrusive_ptr<ExpressionContext>& expCtx,
                                 std::unique_ptr<PlanExecutor> exec,
                                 DepsTracker deps,
diff --git a/src/mongo/dbtests/documentsourcetests.cpp b/src/mongo/dbtests/documentsourcetests.cpp
index 51d34334ecd..a5add096d27 100644
--- a/src/mongo/dbtests/documentsourcetests.cpp
+++ b/src/mongo/dbtests/documentsourcetests.cpp
@@ -110,7 +110,8 @@ protected:
         exec->saveState();
         exec->registerExec(ctx.getCollection());
 
-        _source = DocumentSourceCursor::create(nss.ns(), std::move(exec), _ctx);
+        _source =
+            DocumentSourceCursor::create(ctx.getCollection(), nss.ns(), std::move(exec), _ctx);
     }
 
     intrusive_ptr<ExpressionContextForTest> ctx() {
diff --git a/src/mongo/dbtests/query_plan_executor.cpp b/src/mongo/dbtests/query_plan_executor.cpp
index e3e19873228..f91d8fce94a 100644
--- a/src/mongo/dbtests/query_plan_executor.cpp
+++ b/src/mongo/dbtests/query_plan_executor.cpp
@@ -292,7 +292,8 @@ public:
         // Wrap the "inner" plan executor in a DocumentSourceCursor and add it as the first source
         // in the pipeline.
         innerExec->saveState();
-        auto cursorSource = DocumentSourceCursor::create(nss.ns(), std::move(innerExec), expCtx);
+        auto cursorSource =
+            DocumentSourceCursor::create(collection, nss.ns(), std::move(innerExec), expCtx);
         auto pipeline = assertGet(Pipeline::create({cursorSource}, expCtx));
 
         // Create the output PlanExecutor that pulls results from the pipeline.