SERVER-50644, SERVER-50479: Add resumable index build support for ESE by using persistent key for Sorter temp file encryption

11 files changed, 109 insertions, 55 deletions

diff --git a/buildscripts/resmokeconfig/suites/replica_sets_ese.yml b/buildscripts/resmokeconfig/suites/replica_sets_ese.yml
index 9ced87dae15..3e202ae3c1b 100644
--- a/buildscripts/resmokeconfig/suites/replica_sets_ese.yml
+++ b/buildscripts/resmokeconfig/suites/replica_sets_ese.yml
@@ -7,11 +7,6 @@ test_kind: js_test
 selector:
   roots:
   - jstests/replsets/*.js
-  exclude_files:
-  # TODO (SERVER-50479): Enable resumable index build rollback tests once resumable index builds
-  # support ESE.
-  - jstests/replsets/rollback_resumable_index_build_*.js
-  - jstests/replsets/restart_index_build_if_resume_interrupted_by_rollback.js
 
 executor:
   config:
diff --git a/buildscripts/resmokeconfig/suites/replica_sets_ese_gcm.yml b/buildscripts/resmokeconfig/suites/replica_sets_ese_gcm.yml
index 74ccdaabb5d..b44d14eb841 100644
--- a/buildscripts/resmokeconfig/suites/replica_sets_ese_gcm.yml
+++ b/buildscripts/resmokeconfig/suites/replica_sets_ese_gcm.yml
@@ -7,11 +7,6 @@ test_kind: js_test
 selector:
   roots:
   - jstests/replsets/*.js
-  exclude_files:
-  # TODO (SERVER-50479): Enable resumable index build rollback tests once resumable index builds
-  # support ESE.
-  - jstests/replsets/rollback_resumable_index_build_*.js
-  - jstests/replsets/restart_index_build_if_resume_interrupted_by_rollback.js
 
 executor:
   config:
diff --git a/jstests/noPassthrough/libs/index_build.js b/jstests/noPassthrough/libs/index_build.js
index 26e08901999..b3ef206858d 100644
--- a/jstests/noPassthrough/libs/index_build.js
+++ b/jstests/noPassthrough/libs/index_build.js
@@ -447,7 +447,8 @@ const ResumableIndexBuildTest = class {
                    failPointsIteration,
                    shouldComplete = true,
                    failPointAfterStartup,
-                   runBeforeStartup) {
+                   runBeforeStartup,
+                   options) {
         clearRawMongoProgramOutput();
 
         const buildUUIDs = ResumableIndexBuildTest.generateFailPointsData(
@@ -507,8 +508,8 @@ const ResumableIndexBuildTest = class {
                 ["failpoint." + failPointAfterStartup]: tojson({mode: "alwaysOn"}),
             });
         }
-
-        rst.start(conn, {noCleanData: true, setParameter: setParameter});
+        const defaultOptions = {noCleanData: true, setParameter: setParameter};
+        rst.start(conn, Object.assign(defaultOptions, options || {}));
 
         if (shouldComplete) {
             // Ensure that the index builds were completed upon the node starting back up.
@@ -629,7 +630,8 @@ const ResumableIndexBuildTest = class {
                expectedResumePhases,
                resumeChecks,
                sideWrites = [],
-               postIndexBuildInserts = []) {
+               postIndexBuildInserts = [],
+               restartOptions) {
         const primary = rst.getPrimary();
 
         if (!ResumableIndexBuildTest.resumableIndexBuildsEnabled(primary)) {
@@ -646,8 +648,16 @@ const ResumableIndexBuildTest = class {
                 ResumableIndexBuildTest.createIndexesFails(db, collName, indexSpecs, indexNames);
             }, coll, indexSpecs, indexNames, sideWrites, {hangBeforeBuildingIndex: true});
 
-        const buildUUIDs = ResumableIndexBuildTest.restart(
-            rst, primary, coll, indexNames, failPoints, failPointsIteration);
+        const buildUUIDs = ResumableIndexBuildTest.restart(rst,
+                                                           primary,
+                                                           coll,
+                                                           indexNames,
+                                                           failPoints,
+                                                           failPointsIteration,
+                                                           true,
+                                                           undefined,
+                                                           undefined,
+                                                           restartOptions);
 
         for (const awaitCreateIndex of awaitCreateIndexes) {
             awaitCreateIndex();
diff --git a/src/mongo/db/catalog/multi_index_block.cpp b/src/mongo/db/catalog/multi_index_block.cpp
index 59a242e6e55..f0c46930775 100644
--- a/src/mongo/db/catalog/multi_index_block.cpp
+++ b/src/mongo/db/catalog/multi_index_block.cpp
@@ -282,7 +282,8 @@ StatusWith<std::vector<BSONObj>> MultiIndexBlock::init(
             if (!status.isOK())
                 return status;
 
-            index.bulk = index.real->initiateBulk(eachIndexBuildMaxMemoryUsageBytes, stateInfo);
+            index.bulk = index.real->initiateBulk(
+                eachIndexBuildMaxMemoryUsageBytes, stateInfo, collection->ns().db());
 
             const IndexDescriptor* descriptor = indexCatalogEntry->descriptor();
 
@@ -1142,8 +1143,8 @@ Status MultiIndexBlock::_scanReferenceIdxInsertAndCommit(OperationContext* opCtx
     // comes to the child index. As a result, we need to sort each set of keys that differ only in
     // their record IDs. We're calling this set of keys a key class.
     auto refreshSorter = [&]() {
-        _indexes[0].bulk =
-            _indexes[0].real->initiateBulk(_eachIndexBuildMaxMemoryUsageBytes, boost::none);
+        _indexes[0].bulk = _indexes[0].real->initiateBulk(
+            _eachIndexBuildMaxMemoryUsageBytes, boost::none, collection->ns().db());
     };
 
     auto addToSorter = [&](const KeyString::Value& keyString) {
diff --git a/src/mongo/db/index/index_access_method.cpp b/src/mongo/db/index/index_access_method.cpp
index 2d340cbc997..4de61819671 100644
--- a/src/mongo/db/index/index_access_method.cpp
+++ b/src/mongo/db/index/index_access_method.cpp
@@ -84,11 +84,12 @@ bool isMultikeyFromPaths(const MultikeyPaths& multikeyPaths) {
                        [](const MultikeyComponents& components) { return !components.empty(); });
 }
 
-SortOptions makeSortOptions(size_t maxMemoryUsageBytes) {
+SortOptions makeSortOptions(size_t maxMemoryUsageBytes, StringData dbName) {
     return SortOptions()
         .TempDir(storageGlobalParams.dbpath + "/_tmp")
         .ExtSortAllowed()
-        .MaxMemoryUsageBytes(maxMemoryUsageBytes);
+        .MaxMemoryUsageBytes(maxMemoryUsageBytes)
+        .DBName(dbName.toString());
 }
 
 MultikeyPaths createMultikeyPaths(const std::vector<MultikeyPath>& multikeyPathsVec) {
@@ -479,11 +480,14 @@ Status AbstractIndexAccessMethod::compact(OperationContext* opCtx) {
 
 class AbstractIndexAccessMethod::BulkBuilderImpl : public IndexAccessMethod::BulkBuilder {
 public:
-    BulkBuilderImpl(IndexCatalogEntry* indexCatalogEntry, size_t maxMemoryUsageBytes);
+    BulkBuilderImpl(IndexCatalogEntry* indexCatalogEntry,
+                    size_t maxMemoryUsageBytes,
+                    StringData dbName);
 
     BulkBuilderImpl(IndexCatalogEntry* index,
                     size_t maxMemoryUsageBytes,
-                    const IndexStateInfo& stateInfo);
+                    const IndexStateInfo& stateInfo,
+                    StringData dbName);
 
     Status insert(OperationContext* opCtx,
                   const BSONObj& obj,
@@ -513,6 +517,7 @@ private:
 
     Sorter* _makeSorter(
         size_t maxMemoryUsageBytes,
+        StringData dbName,
         boost::optional<StringData> fileName = boost::none,
         const boost::optional<std::vector<SorterRange>>& ranges = boost::none) const;
 
@@ -536,21 +541,27 @@ private:
 };
 
 std::unique_ptr<IndexAccessMethod::BulkBuilder> AbstractIndexAccessMethod::initiateBulk(
-    size_t maxMemoryUsageBytes, const boost::optional<IndexStateInfo>& stateInfo) {
+    size_t maxMemoryUsageBytes,
+    const boost::optional<IndexStateInfo>& stateInfo,
+    StringData dbName) {
     return stateInfo
-        ? std::make_unique<BulkBuilderImpl>(_indexCatalogEntry, maxMemoryUsageBytes, *stateInfo)
-        : std::make_unique<BulkBuilderImpl>(_indexCatalogEntry, maxMemoryUsageBytes);
+        ? std::make_unique<BulkBuilderImpl>(
+              _indexCatalogEntry, maxMemoryUsageBytes, *stateInfo, dbName)
+        : std::make_unique<BulkBuilderImpl>(_indexCatalogEntry, maxMemoryUsageBytes, dbName);
 }
 
 AbstractIndexAccessMethod::BulkBuilderImpl::BulkBuilderImpl(IndexCatalogEntry* index,
-                                                            size_t maxMemoryUsageBytes)
-    : _indexCatalogEntry(index), _sorter(_makeSorter(maxMemoryUsageBytes)) {}
+                                                            size_t maxMemoryUsageBytes,
+                                                            StringData dbName)
+    : _indexCatalogEntry(index), _sorter(_makeSorter(maxMemoryUsageBytes, dbName)) {}
 
 AbstractIndexAccessMethod::BulkBuilderImpl::BulkBuilderImpl(IndexCatalogEntry* index,
                                                             size_t maxMemoryUsageBytes,
-                                                            const IndexStateInfo& stateInfo)
+                                                            const IndexStateInfo& stateInfo,
+                                                            StringData dbName)
     : _indexCatalogEntry(index),
-      _sorter(_makeSorter(maxMemoryUsageBytes, stateInfo.getFileName(), stateInfo.getRanges())),
+      _sorter(
+          _makeSorter(maxMemoryUsageBytes, dbName, stateInfo.getFileName(), stateInfo.getRanges())),
       _keysInserted(stateInfo.getNumKeys().value_or(0)),
       _isMultiKey(stateInfo.getIsMultikey()),
       _indexMultikeyPaths(createMultikeyPaths(stateInfo.getMultikeyPaths())) {}
@@ -663,14 +674,15 @@ AbstractIndexAccessMethod::BulkBuilderImpl::_makeSorterSettings() const {
 AbstractIndexAccessMethod::BulkBuilderImpl::Sorter*
 AbstractIndexAccessMethod::BulkBuilderImpl::_makeSorter(
     size_t maxMemoryUsageBytes,
+    StringData dbName,
     boost::optional<StringData> fileName,
     const boost::optional<std::vector<SorterRange>>& ranges) const {
     return fileName ? Sorter::makeFromExistingRanges(fileName->toString(),
                                                      *ranges,
-                                                     makeSortOptions(maxMemoryUsageBytes),
+                                                     makeSortOptions(maxMemoryUsageBytes, dbName),
                                                      BtreeExternalSortComparison(),
                                                      _makeSorterSettings())
-                    : Sorter::make(makeSortOptions(maxMemoryUsageBytes),
+                    : Sorter::make(makeSortOptions(maxMemoryUsageBytes, dbName),
                                    BtreeExternalSortComparison(),
                                    _makeSorterSettings());
 }
diff --git a/src/mongo/db/index/index_access_method.h b/src/mongo/db/index/index_access_method.h
index 5aee4fad6e4..4653fffa7a3 100644
--- a/src/mongo/db/index/index_access_method.h
+++ b/src/mongo/db/index/index_access_method.h
@@ -290,7 +290,9 @@ public:
      * new index build.
      */
     virtual std::unique_ptr<BulkBuilder> initiateBulk(
-        size_t maxMemoryUsageBytes, const boost::optional<IndexStateInfo>& stateInfo) = 0;
+        size_t maxMemoryUsageBytes,
+        const boost::optional<IndexStateInfo>& stateInfo,
+        StringData dbName) = 0;
 
     /**
      * Call this when you are ready to finish your bulk work.
@@ -541,8 +543,9 @@ public:
                             KeyStringSet multikeyMetadataKeys,
                             MultikeyPaths paths) final;
 
-    std::unique_ptr<BulkBuilder> initiateBulk(
-        size_t maxMemoryUsageBytes, const boost::optional<IndexStateInfo>& stateInfo) final;
+    std::unique_ptr<BulkBuilder> initiateBulk(size_t maxMemoryUsageBytes,
+                                              const boost::optional<IndexStateInfo>& stateInfo,
+                                              StringData dbName) final;
 
     Status commitBulk(OperationContext* opCtx,
                       BulkBuilder* bulk,
diff --git a/src/mongo/db/index_builds_coordinator.cpp b/src/mongo/db/index_builds_coordinator.cpp
index a02fa7c8460..6d0892f1ae6 100644
--- a/src/mongo/db/index_builds_coordinator.cpp
+++ b/src/mongo/db/index_builds_coordinator.cpp
@@ -393,11 +393,6 @@ bool isIndexBuildResumable(OperationContext* opCtx,
         return false;
     }
 
-    // TODO(SERVER-50479): Remove this check when resumable index builds work with ESE in GCM mode.
-    if (EncryptionHooks::get(opCtx->getServiceContext())->enabled()) {
-        return false;
-    }
-
     if (!opCtx->getServiceContext()->getStorageEngine()->supportsResumableIndexBuilds()) {
         return false;
     }
diff --git a/src/mongo/db/sorter/sorter.cpp b/src/mongo/db/sorter/sorter.cpp
index 4e56b733284..7ed19fbe2d8 100644
--- a/src/mongo/db/sorter/sorter.cpp
+++ b/src/mongo/db/sorter/sorter.cpp
@@ -197,12 +197,14 @@ public:
                  std::streampos fileStartOffset,
                  std::streampos fileEndOffset,
                  const Settings& settings,
+                 const boost::optional<std::string>& dbName,
                  const uint32_t checksum)
         : _settings(settings),
           _done(false),
           _fileFullPath(fileFullPath),
           _fileStartOffset(fileStartOffset),
           _fileEndOffset(fileEndOffset),
+          _dbName(dbName),
           _originalChecksum(checksum) {
         uassert(16815,
                 str::stream() << "unexpected empty file: " << _fileFullPath,
@@ -309,11 +311,12 @@ private:
             std::unique_ptr<char[]> out(new char[blockSize]);
             size_t outLen;
             Status status =
-                encryptionHooks->unprotectTmpData(reinterpret_cast<uint8_t*>(_buffer.get()),
+                encryptionHooks->unprotectTmpData(reinterpret_cast<const uint8_t*>(_buffer.get()),
                                                   blockSize,
                                                   reinterpret_cast<uint8_t*>(out.get()),
                                                   blockSize,
-                                                  &outLen);
+                                                  &outLen,
+                                                  _dbName);
             uassert(28841,
                     str::stream() << "Failed to unprotect data: " << status.toString(),
                     status.isOK());
@@ -380,6 +383,7 @@ private:
     std::streampos _fileStartOffset;  // File offset at which the sorted data range starts.
     std::streampos _fileEndOffset;    // File offset at which the sorted data range ends.
     std::ifstream _file;
+    boost::optional<std::string> _dbName;
 
     // Checksum value that is updated with each read of a data object from disk. We can compare
     // this value with _originalChecksum to check for data corruption if and only if the
@@ -574,6 +578,7 @@ public:
                                range.getStartOffset(),
                                range.getEndOffset(),
                                this->_settings,
+                               this->_opts.dbName,
                                range.getChecksum());
                        });
     }
@@ -1025,7 +1030,8 @@ SortedFileWriter<Key, Value>::SortedFileWriter(const SortOptions& opts,
       // The file descriptor is positioned at the end of a file when opened in append mode, but
       // _file.tellp() is not initialized on all systems to reflect this. Therefore, we must also
       // pass in the expected offset to this constructor.
-      _fileStartOffset(fileStartOffset) {
+      _fileStartOffset(fileStartOffset),
+      _dbName(opts.dbName) {
 
     // This should be checked by consumers, but if we get here don't allow writes.
     uassert(
@@ -1096,7 +1102,8 @@ void SortedFileWriter<Key, Value>::spill() {
                                                         size,
                                                         reinterpret_cast<uint8_t*>(out.get()),
                                                         protectedSizeMax,
-                                                        &resultLen);
+                                                        &resultLen,
+                                                        _dbName);
         uassert(28842,
                 str::stream() << "Failed to compress data: " << status.toString(),
                 status.isOK());
@@ -1133,7 +1140,7 @@ SortIteratorInterface<Key, Value>* SortedFileWriter<Key, Value>::done() {
     _file.close();
 
     return new sorter::FileIterator<Key, Value>(
-        _fileFullPath, _fileStartOffset, _fileEndOffset, _settings, _checksum);
+        _fileFullPath, _fileStartOffset, _fileEndOffset, _settings, _dbName, _checksum);
 }
 
 //
diff --git a/src/mongo/db/sorter/sorter.h b/src/mongo/db/sorter/sorter.h
index af3130c7f52..4d59c6b0f5b 100644
--- a/src/mongo/db/sorter/sorter.h
+++ b/src/mongo/db/sorter/sorter.h
@@ -104,6 +104,12 @@ struct SortOptions {
     // maxMemoryUsageBytes, we will uassert.
     bool extSortAllowed;
 
+    // In case the sorter spills encrypted data to disk that must be readable even after process
+    // restarts, it must encrypt with a persistent key. This key is accessed using the database
+    // name that the sorted collection lives in. If encryption is enabled and dbName is boost::none,
+    // a temporary key is used.
+    boost::optional<std::string> dbName;
+
     // Directory into which we place a file when spilling to disk. Must be explicitly set if
     // extSortAllowed is true.
     std::string tempDir;
@@ -131,6 +137,11 @@ struct SortOptions {
         tempDir = newTempDir;
         return *this;
     }
+
+    SortOptions& DBName(std::string newDbName) {
+        dbName = std::move(newDbName);
+        return *this;
+    }
 };
 
 /**
@@ -345,6 +356,8 @@ private:
     // for the next SortedFileWriter instance using the same file.
     std::streampos _fileStartOffset;
     std::streampos _fileEndOffset;
+
+    boost::optional<std::string> _dbName;
 };
 }  // namespace mongo
 
diff --git a/src/mongo/db/storage/encryption_hooks.cpp b/src/mongo/db/storage/encryption_hooks.cpp
index eac49821a2c..786f963e9e1 100644
--- a/src/mongo/db/storage/encryption_hooks.cpp
+++ b/src/mongo/db/storage/encryption_hooks.cpp
@@ -75,14 +75,22 @@ boost::filesystem::path EncryptionHooks::getProtectedPathSuffix() {
     return "";
 }
 
-Status EncryptionHooks::protectTmpData(
-    const uint8_t* in, size_t inLen, uint8_t* out, size_t outLen, size_t* resultLen) {
+Status EncryptionHooks::protectTmpData(const uint8_t* in,
+                                       size_t inlen,
+                                       uint8_t* out,
+                                       size_t outLen,
+                                       size_t* resultLen,
+                                       boost::optional<std::string> dbName) {
     return Status(ErrorCodes::InternalError,
                   "Encryption hooks must be enabled to use preprocessTmpData.");
 }
 
-Status EncryptionHooks::unprotectTmpData(
-    const uint8_t* in, size_t inLen, uint8_t* out, size_t outLen, size_t* resultLen) {
+Status EncryptionHooks::unprotectTmpData(const uint8_t* in,
+                                         size_t inLen,
+                                         uint8_t* out,
+                                         size_t outLen,
+                                         size_t* resultLen,
+                                         boost::optional<std::string> dbName) {
     return Status(ErrorCodes::InternalError,
                   "Encryption hooks must be enabled to use postprocessTmpData.");
 }
diff --git a/src/mongo/db/storage/encryption_hooks.h b/src/mongo/db/storage/encryption_hooks.h
index f5a5db60923..c2c372d6533 100644
--- a/src/mongo/db/storage/encryption_hooks.h
+++ b/src/mongo/db/storage/encryption_hooks.h
@@ -84,16 +84,31 @@ public:
     virtual boost::filesystem::path getProtectedPathSuffix();
 
     /**
-     * Transform temp data to non-readable form before writing it to disk.
+     * Transform temporary data that has been spilled to disk into non-readable form. If dbName
+     * is specified, the database key corresponding to dbName will be used to encrypt the data.
+     * This key is persistent across process restarts. Otherwise, an ephemeral key that is only
+     * consistent for the duration of the process will be generated and used for encryption.
      */
-    virtual Status protectTmpData(
-        const uint8_t* in, size_t inLen, uint8_t* out, size_t outLen, size_t* resultLen);
+    virtual Status protectTmpData(const uint8_t* in,
+                                  size_t inLen,
+                                  uint8_t* out,
+                                  size_t outLen,
+                                  size_t* resultLen,
+                                  boost::optional<std::string> dbName);
 
     /**
-     * Tranforms temp data back to readable form, after reading from disk.
+     * Transform temporary data that has been spilled to disk back into readable form. If dbName
+     * is specified, the database key corresponding to dbName will be used to decrypt the data.
+     * This key is persistent across process restarts, so decryption will be successful even if a
+     * restart had occurred after encryption. Otherwise, an ephemeral key that can only decrypt data
+     * encrypted earlier in the current process's lifetime will be used.
      */
-    virtual Status unprotectTmpData(
-        const uint8_t* in, size_t inLen, uint8_t* out, size_t outLen, size_t* resultLen);
+    virtual Status unprotectTmpData(const uint8_t* in,
+                                    size_t inLen,
+                                    uint8_t* out,
+                                    size_t outLen,
+                                    size_t* resultLen,
+                                    boost::optional<std::string> dbName);
 
     /**
      * Inform the encryption storage system to prepare its data such that its files can be copied