1 //  Copyright (c) 2011-present, Facebook, Inc.  All rights reserved.
2 //  This source code is licensed under both the GPLv2 (found in the
3 //  COPYING file in the root directory) and Apache 2.0 License
4 //  (found in the LICENSE.Apache file in the root directory).
5 
6 #pragma once
7 
8 #include <atomic>
9 #include <vector>
10 
11 #include "port/jemalloc_helper.h"
12 #include "port/port.h"
13 #include "rocksdb/memory_allocator.h"
14 #include "util/thread_local.h"
15 
16 #if defined(ROCKSDB_JEMALLOC) && defined(ROCKSDB_PLATFORM_POSIX)
17 
18 #include <sys/mman.h>
19 
20 #if (JEMALLOC_VERSION_MAJOR >= 5) && defined(MADV_DONTDUMP)
21 #define ROCKSDB_JEMALLOC_NODUMP_ALLOCATOR
22 
23 namespace ROCKSDB_NAMESPACE {
24 
25 class JemallocNodumpAllocator : public MemoryAllocator {
26  public:
27   JemallocNodumpAllocator(JemallocAllocatorOptions& options,
28                           std::unique_ptr<extent_hooks_t>&& arena_hooks,
29                           unsigned arena_index);
30   ~JemallocNodumpAllocator();
31 
Name()32   const char* Name() const override { return "JemallocNodumpAllocator"; }
33   void* Allocate(size_t size) override;
34   void Deallocate(void* p) override;
35   size_t UsableSize(void* p, size_t allocation_size) const override;
36 
37  private:
38   friend Status NewJemallocNodumpAllocator(
39       JemallocAllocatorOptions& options,
40       std::shared_ptr<MemoryAllocator>* memory_allocator);
41 
42   // Custom alloc hook to replace jemalloc default alloc.
43   static void* Alloc(extent_hooks_t* extent, void* new_addr, size_t size,
44                      size_t alignment, bool* zero, bool* commit,
45                      unsigned arena_ind);
46 
47   // Destroy arena on destruction of the allocator, or on failure.
48   static Status DestroyArena(unsigned arena_index);
49 
50   // Destroy tcache on destruction of the allocator, or thread exit.
51   static void DestroyThreadSpecificCache(void* ptr);
52 
53   // Get or create tcache. Return flag suitable to use with `mallocx`:
54   // either MALLOCX_TCACHE_NONE or MALLOCX_TCACHE(tc).
55   int GetThreadSpecificCache(size_t size);
56 
57   // A function pointer to jemalloc default alloc. Use atomic to make sure
58   // NewJemallocNodumpAllocator is thread-safe.
59   //
60   // Hack: original_alloc_ needs to be static for Alloc() to access it.
61   // alloc needs to be static to pass to jemalloc as function pointer.
62   static std::atomic<extent_alloc_t*> original_alloc_;
63 
64   const JemallocAllocatorOptions options_;
65 
66   // Custom hooks has to outlive corresponding arena.
67   const std::unique_ptr<extent_hooks_t> arena_hooks_;
68 
69   // Arena index.
70   const unsigned arena_index_;
71 
72   // Hold thread-local tcache index.
73   ThreadLocalPtr tcache_;
74 };
75 
76 }  // namespace ROCKSDB_NAMESPACE
77 #endif  // (JEMALLOC_VERSION_MAJOR >= 5) && MADV_DONTDUMP
78 #endif  // ROCKSDB_JEMALLOC && ROCKSDB_PLATFORM_POSIX
79