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 <unistd.h>
18 #include <sys/types.h>
19 #include <cstring>
20 #include <cstdlib>
21 #include <iostream>
22 #include "tscore/ink_memory.h"
23 #include "tscore/ink_error.h"
24 #include "tscore/ink_assert.h"
25 #include "tscore/ink_align.h"
26 #include "tscore/JeAllocator.h"
27 #include "tscore/Diags.h"
28 
29 namespace jearena
30 {
31 #if JEMALLOC_NODUMP_ALLOCATOR_SUPPORTED
32 thread_local extent_alloc_t *JemallocNodumpAllocator::original_alloc = nullptr;
33 thread_local extent_hooks_t JemallocNodumpAllocator::extent_hooks;
34 thread_local int JemallocNodumpAllocator::arena_flags = 0;
35 
36 void *
alloc(extent_hooks_t * extent,void * new_addr,size_t size,size_t alignment,bool * zero,bool * commit,unsigned int arena_id)37 JemallocNodumpAllocator::alloc(extent_hooks_t *extent, void *new_addr, size_t size, size_t alignment, bool *zero, bool *commit,
38                                unsigned int arena_id)
39 {
40   void *result = original_alloc(extent, new_addr, size, alignment, zero, commit, arena_id);
41 
42   if (result != nullptr) {
43     // Seems like we don't really care if the advice went through
44     // in the original code, so just keeping it the same here.
45     ats_madvise((caddr_t)result, size, MADV_DONTDUMP);
46   }
47 
48   return result;
49 }
50 
51 int
extend_and_setup_arena()52 JemallocNodumpAllocator::extend_and_setup_arena()
53 {
54   unsigned int arena_id;
55   size_t arena_id_len = sizeof(arena_id);
56   if (auto ret = mallctl("arenas.create", &arena_id, &arena_id_len, nullptr, 0)) {
57     ink_abort("Unable to extend arena: %s", std::strerror(ret));
58   }
59 
60   int flags = MALLOCX_ARENA(arena_id) | MALLOCX_TCACHE_NONE;
61 
62   // Read the existing hooks
63   const auto key = "arena." + std::to_string(arena_id) + ".extent_hooks";
64   extent_hooks_t *hooks;
65   size_t hooks_len = sizeof(hooks);
66   if (auto ret = mallctl(key.c_str(), &hooks, &hooks_len, nullptr, 0)) {
67     ink_abort("Unable to get the hooks: %s", std::strerror(ret));
68   }
69 
70   // Set the custom hook
71   original_alloc     = hooks->alloc;
72   extent_hooks       = *hooks;
73   extent_hooks.alloc = &JemallocNodumpAllocator::alloc;
74 
75   extent_hooks_t *new_hooks = &extent_hooks;
76   if (auto ret = mallctl(key.c_str(), nullptr, nullptr, &new_hooks, sizeof(new_hooks))) {
77     ink_abort("Unable to set the hooks: %s", std::strerror(ret));
78   }
79 
80   Debug("JeAllocator", "arena \"%ud\" created with flags \"%d\"", arena_id, flags);
81 
82   arena_flags = flags;
83 
84   return flags;
85 }
86 #endif /* JEMALLOC_NODUMP_ALLOCATOR_SUPPORTED */
87 
88 /**
89  * This will retain the original functionality if
90  * !defined(JEMALLOC_NODUMP_ALLOCATOR_SUPPORTED)
91  */
92 void *
allocate(InkFreeList * f)93 JemallocNodumpAllocator::allocate(InkFreeList *f)
94 {
95 #if JEMALLOC_NODUMP_ALLOCATOR_SUPPORTED
96   int flags = arena_flags;
97   if (unlikely(flags == 0)) {
98     flags = extend_and_setup_arena();
99   }
100 #endif
101 
102   void *newp = nullptr;
103   if (f->advice) {
104 #if JEMALLOC_NODUMP_ALLOCATOR_SUPPORTED
105     if (likely(f->type_size > 0)) {
106       flags |= MALLOCX_ALIGN(f->alignment);
107       if (unlikely((newp = mallocx(f->type_size, flags)) == nullptr)) {
108         ink_abort("couldn't allocate %u bytes", f->type_size);
109       }
110     }
111 #else
112     newp = ats_memalign(f->alignment, f->type_size);
113     if (INK_ALIGN((uint64_t)newp, ats_pagesize()) == (uint64_t)newp) {
114       ats_madvise(static_cast<caddr_t>(newp), INK_ALIGN(f->type_size, f->alignment), f->advice);
115     }
116 #endif
117   } else {
118     newp = ats_memalign(f->alignment, f->type_size);
119   }
120   return newp;
121 }
122 
123 JemallocNodumpAllocator &
globalJemallocNodumpAllocator()124 globalJemallocNodumpAllocator()
125 {
126   static auto instance = new JemallocNodumpAllocator();
127   return *instance;
128 }
129 } // namespace jearena
130