94db0e7204c203940c8281d1657d1c5f0a253c9a
[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("arenas.extend", &arena_index_, &len, nullptr, 0)) {
41     LOG(FATAL) << "Unable to extend arena: " << errnoStr(ret);
42   }
43   flags_ = MALLOCX_ARENA(arena_index_) | MALLOCX_TCACHE_NONE;
44
45   // Set the custom alloc hook
46   const auto key =
47       folly::to<std::string>("arena.", arena_index_, ".chunk_hooks");
48   chunk_hooks_t hooks;
49   len = sizeof(hooks);
50   // Read the existing hooks
51   if (auto ret = mallctl(key.c_str(), &hooks, &len, nullptr, 0)) {
52     LOG(FATAL) << "Unable to get the hooks: " << errnoStr(ret);
53   }
54   if (original_chunk_alloc_ == nullptr) {
55     original_chunk_alloc_ = hooks.alloc;
56   } else {
57     DCHECK_EQ(original_chunk_alloc_, hooks.alloc);
58   }
59
60   // Set the custom hook
61   hooks.alloc = &JemallocNodumpAllocator::chunk_alloc;
62   if (auto ret =
63           mallctl(key.c_str(), nullptr, nullptr, &hooks, sizeof(hooks))) {
64     LOG(FATAL) << "Unable to set the hooks: " << errnoStr(ret);
65   }
66
67   return true;
68 #else // FOLLY_JEMALLOC_NODUMP_ALLOCATOR_SUPPORTED
69   return false;
70 #endif // FOLLY_JEMALLOC_NODUMP_ALLOCATOR_SUPPORTED
71 }
72
73 void* JemallocNodumpAllocator::allocate(size_t size) {
74   return mallocx != nullptr ? mallocx(size, flags_) : malloc(size);
75 }
76
77 void* JemallocNodumpAllocator::reallocate(void* p, size_t size) {
78   return rallocx != nullptr ? rallocx(p, size, flags_) : realloc(p, size);
79 }
80
81 #ifdef FOLLY_JEMALLOC_NODUMP_ALLOCATOR_SUPPORTED
82
83 chunk_alloc_t* JemallocNodumpAllocator::original_chunk_alloc_ = nullptr;
84
85 void* JemallocNodumpAllocator::chunk_alloc(
86     void* chunk,
87     size_t size,
88     size_t alignment,
89     bool* zero,
90     bool* commit,
91     unsigned arena_ind) {
92   void* result =
93       original_chunk_alloc_(chunk, size, alignment, zero, commit, arena_ind);
94   if (result != nullptr) {
95     if (auto ret = madvise(result, size, MADV_DONTDUMP)) {
96       VLOG(1) << "Unable to madvise(MADV_DONTDUMP): " << errnoStr(ret);
97     }
98   }
99
100   return result;
101 }
102
103 #endif // FOLLY_JEMALLOC_NODUMP_ALLOCATOR_SUPPORTED
104
105 void JemallocNodumpAllocator::deallocate(void* p) {
106   dallocx != nullptr ? dallocx(p, flags_) : free(p);
107 }
108
109 void JemallocNodumpAllocator::deallocate(void* p, void* userData) {
110   const uint64_t flags = reinterpret_cast<uint64_t>(userData);
111   dallocx != nullptr ? dallocx(p, flags) : free(p);
112 }
113
114 JemallocNodumpAllocator& globalJemallocNodumpAllocator() {
115   static auto instance = new JemallocNodumpAllocator();
116   return *instance;
117 }
118
119 } // folly