SERVER-59333 Add simple implementation of data structures that form the key and value types for the sbe::PlanCache

7 files changed, 211 insertions, 15 deletions

diff --git a/src/mongo/db/SConscript b/src/mongo/db/SConscript
index 5beccabd31f..a128a885280 100644
--- a/src/mongo/db/SConscript
+++ b/src/mongo/db/SConscript
@@ -1267,6 +1267,7 @@ env.Library(
         'query/sbe_cached_solution_planner.cpp',
         'query/sbe_multi_planner.cpp',
         'query/sbe_plan_ranker.cpp',
+        'query/sbe_plan_cache.cpp',
         'query/sbe_runtime_planner.cpp',
         'query/sbe_stage_builder.cpp',
         'query/sbe_stage_builder_accumulator.cpp',
diff --git a/src/mongo/db/exec/sbe/expressions/expression.cpp b/src/mongo/db/exec/sbe/expressions/expression.cpp
index 52b8a6a3b42..d939053b9fb 100644
--- a/src/mongo/db/exec/sbe/expressions/expression.cpp
+++ b/src/mongo/db/exec/sbe/expressions/expression.cpp
@@ -931,14 +931,17 @@ RuntimeEnvironment::Accessor* RuntimeEnvironment::getAccessor(value::SlotId slot
     uasserted(4946301, str::stream() << "undefined slot accessor:" << slot);
 }
 
-std::unique_ptr<RuntimeEnvironment> RuntimeEnvironment::makeCopy(bool isSmp) {
+
+std::unique_ptr<RuntimeEnvironment> RuntimeEnvironment::makeCopy() const {
+    return std::unique_ptr<RuntimeEnvironment>(new RuntimeEnvironment(*this));
+}
+
+std::unique_ptr<RuntimeEnvironment> RuntimeEnvironment::makeCopyForParallelUse() {
     // Once this environment is used to create a copy for a parallel plan execution, it becomes
     // a parallel environment itself.
-    if (isSmp) {
-        _isSmp = isSmp;
-    }
+    _isSmp = true;
 
-    return std::unique_ptr<RuntimeEnvironment>(new RuntimeEnvironment(*this));
+    return makeCopy();
 }
 
 void RuntimeEnvironment::debugString(StringBuilder* builder) {
@@ -995,8 +998,12 @@ void CompileCtx::popCorrelated() {
     correlated.pop_back();
 }
 
-CompileCtx CompileCtx::makeCopy(bool isSmp) {
-    return {env->makeCopy(isSmp)};
+CompileCtx CompileCtx::makeCopyForParallelUse() {
+    return {env->makeCopyForParallelUse()};
+}
+
+CompileCtx CompileCtx::makeCopy() const {
+    return {env->makeCopy()};
 }
 }  // namespace sbe
 }  // namespace mongo
diff --git a/src/mongo/db/exec/sbe/expressions/expression.h b/src/mongo/db/exec/sbe/expressions/expression.h
index 77728b81965..8ecb3ecf737 100644
--- a/src/mongo/db/exec/sbe/expressions/expression.h
+++ b/src/mongo/db/exec/sbe/expressions/expression.h
@@ -156,14 +156,15 @@ public:
     Accessor* getAccessor(value::SlotId slot);
 
     /**
-     * Make a copy of his environment. The new environment will have its own set of SlotAccessors
+     * Make a copy of this environment. The new environment will have its own set of SlotAccessors
      * pointing to the same shared data holding slot values.
      *
-     * To create a copy of the runtime environment for a parallel execution plan, the 'isSmp' flag
-     * must be set to 'true'. This will result in this environment being unconverted to a parallel
-     * environment, as well as the newly created copy.
+     * To create a copy of the runtime environment for a parallel execution plan, please use
+     * makeCopyForParallelUse() method. This will result in this environment being converted to a
+     * parallel environment, as well as the newly created copy.
      */
-    std::unique_ptr<RuntimeEnvironment> makeCopy(bool isSmp);
+    std::unique_ptr<RuntimeEnvironment> makeCopyForParallelUse();
+    std::unique_ptr<RuntimeEnvironment> makeCopy() const;
 
     /**
      * Dumps all the slots currently defined in this environment into the given string builder.
@@ -226,7 +227,15 @@ struct CompileCtx {
     void pushCorrelated(value::SlotId slot, value::SlotAccessor* accessor);
     void popCorrelated();
 
-    CompileCtx makeCopy(bool isSmp);
+    /**
+     * Make a copy of this CompileCtx. The underlying RuntimeEnvironment will also be copied.
+     *
+     * To create a copy of the underlying runtime environment for a parallel execution plan, please
+     * use makeCopyForParallelUse() method. This will result in the environment in this CompileCtx
+     * being converted to a parallel environment, as well as the newly created copy.
+     */
+    CompileCtx makeCopyForParallelUse();
+    CompileCtx makeCopy() const;
 
     PlanStage* root{nullptr};
     value::SlotAccessor* accumulator{nullptr};
diff --git a/src/mongo/db/exec/sbe/stages/exchange.cpp b/src/mongo/db/exec/sbe/stages/exchange.cpp
index 2be1c108aa2..baeba0fdc90 100644
--- a/src/mongo/db/exec/sbe/stages/exchange.cpp
+++ b/src/mongo/db/exec/sbe/stages/exchange.cpp
@@ -191,7 +191,7 @@ void ExchangeConsumer::prepare(CompileCtx& ctx) {
         // Only consumer ID 0 prepares (copies) the compilation context. Note that we do not have to
         // lock the shared state here - it is accessed only by consmer ID 0 again in the future.
         for (size_t idx = 0; idx < _state->numOfProducers(); ++idx) {
-            _state->producerCompileCtxs().push_back(ctx.makeCopy(true));
+            _state->producerCompileCtxs().push_back(ctx.makeCopyForParallelUse());
         }
     }
 
diff --git a/src/mongo/db/query/sbe_plan_cache.cpp b/src/mongo/db/query/sbe_plan_cache.cpp
new file mode 100644
index 00000000000..1ad8c0b1094
--- /dev/null
+++ b/src/mongo/db/query/sbe_plan_cache.cpp
@@ -0,0 +1,44 @@
+/**
+ *    Copyright (C) 2021-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/db/query/sbe_plan_cache.h"
+
+namespace mongo {
+namespace sbe {
+
+uint32_t PlanCacheKey::queryHash() const {
+    return static_cast<uint32_t>(PlanCacheKeyHasher{}(*this));
+}
+
+uint32_t PlanCacheKey::planCacheKeyHash() const {
+    return static_cast<uint32_t>(PlanCacheKeyHasher{}(*this));
+}
+
+}  // namespace sbe
+}  // namespace mongo
diff --git a/src/mongo/db/query/sbe_plan_cache.h b/src/mongo/db/query/sbe_plan_cache.h
new file mode 100644
index 00000000000..6c7e7c5956c
--- /dev/null
+++ b/src/mongo/db/query/sbe_plan_cache.h
@@ -0,0 +1,103 @@
+/**
+ *    Copyright (C) 2021-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/exec/sbe/stages/stages.h"
+#include "mongo/db/hasher.h"
+#include "mongo/db/operation_context.h"
+#include "mongo/db/query/plan_cache.h"
+#include "mongo/db/query/sbe_stage_builder.h"
+#include "mongo/db/service_context.h"
+
+namespace mongo {
+namespace sbe {
+
+/**
+ * Represents the key used to look up entries in the SBE PlanCache.
+ */
+class PlanCacheKey {
+public:
+    explicit PlanCacheKey(BSONObj filter) : _filter(filter.getOwned()) {}
+
+    bool operator==(const PlanCacheKey& other) const {
+        return other._filter.binaryEqual(_filter);
+    }
+
+    bool operator!=(const PlanCacheKey& other) const {
+        return !(*this == other);
+    }
+
+    const BSONObj& getFilter() const {
+        return _filter;
+    }
+
+    uint32_t queryHash() const;
+
+    uint32_t planCacheKeyHash() const;
+
+private:
+    const BSONObj _filter;
+};
+
+/**
+ * Provides hash function to hash a 'PlanCacheKey'.
+ */
+class PlanCacheKeyHasher {
+public:
+    std::size_t operator()(const PlanCacheKey& key) const {
+        return SimpleBSONObjComparator::kInstance.hash(key.getFilter());
+    }
+};
+
+/**
+ * Represents the data cached in the SBE plan cache. This data holds an execution plan and necessary
+ * auxiliary data for preparing and executing the PlanStage tree.
+ */
+struct CachedSbePlan {
+    CachedSbePlan(std::unique_ptr<sbe::PlanStage> root, stage_builder::PlanStageData data)
+        : root(std::move(root)), planStageData(std::move(data)) {}
+
+    std::unique_ptr<CachedSbePlan> clone() const {
+        return std::make_unique<CachedSbePlan>(root->clone(), planStageData);
+    }
+
+    uint64_t estimateObjectSizeInBytes() const {
+        // TODO SERVER-59331: handle size estimation.
+        return 0;
+    }
+
+    std::unique_ptr<sbe::PlanStage> root;
+    stage_builder::PlanStageData planStageData;
+};
+
+using PlanCache = PlanCacheBase<sbe::PlanCacheKey, CachedSbePlan, sbe::PlanCacheKeyHasher>;
+
+}  // namespace sbe
+}  // namespace mongo
diff --git a/src/mongo/db/query/sbe_stage_builder.h b/src/mongo/db/query/sbe_stage_builder.h
index f9dd9a79389..ba0d1ae5b27 100644
--- a/src/mongo/db/query/sbe_stage_builder.h
+++ b/src/mongo/db/query/sbe_stage_builder.h
@@ -222,11 +222,26 @@ void PlanStageSlots::forEachSlot(const PlanStageReqs& reqs,
  * is needed to execute the PlanStage tree.
  */
 struct PlanStageData {
-    PlanStageData() = default;
+    PlanStageData(PlanStageData&&) = default;
+    PlanStageData& operator=(PlanStageData&&) = default;
 
     explicit PlanStageData(std::unique_ptr<sbe::RuntimeEnvironment> env)
         : env(env.get()), ctx(std::move(env)) {}
 
+    PlanStageData(const PlanStageData& other) : PlanStageData(other.env->makeCopy()) {
+        copyFrom(other);
+    }
+
+    PlanStageData& operator=(const PlanStageData& other) {
+        if (this != &other) {
+            auto envCopy = other.env->makeCopy();
+            env = envCopy.get();
+            ctx = sbe::CompileCtx(std::move(envCopy));
+            copyFrom(other);
+        }
+        return *this;
+    }
+
     std::string debugString() const;
 
     // This holds the output slots produced by SBE plan (resultSlot, recordIdSlot, etc).
@@ -251,6 +266,23 @@ struct PlanStageData {
     // If this candidate plan has completed the trial run early by achieving one of the trial run
     // metrics, the stats are cached in here.
     std::unique_ptr<sbe::PlanStageStats> savedStatsOnEarlyExit{nullptr};
+
+private:
+    // This copy function copies data from 'other' but will not create a copy of its
+    // RuntimeEnvironment and CompileCtx.
+    void copyFrom(const PlanStageData& other) {
+        outputs = other.outputs;
+        iamMap = other.iamMap;
+        shouldTrackLatestOplogTimestamp = other.shouldTrackLatestOplogTimestamp;
+        shouldTrackResumeToken = other.shouldTrackResumeToken;
+        shouldUseTailableScan = other.shouldUseTailableScan;
+        replanReason = other.replanReason;
+        if (other.savedStatsOnEarlyExit) {
+            savedStatsOnEarlyExit.reset(other.savedStatsOnEarlyExit->clone());
+        } else {
+            savedStatsOnEarlyExit.reset();
+        }
+    }
 };
 
 /**