Heap allocation failure is something that is hard or impossible to account for in every case in most programming languages. There are either hidden memory allocations that can't be handled, or it's seen as too inconvenient to handle every possible allocation failure so the possibility is ignored.
For example, when concatenating two strings with the + operator (where there is an implicit allocation that's needed to store the result of the concatenation):
std::bad_alloc, but it's easy to ignore or mishandle (or not be aware of the possibility of allocation failure in every case)Even in C, where the return of malloc can be checked against NULL to detect allocation failure, it's pretty common to see unchecked malloc calls in C code (and C compilers let you ignore the possibility of allocation failure without complaint).
One of the unique features of Zig is that it "cares about allocation failure". That is:
error.OutOfMemoryTogether, these conditions make it so that the code you naturally write in Zig will include handling of OutOfMemory errors. However, because actually running into OutOfMemory organically is rare, it's not easy to be sure that you're handling the error correctly in all cases. Additionally, because there are many functions that have OutOfMemory as their only possible error, the error handling of those function calls are not exercised in a typical test environment.
OutOfMemory errors🔗Luckily, though, allocators in Zig also have some unique properties that lend themselves to potential remedies:
The first point means that it's easy/normal to write custom allocators, while the second means that it's easy/normal to swap out allocators during tests. In order to help test OutOfMemory errors, Zig's standard library contains std.testing.FailingAllocator, which will artificially induce an OutOfMemory error once it hits its user-defined number of allocations. Here's a simple example:
test {
// Create an allocator that will fail on the 0th allocation
var failing_allocator = std.testing.FailingAllocator.init(std.testing.allocator, 0);
// Try to allocate 8 bytes
var allocation = failing_allocator.allocator().alloc(u8, 8);
// Confirm that the allocation failed and gave OutOfMemory
try std.testing.expectError(error.OutOfMemory, allocation);
}
This FailingAllocator lays the groundwork for a strategy that allows inducing OutOfMemory for all allocations within a chunk of code. The strategy goes like this:
As long as the number of memory allocations is deterministic, this strategy works, and is the strategy that the Zig parser tests have employed for years (since 2017) to ensure that the parser handles memory allocation failures without introducing memory leaks (interestingly enough, the implementation of this strategy for the Zig parser tests also happens to be the reason that FailingAllocator was created).
Recently, I went ahead and turned the strategy used by the Zig parser tests into something more re-usable---std.testing.checkAllAllocationFailures---which will be available in the next release of Zig (0.10.0), or can be used now in the latest master version of Zig.
checkAllAllocationFailures🔗Here's some code that parses a newline-separated list of key=value pairs, e.g.
something=other
equals=equals
and returns it as a std.BufMap:
const std = @import("std");
pub fn parse(allocator: std.mem.Allocator, stream_source: *std.io.StreamSource) !std.BufMap {
var map = std.BufMap.init(allocator);
errdefer map.deinit();
const reader = stream_source.reader();
const end_pos = try stream_source.getEndPos();
while ((try stream_source.getPos()) < end_pos) {
var key = try reader.readUntilDelimiterAlloc(allocator, '=', std.math.maxInt(usize));
var value = (try reader.readUntilDelimiterOrEofAlloc(allocator, '\n', std.math.maxInt(usize))) orelse return error.EndOfStream;
try map.putMove(key, value);
}
return map;
}
There are some problems lurking in the function that you might be able to spot, but we'll get to them later. Here's a simple test case that passes just fine:
test {
const data =
\\foo=bar
\\baz=qux
;
var stream_source = std.io.StreamSource{ .const_buffer = std.io.fixedBufferStream(data) };
var parsed = try parse(std.testing.allocator, &stream_source);
defer parsed.deinit();
try std.testing.expectEqual(@as(usize, 2), parsed.count());
try std.testing.expectEqualStrings("bar", parsed.get("foo").?);
try std.testing.expectEqualStrings("qux", parsed.get("baz").?);
}
In order to be able to use checkAllAllocationFailures for this test, we'll need to make some changes to it. For reference, here's the signature of std.testing.checkAllAllocationFailures along with a small portion of its doc comment:
/// The provided `test_fn` must have a `std.mem.Allocator` as its first argument,
/// and must have a return type of `!void`. Any extra arguments of `test_fn` can
/// be provided via the `extra_args` tuple.
pub fn checkAllAllocationFailures(
backing_allocator: std.mem.Allocator,
comptime test_fn: anytype,
extra_args: anytype,
) !void
So, we'll need to move our test code into an appropriately constructed function that we can provide to checkAllAllocationFailures:
!void.StreamSource).In this case, this ends up looking something like this:
fn parseTest(allocator: std.mem.Allocator, stream_source: *std.io.StreamSource, expected: std.BufMap) !void {
var parsed = try parse(allocator, stream_source);
defer parsed.deinit();
try std.testing.expectEqual(expected.count(), parsed.count());
var expected_it = expected.iterator();
while (expected_it.next()) |expected_entry| {
const actual_value = parsed.get(expected_entry.key_ptr.*).?;
try std.testing.expectEqualStrings(expected_entry.value_ptr.*, actual_value);
}
}
with a test block like so:
test {
const data =
\\foo=bar
\\baz=qux
;
var stream_source = std.io.StreamSource{ .const_buffer = std.io.fixedBufferStream(data) };
var expected = expected: {
var map = std.BufMap.init(std.testing.allocator);
errdefer map.deinit();
try map.put("foo", "bar");
try map.put("baz", "qux");
break :expected map;
};
defer expected.deinit();
try parseTest(std.testing.allocator, &stream_source, expected);
}
This still passes just fine. Now let's replace the direct parseTest call with a checkAllAllocationFailures call:
std.testing.allocatorparseTest needs to be called with, we can use an anonymous list literal/tuple (the types within the tuple are checked at compile-time to ensure they match the signature of the test_fn)- try parseTest(std.testing.allocator, &stream_source, expected);
+ try std.testing.checkAllAllocationFailures(
+ std.testing.allocator,
+ parseTest,
+ .{ &stream_source, expected },
+ );
Before running this, though, we'll need to make one last change to the parseTest function. Since checkAllAllocationFailures will now be calling parseTest multiple times (one initial call and then another for each induced allocation failure), we need to make sure that any relevant state is reset at the start of every call. From the checkAllAllocationsFailures doc comment:
/// Any relevant state shared between runs of `test_fn` *must* be reset within `test_fn`.
In this case, the cursor of the StreamSource needs to be reset, as otherwise, after the first run, the cursor will remain at the end of the stream and the next run will immediately fail with EndOfStream (instead of the induced OutOfMemory that we'd expect). To fix this, we need to add this to the beginning of parseTest:
try stream_source.seekTo(0);
Now when we run the test, it will induce allocation failures and report any problems it finds. Here are the results from the first run (heavily truncated to only include the relevant portions of the stack traces):
fail_index: 1/5
allocated bytes: 8
freed bytes: 5
allocations: 1
deallocations: 0
allocation that was made to fail:
src/main.zig:12:61: 0x20f2b6 in parse (test)
var value = (try reader.readUntilDelimiterOrEofAlloc(allocator, '\n', std.math.maxInt(usize))) orelse return error.EndOfStream;
^
src/main.zig:24:27: 0x20edc9 in parseTest (test)
var parsed = try parse(allocator, stream_source);
^
Test [1/1] test ""... FAIL (MemoryLeakDetected)
zig/lib/std/testing.zig:713:21: 0x20a5d8 in std.testing.checkAllAllocationFailures (test)
return error.MemoryLeakDetected;
^
src/main.zig:50:5: 0x209ac1 in test "" (test)
try std.testing.checkAllAllocationFailures(
^
[gpa] (err): memory address 0x7fda4c422000 leaked:
src/main.zig:10:53: 0x20f139 in parse (test)
var key = try reader.readUntilDelimiterAlloc(allocator, '=', std.math.maxInt(usize));
^
src/main.zig:24:27: 0x20ece9 in parseTest (test)
var parsed = try parse(allocator, stream_source);
^
From this, we can see a few things (starting from the top):
FailingAllocator including a stack trace of the allocation that was made to fail (note: the stack trace reporting part was only recently merged).error.MemoryLeakDetected returned from checkAllAllocationFailures.std.testing.allocator that we passed as the backing allocator to checkAllAllocationFailures printed the memory address of the leaked allocation along with a stack trace of where the leaked memory was allocated.In particular, this is the problematic code:
var key = try reader.readUntilDelimiterAlloc(allocator, '=', std.math.maxInt(usize));
var value = (try reader.readUntilDelimiterOrEofAlloc(allocator, '\n', std.math.maxInt(usize))) orelse return error.EndOfStream;
try map.putMove(key, value);
That is, we're leaking the allocation for key when the allocation for value fails. This wasn't a problem before the introduction of checkAllAllocationFailures because normally (if all the allocations in the test succeed), the map.putMove would take ownership of the allocated memory of both key and value and then they'd get cleaned up along with the BufMap later on.
The simplest fix here would be to put in an errdefer that will free key (errdefer instead of defer so that it runs only if the value allocation fails or the putMove call fails) like so:
var key = try reader.readUntilDelimiterAlloc(allocator, '=', std.math.maxInt(usize));
+ errdefer allocator.free(key);
If you happen to be thinking that we'll need the same fix for value, you'd be correct. However, for the sake of completeness let's try running the test again with only the errdefer for key added. Here's the result:
fail_index: 2/5
allocated bytes: 16
freed bytes: 13
allocations: 2
deallocations: 1
allocation that was made to fail:
src/main.zig:15:24: 0x20f3ed in parse (test)
try map.putMove(key, value);
^
Test [1/1] test ""... FAIL (MemoryLeakDetected)
zig/lib/std/testing.zig:714:21: 0x20a6b3 in std.testing.checkAllAllocationFailures (test)
return error.MemoryLeakDetected;
^
src/main.zig:50:5: 0x209b81 in test "" (test)
try std.testing.checkAllAllocationFailures(
^
[gpa] (err): memory address 0x7f6ff57de008 leaked:
src/main.zig:12:61: 0x20f2b6 in parse (test)
var value = (try reader.readUntilDelimiterOrEofAlloc(allocator, '\n', std.math.maxInt(usize))) orelse return error.EndOfStream;
^
src/main.zig:24:27: 0x20edc9 in parseTest (test)
var parsed = try parse(allocator, stream_source);
^
It is similar to before, but now we can see that value is being leaked when the map.putMove call fails. Now let's put in the errdefer for value:
var value = (try reader.readUntilDelimiterOrEofAlloc(allocator, '\n', std.math.maxInt(usize))) orelse return error.EndOfStream;
+ errdefer allocator.free(value);
And run the test again:
All 1 tests passed.
With this, we can be reasonably confident that if any of the allocations that occur within the test fail, we handle the OutOfMemory error without introducing more problems.
We might not always want to run our test code N+1 times (where N is the number of allocations that occur within the test). Luckily, once a checkAllAllocationFailures-compatible function is written, it's easy to switch between using it with checkAllAllocationFailures and calling it directly:
const check_allocation_failures = true;
test {
// (omitted, same as previous example)
if (check_allocation_failures) {
try std.testing.checkAllAllocationFailures(std.testing.allocator, parseTest, .{ &stream_source, expected });
} else {
try parseTest(testing.allocator, &stream_source, expected);
}
}
build.zig🔗To make this nicer, we can make the check_allocation_failures constant an option within build.zig so that we can disable it by doing something like zig build test -Dcheck-allocation-failures=false.
pub fn build(b: *std.build.Builder) void {
// ...
// Create the test step
const main_tests = b.addTest("src/main.zig");
// Create the command line option (with a default of true)
const check_allocation_failures = b.option(bool, "check-allocation-failures", "Run tests with checkAllAllocationFailures (default: true)") orelse true;
// Create the option using the value gotten from the command line
const test_options = b.addOptions();
test_options.addOption(bool, "check_allocation_failures", check_allocation_failures);
// Add the options as "test_options" to the main_tests step
// Our option can then be accessed via `@import("test_options").check_allocation_failures`
main_tests.addOptions("test_options", test_options);
// ...
}
which then can be used like so:
test {
// ...
if (@import("test_options").check_allocation_failures) {
try std.testing.checkAllAllocationFailures(std.testing.allocator, parseTest, .{ &stream_source, expected });
} else {
try parseTest(std.testing.allocator, &stream_source, expected);
}
}
Earlier, I said that:
As long as the number of memory allocations is deterministic, this strategy works
However, it turns out that the code we've been testing can actually have non-deterministic memory usage of a sort (at least with the implementation of std.BufMap as of this article's writing). For example, if we use the following input for our parse function:
foo=bar
baz=dup
a=b
b=c
c=d
d=e
baz=qux
then when running with checkAllAllocationFailures, we hit a scenario in which:
BufMap are inserted for foo, baz, a, b, c, and d successfullyBufMap would need to grow to be able to insert another entrytry map.putMove(key, value) is called for baz=qux, but the allocation for trying to grow the map is made to fail by checkAllAllocationFailuresstd.BufMap.putMove can recover from OutOfMemory if the key that is trying to be inserted is found in the map, which is the case here (the baz key was inserted previously, so it can return that previous entry)putMove succeeds and therefore OutOfMemory is not returnedThis means that, although an OutOfMemory error was induced, our parseTest call will succeed, which triggers checkAllAllocationFailures to return error.NondeterministicMemoryUsage and fail the test, as it assumes that all calls of the function with an induced allocation failure will have a return of error.OutOfMemory.
This is something of a false positive in terms of non-determinism, though, as the above scenario is still deterministic, but the OutOfMemory in one particular case is handled without bubbling up the error.
Since we know that this is a false-positive, we can ignore error.NondeterministicMemoryUsage by catching it like so:
std.testing.checkAllAllocationFailures(
std.testing.allocator,
parseTest,
.{ &stream_source, expected },
) catch |err| switch (err) {
error.NondeterministicMemoryUsage => {},
else => |e| return e,
};
This should generally be avoided, though, as treating error.NondeterministicMemoryUsage as a bug by default makes sense. Unless you know that part of the code you're testing has OutOfMemory recovery in place somewhere (like std.BufMap.putMove), then it's generally a good idea to ensure that the code under test doesn't erroneously/unexpectedly 'swallow' OutOfMemory errors.
If your code's memory allocation is truly non-deterministic in the sense that subsequent runs could have more points of allocation than the initial run, then ignoring the error.NondeterministicMemoryUsage is inadvisable, as the strategy used by checkAllAllocationFailures would no longer be guaranteed to provide full coverage of all possible points of allocation failure.
For projects where fuzz testing makes sense, it's possible to use checkAllAllocationFailures alongside fuzz testing to find bugs related to OutOfMemory error handling that are not (yet) covered by your test cases.
For this, we'll need to create a modified version of our parseTest function from before where:
expected, as we don't have that information for fuzzed inputsOutOfMemory when calling parse, since we want to allow invalid inputsfn parseTest(allocator: std.mem.Allocator, stream_source: *std.io.StreamSource) !void {
try stream_source.seekTo(0);
if (parse(allocator, stream_source)) |*parsed| {
parsed.deinit();
} else |err| {
switch (err) {
// We only want to return the error if it's OutOfMemory
error.OutOfMemory => return error.OutOfMemory,
// Any other error is fine since not all inputs will be valid
else => {},
}
}
}
And then we'll need a fuzzing-compatible main which looks something like this (again, see here for more info):
const std = @import("std");
const parse = @import("main.zig").parse;
pub export fn main() void {
zigMain() catch unreachable;
}
pub fn zigMain() !void {
// Set up a GeneralPurposeAllocator so that we can also catch double-frees, etc
var gpa = std.heap.GeneralPurposeAllocator(.{}){};
defer std.debug.assert(gpa.deinit() == false);
const allocator = gpa.allocator();
// Get the fuzzed input form stdin and create a StreamSource with it so we can
// pass that to parseTest via checkAllAllocationFailures
const stdin = std.io.getStdIn();
const data = try stdin.readToEndAlloc(allocator, std.math.maxInt(usize));
defer allocator.free(data);
var stream_source = std.io.StreamSource{ .buffer = std.io.fixedBufferStream(data) };
// Call checkAllAllocationFailures, but ignore error.NondeterministicMemoryUsage
// (normally you wouldn't ignore NondeterministicMemoryUsage, but it's necessary in our
// case because we use `std.BufMap.putMove` which has an OutOfMemory recovery strategy)
std.testing.checkAllAllocationFailures(allocator, parseTest, .{&stream_source}) catch |err| switch (err) {
error.NondeterministicMemoryUsage => {},
else => |e| return e,
};
}
The simple parse function used as an example in this post is not very exciting in terms of fuzzing, unfortunately. Besides the error.NondeterministicMemoryUsage caveat, there's nothing more to be found once we've added in the errdefer's mentioned previously (and the version without the errdefer's would trigger a crash with any reasonable seed input, so afl-fuzz would refuse to fuzz until that is fixed). In more complex projects, though, fuzzing can be very helpful in finding novel OutOfMemory-related bugs.
A code search on GitHub for checkAllAllocationFailures comes up with a few projects that have already started using it:
squeek502/audiometa (my project)
checkAllAllocationFailures is now used to ensure nothing regresses."It already helped to find a dangling pointer error!"
It seems possible that this strategy could be integrated into the test runner itself, which would remove having to manually add integration on a test-by-test basis. A similar type of integration is already included for leak checking via std.testing.allocator, as the test runner initializes and deinits a new GeneralPurposeAllocator for you for each test and reports the results.
If this is done for checking allocation failures, then that'd allow anyone to run all their tests (presumably only when a command line flag is set) with the checkAllAllocationFailures strategy (given they use std.testing.allocator, although that might depend on the implementation).