mi_malloc_source

Posted on | In , | Views:

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中移除。

1
2
3
4
5
6
7
8
9
10
11
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。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
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来看一下内存分配的流程，其流程仅有两部，获取该线程拥有的堆，然后从这个堆上分配一块内存。

1
2
3
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并未初始化则会将其初始化为一个新的堆。

1
2
3
4
5
6
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来进行新的内存的分配。

1
2
3
4
5
6
7
8
9
10
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流程图
mi_heap_malloc流程图

分配Small类型的内存块

我们先来看一下mi_heap_malloc_small，其首先从堆的pages_free_direct数组中找到负责分配对应大小内存块的页，之后调用_mi_page_malloc从该页的Free List中分配一块内存，如果该页的Free List为空则调用_mi_malloc_generic来进行内存的分配。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
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)
  • 找到或分配新的页
  • 从页中分配内存
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
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的堆。

1
2
3
4
5
6
7
8
9
10
11
12
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有关的常量，例如页面大小等。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
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。

1
2
3
4
5
6
7
8
9
10
11
12
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来回收页。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
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
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
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_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。

1
2
3
4
5
6
7
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的指向。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
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;
}