dataloader.zig

const std = @import("std");
const xev = @import("xev");
const concurrent_ring = @import("concurrent_ring.zig");

const logger = std.log.scoped(.dataloader);
const wlog = std.log.scoped(.dataloader_io_thread);

// kqueue (macOS) dispatches regular-file I/O to a thread pool; without one
// every read returns EPERM. io_uring (Linux) handles file I/O in-kernel,
// so we skip the pool there to avoid idle worker threads.
const needs_thread_pool = @import("builtin").os.tag != .linux;
const ThreadPoolStorage = if (needs_thread_pool) xev.ThreadPool else void;

pub const FileHandle = packed struct {
    idx: u20,
    generation: u20,
    path_checksum: u8,
    _: u16 = 0,
};

const max_file_slots = std.math.maxInt(u20);
const max_generation = std.math.maxInt(u20);

pub const ReadBlockReq = struct {
    base: u64,
    file: FileHandle,
    result_buffer: []u8,
};

pub const Request = union(enum) {
    open_file: struct {
        file_path: []const u8,
    },
    close_file: FileHandle,
    read_block: ReadBlockReq,
    drain: struct {},
};

pub const ResponsePayload = union(enum) {
    open_file: FileHandle,
    read_block: struct {},
};

pub const Response = struct {
    request_id: u64,
    payload: LoaderError!ResponsePayload,
};

pub const LoaderError = error{
    TooManyOpenFiles,
    InvalidFileHandle,
    ReadError,
} || std.Io.File.OpenError || std.mem.Allocator.Error;

fn path_checksum(path: []const u8) u8 {
    var checksum: u8 = 0;
    for (path) |c| {
        checksum +%= c;
    }
    return checksum;
}

const XevReq = struct {
    req: ReadBlockReq,
    c: xev.Completion = undefined,
    request_id: u64 = 0,
};

pub const LoaderCtx = struct {
    const Self = @This();

    // Result ring is larger than request ring: the loader thread may be shared
    // with Python and stall on dequeue, while the IO thread polls tightly.
    const IdxedRequest = struct {
        payload: Request,
        request_id: u64,
    };
    const ReqRing = concurrent_ring.SPSCRing(7, IdxedRequest);
    const ResultRing = concurrent_ring.SPSCRing(63, Response);

    alloc: std.mem.Allocator,
    io: std.Io = undefined,

    request_ring: ReqRing = ReqRing.init(),
    result_ring: ResultRing = ResultRing.init(),

    file_slots: [max_file_slots]?std.Io.File = @splat(null),
    xfile_slots: [max_file_slots]xev.File = undefined,
    file_refcount: [max_file_slots]u32 = @splat(0),
    file_slots_generation: [max_file_slots]u32 = std.mem.zeroes([max_file_slots]u32),
    file_slots_checksum: [max_file_slots]u8 = undefined,

    // Owned thread pool for libxev's kqueue backend (macOS): regular-file
    // I/O is dispatched here, otherwise reads fail EPERM. io_uring (Linux)
    // does file I/O in-kernel and doesn't need it, so this is `void` there.
    thread_pool: ThreadPoolStorage = if (needs_thread_pool) undefined else {},
    loop: xev.Loop,
    worker_thread: ?std.Thread = null,
    worker_done: std.atomic.Value(bool) = std.atomic.Value(bool).init(false),

    req_cnt: u64 = 0,
    is_running: bool = false,
    is_draining: bool = false,
    last_slot: usize = 0,

    debug_max_req_id: u64 = std.math.maxInt(u64),
    tick: u64 = 0,
    debug_max_tick: u64 = std.math.maxInt(u64),

    req_mem_pool: std.heap.MemoryPool(XevReq),

    fn findFreeFileSlot(self: *Self) !FileHandle {
        // Linear scan: open/close is rare relative to reads.
        for (0..max_file_slots) |offset| {
            const i = (self.last_slot + offset) % max_file_slots;
            const f = self.file_slots[0..][i];
            if (f == null) {
                self.last_slot = (i + 1) % max_file_slots;
                return .{ .idx = @intCast(i), .generation = 0, .path_checksum = 0 };
            }
        }

        return LoaderError.TooManyOpenFiles;
    }

    fn checkFilehandle(self: *Self, file: FileHandle) !void {
        const slot: usize = file.idx;
        const file_slots = self.file_slots[0..];
        const generations = self.file_slots_generation[0..];
        const checksums = self.file_slots_checksum[0..];

        if (slot >= file_slots.len) return LoaderError.InvalidFileHandle;
        if (file_slots[slot] == null) return LoaderError.InvalidFileHandle;
        if (generations[slot] != file.generation or checksums[slot] != file.path_checksum) {
            logger.warn("File handle {} is corrupted, current generation: {}, checksum: {}", .{ file, generations[slot], checksums[slot] });
            return LoaderError.InvalidFileHandle;
        }
    }

    fn sendResponseSynced(self: *Self, req_id: u64, resp: LoaderError!ResponsePayload) void {
        while (true) {
            self.result_ring.enqueue(.{
                .request_id = req_id,
                .payload = resp,
            }) catch {
                std.Thread.yield() catch {};
                std.atomic.spinLoopHint();
                continue;
            };

            break;
        }
    }

    fn xevReadCb(
        ud: ?*Self,
        _: *xev.Loop,
        c: *xev.Completion,
        _: xev.File,
        _: xev.ReadBuffer,
        r: xev.ReadError!usize,
    ) xev.CallbackAction {
        const self = ud orelse unreachable;

        const xreq: *XevReq = @fieldParentPtr("c", c);
        const slot: usize = xreq.req.file.idx;
        const request_id = xreq.request_id;

        const actual = r catch {
            self.fileDecRef(slot);
            self.req_mem_pool.destroy(xreq);
            self.sendResponseSynced(request_id, LoaderError.ReadError);
            return .disarm;
        };
        self.fileDecRef(slot);

        if (actual != xreq.req.result_buffer.len) {
            std.debug.panic("Unhandled error: {} != {}", .{ actual, xreq.req.result_buffer.len });
        }

        self.sendResponseSynced(xreq.request_id, .{ .read_block = .{} });

        self.req_mem_pool.destroy(xreq);
        return .disarm;
    }

    fn fileAddRef(self: *Self, slot: usize) void {
        std.debug.assert(self.file_refcount[0..][slot] > 0);
        self.file_refcount[0..][slot] += 1;
    }

    fn fileDecRef(self: *Self, slot: usize) void {
        std.debug.assert(self.file_refcount[0..][slot] > 0);
        self.file_refcount[0..][slot] -= 1;
        if (self.file_refcount[0..][slot] == 0) {
            const f = self.file_slots[0..][slot] orelse unreachable;
            f.close(self.io);
            self.file_slots[0..][slot] = null;
        }
    }

    fn handleReq(self: *Self, req_id: u64, req: Request) void {
        switch (req) {
            .open_file => |open_req| {
                wlog.debug("Req {}: open_file: file = {s}", .{ req_id, open_req.file_path });

                var h = self.findFreeFileSlot() catch |err| {
                    self.sendResponseSynced(req_id, err);
                    return;
                };

                const f = std.Io.Dir.cwd().openFile(self.io, open_req.file_path, .{ .mode = .read_only }) catch |err| {
                    self.sendResponseSynced(req_id, err);
                    return;
                };
                const xf = xev.File.init(f) catch unreachable;

                const slot: usize = h.idx;
                const checksum = path_checksum(open_req.file_path);
                self.file_slots[slot] = f;
                self.xfile_slots[slot] = xf;
                self.file_refcount[slot] = 1;
                self.file_slots_generation[slot] += 1; // gen 0 is reserved to catch errors
                self.file_slots_checksum[slot] = checksum;
                const gen = self.file_slots_generation[slot];
                if (gen > max_generation) @panic("Open file generation overflow");
                h = .{ .idx = @intCast(slot), .generation = @intCast(gen), .path_checksum = checksum };

                self.sendResponseSynced(req_id, .{ .open_file = h });
            },

            .close_file => |file_handle| {
                wlog.debug("Req {}: close_file: {}", .{ req_id, file_handle });

                self.checkFilehandle(file_handle) catch |err| {
                    self.sendResponseSynced(req_id, err);
                    return;
                };

                const slot: usize = file_handle.idx;

                self.fileDecRef(slot);
            },

            .read_block => |read_req| {
                wlog.debug("Req {}: read_block: file = {}, base = {}, size = {}", .{ req_id, read_req.file, read_req.base, read_req.result_buffer.len });

                self.checkFilehandle(read_req.file) catch |err| {
                    self.sendResponseSynced(req_id, err);
                    return;
                };

                const slot: usize = read_req.file.idx;
                const xf = self.xfile_slots[slot];

                var xreq = self.req_mem_pool.create(self.alloc) catch |err| {
                    self.sendResponseSynced(req_id, err);
                    return;
                };
                xreq.request_id = req_id;
                xreq.req = read_req;

                self.fileAddRef(slot);

                xf.pread(
                    &self.loop,
                    &xreq.c,
                    .{ .slice = read_req.result_buffer },
                    read_req.base,
                    Self,
                    self,
                    Self.xevReadCb,
                );
            },

            .drain => {
                self.is_draining = true;
            },
        }
    }

    /// Worker loop. Caller must keep the backing Threaded alive for the worker's lifetime.
    fn run_worker(self: *Self, io: std.Io) void {
        self.io = io;
        while (self.is_running) {
            self.tick += 1;
            if (self.tick > self.debug_max_tick) {
                @panic("Tick overflow");
            }

            if (self.request_ring.dequeue()) |req| {
                self.handleReq(req.request_id, req.payload);
            }

            self.loop.run(.no_wait) catch |err| {
                std.debug.panic("Error in event loop: {}", .{err});
            };
            if (self.loop.active == 0) {
                if (self.is_draining) {
                    wlog.debug("No more requests & IO loop drained.", .{});
                    self.is_running = false;
                    break;
                }
                std.Thread.yield() catch {};
            }
        }

        self.worker_done.store(true, .release);
    }

    fn drainResponse(self: *Self) void {
        var count: usize = 0;
        while (self.result_ring.dequeue()) |_| {
            std.atomic.spinLoopHint();
            count += 1;
        }
        if (count > 0) {
            logger.debug("Drained {} responses", .{count});
        }
    }

    pub fn sendSynced(self: *Self, req: Request) u64 {
        while (true) {
            if (self.trySend(req)) |rid| {
                return rid;
            }
            std.atomic.spinLoopHint();
            std.Thread.yield() catch {};
        }
    }

    pub fn recvSynced(self: *Self) Response {
        while (true) {
            if (self.result_ring.dequeue()) |r| {
                return r;
            }

            std.Thread.yield() catch {};
        }
    }

    /// Non-blocking enqueue. Returns request id (> 0) on success, null if the ring is full.
    pub fn trySend(self: *Self, req: Request) ?u64 {
        self.req_cnt += 1;
        const req_id = self.req_cnt;
        self.request_ring.enqueue(.{ .payload = req, .request_id = req_id }) catch {
            self.req_cnt -= 1; // Safe rollback: SPSC, only this thread writes req_cnt.
            return null;
        };
        return req_id;
    }

    pub fn tryRecv(self: *Self) ?Response {
        return self.result_ring.dequeue();
    }

    pub fn join(self: *Self) void {
        const thread = self.worker_thread orelse @panic("Worker thread not started");
        logger.debug("Joining worker thread", .{});
        // Drain results while sending drain: a full result ring stalls the worker,
        // which would otherwise leave the request ring full and block trySend.
        while (self.trySend(.{ .drain = .{} }) == null) {
            self.drainResponse();
            std.Thread.yield() catch {};
        }
        while (!self.worker_done.load(.acquire)) {
            self.drainResponse();
            std.Thread.yield() catch {};
        }
        self.drainResponse();
        thread.join();
        self.worker_thread = null;
    }

    pub fn start(self: *Self, io: std.Io) !void {
        if (self.worker_thread) |_| {
            @panic("Worker thread already started");
        }

        self.io = io;
        self.is_running = true;
        self.worker_done.store(false, .monotonic);

        self.worker_thread = try std.Thread.spawn(.{
            .allocator = self.alloc,
        }, Self.run_worker, .{ self, io });
        self.worker_thread.?.setName(io, "dataloader_io_worker") catch {};

        logger.debug("Worker thread started", .{});
    }

    pub fn initInPlace(self: *Self, alloc: std.mem.Allocator) !void {
        self.alloc = alloc;
        self.io = undefined;
        self.request_ring = ReqRing.init();
        self.result_ring = ResultRing.init();
        // These arrays are several MB each; assigning a literal RHS (e.g.
        // `@splat(0)` / `std.mem.zeroes(...)`) materializes the value on the
        // stack first, which blows the 8 MB stack on macOS at Python call
        // depth. Zero / set in place instead.
        @memset(&self.file_slots, null);
        self.xfile_slots = undefined;
        @memset(&self.file_refcount, 0);
        @memset(&self.file_slots_generation, 0);
        self.file_slots_checksum = undefined;
        if (needs_thread_pool) {
            self.thread_pool = .init(.{});
        }
        errdefer if (needs_thread_pool) {
            self.thread_pool.shutdown();
            self.thread_pool.deinit();
        };
        self.loop = try xev.Loop.init(.{
            .thread_pool = if (needs_thread_pool) &self.thread_pool else null,
        });
        errdefer self.loop.deinit();
        self.worker_thread = null;
        self.worker_done = std.atomic.Value(bool).init(false);
        self.req_cnt = 0;
        self.is_running = false;
        self.is_draining = false;
        self.last_slot = 0;
        self.debug_max_req_id = std.math.maxInt(u64);
        self.tick = 0;
        self.debug_max_tick = std.math.maxInt(u64);
        self.req_mem_pool = try std.heap.MemoryPool(XevReq).initCapacity(alloc, 16);
    }

    pub fn deinit(self: *Self) void {
        if (self.worker_thread) |_| {
            self.join();
        }
        self.loop.deinit();
        if (needs_thread_pool) {
            self.thread_pool.shutdown();
            self.thread_pool.deinit();
        }
        self.req_mem_pool.deinit(self.alloc);
    }
};

test "test dataloader" {
    const builtin = @import("builtin");

    if (builtin.os.tag != .linux) {
        return error.SkipZigTest;
    }

    const io = std.testing.io;

    const f = try std.Io.Dir.cwd().createFile(io, "testfile.tar", .{ .truncate = true });
    var rng = std.Random.DefaultPrng.init(42);
    var ref: [256]u8 = undefined;
    rng.fill(&ref);
    try f.writeStreamingAll(io, &ref);
    f.close(io);

    var debug_alloc = std.heap.DebugAllocator(.{}).init;
    defer _ = debug_alloc.deinit();

    const ctx = try debug_alloc.allocator().create(LoaderCtx);
    defer debug_alloc.allocator().destroy(ctx);
    try ctx.initInPlace(debug_alloc.allocator());
    ctx.debug_max_req_id = 5;
    ctx.debug_max_tick = 1000;
    try ctx.start(io);
    defer ctx.deinit();

    const open_rid = ctx.sendSynced(.{ .open_file = .{ .file_path = "testfile.tar" } });
    const resp = ctx.recvSynced();
    try std.testing.expect(open_rid == resp.request_id);
    const file = (try resp.payload).open_file;
    var buf: [256 - 42]u8 = undefined;
    const read_rid = ctx.sendSynced(.{ .read_block = .{ .file = file, .base = 42, .result_buffer = &buf } });
    const read_resp = ctx.recvSynced();
    try std.testing.expect(read_rid == read_resp.request_id);
    _ = try read_resp.payload;
    _ = ctx.sendSynced(.{ .close_file = file });
    ctx.join();

    try std.testing.expectEqualSlices(u8, ref[42..], &buf);
    try std.testing.expect(ctx.is_running == false);
    try std.testing.expectEqual(null, ctx.result_ring.dequeue());
}

test "test dataloader blocked join" {
    const io = std.testing.io;

    var debug_alloc = std.heap.DebugAllocator(.{}).init;
    defer _ = debug_alloc.deinit();

    const ctx = try debug_alloc.allocator().create(LoaderCtx);
    defer debug_alloc.allocator().destroy(ctx);
    try ctx.initInPlace(debug_alloc.allocator());
    ctx.debug_max_req_id = 100;
    ctx.debug_max_tick = 1000;
    try ctx.start(io);
    defer ctx.deinit();

    // Fill rings until trySend fails (worker stalled on full result ring).
    while (ctx.trySend(.{ .open_file = .{ .file_path = "/sys/does_not_exist" } })) |_| {}
    std.Io.sleep(io, .fromNanoseconds(1000), .awake) catch {};
    // join must drain results to unblock the worker.
    ctx.join();

    try std.testing.expect(ctx.is_running == false);
    try std.testing.expectEqual(null, ctx.result_ring.dequeue());
}

test "invalid file handle index is rejected" {
    var debug_alloc = std.heap.DebugAllocator(.{}).init;
    defer _ = debug_alloc.deinit();

    const ctx = try debug_alloc.allocator().create(LoaderCtx);
    defer debug_alloc.allocator().destroy(ctx);
    try ctx.initInPlace(debug_alloc.allocator());
    defer ctx.deinit();

    try std.testing.expectError(
        LoaderError.InvalidFileHandle,
        ctx.checkFilehandle(.{
            .idx = std.math.maxInt(u20),
            .generation = 0,
            .path_checksum = 0,
        }),
    );
}