Implement jemalloc extent hooks memory tracker (#1250)

Should improve/fix memory usage exceeds --memory-limit issues
This commit is contained in:
Antonio Filipovic
2023-10-23 12:48:26 +02:00
committed by GitHub
parent 26e31ca06f
commit 7f7f3adfcb
17 changed files with 988 additions and 100 deletions

View File

@@ -23,6 +23,7 @@
#include "glue/run_id.hpp"
#include "helpers.hpp"
#include "license/license_sender.hpp"
#include "memory/memory_control.hpp"
#include "query/config.hpp"
#include "query/discard_value_stream.hpp"
#include "query/interpreter.hpp"
@@ -108,6 +109,7 @@ void InitSignalHandlers(const std::function<void()> &shutdown_fun) {
}
int main(int argc, char **argv) {
memgraph::memory::SetHooks();
google::SetUsageMessage("Memgraph database server");
gflags::SetVersionString(version_string);
@@ -199,7 +201,6 @@ int main(int argc, char **argv) {
"won't be available.");
}
}
std::cout << "You are running Memgraph v" << gflags::VersionString() << std::endl;
std::cout << "To get started with Memgraph, visit https://memgr.ph/start" << std::endl;

View File

@@ -2,7 +2,9 @@ set(memory_src_files
new_delete.cpp
memory_control.cpp)
find_package(Jemalloc REQUIRED)
find_package(jemalloc REQUIRED)
add_library(mg-memory STATIC ${memory_src_files})
target_link_libraries(mg-memory mg-utils fmt)

View File

@@ -1,4 +1,4 @@
// Copyright 2022 Memgraph Ltd.
// Copyright 2023 Memgraph Ltd.
//
// Use of this software is governed by the Business Source License
// included in the file licenses/BSL.txt; by using this file, you agree to be bound by the terms of the Business Source
@@ -10,6 +10,8 @@
// licenses/APL.txt.
#include "memory_control.hpp"
#include "utils/logging.hpp"
#include "utils/memory_tracker.hpp"
#if USE_JEMALLOC
#include <jemalloc/jemalloc.h>
@@ -22,6 +24,179 @@ namespace memgraph::memory {
// NOLINTNEXTLINE(cppcoreguidelines-macro-usage)
#define STRINGIFY(x) STRINGIFY_HELPER(x)
#if USE_JEMALLOC
static void *my_alloc(extent_hooks_t *extent_hooks, void *new_addr, size_t size, size_t alignment, bool *zero,
bool *commit, unsigned arena_ind);
static bool my_dalloc(extent_hooks_t *extent_hooks, void *addr, size_t size, bool committed, unsigned arena_ind);
static void my_destroy(extent_hooks_t *extent_hooks, void *addr, size_t size, bool committed, unsigned arena_ind);
static bool my_commit(extent_hooks_t *extent_hooks, void *addr, size_t size, size_t offset, size_t length,
unsigned arena_ind);
static bool my_decommit(extent_hooks_t *extent_hooks, void *addr, size_t size, size_t offset, size_t length,
unsigned arena_ind);
static bool my_purge_forced(extent_hooks_t *extent_hooks, void *addr, size_t size, size_t offset, size_t length,
unsigned arena_ind);
extent_hooks_t *old_hooks = nullptr;
static extent_hooks_t custom_hooks = {
.alloc = &my_alloc,
.dalloc = &my_dalloc,
.destroy = &my_destroy,
.commit = &my_commit,
.decommit = &my_decommit,
.purge_lazy = nullptr,
.purge_forced = &my_purge_forced,
.split = nullptr,
.merge = nullptr,
};
static const extent_hooks_t *new_hooks = &custom_hooks;
void *my_alloc(extent_hooks_t *extent_hooks, void *new_addr, size_t size, size_t alignment, bool *zero, bool *commit,
unsigned arena_ind) {
// This needs to be before, to throw exception in case of too big alloc
if (*commit) [[likely]] {
memgraph::utils::total_memory_tracker.Alloc(static_cast<int64_t>(size));
}
auto *ptr = old_hooks->alloc(extent_hooks, new_addr, size, alignment, zero, commit, arena_ind);
if (ptr == nullptr) [[unlikely]] {
if (*commit) {
memgraph::utils::total_memory_tracker.Free(static_cast<int64_t>(size));
}
return ptr;
}
return ptr;
}
static bool my_dalloc(extent_hooks_t *extent_hooks, void *addr, size_t size, bool committed, unsigned arena_ind) {
auto err = old_hooks->dalloc(extent_hooks, addr, size, committed, arena_ind);
if (err) [[unlikely]] {
return err;
}
if (committed) [[likely]] {
memgraph::utils::total_memory_tracker.Free(static_cast<int64_t>(size));
}
return false;
}
static void my_destroy(extent_hooks_t *extent_hooks, void *addr, size_t size, bool committed, unsigned arena_ind) {
if (committed) [[likely]] {
memgraph::utils::total_memory_tracker.Free(static_cast<int64_t>(size));
}
old_hooks->destroy(extent_hooks, addr, size, committed, arena_ind);
}
static bool my_commit(extent_hooks_t *extent_hooks, void *addr, size_t size, size_t offset, size_t length,
unsigned arena_ind) {
auto err = old_hooks->commit(extent_hooks, addr, size, offset, length, arena_ind);
if (err) {
return err;
}
memgraph::utils::total_memory_tracker.Alloc(static_cast<int64_t>(length));
return false;
}
static bool my_decommit(extent_hooks_t *extent_hooks, void *addr, size_t size, size_t offset, size_t length,
unsigned arena_ind) {
MG_ASSERT(old_hooks && old_hooks->decommit);
auto err = old_hooks->decommit(extent_hooks, addr, size, offset, length, arena_ind);
if (err) {
return err;
}
memgraph::utils::total_memory_tracker.Free(static_cast<int64_t>(length));
return false;
}
static bool my_purge_forced(extent_hooks_t *extent_hooks, void *addr, size_t size, size_t offset, size_t length,
unsigned arena_ind) {
MG_ASSERT(old_hooks && old_hooks->purge_forced);
auto err = old_hooks->purge_forced(extent_hooks, addr, size, offset, length, arena_ind);
if (err) [[unlikely]] {
return err;
}
memgraph::utils::total_memory_tracker.Free(static_cast<int64_t>(length));
return false;
}
#endif
void SetHooks() {
#if USE_JEMALLOC
uint64_t allocated{0};
uint64_t sz{sizeof(allocated)};
sz = sizeof(unsigned);
unsigned n_arenas{0};
int err = mallctl("opt.narenas", (void *)&n_arenas, &sz, nullptr, 0);
if (err) {
return;
}
if (nullptr != old_hooks) {
return;
}
for (int i = 0; i < n_arenas; i++) {
std::string func_name = "arena." + std::to_string(i) + ".extent_hooks";
size_t hooks_len = sizeof(old_hooks);
int err = mallctl(func_name.c_str(), &old_hooks, &hooks_len, nullptr, 0);
if (err) {
LOG_FATAL("Error getting hooks for jemalloc arena {}", i);
}
// Due to the way jemalloc works, we need first to set their hooks
// which will trigger creating arena, then we can set our custom hook wrappers
err = mallctl(func_name.c_str(), nullptr, nullptr, &old_hooks, sizeof(old_hooks));
MG_ASSERT(old_hooks);
MG_ASSERT(old_hooks->alloc);
MG_ASSERT(old_hooks->dalloc);
MG_ASSERT(old_hooks->destroy);
MG_ASSERT(old_hooks->commit);
MG_ASSERT(old_hooks->decommit);
MG_ASSERT(old_hooks->purge_forced);
MG_ASSERT(old_hooks->purge_lazy);
MG_ASSERT(old_hooks->split);
MG_ASSERT(old_hooks->merge);
custom_hooks.purge_lazy = old_hooks->purge_lazy;
custom_hooks.split = old_hooks->split;
custom_hooks.merge = old_hooks->merge;
if (err) {
LOG_FATAL("Error setting jemalloc hooks for jemalloc arena {}", i);
}
err = mallctl(func_name.c_str(), nullptr, nullptr, &new_hooks, sizeof(new_hooks));
if (err) {
LOG_FATAL("Error setting custom hooks for jemalloc arena {}", i);
}
}
#endif
}
void PurgeUnusedMemory() {
#if USE_JEMALLOC
mallctl("arena." STRINGIFY(MALLCTL_ARENAS_ALL) ".purge", nullptr, nullptr, nullptr, 0);
@@ -30,4 +205,5 @@ void PurgeUnusedMemory() {
#undef STRINGIFY
#undef STRINGIFY_HELPER
} // namespace memgraph::memory

View File

@@ -1,4 +1,4 @@
// Copyright 2022 Memgraph Ltd.
// Copyright 2023 Memgraph Ltd.
//
// Use of this software is governed by the Business Source License
// included in the file licenses/BSL.txt; by using this file, you agree to be bound by the terms of the Business Source
@@ -11,6 +11,11 @@
#pragma once
#include <cstddef>
#include "utils/logging.hpp"
namespace memgraph::memory {
void PurgeUnusedMemory();
void SetHooks();
} // namespace memgraph::memory

View File

@@ -1,4 +1,4 @@
// Copyright 2022 Memgraph Ltd.
// Copyright 2023 Memgraph Ltd.
//
// Use of this software is governed by the Business Source License
// included in the file licenses/BSL.txt; by using this file, you agree to be bound by the terms of the Business Source
@@ -87,21 +87,15 @@ void deleteSized(void *ptr, const std::size_t /*unused*/, const std::align_val_t
#endif
void TrackMemory(std::size_t size) {
#if USE_JEMALLOC
if (size != 0) [[likely]] {
size = nallocx(size, 0);
}
#endif
#if !USE_JEMALLOC
memgraph::utils::total_memory_tracker.Alloc(static_cast<int64_t>(size));
#endif
}
void TrackMemory(std::size_t size, const std::align_val_t align) {
#if USE_JEMALLOC
if (size != 0) [[likely]] {
size = nallocx(size, MALLOCX_ALIGN(align)); // NOLINT(hicpp-signed-bitwise)
}
#endif
#if !USE_JEMALLOC
memgraph::utils::total_memory_tracker.Alloc(static_cast<int64_t>(size));
#endif
}
bool TrackMemoryNoExcept(const std::size_t size) {
@@ -126,11 +120,7 @@ bool TrackMemoryNoExcept(const std::size_t size, const std::align_val_t align) {
void UntrackMemory([[maybe_unused]] void *ptr, [[maybe_unused]] std::size_t size = 0) noexcept {
try {
#if USE_JEMALLOC
if (ptr != nullptr) [[likely]] {
memgraph::utils::total_memory_tracker.Free(sallocx(ptr, 0));
}
#else
#if !USE_JEMALLOC
if (size) {
memgraph::utils::total_memory_tracker.Free(static_cast<int64_t>(size));
} else {
@@ -144,11 +134,7 @@ void UntrackMemory([[maybe_unused]] void *ptr, [[maybe_unused]] std::size_t size
void UntrackMemory(void *ptr, const std::align_val_t align, [[maybe_unused]] std::size_t size = 0) noexcept {
try {
#if USE_JEMALLOC
if (ptr != nullptr) [[likely]] {
memgraph::utils::total_memory_tracker.Free(sallocx(ptr, MALLOCX_ALIGN(align))); // NOLINT(hicpp-signed-bitwise)
}
#else
#if !USE_JEMALLOC
if (size) {
memgraph::utils::total_memory_tracker.Free(static_cast<int64_t>(size));
} else {