1. 程式人生 > >mimalloc記憶體分配程式碼分析

mimalloc記憶體分配程式碼分析

這篇文章中我們會介紹一下mimalloc的實現,其中可能涉及上一篇文章提到的內容,如果不瞭解的可以先看下這篇mimalloc剖析。首先我們需要了解的是其整體結構,mimalloc的結構如下圖所示

 
mimalloc整體結構
在mimalloc中,每個執行緒都有一個Thread Local的堆,每個執行緒在進行記憶體的分配時均從該執行緒對應的堆上進行分配。在一個堆中會有一個或多個segment,一個segment會對應一個或多個頁,而記憶體的分配就是在這些頁上進行。mimalloc將頁分為三類:
  • small型別的segment的大小為4M,其負責分配大小小於MI_SMALL_SIZE_MAX的記憶體塊,該segment中一個頁的大小均為64KB,因此在一個segment中會包含多個頁,每個頁中會有多個塊
  • large型別的segment的大小為4M,其負責分配大小處於MI_SMALL_SIZE_MAX與MI_LARGE_SIZE_MAX之間的記憶體塊,該segment中僅會有一個頁,該頁佔據該segment的剩餘所有空間,該頁中會有多個塊
  • huge型別的segment,該類segment的負責分配大小大於MI_LARGE_SIZE_MAX的記憶體塊,該類segment的大小取決於需要分配的記憶體的大小,該segment中也僅包含一個頁,該頁中僅會有一個塊

根據heap的定義我們可以看到其有pages_free_direct陣列、pages陣列、Thread Delayed Free List以及一些元資訊。其中pages_free_direct陣列中每個元素對應一個記憶體塊大小的類別,其內容為一個指標,指向一個負責分配對應大小記憶體塊的頁,mimalloc在分配比較小的記憶體時可以通過該陣列直接找到對應的頁,然後試圖從該頁上分配記憶體,從而提升效率。pages陣列中每個元素為一個佇列,該佇列中所有的頁大小均相同,這些頁可能來自不同的segment,其中陣列的最後一個元素(即pages[MI_BIN_FULL])就是前文提到的Full List,倒數第二個元素(即pages[MIN_BIN_HUGE])包含了所有的huge型別的頁。thread_delayed_free就是前文提到的Thread Delayed Free List,用來讓執行緒的擁有者能夠將頁面從Full List中移除。  
struct mi_heap_s {
  mi_tld_t*             tld;
  mi_page_t*            pages_free_direct[MI_SMALL_WSIZE_MAX + 2];
  mi_page_queue_t       pages[MI_BIN_FULL + 1];
  volatile mi_block_t*  thread_delayed_free;
  uintptr_t             thread_id;
  uintptr_t             cookie;
  uintptr_t             random;
  size_t                page_count;
  bool                  no_reclaim;
};
在heap的定義中我們需要特別注意的一個成員是tld(即Thread Local Data)。其成員包括指向對應堆的heap_backing,以及用於segment分配的segment tld以及os tld。
struct mi_tld_s {
  unsigned long long  heartbeat;
  mi_heap_t*          heap_backing;
  mi_segments_tld_t   segments;
  mi_os_tld_t         os;
  mi_stats_t          stats;
};

typedef struct mi_segments_tld_s {
  // 該佇列中所有的segment均有空閒頁,由於large與huge型別的segment僅有一個頁,因此該佇列中所有segment均為small型別
  mi_segment_queue_t  small_free;
  size_t              current_size;
  size_t              peak_size;
  size_t              cache_count;
  size_t              cache_size;
  // segment的快取
  mi_segment_queue_t  cache;
  mi_stats_t*         stats;
} mi_segments_tld_t;

typedef struct mi_os_tld_s {
  uintptr_t           mmap_next_probable;
  void*               mmap_previous;
  uint8_t*            pool;
  size_t              pool_available;
  mi_stats_t*         stats;
} mi_os_tld_t;

mi_malloc

首先要說明一下,所有貼出的原始碼都可能會有一定程度的刪減,例如一些平臺相關的程式碼,一些用於資訊統計的程式碼都可能被刪去。接下來我們跟著mi_malloc來看一下記憶體分配的流程,其流程僅有兩部,獲取該執行緒擁有的堆,然後從這個堆上分配一塊記憶體。
extern inline void* mi_malloc(size_t size) mi_attr_noexcept {
  return mi_heap_malloc(mi_get_default_heap(), size);
}

獲取執行緒擁有的堆

首先介紹一下mimalloc有哪些堆,mimalloc會為每個執行緒保留一個Thread Local的堆,每個執行緒均使用該堆進行記憶體分配,除此之外還有一個全域性變數_mi_heap_main,該堆會被主執行緒視為Thread Local的堆,由於某些OS會用malloc來進行Thread Local的記憶體分配,因此_mi_heap_main在mimalloc尚未初始化時也會被視作預設的堆來進行記憶體分配。
我們先來看一下mi_get_default_heap,該函式會直接返回一個Thread Local的_mi_heap_default,但是該Thread Local預設是被初始化為_mi_heap_empty,之後在呼叫mi_heap_malloc時如果發現該Thread Local並未初始化則會將其初始化為一個新的堆。
static inline mi_heap_t* mi_get_default_heap(void) {
#ifdef MI_TLS_RECURSE_GUARD
  if (!_mi_process_is_initialized) return &_mi_heap_main;
#endif
  return _mi_heap_default;
}

從堆上分配記憶體

由於mimalloc的堆維護了pages_free_direct陣列,可以直接通過該陣列來找到所有針對對應大小的small型別的頁,因此我們可以看到當需要分配的記憶體塊大小小於等於MI_SMALL_SIZE_MAX會呼叫mi_heap_malloc_small從堆上進行記憶體的分配,否則呼叫_mi_malloc_generic從堆上分配記憶體。當然由於pages_free_direct中指向的頁可能Free List已經為空了,那麼其最終還是會呼叫_mi_malloc_generic來進行新的記憶體的分配。
extern inline void* mi_heap_malloc(mi_heap_t* heap, size_t size) mi_attr_noexcept {
  void* p;
  if (mi_likely(size <= MI_SMALL_SIZE_MAX)) {
    p = mi_heap_malloc_small(heap, size);
  }
  else {
    p = _mi_malloc_generic(heap, size);
  }
  return p;
}
先貼一張從堆上分配記憶體的總體流程圖,接下來我們仔細介紹一下這兩個函式具體的呼叫。 mi_heap_malloc流程圖

分配Small型別的記憶體塊

我們先來看一下mi_heap_malloc_small,其首先從堆的pages_free_direct陣列中找到負責分配對應大小記憶體塊的頁,之後呼叫_mi_page_malloc從該頁的Free List中分配一塊記憶體,如果該頁的Free List為空則呼叫_mi_malloc_generic來進行記憶體的分配。
extern inline void* mi_heap_malloc_small(mi_heap_t* heap, size_t size) mi_attr_noexcept {
  mi_page_t* page = _mi_heap_get_free_small_page(heap,size);
  return _mi_page_malloc(heap, page, size);
}

extern inline void* _mi_page_malloc(mi_heap_t* heap, mi_page_t* page, size_t size) mi_attr_noexcept {
  mi_block_t* block = page->free;
  if (mi_unlikely(block == NULL)) {
    return _mi_malloc_generic(heap, size);
  }
  page->free = mi_block_next(page,block);
  page->used++;

  ...

  return block;
}

分配Large或者Huge型別的記憶體塊

接下來我們看一下_mi_malloc_generic,該函式呼叫的原因可能有如下兩種:
  • 需要分配small型別的記憶體塊,但是由pages_free_direct獲得的頁的Free List已經為空
  • 需要分配large或者huge型別的記憶體塊

我們可以看到_mi_malloc_generic的流程可以歸納為:
  • 如果需要的話進行全域性資料/執行緒相關的資料/堆的初始化
  • 呼叫回撥函式(即實現前文所說的deferred free)
  • 找到或分配新的頁
  • 從頁中分配記憶體
void* _mi_malloc_generic(mi_heap_t* heap, size_t size) mi_attr_noexcept
{
  if (mi_unlikely(!mi_heap_is_initialized(heap))) {
    mi_thread_init();
    heap = mi_get_default_heap();
  }

  _mi_deferred_free(heap, false);

  mi_page_t* page;
  if (mi_unlikely(size > MI_LARGE_SIZE_MAX)) {
    if (mi_unlikely(size >= (SIZE_MAX - MI_MAX_ALIGN_SIZE))) {
      page = NULL;
    }
    else {
      page = mi_huge_page_alloc(heap,size);
    }
  }
  else {
    page = mi_find_free_page(heap,size);
  }
  if (page == NULL) return NULL;

  return _mi_page_malloc(heap, page, size);
}

初始化

前面我們提到過每個執行緒都有一個Thread Local的堆,該堆預設被設為_mi_heap_empty。如果呼叫_mi_malloc_generic時發現該執行緒的堆為_mi_heap_empty則進行初始化。mi_thread_init會首先呼叫mi_process_init來進行程序相關資料的初始化,之後初始化Thread Local的堆。
void mi_thread_init(void) mi_attr_noexcept
{
  // ensure our process has started already
  mi_process_init();

  // initialize the thread local default heap
  if (_mi_heap_init()) return;  // returns true if already initialized

  ...

  #endif
}
我們可以看到mi_process_init僅會被呼叫一次,其初始化了_mi_heap_main,其會被設為主執行緒的Thread Local的堆。其註冊了mi_process_done為執行緒結束的回撥函式,並呼叫mi_process_setup_auto_thread_done來設定mi_thread_done為執行緒結束時的回撥函式,而_mi_os_init則是用來設定一些與OS有關的常量,例如頁面大小等。
void mi_process_init(void) mi_attr_noexcept {
  // ensure we are called once
  if (_mi_process_is_initialized) return;
  // access _mi_heap_default before setting _mi_process_is_initialized to ensure
  // that the TLS slot is allocated without getting into recursion on macOS
  // when using dynamic linking with interpose.
  mi_heap_t* h = _mi_heap_default;
  _mi_process_is_initialized = true;

  _mi_heap_main.thread_id = _mi_thread_id();
  uintptr_t random = _mi_random_init(_mi_heap_main.thread_id)  ^ (uintptr_t)h;
  #ifndef __APPLE__
  _mi_heap_main.cookie = (uintptr_t)&_mi_heap_main ^ random;
  #endif
  _mi_heap_main.random = _mi_random_shuffle(random);

  atexit(&mi_process_done);
  mi_process_setup_auto_thread_done();
  mi_stats_reset();
  _mi_os_init();
}
我們來看一下mi_process_done與mi_thread_done分別做了什麼。 我們可以看到mi_process_done主要是呼叫了mi_collect來回收已經分配的記憶體,該函式呼叫的也是mi_heap_collect_ex,不過由於其呼叫的引數不同,行為會稍有不同,在此處的呼叫會收集abandon segment,然後釋放這些segment。
static void mi_process_done(void) {
  // only shutdown if we were initialized
  if (!_mi_process_is_initialized) return;
  // ensure we are called once
  static bool process_done = false;
  if (process_done) return;
  process_done = true;

  #ifndef NDEBUG
  mi_collect(true);
  #endif
}
mi_thread_done則主要是呼叫_mi_heap_done來回收部分資源。該函式會先把_mi_heap_default重新設為預設值,如果是主執行緒就設為_mi_heap_main,否則設為_mi_heap_empty。如果該執行緒不是主執行緒的話則呼叫_mi_heap_collect_abandon來回收記憶體並釋放動態分配的heap,如果是主執行緒的話會呼叫_mi_heap_destroy_pages來回收頁。
static bool _mi_heap_done(void) {
  mi_heap_t* heap = _mi_heap_default;
  if (!mi_heap_is_initialized(heap)) return true;

  // reset default heap
  _mi_heap_default = (_mi_is_main_thread() ? &_mi_heap_main : (mi_heap_t*)&_mi_heap_empty);

  // todo: delete all non-backing heaps?

  // switch to backing heap and free it
  heap = heap->tld->heap_backing;
  if (!mi_heap_is_initialized(heap)) return false;

  // collect if not the main thread 
  if (heap != &_mi_heap_main) {
    _mi_heap_collect_abandon(heap);
  }

  // merge stats
  _mi_stats_done(&heap->tld->stats);

  // free if not the main thread
  if (heap != &_mi_heap_main) {
    _mi_os_free(heap, sizeof(mi_thread_data_t), &_mi_stats_main);
  }
#if (MI_DEBUG > 0)
  else {
    _mi_heap_destroy_pages(heap);
  }
#endif
  return false;
}
在mimalloc中,如果一個執行緒結束了,那麼其對應的Thread Local的堆就可以釋放了,但是在該堆中還可能存在有一些記憶體塊正在被使用,且此時會將對應的segment設定為ABANDON,之後由其他執行緒來獲取該segment,之後利用該segment進行對應的記憶體分配與釋放(mimalloc也有一個no_reclaim的選項,設定了該選項的堆不會主動獲取其他執行緒ABANDON的segment)。
接下來我們來看一下_mi_heap_collect_abandon,其實際呼叫了mi_heap_collect_ex,下面的程式碼中略去了部分不會被_mi_heap_done使用到的分支。該函式的流程如下:
  • 呼叫deferred free回撥函式
  • 標記當前堆的Full List中的所有頁面為Normal,從而讓其在釋放時加入Thread Free List,因為該segment之後可能會被其他執行緒接收
  • 釋放該堆的Thread Delayed Free List中的記憶體塊(不是每頁一個的Thread Free List)
  • 遍歷該堆所擁有的所有頁,對每個頁呼叫一次mi_heap_page_collect
  • 呼叫_mi_page_free_collect將頁中的Local Free List以及Thread Free List追加到Free List之後
  • 如果該頁沒有正在使用的塊則呼叫_mi_page_free將該頁釋放回對應的segment中,如果segment中所有的空閒頁均被釋放則可能直接釋放對應的segment回OS或加入堆的快取中
  • 如果該頁尚有正在使用的塊則將該頁標記為abandon,當某個segment中所有的頁均被標記為abandon後會將對應的segment加入全域性的abandon segment list中(堆中並未保留有哪些segment的資訊,因此需要遍歷所有頁來完成這一操作)
  • 釋放堆中所有快取的segment
static void mi_heap_collect_ex(mi_heap_t* heap, mi_collect_t collect)
{
  _mi_deferred_free(heap,collect > NORMAL);
  if (!mi_heap_is_initialized(heap)) return;

  // 一些接收abandon list中的segment的程式碼
  ...

  // if abandoning, mark all full pages to no longer add to delayed_free
  if (collect == ABANDON) {
    for (mi_page_t* page = heap->pages[MI_BIN_FULL].first; page != NULL; page = page->next) {
      _mi_page_use_delayed_free(page, false);  // set thread_free.delayed to MI_NO_DELAYED_FREE      
    }
  }

  // free thread delayed blocks. 
  // (if abandoning, after this there are no more local references into the pages.)
  _mi_heap_delayed_free(heap);

  // collect all pages owned by this thread
  mi_heap_visit_pages(heap, &mi_heap_page_collect, &collect, NULL);
  mi_assert_internal( collect != ABANDON || heap->thread_delayed_free == NULL );
  
  // collect segment caches
  if (collect >= FORCE) {
    _mi_segment_thread_collect(&heap->tld->segments);
  }
}

Huge型別頁面的分配

由於huge型別的頁面對應的segment中僅有一個頁,且該頁僅能分配一個塊,因此其會重新分配一個segment,從中建立新的頁面。mi_huge_page_alloc會呼叫mi_page_fresh_alloc分配一個頁面,然後將其插入堆對應的BIN中(即heap->pages[MI_BIN_HUGE])。由下圖可以看到Small與Large型別頁面分配時所呼叫的mi_find_free_page也會呼叫該函式來進行頁面的分配,接下來我們就介紹一下mi_page_fresh_alloc。 _mi_malloc_generic函式呼叫關係 我們可以看到mi_page_fresh_alloc主要做了三件事,先從堆中分配一個新的頁,並對該頁進行初始化,最後將該頁加入對應的BIN中。其中mi_segment_page_alloc就是從堆中找到一個足夠容納新頁的segment並分配一個新的頁,其會根據需要分配的記憶體塊的大小呼叫mi_segment_small_page_alloc/mi_segment_large_page_alloc/mi_segment_huge_page_alloc。
static mi_page_t* mi_page_fresh_alloc(mi_heap_t* heap, mi_page_queue_t* pq, size_t block_size) {
  mi_page_t* page = _mi_segment_page_alloc(block_size, &heap->tld->segments, &heap->tld->os);
  if (page == NULL) return NULL;
  mi_page_init(heap, page, block_size, &heap->tld->stats);
  mi_page_queue_push(heap, pq, page);
  return page;
}
mi_huge_page_alloc/mi_large_page_alloc
mi_huge_page_alloc與mi_large_page_alloc非常類似,因為這兩種型別記憶體塊對應的segment都僅有一個頁,稍有區別的是large型別的segment的大小為4M,而huge型別的segment大小取決於需要的記憶體塊的大小。因為這兩種型別的塊的分配必須獲取新的segment,因此其均呼叫mi_segment_alloc獲取一個新的segment,然後在新獲取的segment中建立一個新的頁並標記該頁為正在使用。
接下來介紹一下其中用於分配新segment的函式mi_segment_alloc的流程:
  • 計算segment的大小,頁的大小
  • 從cache中試圖找到一個足夠大的segment,如果segment中有較多未使用的空間則會將部分空間釋放回OS
  • 設定segment的元資訊
mi_segment_small_page_alloc
Small型別的頁的分配稍微有些不同,因為large與huge型別的記憶體塊其對應的segment中均只有一個頁,而small型別的segment中每個頁均有多個頁,因此mimalloc在堆中儲存了一個segment small free list,該佇列中所有的segment均為small型別且均有空閒的頁。mi_segment_small_page_alloc會首先從該列表中試圖找到一個有空閒頁的segment,然後從該segment中分配一頁,如果分配完成後該segment中已經沒有空閒頁了則將其移出該列表,如果沒有找到則會呼叫mi_segment_alloc新分配一個segment並將其加入該列表中。

Small/Large型別頁面的分配

由於Small與Large型別的頁面中均可以包含多個塊,因此分配這兩種型別的記憶體塊時需要查詢已有的頁面,檢視其中是否有頁中有尚未分配的記憶體塊。因此其會首先找到對應的BIN,遍歷其中的所有頁面,試圖擴充套件Free List(包括利用尚未使用的空間、合併Local Free List與Thread Free List)從而找到一個有足夠空間的頁。由於在該過程中會進行Free List的合併,因此其還會釋放一些完全空閒的頁,進而可能導致segment的釋放。如果在遍歷完BIN後仍舊沒有找到空閒頁則會mi_page_fresh來分配一個新的頁,在該過程中會呼叫_mi_segment_try_reclaim_abandoned來試圖獲取一個abandon的segment,但是要注意的是重新獲取一個segment並不一定會帶來新的頁,因為可能接收的segment為large或huge型別或者其已經沒有空閒頁了,在這種情況下會去呼叫mi_page_fresh_alloc去獲取新的segment和頁或者從已有的segment中分配新的頁。

從頁中分配記憶體塊

此時我們終於獲得了一個空閒頁,我們可以從該頁中分配一個記憶體塊了,其程式碼如下。我們可以看到其首先檢查了一下當前頁的Free List是否為空,如果為空則呼叫_mi_malloc_generic,這是因為該函式的呼叫入口有兩種,第一種是分配small型別的記憶體塊時呼叫的mi_heap_malloc_small,第二種才是_mi_malloc_generic。
這裡需要介紹一下mimalloc更新pages_free_direct的機制,mimalloc通過在將一個頁向BIN中新增或者移除頁時更新對應的pages free direct陣列,由於對齊的問題,因此一個頁面的分配可能需要改變多個pages_free_direct的指向。
extern inline void* _mi_page_malloc(mi_heap_t* heap, mi_page_t* page, size_t size) mi_attr_noexcept {
  mi_block_t* block = page->free;
  if (mi_unlikely(block == NULL)) {
    return _mi_malloc_generic(heap, size); // slow path
  }
  mi_assert_internal(block != NULL && _mi_ptr_page(block) == page);
  // pop from the free list
  page->free = mi_block_next(page,block);
  page->used++;

  ...

  return block;
}

總結

以上就是mimalloc中用於記憶體分配部分的程式碼的解析了,其中還有很多沒有講到的地方,例如其向OS請求記憶體部分的程式碼等等。文章如果有哪裡有問題,歡迎提出,對該專案感興趣的可以去看一下其倉庫1,或者參考這篇文章2

引用

[1] https://github.com/microsoft/mimalloc [2] https://www.microsoft.com/en-us/research/publication/mimalloc-free-list-sharding-in-action/