1/*
2 * Copyright 2016-present Facebook, Inc.
3 *
4 * Licensed under the Apache License, Version 2.0 (the "License");
5 * you may not use this file except in compliance with the License.
6 * You may obtain a copy of the License at
7 *
8 * http://www.apache.org/licenses/LICENSE-2.0
9 *
10 * Unless required by applicable law or agreed to in writing, software
11 * distributed under the License is distributed on an "AS IS" BASIS,
12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 * See the License for the specific language governing permissions and
14 * limitations under the License.
15 */
16
17#include <folly/experimental/JemallocNodumpAllocator.h>
18
19#include <folly/Conv.h>
20#include <folly/String.h>
21#include <folly/memory/Malloc.h>
22#include <glog/logging.h>
23
24namespace folly {
25
26JemallocNodumpAllocator::JemallocNodumpAllocator(State state) {
27 if (state == State::ENABLED && extend_and_setup_arena()) {
28 LOG(INFO) << "Set up arena: " << arena_index_;
29 }
30}
31
32bool JemallocNodumpAllocator::extend_and_setup_arena() {
33#ifdef FOLLY_JEMALLOC_NODUMP_ALLOCATOR_SUPPORTED
34 if (mallctl == nullptr) {
35 // Not linked with jemalloc.
36 return false;
37 }
38
39 size_t len = sizeof(arena_index_);
40 if (auto ret = mallctl(
41#ifdef FOLLY_JEMALLOC_NODUMP_ALLOCATOR_CHUNK
42 "arenas.extend"
43#else
44 "arenas.create"
45#endif
46 ,
47 &arena_index_,
48 &len,
49 nullptr,
50 0)) {
51 LOG(FATAL) << "Unable to extend arena: " << errnoStr(ret);
52 }
53 flags_ = MALLOCX_ARENA(arena_index_) | MALLOCX_TCACHE_NONE;
54
55#ifdef FOLLY_JEMALLOC_NODUMP_ALLOCATOR_CHUNK
56 const auto key =
57 folly::to<std::string>("arena.", arena_index_, ".chunk_hooks");
58 chunk_hooks_t hooks;
59 len = sizeof(hooks);
60 // Read the existing hooks
61 if (auto ret = mallctl(key.c_str(), &hooks, &len, nullptr, 0)) {
62 LOG(FATAL) << "Unable to get the hooks: " << errnoStr(ret);
63 }
64 if (original_alloc_ == nullptr) {
65 original_alloc_ = hooks.alloc;
66 } else {
67 DCHECK_EQ(original_alloc_, hooks.alloc);
68 }
69
70 // Set the custom hook
71 hooks.alloc = &JemallocNodumpAllocator::alloc;
72 if (auto ret =
73 mallctl(key.c_str(), nullptr, nullptr, &hooks, sizeof(hooks))) {
74 LOG(FATAL) << "Unable to set the hooks: " << errnoStr(ret);
75 }
76#else
77 const auto key =
78 folly::to<std::string>("arena.", arena_index_, ".extent_hooks");
79 extent_hooks_t* hooks;
80 len = sizeof(hooks);
81 // Read the existing hooks
82 if (auto ret = mallctl(key.c_str(), &hooks, &len, nullptr, 0)) {
83 LOG(FATAL) << "Unable to get the hooks: " << errnoStr(ret);
84 }
85 if (original_alloc_ == nullptr) {
86 original_alloc_ = hooks->alloc;
87 } else {
88 DCHECK_EQ(original_alloc_, hooks->alloc);
89 }
90
91 // Set the custom hook
92 extent_hooks_ = *hooks;
93 extent_hooks_.alloc = &JemallocNodumpAllocator::alloc;
94 extent_hooks_t* new_hooks = &extent_hooks_;
95 if (auto ret = mallctl(
96 key.c_str(), nullptr, nullptr, &new_hooks, sizeof(new_hooks))) {
97 LOG(FATAL) << "Unable to set the hooks: " << errnoStr(ret);
98 }
99#endif
100
101 return true;
102#else // FOLLY_JEMALLOC_NODUMP_ALLOCATOR_SUPPORTED
103 return false;
104#endif // FOLLY_JEMALLOC_NODUMP_ALLOCATOR_SUPPORTED
105}
106
107void* JemallocNodumpAllocator::allocate(size_t size) {
108 return mallocx != nullptr ? mallocx(size, flags_) : malloc(size);
109}
110
111void* JemallocNodumpAllocator::reallocate(void* p, size_t size) {
112 return rallocx != nullptr ? rallocx(p, size, flags_) : realloc(p, size);
113}
114
115#ifdef FOLLY_JEMALLOC_NODUMP_ALLOCATOR_SUPPORTED
116
117#ifdef FOLLY_JEMALLOC_NODUMP_ALLOCATOR_CHUNK
118chunk_alloc_t* JemallocNodumpAllocator::original_alloc_ = nullptr;
119void* JemallocNodumpAllocator::alloc(
120 void* chunk,
121#else
122extent_hooks_t JemallocNodumpAllocator::extent_hooks_;
123extent_alloc_t* JemallocNodumpAllocator::original_alloc_ = nullptr;
124void* JemallocNodumpAllocator::alloc(
125 extent_hooks_t* extent,
126 void* new_addr,
127#endif
128 size_t size,
129 size_t alignment,
130 bool* zero,
131 bool* commit,
132 unsigned arena_ind) {
133 void* result = original_alloc_(
134 JEMALLOC_CHUNK_OR_EXTENT,
135#ifdef FOLLY_JEMALLOC_NODUMP_ALLOCATOR_EXTENT
136 new_addr,
137#endif
138 size,
139 alignment,
140 zero,
141 commit,
142 arena_ind);
143 if (result != nullptr) {
144 if (auto ret = madvise(result, size, MADV_DONTDUMP)) {
145 VLOG(1) << "Unable to madvise(MADV_DONTDUMP): " << errnoStr(ret);
146 }
147 }
148
149 return result;
150}
151
152#endif // FOLLY_JEMALLOC_NODUMP_ALLOCATOR_SUPPORTED
153
154void JemallocNodumpAllocator::deallocate(void* p, size_t) {
155 dallocx != nullptr ? dallocx(p, flags_) : free(p);
156}
157
158void JemallocNodumpAllocator::deallocate(void* p, void* userData) {
159 const uint64_t flags = reinterpret_cast<uint64_t>(userData);
160 dallocx != nullptr ? dallocx(p, static_cast<int>(flags)) : free(p);
161}
162
163JemallocNodumpAllocator& globalJemallocNodumpAllocator() {
164 static auto instance = new JemallocNodumpAllocator();
165 return *instance;
166}
167
168} // namespace folly
169