Apache Arrow 中使用 jemalloc 实现更快、可扩展的内存分配

发布 2018 年 7 月 20 日
作者 Uwe Korn (uwe)

随着 Apache Arrow 0.9 版本的发布，我们已将 OSX 和 Linux 上数组缓冲区的默认分配器从系统分配器切换到 jemalloc。这适用于 Arrow 的 C++/GLib/Python 实现。在大多数情况下，更改默认分配器通常是为了避免由于频繁的小规模（解）分配而出现的问题。相比之下，在 Arrow 中，我们通常处理大型内存数据集。虽然 jemalloc 为避免小于内存页（4KB）的分配的 RAM 碎片提供了良好的策略，但它也提供了提高跨多个内存页的分配性能的功能。

在 Apache Arrow 之外，jemalloc 为 Facebook 的基础设施提供支持（这也是其大部分开发发生的地方）。它也被用作 Rust 中的默认分配器，并帮助 Redis 减少 Linux 上的内存碎片（“Allocator”）。

我们在 Arrow 中需要的一种特殊分配是内存应为 64 字节对齐。这是为了充分发挥 AVX 等 SIMD 指令集的性能。虽然大多数现代 SIMD 指令也可以在未对齐的内存上工作，但它们在对齐内存上的性能要好得多。为了获得我们分析应用程序的最佳性能，我们希望所有内存都进行分配，以最大限度地提高 SIMD 性能。

对于对齐分配，POSIX API 只提供 aligned_alloc(void** ptr, size_t alignment, size_t size) 函数来分配对齐内存。还有 posix_memalign(void **ptr, size_t alignment, size_t size) 用于将分配修改为首选对齐方式。但两者都无法满足分配的扩展。虽然 realloc 函数通常可以在不物理移动的情况下扩展分配，但它不能确保在分配移动的情况下保持对齐。

如果 Arrow 在未启用 jemalloc 的情况下构建，则会导致每次分配新扩展时都复制数据。为了减少内存复制的数量，我们使用 jemalloc 的 *allocx()-API 来创建、修改和释放对齐分配。其中一个典型任务，它给我们带来了显著的速度提升，那就是增量构建由多个列组成的 Arrow 表。我们通常不知道表的预先大小，需要随着数据加载而扩展我们的分配。

要使用内存扩展因子为 2 的方式增量构建向量，我们将使用带有标准 POSIX API 的以下 C 代码

size_t size = 128 * 1024;
void* ptr = aligned_alloc(64, size);
for (int i = 0; i < 10; i++) {
  size_t new_size = size * 2;
  void* ptr2 = aligned_alloc(64, new_size);
  memcpy(ptr2, ptr, size);
  free(ptr);
  ptr = ptr2;
  size = new_size;
}
free(ptr);

使用 jemalloc 的特殊 API，我们能够省略对 memcpy 的显式调用。在内存扩展无法就地完成的情况下，它仍然由分配器调用，但并非所有情况下都需要。这简化了我们的用户代码为

size_t size = 128 * 1024;
void* ptr = mallocx(size, MALLOCX_ALIGN(64));
for (int i = 0; i < 10; i++) {
  size *= 2;
  ptr = rallocx(ptr, size, MALLOCX_ALIGN(64));
}
dallocx(ptr, MALLOCX_ALIGN(64));

为了了解使用 jemalloc 的实际好处，我们查看了 Arrow C++ 中的基准测试。我们在其中模拟了一个增量构建原始值数组的典型用例。对于数组的构建，我们不知道最终数组中的元素数量，因此我们需要不断扩展存储数据的内存区域。此基准测试的代码是 Arrow C++ 源代码中 builder-benchmark 的一部分，名为 BuildPrimitiveArrayNoNulls。

不使用 jemalloc 的运行时

BM_BuildPrimitiveArrayNoNulls/repeats:3                 636726 us   804.114MB/s
BM_BuildPrimitiveArrayNoNulls/repeats:3                 621345 us   824.019MB/s
BM_BuildPrimitiveArrayNoNulls/repeats:3                 625008 us    819.19MB/s
BM_BuildPrimitiveArrayNoNulls/repeats:3_mean            627693 us   815.774MB/s
BM_BuildPrimitiveArrayNoNulls/repeats:3_median          625008 us    819.19MB/s
BM_BuildPrimitiveArrayNoNulls/repeats:3_stddev            8034 us   10.3829MB/s

使用 jemalloc 的运行时

BM_BuildPrimitiveArrayNoNulls/repeats:3                 630881 us   811.563MB/s
BM_BuildPrimitiveArrayNoNulls/repeats:3                 352891 us   1.41687GB/s
BM_BuildPrimitiveArrayNoNulls/repeats:3                 351039 us   1.42434GB/s
BM_BuildPrimitiveArrayNoNulls/repeats:3_mean            444937 us   1.21125GB/s
BM_BuildPrimitiveArrayNoNulls/repeats:3_median          352891 us   1.41687GB/s
BM_BuildPrimitiveArrayNoNulls/repeats:3_stddev          161035 us   371.335MB/s

对每种配置运行了三次基准测试，以查看性能差异。每种配置的第一次运行都产生了相同的性能，但在所有后续运行中，使用 jemalloc 的版本速度大约快了一倍。在这些情况下，用于构建数组的内存区域可以就地扩展，而无需移动数据。这是可能的，因为有未使用的内存页分配给进程，但未被操作系统回收。如果没有 jemalloc，我们无法利用它们，仅仅因为默认分配器没有提供对齐重新分配的 API。