FuzzMultiThreadedAllocator

const FuzzMultiThreadedAllocator = struct

File

Code

const FuzzMultiThreadedAllocator = struct {
    gpa: mem.Allocator,

    fill: usize,
    active_allocs: usize,
    fail_i: usize,
    fixed_remap_i: usize,

    // The below are assumed to be externally synchronized
    // i.e. each thread has an acquire fence before **first** using the allocator
    buf: []u8,
    fails: []const bool,
    fixed_remaps: []const bool,

    fn allocator(f: *FuzzMultiThreadedAllocator) mem.Allocator {
        @disableInstrumentation();
        return .{ .ptr = f, .vtable = &.{
            .alloc = FuzzMultiThreadedAllocator.alloc,
            .free = FuzzMultiThreadedAllocator.free,
            .resize = FuzzMultiThreadedAllocator.resize,
            .remap = FuzzMultiThreadedAllocator.remap,
        } };
    }

    fn maybeFail(f: *FuzzMultiThreadedAllocator) bool {
        @disableInstrumentation();
        const i = @atomicRmw(usize, &f.fail_i, .Add, 1, .monotonic);
        return i < f.fails.len and f.fails[i];
    }

    fn maybeFixedRemap(f: *FuzzMultiThreadedAllocator) bool {
        @disableInstrumentation();
        const i = @atomicRmw(usize, &f.fixed_remap_i, .Add, 1, .monotonic);
        return i < f.fixed_remaps.len and f.fixed_remaps[i];
    }

    fn alloc(ctx: *anyopaque, len: usize, alignment: Alignment, _: usize) ?[*]u8 {
        @disableInstrumentation();

        const f: *FuzzMultiThreadedAllocator = @ptrCast(@alignCast(ctx));
        const memory = f.allocInner(len, alignment) orelse return null;
        _ = @atomicRmw(usize, &f.active_allocs, .Add, 1, .monotonic);
        return memory;
    }

    fn allocInner(f: *FuzzMultiThreadedAllocator, len: usize, alignment: Alignment) ?[*]u8 {
        var prev_fill = @atomicLoad(usize, &f.fill, .monotonic);
        var start: usize = undefined;
        while (true) {
            const start_addr = alignment.forward(@intFromPtr(f.buf[prev_fill..].ptr));
            start = @as([*]u8, @ptrFromInt(start_addr)) - f.buf.ptr;
            if (start +| len > f.buf.len or f.maybeFail()) return null;
            prev_fill = @cmpxchgStrong(
                usize,
                &f.fill,
                prev_fill,
                start + len,
                .monotonic,
                .monotonic,
            ) orelse {
                @branchHint(.likely);
                break;
            };
        }
        return f.buf[start..][0..len].ptr;
    }

    fn free(ctx: *anyopaque, _: []u8, _: Alignment, _: usize) void {
        @disableInstrumentation();

        const f: *FuzzMultiThreadedAllocator = @ptrCast(@alignCast(ctx));
        assert(@atomicRmw(usize, &f.active_allocs, .Sub, 1, .monotonic) != 0);
    }

    fn resize(ctx: *anyopaque, memory: []u8, _: Alignment, new_len: usize, _: usize) bool {
        @disableInstrumentation();

        const f: *FuzzMultiThreadedAllocator = @ptrCast(@alignCast(ctx));
        const start = memory.ptr - f.buf.ptr;
        const old_end = start + memory.len;
        const new_end = start +| new_len;
        if (new_end > f.buf.len or f.maybeFail()) {
            return false;
        }

        if (new_len <= memory.len) {
            // The fill is not decreased so memory is not reused.
            return true;
        }

        return @cmpxchgStrong(usize, &f.fill, old_end, new_end, .monotonic, .monotonic) == null;
    }

    fn remap(ctx: *anyopaque, memory: []u8, alignment: Alignment, new_len: usize, _: usize) ?[*]u8 {
        @disableInstrumentation();

        if (maybeFixedRemap(@ptrCast(@alignCast(ctx)))) {
            const resized = FuzzMultiThreadedAllocator.resize(
                ctx,
                memory,
                alignment,
                new_len,
                undefined,
            );
            return if (resized) memory.ptr else null;
        }

        const f: *FuzzMultiThreadedAllocator = @ptrCast(@alignCast(ctx));
        const new_ptr = f.allocInner(new_len, alignment) orelse return null;
        const copy_len = @min(memory.len, new_len);
        @memcpy(new_ptr[0..copy_len], memory[0..copy_len]);
        return new_ptr;
    }
}