abcdefGets

ゲッツ!

mallocを再実装した話

C++ AdventCalendarの12日目

普段私はWEBのフロントエンドを仕事にしている。
つまり使う言語はjavascript/typescript等のScript言語だ。
ただ前職や趣味、OSS等でC++によく触っていたので昔実装したmallocの話をすることにした。

mallocとは

mallocとはC言語のstdlib.hに含まれるメモリ割り当て関数のことで、
C++やその他の多くの言語で内部的に利用されている。
ヒープを割り当てる方法はいくつかあるが、このmallocがもっともメジャーといえるだろう。

mallocを再実装した

今回はmallocを自分で再実装してちょっと早くした話を書く。
再実装した理由は色々あるが最も大きな理由はただの好奇心。
yatscというtypescriptコンパイラC++で書こうと思って実装を始めたときに作った。
ただしyatsc自体は未完で飽きて終了。
作ってみて思ったが、実際にプロダクトで使う場合はシンプルなメモリプールとか作ったほうがよっぽど利口だと思う。

Inside malloc

malloc関数は現在のモダンな環境ではmmapを利用してメモリを確保、それを各チャンクに分けて管理しているはず。
細かいことは以下のスライドに詳しい。

www.slideshare.net

Direction of implementations

今回はGoogleのTLMallocの実装と上記のスライドをベースにjust-fitアロケータを作成することにした。

基本的な構成は以下の通り。
メインスレッドにあるCentralArenaが各スレッドのLocalArenaを管理しており、
それぞれのArenaはTls(Thread Local Storage)に格納されている。
また、各LocalArenaはスレッドが終了すると、free-listに格納されて空き状態となり、
新たにスレッドが生成された場合に再利用される。
LocalArena自体もlinked-listになっており、nextポインタで各スレッドのLocalArenaはつながっている。
さらにLocalArenaは内部にChunkを持ち、それらはChunkHeaderと後続の複数のHeapHeaderによって管理されている。

Tlsにヒープを格納することでスレッド毎にヒープ領域を分割し、スレッド競合を防ぎつつjust-fitアロケータでメモリの無駄も防ぐ。
というのが、このmalloc実装のキーポイントになる。
ちなみにTlsを実装するのは面倒だったので、AndroidTls実装をコピーした。
Tlsのコールバックまで実装されており非常に便利。サンキューGoogle

以下はCentralArenaからLocalArenaまでの図

+----------------+  +---------------+
|  CentralArena  |--|  Main Thread  |
+----------------+  +---------------+
      |     |  +------------+  +--------------+
      |     +--|  Thread-1  |--|  LocalArena  |
      |     |  +------------+  +--------------+
      |     |  +------------+  +--------------+
      |     +--|  Thread-2  |--|  LocalArena  |
      |     |  +------------+  +--------------+
      |     |  +------------+  +--------------+
      |     +--|  Thread-3  |--|  LocalArena  |
      |        +------------+  +--------------+
      +------------+  +--------------+  +--------------+
      |  free list |--|  LocalArena  |--|  LocalArena  |-- ...
      +------------+  +--------------+  +--------------+

CentralArena

CentralArenaはメモリのアロケーションを要求されると、 GetLocalArena関数を呼び出して、TlsからLocalArenaを取り出すか空いているLocalArenaをfree-listから探し出して割り当てる。

以下のような感じでAtomicにlockを行ってスレッド競合を防いでいる。

LocalArena* CentralArena::FindUnlockedArena() YATSC_NOEXCEPT {
  // If released_arena_count_ is zero,
  // that mean free arena is not exists.
  // So we simply allocate new LocalArena.
  if (released_arena_count_.load() == 0) {
    return StoreNewLocalArena();
  }

  LocalArena* current = local_arena_.load(std::memory_order_relaxed);

  while (current != nullptr) {
    // Released arena is found.
    if (current->AcquireLock()) {
      return current;
    }
    current = current->next();
  }

  // If released arena is not found,
  // allocate new LocalArena.
  return StoreNewLocalArena();
}

ちなみにAcquireLockは以下のような感じ。Atomicにロックしている。

YATSC_INLINE bool AcquireLock() YATSC_NOEXCEPT {
  bool ret = !lock_.test_and_set();
  if (ret) {
    central_arena_->NotifyLockAcquired();
  }
  return ret;
}

さて上記FindUnlockedArenaを使って空いたLocalArenaを探し出すのに失敗した場合、 そこで初めてメモリを割り当てることになる。
このメモリ割り当ては以下のような感じで行う。

LocalArena* CentralArena::StoreNewLocalArena() YATSC_NOEXCEPT {
  Byte* block = reinterpret_cast<Byte*>(
      GetInternalHeap(sizeof(LocalArena) + sizeof(InternalHeap)));

  LocalArena* arena = new (block) LocalArena(
      this, new(block + sizeof(LocalArena)) InternalHeap(
          block + sizeof(LocalArena) + sizeof(InternalHeap)));
  arena->AcquireLock();

  LocalArena* arena_head = local_arena_.load(std::memory_order_acquire);

  // Run CAS operation.
  // This case ABA problem is not the matter,
  // because the local_arena_ allowed only one-way operation
  // that is append new LocalArena.
  // So head of local_arena_ is change only when new local arena is appended.
  do {
    arena->set_next(arena_head);
  } while (!local_arena_.compare_exchange_weak(arena_head, arena));

  return arena;
}

// Get heap space that used to allocate LocalArena and ChunkHeader.
YATSC_INLINE void* CentralArena::GetInternalHeap(size_t additional) {
  return VirtualHeapAllocator::Map(
      nullptr, sizeof(ChunkHeader) * kChunkHeaderAllocatableCount + additional,
      VirtualHeapAllocator::Prot::WRITE | VirtualHeapAllocator::Prot::READ,
      VirtualHeapAllocator::Flags::ANONYMOUS
      | VirtualHeapAllocator::Flags::PRIVATE);
}

解説すると、まず最初にVirtualHeapAllocatorを利用してメモリを割り当てる。
このVirtualHeapAllocatormmapのプラットフォーム毎の差分を吸収したクラス。
このVirtualHeapAllocatorからLocalArenaInternalHeapという内部向けの管理ヘッダ分のメモリを割り当てる。
それをinplacement newを利用してLocalArenaに割り当てる。
つまりLocalArenaは以下のようなメモリレイアウトになる。

  sizeof(InternalHeap)            + sizeof(LocalArena)
+-------------------------------+   +--------------+
|  InternalHeap(hidden header)  |---|  LocalArena  |
+-------------------------------+   +--------------+

あとは割り当てたLocalArenaのロックを獲得して、local-arenaのリストに格納する。
このリストに追加するのも当然atomicに行わなければならないので、CAS(Compare And Swap)を利用してリストに追加する。
CASを使うとABA問題を気にしなければならないが、コメントにある通りlocal_arena_のリストは追加しか行わないのでABA問題は気にしなくて良い。
これでめでたくCentralArenaからロックフリーでLocalArenaを確保することに成功した。
次はLocalArenaを実装する。

LocalArena

LocalArenaこそがメモリ割り当てのコアになっていて、この部分が一番面倒くさい。
LocalArenaはメモリ割り当て要求を受けるとsmall_bin_とよばれる小さなメモリ専用のツリーにメモリを割り当てていく。
今回はこのツリーはRedBlackTreeを利用した。
このsmall_bin_に要求メモリサイズ毎のリストを作りそこに割り当てていく。

+--------------+
|  LocalArena  |
+--------------+
       |
       |        +--------------+
       +--------|  small_bin_  |
                +--------------+
                        |
            +-----------------------+
            |                       |
        +-------+               +-------+
        | 1byte |               | 2byte |
        +-------+               +-------+
            |                       |          
      +-----------+           +-----------+    
      |           |           |           |    
  +-------+   +-------+   +-------+   +-------+
  | 3byte |   | 4byte |   | 5byte |   | 6byte |
  +-------+   +-------+   +-------+   +-------+

こんな感じでsmall_bin_には各サイズ毎のchunkのリストが格納される。
これでjust-fitなアロケータになる。
ただし面倒だったのは、このsmall_bin_をstd::mapとかで実装するのはちょっと難しく、
結局自分でRBTreeを実装する羽目になった。
何故かと言うと、今はメモリアロケーションの途中なのでヒープから値を確保することは不可能で、
各Chunkに追加のメモリ割り当てで格納できるコンテナが必要だったからである。
自分で実装したRBTreeはIntrusiveRBTreeというクラスで、ある条件を満たせばヒープからメモリのアロケーションをしなくてもツリーを構成することができる。 その条件とは格納される値自身がコンテナの実装をしていること。 つまり格納される値がRbTreeNodeというRedBlackTreeのNodeを継承していれば、
RBTree自身はそれをつなぎ合わせるアルゴリズムの実装のみでよく、
メモリ割り当てを必要としないで済む。

ただし、small_bin_は名の通り小さなメモリのみを格納する場所なので、64KBを制限とし、
それを超える場合には直接CentralArenaから巨大なメモリ領域をmmapで割り当てる。
そのときにはSpinLockでロックをかけるので当然遅くなる。

これでjust-fitなbinが実装できた。
さてsmall_bin_に格納されるメモリを実装していく。
small_bin_に格納されるメモリはChunkHeaderという管理ヘッダと実際に値を割り当てる後続のブロックで構成される。

ChunkHeader

ChunkHeaderクラスは実際に値を割り当てるヒープを管理している。
ChunkHeaderの後ろには各Heapを管理するHeapHeaderがつながっている。 確保されているヒープはheap_list_に格納され、
空いているヒープはfree_list_に格納される。

+---------------+
|  ChunkHeader  |
+---------------+
     |     |
     |     +--------------+  +------------+  +------------+  +------------+
     |     |  heap_list_  |--| HeapHeader |--| HeapHeader |--| HeapHeader |
     |     +--------------+  +------------+  +------------+  +------------+
     | 
     |
     +--------------+  +------------+  +------------+  +------------+
     |  free_list_  |--| HeapHeader |--| HeapHeader |--| HeapHeader |
     +--------------+  +------------+  +------------+  +------------+

そしてChunkHeaderはメモリ確保の要求が来ると、最初にfree_list_を探索する。
もしfree_list_にHeapHeaderがあればそれをheap_list_にAtomicに繋いで返す。
空きがなければ、最後につながれたHeapHeaderの使用量を確認して、まだ格納できるのであればHeapHeaderに値を追加する。
もしHeapHeaderに空きが無ければ、新たにHeapHeaderを割り当てる。

YATSC_INLINE void* ChunkHeader::Distribute() {
  auto free_list_head = free_list_.load(std::memory_order_acquire);
  if (free_list_head == nullptr) {
    // If heap is not filled.
    if (heap_list_->used() < max_allocatable_size_) {
      // Calc next positon.
      Byte* block = reinterpret_cast<Byte*>(heap_list_) + heap_list_->used();
      // Update heap used value.
      heap_list_->set_used(heap_list_->used() + size_class_);
      return block;
    } else {
      // If heap is exhausted allocate new memory.
      Byte* block = InitHeap(AllocateBlock());
      auto heap_header = reinterpret_cast<HeapHeader*>(
          block - sizeof(HeapHeader));
      heap_header->set_used(heap_header->used() + size_class_);
      return block;
    }
  }

  // Swap free list.
  while (!free_list_.compare_exchange_weak(
             free_list_head, free_list_head->next())) {}
  return reinterpret_cast<void*>(free_list_head);
}

HeapHeader

HeapHeaderクラスは実際にヒープの値を直接管理しているクラス。
このHeapHeaderChunkHeaderから割り当てられる、1 MBでアライメントされたメモリブロックとなっている。
なぜ1 MBでアライメントするかというと、メモリの節約のためにHeapHeaderが確保しているメモリブロックは管理ヘッダを持っていない。
本来はメモリを削除する場合にHeapHeaderを参照するための情報が必要なのだが、1 MBでアライメントすることによって、
メモリの下位16ビットをマスクするだけで、HeapHeaderの先頭アドレスを取得することができるようになる。
こうすることで1割当ごとにX64なら最低でも8 byte、x86なら4 byteの節約になる。

CentralArenaにメモリのfree要求が来た場合には以下のようにポインタのアドレス下位16ビットをマスクして、
HeapHeaderを取得する。

ChunkHeader::kAddrMask~0xFFFF

auto h = reinterpret_cast<HeapHeader*>(
    reinterpret_cast<uintptr_t>(ptr) & ChunkHeader::kAddrMask);
ASSERT(true, h != nullptr);
ChunkHeader* chunk_header = h->chunk_header();
ASSERT(true, chunk_header != nullptr);
chunk_header->Dealloc(ptr);

構造は以下の図の通り

+----------------------------+  +----------------------+  +-----------------------+
|  HeapHeader(0x103e600000)  |--|  value(0x103e600008) |--|  value(0x103e000010)  | 
+----------------------------+  +----------------------+  +-----------------------+

削除する場合

+----------------------+                                +----------------------------+
|  value(0x103e600008) | => 0x103e600008             => |  HeapHeader(0x103e600000)  |
+----------------------+           ^^^^^ Mask heare.    +----------------------------+

概ねこんな感じでjust-fitでロックを必要としないメモリアロケータを実装することができた。
パフォーマンスだが、シングルスレッドではmallocの約5倍ほど高速化することができたが、
マルチスレッドでは1.2 ~ 1.1倍ほどしか高速化せずちょっとなぁ〜という感じ。
また64 KB以上のビックな割当を行うとほとんどmallocと変わらないのでここは改善の余地がありそう。

という感じでmallocを再実装した話でした。

全体像は以下のようになる。

+----------------+  +---------------+
|  CentralArena  |--|  Main Thread  |
+----------------+  +---------------+
      |     |  +------------+  +--------------+
      |     +--|  Thread-1  |--|  LocalArena  |
      |     |  +------------+  +--------------+
      |     |                         |                                                  
      |     |                         |        +--------------+                
      |     |                         ---------|  small_bin_  |                
      |     |                                  +--------------+                
      |     |                                          |                       
      |     |                              ------------+------------           
      |     |                              |                       |           
      |     |                          +-------+               +-------+       
      |     |                          | 1byte |               | 2byte |       
      |     |                          +-------+               +-------+       
      |     |                              |                       |           
      |     |                        ------+------           ------+------     
      |     |                        |           |           |           |     
      |     |                    +-------+   +-------+   +-------+   +-------+
      |     |                    | 3byte |   | 4byte |   | 5byte |   | 6byte |
      |     |                    +-------+   +-------+   +-------+   +-------+
      |     |                    +---------------+                                                             
      |     |                    |  ChunkHeader  |                                                             
      |     |                    +---------------+                                                             
      |     |                         |     |                                                                  
      |     |                         |     +--------------+  +------------+  +------------+  +------------+   
      |     |                         |     |  heap_list_  |--| HeapHeader |--| HeapHeader |--| HeapHeader |   
      |     |                         |     +--------------+  +------------+  +------------+  +------------+   
      |     |                         |                           +----------------------------+  +----------------------+  +-----------------------+  
      |     |                         |                           |  HeapHeader(0x103e600000)  |--|  value(0x103e600008) |--|  value(0x103e000010)  |  
      |     |                         |                           +----------------------------+  +----------------------+  +-----------------------+  
      |     |                         |                                                                        
      |     |                         +--------------+  +------------+  +------------+  +------------+         
      |     |                         |  free_list_  |--| HeapHeader |--| HeapHeader |--| HeapHeader |         
      |     |                         +--------------+  +------------+  +------------+  +------------+         
      |     |  +------------+  +--------------+
      |     +--|  Thread-2  |--|  LocalArena  |
      |     |  +------------+  +--------------+
      |     |  +------------+  +--------------+
      |     +--|  Thread-3  |--|  LocalArena  |
      |        +------------+  +--------------+
      +------------+  +--------------+  +--------------+ 
      |  free list |--|  LocalArena  |--|  LocalArena  |-- ...
      +------------+  +--------------+  +--------------+ 

まとめ

結果得られたのはロックフリーな実装の難しさとメモリ割当関連のバグは発見が不可能に近いという教訓であった。
意味不明なSegvがきつい。

参考

https://www.slideshare.net/kosaki55tea/glibc-malloc

ソースは https://github.com/brn/yatsc/tree/master/src/memory どす〜