SERVER-5905 Add collStats aggregation stage

8 files changed, 392 insertions, 1 deletions

diff --git a/jstests/auth/lib/commands_lib.js b/jstests/auth/lib/commands_lib.js
index 2527e679c66..f92d2664cca 100644
--- a/jstests/auth/lib/commands_lib.js
+++ b/jstests/auth/lib/commands_lib.js
@@ -358,6 +358,52 @@ var authCommandsLib = {
           ]
         },
         {
+          testname: "aggregate_collStats",
+          command: {aggregate: "foo", pipeline: [{$collStats: {latencyStats: {}}}]},
+          setup: function(db) {
+              db.createCollection("foo");
+          },
+          teardown: function(db) {
+              db.foo.drop();
+          },
+          testcases: [
+              {
+                runOnDb: firstDbName,
+                roles: {
+                    read: 1,
+                    readAnyDatabase: 1,
+                    readWrite: 1,
+                    readWriteAnyDatabase: 1,
+                    dbAdmin: 1,
+                    dbAdminAnyDatabase: 1,
+                    dbOwner: 1,
+                    clusterMonitor: 1,
+                    clusterAdmin: 1,
+                    backup: 1,
+                    root: 1,
+                    __system: 1
+                },
+                privileges:
+                    [{resource: {db: firstDbName, collection: "foo"}, actions: ["collStats"]}]
+              },
+              {
+                runOnDb: secondDbName,
+                roles: {
+                    readAnyDatabase: 1,
+                    readWriteAnyDatabase: 1,
+                    dbAdminAnyDatabase: 1,
+                    clusterMonitor: 1,
+                    clusterAdmin: 1,
+                    backup: 1,
+                    root: 1,
+                    __system: 1
+                },
+                privileges:
+                    [{resource: {db: secondDbName, collection: "foo"}, actions: ["collStats"]}]
+              }
+          ]
+        },
+        {
           testname: "appendOplogNote",
           command: {appendOplogNote: 1, data: {a: 1}},
           skipSharded: true,
diff --git a/jstests/core/operation_latency_histogram.js b/jstests/core/operation_latency_histogram.js
new file mode 100644
index 00000000000..4ee4b2f5946
--- /dev/null
+++ b/jstests/core/operation_latency_histogram.js
@@ -0,0 +1,200 @@
+// Checks that histogram counters for collections are updated as we expect.
+
+(function() {
+    "use strict";
+    var name = "operationalLatencyHistogramTest";
+
+    var testDB = db.getSiblingDB(name);
+    var testColl = testDB[name + "coll"];
+    var inShardedCollection = (db.runCommand({isMaster: 1}).msg == "isdbgrid");
+
+    testColl.drop();
+
+    // Insert a document initially so sharding passthrough gets a non-empty response
+    // from latencyStats.
+    testColl.insert({x: -1});
+
+    // Test aggregation command output format.
+    var commandResult = testDB.runCommand(
+        {aggregate: testColl.getName(), pipeline: [{$collStats: {latencyStats: {}}}]});
+    assert.commandWorked(commandResult);
+    assert(commandResult.result.length == 1);
+
+    var stats = commandResult.result[0];
+    var histogramTypes = ["reads", "writes", "commands"];
+
+    assert(stats.hasOwnProperty("localTime"));
+    assert(stats.hasOwnProperty("latencyStats"));
+
+    histogramTypes.forEach(function(key) {
+        assert(stats.latencyStats.hasOwnProperty(key));
+        assert(stats.latencyStats[key].hasOwnProperty("ops"));
+        assert(stats.latencyStats[key].hasOwnProperty("latency"));
+    });
+
+    function getHistogramStats() {
+        return testColl.latencyStats().toArray()[0].latencyStats;
+    }
+
+    var lastHistogram = getHistogramStats();
+
+    // Checks that the difference in the histogram is what we expect, and also accounts for the
+    // $collStats aggregation stage itself.
+    function checkHistogramDiff(reads, writes, commands) {
+        var thisHistogram = getHistogramStats();
+        // Running the aggregation itself will increment read stats by one.
+        assert.eq(thisHistogram.reads.ops - lastHistogram.reads.ops, reads + 1);
+        assert.eq(thisHistogram.writes.ops - lastHistogram.writes.ops, writes);
+        assert.eq(thisHistogram.commands.ops - lastHistogram.commands.ops, commands);
+        return thisHistogram;
+    }
+
+    // Insert
+    var numRecords = 100;
+    for (var i = 0; i < numRecords; i++) {
+        assert.writeOK(testColl.insert({_id: i}));
+    }
+    lastHistogram = checkHistogramDiff(0, numRecords, 0);
+
+    // Update
+    for (var i = 0; i < numRecords; i++) {
+        assert.writeOK(testColl.update({_id: i}, {x: i}));
+    }
+    lastHistogram = checkHistogramDiff(0, numRecords, 0);
+
+    // Find
+    var cursors = [];
+    for (var i = 0; i < numRecords; i++) {
+        cursors[i] = testColl.find({x: {$gte: i}}).batchSize(2);
+        assert.eq(cursors[i].next()._id, i);
+    }
+    lastHistogram = checkHistogramDiff(numRecords, 0, 0);
+
+    // GetMore
+    for (var i = 0; i < numRecords / 2; i++) {
+        // Trigger two getmore commands.
+        assert.eq(cursors[i].next()._id, i + 1);
+        assert.eq(cursors[i].next()._id, i + 2);
+        assert.eq(cursors[i].next()._id, i + 3);
+        assert.eq(cursors[i].next()._id, i + 4);
+    }
+    lastHistogram = checkHistogramDiff(numRecords, 0, 0);
+
+    if (!inShardedCollection) {
+        // KillCursors
+        // The last cursor has no additional results, hence does not need to be closed.
+        for (var i = 0; i < numRecords - 1; i++) {
+            cursors[i].close();
+        }
+        lastHistogram = checkHistogramDiff(0, 0, numRecords - 1);
+    }
+
+    // Remove
+    for (var i = 0; i < numRecords; i++) {
+        assert.writeOK(testColl.remove({_id: i}));
+    }
+    lastHistogram = checkHistogramDiff(0, numRecords, 0);
+
+    // Upsert
+    for (var i = 0; i < numRecords; i++) {
+        assert.writeOK(testColl.update({_id: i}, {x: i}, {upsert: 1}));
+    }
+    lastHistogram = checkHistogramDiff(0, numRecords, 0);
+
+    // Aggregate
+    for (var i = 0; i < numRecords; i++) {
+        testColl.aggregate([{$match: {x: i}}, {$group: {_id: "$x"}}]);
+    }
+    // TODO SERVER-24704: Agg is currently counted by Top as two operations, but should be counted
+    // as one.
+    lastHistogram = checkHistogramDiff(2 * numRecords, 0, 0);
+
+    // Count
+    for (var i = 0; i < numRecords; i++) {
+        testColl.count({x: i});
+    }
+    lastHistogram = checkHistogramDiff(numRecords, 0, 0);
+
+    // Group
+    testColl.group({initial: {}, reduce: function() {}, key: {a: 1}});
+    lastHistogram = checkHistogramDiff(1, 0, 0);
+
+    if (!inShardedCollection) {
+        // ParallelCollectionScan
+        testDB.runCommand({parallelCollectionScan: testColl.getName(), numCursors: 5});
+        lastHistogram = checkHistogramDiff(0, 0, 1);
+    }
+
+    // FindAndModify
+    testColl.findAndModify({query: {}, update: {pt: {type: "Point", coordinates: [0, 0]}}});
+    // TODO SERVER-24462: findAndModify is not currently counted in Top.
+    lastHistogram = checkHistogramDiff(0, 0, 0);
+
+    // CreateIndex
+    assert.commandWorked(testColl.createIndex({pt: "2dsphere"}));
+    // TODO SERVER-24705: createIndex is not currently counted in Top.
+    lastHistogram = checkHistogramDiff(0, 0, 0);
+
+    // GeoNear
+    assert.commandWorked(testDB.runCommand({
+        geoNear: testColl.getName(),
+        near: {type: "Point", coordinates: [0, 0]},
+        spherical: true
+    }));
+    lastHistogram = checkHistogramDiff(1, 0, 0);
+
+    // GetIndexes
+    testColl.getIndexes();
+    lastHistogram = checkHistogramDiff(0, 0, 1);
+
+    // Reindex
+    assert.commandWorked(testColl.reIndex());
+    lastHistogram = checkHistogramDiff(0, 0, 1);
+
+    // DropIndex
+    assert.commandWorked(testColl.dropIndex({pt: "2dsphere"}));
+    lastHistogram = checkHistogramDiff(0, 0, 1);
+
+    // Explain
+    testColl.explain().find().next();
+    lastHistogram = checkHistogramDiff(0, 0, 1);
+
+    // CollStats
+    assert.commandWorked(testDB.runCommand({collStats: testColl.getName()}));
+    lastHistogram = checkHistogramDiff(0, 0, 1);
+
+    // CollMod
+    assert.commandWorked(
+        testDB.runCommand({collStats: testColl.getName(), validationLevel: "off"}));
+    lastHistogram = checkHistogramDiff(0, 0, 1);
+
+    if (!inShardedCollection) {
+        // Compact
+        // Use force:true in case we're in replset.
+        var commandResult = testDB.runCommand({compact: testColl.getName(), force: true});
+        // If storage engine supports compact, it should count as a command.
+        if (!commandResult.ok) {
+            assert.commandFailedWithCode(commandResult, ErrorCodes.CommandNotSupported);
+        }
+        lastHistogram = checkHistogramDiff(0, 0, 1);
+    }
+
+    // DataSize
+    testColl.dataSize();
+    lastHistogram = checkHistogramDiff(0, 0, 1);
+
+    // PlanCache
+    testColl.getPlanCache().listQueryShapes();
+    lastHistogram = checkHistogramDiff(0, 0, 1);
+
+    // Commands which occur on the database only should not effect the collection stats.
+    assert.commandWorked(testDB.serverStatus());
+    lastHistogram = checkHistogramDiff(0, 0, 0);
+
+    assert.commandWorked(testColl.runCommand("whatsmyuri"));
+    lastHistogram = checkHistogramDiff(0, 0, 0);
+
+    // Test non-command.
+    assert.commandFailed(testColl.runCommand("IHopeNobodyEverMakesThisACommand"));
+    lastHistogram = checkHistogramDiff(0, 0, 0);
+}());
\ No newline at end of file
diff --git a/src/mongo/db/pipeline/SConscript b/src/mongo/db/pipeline/SConscript
index fc77727d588..32c0f82eae8 100644
--- a/src/mongo/db/pipeline/SConscript
+++ b/src/mongo/db/pipeline/SConscript
@@ -154,6 +154,7 @@ docSourceEnv.Library(
     target='document_source',
     source=[
         'document_source.cpp',
+        'document_source_coll_stats.cpp',
         'document_source_count.cpp',
         'document_source_geo_near.cpp',
         'document_source_graph_lookup.cpp',
@@ -185,6 +186,7 @@ docSourceEnv.Library(
         '$BUILD_DIR/mongo/db/matcher/expressions',
         '$BUILD_DIR/mongo/db/matcher/expression_algo',
         '$BUILD_DIR/mongo/db/service_context',
+        '$BUILD_DIR/mongo/db/stats/top',
         '$BUILD_DIR/mongo/db/storage/storage_options',
         '$BUILD_DIR/mongo/db/storage/wiredtiger/storage_wiredtiger_customization_hooks',
         '$BUILD_DIR/third_party/shim_snappy',
diff --git a/src/mongo/db/pipeline/document_source.h b/src/mongo/db/pipeline/document_source.h
index 49466a9df50..df3f4eecaf5 100644
--- a/src/mongo/db/pipeline/document_source.h
+++ b/src/mongo/db/pipeline/document_source.h
@@ -348,6 +348,12 @@ public:
 
         virtual bool hasUniqueIdIndex(const NamespaceString& ns) const = 0;
 
+        /**
+         * Appends operation latency statistics for collection "nss" to "builder"
+         */
+        virtual void appendLatencyStats(const NamespaceString& nss,
+                                        BSONObjBuilder* builder) const = 0;
+
         // Add new methods as needed.
     };
 
@@ -1769,4 +1775,29 @@ public:
 private:
     DocumentSourceCount() = default;
 };
-}
+
+/**
+ * Provides a document source interface to retrieve collection-level statistics for a given
+ * collection.
+ */
+class DocumentSourceCollStats : public DocumentSourceNeedsMongod {
+public:
+    DocumentSourceCollStats(const boost::intrusive_ptr<ExpressionContext>& pExpCtx)
+        : DocumentSourceNeedsMongod(pExpCtx) {}
+
+    boost::optional<Document> getNext() final;
+
+    const char* getSourceName() const final;
+
+    bool isValidInitialSource() const final;
+
+    Value serialize(bool explain = false) const;
+
+    static boost::intrusive_ptr<DocumentSource> createFromBson(
+        BSONElement elem, const boost::intrusive_ptr<ExpressionContext>& pExpCtx);
+
+private:
+    bool _latencySpecified = false;
+    bool _finished = false;
+};
+}  // namespace mongo
diff --git a/src/mongo/db/pipeline/document_source_coll_stats.cpp b/src/mongo/db/pipeline/document_source_coll_stats.cpp
new file mode 100644
index 00000000000..ab95e1d08d3
--- /dev/null
+++ b/src/mongo/db/pipeline/document_source_coll_stats.cpp
@@ -0,0 +1,98 @@
+/**
+ * Copyright (C) 2016 MongoDB Inc.
+ *
+ * This program is free software: you can redistribute it and/or  modify
+ * it under the terms of the GNU Affero General Public License, version 3,
+ * as published by the Free Software Foundation.
+ *
+ * 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
+ * GNU Affero General Public License for more details.
+ *
+ * You should have received a copy of the GNU Affero General Public License
+ * along with this program.  If not, see <http://www.gnu.org/licenses/>.
+ *
+ * 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 GNU Affero General 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/platform/basic.h"
+
+#include "mongo/db/pipeline/document_source.h"
+
+#include "mongo/bson/bsonobj.h"
+#include "mongo/db/stats/top.h"
+#include "mongo/util/time_support.h"
+
+using boost::intrusive_ptr;
+
+namespace mongo {
+
+REGISTER_DOCUMENT_SOURCE(collStats, DocumentSourceCollStats::createFromBson);
+
+const char* DocumentSourceCollStats::getSourceName() const {
+    return "$collStats";
+}
+
+intrusive_ptr<DocumentSource> DocumentSourceCollStats::createFromBson(
+    BSONElement specElem, const intrusive_ptr<ExpressionContext>& pExpCtx) {
+    uassert(40166,
+            str::stream() << "$collStats must take a nested object but found: " << specElem,
+            specElem.type() == BSONType::Object);
+    intrusive_ptr<DocumentSourceCollStats> collStats(new DocumentSourceCollStats(pExpCtx));
+
+    for (const auto& elem : specElem.embeddedObject()) {
+        StringData fieldName = elem.fieldNameStringData();
+
+        if (fieldName == "latencyStats") {
+            uassert(40167,
+                    str::stream() << "latencyStats argument must be an object, but found: " << elem,
+                    elem.type() == BSONType::Object);
+            collStats->_latencySpecified = true;
+        } else {
+            uasserted(40168, str::stream() << "unrecognized option to $collStats: " << fieldName);
+        }
+    }
+
+    return collStats;
+}
+
+boost::optional<Document> DocumentSourceCollStats::getNext() {
+    if (_finished) {
+        return boost::none;
+    }
+
+    _finished = true;
+
+    BSONObjBuilder builder;
+
+    builder.appendDate("localTime", jsTime());
+    if (_latencySpecified) {
+        _mongod->appendLatencyStats(pExpCtx->ns, &builder);
+    }
+
+    return Document(builder.obj());
+}
+
+bool DocumentSourceCollStats::isValidInitialSource() const {
+    return true;
+}
+
+Value DocumentSourceCollStats::serialize(bool explain) const {
+    if (_latencySpecified) {
+        return Value(DOC(getSourceName() << DOC("latencyStats" << Document())));
+    }
+    return Value(DOC(getSourceName() << Document()));
+}
+
+}  // namespace mongo
diff --git a/src/mongo/db/pipeline/pipeline.cpp b/src/mongo/db/pipeline/pipeline.cpp
index 42961d4e174..3b76d574cc7 100644
--- a/src/mongo/db/pipeline/pipeline.cpp
+++ b/src/mongo/db/pipeline/pipeline.cpp
@@ -143,6 +143,9 @@ Status Pipeline::checkAuthForCommand(ClientBasic* client,
         Privilege::addPrivilegeToPrivilegeVector(
             &privileges,
             Privilege(ResourcePattern::forAnyNormalResource(), ActionType::indexStats));
+    } else if (dps::extractElementAtPath(cmdObj, "pipeline.0.$collStats")) {
+        Privilege::addPrivilegeToPrivilegeVector(&privileges,
+                                                 Privilege(inputResource, ActionType::collStats));
     } else {
         // If no source requiring an alternative permission scheme is specified then default to
         // requiring find() privileges on the given namespace.
diff --git a/src/mongo/db/pipeline/pipeline_d.cpp b/src/mongo/db/pipeline/pipeline_d.cpp
index 7d74151b98c..c26333245bd 100644
--- a/src/mongo/db/pipeline/pipeline_d.cpp
+++ b/src/mongo/db/pipeline/pipeline_d.cpp
@@ -55,6 +55,7 @@
 #include "mongo/db/s/sharded_connection_info.h"
 #include "mongo/db/s/sharding_state.h"
 #include "mongo/db/service_context.h"
+#include "mongo/db/stats/top.h"
 #include "mongo/db/storage/record_store.h"
 #include "mongo/db/storage/sorted_data_interface.h"
 #include "mongo/s/chunk_version.h"
@@ -129,6 +130,10 @@ public:
         return collection->getIndexCatalog()->findIdIndex(_ctx->opCtx);
     }
 
+    void appendLatencyStats(const NamespaceString& nss, BSONObjBuilder* builder) const {
+        Top::get(_ctx->opCtx->getServiceContext()).appendLatencyStats(nss.ns(), builder);
+    }
+
 private:
     intrusive_ptr<ExpressionContext> _ctx;
     DBDirectClient _client;
diff --git a/src/mongo/shell/collection.js b/src/mongo/shell/collection.js
index 67f4d8327dc..7d7c79871c9 100644
--- a/src/mongo/shell/collection.js
+++ b/src/mongo/shell/collection.js
@@ -141,6 +141,8 @@ DBCollection.prototype.help = function() {
     print(
         "\tdb." + shortName +
         ".unsetWriteConcern( <write concern doc> ) - unsets the write concern for writes to the collection");
+    print("\tdb." + shortName +
+          ".latencyStats() - display operation latency histograms for this collection");
     // print("\tdb." + shortName + ".getDiskStorageStats({...}) - prints a summary of disk usage
     // statistics");
     // print("\tdb." + shortName + ".getPagesInRAM({...}) - prints a summary of storage pages
@@ -1752,6 +1754,10 @@ DBCollection.prototype._distinct = function(keyString, query) {
     return this._dbReadCommand({distinct: this._shortName, key: keyString, query: query || {}});
 };
 
+DBCollection.prototype.latencyStats = function() {
+    return this.aggregate([{$collStats: {latencyStats: {}}}]);
+};
+
 /**
  * PlanCache
  * Holds a reference to the collection.