Testing memory allocation failures with Zig

Zig is a relatively new language that describes itself as "a general-purpose programming language and toolchain for maintaining robust, optimal, and reusable software."  One interesting feature of Zig is that all memory management is manual. There are no implicit heap allocations anywhere, even in the standard library (compare to, for example, C's strcpy/strncpy). Any function that wishes to allocate memory must take an allocator parameter. This means that the same code can run on machines with dramatically different memory availability, even a freestanding (embedded) target, simply by choosing an appropriate allocator.

Another important feature of Zig is that errors cannot be implicitly ignored. Errors can easily be passed up the stack with "try", or explicitly handled or ignored, but gone are the days of forgetting (or being too lazy) to check the return code of a function. Since the default behavior is for errors to be checked or propagated, the language naturally guides developers into properly handling errors. It takes special effort to ignore an error!

With these two features, memory allocation failures become "just another failure" mode that must be checked - compare to the common but regrettable C practice of ignoring malloc errors - and it is also very easy to plug in different allocators that fail in different or interesting ways, to see how your code will behave. In this post we'll demonstrate an allocator that fails after N allocations, one that fails on allocations of a specific size, and one that fails randomly with a certain probability (using a PRNG from the Zig standard library).

For these examples we'll use a very simple function that allocates and immediately frees {% c-line %}size{% c-line-end %} byte blocks of memory {% c-line %}times{% c-line-end %} times:

{% c-block language="zig" %}
const std = @import("std");
pub fn functionThatUsesAnAllocator(allocator: *std.mem.Allocator, size: usize, times: usize) !usize {
   var i : usize = 0;
   while (i < times) {
       const memory = try allocator.alloc(u8, size);
       allocator.free(memory);
       i += 1;
   }
   return size;
}
{% c-block-end %}

Now, suppose we want to test what happens to our code if allocations start failing after the Nth request for memory. Luckily we can use the built-in {% c-line %}FailingAllocator{% c-line-end %} in the {% c-line %}std.testing package{% c-line-end %}:

{% c-block language="zig" %}
const std = @import("std");

test "fail after N allocations" {
   var failing_allocator = std.testing.FailingAllocator.init(std.testing.allocator, 5);
   const result = functionThatUsesAnAllocator(&failing_allocator.allocator, 32, 10);
   std.testing.expectError(error.OutOfMemory, result);
}
{% c-block-end %}

However, what if we want to check other kinds of failures? For those cases not covered by the built-in {% c-line %}FailingAllocator{% c-line-end %}, we'll have to provide our own. Luckily this is fairly simple, and we can even use the standard failing allocator as the basis for our code. The relevant portion is the alloc function, shown below with a comment in place of the code that implements failing on the Nth request:

{% c-block language="zig" %}
fn alloc(
   allocator: *std.mem.Allocator,
   len: usize,
   ptr_align: u29,
   len_align: u29,
   return_address: usize,
) error{OutOfMemory}![]u8 {
   const self = @fieldParentPtr(MyAllocator, "allocator", allocator);
   // implement custom logic here
   return self.internal_allocator.allocFn(self.internal_allocator, len, ptr_align, len_align, return_address);
}
{% c-block-end %}

With this in hand, we can start implementing custom failure modes. For example, maybe we want allocations of a specific size to fail. We'll replace the internal {% c-line %}fail_index{% c-line-end %} with {% c-line %}fail_size{% c-line-end %}, and return a memory failure if the requested size matches:

{% c-block language="zig" %}
fn alloc(
   allocator: *std.mem.Allocator,
   len: usize,
   ptr_align: u29,
   len_align: u29,
   return_address: usize,
) error{OutOfMemory}![]u8 {
   const self = @fieldParentPtr(FailingAllocator, "allocator", allocator);
   if (self.fail_size == len) {
       return error.OutOfMemory;
   }
   return self.internal_allocator.allocFn(self.internal_allocator, len, ptr_align, len_align, return_address);
}
{% c-block-end %}

Alternatively, maybe we want an allocator that fails randomly with a certain probability (perhaps as part of fuzz testing to explore a large set of paths through the code before memory failure). In this case we're using a pseudorandom number generator from the {% c-line %}std.rand{% c-line-end %} package, but in real code we would make the seed a parameter instead of hardcoding it to zero.

{% c-block language="zig" %}
fn alloc(
   allocator: *std.mem.Allocator,
   len: usize,
   ptr_align: u29,
   len_align: u29,
   return_address: usize,
) error{OutOfMemory}![]u8 {
   const self = @fieldParentPtr(FailingAllocator, "allocator", allocator);
   const val = self.prng.float(f64);
   if (val <= self.fail_chance) {
       return error.OutOfMemory;
   }
   return self.internal_allocator.allocFn(self.internal_allocator, len, ptr_align, len_align, return_address);
}
{% c-block-end %}

Notice that no special work was needed to adapt the code for testing it with different allocators. Furthermore, any standard Zig code will follow this pattern of taking an allocator as a parameter if it needs to allocate memory. This means that code from the standard library or a third party will be amenable to this testing strategy, without requiring any more work than providing an appropriate custom allocator. To see the full example code, see our repo.

‍