SERVER-18940 Optimise sharded aggregations that are targeted to a single shard

1 files changed, 76 insertions, 22 deletions

diff --git a/jstests/libs/profiler.js b/jstests/libs/profiler.js
index e8282fecba6..d8b5f45bd06 100644
--- a/jstests/libs/profiler.js
+++ b/jstests/libs/profiler.js
@@ -1,11 +1,11 @@
 // Provides convenience methods for confirming system.profile content.
 
 // Retrieve latest system.profile entry.
-function getLatestProfilerEntry(inputDb, filter) {
+function getLatestProfilerEntry(profileDB, filter) {
     if (filter === null) {
         filter = {};
     }
-    var cursor = inputDb.system.profile.find(filter).sort({$natural: -1});
+    var cursor = profileDB.system.profile.find(filter).sort({$natural: -1});
     assert(
         cursor.hasNext(),
         "could not find any entries in the profile collection matching filter: " + tojson(filter));
@@ -25,29 +25,83 @@ function getProfilerProtocolStringForCommand(conn) {
     doassert(`Unknown prototocol string ${protocols}`);
 }
 
-// Throws an assertion if the profiler does not contain *exactly one* entry matching <filter>.
-// Optional arguments <errorMsgFilter> and <errorMsgProj> limit profiler output if this asserts.
-function profilerHasSingleMatchingEntryOrThrow(inputDb, filter, errorMsgFilter, errorMsgProj) {
-    assert.eq(inputDb.system.profile.find(filter).itcount(),
-              1,
-              "Expected exactly one op matching: " + tojson(filter) + " in profiler " +
-                  tojson(inputDb.system.profile.find(errorMsgFilter, errorMsgProj).toArray()));
+/**
+ * Throws an assertion if the profiler contains more than 'maxExpectedMatches' entries matching
+ * "filter", or if there are no matches. Optional arguments "errorMsgFilter" and "errorMsgProj"
+ * limit profiler output if this asserts.
+ */
+function profilerHasAtLeastOneAtMostNumMatchingEntriesOrThrow(
+    {profileDB, filter, maxExpectedMatches, errorMsgFilter, errorMsgProj}) {
+    assert(typeof maxExpectedMatches === 'number' && maxExpectedMatches > 0,
+           "'maxExpectedMatches' must be a number > 0");
+
+    const numMatches = profileDB.system.profile.find(filter).itcount();
+
+    assert.gt(numMatches,
+              0,
+              "Expected at least 1 op matching: " + tojson(filter) + " in profiler " +
+                  tojson(profileDB.system.profile.find(errorMsgFilter, errorMsgProj).toArray()));
+
+    assert.lte(numMatches,
+               maxExpectedMatches,
+               "Expected at most " + maxExpectedMatches + " op(s) matching: " + tojson(filter) +
+                   " in profiler " +
+                   tojson(profileDB.system.profile.find(errorMsgFilter, errorMsgProj).toArray()));
+}
+
+/**
+ * Throws an assertion if the profiler does not contain exactly 'numExpectedMatches' entries
+ * matching "filter". Optional arguments "errorMsgFilter" and "errorMsgProj" limit profiler output
+ * if this asserts.
+ */
+function profilerHasNumMatchingEntriesOrThrow(
+    {profileDB, filter, numExpectedMatches, errorMsgFilter, errorMsgProj}) {
+    assert(typeof numExpectedMatches === 'number' && numExpectedMatches >= 0,
+           "'numExpectedMatches' must be a number >= 0");
+
+    assert.eq(profileDB.system.profile.find(filter).itcount(),
+              numExpectedMatches,
+              "Expected exactly " + numExpectedMatches + " op(s) matching: " + tojson(filter) +
+                  " in profiler " +
+                  tojson(profileDB.system.profile.find(errorMsgFilter, errorMsgProj).toArray()));
 }
 
-// Throws an assertion if the profiler does not contain *at least one* entry matching <filter>.
-// Optional arguments <errorMsgFilter> and <errorMsgProj> limit profiler output if this asserts.
-function profilerHasMatchingEntryOrThrow(inputDb, filter, errorMsgFilter, errorMsgProj) {
-    assert.gte(inputDb.system.profile.find(filter).itcount(),
+/**
+ * Throws an assertion if the profiler does not contain any entries matching "filter". Optional
+ * arguments "errorMsgFilter" and "errorMsgProj" limit profiler output if this asserts.
+ */
+function profilerHasAtLeastOneMatchingEntryOrThrow(
+    {profileDB, filter, errorMsgFilter, errorMsgProj}) {
+    assert.gte(profileDB.system.profile.find(filter).itcount(),
                1,
-               "Expected at least one op matching: " + tojson(filter) + " in profiler " +
-                   tojson(inputDb.system.profile.find(errorMsgFilter, errorMsgProj).toArray()));
+               "Expected at least 1 op matching: " + tojson(filter) + " in profiler " +
+                   tojson(profileDB.system.profile.find(errorMsgFilter, errorMsgProj).toArray()));
 }
 
-// Throws an assertion if the profiler contains an entry matching <filter>.
-// Optional arguments <errorMsgFilter> and <errorMsgProj> limit profiler output if this asserts.
-function profilerDoesNotHaveMatchingEntryOrThrow(inputDb, filter, errorMsgFilter, errorMsgProj) {
-    assert.eq(inputDb.system.profile.find(filter).itcount(),
-              0,
-              "Expected no ops matching: " + tojson(filter) + " in profiler " +
-                  tojson(inputDb.system.profile.find(errorMsgFilter, errorMsgProj).toArray()));
+/**
+ * Throws an assertion if the profiler does not contain exactly one entry matching "filter".
+ * Optional arguments "errorMsgFilter" and "errorMsgProj" limit profiler output if this asserts.
+ */
+function profilerHasSingleMatchingEntryOrThrow({profileDB, filter, errorMsgFilter, errorMsgProj}) {
+    profilerHasNumMatchingEntriesOrThrow({
+        profileDB: profileDB,
+        filter: filter,
+        numExpectedMatches: 1,
+        errorMsgFilter: errorMsgFilter,
+        errorMsgProj: errorMsgProj
+    });
+}
+
+/**
+ * Throws an assertion if the profiler contains any entries matching "filter". Optional arguments
+ * "errorMsgFilter" and "errorMsgProj" limit profiler output if this asserts.
+ */
+function profilerHasZeroMatchingEntriesOrThrow({profileDB, filter, errorMsgFilter, errorMsgProj}) {
+    profilerHasNumMatchingEntriesOrThrow({
+        profileDB: profileDB,
+        filter: filter,
+        numExpectedMatches: 0,
+        errorMsgFilter: errorMsgFilter,
+        errorMsgProj: errorMsgProj
+    });
 }