深入解析TCMalloc中的可重启序列机制

深入解析TCMalloc中的可重启序列机制

tcmalloc 项目地址: https://gitcode.com/gh_mirrors/tc/tcmalloc

前言

在现代内存分配器中，性能优化是一个永恒的话题。TCMalloc作为Google开发的高性能内存分配器，其核心设计理念之一就是通过减少锁竞争来提高性能。本文将重点解析TCMalloc中使用的可重启序列(Restartable Sequences, rseq)机制，这是实现高效每CPU缓存的关键技术。

什么是可重启序列

可重启序列是Linux内核提供的一种机制，它允许用户空间程序执行一段"原子"代码序列，这里的原子性是指相对于同一CPU上的其他线程而言。如果在执行过程中被内核抢占、中断或信号处理打断，内核会确保这段代码要么完整执行完毕，要么完全重新开始执行。

这种机制由Google的Paul Turner、Andrew Hunter和EfficiOS的Mathieu Desnoyers共同开发，最初目的是为了优化每CPU原子操作。

TCMalloc为什么需要rseq

TCMalloc使用每CPU缓存来避免多线程环境下的锁竞争。传统实现可能需要：

1. 使用原子操作来保证一致性，但这会带来性能开销
2. 使用互斥锁，这会导致线程阻塞

rseq提供了一种更高效的替代方案：在大多数情况下(线程保持在同一个CPU上)，可以完全避免原子操作；只有在少数情况下(线程被迁移到其他CPU)，才需要重新执行操作。

TCMallocSlab数据结构

在每CPU模式下，TCMalloc为每个逻辑CPU分配一个TcmallocSlab::Slabs数组。每个slab包含：

1. 控制数据的头部区域(每个大小类一个8字节的头部)
2. 指向空闲对象的指针数组

内存布局如下所示：

+---------------------+---------------------+-----+---------------------+
| Header for class 0  | Header for class 1  | ... | Header for class N  |
+---------------------+---------------------+-----+---------------------+
| Pointer to object 0 | Pointer to object 1 | ... | Pointer to object M |
+---------------------+---------------------+-----+---------------------+

头部结构包含两个关键字段：

• current: 当前已占用槽位的结束偏移量
• end: 该大小类槽位数组的结束偏移量

分配操作实现

让我们看一个分配操作的伪代码实现：

void* TcmallocSlab_Pop(void *slabs, size_t size_class) {
restart:
  // 设置rseq上下文
  __rseq_abi.rseq_cs = &__rseq_cs_TcmallocSlab_Pop;
  
start:
  // 获取当前CPU ID
  uint64_t cpu_id = __rseq_abi.cpu_id;
  
  // 获取头部指针
  Header* hdr = &slabs[cpu_id].header[size_class];
  uint64_t current = hdr->current;
  
  // 获取要返回的对象
  void* ret = *(&slabs[cpu_id] + current * sizeof(void*) - sizeof(void*));
  
  // 检查是否下溢
  if ((uintptr_t)ret & 1) goto underflow;
  
  // 预取下一个可能分配的对象
  void* next = *(&slabs[cpu_id] + current * sizeof(void*) - 2 * sizeof(void*));
  
  // 提交点：更新current
  --current;
  hdr->current = current;
  
commit:
  prefetcht0(next);  // 预取优化
  return ret;
  
underflow:
  return nullptr;
}

关键点：

1. 整个操作只有最后的hdr->current = current是提交点
2. 如果在此之前的任何点被中断，操作会完全重启
3. 预取优化利用了空间局部性原理

释放操作实现

释放操作的伪代码如下：

int TcmallocSlab_Push(void *slab, size_t size_class, void* item) {
restart:
  // 设置rseq上下文
  __rseq_abi.rseq_cs = &__rseq_cs_TcmallocSlab_Push;
  
start:
  // 获取当前CPU ID
  uint64_t cpu_id = __rseq_abi.cpu_id;
  
  // 获取头部信息
  Header* hdr = &slabs[cpu_id].header[size_class];
  uint64_t current = hdr->current;
  uint64_t end = hdr->end;
  
  // 检查是否上溢
  if (current >= end) goto overflow;
  
  // 存储释放的对象
  *(&slabs[cpu_id] + current * sizeof(void*)) = item;
  
  // 提交点：更新current
  current++;
  hdr->current = current;
  
commit:
  return;
  
overflow:
  return overflow_handler(cpu_id, size_class, item);
}

虽然释放操作有两个存储操作(存储对象和更新current)，但只有更新current是提交点。如果被中断，之前的存储会被后续的重试覆盖。

重启处理机制

当内核检测到在rseq关键段内发生上下文切换时，它会：

1. 检查当前指令指针是否在[start, commit)范围内
2. 如果是，将控制流转到指定的abort处理程序
3. abort处理程序跳转回restart标签重新执行

在x86架构上，abort处理程序前需要有4字节的签名(通常是一个特殊编码的nop指令)：

.byte 0x0f, 0x1f, 0x05  ; nop指令
.long RSEQ_SIGNATURE     ; 签名
TcmallocSlab_Push_trampoline:
abort:
  jmp restart            ; 跳回重启

内核实现细节

在Linux内核中，rseq机制的关键部分包括：

1. sys_rseq系统调用：负责注册/注销rseq
2. rseq_ip_fixup：处理关键段中断后的重启逻辑
3. CPU ID管理：维护用户空间可见的CPU ID信息

内核通过以下方式保证安全性：

• 每次重启都重新验证rseq上下文
• 要求abort处理程序前有预注册的签名，防止代码注入攻击
• 在信号处理和任务切换时清除rseq上下文

性能优化技巧

TCMalloc在使用rseq时采用了几个重要的优化：

1. slab指针缓存：将当前CPU的slab指针缓存在线程本地存储中，减少指针计算开销
2. 懒初始化：延迟初始化每个CPU的slab，减少启动开销
3. 批量操作：在缓存上溢/下溢时处理多个对象，分摊锁开销
4. 预取优化：在分配路径预取下一个可能分配的对象

总结

TCMalloc通过可重启序列机制实现了高效的每CPU缓存，这种设计：

1. 在常见路径(无迁移)上完全避免了原子操作
2. 在罕见路径(有迁移)上通过重启保证正确性
3. 通过精心设计的提交点和重启逻辑确保一致性
4. 结合多种优化技术最大化性能

理解这一机制不仅有助于深入掌握TCMalloc的工作原理，也为设计高性能并发数据结构提供了宝贵思路。rseq作为一种新兴的底层原语，在未来可能会有更广泛的应用场景。

tcmalloc 项目地址: https://gitcode.com/gh_mirrors/tc/tcmalloc