Arena comptime lock

Help
,
1

I’m using several arenas and I lock them when I use them, so that I don’t accidentally reset them further down the stack, and can only get/reset them in one scope:

code 
//! Arenas used by the editor.

/// Temporary allocations that don't survive function scope.
temp: Arena = .{},

/// Used when drawing rows.
render: Arena = .{},

/// Used when building edit targets.
targets: Arena = .{},

///////////////////////////////////////////////////////////////////////////////
//
//                              Types
//
///////////////////////////////////////////////////////////////////////////////

const Kind = enum { temp, targets, render };

const Arena = struct {
    /// Underlying arena instance.
    instance: std.heap.ArenaAllocator = undefined,

    /// True when an arena is already being used and shouldn't be reset.
    ///
    /// For example, if an outer scope starts using an arena, resetting it on
    /// scope exit, but it calls another function that does the same, if the
    /// second function resets the arena it invalidates all memory.
    ///
    /// To avoid this, arenas can only be used in one scope, and if they are
    /// needed further down their allocator must be passed explicitly.
    busy: bool = false,
};

///////////////////////////////////////////////////////////////////////////////
//
//                              Init/deinit
//
///////////////////////////////////////////////////////////////////////////////

pub fn init(gpa: mem.Allocator) Arenas {
    return .{
        .temp = .{ .instance = .init(gpa) },
        .render = .{ .instance = .init(gpa) },
        .targets = .{ .instance = .init(gpa) },
    };
}

pub fn deinit(self: *const Arenas) void {
    self.temp.instance.deinit();
    self.render.instance.deinit();
    self.targets.instance.deinit();
}

///////////////////////////////////////////////////////////////////////////////
//
//                              Methods
//
///////////////////////////////////////////////////////////////////////////////

/// Resets arenas with different strategies.
pub fn reset(self: *Arenas, comptime which: Kind) void {
    switch (which) {
        .temp => {
            self.temp.busy = false;
            _ = self.temp.instance.reset(.{ .retain_with_limit = 1024 * 1024 });
        },
        .targets => {
            self.targets.busy = false;
            _ = self.targets.instance.reset(.{ .retain_with_limit = 1024 * 1024 });
        },
        .render => {
            self.render.busy = false;
            _ = self.render.instance.reset(.retain_capacity);
        },
    }
}

/// Returns an arena allocator.
pub fn get(self: *Arenas, comptime which: Kind) !mem.Allocator {
    switch (which) {
        .temp => {
            if (self.temp.busy) {
                return error.ArenaIsBusy;
            }
            else {
                self.temp.busy = true;
                return self.temp.instance.allocator();
            }
        },

        .targets => {
            if (self.targets.busy) {
                return error.ArenaIsBusy;
            }
            else {
                self.targets.busy = true;
                return self.targets.instance.allocator();
            }
        },

        .render => {
            if (self.render.busy) {
                return error.ArenaIsBusy;
            }
            else {
                self.render.busy = true;
                return self.render.instance.allocator();
            }
        },
    }
}

///////////////////////////////////////////////////////////////////////////////
//
//                              Tests
//
///////////////////////////////////////////////////////////////////////////////

test "Arenas.get rejects a busy arena until reset" {
    inline for ([_]Kind{ .temp, .targets, .render }) |kind|
    {
        var arenas = Arenas.init(std.testing.allocator);
        defer arenas.deinit();

        _ = try arenas.get(kind);
        try testing.expectError(error.ArenaIsBusy, arenas.get(kind));

        arenas.reset(kind);
        _ = try arenas.get(kind);
    }
}

test "Arenas.reset only clears the requested arena" {
    var arenas = Arenas.init(std.testing.allocator);
    defer arenas.deinit();

    _ = try arenas.get(.temp);
    arenas.reset(.targets);

    try testing.expectError(error.ArenaIsBusy, arenas.get(.temp));
    _ = try arenas.get(.targets);
}

///////////////////////////////////////////////////////////////////////////////
//
//                              Constants
//
///////////////////////////////////////////////////////////////////////////////

const Arenas = @This();

const std = @import("std");
const mem = std.mem;
const testing = std.testing;

But I’m wondering if these busy checks can be moved to comptime, because from the functions flow it should be evident already at comptime if an arena is being used where it shouldn’t (because it’s being used somewhere up in the stack). Is this somehow possible?

2

This sort of comptime property enforcement / tracking seems difficult and not very ergonomic in Zig, you are basically adhoc re-inventing rust style borrow mechanics[1] where you try to statically ensure that some resource is used at most once.

I think it is better to stay in Zig style and just write normal code being aware of where something is allowed or not. If you really want safety checks you could for example use std.debug.SafetyLocks similar to std.ArrayHashMapUnmanaged’s (un-/lockPointers).

You would use one safety lock per arena, then you could replace the get and reset functions with two functions acquire and release the former one just locks the lock and returns the allocator, while the latter unlocks the lock and then calls reset. (You also could put the call to reset in acquire but then you have one unnecessary reset of the initially fresh arena)

That way the two functions create a span of time between which any nested calls to acquire or release cause an assertion failure in safe modes.

In unsafe modes all the safety locks change to zero sized types and get optimized out to noops.
This also means that you don’t need any booleans to track the state.

  1. or trying to use linear types, which isn’t a good fit for what comptime supports (it would be possible to model but tends to be very verbose because you are basically emulating a more complex type system with a simpler one) ↩︎

4 Likes
3

Thanks, this is useful (or I could actually just keep busy field in safe builds only, void otherwise, and gate their use to runtime safety).

It’s what get/reset are currently doing. but “span of time”? I’m not using timers, the danger comes only from nested calls, but it’s easy to avoid if I just pass down the arena allocator where needed, so I’ll just keep this system in safe builds.

4

I guess “span of time” is a bad way to phrase it but everything that happens between acquire and release.

Yeah re-reflecting on your code it seems pretty equivalent to what I propose, but I find acquire and release more obvious, that there is something lock like here.

With a name like get I expect that I can call it whenever I want and that it would possible return an optional. And with arenas reset is also something you can call whenever you want.

That is already the equivalent of what SafetyLock does just that it uses a two state enum in safe modes and a one state enum in unsafe modes (which results in a zero sized type).

5

Thanks, acquire and release are more specific (I just liked the brevity at call site in reality for get), maybe I’ll change the name.

Updated code 
//! Arenas used by the editor.

/// Temporary allocations that don't survive function scope.
temp: Arena = .{},

/// Used when drawing rows.
render: Arena = .{},

/// Used when building edit targets.
targets: Arena = .{},

///////////////////////////////////////////////////////////////////////////////
//
//                              Types
//
///////////////////////////////////////////////////////////////////////////////

const Kind = enum { temp, targets, render };

const Arena = struct {
    /// Underlying arena instance.
    instance: std.heap.ArenaAllocator = undefined,

    /// .locked when an arena is already being used and shouldn't be reset.
    ///
    /// For example, if an outer scope starts using an arena, resetting it on
    /// scope exit, but it calls another function that does the same, if the
    /// second function resets the arena it invalidates all memory.
    ///
    /// To avoid this, arenas can only be used in one scope, and if they are
    /// needed further down their allocator must be passed explicitly.
    lock: std.debug.SafetyLock = .{},
};

///////////////////////////////////////////////////////////////////////////////
//
//                              Init/deinit
//
///////////////////////////////////////////////////////////////////////////////

pub fn init(gpa: mem.Allocator) Arenas {
    return .{
        .temp = .{ .instance = .init(gpa) },
        .render = .{ .instance = .init(gpa) },
        .targets = .{ .instance = .init(gpa) },
    };
}

pub fn deinit(self: *const Arenas) void {
    self.temp.instance.deinit();
    self.render.instance.deinit();
    self.targets.instance.deinit();
}

///////////////////////////////////////////////////////////////////////////////
//
//                              Methods
//
///////////////////////////////////////////////////////////////////////////////

/// Resets arenas with different strategies.
pub fn reset(self: *Arenas, comptime which: Kind) void {
    switch (which) {
        .temp => {
            self.temp.lock.unlock();
            _ = self.temp.instance.reset(.{ .retain_with_limit = 1024 * 1024 });
        },
        .targets => {
            self.targets.lock.unlock();
            _ = self.targets.instance.reset(.{ .retain_with_limit = 1024 * 1024 });
        },
        .render => {
            self.render.lock.unlock();
            _ = self.render.instance.reset(.retain_capacity);
        },
    }
}

/// Returns an arena allocator.
pub fn get(self: *Arenas, comptime which: Kind) mem.Allocator {
    switch (which) {
        .temp => {
            self.temp.lock.lock();
            return self.temp.instance.allocator();
        },

        .targets => {
            self.targets.lock.lock();
            return self.targets.instance.allocator();
        },

        .render => {
            self.render.lock.lock();
            return self.render.instance.allocator();
        },
    }
}

///////////////////////////////////////////////////////////////////////////////
//
//                              Constants
//
///////////////////////////////////////////////////////////////////////////////

const Arenas = @This();

const std = @import("std");
const mem = std.mem;
const testing = std.testing;
const builtin = @import("builtin");

const assert = std.debug.assert;