Pony's Arena Allocator

For a prolonged period, stress tests on the TCP system in Pony consistently passed without issue. However, the author felt the tests were insufficient, as they weren't covering the entirety of the codebase. To address this, the author introduced "swarm testing goodness" to the tests, which involved randomizing the tests' actions and execution methods. This change uncovered several bugs, one of which was that memory usage would increase indefinitely under specific usage patterns.

Upon investigation, it was discovered that the issue was not within the TCP code itself, but rather in the Pony runtime's allocator. There were three reasons why the old allocator could retain memory in a manner that could lead to the observed degenerate case. Firstly, a message in Pony is allocated by its sender and freed by its receiver, and about half of all frees occur on a different thread than the one that did the allocation.

Secondly, a large block freed on the wrong thread remained reserved, never reclaimed. Lastly, a workload that transferred blocks between threads reserved fresh address space for every block it allocated and freed, resulting in unbounded memory usage.

As the workload increased, memory consumption grew significantly. For example, when around ten thousand blocks were passed between threads, memory usage climbed by approximately 4.3 MiB per block. Once memory was divided into 32-byte slots, those 32-byte objects remained allocated for the duration of the program. The allocator failed to return these freed blocks of memory for different sizes.

Additionally, the old allocator maintained a global pool with per-size-class free lists accessible by every thread. No thread was responsible for tracking the memory it had allocated. Consequently, the fix for all three issues lay in making each thread responsible for tracking the memory it allocated, including what was in use, what was free, and when a region was empty.

To tackle this problem, the author developed a new allocator inspired by snmalloc's region-based design. Memory is divided into two tiers: large regions requested from the operating system, which are subsequently split into smaller arenas. Each arena is assigned to a single thread, and that thread is responsible for managing all the bookkeeping within its arena, such as tracking allocated and free memory, as well as returning empty regions to the operating system when appropriate.

Cross-thread frees are batched and routed to the appropriate owning thread instead of being handled globally.

The allocator requests memory from the operating system in large aligned chunks called regions - 256 MiB on 64-bit machines and 64 MiB on 32-bit machines. These regions are shared among all threads and are never unmapped. This design ensures that when a thread walks the region list, it can safely assume that the rest of the design is in place, providing a memory safety guarantee.

Threads acquire arenas from these regions, with each arena being 8 MiB on 64-bit machines and 2 MiB on 32-bit machines. The owning thread is responsible for maintaining the arena's bookkeeping. When an arena becomes empty, its physical pages are returned to the operating system, but its address space remains allocated in the region for potential reuse.

Finding the arena that a piece of freed memory belongs to is streamlined using the fact that each arena's starting address is aligned to its own size. On 64-bit machines, every arena starts on an 8 MiB boundary, and masking the low bits of any pointer will reveal the arena's base address. This process takes just one instruction and avoids any memory reads, resulting in efficient performance.

Memory is organized into 16 KiB units within each arena, with each unit serving a specific size class. There are 16 size classes, ranging from 32 bytes up to 1 MiB, with each class being a power of two. Small classes contain multiple objects per unit - for instance, a 32-byte class holds 512 objects per unit. In contrast, large classes occupy multiple units.

Free or used bits are stored as individual bits per unit within a bitmap. For example, an 8 MiB arena has 512 units, which can be accommodated in 8 bitmap words on a 64-bit machine.

The bitmap representation allows for contiguous free units to be identified as spans of zero bits, eliminating the need for a sorted free list or merging adjacent free blocks. This aspect eliminates the potential for adjacent free blocks to remain separate, which was a common issue in the old allocator. The bitmap also enables fast scanning for N consecutive zero bits, signifying N free units, without the need for a complex merge code or O(n) walk.

Cross-thread frees are handled by having the freeing thread collect freed objects and send them to the owning thread in batches, rather than sending them individually. This batching approach reduces the number of atomic operations required, resulting in significant performance improvements. The owning thread processes the freed memory and can subsequently reuse it. While sending the batch back to the owning thread incurs some overhead, it is more efficient than sending individual atomics for each freed object.