Fiber

const Fiber = struct

File

Code

const Fiber = struct {
    required_align: void align(4),
    context: Io.fiber.Context,
    link: union {
        awaiter: ?*Fiber,
        group: struct { prev: ?*Fiber, next: ?*Fiber },
    },
    status: union(enum) {
        queue_next: ?*Fiber,
        awaiting_group: Group,
        free_next: ?*Fiber,
    },
    cancel_status: CancelStatus,
    cancel_protection: CancelProtection,
    name: if (tracy.enable) [*:0]const u8 else void,

    var next_name: u64 = 0;

    const CancelStatus = packed struct(u32) {
        requested: bool,
        awaiting: Awaiting,

        const unrequested: CancelStatus = .{ .requested = false, .awaiting = .nothing };

        const Awaiting = enum(u31) {
            nothing = std.math.maxInt(u31),
            group = std.math.maxInt(u31) - 1,
            /// An io_uring fd.
            _,

            fn subWrap(lhs: Awaiting, rhs: Awaiting) Awaiting {
                return @fromBackingInt(@intCast(@backingInt(lhs) -% @backingInt(rhs)));
            }

            fn fromIoUringFd(fd: fd_t) Awaiting {
                const awaiting: Awaiting = @fromBackingInt(@intCast(fd));
                switch (awaiting) {
                    .nothing, .group => unreachable,
                    _ => return awaiting,
                }
            }

            fn toIoUringFd(awaiting: Awaiting) fd_t {
                switch (awaiting) {
                    .nothing, .group => unreachable,
                    _ => return @backingInt(awaiting),
                }
            }
        };

        fn changeAwaiting(
            cancel_status: *CancelStatus,
            old_awaiting: Awaiting,
            new_awaiting: Awaiting,
        ) bool {
            const old_cancel_status = @atomicRmw(CancelStatus, cancel_status, .Add, .{
                .requested = false,
                .awaiting = new_awaiting.subWrap(old_awaiting),
            }, .monotonic);
            assert(old_cancel_status.awaiting == old_awaiting);
            return old_cancel_status.requested;
        }
    };

    const CancelProtection = packed struct {
        user: Io.CancelProtection,
        acknowledged: bool,

        const unblocked: CancelProtection = .{ .user = .unblocked, .acknowledged = false };

        fn check(cancel_protection: CancelProtection) Io.CancelProtection {
            return @fromBackingInt(@intCast(@intFromBool(cancel_protection != unblocked)));
        }

        fn acknowledge(cancel_protection: *CancelProtection) void {
            assert(!cancel_protection.acknowledged);
            cancel_protection.acknowledged = true;
        }

        fn recancel(cancel_protection: *CancelProtection) void {
            assert(cancel_protection.acknowledged);
            cancel_protection.acknowledged = false;
        }

        test check {
            try std.testing.expectEqual(Io.CancelProtection.unblocked, check(.unblocked));
            try std.testing.expectEqual(Io.CancelProtection.blocked, check(.{
                .user = .unblocked,
                .acknowledged = true,
            }));
            try std.testing.expectEqual(Io.CancelProtection.blocked, check(.{
                .user = .blocked,
                .acknowledged = false,
            }));
            try std.testing.expectEqual(Io.CancelProtection.blocked, check(.{
                .user = .blocked,
                .acknowledged = true,
            }));
        }
    };

    const finished: ?*Fiber = @ptrFromInt(@alignOf(Fiber));

    const max_result_align: Alignment = .@"16";
    const max_result_size = max_result_align.forward(512);
    /// This includes any stack realignments that need to happen, and also the
    /// initial frame return address slot and argument frame, depending on target.
    const min_stack_size = 60 * 1024 * 1024;
    const max_context_align: Alignment = .@"16";
    const max_context_size = max_context_align.forward(1024);
    const max_closure_size: usize = @sizeOf(AsyncClosure);
    const max_closure_align: Alignment = .of(AsyncClosure);
    const allocation_size = std.mem.alignForward(
        usize,
        max_closure_align.max(max_context_align).forward(
            max_result_align.forward(@sizeOf(Fiber)) + max_result_size + min_stack_size,
        ) + max_closure_size + max_context_size,
        std.heap.page_size_max,
    );
    comptime {
        assert(max_result_align.compare(.gte, .of(Completion)));
        assert(max_result_size >= @sizeOf(Completion));
    }

    fn create(ev: *Evented) error{OutOfMemory}!*Fiber {
        const thread: *Thread = .current();
        if (@atomicRmw(?*Fiber, &thread.free_queue, .Xchg, finished, .acquire)) |free_fiber| {
            assert(free_fiber != finished);
            @atomicStore(?*Fiber, &thread.free_queue, free_fiber.status.free_next, .release);
            return free_fiber;
        }
        const active_threads = @atomicLoad(u32, &ev.threads.active, .acquire);
        for (0..@min(max_steal_free_search, active_threads)) |_| {
            defer thread.steal_free_search_index += 1;
            if (thread.steal_free_search_index == active_threads) thread.steal_free_search_index = 0;
            const steal_free_search_thread =
                &ev.threads.allocated[0..active_threads][thread.steal_free_search_index];
            if (steal_free_search_thread == thread) continue;
            const free_fiber =
                @atomicLoad(?*Fiber, &steal_free_search_thread.free_queue, .monotonic) orelse continue;
            if (free_fiber == finished) continue;
            if (@cmpxchgWeak(
                ?*Fiber,
                &steal_free_search_thread.free_queue,
                free_fiber,
                null,
                .acquire,
                .monotonic,
            )) |_| continue;
            @atomicStore(?*Fiber, &thread.free_queue, free_fiber.status.free_next, .release);
            return free_fiber;
        }
        @atomicStore(?*Fiber, &thread.free_queue, null, .monotonic);
        return @ptrCast(try ev.allocator().alignedAlloc(u8, .of(Fiber), allocation_size));
    }

    fn destroy(fiber: *Fiber) void {
        const thread: *Thread = .current();
        assert(fiber.status.queue_next == null);
        fiber.status = .{ .free_next = @atomicLoad(?*Fiber, &thread.free_queue, .acquire) };
        while (true) fiber.status.free_next = @cmpxchgWeak(
            ?*Fiber,
            &thread.free_queue,
            fiber.status.free_next,
            fiber,
            .acq_rel,
            .acquire,
        ) orelse break;
    }

    fn allocatedSlice(f: *Fiber) []align(@alignOf(Fiber)) u8 {
        return @as([*]align(@alignOf(Fiber)) u8, @ptrCast(f))[0..allocation_size];
    }

    fn allocatedEnd(f: *Fiber) [*]u8 {
        const allocated_slice = f.allocatedSlice();
        return allocated_slice[allocated_slice.len..].ptr;
    }

    fn resultPointer(f: *Fiber, comptime Result: type) *Result {
        return @ptrCast(@alignCast(f.resultBytes(.of(Result))));
    }

    fn resultBytes(f: *Fiber, alignment: Alignment) [*]u8 {
        return @ptrFromInt(alignment.forward(@intFromPtr(f) + @sizeOf(Fiber)));
    }

    const Queue = struct { head: *Fiber, tail: *Fiber };

    /// Like a `*Fiber`, but 2 bits smaller than a pointer (because the LSBs are always 0 due to
    /// alignment) so that those two bits can be used in a `packed struct`.
    const PackedPtr = enum(@Int(.unsigned, @bitSizeOf(usize) - 2)) {
        null = 0,
        all_ones = std.math.maxInt(@Int(.unsigned, @bitSizeOf(usize) - 2)),
        _,

        const Split = packed struct(usize) { low: u2, high: PackedPtr };
        fn pack(ptr: ?*Fiber) PackedPtr {
            const split: Split = @bitCast(@intFromPtr(ptr));
            assert(split.low == 0);
            return split.high;
        }
        fn unpack(ptr: PackedPtr) ?*Fiber {
            const split: Split = .{ .low = 0, .high = ptr };
            return @ptrFromInt(@as(usize, @bitCast(split)));
        }
    };

    fn requestCancel(fiber: *Fiber, ev: *Evented) void {
        const cancel_status = @atomicRmw(
            Fiber.CancelStatus,
            &fiber.cancel_status,
            .Or,
            .{ .requested = true, .awaiting = @fromBackingInt(@intCast(0)) },
            .acquire,
        );
        assert(!cancel_status.requested);
        switch (cancel_status.awaiting) {
            .nothing => {},
            .group => {
                // The awaiter received a cancelation request while awaiting a group,
                // so propagate the cancelation to the group.
                if (fiber.status.awaiting_group.cancel(ev, null)) {
                    fiber.status = .{ .queue_next = null };
                    _ = ev.schedule(.current(), .{ .head = fiber, .tail = fiber });
                }
            },
            _ => |awaiting| {
                const awaiting_io_uring_fd = awaiting.toIoUringFd();
                const thread: *Thread = .current();
                thread.enqueue().* = if (thread.io_uring.fd == awaiting_io_uring_fd) .{
                    .opcode = .ASYNC_CANCEL,
                    .flags = linux.IOSQE_CQE_SKIP_SUCCESS,
                    .ioprio = 0,
                    .fd = 0,
                    .off = 0,
                    .addr = @intFromPtr(fiber),
                    .len = 0,
                    .rw_flags = 0,
                    .user_data = @backingInt(Completion.Userdata.wakeup),
                    .buf_index = 0,
                    .personality = 0,
                    .splice_fd_in = 0,
                    .addr3 = 0,
                    .resv = 0,
                } else .{
                    .opcode = .MSG_RING,
                    .flags = linux.IOSQE_CQE_SKIP_SUCCESS,
                    .ioprio = 0,
                    .fd = awaiting_io_uring_fd,
                    .off = @intFromPtr(fiber) | 0b01,
                    .addr = @backingInt(linux.IORING_MSG_RING_COMMAND.DATA),
                    .len = 0,
                    .rw_flags = 0,
                    .user_data = @backingInt(Completion.Userdata.cleanup),
                    .buf_index = 0,
                    .personality = 0,
                    .splice_fd_in = 0,
                    .addr3 = 0,
                    .resv = 0,
                };
            },
        }
    }
}