SERVER-22048 Add index stats test covering multiple potential plan match

(cherry picked from commit c8bc89143725e9c88bb80ca84a261024e170dc50)

2 files changed, 118 insertions, 22 deletions

diff --git a/jstests/core/index_stats.js b/jstests/core/index_stats.js
index acf2951213b..8d2f9fd4d60 100644
--- a/jstests/core/index_stats.js
+++ b/jstests/core/index_stats.js
@@ -1,10 +1,13 @@
 (function() {
     "use strict";
+
+    load("jstests/libs/analyze_plan.js");
+
     var colName = "jstests_index_stats";
     var col = db[colName];
     col.drop();
 
-    var getUsageCount = function (indexName) {
+    var getUsageCount = function(indexName) {
         var cursor = col.aggregate([{$indexStats: {}}]);
         while (cursor.hasNext()) {
             var doc = cursor.next();
@@ -17,7 +20,7 @@
         return undefined;
     }
 
-    var getIndexKey = function (indexName) {
+    var getIndexKey = function(indexName) {
         var cursor = col.aggregate([{$indexStats: {}}]);
         while (cursor.hasNext()) {
             var doc = cursor.next();
@@ -30,31 +33,51 @@
         return undefined;
     }
 
+    var getIndexNamesForWinningPlan = function(explain) {
+        var indexNameList = [];
+        var winningStages = getPlanStages(explain.queryPlanner.winningPlan, "IXSCAN");
+        for (var i = 0; i < winningStages.length; ++i) {
+            indexNameList.push(winningStages[i].indexName);
+        }
+
+        return indexNameList;
+    }
+
     assert.writeOK(col.insert({a: 1, b: 1, c: 1}));
     assert.writeOK(col.insert({a: 2, b: 2, c: 2}));
     assert.writeOK(col.insert({a: 3, b: 3, c: 3}));
 
+    //
     // Confirm no index stats object exists prior to index creation.
+    //
     col.findOne({a: 1});
     assert.eq(undefined, getUsageCount("a_1"));
 
+    //
     // Create indexes.
+    //
     assert.commandWorked(col.createIndex({a: 1}, {name: "a_1"}));
     assert.commandWorked(col.createIndex({b: 1, c: 1}, {name: "b_1_c_1"}));
-    var countA = 0;
-    var countB = 0;
+    var countA = 0;  // Tracks expected index access for "a_1".
+    var countB = 0;  // Tracks expected index access for "b_1_c_1".
 
+    //
     // Confirm a stats object exists post index creation (with 0 count).
+    //
     assert.eq(countA, getUsageCount("a_1"));
     assert.eq({a: 1}, getIndexKey("a_1"));
 
+    //
     // Confirm index stats tick on find().
+    //
     col.findOne({a: 1});
     countA++;
 
     assert.eq(countA, getUsageCount("a_1"));
 
+    //
     // Confirm index stats tick on findAndModify() update.
+    //
     var res = db.runCommand({findAndModify: colName,
                              query: {a: 1},
                              update: {$set: {d: 1}},
@@ -63,7 +86,9 @@
     countA++;
     assert.eq(countA, getUsageCount("a_1"));
 
+    //
     // Confirm index stats tick on findAndModify() delete.
+    //
     res = db.runCommand({findAndModify: colName,
                          query: {a: 2},
                          remove: true});
@@ -72,30 +97,67 @@
     assert.eq(countA, getUsageCount("a_1"));
     assert.writeOK(col.insert(res.value));
 
+    //
+    // Confirm $and operation ticks indexes for winning plan, but not rejected plans.
+    //
+
+    // Run explain to determine which indexes would be used for this query. Note that index
+    // access counters are not incremented for explain execution.
+    var explain = col.find({a: 2, b: 2}).explain("queryPlanner");
+    var indexNameList = getIndexNamesForWinningPlan(explain);
+    assert.gte(indexNameList.length, 1);
+
+    for (var i = 0; i < indexNameList.length; ++i) {
+        // Increment the expected $indexStats count for each index used.
+        var name = indexNameList[i];
+        if (name === "a_1") {
+            countA++;
+        }
+        else {
+            assert(name === "b_1_c_1");
+            countB++;
+        }
+    }
+
+    // Run the query again without explain to increment index access counters.
+    col.findOne({a: 2, b: 2});
+    // Check all indexes for proper count.
+    assert.eq(countA, getUsageCount("a_1"));
+    assert.eq(countB, getUsageCount("b_1_c_1"));
+    assert.eq(0, getUsageCount("_id_"));
+
+    //
     // Confirm index stats tick on distinct().
+    //
     res = db.runCommand({distinct: colName, key: "b", query: {b: 1}});
     assert.commandWorked(res);
     countB++;
     assert.eq(countB, getUsageCount("b_1_c_1"));
 
+    //
     // Confirm index stats tick on group().
+    //
     res = db.runCommand({group: {ns: colName,
                          key: {b: 1, c: 1},
                          cond: {b: {$gt: 0}},
-                         $reduce: function(curr, result) {}, 
+                         $reduce: function(curr, result) {},
                          initial: {}}});
     assert.commandWorked(res);
     countB++;
     assert.eq(countB, getUsageCount("b_1_c_1"));
 
+    //
     // Confirm index stats tick on aggregate w/ match.
+    //
     res = db.runCommand({aggregate: colName,
                          pipeline: [{$match: {b: 1}}]});
     assert.commandWorked(res);
     countB++;
     assert.eq(countB, getUsageCount("b_1_c_1"));
 
+    //
     // Confirm index stats tick on mapReduce with query.
+    //
     res = db.runCommand({mapReduce: colName,
                          map: function() {emit(this.b, this.c);},
                          reduce: function(key, val) {return val;},
@@ -105,32 +167,44 @@
     countB++;
     assert.eq(countB, getUsageCount("b_1_c_1"));
 
+    //
     // Confirm index stats tick on update().
+    //
     assert.writeOK(col.update({a: 2}, {$set: {d: 2}}));
     countA++;
     assert.eq(countA, getUsageCount("a_1"));
 
+    //
     // Confirm index stats tick on remove().
+    //
     assert.writeOK(col.remove({a: 2}));
     countA++;
     assert.eq(countA, getUsageCount("a_1"));
 
+    //
     // Confirm multiple index $or operation ticks all involved indexes.
+    //
     col.findOne({$or: [{a: 1}, {b: 1, c: 1}]});
     countA++;
     countB++;
     assert.eq(countA, getUsageCount("a_1"));
     assert.eq(countB, getUsageCount("b_1_c_1"));
 
+    //
     // Confirm index stats object does not exist post index drop.
+    //
     assert.commandWorked(col.dropIndex("b_1_c_1"));
     countB = 0;
     assert.eq(undefined, getUsageCount("b_1_c_1"));
 
+    //
     // Confirm index stats object exists with count 0 once index is recreated.
+    //
     assert.commandWorked(col.createIndex({b: 1, c: 1}, {name: "b_1_c_1"}));
     assert.eq(countB, getUsageCount("b_1_c_1"));
 
+    //
     // Confirm that retrieval fails if $indexStats is not in the first pipeline position.
+    //
     assert.throws(function() { col.aggregate([{$match: {}}, {$indexStats: {}}]) });
 })();
diff --git a/jstests/libs/analyze_plan.js b/jstests/libs/analyze_plan.js
index 25341f5995f..12a28671866 100644
--- a/jstests/libs/analyze_plan.js
+++ b/jstests/libs/analyze_plan.js
@@ -3,31 +3,53 @@
 // scan or whether the plan is covered (index only).
 
 /**
- * Given the root stage of explain's JSON representation of a query plan ('root'), returns the
- * subdocument with its stage as 'stage'. Returns null if the plan does not have such a stage.
+ * Given the root stage of explain's JSON representation of a query plan ('root'), returns all
+ * subdocuments whose stage is 'stage'. Returns an empty array if the plan does not have the
+ * requested stage.
  */
-function getPlanStage(root, stage) {
+function getPlanStages(root, stage) {
+    var results = [];
+
     if (root.stage === stage) {
-        return root;
-    } else if ("inputStage" in root) {
-        return getPlanStage(root.inputStage, stage);
-    } else if ("inputStages" in root) {
+        results.push(root);
+    }
+
+    if ("inputStage" in root) {
+        results = results.concat(getPlanStages(root.inputStage, stage));
+    }
+
+    if ("inputStages" in root) {
         for (var i = 0; i < root.inputStages.length; i++) {
-            var stage = getPlanStage(root.inputStages[i], stage);
-            if (stage !== null) {
-                return stage;
-            }
+            results = results.concat(getPlanStages(root.inputStages[i], stage));
         }
-    } else if ("shards" in root) {
+    }
+
+    if ("shards" in root) {
         for (var i = 0; i < root.shards.length; i++) {
-            var stage = getPlanStage(root.shards[i].winningPlan, stage);
-            if (stage !== null) {
-                return stage;
-            }
+            results = results.concat(getPlanStages(root.shards[i].winningPlan, stage));
         }
     }
 
-    return null;
+    return results;
+}
+
+/**
+ * Given the root stage of explain's JSON representation of a query plan ('root'), returns the
+ * subdocument with its stage as 'stage'. Returns null if the plan does not have such a stage.
+ * Asserts that no more than one stage is a match.
+ */
+function getPlanStage(root, stage) {
+    var planStageList = getPlanStages(root, stage);
+
+    if (planStageList.length === 0) {
+        return null;
+    }
+    else {
+        assert(planStageList.length === 1,
+               "getPlanStage expects to find 0 or 1 matching stages. planStageList: "
+               + tojson(planStageList));
+        return planStageList[0];
+    }
 }
 
 /**