1 files changed, 25 insertions, 117 deletions

diff --git a/src/mongo/db/exec/sbe/util/spilling.cpp b/src/mongo/db/exec/sbe/util/spilling.cpp
index 7675fad6846..c54f3bfe956 100644
--- a/src/mongo/db/exec/sbe/util/spilling.cpp
+++ b/src/mongo/db/exec/sbe/util/spilling.cpp
@@ -29,18 +29,6 @@
 
 #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 {
 
@@ -69,136 +57,56 @@ KeyString::Value decodeKeyString(const RecordId& rid, KeyString::TypeBits typeBi
     return kb.getValueCopy();
 }
 
-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;
-}
-
-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);
+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;
 }
 
-int SpillingStore::upsertToRecordStore(
-    OperationContext* opCtx,
-    const RecordId& key,
-    const value::MaterializedRow& val,
-    const KeyString::TypeBits& typeBits,  // recover type of value.
-    bool update) {
+int upsertToRecordStore(OperationContext* opCtx,
+                        RecordStore* rs,
+                        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);
+    return upsertToRecordStore(opCtx, rs, key, bufValue, typeBits, update);
 }
 
-int SpillingStore::upsertToRecordStore(
-    OperationContext* opCtx,
-    const RecordId& key,
-    BufBuilder& buf,
-    const KeyString::TypeBits& typeBits,  // recover type of value.
-    bool update) {
+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, 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;
     }
     return buf.len();
 }
-
-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;
-}
-
-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;
-}
-
-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