在 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 上的内存碎片 (“分配器”)。

我们在 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。