e2ce6d8c5061c48ff5fd6ebd73d2322618b778be
[folly.git] / folly / experimental / JemallocNodumpAllocator.cpp
1 /*
2  * Copyright 2017 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/Malloc.h>
21 #include <folly/String.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   hooks->alloc = &JemallocNodumpAllocator::alloc;
93 #endif
94
95   return true;
96 #else // FOLLY_JEMALLOC_NODUMP_ALLOCATOR_SUPPORTED
97   return false;
98 #endif // FOLLY_JEMALLOC_NODUMP_ALLOCATOR_SUPPORTED
99 }
100
101 void* JemallocNodumpAllocator::allocate(size_t size) {
102   return mallocx != nullptr ? mallocx(size, flags_) : malloc(size);
103 }
104
105 void* JemallocNodumpAllocator::reallocate(void* p, size_t size) {
106   return rallocx != nullptr ? rallocx(p, size, flags_) : realloc(p, size);
107 }
108
109 #ifdef FOLLY_JEMALLOC_NODUMP_ALLOCATOR_SUPPORTED
110
111 #ifdef FOLLY_JEMALLOC_NODUMP_ALLOCATOR_CHUNK
112 chunk_alloc_t* JemallocNodumpAllocator::original_alloc_ = nullptr;
113 void* JemallocNodumpAllocator::alloc(
114     void* chunk,
115 #else
116 extent_alloc_t* JemallocNodumpAllocator::original_alloc_ = nullptr;
117 void* JemallocNodumpAllocator::alloc(
118     extent_hooks_t* extent,
119     void* new_addr,
120 #endif
121     size_t size,
122     size_t alignment,
123     bool* zero,
124     bool* commit,
125     unsigned arena_ind) {
126   void* result = original_alloc_(
127       JEMALLOC_CHUNK_OR_EXTENT,
128 #ifdef FOLLY_JEMALLOC_NODUMP_ALLOCATOR_EXTENT
129       new_addr,
130 #endif
131       size,
132       alignment,
133       zero,
134       commit,
135       arena_ind);
136   if (result != nullptr) {
137     if (auto ret = madvise(result, size, MADV_DONTDUMP)) {
138       VLOG(1) << "Unable to madvise(MADV_DONTDUMP): " << errnoStr(ret);
139     }
140   }
141
142   return result;
143 }
144
145 #endif // FOLLY_JEMALLOC_NODUMP_ALLOCATOR_SUPPORTED
146
147 void JemallocNodumpAllocator::deallocate(void* p) {
148   dallocx != nullptr ? dallocx(p, flags_) : free(p);
149 }
150
151 void JemallocNodumpAllocator::deallocate(void* p, void* userData) {
152   const uint64_t flags = reinterpret_cast<uint64_t>(userData);
153   dallocx != nullptr ? dallocx(p, static_cast<int>(flags)) : free(p);
154 }
155
156 JemallocNodumpAllocator& globalJemallocNodumpAllocator() {
157   static auto instance = new JemallocNodumpAllocator();
158   return *instance;
159 }
160
161 } // folly