SERVER-70352 Create statistics cache at the startup

17 files changed, 219 insertions, 57 deletions

diff --git a/src/mongo/db/SConscript b/src/mongo/db/SConscript
index 0b46b67c425..6f397fba70c 100644
--- a/src/mongo/db/SConscript
+++ b/src/mongo/db/SConscript
@@ -2317,6 +2317,7 @@ env.Library(
         '$BUILD_DIR/mongo/db/change_stream_options_manager',
         '$BUILD_DIR/mongo/db/change_streams_cluster_parameter',
         '$BUILD_DIR/mongo/db/pipeline/change_stream_expired_pre_image_remover',
+        '$BUILD_DIR/mongo/db/query/ce/query_ce',
         '$BUILD_DIR/mongo/db/set_change_stream_state_coordinator',
         '$BUILD_DIR/mongo/idl/cluster_server_parameter',
         '$BUILD_DIR/mongo/idl/cluster_server_parameter_op_observer',
diff --git a/src/mongo/db/exec/batched_delete_stage.h b/src/mongo/db/exec/batched_delete_stage.h
index 84c166ba0dd..511eacee901 100644
--- a/src/mongo/db/exec/batched_delete_stage.h
+++ b/src/mongo/db/exec/batched_delete_stage.h
@@ -114,7 +114,7 @@ public:
     // Returns true when no more work can be done (there are no more deletes to commit).
     bool isEOF() final;
 
-    std::unique_ptr<PlanStageStats> getStats() final;
+    std::unique_ptr<mongo::PlanStageStats> getStats() final;
 
     const SpecificStats* getSpecificStats() const final;
 
diff --git a/src/mongo/db/exec/delete_stage.h b/src/mongo/db/exec/delete_stage.h
index 75b5f7cb2fe..bf9b1e9cc52 100644
--- a/src/mongo/db/exec/delete_stage.h
+++ b/src/mongo/db/exec/delete_stage.h
@@ -128,7 +128,7 @@ public:
         return STAGE_DELETE;
     }
 
-    std::unique_ptr<PlanStageStats> getStats();
+    std::unique_ptr<mongo::PlanStageStats> getStats();
 
     const SpecificStats* getSpecificStats() const;
 
diff --git a/src/mongo/db/exec/plan_stage.h b/src/mongo/db/exec/plan_stage.h
index a5b894e529f..e4816cf943d 100644
--- a/src/mongo/db/exec/plan_stage.h
+++ b/src/mongo/db/exec/plan_stage.h
@@ -320,7 +320,7 @@ public:
      * Creates plan stats tree which has the same topology as the original execution tree,
      * but has a separate lifetime.
      */
-    virtual std::unique_ptr<PlanStageStats> getStats() = 0;
+    virtual std::unique_ptr<mongo::PlanStageStats> getStats() = 0;
 
     /**
      * Get the CommonStats for this stage. The pointer is *not* owned by the caller.
@@ -329,7 +329,7 @@ public:
      * It must not exist past the stage. If you need the stats to outlive the stage,
      * use the getStats(...) method above.
      */
-    const CommonStats* getCommonStats() const {
+    const mongo::CommonStats* getCommonStats() const {
         return &_commonStats;
     }
 
@@ -413,7 +413,7 @@ protected:
     }
 
     Children _children;
-    CommonStats _commonStats;
+    mongo::CommonStats _commonStats;
 
 private:
     OperationContext* _opCtx;
diff --git a/src/mongo/db/exec/plan_stats.h b/src/mongo/db/exec/plan_stats.h
index d7336e79a31..8bf785a5b82 100644
--- a/src/mongo/db/exec/plan_stats.h
+++ b/src/mongo/db/exec/plan_stats.h
@@ -180,7 +180,7 @@ private:
     BasePlanStageStats& operator=(const BasePlanStageStats<C, T>&) = delete;
 };
 
-using PlanStageStats = BasePlanStageStats<CommonStats, StageType>;
+using PlanStageStats = BasePlanStageStats<mongo::CommonStats, StageType>;
 
 struct AndHashStats : public SpecificStats {
     AndHashStats() = default;
diff --git a/src/mongo/db/mongod_main.cpp b/src/mongo/db/mongod_main.cpp
index cf5deb4ef0a..397e01af848 100644
--- a/src/mongo/db/mongod_main.cpp
+++ b/src/mongo/db/mongod_main.cpp
@@ -107,6 +107,8 @@
 #include "mongo/db/periodic_runner_job_abort_expired_transactions.h"
 #include "mongo/db/pipeline/change_stream_expired_pre_image_remover.h"
 #include "mongo/db/pipeline/process_interface/replica_set_node_process_interface.h"
+#include "mongo/db/query/ce/stats_cache_loader_impl.h"
+#include "mongo/db/query/ce/stats_catalog.h"
 #include "mongo/db/query/internal_plans.h"
 #include "mongo/db/read_write_concern_defaults_cache_lookup_mongod.h"
 #include "mongo/db/repl/drop_pending_collection_reaper.h"
@@ -822,6 +824,10 @@ ExitCode _initAndListen(ServiceContext* serviceContext, int listenPort) {
 
     LogicalSessionCache::set(serviceContext, makeLogicalSessionCacheD(kind));
 
+    auto cacheLoader = std::make_unique<StatsCacheLoaderImpl>();
+    auto catalog = std::make_unique<StatsCatalog>(serviceContext, std::move(cacheLoader));
+    StatsCatalog::set(serviceContext, std::move(catalog));
+
     // MessageServer::run will return when exit code closes its socket and we don't need the
     // operation context anymore
     startupOpCtx.reset();
diff --git a/src/mongo/db/query/ce/SConscript b/src/mongo/db/query/ce/SConscript
index 7529c1be714..aa7dee358f5 100644
--- a/src/mongo/db/query/ce/SConscript
+++ b/src/mongo/db/query/ce/SConscript
@@ -11,6 +11,7 @@ env.Library(
         'ce_sampling.cpp',
         'collection_statistics_impl.cpp',
         'histogram_estimation.cpp',
+        'stats_catalog.cpp',
         'stats_cache.cpp',
         'stats_cache_loader_impl.cpp',
     ],
diff --git a/src/mongo/db/query/ce/collection_statistics_impl.cpp b/src/mongo/db/query/ce/collection_statistics_impl.cpp
index a0a0bab0eb3..7bf6b4e7a11 100644
--- a/src/mongo/db/query/ce/collection_statistics_impl.cpp
+++ b/src/mongo/db/query/ce/collection_statistics_impl.cpp
@@ -29,7 +29,7 @@
 
 #include "mongo/db/query/ce/collection_statistics_impl.h"
 #include "mongo/db/client.h"
-#include "mongo/db/query/ce/stats_cache.h"
+#include "mongo/db/query/ce/stats_catalog.h"
 
 namespace mongo::ce {
 
@@ -52,13 +52,17 @@ const ArrayHistogram* CollectionStatisticsImpl::getHistogram(const std::string&
         uassert(8423368, "no current client", Client::getCurrent());
         auto opCtx = Client::getCurrent()->getOperationContext();
         uassert(8423367, "no operation context", opCtx);
-        StatsCache& cache = StatsCache::get(opCtx);
-        auto handle = cache.acquire(opCtx, std::make_pair(_nss, path));
-        if (!handle) {
+        StatsCatalog& statsCatalog = StatsCatalog::get(opCtx);
+        const auto swHistogram = statsCatalog.getHistogram(opCtx, _nss, path);
+        if (!swHistogram.isOK()) {
+            if (swHistogram != ErrorCodes::NamespaceNotFound) {
+                uasserted(swHistogram.getStatus().code(),
+                          str::stream() << "Error getting histograms for path " << _nss << " : "
+                                        << path << swHistogram.getStatus().reason());
+            }
             return nullptr;
         }
-
-        auto histogram = *(handle.get());
+        const auto histogram = std::move(swHistogram.getValue());
         addHistogram(path, histogram);
         return histogram.get();
     }
diff --git a/src/mongo/db/query/ce/scalar_histogram.cpp b/src/mongo/db/query/ce/scalar_histogram.cpp
index f6dc5c23057..9ac9f088105 100644
--- a/src/mongo/db/query/ce/scalar_histogram.cpp
+++ b/src/mongo/db/query/ce/scalar_histogram.cpp
@@ -59,7 +59,7 @@ std::string Bucket::toString() const {
 
 ScalarHistogram::ScalarHistogram() : ScalarHistogram({}, {}) {}
 
-ScalarHistogram::ScalarHistogram(const Histogram& histogram) {
+ScalarHistogram::ScalarHistogram(const StatsHistogram& histogram) {
 
     for (const auto& bucket : histogram.getBuckets()) {
         Bucket b(bucket.getBoundaryCount(),
diff --git a/src/mongo/db/query/ce/scalar_histogram.h b/src/mongo/db/query/ce/scalar_histogram.h
index 3633cf2ff3f..932f980bce4 100644
--- a/src/mongo/db/query/ce/scalar_histogram.h
+++ b/src/mongo/db/query/ce/scalar_histogram.h
@@ -76,7 +76,7 @@ struct Bucket {
 class ScalarHistogram {
 public:
     ScalarHistogram();
-    ScalarHistogram(const Histogram& histogram);
+    ScalarHistogram(const StatsHistogram& histogram);
     ScalarHistogram(sbe::value::Array bounds, std::vector<Bucket> buckets);
 
     std::string toString() const;
diff --git a/src/mongo/db/query/ce/stats.idl b/src/mongo/db/query/ce/stats.idl
index 684b2fe56ac..a43bc8beba3 100644
--- a/src/mongo/db/query/ce/stats.idl
+++ b/src/mongo/db/query/ce/stats.idl
@@ -46,7 +46,7 @@ structs:
             cumulativeDistincts:
                 type: long
 
-    Histogram:
+    StatsHistogram:
         description: "MaxDiff Histogram"
         fields:
             buckets:
@@ -75,11 +75,11 @@ counts"
         description: "Array Histogram"
         fields:
             minHistogram:
-                type: Histogram 
+                type: StatsHistogram
             maxHistogram:
-                type: Histogram 
+                type: StatsHistogram
             uniqueHistogram:
-                type: Histogram 
+                type: StatsHistogram
             boolCount:
                 type: BoolCount
             typeCount:
@@ -97,7 +97,7 @@ counts"
                 type: array<TypeTag>
                 optional: true
             scalarHistogram:
-                type: Histogram
+                type: StatsHistogram
                 optional: true
             arrayStatistics:
                 type: StatsArrayHistogram
diff --git a/src/mongo/db/query/ce/stats_cache.cpp b/src/mongo/db/query/ce/stats_cache.cpp
index decd5fc1c5c..2fb2be400a6 100644
--- a/src/mongo/db/query/ce/stats_cache.cpp
+++ b/src/mongo/db/query/ce/stats_cache.cpp
@@ -79,28 +79,4 @@ StatsCache::LookupResult StatsCache::_lookupStats(OperationContext* opCtx,
     }
 }
 
-void StatsCache::set(ServiceContext* serviceContext, std::unique_ptr<StatsCache> cache) {
-    auto& statsCache = statsCacheDecoration(serviceContext);
-    invariant(!statsCache);
-
-    statsCache = std::move(cache);
-}
-
-void StatsCache::clearForTests(ServiceContext* serviceContext) {
-    auto& statsCache = statsCacheDecoration(serviceContext);
-    invariant(statsCache);
-
-    statsCache.reset();
-}
-
-StatsCache& StatsCache::get(ServiceContext* serviceContext) {
-    auto& statsCache = statsCacheDecoration(serviceContext);
-    invariant(statsCache);
-
-    return *statsCache;
-}
-
-StatsCache& StatsCache::get(OperationContext* opCtx) {
-    return get(opCtx->getServiceContext());
-}
 }  // namespace mongo
diff --git a/src/mongo/db/query/ce/stats_cache.h b/src/mongo/db/query/ce/stats_cache.h
index f9001adc736..65b5bcd19b4 100644
--- a/src/mongo/db/query/ce/stats_cache.h
+++ b/src/mongo/db/query/ce/stats_cache.h
@@ -45,21 +45,11 @@ using StatsCacheValueHandle = StatsCacheType::ValueHandle;
 
 /**
  * Collectoin statistics read through cache. It reads from the persitent storage but never wrties to
- * it. Stored on the service context.
+ * it.
  */
 class StatsCache : public StatsCacheType {
 public:
     /**
-     * Stores the cache on the specified service context. May only be called once for the lifetime
-     * of the service context.
-     */
-    static void set(ServiceContext* serviceContext, std::unique_ptr<StatsCache> cache);
-    static void clearForTests(ServiceContext* serviceContext);
-
-    static StatsCache& get(ServiceContext* serviceContext);
-    static StatsCache& get(OperationContext* opCtx);
-
-    /**
      * The constructor provides the Service context under which this cache has been instantiated,
      * and a Thread pool to be used for invoking the blocking 'lookup' calls. The size is the number
      * of entries the underlying LRU cache will hold.
diff --git a/src/mongo/db/query/ce/stats_catalog.cpp b/src/mongo/db/query/ce/stats_catalog.cpp
new file mode 100644
index 00000000000..8e37f65459d
--- /dev/null
+++ b/src/mongo/db/query/ce/stats_catalog.cpp
@@ -0,0 +1,107 @@
+/**
+ *    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/platform/basic.h"
+
+#include "mongo/db/query/ce/stats_cache.h"
+#include "mongo/db/query/ce/stats_catalog.h"
+
+#include "mongo/db/query/ce/array_histogram.h"
+#include "mongo/db/query/ce/collection_statistics.h"
+#include "mongo/util/read_through_cache.h"
+
+#include "mongo/logv2/log.h"
+
+#define MONGO_LOGV2_DEFAULT_COMPONENT ::mongo::logv2::LogComponent::kQuery
+
+namespace mongo {
+using namespace mongo::ce;
+
+namespace {
+
+const auto statsCatalogDecoration =
+    ServiceContext::declareDecoration<std::unique_ptr<StatsCatalog>>();
+
+}  // namespace
+
+StatsCatalog::StatsCatalog(ServiceContext* service,
+                           std::unique_ptr<StatsCacheLoader> statsCacheLoader)
+    : _executor(std::make_shared<ThreadPool>([] {
+          ThreadPool::Options options;
+          options.poolName = "StatsCache";
+          options.minThreads = 0;
+          options.maxThreads = 1;
+          return options;
+      }())),
+      _statsCache(service, std::move(statsCacheLoader), *_executor, 1000) {
+    _executor->startup();
+}
+
+StatsCatalog::~StatsCatalog() {
+    // The executor is used by the StatsCatalog, so it must be joined, before this cache is
+    // destroyed, per the contract of ReadThroughCache.
+    _executor->shutdown();
+    _executor->join();
+}
+
+void StatsCatalog::set(ServiceContext* serviceContext, std::unique_ptr<StatsCatalog> cache) {
+    auto& statsCatalog = statsCatalogDecoration(serviceContext);
+    invariant(!statsCatalog);
+
+    statsCatalog = std::move(cache);
+}
+
+StatsCatalog& StatsCatalog::get(ServiceContext* serviceContext) {
+    auto& statsCatalog = statsCatalogDecoration(serviceContext);
+    invariant(statsCatalog);
+
+    return *statsCatalog;
+}
+
+StatsCatalog& StatsCatalog::get(OperationContext* opCtx) {
+    return get(opCtx->getServiceContext());
+}
+
+StatusWith<std::shared_ptr<ArrayHistogram>> StatsCatalog::getHistogram(OperationContext* opCtx,
+                                                                       const NamespaceString& nss,
+                                                                       const std::string& path) {
+    try {
+        auto handle = _statsCache.acquire(opCtx, std::make_pair(nss, path));
+        uassert(ErrorCodes::NamespaceNotFound,
+                str::stream() << "path " << nss << " : " << path << " not found",
+                handle);
+
+        return *(handle.get());
+    } catch (const DBException& ex) {
+        return ex.toStatus();
+    }
+}
+
+}  // namespace mongo
diff --git a/src/mongo/db/query/ce/stats_catalog.h b/src/mongo/db/query/ce/stats_catalog.h
new file mode 100644
index 00000000000..cd5a378d4c5
--- /dev/null
+++ b/src/mongo/db/query/ce/stats_catalog.h
@@ -0,0 +1,78 @@
+/**
+ *    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/base/string_data.h"
+#include "mongo/db/namespace_string.h"
+#include "mongo/db/query/ce/collection_statistics.h"
+#include "mongo/db/query/ce/stats_cache.h"
+#include "mongo/db/query/ce/stats_cache_loader.h"
+#include "mongo/util/concurrency/thread_pool.h"
+
+namespace mongo {
+
+using namespace mongo::ce;
+
+/**
+ * This class owns statsCache and manages executor lifetime.
+ */
+class StatsCatalog {
+public:
+    /**
+     * Stores the catalog on the specified service context. May only be called once for the lifetime
+     * of the service context.
+     */
+    static void set(ServiceContext* serviceContext, std::unique_ptr<StatsCatalog> catalog);
+
+    static StatsCatalog& get(ServiceContext* serviceContext);
+    static StatsCatalog& get(OperationContext* opCtx);
+
+    /**
+     * The constructor provides the Service context under which the cache needs to be instantiated,
+     * and a Thread pool to be used for invoking the blocking 'lookup' calls. The size is the number
+     * of entries the underlying LRU cache will hold.
+     */
+    StatsCatalog(ServiceContext* service, std::unique_ptr<StatsCacheLoader> cacheLoader);
+
+    ~StatsCatalog();
+
+    StatusWith<std::shared_ptr<ArrayHistogram>> getHistogram(OperationContext* opCtx,
+                                                             const NamespaceString& nss,
+                                                             const std::string& path);
+
+private:
+    /**
+     * The executor is used by the cache.
+     */
+    std::shared_ptr<ThreadPool> _executor;
+    StatsCache _statsCache;
+};
+
+}  // namespace mongo
diff --git a/src/mongo/db/query/cqf_get_executor.cpp b/src/mongo/db/query/cqf_get_executor.cpp
index cbe7b9afb1f..311a0e17334 100644
--- a/src/mongo/db/query/cqf_get_executor.cpp
+++ b/src/mongo/db/query/cqf_get_executor.cpp
@@ -539,7 +539,6 @@ static OptPhaseManager createPhaseManager(const CEMode mode,
         }
 
         case CEMode::kHistogram:
-            // TODO SERVER-70352: Take into account statistics cache (maybe fallback to heuristic).
             return {OptPhaseManager::getAllRewritesSet(),
                     prefixId,
                     requireRID,
diff --git a/src/mongo/db/query/explain.h b/src/mongo/db/query/explain.h
index 1dcabdeb7e3..6892f3982ba 100644
--- a/src/mongo/db/query/explain.h
+++ b/src/mongo/db/query/explain.h
@@ -140,8 +140,8 @@ public:
      * intended to be human readable, and useful for debugging query performance problems related to
      * the plan cache.
      */
-    static void planCacheEntryToBSON(const PlanCacheEntry& entry, BSONObjBuilder* out);
-    static void planCacheEntryToBSON(const sbe::PlanCacheEntry& entry, BSONObjBuilder* out);
+    static void planCacheEntryToBSON(const mongo::PlanCacheEntry& entry, BSONObjBuilder* out);
+    static void planCacheEntryToBSON(const mongo::sbe::PlanCacheEntry& entry, BSONObjBuilder* out);
 };
 
 }  // namespace mongo