From da7af1bad1fd6c6f6eab6d5a1774d59ab21b74b5 Mon Sep 17 00:00:00 2001 From: Lee Maguire Date: Wed, 13 May 2026 10:49:12 -0400 Subject: [PATCH] SERVER-116055 Add support for $accumulator in MozJS WASM engine. (#53341) GitOrigin-RevId: 575a1be1e9bf6e8510aab16cdd924428473e7ca2 --- MODULE.bazel | 6 + bazel/crates.lock | 12 +- bazel/wasmtime/BUILD.bazel | 5 + bazel/wasmtime/store_h.patch | 19 +++ bazel/wasmtime/store_hh.patch | 18 +++ bazel/wasmtime/store_rs.patch | 20 +++ buildscripts/resmokelib/configure_resmoke.py | 8 +- .../accumulators/accumulator_js.js | 146 ++++++++++++++++++ .../accumulator_js_size_limits.js | 8 +- jstests/core/query/system_js_access.js | 1 - src/mongo/db/pipeline/BUILD.bazel | 4 + src/mongo/db/query/query_tester/main.cpp | 6 +- .../scripting/mozjs/common/objectwrapper.cpp | 14 +- src/mongo/scripting/mozjs/wasm/BUILD.bazel | 8 +- .../scripting/mozjs/wasm/bridge/bridge.cpp | 6 + .../scripting/mozjs/wasm/engine/engine.cpp | 18 +-- .../scripting/mozjs/wasm/engine/error.cpp | 2 + .../scripting/mozjs/wasm/scope/scope.cpp | 9 +- src/mongo/scripting/mozjs/wasm/scope/scope.h | 1 + .../scripting/mozjs/wasm/scope/scope_test.cpp | 68 ++++---- .../scripting/mozjs/wasm/wasmtime_engine.cpp | 13 +- 21 files changed, 307 insertions(+), 85 deletions(-) create mode 100644 bazel/wasmtime/BUILD.bazel create mode 100644 bazel/wasmtime/store_h.patch create mode 100644 bazel/wasmtime/store_hh.patch create mode 100644 bazel/wasmtime/store_rs.patch diff --git a/MODULE.bazel b/MODULE.bazel index 74bf80232de..28ddf30bc33 100644 --- a/MODULE.bazel +++ b/MODULE.bazel @@ -391,6 +391,12 @@ crate.annotation( crate = "wasmtime-c-api-impl", extra_aliased_targets = {"wasmtime_c": "wasmtime_c"}, gen_build_script = "off", + patch_args = ["-p1"], + patches = [ + "//bazel/wasmtime:store_rs.patch", + "//bazel/wasmtime:store_h.patch", + "//bazel/wasmtime:store_hh.patch", + ], ) # Disable default features and re-enable everything except "objdump/explore" so we don't pull in capstone/capstone-sys. diff --git a/bazel/crates.lock b/bazel/crates.lock index 0743905ac95..c82671f28ba 100644 --- a/bazel/crates.lock +++ b/bazel/crates.lock @@ -1,5 +1,5 @@ { - "checksum": "09e089e4d857eb519291450cae1e356a0e8486203c17893a9ded5d57f1857484", + "checksum": "c5167d945fea0a88114d893d2437bad55ac3b4e295eaf0ed5348d3ba97edf335", "crates": { "addr2line 0.26.1": { "name": "addr2line", @@ -17534,7 +17534,15 @@ "repository": { "Http": { "url": "https://static.crates.io/crates/wasmtime-c-api-impl/44.0.1/download", - "sha256": "af63f5db854133f67cdec5faabff9db2764221968231c42fc18e8a9d63f849d5" + "sha256": "af63f5db854133f67cdec5faabff9db2764221968231c42fc18e8a9d63f849d5", + "patch_args": [ + "-p1" + ], + "patches": [ + "@@//bazel/wasmtime:store_h.patch", + "@@//bazel/wasmtime:store_hh.patch", + "@@//bazel/wasmtime:store_rs.patch" + ] } }, "targets": [ diff --git a/bazel/wasmtime/BUILD.bazel b/bazel/wasmtime/BUILD.bazel new file mode 100644 index 00000000000..a1f98bdb229 --- /dev/null +++ b/bazel/wasmtime/BUILD.bazel @@ -0,0 +1,5 @@ +exports_files([ + "store_rs.patch", + "store_h.patch", + "store_hh.patch", +]) diff --git a/bazel/wasmtime/store_h.patch b/bazel/wasmtime/store_h.patch new file mode 100644 index 00000000000..ddf08e4b86e --- /dev/null +++ b/bazel/wasmtime/store_h.patch @@ -0,0 +1,19 @@ +--- a/include/wasmtime/store.h ++++ b/include/wasmtime/store.h +@@ -164,6 +164,16 @@ + WASM_API_EXTERN wasmtime_error_t * + wasmtime_context_set_fuel(wasmtime_context_t *store, uint64_t fuel); + ++/** ++ * \brief Set the per-Store hostcall fuel cap (component model). ++ * ++ * Bounds how much data wasmtime will copy between host and guest across ++ * component-model calls. Pass `SIZE_MAX` to disable the cap. Mirrors ++ * `Store::set_hostcall_fuel` on the Rust side. ++ */ ++WASM_API_EXTERN void ++wasmtime_context_set_hostcall_fuel(wasmtime_context_t *store, size_t fuel); ++ + /** + * \brief Returns the amount of fuel remaining in this context's store. + * diff --git a/bazel/wasmtime/store_hh.patch b/bazel/wasmtime/store_hh.patch new file mode 100644 index 00000000000..da8d85bb045 --- /dev/null +++ b/bazel/wasmtime/store_hh.patch @@ -0,0 +1,18 @@ +--- a/include/wasmtime/store.hh ++++ b/include/wasmtime/store.hh +@@ -118,6 +118,15 @@ + return std::monostate(); + } + ++ /// Sets the per-Store hostcall fuel cap (component model). ++ /// ++ /// Bounds how much data wasmtime will copy between host and guest across ++ /// component-model calls. Pass `SIZE_MAX` to disable the cap. ++ /// Mirrors `Store::set_hostcall_fuel` on the Rust side. ++ void set_hostcall_fuel(size_t fuel) { ++ wasmtime_context_set_hostcall_fuel(ptr, fuel); ++ } ++ + /// Returns the amount of fuel consumed so far by executing WebAssembly. + /// + /// Returns `std::nullopt` if fuel consumption is not enabled. diff --git a/bazel/wasmtime/store_rs.patch b/bazel/wasmtime/store_rs.patch new file mode 100644 index 00000000000..3ac4096d72f --- /dev/null +++ b/bazel/wasmtime/store_rs.patch @@ -0,0 +1,20 @@ +--- a/src/store.rs ++++ b/src/store.rs +@@ -263,6 +263,17 @@ + crate::handle_result(store.set_fuel(fuel), |()| {}) + } + ++/// Sets the per-Store hostcall fuel cap that bounds how much data wasmtime ++/// will copy between the host and the guest across component-model calls. ++/// Pass `usize::MAX` for "unlimited". ++#[unsafe(no_mangle)] ++pub extern "C" fn wasmtime_context_set_hostcall_fuel( ++ mut store: WasmtimeStoreContextMut<'_>, ++ fuel: usize, ++) { ++ store.set_hostcall_fuel(fuel); ++} ++ + #[unsafe(no_mangle)] + pub extern "C" fn wasmtime_context_get_fuel( + store: WasmtimeStoreContext<'_>, diff --git a/buildscripts/resmokelib/configure_resmoke.py b/buildscripts/resmokelib/configure_resmoke.py index 779ca44d16a..b19ce42a79c 100644 --- a/buildscripts/resmokelib/configure_resmoke.py +++ b/buildscripts/resmokelib/configure_resmoke.py @@ -778,9 +778,9 @@ flags in common: {common_set} _config.MONGOD_EXECUTABLE = _expand_user(config.pop("mongod_executable")) - # TODO SERVER-116054, SERVER-116052, SERVER-116055, SERVER-116053: Remove this js_engine handling - # section and_detect_js_engine once mozjs-wasm supports $where, $function, $accumulator, and - # mapReduce, eliminating the need for this startup-time binary invocation. + # TODO SERVER-116054, SERVER-116052, SERVER-116053: Remove this js_engine handling + # section and _detect_js_engine once mozjs-wasm supports $where, $function, and mapReduce, + # eliminating the need for this startup-time binary invocation. _config.JS_ENGINE = _detect_js_engine(_config.MONGOD_EXECUTABLE) if _config.JS_ENGINE == "mozjs-wasm": _config.EXCLUDE_WITH_ANY_TAGS.append("mozjs_wasm_unsupported") @@ -1198,8 +1198,6 @@ _MOZJS_PATTERNS = ( "mapreduce", # TODO SERVER-116054: Add support for $where. '"$where"', - # TODO SERVER-116055: Add support for $accumulator. - '"$accumulator"', ) diff --git a/jstests/aggregation/accumulators/accumulator_js.js b/jstests/aggregation/accumulators/accumulator_js.js index 1155b554078..de7782d6fed 100644 --- a/jstests/aggregation/accumulators/accumulator_js.js +++ b/jstests/aggregation/accumulators/accumulator_js.js @@ -376,3 +376,149 @@ command.pipeline = [ res = assert.commandWorked(db.runCommand(command)); expectedResults = [{_id: 1, value: {len: 3, types: ["object", "object", "object"], values: [null, null, null]}}]; assert(resultsEq(res.cursor.firstBatch, expectedResults), res.cursor); + +// Test that throwing inside accumulate causes the command to fail. +assert(db.accumulator_js.drop()); +assert.commandWorked(db.accumulator_js.insert({val: 1})); +assert.commandFailedWithCode( + db.runCommand({ + aggregate: "accumulator_js", + cursor: {}, + pipeline: [ + { + $group: { + _id: 1, + value: { + $accumulator: { + init: function () { + return 0; + }, + accumulateArgs: ["$val"], + accumulate: function (state, val) { + throw new Error("accumulate error"); + }, + merge: function (s1, s2) { + return s1 + s2; + }, + lang: "js", + }, + }, + }, + }, + ], + }), + ErrorCodes.JSInterpreterFailure, +); + +// Test that init returning null is handled: accumulate receives null as initial state. +assert(db.accumulator_js.drop()); +assert.commandWorked(db.accumulator_js.insert([{val: 10}, {val: 32}])); +res = assert.commandWorked( + db.runCommand({ + aggregate: "accumulator_js", + cursor: {}, + pipeline: [ + { + $group: { + _id: 1, + value: { + $accumulator: { + init: function () { + return null; + }, + accumulateArgs: ["$val"], + accumulate: function (state, val) { + return (state || 0) + val; + }, + merge: function (s1, s2) { + return (s1 || 0) + (s2 || 0); + }, + lang: "js", + }, + }, + }, + }, + ], + }), +); +assert(resultsEq(res.cursor.firstBatch, [{_id: 1, value: 42}]), res.cursor); + +// Test that string state works: accumulate builds a comma-separated string. +assert(db.accumulator_js.drop()); +assert.commandWorked(db.accumulator_js.insert([{word: "hello"}, {word: "world"}])); +res = assert.commandWorked( + db.runCommand({ + aggregate: "accumulator_js", + cursor: {}, + pipeline: [ + { + $sort: {word: 1}, + }, + { + $group: { + _id: 1, + value: { + $accumulator: { + init: function () { + return ""; + }, + accumulateArgs: ["$word"], + accumulate: function (state, val) { + return state ? state + "," + val : val; + }, + merge: function (s1, s2) { + return s1 ? s1 + "," + s2 : s2; + }, + finalize: function (state) { + return state; + }, + lang: "js", + }, + }, + }, + }, + ], + }), +); +assert.eq(res.cursor.firstBatch.length, 1); +// The two words appear in the result separated by a comma (order may vary across shards). +const words = res.cursor.firstBatch[0].value.split(",").sort(); +assert.sameMembers(words, ["hello", "world"]); + +// Test that when no documents match a group, finalize is called on the init state. +assert(db.accumulator_js.drop()); +assert.commandWorked(db.accumulator_js.insert([{x: 1}, {x: 1}])); +res = assert.commandWorked( + db.runCommand({ + aggregate: "accumulator_js", + cursor: {}, + pipeline: [ + {$match: {x: 999}}, + { + $group: { + _id: 1, + value: { + $accumulator: { + init: function () { + return {count: 0, sum: 0}; + }, + accumulateArgs: ["$x"], + accumulate: function (state, val) { + return {count: state.count + 1, sum: state.sum + val}; + }, + merge: function (s1, s2) { + return {count: s1.count + s2.count, sum: s1.sum + s2.sum}; + }, + finalize: function (state) { + return state.count > 0 ? state.sum / state.count : 0; + }, + lang: "js", + }, + }, + }, + }, + ], + }), +); +// No documents matched, so the group produces no output. +assert.eq(res.cursor.firstBatch.length, 0, res.cursor); diff --git a/jstests/aggregation/accumulators/accumulator_js_size_limits.js b/jstests/aggregation/accumulators/accumulator_js_size_limits.js index 35af5f340dd..75ef111fe0c 100644 --- a/jstests/aggregation/accumulators/accumulator_js_size_limits.js +++ b/jstests/aggregation/accumulators/accumulator_js_size_limits.js @@ -2,8 +2,6 @@ // @tags: [ // requires_scripting, // resource_intensive, -// # TODO SERVER-116055: Add support for $accumulate. -// mozjs_wasm_unsupported, // ] const coll = db.accumulator_js_size_limits; @@ -42,7 +40,9 @@ let res = runExample(1, { }, lang: "js", }); -assert.commandFailedWithCode(res, [ErrorCodes.BSONObjectTooLarge, 10334]); +// WASM path: objectwrapper.cpp throws 17260 ("Object size exceeds limit"). +// Legacy MozJS path: may throw 10334 (BSONObjectTooLarge) from BSONObjBuilder. +assert.commandFailedWithCode(res, [ErrorCodes.BSONObjectTooLarge, 17260, 10334]); // Accumulator tries to return BSON larger than 16MB from JS. assert(coll.drop()); @@ -88,6 +88,8 @@ res = runExample(1, { }, lang: "js", }); +// 4545000 is thrown by accumulator_js_reduce.cpp on the server before JS runs, +// so it surfaces identically on legacy and WASM builds. assert.commandFailedWithCode(res, [4545000]); // $group size limit exceeded, and cannot spill. diff --git a/jstests/core/query/system_js_access.js b/jstests/core/query/system_js_access.js index ec1a0a4bb67..ee657709c9e 100644 --- a/jstests/core/query/system_js_access.js +++ b/jstests/core/query/system_js_access.js @@ -12,7 +12,6 @@ // exclude_from_timeseries_crud_passthrough, // # TODO SERVER-116052: Add support for $function. // # TODO SERVER-116054: Add support for $where. -// # TODO SERVER-116055: Add support for $accumulate. // mozjs_wasm_unsupported, // ] diff --git a/src/mongo/db/pipeline/BUILD.bazel b/src/mongo/db/pipeline/BUILD.bazel index 16a15cf3b21..771a5f2484d 100644 --- a/src/mongo/db/pipeline/BUILD.bazel +++ b/src/mongo/db/pipeline/BUILD.bazel @@ -1190,6 +1190,10 @@ mongo_cc_unit_test( ], "//bazel/config:js_engine_none": [ ], + "//bazel/config:js_engine_wasm": [ + "accumulator_js_test.cpp", + "expression_javascript_test.cpp", + ], }], tags = [ "code_coverage_quarantine", diff --git a/src/mongo/db/query/query_tester/main.cpp b/src/mongo/db/query/query_tester/main.cpp index d6f3f6d3866..ac703642685 100644 --- a/src/mongo/db/query/query_tester/main.cpp +++ b/src/mongo/db/query/query_tester/main.cpp @@ -124,14 +124,12 @@ void exitWithError(const int statusCode, const std::string& msg) { // Operators checked (will be removed from this check in the future when mozjs-wasm supports them): // TODO SERVER-116054: Add support for $where. // TODO SERVER-116052: Add support for $function. -// TODO SERVER-116055: Add support for $accumulator. // TODO SERVER-116053: Add support for mapReduce. bool containsUnsupportedJSWasmOperators(const BSONObj& obj) { for (const auto& elem : obj) { const auto fieldName = elem.fieldNameStringData(); if (fieldName == "$where"_sd || fieldName == "$function"_sd || - fieldName == "$accumulator"_sd || fieldName == "mapReduce"_sd || - fieldName == "mapreduce"_sd) { + fieldName == "mapReduce"_sd || fieldName == "mapreduce"_sd) { return true; } if (elem.type() == BSONType::object || elem.type() == BSONType::array) { @@ -151,7 +149,7 @@ bool shouldSkipFile(const QueryFile& currFile, DBClientConnection* conn) { // If the server is running mozjs-wasm, we need to check if any queries contain MozJS // operators, and if so, skip the file since those queries won't run successfully. - // TODO SERVER-116054, SERVER-116052, SERVER-116055, SERVER-116053: Remove this check once + // TODO SERVER-116054, SERVER-116052, SERVER-116053: Remove this check once // mozjs-wasm supports all MozJS operators used in the test files. static constexpr auto kMozJsWasmEngine = "mozjs-wasm"_sd; auto bob = BSONObjBuilder{}; diff --git a/src/mongo/scripting/mozjs/common/objectwrapper.cpp b/src/mongo/scripting/mozjs/common/objectwrapper.cpp index 7df1a69d595..15e01336e78 100644 --- a/src/mongo/scripting/mozjs/common/objectwrapper.cpp +++ b/src/mongo/scripting/mozjs/common/objectwrapper.cpp @@ -34,6 +34,7 @@ #include "mongo/bson/util/builder.h" #include "mongo/platform/decimal128.h" #include "mongo/scripting/js_regex.h" +#include "mongo/scripting/mozjs/common/exception.h" #include "mongo/scripting/mozjs/common/idwrapper.h" #include "mongo/scripting/mozjs/common/runtime.h" #include "mongo/scripting/mozjs/common/types/bson.h" @@ -673,8 +674,6 @@ BSONObj ObjectWrapper::toBSON() { if (frames.size() == 1) { IdWrapper idw(_context, id); - // TODO: check if it's cheaper to just compare with an interned - // string of "_id" rather than with ascii if (idw.isString() && idw.equalsAscii("_id")) { continue; } @@ -687,11 +686,12 @@ BSONObj ObjectWrapper::toBSON() { } const int sizeWithEOO = b.len() + 1 /*EOO*/ - 4 /*BSONObj::Holder ref count*/; - uassert(17260, - str::stream() << "Converting from JavaScript to BSON failed: " - << "Object size " << sizeWithEOO << " exceeds limit of " - << BSONObjMaxInternalSize << " bytes.", - sizeWithEOO <= BSONObjMaxInternalSize); + if (sizeWithEOO > BSONObjMaxInternalSize) { + std::string msg = str::stream() << "Converting from JavaScript to BSON failed: " + << "Object size " << sizeWithEOO << " exceeds limit of " + << BSONObjMaxInternalSize << " bytes."; + uasserted(17260, msg); + } return b.obj(); } diff --git a/src/mongo/scripting/mozjs/wasm/BUILD.bazel b/src/mongo/scripting/mozjs/wasm/BUILD.bazel index 8f884567e58..1c6a3dd4273 100644 --- a/src/mongo/scripting/mozjs/wasm/BUILD.bazel +++ b/src/mongo/scripting/mozjs/wasm/BUILD.bazel @@ -231,6 +231,7 @@ mongo_cc_unit_test( deps = [ ":wasmtime_engine", "//src/mongo:base", + "//src/mongo/db/query:query_knobs", "//src/mongo/scripting:scripting_common", ], ) @@ -245,9 +246,9 @@ mongo_cc_library( "wasmtime_engine.h", "//src/mongo/scripting/mozjs/wasm/scope:scope.h", ], + private_hdrs = ["embedded_wasm_resource.h"], deps = [ ":bridge", - ":embed_mozjs_wasm_obj", "//src/mongo:base", "//src/mongo/db:server_options", "//src/mongo/db:service_context", @@ -255,5 +256,8 @@ mongo_cc_library( "//src/mongo/scripting:scripting_common", "//src/mongo/util/concurrency:spin_lock", "@crates//:wasmtime_c", - ], + ] + select({ + "@platforms//os:windows": [":embed_mozjs_wasm_rc"], + "//conditions:default": [":embed_mozjs_wasm_obj"], + }), ) diff --git a/src/mongo/scripting/mozjs/wasm/bridge/bridge.cpp b/src/mongo/scripting/mozjs/wasm/bridge/bridge.cpp index 72a47fe719e..87da2e9e2f4 100644 --- a/src/mongo/scripting/mozjs/wasm/bridge/bridge.cpp +++ b/src/mongo/scripting/mozjs/wasm/bridge/bridge.cpp @@ -117,6 +117,12 @@ MozJSWasmBridge::MozJSWasmBridge(std::shared_ptr ctx, Options // This is used to signal process killing. storeCtx.set_epoch_deadline(1); + // The default 128 MiB hostcall fuel cap is exhausted by long-running $accumulator / + // mapReduce pipelines that pass multi-megabyte BSON state across WIT calls. + // Disable it here; resource use is already bounded by the linear-memory limiter + // (opts.linearMemoryLimitMB), internalQueryMaxJsEmitBytes, and BSON 16 MiB per object. + storeCtx.set_hostcall_fuel(SIZE_MAX); + wt::WasiConfig wasiConfig; wasiConfig.inherit_stdout(); wasiConfig.inherit_stderr(); diff --git a/src/mongo/scripting/mozjs/wasm/engine/engine.cpp b/src/mongo/scripting/mozjs/wasm/engine/engine.cpp index a8a095a65ab..32fcf91b46d 100644 --- a/src/mongo/scripting/mozjs/wasm/engine/engine.cpp +++ b/src/mongo/scripting/mozjs/wasm/engine/engine.cpp @@ -397,14 +397,6 @@ err_code_t MozJSScriptEngine::invokeFunction(uint64_t handle, return err ? err->code : SM_E_RUNTIME; } - if (_emitByteLimit > 0 && _emitBytesUsed > _emitByteLimit) { - if (err) { - err->code = SM_E_RUNTIME; - set_string(&err->msg, &err->msg_len, "emit() exceeded memory limit"); - } - return SM_E_RUNTIME; - } - // Store return value on global (same key as implscope) so getReturnValueBson can read it. ObjectWrapper(_cx, _global).setValue(kInvokeResult, out); @@ -693,9 +685,15 @@ err_code_t MozJSScriptEngine::setGlobalValue(const char* name, BSONObj MozJSScriptEngine::_emitCallback(const BSONObj& args, void* data) { auto* engine = static_cast(data); + int nArgs = args.nFields(); + if (nArgs != 2) { + constexpr int kEmitArgCountCode = 31220; + uasserted(ErrorCodes::Error(kEmitArgCountCode), "emit takes 2 args"); + } + BSONObjIterator it(args); - BSONElement keyElem = it.more() ? it.next() : BSONElement(); - BSONElement valElem = it.more() ? it.next() : BSONElement(); + BSONElement keyElem = it.next(); + BSONElement valElem = it.next(); BSONObjBuilder b; if (keyElem.type() == BSONType::undefined || keyElem.eoo()) diff --git a/src/mongo/scripting/mozjs/wasm/engine/error.cpp b/src/mongo/scripting/mozjs/wasm/engine/error.cpp index 58005222b31..87c4d6d28f1 100644 --- a/src/mongo/scripting/mozjs/wasm/engine/error.cpp +++ b/src/mongo/scripting/mozjs/wasm/engine/error.cpp @@ -31,6 +31,8 @@ #include "mongo/scripting/mozjs/common/exception.h" +#include + #include "js/ErrorReport.h" #include "js/Exception.h" diff --git a/src/mongo/scripting/mozjs/wasm/scope/scope.cpp b/src/mongo/scripting/mozjs/wasm/scope/scope.cpp index 1b74b91db4e..8216899c3d9 100644 --- a/src/mongo/scripting/mozjs/wasm/scope/scope.cpp +++ b/src/mongo/scripting/mozjs/wasm/scope/scope.cpp @@ -98,6 +98,7 @@ void WasmtimeImplScope::init(const BSONObj* data) { opts.jsHeapLimitMB = static_cast(*_jsHeapLimitMB); } opts.linearMemoryLimitMB = gWasmtimeStoreMemoryLimitMB.load(); + _storeLinearMemBytes = static_cast(opts.linearMemoryLimitMB) * 1024 * 1024; _bridge = std::make_unique(_wasmEngineCtx, opts); bool initialized = _bridge->initialize(); uassert(ErrorCodes::BadValue, "MozJS WASM bridge failed to initialize", initialized); @@ -188,7 +189,13 @@ void WasmtimeImplScope::injectNative(const char* field, NativeFunction func, voi _emitCallbackData = data; // Margin lets WASM buffer one over-limit doc so the host's EmitState sees // it during drain and can throw, instead of WASM silently dropping it. - _bridge->setupEmit(internalQueryMaxJsEmitBytes.load() + BSONObjMaxInternalSize); + const int64_t emitBufBytes = + static_cast(internalQueryMaxJsEmitBytes.load()) + BSONObjMaxInternalSize; + uassert(ErrorCodes::BadValue, + "internalQueryMaxJsEmitBytes exceeds wasmtimeStoreMemoryLimitMB: the emit buffer " + "must fit within the WASM store's linear memory", + emitBufBytes <= _storeLinearMemBytes); + _bridge->setupEmit(emitBufBytes); } BSONObj WasmtimeImplScope::_resolveGlobal(const char* field) const { diff --git a/src/mongo/scripting/mozjs/wasm/scope/scope.h b/src/mongo/scripting/mozjs/wasm/scope/scope.h index 958debdbc30..d009919510e 100644 --- a/src/mongo/scripting/mozjs/wasm/scope/scope.h +++ b/src/mongo/scripting/mozjs/wasm/scope/scope.h @@ -120,6 +120,7 @@ private: const boost::optional _jsHeapLimitMB; std::unique_ptr _bridge; + int64_t _storeLinearMemBytes = 0; DeadlineMonitor _deadlineMonitor; void _drainEmitToCallback(); void _installHelpers(); diff --git a/src/mongo/scripting/mozjs/wasm/scope/scope_test.cpp b/src/mongo/scripting/mozjs/wasm/scope/scope_test.cpp index 4474cc79160..7857777fece 100644 --- a/src/mongo/scripting/mozjs/wasm/scope/scope_test.cpp +++ b/src/mongo/scripting/mozjs/wasm/scope/scope_test.cpp @@ -31,6 +31,7 @@ #include "mongo/bson/bsontypes.h" #include "mongo/bson/bsontypes_util.h" +#include "mongo/db/query/query_execution_knobs_gen.h" #include "mongo/scripting/config_engine_gen.h" #include "mongo/scripting/js_regex.h" #include "mongo/scripting/mozjs/wasm/wasmtime_engine.h" @@ -111,46 +112,6 @@ TEST(WasmtimeScope, FunctionPattern_WithArgs) { ASSERT_EQ(retVal.getIntField("sum"), 42); } -// $accumulator init/accumulate/merge pattern -TEST(WasmtimeScope, AccumulatorPattern_InitAccumulateMerge) { - WasmtimeScriptEngine engine; - std::unique_ptr scope(engine.createScopeForCurrentThread(boost::none)); - ASSERT(scope); - - ScriptingFunction initFn = scope->createFunction("function() { return 0; }"); - ASSERT(initFn != 0); - - ScriptingFunction accFn = scope->createFunction("function(state, val) { return state + val; }"); - ASSERT(accFn != 0); - - ScriptingFunction mergeFn = scope->createFunction("function(s1, s2) { return s1 + s2; }"); - ASSERT(mergeFn != 0); - - // init - BSONObj emptyArgs; - ASSERT_EQ(0, scope->invoke(initFn, &emptyArgs, nullptr, 0)); - double state = scope->getNumber("__returnValue"); - ASSERT_EQ(state, 0.0); - - // accumulate: state + 10 - BSONObj accArgs1 = BSON("0" << state << "1" << 10); - ASSERT_EQ(0, scope->invoke(accFn, &accArgs1, nullptr, 0)); - state = scope->getNumber("__returnValue"); - ASSERT_EQ(state, 10.0); - - // accumulate: state + 20 - BSONObj accArgs2 = BSON("0" << state << "1" << 20); - ASSERT_EQ(0, scope->invoke(accFn, &accArgs2, nullptr, 0)); - state = scope->getNumber("__returnValue"); - ASSERT_EQ(state, 30.0); - - // merge: 30 + 12 - BSONObj mergeArgs = BSON("0" << state << "1" << 12.0); - ASSERT_EQ(0, scope->invoke(mergeFn, &mergeArgs, nullptr, 0)); - double merged = scope->getNumber("__returnValue"); - ASSERT_EQ(merged, 42.0); -} - // mapReduce.map pattern: injectNative("emit", ...) + invoke(func, nullptr, &doc, timeout, true) TEST(WasmtimeScope, MapReducePattern_EmitAndDrain) { WasmtimeScriptEngine engine; @@ -736,6 +697,33 @@ TEST(WasmtimeScope, MemoryLimit_ResetWithInvalidParamsFails) { ASSERT_THROWS_CODE(scope->reset(), DBException, ErrorCodes::BadValue); } +// Emit buffer must fit within the WASM store's linear memory. +TEST(WasmtimeScope, EmitBufferExceedsStoreLimitFails) { + auto savedHeap = gJSHeapLimitMB.load(); + auto savedStore = gWasmtimeStoreMemoryLimitMB.load(); + auto savedEmit = internalQueryMaxJsEmitBytes.load(); + ON_BLOCK_EXIT([&] { + gJSHeapLimitMB.store(savedHeap); + gWasmtimeStoreMemoryLimitMB.store(savedStore); + internalQueryMaxJsEmitBytes.store(savedEmit); + }); + + // heap=64 MB, overhead=max(64, 6)=64 MB → min store=128 MB for scope init to pass. + // With store=128 MB, emitBuf=128MB+16MB=144MB > 128MB → injectNative throws BadValue. + gJSHeapLimitMB.store(64); + gWasmtimeStoreMemoryLimitMB.store(128); + internalQueryMaxJsEmitBytes.store(128 * 1024 * 1024); + + WasmtimeScriptEngine engine; + std::unique_ptr scope(engine.createScopeForCurrentThread(boost::none)); + ASSERT(scope); + + ASSERT_THROWS_CODE(scope->injectNative( + "emit", [](const BSONObj&, void*) { return BSONObj(); }, nullptr), + DBException, + ErrorCodes::BadValue); +} + // --- OOM detection --- // hasOutOfMemoryException() starts false and is set when an OOM occurs. diff --git a/src/mongo/scripting/mozjs/wasm/wasmtime_engine.cpp b/src/mongo/scripting/mozjs/wasm/wasmtime_engine.cpp index 0324298ce53..0b493d9b27c 100644 --- a/src/mongo/scripting/mozjs/wasm/wasmtime_engine.cpp +++ b/src/mongo/scripting/mozjs/wasm/wasmtime_engine.cpp @@ -36,17 +36,11 @@ #include "mongo/scripting/config_gen.h" #include "mongo/scripting/engine.h" #include "mongo/scripting/mozjs/wasm/bridge/bridge.h" +#include "mongo/scripting/mozjs/wasm/embedded_wasm_resource.h" #include "mongo/scripting/mozjs/wasm/scope/scope.h" #define MONGO_LOGV2_DEFAULT_COMPONENT ::mongo::logv2::LogComponent::kQuery -// Symbols produced by objcopy from the AOT-compiled mozjs_wasm_api.cwasm. -// See embed_mozjs_wasm_obj in BUILD.bazel. -extern "C" { -extern const uint8_t _binary_mozjs_wasm_api_cwasm_start[]; -extern const uint8_t _binary_mozjs_wasm_api_cwasm_end[]; -} - namespace mongo { bool isExternalScriptingEnabled() { @@ -88,9 +82,8 @@ WasmtimeScriptEngine::WasmtimeScriptEngine() {} WasmtimeScriptEngine::~WasmtimeScriptEngine() {} std::shared_ptr WasmtimeScriptEngine::createWasmEngineContext() const { - size_t size = - static_cast(_binary_mozjs_wasm_api_cwasm_end - _binary_mozjs_wasm_api_cwasm_start); - return wasm::WasmEngineContext::createFromPrecompiled(_binary_mozjs_wasm_api_cwasm_start, size); + auto [data, size] = wasm::getEmbeddedWasmResource(); + return wasm::WasmEngineContext::createFromPrecompiled(data, size); } mongo::Scope* WasmtimeScriptEngine::createScope() {