summaryrefslogtreecommitdiff
path: root/src/mongo/db/exec/sbe/util
diff options
context:
space:
mode:
Diffstat (limited to 'src/mongo/db/exec/sbe/util')
-rw-r--r--src/mongo/db/exec/sbe/util/spilling.cpp163
-rw-r--r--src/mongo/db/exec/sbe/util/spilling.h137
2 files changed, 221 insertions, 79 deletions
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
generated by cgit v1.2.3 (git 2.39.1) at 2026-09-11 12:47:54 +0000