SERVER-65085 SBE plan cache entries can be incorrectly reused after a refineCollectionShardKey operation

3 files changed, 34 insertions, 14 deletions

diff --git a/src/mongo/db/catalog/SConscript b/src/mongo/db/catalog/SConscript
index 9be42bf48c1..b116e2ad284 100644
--- a/src/mongo/db/catalog/SConscript
+++ b/src/mongo/db/catalog/SConscript
@@ -546,6 +546,7 @@ env.Library(
         '$BUILD_DIR/mongo/db/concurrency/lock_manager',
         '$BUILD_DIR/mongo/db/curop',
         '$BUILD_DIR/mongo/db/fts/base_fts',
+        '$BUILD_DIR/mongo/db/s/sharding_api_d',
         '$BUILD_DIR/mongo/db/service_context',
         'index_catalog',
     ],
diff --git a/src/mongo/db/query/plan_cache_key_factory.cpp b/src/mongo/db/query/plan_cache_key_factory.cpp
index 01d6a8ef934..d47f1768858 100644
--- a/src/mongo/db/query/plan_cache_key_factory.cpp
+++ b/src/mongo/db/query/plan_cache_key_factory.cpp
@@ -31,6 +31,7 @@
 
 #include "mongo/db/query/collection_query_info.h"
 #include "mongo/db/query/planner_ixselect.h"
+#include "mongo/db/s/operation_sharding_state.h"
 
 namespace mongo {
 namespace plan_cache_detail {
@@ -135,12 +136,17 @@ sbe::PlanCacheKey make(const CanonicalQuery& query,
                        PlanCacheKeyTag<sbe::PlanCacheKey>) {
     OperationContext* opCtx = query.getOpCtx();
     auto collectionVersion = CollectionQueryInfo::get(collection).getPlanCacheInvalidatorVersion();
+    const auto shardVersion{OperationShardingState::get(opCtx).getShardVersion(collection->ns())};
+    const auto keyShardingEpoch = shardVersion
+        ? boost::make_optional(
+              sbe::PlanCacheKeyShardingEpoch{shardVersion->epoch(), shardVersion->getTimestamp()})
+        : boost::none;
 
     return {makePlanCacheKeyInfo(query, collection),
             collection->uuid(),
             collectionVersion,
             computeNewestVisibleIndexTimestamp(opCtx, collection),
-            collection.isSharded()};
+            keyShardingEpoch};
 }
 }  // namespace plan_cache_detail
 }  // namespace mongo
diff --git a/src/mongo/db/query/sbe_plan_cache.h b/src/mongo/db/query/sbe_plan_cache.h
index 3e16371fcbd..227dc104321 100644
--- a/src/mongo/db/query/sbe_plan_cache.h
+++ b/src/mongo/db/query/sbe_plan_cache.h
@@ -43,6 +43,20 @@ namespace mongo {
 namespace sbe {
 
 /**
+ * Represents the sharding epoch of the collection which entry will be stored in the plan cache. The
+ * sharding epoch is not updated on operations like chunk splits and moves but rather on sharding
+ * and refine shard key operations.
+ */
+struct PlanCacheKeyShardingEpoch {
+    bool operator==(const PlanCacheKeyShardingEpoch& other) const {
+        return epoch == other.epoch && ts == other.ts;
+    }
+
+    OID epoch;
+    Timestamp ts;
+};
+
+/**
  * Represents the "key" used in the PlanCache mapping from query shape -> query plan.
  */
 class PlanCacheKey {
@@ -51,12 +65,12 @@ public:
                  UUID collectionUuid,
                  size_t collectionVersion,
                  boost::optional<Timestamp> newestVisibleIndexTimestamp,
-                 bool isShardedCollection)
+                 boost::optional<PlanCacheKeyShardingEpoch> shardVersion)
         : _info{std::move(info)},
           _collectionUuid{collectionUuid},
           _collectionVersion{collectionVersion},
           _newestVisibleIndexTimestamp{newestVisibleIndexTimestamp},
-          _isShardedCollection{isShardedCollection} {}
+          _shardVersion{shardVersion} {}
 
     const UUID& getCollectionUuid() const {
         return _collectionUuid;
@@ -66,16 +80,11 @@ public:
         return _collectionVersion;
     }
 
-    bool isShardedCollection() const {
-        return _isShardedCollection;
-    }
-
     bool operator==(const PlanCacheKey& other) const {
         return other._collectionVersion == _collectionVersion &&
-            other._isShardedCollection == _isShardedCollection &&
             other._collectionUuid == _collectionUuid &&
             other._newestVisibleIndexTimestamp == _newestVisibleIndexTimestamp &&
-            other._info == _info;
+            other._info == _info && other._shardVersion == _shardVersion;
     }
 
     bool operator!=(const PlanCacheKey& other) const {
@@ -93,7 +102,10 @@ public:
         if (_newestVisibleIndexTimestamp) {
             boost::hash_combine(hash, _newestVisibleIndexTimestamp->asULL());
         }
-        boost::hash_combine(hash, _isShardedCollection);
+        if (_shardVersion) {
+            _shardVersion->epoch.hash_combine(hash);
+            boost::hash_combine(hash, _shardVersion->ts.asULL());
+        }
         return hash;
     }
 
@@ -131,10 +143,11 @@ private:
     // In the future, this could instead be solved with point-in-time catalog lookups.
     const boost::optional<Timestamp> _newestVisibleIndexTimestamp;
 
-    // Ensures that a cached SBE plan cannot be reused if the collection has since become sharded.
-    // The cached plan may no longer be valid after sharding, since orphan filtering is necessary
-    // for sharded collections.
-    const bool _isShardedCollection;
+    // Ensures that a cached SBE plan cannot be reused if the collection has since become sharded or
+    // changed its shard key. The cached plan may no longer be valid after sharding or shard key
+    // refining since the structure of the plan depends on whether the collection is sharded, and if
+    // sharded depends on the shard key.
+    const boost::optional<PlanCacheKeyShardingEpoch> _shardVersion;
 };
 
 class PlanCacheKeyHasher {