SERVER-67996 implement collection statistics read through cache

7 files changed, 438 insertions, 1 deletions

diff --git a/src/mongo/db/query/ce/SConscript b/src/mongo/db/query/ce/SConscript
index 6926b702482..fac1c5a01ab 100644
--- a/src/mongo/db/query/ce/SConscript
+++ b/src/mongo/db/query/ce/SConscript
@@ -13,6 +13,8 @@ env.Library(
         'collection_statistics.cpp',
         'histogram_estimation.cpp',
         'scalar_histogram.cpp',
+        'stats_cache.cpp',
+        'stats_cache_loader.cpp',
     ],
     LIBDEPS_PRIVATE=[
         '$BUILD_DIR/mongo/db/exec/sbe/query_sbe_abt',
@@ -54,3 +56,16 @@ env.CppUnitTest(
         'ce_test_utils',
     ],
 )
+
+env.CppUnitTest(
+    target="stats_cache_test",
+    source=[
+        "stats_cache_test.cpp",
+    ],
+    LIBDEPS=[
+        '$BUILD_DIR/mongo/base',
+        '$BUILD_DIR/mongo/db/service_context',
+        'ce_test_utils',
+        'query_ce',
+    ],
+)
diff --git a/src/mongo/db/query/ce/collection_statistics.h b/src/mongo/db/query/ce/collection_statistics.h
index 4ac133129cf..e92c8c05eef 100644
--- a/src/mongo/db/query/ce/collection_statistics.h
+++ b/src/mongo/db/query/ce/collection_statistics.h
@@ -34,7 +34,7 @@
 
 namespace mongo::ce {
 
-using Histograms = std::map<std::string, std::unique_ptr<ArrayHistogram>>;
+using Histograms = std::map<std::string, std::shared_ptr<ArrayHistogram>>;
 
 class CollectionStatistics {
 public:
diff --git a/src/mongo/db/query/ce/stats_cache.cpp b/src/mongo/db/query/ce/stats_cache.cpp
new file mode 100644
index 00000000000..ee8fc079df4
--- /dev/null
+++ b/src/mongo/db/query/ce/stats_cache.cpp
@@ -0,0 +1,100 @@
+/**
+ *    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/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 statsCacheDecoration = ServiceContext::declareDecoration<std::unique_ptr<StatsCache>>();
+
+}  // namespace
+
+StatsCache::StatsCache(ServiceContext* service, ThreadPoolInterface& threadPool, int size)
+    : ReadThroughCache(_mutex,
+                       service,
+                       threadPool,
+                       [this](OperationContext* opCtx,
+                              const NamespaceString& nss,
+                              const ValueHandle& stats) { return _lookupStats(opCtx, nss, stats); },
+                       size) {}
+
+StatsCache::LookupResult StatsCache::_lookupStats(OperationContext* opCtx,
+                                                  const NamespaceString& nss,
+                                                  const StatsCacheValueHandle& stats) {
+
+    try {
+        auto newStats = _statsCacheLoader.getStats(nss).get();
+        return LookupResult(std::move(newStats));
+    } catch (const DBException& ex) {
+        if (ex.code() == ErrorCodes::NamespaceNotFound) {
+            return StatsCache::LookupResult(boost::none);
+        }
+        throw;
+    }
+}
+
+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
new file mode 100644
index 00000000000..31bf6802b5c
--- /dev/null
+++ b/src/mongo/db/query/ce/stats_cache.h
@@ -0,0 +1,89 @@
+/**
+ *    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_loader.h"
+#include "mongo/util/concurrency/thread_pool.h"
+#include "mongo/util/read_through_cache.h"
+
+namespace mongo {
+
+using namespace mongo::ce;
+
+using StatsCacheType = ReadThroughCache<NamespaceString, CollectionStatistics>;
+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.
+ */
+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.
+     */
+    StatsCache(ServiceContext* service, ThreadPoolInterface& threadPool, int size);
+
+    /**
+     *  Returns statsCacheLoader currently used for testing only.
+     */
+    StatsCacheLoader& getStatsCacheLoader() {
+        return _statsCacheLoader;
+    }
+
+private:
+    /**
+     * Reads collection stats from the underlying storage if its not found in the in memory cache.
+     */
+    LookupResult _lookupStats(OperationContext* opCtx,
+                              const NamespaceString& nss,
+                              const ValueHandle& stats);
+
+    Mutex _mutex = MONGO_MAKE_LATCH("StatsCache::_mutex");
+
+    StatsCacheLoader _statsCacheLoader;
+};
+
+}  // namespace mongo
diff --git a/src/mongo/db/query/ce/stats_cache_loader.cpp b/src/mongo/db/query/ce/stats_cache_loader.cpp
new file mode 100644
index 00000000000..0a2b265b25f
--- /dev/null
+++ b/src/mongo/db/query/ce/stats_cache_loader.cpp
@@ -0,0 +1,50 @@
+/**
+ *    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_loader.h"
+
+#include "mongo/db/query/ce/collection_statistics.h"
+#include "mongo/stdx/thread.h"
+
+namespace mongo {
+
+const Status StatsCacheLoader::kInternalErrorStatus = {
+    ErrorCodes::InternalError, "Stats cache loader received unexpected request"};
+
+SemiFuture<CollectionStatistics> StatsCacheLoader::getStats(const NamespaceString& nss) {
+    return makeReadyFutureWith([this] { return _swStatsReturnValueForTest; }).semi();
+}
+
+void StatsCacheLoader::setStatsReturnValueForTest(StatusWith<CollectionStatistics> swStats) {
+    _swStatsReturnValueForTest = std::move(swStats);
+}
+}  // namespace mongo
diff --git a/src/mongo/db/query/ce/stats_cache_loader.h b/src/mongo/db/query/ce/stats_cache_loader.h
new file mode 100644
index 00000000000..cdabd4acffb
--- /dev/null
+++ b/src/mongo/db/query/ce/stats_cache_loader.h
@@ -0,0 +1,58 @@
+/**
+ *    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/db/namespace_string.h"
+#include "mongo/db/query/ce/collection_statistics.h"
+#include "mongo/stdx/thread.h"
+
+namespace mongo {
+
+using namespace mongo::ce;
+
+class StatsCacheLoader {
+public:
+    /**
+     * Non-blocking call, which returns CollectionStatistics from the the persistent metadata store.
+     *
+     * If for some reason the asynchronous fetch operation cannot be dispatched (for example on
+     * shutdown), throws a DBException.
+     */
+    // TODO: SERVER-68459 implement statsCacheLoader
+    SemiFuture<CollectionStatistics> getStats(const NamespaceString& nss);
+
+    void setStatsReturnValueForTest(StatusWith<CollectionStatistics> swStats);
+    static const Status kInternalErrorStatus;
+
+private:
+    StatusWith<CollectionStatistics> _swStatsReturnValueForTest{kInternalErrorStatus};
+};
+
+}  // namespace mongo
diff --git a/src/mongo/db/query/ce/stats_cache_test.cpp b/src/mongo/db/query/ce/stats_cache_test.cpp
new file mode 100644
index 00000000000..6111ea02dae
--- /dev/null
+++ b/src/mongo/db/query/ce/stats_cache_test.cpp
@@ -0,0 +1,125 @@
+/**
+ *    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 <string>
+
+#include "mongo/db/client.h"
+#include "mongo/db/concurrency/locker_noop_service_context_test_fixture.h"
+#include "mongo/db/operation_context.h"
+#include "mongo/db/query/ce/stats_cache.h"
+#include "mongo/unittest/barrier.h"
+#include "mongo/unittest/unittest.h"
+#include "mongo/util/concurrency/thread_pool.h"
+#include "mongo/util/read_through_cache.h"
+#include "mongo/util/scopeguard.h"
+
+#define MONGO_LOGV2_DEFAULT_COMPONENT ::mongo::logv2::LogComponent::kDefault
+
+namespace mongo {
+namespace {
+
+using unittest::assertGet;
+
+/**
+ * Fixture for tests, which do not need to exercise the multi-threading capabilities of the cache
+ * and as such do not require control over the creation/destruction of their operation contexts.
+ */
+class StatsCacheTest : public LockerNoopServiceContextTest {
+protected:
+    // Extends StatsCache and automatically provides it with a thread  pool, which will be
+    // shutdown and joined before the StatsCache is destroyed (which is part of the  contract of
+    // ReadThroughCache)
+    class CacheWithThreadPool : public StatsCache {
+    public:
+        CacheWithThreadPool(ServiceContext* service, size_t size)
+            : StatsCache(service, _threadPool, size) {
+            _threadPool.startup();
+        }
+
+    private:
+        ThreadPool _threadPool{[] {
+            ThreadPool::Options options;
+            options.poolName = "StatsCacheTest";
+            options.minThreads = 1;
+            options.maxThreads = 1;
+            return options;
+        }()};
+    };
+
+    const ServiceContext::UniqueOperationContext _opCtxHolder{makeOperationContext()};
+    OperationContext* const _opCtx{_opCtxHolder.get()};
+};
+
+TEST(StatsCacheTest, StandaloneValueHandle) {
+    StatsCache::ValueHandle standaloneHandle(CollectionStatistics(100));
+    ASSERT(standaloneHandle.isValid());
+    ASSERT_EQ(100, standaloneHandle->getCardinality());
+}
+
+TEST_F(StatsCacheTest, KeyDoesNotExist) {
+    Status namespaceNotFoundErrorStatus = {ErrorCodes::NamespaceNotFound,
+                                           "The key does not exists"};
+    auto cache = CacheWithThreadPool(getServiceContext(), 1);
+    cache.getStatsCacheLoader().setStatsReturnValueForTest(std::move(namespaceNotFoundErrorStatus));
+    auto handle = cache.acquire(_opCtx, NamespaceString("db", "coll"));
+    ASSERT(!handle);
+}
+
+TEST_F(StatsCacheTest, LoadStats) {
+    auto cache = CacheWithThreadPool(getServiceContext(), 1);
+
+    auto stats1 = CollectionStatistics(1);
+    auto stats2 = CollectionStatistics(2);
+
+    cache.getStatsCacheLoader().setStatsReturnValueForTest(std::move(stats1));
+
+    auto handle = cache.acquire(_opCtx, NamespaceString("db", "coll1"));
+    ASSERT(handle.isValid());
+    ASSERT_EQ(1, handle->getCardinality());
+
+    // Make all requests to StatsCacheLoader to throw an exception to ensre that test returns value
+    // from cache.
+    Status internalErrorStatus = {ErrorCodes::InternalError,
+                                  "Stats cache loader received unexpected request"};
+    cache.getStatsCacheLoader().setStatsReturnValueForTest(std::move(internalErrorStatus));
+
+    handle = cache.acquire(_opCtx, NamespaceString("db", "coll1"));
+    ASSERT(handle.isValid());
+    ASSERT_EQ(1, handle->getCardinality());
+
+    cache.getStatsCacheLoader().setStatsReturnValueForTest(std::move(stats2));
+    handle = cache.acquire(_opCtx, NamespaceString("db", "coll2"));
+    ASSERT(handle.isValid());
+    ASSERT_EQ(2, handle->getCardinality());
+}
+
+}  // namespace
+}  // namespace mongo