Looking for a Review. Kernel Page Allocator, Atomics Inside. Beware

[proper nimbus]

This is my first real zig code. So style, beginner stuff are all good suggestions.

I am looking for a review for on my kernel page allocator, it's a 2mib page allocator over raw heapspace. I want to be sure I have the atomics and logic correct before I put another layer over it for the kernel itself. This allocator at its core will be used for guest pages, the layer above will implement the zig allocator interface for the kernels own heap memory.

https://gitlab.com/vm-kit/vm-kit/-/blob/main/src/kalloc/root.zig

[safe rivet]

style comments:

• instead of a kalloc dir with a single root.zig inside, just have a kalloc.zig. even if you had more files, the convention would be a kalloc.zig that imports kalloc/whatever.zig.
  the zig init template calls the root of the module root.zig, but even then i personally prefer naming it <library name>.zig, which is what std does as well
• constants should be snake_case, not SCREAMING_SNAKE_CASE. global variables should especially not be screaming, since in eg. C libraries' typical conventions that means "global constant"
• consider putting the allocator state into a struct. even if you only have one global instance of that struct, it helps to compartmentalize it and makes things like testing easier
• consider implementing the std.mem.Allocator API
• return KallocError.SpaceNotFoundInPass; is weird for two reasons:
  1. you'd typically write this as return error.SpaceNotFoundInPass;. no need to specify the error set, since that's already part of the function's type
  2. the standard name for allocation failure errors is error.OutOfMemory
• instead of var SCATTER_CHECKER = std.atomic.Value(u32).init(0);, consider var scatter_checker: std.atomic.Value(u32) = .init(0);
• instead of var INUSE_PAGES: [PAGE_SIZE >> 6]std.atomic.Value(u64) = [_]std.atomic.Value(u64){std.atomic.Value(u64).init(0)} ** (PAGE_SIZE >> 6);, consider var inuse_pages: [page_size >> 6]std.atomic.Value(u64) = @splat(.init(0);
• instead of std.atomic.Value(u32).load(&SCATTER_CHECKER, .acquire); consider scatter_checker.load(.acquire) (and same for all the other places you're using functions from std.atomic.Value)
• return @as(*[PAGE_SIZE]u8, @ptrFromInt(address)); the @as here is redundant, simply return @ptrFromInt(address);

oh, also you're using snake_case for function names - that's fine if you wanna make that your project's style, but do be aware that the standard zig style uses camelCase for function names :)

[safe rivet]

questions:

• why is inuse_pages the length that it is? you're effectively doing @divExact(page_size, @bitSizeOf(u64)) i think, which would allocate a bit for each byte in a single page, which doesn't make any sense to me
• what's the deal with scatter_checker? it seems like the idea is to start the search after the page that was most recently allocated? but it'd be much better to do that with a thread local; sharing it between threads just means you add a ton of unnecessary contention
• what made you choose the bitset approach? a freelist will be simpler, faster, and use less memory, but they have the limitation that you can only allocate fixed size blocks - do you want to extend this to support allocating multiple sequential pages in future?

this logic seems wrong: https://gitlab.com/vm-kit/vm-kit/-/blob/main/src/kalloc/root.zig#L71-77
i think you meant to check (load & mask) != 0? currently it'll just infinite loop if there's another page allocated in that bitset chunk.
however, if you make that change, the loop does this:

• free the page
• check if the page is still free
• if not, free it again
  which is really really bad, for what should hopefully be obvious reasons ^-^

i'm not entirely sure what this loop is supposed to do, but to me it looks like it can be safely removed and replaced with just the single fetchAnd line

https://gitlab.com/vm-kit/vm-kit/-/blob/main/src/kalloc/root.zig#L64
zig convention would be to memset freed memory to undefined, not zero. unless you're doing this for security reasons ofc

though actually, i think zig's allocators do that memset on allocation, rather than on free

[proper nimbus]

wow. this is very helpful. thank you @safe rivet ! I have a lot to change, and learned a lot here. I can make pretty much all the style changes, but the soundness/logic things you pointed out. Let me ask some more questions or explain some of them. maybe that will help point out what I should do.

• why is inuse_pages the length that it is - each bit in inuse_pages is supposed to indicate if an index of the page space is in use. riscv has 512GiB guest mappable address space, so inuse_pages is that long page_size >> 6 to account for every possible index rolled into a u64. I plan to come back to 32 bit, so i wanted to use usize. However, INDEX_AVAIL indicates how many pages should be used from the START and END. I think what you are saying is then my logic is actually wrong. and ur dead on about making those changes for testing. although to me, it seemed hard to test all this. It's clear I need at least some testing. Please let me know if the logic there is wrong

• whats the deal with scatter_checker - The concept of threads don't really exist at this point, so im not sure what thread local does here (ill go look it up). The scatter checker starts basically where it is in alloc, but if it ran into a contention (someone else actively in the current scatter checker range of indexes) I am setting its index to a modulus of hopefully some random bits to get them away from each other.

• freelist - I'm only just dipping my toes in low level stuff, this is what megpt came up with. This kalloc is going to be the root of one other allocator for the rest of the kernel code. This current kalloc is just getting pages ready to hand to guest memory, but will use these pages internally. I would like sequential pages at some point for the kernel itself

• atomic stuff yea I think I don't understand atomics. my understanding is i need to check the previous value, to be sure it wasn't changed since I last read the value? otherwise i need to account for it in code, since there are no locks

[proper nimbus]

const ADDRESSABLE: u64 = 0x8000000000;
// possible improvement reduce size stored for fewer collisions
var inuse_pages: [(ADDRESSABLE / PAGE_SIZE) >> 6]std.atomic.Value(u64) = @splat(.init(0));

I added this. I think thats what you were pointing out

it looks like that number i had for 512GiB is one specific guest can have at max 512GiB using the paging system im going to use

im fine with my whole hypervisor having that much for now.

as for the atomic stuff, yea Im not seeing how that fetchAnd on the freeing is supposed to work. I have to be sure I don't accidentally set or unset bits.

[safe rivet]

riscv has 512GiB guest mappable address space
doesn't this mean the number of pages would be 512GiB / 2MiB, so the length of the bitset should be (512 << 30) / (2 << 20) / 64 = 4096?
your calculation is (2 << 20) / 64 = 32768, which would be the value to use if you wanted one bit per byte in a page, not one bit per page in the address space

[proper nimbus]

start_addr and end_addr is the actual space

inuse_pages is the tracker

for the index which offsets into that actual space

but inuse_pages is the whole possibility

index_avail is the number used

[safe rivet]

if it ran into a contention
any contention between threads will be resolved with a single atomic op. you're actually increasing contention by doing this because now every thread is constantly accessing the scatter_checker variable

[safe rivet]

i need to check the previous value, to be sure it wasn't changed since I last read the value
well, it depends. when you do fetchOr in alloc, you need to check the return value to make sure nobody else allocated the page you were trying to allocate. but when you do fetchAnd in free, you don't need to do any checking because you already know that the page is allocated, because you have ownership of it

[proper nimbus]

I think i see what you are saying on the scatter checker.

so with the scatter checker, i have it shared across actual hardware threads. it would be better to use my separate hart contexts and hav their own iterator which can also use entropy to set it for itself

[safe rivet]

yeah that looks a bit better :)
I'd write it like this:

const address_space = 512 << 30; // 512GiB
const page_bitset_len = @divExact(@divExact(address_space, page_size), @sizeOf(u64))
var inuse_pages: [page_bitset_len]std.atomic.Value(u64) = @splat(.init(0));

[proper nimbus]

screaming snake case is here to stay for global const variables 🙂

[safe rivet]

😭

i'll be the one screaming if i see another one /j

[proper nimbus]

alright alright

lowercase snake case. but the functions being snake case are here to stay as well

that just means that everything is lower case snake case haha

[safe rivet]

yeah that's fine, i don't have a particular preference on snake vs camel for functions :)

but for constants, it doesn't really make much sense to differentiate them from variables in zig because they're used for basically the exact same things

in other languages, constants often have extra restrictions on them, but in zig the only real difference is that variables are mutable

there's also basically no difference between a local and a global constant, so using different naming for them doesn't really make much sense

[proper nimbus]

is that because of the everything is a struct thingy

[safe rivet]

it's because of comptime, mostly

[safe rivet]

yeah, if you wanna store iterator state, it should be per-thread (i'm guessing "hart" is the riscv name for a hardware thread?). that also means it doesn't need to be atomic :)

[proper nimbus]

i had considered just doing the scatter checker without atomics, i guess that should have made it clear

[safe rivet]

i still don't think i fully understand what the scatter checker is actually supposed to do tbh

[proper nimbus]

move the index into inuse_pages around

[safe rivet]

also fwiw, this allocator is not going to be particularly efficient ^^'
hopefully ur okay with that :)

[proper nimbus]

so there aren't multiple contentions

[safe rivet]

contention between threads? you don't really have a way to detect that

[proper nimbus]

if i did a freelist i would need full on locks

[safe rivet]

nah, you can do atomic freelists

also, locks are fast and easy to implement, as long as you use them right

[proper nimbus]

where as this, the space is pretty large, each hart asking for pages, probably doesn't really run over with the same index

im now very confused about how to give each hart an instance of this, but keep the same atomics

[safe rivet]

const PageAllocator = struct {
    // thread-local state
    index: usize = 0,
    global: *GlobalState,

    const GlobalState = struct {
        // global state
        inuse_pages: [page_bitset_len]std.atomic.Value(u64) = @splat(.init(0)),
    }l
};

then you create one PageAllocator.GlobalState on startup, and create one PageAllocator per thread, pointing to that single GlobalState

[proper nimbus]

okay that is perfect. Thanks for all your help. when I get back on tonight, ill finish going over the atomics. I read Protty's dev.to article on atomics and I was hoping I got the .acquire, rel, unordered parts right

but that last one seems like a lapse in the and logic. I didn't want to accidentally set a bit that wasn't set, but fetchAnd with the inverted mask should work

[proper nimbus]

also, im not sure how. this has less atomic access than a freelist. maybe you can enlighten me here. especially since we got rid of the scatter checker atomic

[safe rivet]

i think they're mostly right, though using unordered is a bit weird (but should be fine in the context you used it, i think)

[safe rivet]

atomic linked lists are pretty fast. there are some finicky implementation details to do with ABA race conditions, but it's not too bad

with a freelist, you do maybe 3 or 4 atomic ops to pop one item, as opposed to your iteration approach where you could get stuck in a bad spot and have to iterate through potentially hundreds or thousands of used pages

[proper nimbus]

ahh, yea u are right

especially when close to full

[safe rivet]

mhm, though filling up 512GiB of ram would be impressive ^-^

[proper nimbus]

some of the larger cloud computers at amazon have 32TiB

[safe rivet]

😳

[proper nimbus]

and its vms all the way down

[safe rivet]

das a lotta ram

[proper nimbus]

amazing review thank you again @safe rivet this is my first foray into zig. and i really like it. very ergonomic

[safe rivet]

happy to help! ^-^

glad you're enjoying it :)

[dapper hamlet]

I'd prefer SCREAMING_SNAKE_CASE for all compile-time constants over that.

[proper nimbus]

hehe

[proper nimbus]

do you have an example of this somewhere? This syntax is driving me crazy coming from any other language. global is a pointer to GlobalState, but the GlobalState type is defined in the PageAllocator

[safe rivet]

its just a namespacing thing

a struct's fields can depend on other things defined inside that struct's namespace just fine, as long as it doesn't cause any dependency loops

[proper nimbus]

I think i managed 🙂

const page_size: u32 = 2 * 1024 * 1024;
const address_space: u64 = 512 << 30;
const page_bitset_len = @divExact(@divExact(address_space, page_size), @sizeOf(u64));

var global_state: Kallocator.GlobalState = .{};

pub const Kallocator = struct {
    start_addr: usize = 0,
    end_addr: usize = 0,
    global: *GlobalState,
    scatter_checker: u32 = 0,
    index_avail: u16 = 0,

    pub const GlobalState = struct {
        start_addr: usize = 0,
        end_addr: usize = 0,
        inuse_pages: [page_bitset_len]std.atomic.Value(u64) = @splat(.init(0)),
        index_avail: u16 = 0,
        pub fn early_init(start_addr: usize, end_addr: usize) void {
            global_state.start_addr = start_addr;
            global_state.end_addr = end_addr;
            global_state.index_avail = @intCast(((end_addr - start_addr) >> 21) - 1);
        }
    };

this way I can have locality for everything but that global inuse_pages