SERVER-113893: Improve memory handling in mozjs allocator (#45077)

GitOrigin-RevId: 453ee2003fd56823649f6d16f15fe5edbaadd2a9
This commit is contained in:
Anna Veselova 2025-12-15 09:27:19 -06:00 committed by MongoDB Bot
parent 9a7c840349
commit 61d51036f9
3 changed files with 247 additions and 80 deletions

View File

@ -250,6 +250,7 @@ mongo_cc_unit_test(
srcs = [
"//src/mongo/scripting/mozjs:asan_handles_test.cpp",
"//src/mongo/scripting/mozjs:implscope_test.cpp",
"//src/mongo/scripting/mozjs:jscustomallocator_test.cpp",
"//src/mongo/scripting/mozjs:module_loader_test.cpp",
],
copts = select({

View File

@ -46,7 +46,7 @@
#elif defined(__FreeBSD__)
#include <malloc_np.h>
#else
#define MONGO_NO_MALLOC_USABLE_SIZE
#error "Unsupported platform"
#endif
/**
@ -74,19 +74,6 @@ namespace {
thread_local size_t total_bytes = 0;
thread_local size_t max_bytes = 0;
/**
* When we don't have malloc_usable_size, we manage by adjusting our pointer by
* kMaxAlign bytes and storing the size of the allocation kMaxAlign bytes
* behind the pointer we hand back. That let's us get to the value at runtime.
* We know kMaxAlign is enough (generally 8 or 16 bytes), because that's
* literally the contract between malloc and std::max_align_t.
*
* This is commented out right now because std::max_align_t didn't seem to be
* available on our solaris builder. TODO: revisit in the future to see if that
* still holds.
*/
// const size_t kMaxAlign = std::alignment_of<std::max_align_t>::value;
const size_t kMaxAlign = 16;
} // namespace
size_t get_total_bytes() {
@ -102,13 +89,12 @@ size_t get_max_bytes() {
return max_bytes;
}
size_t get_current(void* ptr);
/**
* Wraps std::Xalloc functions
*
* The idea here is to abstract soft limits on allocations, as well as possibly
* necessary pointer adjustment (if we don't have a malloc_usable_size
* replacement).
*
* The idea here is to abstract soft limits on allocations.
*/
template <typename T>
void* wrap_alloc(T&& func, void* ptr, size_t bytes) {
@ -123,26 +109,16 @@ void* wrap_alloc(T&& func, void* ptr, size_t bytes) {
// for the SharedImmutableStringsCache (order 600). This triggered a failure of a MOZ_ASSERT
// which enforces correct lock ordering in the JS engine. For this reason, we avoid checking
// for an OOM here if we are requesting zero bytes (i.e freeing memory).
if (mb && bytes && (tb + bytes > mb)) {
auto scope = mongo::mozjs::MozJSImplScope::getThreadScope();
if (scope) {
scope->setOOM();
return nullptr;
}
// We fall through here because we want to let spidermonkey continue
// with whatever it was doing. Calling setOOM will fail the top level
// operation as soon as possible.
return nullptr;
}
#ifdef MONGO_NO_MALLOC_USABLE_SIZE
ptr = ptr ? static_cast<char*>(ptr) - kMaxAlign : nullptr;
#endif
#ifdef MONGO_NO_MALLOC_USABLE_SIZE
void* p = func(ptr, bytes + kMaxAlign);
#else
void* p = func(ptr, bytes);
#endif
#if __has_feature(address_sanitizer)
{
@ -171,23 +147,16 @@ void* wrap_alloc(T&& func, void* ptr, size_t bytes) {
return nullptr;
}
#ifdef MONGO_NO_MALLOC_USABLE_SIZE
*reinterpret_cast<size_t*>(p) = bytes;
p = static_cast<char*>(p) + kMaxAlign;
#endif
total_bytes = tb + bytes;
total_bytes += mongo::sm::get_current(p);
return p;
}
/*
* gets the current available size at this pointer, which may be larger than the allocated size.
* this size is not valid to access unless realloc() is called
*/
size_t get_current(void* ptr) {
#ifdef MONGO_NO_MALLOC_USABLE_SIZE
if (!ptr)
return 0;
return *reinterpret_cast<size_t*>(static_cast<char*>(ptr) - kMaxAlign);
#elif defined(__linux__) || defined(__FreeBSD__)
#if defined(__linux__) || defined(__FreeBSD__)
return malloc_usable_size(ptr);
#elif defined(__APPLE__)
return malloc_size(ptr);
@ -205,44 +174,51 @@ JS_PUBLIC_DATA arena_id_t js::MallocArena;
JS_PUBLIC_DATA arena_id_t js::ArrayBufferContentsArena;
JS_PUBLIC_DATA arena_id_t js::StringBufferArena;
void* mongo_arena_malloc(arena_id_t arena, size_t bytes) {
void* mongo_arena_malloc(size_t bytes) {
return std::malloc(bytes);
}
void* mongo_arena_calloc(arena_id_t arena, size_t nmemb, size_t size) {
return std::calloc(nmemb, size);
void* mongo_arena_calloc(size_t bytes) {
return std::calloc(bytes, 1);
}
void* mongo_arena_realloc(arena_id_t arena, void* p, size_t bytes) {
void* mongo_arena_realloc(void* p, size_t bytes) {
MOZ_ASSERT(bytes != 0); // realloc() with zero size is unsupported
if (!p) {
return mongo_arena_malloc(arena, bytes);
return mongo_arena_malloc(bytes);
}
if (!bytes) {
js_free(p);
// Like in mongo_free, this count is imprecise because get_current calls malloc_usable_size
// which can return a larger value than the initial allocation
size_t current = mongo::sm::get_current(p);
void* ptr = std::realloc(p, bytes);
if (!ptr) {
// on failure the old ptr isn't freed, no need to adjust total_bytes
return nullptr;
}
size_t current = mongo::sm::get_current(p);
mongo::sm::total_bytes -= std::min(mongo::sm::total_bytes, current);
return ptr;
}
if (current >= bytes) {
return p;
}
void* mongo_free(void* ptr) {
// Note that malloc_usable_size and equivalents can return a larger size than the allocated
// buffer, so this may result in undercounting
size_t current = mongo::sm::get_current(ptr);
size_t tb = mongo::sm::total_bytes;
mongo::sm::total_bytes -= std::min(mongo::sm::total_bytes, current);
if (tb >= current) {
mongo::sm::total_bytes = tb - current;
}
return std::realloc(p, bytes);
std::free(ptr);
return nullptr;
}
void* js_arena_malloc(size_t arena, size_t bytes) {
JS_OOM_POSSIBLY_FAIL();
JS_CHECK_LARGE_ALLOC(bytes);
return mongo::sm::wrap_alloc(
[&](void* ptr, size_t b) { return mongo_arena_malloc(arena, bytes); }, nullptr, bytes);
[](void* ptr, size_t b) { return mongo_arena_malloc(b); }, nullptr, bytes);
}
void* js_malloc(size_t bytes) {
@ -253,16 +229,14 @@ void* js_arena_calloc(arena_id_t arena, size_t bytes) {
JS_OOM_POSSIBLY_FAIL();
JS_CHECK_LARGE_ALLOC(bytes);
return mongo::sm::wrap_alloc(
[&](void* ptr, size_t b) { return mongo_arena_calloc(arena, 1, b); }, nullptr, bytes);
[](void* ptr, size_t b) { return mongo_arena_calloc(b); }, nullptr, bytes);
}
void* js_arena_calloc(arena_id_t arena, size_t nmemb, size_t size) {
JS_OOM_POSSIBLY_FAIL();
JS_CHECK_LARGE_ALLOC(size);
return mongo::sm::wrap_alloc(
[&](void* ptr, size_t b) { return mongo_arena_calloc(arena, nmemb, size); },
nullptr,
size * nmemb);
[](void* ptr, size_t b) { return mongo_arena_calloc(b); }, nullptr, size * nmemb);
}
void* js_calloc(size_t bytes) {
@ -282,7 +256,7 @@ void* js_arena_realloc(arena_id_t arena, void* p, size_t bytes) {
JS_OOM_POSSIBLY_FAIL();
JS_CHECK_LARGE_ALLOC(bytes);
return mongo::sm::wrap_alloc(
[&](void* ptr, size_t b) { return mongo_arena_realloc(arena, ptr, b); }, p, bytes);
[](void* ptr, size_t b) { return mongo_arena_realloc(ptr, b); }, p, bytes);
}
void* js_realloc(void* p, size_t bytes) {
@ -293,20 +267,7 @@ void js_free(void* p) {
if (!p)
return;
size_t current = mongo::sm::get_current(p);
size_t tb = mongo::sm::get_total_bytes();
if (tb >= current) {
mongo::sm::total_bytes = tb - current;
}
mongo::sm::wrap_alloc(
[](void* ptr, size_t b) {
std::free(ptr);
return nullptr;
},
p,
0);
mongo::sm::wrap_alloc([](void* ptr, size_t b) { return mongo_free(ptr); }, p, 0);
}
void js::InitMallocAllocator() {

View File

@ -0,0 +1,205 @@
/**
* Copyright (C) 2022-present MongoDB, Inc.
*
* This program is free software: you can redistribute it and/or modify
* it under the terms of the Server Side Public License, version 1,
* as published by MongoDB, Inc.
*
* This program is distributed in the hope that it will be useful,
* but WITHOUT ANY WARRANTY; without even the implied warranty of
* MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
* Server Side Public License for more details.
*
* You should have received a copy of the Server Side Public License
* along with this program. If not, see
* <http://www.mongodb.com/licensing/server-side-public-license>.
*
* As a special exception, the copyright holders give permission to link the
* code of portions of this program with the OpenSSL library under certain
* conditions as described in each individual source file and distribute
* linked combinations including the program with the OpenSSL library. You
* must comply with the Server Side Public License in all respects for
* all of the code used other than as permitted herein. If you modify file(s)
* with this exception, you may extend this exception to your version of the
* file(s), but you are not obligated to do so. If you do not wish to do so,
* delete this exception statement from your version. If you delete this
* exception statement from all source files in the program, then also delete
* it in the license file.
*/
#include "mongo/base/string_data.h"
#include "mongo/scripting/engine.h"
#include "mongo/unittest/unittest.h"
#include <memory>
#include <string>
#include <jscustomallocator.h>
#include <boost/filesystem/operations.hpp>
#include <boost/filesystem/path.hpp>
#include <fmt/format.h>
namespace mongo {
namespace mozjs {
class JSCustomAllocatorTest : public unittest::Test {
protected:
void setUp() override {
mongo::sm::reset(0);
// Note: get_total_bytes() is an estimate and won't exactly count allocated bytes,
// test should only compare to 0
ASSERT_EQUALS(mongo::sm::get_total_bytes(), 0);
}
void tearDown() override {
mongo::sm::reset(0);
}
};
TEST_F(JSCustomAllocatorTest, MallocUpToLimit) {
mongo::sm::reset(100);
ASSERT_EQUALS(mongo::sm::get_total_bytes(), 0);
void* ptr1 = js_malloc(20);
ASSERT_NOT_EQUALS(ptr1, nullptr);
ASSERT_GREATER_THAN(mongo::sm::get_total_bytes(), 0);
void* ptr2 = js_malloc(60);
ASSERT_NOT_EQUALS(ptr2, nullptr);
ASSERT_EQUALS(js_malloc(50), nullptr);
js_free(ptr2);
ptr2 = nullptr;
void* ptr3 = js_malloc(50);
ASSERT_NOT_EQUALS(ptr3, nullptr);
js_free(ptr1);
js_free(ptr3);
ASSERT_EQUALS(mongo::sm::get_total_bytes(), 0);
}
TEST_F(JSCustomAllocatorTest, ReallocUpToLimit) {
mongo::sm::reset(100);
ASSERT_EQUALS(mongo::sm::get_total_bytes(), 0);
void* ptr1 = js_malloc(20);
ASSERT_NOT_EQUALS(ptr1, nullptr);
ASSERT_GREATER_THAN(mongo::sm::get_total_bytes(), 0);
void* ptr2 = js_realloc(ptr1, 40);
ASSERT_NOT_EQUALS(ptr2, nullptr);
void* ptr3 = js_realloc(ptr2, 200);
ASSERT_EQUALS(ptr3, nullptr);
js_free(ptr2);
ASSERT_EQUALS(mongo::sm::get_total_bytes(), 0);
}
TEST_F(JSCustomAllocatorTest, CallocUpToLimit) {
mongo::sm::reset(100);
ASSERT_EQUALS(mongo::sm::get_total_bytes(), 0);
void* ptr1 = js_calloc(10, 2);
ASSERT_NOT_EQUALS(ptr1, nullptr);
ASSERT_GREATER_THAN(mongo::sm::get_total_bytes(), 0);
void* ptr2 = js_calloc(10, 6);
ASSERT_NOT_EQUALS(ptr2, nullptr);
ASSERT_EQUALS(js_calloc(50), nullptr);
js_free(ptr2);
ptr2 = nullptr;
void* ptr3 = js_calloc(50);
ASSERT_NOT_EQUALS(ptr3, nullptr);
js_free(ptr1);
js_free(ptr3);
ASSERT_EQUALS(mongo::sm::get_total_bytes(), 0);
}
// control for whether allocations are happening in a query or just during setup/teardown
TEST_F(JSCustomAllocatorTest, SetupScriptEngine) {
mongo::ScriptEngine::setup(ExecutionEnvironment::TestRunner);
std::unique_ptr<mongo::Scope> scope(mongo::getGlobalScriptEngine()->newScope());
scope.reset();
setGlobalScriptEngine(nullptr);
}
// SERVER-113893 trigger js_arena_malloc and verify that entire buffer is valid
TEST_F(JSCustomAllocatorTest, SingleAlloc) {
mongo::ScriptEngine::setup(ExecutionEnvironment::TestRunner);
std::unique_ptr<mongo::Scope> scope(mongo::getGlobalScriptEngine()->newScope());
std::string codeStr = R"(
const buf = new ArrayBuffer(128);
const view = new Uint8Array(buf);
for (let i = 0; i < 128; i++) view[i] = i;
)";
StringData code(codeStr);
ASSERT_DOES_NOT_THROW(scope->exec(code,
"root_module",
true /* printResult */,
true /* reportError */,
true /* assertOnError , timeout*/));
scope.reset();
setGlobalScriptEngine(nullptr);
}
TEST_F(JSCustomAllocatorTest, ResizeMany) {
mongo::ScriptEngine::setup(ExecutionEnvironment::TestRunner);
std::unique_ptr<mongo::Scope> scope(mongo::getGlobalScriptEngine()->newScope());
std::string codeStr = R"(
let buffers = [];
const numBuffers = 100;
for (let i = 0; i < numBuffers; i++) {
const bufferSize = i * i;
const buf = new ArrayBuffer(bufferSize);
buffers.push(buf);
// access entire buffer to make sure memory is valid
const view = new Uint8Array(buf);
for (let j = 0; j < bufferSize; j++) {
view[j] = j;
}
}
// resize in reverse order so half get bigger and half get smaller
for (let i = numBuffers - 1; i <= 0; i--) {
const bufferSize = (numBuffers - i) * (numBuffers - i) + 3;
buffers[i].resize(bufferSize);
// access entire buffer to make sure memory is valid
const view = new Uint8Array(buf);
for (let j = 0; j < bufferSize; j++) {
view[j] = j;
}
}
)";
StringData code(codeStr);
ASSERT_DOES_NOT_THROW(scope->exec(code,
"root_module",
true /* printResult */,
true /* reportError */,
true /* assertOnError , timeout*/));
scope.reset();
setGlobalScriptEngine(nullptr);
}
} // namespace mozjs
} // namespace mongo