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 | |
24 | namespace folly { |
25 | |
26 | JemallocNodumpAllocator::JemallocNodumpAllocator(State state) { |
27 | if (state == State::ENABLED && extend_and_setup_arena()) { |
28 | LOG(INFO) << "Set up arena: " << arena_index_; |
29 | } |
30 | } |
31 | |
32 | bool 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 | |
107 | void* JemallocNodumpAllocator::allocate(size_t size) { |
108 | return mallocx != nullptr ? mallocx(size, flags_) : malloc(size); |
109 | } |
110 | |
111 | void* 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 |
118 | chunk_alloc_t* JemallocNodumpAllocator::original_alloc_ = nullptr; |
119 | void* JemallocNodumpAllocator::alloc( |
120 | void* chunk, |
121 | #else |
122 | extent_hooks_t JemallocNodumpAllocator::extent_hooks_; |
123 | extent_alloc_t* JemallocNodumpAllocator::original_alloc_ = nullptr; |
124 | void* 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 | |
154 | void JemallocNodumpAllocator::deallocate(void* p, size_t) { |
155 | dallocx != nullptr ? dallocx(p, flags_) : free(p); |
156 | } |
157 | |
158 | void 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 | |
163 | JemallocNodumpAllocator& globalJemallocNodumpAllocator() { |
164 | static auto instance = new JemallocNodumpAllocator(); |
165 | return *instance; |
166 | } |
167 | |
168 | } // namespace folly |
169 | |