Guest User

Untitled

a guest
Feb 27th, 2020
151
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
  1. #include "toyalloc.hpp"
  2.  
  3. #include <twist/stdlike/atomic.hpp>
  4. #include <twist/support/compiler.hpp>
  5. #include <twist/support/locking.hpp>
  6.  
  7. #include <pthread.h>
  8. #include <sys/types.h>
  9. #include <unistd.h>
  10.  
  11. using twist::MemSpan;
  12. using twist::MmapAllocation;
  13.  
  14. namespace toyalloc {
  15.  
  16. static const size_t kBlockSize = 4096;
  17.  
  18. struct BlockNode {
  19.   BlockNode* next_;
  20. };
  21.  
  22. class Allocator {
  23.  public:
  24.   friend void ForkPrepare(void);
  25.   friend void ChildAfterFork(void);
  26.   friend void ParentAfterFork(void);
  27.   void Init(MmapAllocation arena) {
  28.     arena_ = std::move(arena);
  29.     list_head = arena_.Start();
  30.     char* current_ptr = reinterpret_cast<char*>(list_head);
  31.     for (size_t i = 0; i < arena_.Size() / kBlockSize; ++i) {
  32.       void** next_ptr = reinterpret_cast<void**>(current_ptr);
  33.       current_ptr += kBlockSize;
  34.       if (current_ptr < arena_.End()) {
  35.         *next_ptr = current_ptr;
  36.       } else {
  37.         *next_ptr = nullptr;
  38.       }
  39.     }
  40.   }
  41.  
  42.   MemSpan GetArena() const {
  43.     return arena_.AsMemSpan();
  44.   }
  45.  
  46.   void* Allocate() {
  47. //    printf("%s %d ", "allocate begin", getpid());
  48. //    fflush(stdout);
  49. //    auto lck = twist::LockUnique(mutex);
  50.       mutex.lock();
  51. //    printf("%s %d ", "mutex lock", getpid());
  52. //    fflush(stdout);
  53.     void* result = nullptr;
  54.  
  55.     if (list_head != nullptr) {
  56.       result = list_head;
  57.       list_head = *reinterpret_cast<void**>(list_head);
  58.     }
  59. //    printf("%s %d \n", "finish", getpid());
  60. //    fflush(stdout);
  61. mutex.unlock();
  62.     return result;
  63.   }
  64.  
  65.   void Free(void* addr) {
  66. //    printf("%s %d ", "free begin", getpid());
  67. //    fflush(stdout);
  68. //    auto lck = twist::LockUnique(mutex);
  69. //    printf("%s %d ", "mutex lock", getpid());
  70. //    fflush(stdout);
  71. mutex.lock();
  72.     *reinterpret_cast<void**>(addr) = list_head;
  73.     list_head = addr;
  74.     TWIST_UNUSED(addr);
  75.     mutex.unlock();
  76. //    printf("%s %d \n", "finish", getpid());
  77. //    fflush(stdout);
  78.   }
  79.  
  80.  private:
  81.   MmapAllocation arena_;
  82.   void* list_head = nullptr;
  83.   std::mutex mutex;
  84. };
  85.  
  86. /////////////////////////////////////////////////////////////////////
  87.  
  88. static Allocator allocator;
  89.  
  90. /////////////////////////////////////////////////////////////////////
  91.  
  92. void ForkPrepare(void) {
  93. //  auto lck = twist::LockUnique(allocator.mutex);
  94.     allocator.mutex.lock();
  95. //  printf("%s %d \n", "prepare", getpid());
  96. //  fflush(stdout);
  97. }
  98.  
  99. void ChildAfterFork(void) {
  100.   allocator.mutex.unlock();
  101. //  printf("%s %d \n", "after", getpid());
  102. //  fflush(stdout);
  103. }
  104.  
  105. void ParentAfterFork(void) {
  106.   allocator.mutex.unlock();
  107.   //  printf("%s %d \n", "after", getpid());
  108.   //  fflush(stdout);
  109. }
  110.  
  111. void Init(MmapAllocation arena) {
  112.   allocator.Init(std::move(arena));
  113.   pthread_atfork(&ForkPrepare, &ParentAfterFork,
  114.                  &ChildAfterFork);
  115. }
  116.  
  117. MemSpan GetArena() {
  118.   return allocator.GetArena();
  119. }
  120.  
  121. size_t GetBlockSize() {
  122.   return kBlockSize;
  123. }
  124.  
  125. void* Allocate() {
  126.   return allocator.Allocate();
  127. }
  128.  
  129. void Free(void* addr) {
  130.   allocator.Free(addr);
  131. }
  132.  
  133. }  // namespace toyalloc
RAW Paste Data