New upstream version 6.0.20upstream

10 files changed, 315 insertions, 140 deletions

diff --git a/src/mongo/db/exec/sbe/stages/hash_agg.cpp b/src/mongo/db/exec/sbe/stages/hash_agg.cpp
index 99bcc9f11c3..2514dedf1b0 100644
--- a/src/mongo/db/exec/sbe/stages/hash_agg.cpp
+++ b/src/mongo/db/exec/sbe/stages/hash_agg.cpp
@@ -103,18 +103,26 @@ std::unique_ptr<PlanStage> HashAggStage::clone() const {
 void HashAggStage::doSaveState(bool relinquishCursor) {
     if (relinquishCursor) {
         if (_rsCursor) {
-            _rsCursor->save();
+            _recordStore->saveCursor(_opCtx, _rsCursor);
         }
     }
     if (_rsCursor) {
         _rsCursor->setSaveStorageCursorOnDetachFromOperationContext(!relinquishCursor);
     }
+
+    if (_recordStore) {
+        _recordStore->saveState();
+    }
 }
 
 void HashAggStage::doRestoreState(bool relinquishCursor) {
     invariant(_opCtx);
+    if (_recordStore) {
+        _recordStore->restoreState();
+    }
+
     if (_rsCursor && relinquishCursor) {
-        auto couldRestore = _rsCursor->restore();
+        auto couldRestore = _recordStore->restoreCursor(_opCtx, _rsCursor);
         uassert(6196500, "HashAggStage could not restore cursor", couldRestore);
     }
 }
@@ -262,8 +270,7 @@ void HashAggStage::makeTemporaryRecordStore() {
             "No storage engine so HashAggStage cannot spill to disk",
             _opCtx->getServiceContext()->getStorageEngine());
     assertIgnorePrepareConflictsBehavior(_opCtx);
-    _recordStore = _opCtx->getServiceContext()->getStorageEngine()->makeTemporaryRecordStore(
-        _opCtx, KeyFormat::String);
+    _recordStore = std::make_unique<SpillingStore>(_opCtx);
 
     _specificStats.usedDisk = true;
 }
@@ -291,10 +298,10 @@ void HashAggStage::spillRowToDisk(const value::MaterializedRow& key,
     if (collator) {
         // The keystring cannot always be deserialized back to the original keys when a collation is
         // in use, so we also store the unmodified key in the data part of the spilled record.
-        upsertToRecordStore(_opCtx, _recordStore->rs(), rid, key, val, false /*update*/);
+        _recordStore->upsertToRecordStore(_opCtx, rid, key, val, false /*update*/);
     } else {
         auto typeBits = kb.getTypeBits();
-        upsertToRecordStore(_opCtx, _recordStore->rs(), rid, val, typeBits, false /*update*/);
+        _recordStore->upsertToRecordStore(_opCtx, rid, val, typeBits, false /*update*/);
     }
 
     _specificStats.spilledRecords++;
@@ -414,7 +421,9 @@ void HashAggStage::open(bool reOpen) {
         for (auto&& accessor : _outAggAccessors) {
             accessor->setIndex(0);
         }
-        _rsCursor.reset();
+        if (_recordStore) {
+            _recordStore->resetCursor(_opCtx, _rsCursor);
+        }
         _recordStore.reset();
         _outKeyRowRecordStore = {0};
         _outAggRowRecordStore = {0};
@@ -490,7 +499,7 @@ void HashAggStage::open(bool reOpen) {
             _specificStats.spilledDataStorageSize = _recordStore->rs()->storageSize(_opCtx);
 
             // Establish a cursor, positioned at the beginning of the record store.
-            _rsCursor = _recordStore->rs()->getCursor(_opCtx);
+            _rsCursor = _recordStore->getCursor(_opCtx);
 
             // Callers will be obtaining the results from the spill table, so set the
             // 'SwitchAccessors' so that they refer to the rows recovered from the record store
@@ -675,6 +684,9 @@ void HashAggStage::close() {
 
     trackClose();
     _ht = boost::none;
+    if (_recordStore && _opCtx) {
+        _recordStore->resetCursor(_opCtx, _rsCursor);
+    }
     _rsCursor.reset();
     _recordStore.reset();
     _outKeyRowRecordStore = {0};
diff --git a/src/mongo/db/exec/sbe/stages/hash_agg.h b/src/mongo/db/exec/sbe/stages/hash_agg.h
index 2f77e445883..91f91051363 100644
--- a/src/mongo/db/exec/sbe/stages/hash_agg.h
+++ b/src/mongo/db/exec/sbe/stages/hash_agg.h
@@ -31,6 +31,7 @@
 
 #include "mongo/db/exec/sbe/expressions/expression.h"
 #include "mongo/db/exec/sbe/stages/stages.h"
+#include "mongo/db/exec/sbe/util/spilling.h"
 #include "mongo/db/exec/sbe/vm/vm.h"
 #include "mongo/db/query/query_knobs_gen.h"
 #include "mongo/db/storage/temporary_record_store.h"
@@ -277,7 +278,7 @@ private:
         internalQuerySBEAggApproxMemoryUseInBytesBeforeSpill.load();
 
     // A record store which is instantiated and written to in the case of spilling.
-    std::unique_ptr<TemporaryRecordStore> _recordStore;
+    std::unique_ptr<SpillingStore> _recordStore;
     std::unique_ptr<SeekableRecordCursor> _rsCursor;
 
     // A monotically increasing counter used to ensure uniqueness of 'RecordId' values. When
diff --git a/src/mongo/db/exec/sbe/stages/hash_lookup.cpp b/src/mongo/db/exec/sbe/stages/hash_lookup.cpp
index fbc8ff73058..02e95307c4b 100644
--- a/src/mongo/db/exec/sbe/stages/hash_lookup.cpp
+++ b/src/mongo/db/exec/sbe/stages/hash_lookup.cpp
@@ -187,6 +187,22 @@ value::SlotAccessor* HashLookupStage::getAccessor(CompileCtx& ctx, value::SlotId
         return outerChild()->getAccessor(ctx, slot);
     }
 }
+void HashLookupStage::doSaveState(bool relinquishCursor) {
+    if (_recordStoreHt) {
+        _recordStoreHt->saveState();
+    }
+    if (_recordStoreBuf) {
+        _recordStoreBuf->saveState();
+    }
+}
+void HashLookupStage::doRestoreState(bool relinquishCursor) {
+    if (_recordStoreHt) {
+        _recordStoreHt->restoreState();
+    }
+    if (_recordStoreBuf) {
+        _recordStoreBuf->restoreState();
+    }
+}
 
 void HashLookupStage::reset() {
     _ht = boost::none;
@@ -259,7 +275,7 @@ void HashLookupStage::addHashTableEntry(value::SlotAccessor* keyAccessor, size_t
 
             auto val = std::vector<size_t>{valueIndex};
             auto [tagKey, valKey] = keyAccessor->getViewOfValue();
-            spillIndicesToRecordStore(_recordStoreHt->rs(), tagKey, valKey, val);
+            spillIndicesToRecordStore(_recordStoreHt.get(), tagKey, valKey, val);
         }
     } else {
         // The key is already present in '_ht' so the memory will only grow by one size_t. If we
@@ -281,7 +297,7 @@ void HashLookupStage::addHashTableEntry(value::SlotAccessor* keyAccessor, size_t
             // Evict the hash table value.
             _computedTotalMemUsage -= htIt->second.size() * sizeof(size_t);
             htIt->second.push_back(valueIndex);
-            spillIndicesToRecordStore(_recordStoreHt->rs(), tagKeyView, valKeyView, htIt->second);
+            spillIndicesToRecordStore(_recordStoreHt.get(), tagKeyView, valKeyView, htIt->second);
             _ht->erase(htIt);
         }
     }
@@ -297,17 +313,15 @@ void HashLookupStage::makeTemporaryRecordStore() {
             _opCtx->getServiceContext()->getStorageEngine());
     assertIgnorePrepareConflictsBehavior(_opCtx);
 
-    _recordStoreBuf = _opCtx->getServiceContext()->getStorageEngine()->makeTemporaryRecordStore(
-        _opCtx, KeyFormat::Long);
+    _recordStoreBuf = std::make_unique<SpillingStore>(_opCtx, KeyFormat::Long);
 
-    _recordStoreHt = _opCtx->getServiceContext()->getStorageEngine()->makeTemporaryRecordStore(
-        _opCtx, KeyFormat::String);
+    _recordStoreHt = std::make_unique<SpillingStore>(_opCtx, KeyFormat::String);
 
     _specificStats.usedDisk = true;
 }
 
 void HashLookupStage::spillBufferedValueToDisk(OperationContext* opCtx,
-                                               RecordStore* rs,
+                                               SpillingStore* rs,
                                                size_t bufferIdx,
                                                const value::MaterializedRow& val) {
     auto rid = getValueRecordId(bufferIdx);
@@ -315,15 +329,7 @@ void HashLookupStage::spillBufferedValueToDisk(OperationContext* opCtx,
     BufBuilder buf;
     val.serializeForSorter(buf);
 
-    assertIgnorePrepareConflictsBehavior(opCtx);
-    WriteUnitOfWork wuow(opCtx);
-
-    auto status = rs->insertRecord(opCtx, rid, buf.buf(), buf.len(), Timestamp{});
-    wuow.commit();
-
-    tassert(6373906,
-            str::stream() << "Failed to write to disk because " << status.getStatus().reason(),
-            status.isOK());
+    rs->upsertToRecordStore(opCtx, rid, buf, false);
 
     _specificStats.spilledBuffRecords++;
     // Add size of record ID + size of buffer.
@@ -341,7 +347,7 @@ size_t HashLookupStage::bufferValueOrSpill(value::MaterializedRow& value) {
         if (!hasSpilledBufToDisk()) {
             makeTemporaryRecordStore();
         }
-        spillBufferedValueToDisk(_opCtx, _recordStoreBuf->rs(), bufferIndex, value);
+        spillBufferedValueToDisk(_opCtx, _recordStoreBuf.get(), bufferIndex, value);
     }
     _valueId++;
     return bufferIndex;
@@ -427,7 +433,7 @@ void HashLookupStage::accumulateFromValueIndices(const C& bufferIndices) {
             // We must shift the '_bufferIt' index by one when using it as a RecordId because a
             // RecordId of 0 is invalid.
             auto rid = getValueRecordId(_bufferIt);
-            auto rsValue = readFromRecordStore(_opCtx, _recordStoreBuf->rs(), rid);
+            auto rsValue = _recordStoreBuf->readFromRecordStore(_opCtx, rid);
             if (!rsValue) {
                 tasserted(6373900, "bufferIdx not found in record store");
             }
@@ -443,7 +449,7 @@ void HashLookupStage::accumulateFromValueIndices(const C& bufferIndices) {
     }
 }
 
-void HashLookupStage::writeIndicesToRecordStore(RecordStore* rs,
+void HashLookupStage::writeIndicesToRecordStore(SpillingStore* rs,
                                                 value::TypeTags tagKey,
                                                 value::Value valKey,
                                                 const std::vector<size_t>& value,
@@ -458,7 +464,7 @@ void HashLookupStage::writeIndicesToRecordStore(RecordStore* rs,
     key.reset(0, false, tagKey, valKey);
     auto [rid, typeBits] = serializeKeyForRecordStore(key);
 
-    upsertToRecordStore(_opCtx, rs, rid, buf, typeBits, update);
+    rs->upsertToRecordStore(_opCtx, rid, buf, typeBits, update);
     if (!update) {
         _specificStats.spilledHtRecords++;
         // Add the size of key (which comprises of the memory usage for the key + its type bits),
@@ -471,7 +477,7 @@ void HashLookupStage::writeIndicesToRecordStore(RecordStore* rs,
 }
 
 boost::optional<std::vector<size_t>> HashLookupStage::readIndicesFromRecordStore(
-    RecordStore* rs, value::TypeTags tagKey, value::Value valKey) {
+    SpillingStore* rs, value::TypeTags tagKey, value::Value valKey) {
     _probeKey.reset(0, false, tagKey, valKey);
 
     auto [rid, _] = serializeKeyForRecordStore(_probeKey);
@@ -490,7 +496,7 @@ boost::optional<std::vector<size_t>> HashLookupStage::readIndicesFromRecordStore
     return boost::none;
 }
 
-void HashLookupStage::spillIndicesToRecordStore(RecordStore* rs,
+void HashLookupStage::spillIndicesToRecordStore(SpillingStore* rs,
                                                 value::TypeTags tagKey,
                                                 value::Value valKey,
                                                 const std::vector<size_t>& value) {
@@ -545,7 +551,7 @@ PlanState HashLookupStage::getNext() {
                         normalizeStringIfCollator(tagElemView, valElemView);
 
                     auto indicesFromRS = readIndicesFromRecordStore(
-                        _recordStoreHt->rs(), tagElemCollView, valElemCollView);
+                        _recordStoreHt.get(), tagElemCollView, valElemCollView);
                     if (indicesFromRS) {
                         indices.insert(indicesFromRS->begin(), indicesFromRS->end());
                     }
@@ -567,7 +573,7 @@ PlanState HashLookupStage::getNext() {
                     normalizeStringIfCollator(tagKeyView, valKeyView);
 
                 auto indicesFromRS = readIndicesFromRecordStore(
-                    _recordStoreHt->rs(), tagKeyCollView, valKeyCollView);
+                    _recordStoreHt.get(), tagKeyCollView, valKeyCollView);
                 if (indicesFromRS) {
                     accumulateFromValueIndices(*indicesFromRS);
                 }
diff --git a/src/mongo/db/exec/sbe/stages/hash_lookup.h b/src/mongo/db/exec/sbe/stages/hash_lookup.h
index 2e3f0b34816..b312e0a68f4 100644
--- a/src/mongo/db/exec/sbe/stages/hash_lookup.h
+++ b/src/mongo/db/exec/sbe/stages/hash_lookup.h
@@ -33,6 +33,7 @@
 
 #include "mongo/db/exec/sbe/expressions/expression.h"
 #include "mongo/db/exec/sbe/stages/stages.h"
+#include "mongo/db/exec/sbe/util/spilling.h"
 #include "mongo/db/exec/sbe/vm/vm.h"
 #include "mongo/db/query/query_knobs_gen.h"
 
@@ -101,6 +102,10 @@ public:
     std::vector<DebugPrinter::Block> debugPrint() const final;
     size_t estimateCompileTimeSize() const final;
 
+protected:
+    void doSaveState(bool relinquishCursor) override;
+    void doRestoreState(bool relinquishCursor) override;
+
 private:
     using HashTableType = std::unordered_map<value::MaterializedRow,  // NOLINT
                                              std::vector<size_t>,
@@ -119,23 +124,23 @@ private:
     // Spilling helpers.
     void addHashTableEntry(value::SlotAccessor* keyAccessor, size_t valueIndex);
     void spillBufferedValueToDisk(OperationContext* opCtx,
-                                  RecordStore* rs,
+                                  SpillingStore* rs,
                                   size_t bufferIdx,
                                   const value::MaterializedRow&);
     size_t bufferValueOrSpill(value::MaterializedRow& value);
     void setInnerProjectSwitchAccessor(int idx);
 
-    boost::optional<std::vector<size_t>> readIndicesFromRecordStore(RecordStore* rs,
+    boost::optional<std::vector<size_t>> readIndicesFromRecordStore(SpillingStore* rs,
                                                                     value::TypeTags tagKey,
                                                                     value::Value valKey);
 
-    void writeIndicesToRecordStore(RecordStore* rs,
+    void writeIndicesToRecordStore(SpillingStore* rs,
                                    value::TypeTags tagKey,
                                    value::Value valKey,
                                    const std::vector<size_t>& value,
                                    bool update);
 
-    void spillIndicesToRecordStore(RecordStore* rs,
+    void spillIndicesToRecordStore(SpillingStore* rs,
                                    value::TypeTags tagKey,
                                    value::Value valKey,
                                    const std::vector<size_t>& value);
@@ -229,8 +234,8 @@ private:
     // rows in '_buffer'.
     long long _computedTotalMemUsage = 0;
 
-    std::unique_ptr<TemporaryRecordStore> _recordStoreHt;
-    std::unique_ptr<TemporaryRecordStore> _recordStoreBuf;
+    std::unique_ptr<SpillingStore> _recordStoreHt;
+    std::unique_ptr<SpillingStore> _recordStoreBuf;
 
     HashLookupStats _specificStats;
 };
diff --git a/src/mongo/db/exec/sbe/util/spilling.cpp b/src/mongo/db/exec/sbe/util/spilling.cpp
index 0f0cbb93d94..7675fad6846 100644
--- a/src/mongo/db/exec/sbe/util/spilling.cpp
+++ b/src/mongo/db/exec/sbe/util/spilling.cpp
@@ -29,6 +29,18 @@
 
 #include "mongo/db/exec/sbe/util/spilling.h"
 
+#include "mongo/base/status.h"
+#include "mongo/base/status_with.h"
+#include "mongo/base/string_data.h"
+#include "mongo/bson/timestamp.h"
+#include "mongo/db/query/query_knobs_gen.h"
+#include "mongo/db/storage/record_data.h"
+#include "mongo/db/storage/recovery_unit.h"
+#include "mongo/db/storage/write_unit_of_work.h"
+#include "mongo/util/assert_util.h"
+#include "mongo/util/bufreader.h"
+#include "mongo/util/str.h"
+
 namespace mongo {
 namespace sbe {
 
@@ -57,32 +69,76 @@ KeyString::Value decodeKeyString(const RecordId& rid, KeyString::TypeBits typeBi
     return kb.getValueCopy();
 }
 
-boost::optional<value::MaterializedRow> readFromRecordStore(OperationContext* opCtx,
-                                                            RecordStore* rs,
-                                                            const RecordId& rid) {
-    RecordData record;
-    if (rs->findRecord(opCtx, rid, &record)) {
-        auto valueReader = BufReader(record.data(), record.size());
-        return value::MaterializedRow::deserializeForSorter(valueReader, {});
-    }
-    return boost::none;
+SpillingStore::SpillingStore(OperationContext* opCtx, KeyFormat format) {
+    _recordStore =
+        opCtx->getServiceContext()->getStorageEngine()->makeTemporaryRecordStore(opCtx, format);
+
+    _spillingUnit = std::unique_ptr<RecoveryUnit>(
+        opCtx->getServiceContext()->getStorageEngine()->newRecoveryUnit());
+    _spillingUnit->setCacheMaxWaitTimeout(Milliseconds(internalQuerySpillingMaxWaitTimeout.load()));
+    _spillingState = WriteUnitOfWork::RecoveryUnitState::kNotInUnitOfWork;
 }
 
-static int upsertToRecordStore(
-    OperationContext* opCtx, RecordStore* rs, const RecordId& key, BufBuilder& buf, bool update) {
+SpillingStore::~SpillingStore() {}
 
+int SpillingStore::upsertToRecordStore(OperationContext* opCtx,
+                                       const RecordId& recordKey,
+                                       const value::MaterializedRow& key,
+                                       const value::MaterializedRow& val,
+                                       bool update) {
+    BufBuilder buf;
+    key.serializeForSorter(buf);
+    val.serializeForSorter(buf);
+    return upsertToRecordStore(opCtx, recordKey, buf, update);
+}
+
+int SpillingStore::upsertToRecordStore(
+    OperationContext* opCtx,
+    const RecordId& key,
+    const value::MaterializedRow& val,
+    const KeyString::TypeBits& typeBits,  // recover type of value.
+    bool update) {
+    BufBuilder bufValue;
+    val.serializeForSorter(bufValue);
+    // Append the 'typeBits' to the end of the val's buffer so the 'key' can be reconstructed when
+    // draining HashAgg.
+    bufValue.appendBuf(typeBits.getBuffer(), typeBits.getSize());
+
+    return upsertToRecordStore(opCtx, key, bufValue, update);
+}
+
+int SpillingStore::upsertToRecordStore(
+    OperationContext* opCtx,
+    const RecordId& key,
+    BufBuilder& buf,
+    const KeyString::TypeBits& typeBits,  // recover type of value.
+    bool update) {
+    // Append the 'typeBits' to the end of the val's buffer so the 'key' can be reconstructed when
+    // draining HashAgg.
+    buf.appendBuf(typeBits.getBuffer(), typeBits.getSize());
+
+    return upsertToRecordStore(opCtx, key, buf, update);
+}
+
+int SpillingStore::upsertToRecordStore(OperationContext* opCtx,
+                                       const RecordId& key,
+                                       BufBuilder& buf,
+                                       bool update) {
     assertIgnorePrepareConflictsBehavior(opCtx);
 
+    switchToSpilling(opCtx);
+    ON_BLOCK_EXIT([&] { switchToOriginal(opCtx); });
     WriteUnitOfWork wuow(opCtx);
 
     auto result = mongo::Status::OK();
     if (update) {
-        result = rs->updateRecord(opCtx, key, buf.buf(), buf.len());
+        result = rs()->updateRecord(opCtx, key, buf.buf(), buf.len());
     } else {
-        auto status = rs->insertRecord(opCtx, key, buf.buf(), buf.len(), Timestamp{});
+        auto status = rs()->insertRecord(opCtx, key, buf.buf(), buf.len(), Timestamp{});
         result = status.getStatus();
     }
     wuow.commit();
+
     if (!result.isOK()) {
         tasserted(5843600, str::stream() << "Failed to write to disk because " << result.reason());
         return 0;
@@ -90,42 +146,59 @@ static int upsertToRecordStore(
     return buf.len();
 }
 
-int upsertToRecordStore(OperationContext* opCtx,
-                        RecordStore* rs,
-                        const RecordId& key,
-                        const value::MaterializedRow& val,
-                        const KeyString::TypeBits& typeBits,  // recover type of value.
-                        bool update) {
-    BufBuilder buf;
-    val.serializeForSorter(buf);
-    // Append the 'typeBits' to the end of the val's buffer so the 'key' can be reconstructed when
-    // draining HashAgg.
-    buf.appendBuf(typeBits.getBuffer(), typeBits.getSize());
-    return upsertToRecordStore(opCtx, rs, key, buf, update);
+Status SpillingStore::insertRecords(OperationContext* opCtx,
+                                    std::vector<Record>* inOutRecords,
+                                    const std::vector<Timestamp>& timestamps) {
+    assertIgnorePrepareConflictsBehavior(opCtx);
+
+    switchToSpilling(opCtx);
+    ON_BLOCK_EXIT([&] { switchToOriginal(opCtx); });
+    WriteUnitOfWork wuow(opCtx);
+    auto status = rs()->insertRecords(opCtx, inOutRecords, timestamps);
+    wuow.commit();
+
+    return status;
 }
 
-int upsertToRecordStore(OperationContext* opCtx,
-                        RecordStore* rs,
-                        const RecordId& recordKey,
-                        const value::MaterializedRow& key,
-                        const value::MaterializedRow& val,
-                        bool update) {
-    BufBuilder buf;
-    key.serializeForSorter(buf);
-    val.serializeForSorter(buf);
-    return upsertToRecordStore(opCtx, rs, recordKey, buf, update);
+boost::optional<value::MaterializedRow> SpillingStore::readFromRecordStore(OperationContext* opCtx,
+                                                                           const RecordId& rid) {
+    switchToSpilling(opCtx);
+    ON_BLOCK_EXIT([&] { switchToOriginal(opCtx); });
+
+    RecordData record;
+    if (rs()->findRecord(opCtx, rid, &record)) {
+        auto valueReader = BufReader(record.data(), record.size());
+        return value::MaterializedRow::deserializeForSorter(valueReader, {});
+    }
+    return boost::none;
 }
 
-int upsertToRecordStore(OperationContext* opCtx,
-                        RecordStore* rs,
-                        const RecordId& key,
-                        BufBuilder& buf,
-                        const KeyString::TypeBits& typeBits,  // recover type of value.
-                        bool update) {
-    // Append the 'typeBits' to the end of the val's buffer so the 'key' can be reconstructed when
-    // draining HashAgg.
-    buf.appendBuf(typeBits.getBuffer(), typeBits.getSize());
-    return upsertToRecordStore(opCtx, rs, key, buf, update);
+bool SpillingStore::findRecord(OperationContext* opCtx, const RecordId& loc, RecordData* out) {
+    switchToSpilling(opCtx);
+    ON_BLOCK_EXIT([&] { switchToOriginal(opCtx); });
+
+    return rs()->findRecord(opCtx, loc, out);
 }
+
+void SpillingStore::switchToSpilling(OperationContext* opCtx) {
+    invariant(!_originalUnit);
+    _originalUnit = opCtx->releaseRecoveryUnit();
+    _originalState = opCtx->setRecoveryUnit(std::move(_spillingUnit), _spillingState);
+}
+void SpillingStore::switchToOriginal(OperationContext* opCtx) {
+    invariant(!_spillingUnit);
+    _spillingUnit = opCtx->releaseRecoveryUnit();
+    _spillingState = opCtx->setRecoveryUnit(std::move(_originalUnit), _originalState);
+    invariant(!(_spillingUnit->getState() == RecoveryUnit::State::kInactiveInUnitOfWork ||
+                _spillingUnit->getState() == RecoveryUnit::State::kActive));
+}
+
+void SpillingStore::saveState() {
+    _spillingUnit->abandonSnapshot();
+}
+void SpillingStore::restoreState() {
+    // We do not have to do anything.
+}
+
 }  // namespace sbe
 }  // namespace mongo
diff --git a/src/mongo/db/exec/sbe/util/spilling.h b/src/mongo/db/exec/sbe/util/spilling.h
index 2d0eb98ec88..205d6f1a031 100644
--- a/src/mongo/db/exec/sbe/util/spilling.h
+++ b/src/mongo/db/exec/sbe/util/spilling.h
@@ -29,9 +29,14 @@
 
 #pragma once
 
-#include "mongo/platform/basic.h"
+#include <boost/optional/optional.hpp>
+#include <utility>
 
+#include "mongo/bson/util/builder.h"
 #include "mongo/db/exec/sbe/values/slot.h"
+#include "mongo/db/operation_context.h"
+#include "mongo/db/record_id.h"
+#include "mongo/db/storage/record_store.h"
 #include "mongo/db/storage/temporary_record_store.h"
 
 namespace mongo {
@@ -50,40 +55,104 @@ std::pair<RecordId, KeyString::TypeBits> encodeKeyString(KeyString::Builder&,
 // Reconstructs the KeyString carried in RecordId using 'typeBits'.
 KeyString::Value decodeKeyString(const RecordId& rid, KeyString::TypeBits typeBits);
 
-// Reads a materialized row from the record store.
-boost::optional<value::MaterializedRow> readFromRecordStore(OperationContext* opCtx,
-                                                            RecordStore* rs,
-                                                            const RecordId& rid);
-
-/**
- * Inserts or updates a key/value into 'rs'. The 'update' flag controls whether or not an update
- * will be performed. If a key/value pair is inserted into the 'rs' that already exists and
- * 'update' is false, this function will tassert.
- *
- * Returns the size of the new record in bytes, including the record id and value portions.
- */
-int upsertToRecordStore(OperationContext* opCtx,
-                        RecordStore* rs,
-                        const RecordId& key,
-                        const value::MaterializedRow& val,
-                        const KeyString::TypeBits& typeBits,
-                        bool update);
 /**
- * When a collator is provided, the key is encoded using the collator before being converted to a
- * record id. In this case, it is not possible to recover the key from the record id, thus we need
- * to store the original value of the key as well.
+ * SpillingStore is a wrapper around a temporary record store than maintains its own transaction as
+ * we do not want to intermingle operations running in the main query with spill reads and writes.
  */
-int upsertToRecordStore(OperationContext* opCtx,
-                        RecordStore* rs,
-                        const RecordId& recordKey,
-                        const value::MaterializedRow& key,
-                        const value::MaterializedRow& val,
-                        bool update);
-int upsertToRecordStore(OperationContext* opCtx,
-                        RecordStore* rs,
-                        const RecordId& key,
-                        BufBuilder& buf,
-                        const KeyString::TypeBits& typeBits,  // recover type of value.
-                        bool update);
+class SpillingStore {
+public:
+    SpillingStore(OperationContext* opCtx, KeyFormat format = KeyFormat::String);
+    ~SpillingStore();
+
+    /**
+     * When a collator is provided, the key is encoded using the collator before being converted to
+     * a record id. In this case, it is not possible to recover the key from the record id, thus we
+     * need to store the original value of the key as well.
+     */
+    int upsertToRecordStore(OperationContext* opCtx,
+                            const RecordId& recordKey,
+                            const value::MaterializedRow& key,
+                            const value::MaterializedRow& val,
+                            bool update);
+    /**
+     * Inserts or updates a key/value into 'rs'. The 'update' flag controls whether or not an update
+     * will be performed. If a key/value pair is inserted into the 'rs' that already exists and
+     * 'update' is false, this function will tassert.
+     *
+     * Returns the size of the new record in bytes, including the record id and value portions.
+     */
+    int upsertToRecordStore(OperationContext* opCtx,
+                            const RecordId& key,
+                            const value::MaterializedRow& val,
+                            const KeyString::TypeBits& typeBits,
+                            bool update);
+    int upsertToRecordStore(OperationContext* opCtx,
+                            const RecordId& key,
+                            BufBuilder& buf,
+                            const KeyString::TypeBits& typeBits,  // recover type of value.
+                            bool update);
+    int upsertToRecordStore(OperationContext* opCtx,
+                            const RecordId& key,
+                            BufBuilder& buf,
+                            bool update);
+
+
+    Status insertRecords(OperationContext* opCtx,
+                         std::vector<Record>* inOutRecords,
+                         const std::vector<Timestamp>& timestamps);
+
+    // Reads a materialized row from the record store.
+    boost::optional<value::MaterializedRow> readFromRecordStore(OperationContext* opCtx,
+                                                                const RecordId& rid);
+
+    bool findRecord(OperationContext* opCtx, const RecordId& loc, RecordData* out);
+
+    auto rs() {
+        return _recordStore->rs();
+    }
+
+    auto getCursor(OperationContext* opCtx) {
+        switchToSpilling(opCtx);
+        ON_BLOCK_EXIT([&] { switchToOriginal(opCtx); });
+        return rs()->getCursor(opCtx);
+    }
+
+    void resetCursor(OperationContext* opCtx, std::unique_ptr<SeekableRecordCursor>& cursor) {
+        switchToSpilling(opCtx);
+        ON_BLOCK_EXIT([&] { switchToOriginal(opCtx); });
+        cursor.reset();
+    }
+
+    auto saveCursor(OperationContext* opCtx, std::unique_ptr<SeekableRecordCursor>& cursor) {
+        switchToSpilling(opCtx);
+        ON_BLOCK_EXIT([&] { switchToOriginal(opCtx); });
+
+        return cursor->save();
+    }
+
+    auto restoreCursor(OperationContext* opCtx, std::unique_ptr<SeekableRecordCursor>& cursor) {
+        switchToSpilling(opCtx);
+        ON_BLOCK_EXIT([&] { switchToOriginal(opCtx); });
+
+        return cursor->restore();
+    }
+
+    void saveState();
+    void restoreState();
+
+private:
+    void switchToSpilling(OperationContext* opCtx);
+    void switchToOriginal(OperationContext* opCtx);
+
+    std::unique_ptr<TemporaryRecordStore> _recordStore;
+
+    std::unique_ptr<RecoveryUnit> _originalUnit;
+    WriteUnitOfWork::RecoveryUnitState _originalState;
+
+    std::unique_ptr<RecoveryUnit> _spillingUnit;
+    WriteUnitOfWork::RecoveryUnitState _spillingState;
+
+    size_t _counter{0};
+};
 }  // namespace sbe
 }  // namespace mongo
diff --git a/src/mongo/db/exec/sbe/values/slot.cpp b/src/mongo/db/exec/sbe/values/slot.cpp
index 2dd622fcecb..cb53de849ac 100644
--- a/src/mongo/db/exec/sbe/values/slot.cpp
+++ b/src/mongo/db/exec/sbe/values/slot.cpp
@@ -271,7 +271,7 @@ static void serializeValue(BufBuilder& buf, TypeTags tag, Value val) {
         case TypeTags::StringSmall: {
             // Small strings cannot contain null bytes, so it is safe to serialize them as plain
             // C-strings with a null terminator.
-            buf.appendStr(getStringView(tag, val), true /* includeEndingNull */);
+            buf.appendCStr(getStringView(tag, val));
             break;
         }
         case TypeTags::StringBig:
@@ -279,7 +279,7 @@ static void serializeValue(BufBuilder& buf, TypeTags tag, Value val) {
         case TypeTags::bsonSymbol: {
             auto sv = getStringOrSymbolView(tag, val);
             buf.appendNum(static_cast<uint32_t>(sv.size()));
-            buf.appendStr(sv, false /* includeEndingNull */);
+            buf.appendStrBytes(sv);
             break;
         }
         case TypeTags::Array: {
@@ -309,7 +309,7 @@ static void serializeValue(BufBuilder& buf, TypeTags tag, Value val) {
             auto obj = getObjectView(val);
             buf.appendNum(obj->size());
             for (size_t idx = 0; idx < obj->size(); ++idx) {
-                buf.appendStr(obj->field(idx), true /* includeEndingNull */);
+                buf.appendCStr(obj->field(idx));
                 auto [tag, val] = obj->getAt(idx);
                 serializeValue(buf, tag, val);
             }
@@ -352,27 +352,27 @@ static void serializeValue(BufBuilder& buf, TypeTags tag, Value val) {
         }
         case TypeTags::bsonRegex: {
             auto regex = getBsonRegexView(val);
-            buf.appendStr(regex.pattern, true /* includeEndingNull */);
-            buf.appendStr(regex.flags, true /* includeEndingNull */);
+            buf.appendCStr(regex.pattern);
+            buf.appendCStr(regex.flags);
             break;
         }
         case TypeTags::bsonJavascript: {
             auto javascriptCode = getBsonJavascriptView(val);
             buf.appendNum(static_cast<uint32_t>(javascriptCode.size()));
-            buf.appendStr(javascriptCode, false /* includeEndingNull */);
+            buf.appendStrBytes(javascriptCode);
             break;
         }
         case TypeTags::bsonDBPointer: {
             auto dbptr = getBsonDBPointerView(val);
             buf.appendNum(static_cast<uint32_t>(dbptr.ns.size()));
-            buf.appendStr(dbptr.ns, false /* includeEndingNull */);
+            buf.appendStrBytes(dbptr.ns);
             buf.appendBuf(dbptr.id, sizeof(ObjectIdType));
             break;
         }
         case TypeTags::bsonCodeWScope: {
             auto cws = getBsonCodeWScopeView(val);
             buf.appendNum(static_cast<uint32_t>(cws.code.size()));
-            buf.appendStr(cws.code, false /* includeEndingNull */);
+            buf.appendStrBytes(cws.code);
             auto scopeLen = ConstDataView(cws.scope).read<LittleEndian<uint32_t>>();
             buf.appendBuf(cws.scope, scopeLen);
             break;
@@ -507,9 +507,10 @@ static void serializeValueIntoKeyString(KeyString::Builder& buf,
             }
             break;
         }
+        case TypeTags::bsonObjectId:
         case TypeTags::ObjectId: {
             buf.appendBool(true);
-            buf.appendBytes(getObjectIdView(val), sizeof(ObjectIdType));
+            buf.appendOID(OID::from(getRawPointerView(val)));
             break;
         }
         case TypeTags::bsonObject: {
@@ -532,11 +533,6 @@ static void serializeValueIntoKeyString(KeyString::Builder& buf,
             }
             break;
         }
-        case TypeTags::bsonObjectId: {
-            buf.appendBool(true);
-            buf.appendOID(OID::from(getRawPointerView(val)));
-            break;
-        }
         case TypeTags::bsonBinData: {
             BufBuilder innerBinDataBuf;
             innerBinDataBuf.appendUChar(static_cast<uint8_t>(tag));
diff --git a/src/mongo/db/exec/sbe/values/value_builder.h b/src/mongo/db/exec/sbe/values/value_builder.h
index 00333e9f824..53748519723 100644
--- a/src/mongo/db/exec/sbe/values/value_builder.h
+++ b/src/mongo/db/exec/sbe/values/value_builder.h
@@ -112,21 +112,21 @@ public:
         } else {
             appendValueBufferOffset(TypeTags::StringBig);
             _valueBufferBuilder->appendNum(static_cast<int32_t>(in.size() + 1));
-            _valueBufferBuilder->appendStr(in, true /* includeEndingNull */);
+            _valueBufferBuilder->appendStrBytesAndNul(in);
         }
     }
 
     void append(const BSONSymbol& in) {
         appendValueBufferOffset(TypeTags::bsonSymbol);
         _valueBufferBuilder->appendNum(static_cast<int32_t>(in.symbol.size() + 1));
-        _valueBufferBuilder->appendStr(in.symbol, true /* includeEndingNull */);
+        _valueBufferBuilder->appendStrBytesAndNul(in.symbol);
     }
 
     void append(const BSONCode& in) {
         appendValueBufferOffset(TypeTags::bsonJavascript);
         // Add one to account null byte at the end.
         _valueBufferBuilder->appendNum(static_cast<uint32_t>(in.code.size() + 1));
-        _valueBufferBuilder->appendStr(in.code, true /* includeEndingNull */);
+        _valueBufferBuilder->appendStrBytesAndNul(in.code);
     }
 
     void append(const BSONCodeWScope& in) {
@@ -134,7 +134,7 @@ public:
         _valueBufferBuilder->appendNum(
             static_cast<uint32_t>(4 + in.code.size() + 1 + in.scope.objsize()));
         _valueBufferBuilder->appendNum(static_cast<int32_t>(in.code.size() + 1));
-        _valueBufferBuilder->appendStr(in.code, true /* includeEndingNull */);
+        _valueBufferBuilder->appendStrBytesAndNul(in.code);
         _valueBufferBuilder->appendBuf(in.scope.objdata(), in.scope.objsize());
     }
 
@@ -147,14 +147,14 @@ public:
 
     void append(const BSONRegEx& in) {
         appendValueBufferOffset(TypeTags::bsonRegex);
-        _valueBufferBuilder->appendStr(in.pattern, true /* includeEndingNull */);
-        _valueBufferBuilder->appendStr(in.flags, true /* includeEndingNull */);
+        _valueBufferBuilder->appendCStr(in.pattern);
+        _valueBufferBuilder->appendCStr(in.flags);
     }
 
     void append(const BSONDBRef& in) {
         appendValueBufferOffset(TypeTags::bsonDBPointer);
         _valueBufferBuilder->appendNum(static_cast<int32_t>(in.ns.size() + 1));
-        _valueBufferBuilder->appendStr(in.ns, true /* includeEndingNull */);
+        _valueBufferBuilder->appendStrBytesAndNul(in.ns);
         _valueBufferBuilder->appendBuf(in.oid.view().view(), OID::kOIDSize);
     }
 
diff --git a/src/mongo/db/exec/sbe/values/value_serialization_test.cpp b/src/mongo/db/exec/sbe/values/value_serialization_test.cpp
index 3be3212c627..ceacf610a4a 100644
--- a/src/mongo/db/exec/sbe/values/value_serialization_test.cpp
+++ b/src/mongo/db/exec/sbe/values/value_serialization_test.cpp
@@ -467,4 +467,17 @@ TEST_F(ValueSerializeForKeyString, RoundtripWideRow) {
     }
     runTest(row);
 }
+
+// Test that roundtripping through KeyString works for ObjectIdType: ObjectId; bsonObjectId.
+TEST_F(ValueSerializeForKeyString, RoundtripObjectIdType) {
+    auto [objectIdTag, objectIdVal] = value::makeNewObjectId();
+
+    auto oid = OID::gen();
+    auto obj = BSON("" << oid);
+    auto oidStorage = obj.firstElement().value();
+
+    sbe::value::ValueGuard testDataGuard{objectIdTag, objectIdVal};
+    runTest({{objectIdTag, objectIdVal},
+             {value::TypeTags::bsonObjectId, value::bitcastFrom<const char*>(oidStorage)}});
+}
 }  // namespace mongo::sbe
diff --git a/src/mongo/db/exec/sbe/vm/arith.cpp b/src/mongo/db/exec/sbe/vm/arith.cpp
index 1d41ff59ce1..e6c9d380ffd 100644
--- a/src/mongo/db/exec/sbe/vm/arith.cpp
+++ b/src/mongo/db/exec/sbe/vm/arith.cpp
@@ -1027,7 +1027,7 @@ std::tuple<bool, value::TypeTags, value::Value> ByteCode::genericLn(value::TypeT
             if (!operand.isGreater(Decimal128::kNormalizedZero) && !operand.isNaN()) {
                 return {false, value::TypeTags::Nothing, 0};
             }
-            auto operandLn = operand.logarithm();
+            auto operandLn = operand.naturalLogarithm();
 
             auto [tag, value] = value::makeCopyDecimal(operandLn);
             return {true, tag, value};