SERVER-64057 Implement a scoped object for setting the expected shard/database versions

1 files changed, 48 insertions, 14 deletions

diff --git a/src/mongo/db/s/operation_sharding_state.h b/src/mongo/db/s/operation_sharding_state.h
index 2a39b630658..63c726fcc0b 100644
--- a/src/mongo/db/s/operation_sharding_state.h
+++ b/src/mongo/db/s/operation_sharding_state.h
@@ -41,6 +41,28 @@
 namespace mongo {
 
 /**
+ * Marks the opCtx during scope in which it has been instantiated as running in the shard role for
+ * the specified collection. This indicates to the underlying storage system that the caller has
+ * performed 'routing', in the sense that it is aware of what data is located on this node.
+ */
+class ScopedSetShardRole {
+public:
+    ScopedSetShardRole(OperationContext* opCtx,
+                       NamespaceString nss,
+                       boost::optional<ChunkVersion> shardVersion,
+                       boost::optional<DatabaseVersion> databaseVersion);
+    ~ScopedSetShardRole();
+
+private:
+    OperationContext* const _opCtx;
+
+    NamespaceString _nss;
+
+    boost::optional<ChunkVersion> _shardVersion;
+    boost::optional<DatabaseVersion> _databaseVersion;
+};
+
+/**
  * A decoration on OperationContext representing per-operation shard version metadata sent to mongod
  * from mongos as a command parameter.
  *
@@ -86,13 +108,13 @@ public:
     };
 
     /**
-     * Stores the given shardVersion and databaseVersion for the given namespace. Note: The shard
-     * version for the given namespace stored in the OperationShardingState can be overwritten if it
-     * has not been checked yet.
+     * Same semantics as ScopedSetShardRole above, but the role remains set for the lifetime of the
+     * opCtx.
      */
-    void initializeClientRoutingVersions(NamespaceString nss,
-                                         const boost::optional<ChunkVersion>& shardVersion,
-                                         const boost::optional<DatabaseVersion>& dbVersion);
+    static void setShardRole(OperationContext* opCtx,
+                             const NamespaceString& nss,
+                             const boost::optional<ChunkVersion>& shardVersion,
+                             const boost::optional<DatabaseVersion>& dbVersion);
 
     /**
      * Removes the databaseVersion stored for the given namespace.
@@ -173,19 +195,31 @@ public:
     boost::optional<Status> resetShardingOperationFailedStatus();
 
 private:
+    friend class ScopedSetShardRole;
     friend class ShardServerOpObserver;  // For access to _allowCollectionCreation below
 
     // Specifies whether the request is allowed to create database/collection implicitly
     bool _allowCollectionCreation{false};
 
-    // The OperationShardingState class supports storing shardVersions for multiple namespaces (and
-    // databaseVersions for multiple databases), even though client code has not been written yet to
-    // *send* multiple shardVersions or databaseVersions.
-    StringMap<ChunkVersion> _shardVersions;
-    StringMap<DatabaseVersion> _databaseVersions;
-
-    // Stores shards that have undergone a version check.
-    StringSet _shardVersionsChecked;
+    // Stores the shard version expected for each collection that will be accessed
+    struct ShardVersionTracker {
+        ShardVersionTracker(ChunkVersion v) : v(v) {}
+        ShardVersionTracker(ShardVersionTracker&&) = default;
+        ShardVersionTracker(const ShardVersionTracker&) = delete;
+        ChunkVersion v;
+        int recursion{0};
+    };
+    StringMap<ShardVersionTracker> _shardVersions;
+
+    // Stores the database version expected for each database that will be accessed
+    struct DatabaseVersionTracker {
+        DatabaseVersionTracker(DatabaseVersion v) : v(v) {}
+        DatabaseVersionTracker(DatabaseVersionTracker&&) = default;
+        DatabaseVersionTracker(const DatabaseVersionTracker&) = delete;
+        DatabaseVersion v;
+        int recursion{0};
+    };
+    StringMap<DatabaseVersionTracker> _databaseVersions;
 
     // This value will only be non-null if version check during the operation execution failed due
     // to stale version and there was a migration for that namespace, which was in critical section.