ParkingMutex

const ParkingMutex = struct

File

Code

const ParkingMutex = struct {
    state: std.atomic.Value(State),

    const init: ParkingMutex = .{ .state = .init(.unlocked) };

    comptime {
        assert(use_parking_futex);
    }

    const State = enum(usize) {
        unlocked = 1,
        /// This value is intentionally 0 so that `waiter` returns `null`.
        locked_once = 0,
        /// Contended; value is a `*Waiter`.
        _,
        /// Returns the head of the waiter list. Illegal to call if `s == .unlocked`.
        fn waiter(s: State) ?*Waiter {
            return @ptrFromInt(@backingInt(s));
        }
        /// Returns a locked state where `w` is contending the lock.
        /// If `w` is `null`, returns `.locked_once`.
        fn fromWaiter(w: ?*Waiter) State {
            return @fromBackingInt(@intCast(@intFromPtr(w)));
        }
    };
    const Waiter = struct {
        unpark_flag: UnparkFlag,
        /// Never modified once the `Waiter` is in the linked list.
        next: ?*Waiter,
        /// Never modified once the `Waiter` is in the linked list.
        tid: std.Thread.Id,
    };
    fn lock(m: *ParkingMutex) void {
        state: switch (State.unlocked) { // assume 'unlocked' to optimize for uncontended case
            .unlocked => continue :state m.state.cmpxchgWeak(
                .unlocked,
                .locked_once,
                .acquire, // acquire lock
                .monotonic,
            ) orelse {
                @branchHint(.likely);
                return;
            },

            .locked_once, _ => |last_state| {
                const old_waiter = last_state.waiter();
                const self_tid = if (Thread.current) |t| t.id else std.Thread.getCurrentId();
                var waiter: Waiter = .{
                    .next = old_waiter,
                    .unpark_flag = unpark_flag_init,
                    .tid = self_tid,
                };
                if (m.state.cmpxchgWeak(
                    .fromWaiter(old_waiter),
                    .fromWaiter(&waiter),
                    .release, // release `waiter`
                    .monotonic,
                )) |new_state| {
                    continue :state new_state;
                }
                // We're now in the list of waiters---park until we're given the lock.
                park(.none, m, if (need_unpark_flag) &waiter.unpark_flag) catch |err| switch (err) {
                    error.Timeout => unreachable,
                };
                return;
            },
        }
    }
    fn unlock(m: *ParkingMutex) void {
        state: switch (State.locked_once) { // assume 'locked_once' to optimize for uncontended case
            .unlocked => unreachable, // we hold the lock

            .locked_once => continue :state m.state.cmpxchgWeak(
                .locked_once,
                .unlocked,
                .release, // release lock
                .acquire, // acquire any `Waiter` memory
            ) orelse {
                @branchHint(.likely);
                return;
            },

            _ => |last_state| {
                // The logic here does not have ABA problems, and does some accesses non-atomically,
                // because `Waiter.next` is owned by the lock holder (that's us!) once the waiter is
                // in the linked list, up until we unpark the waiter.

                // Run through the waiter list to the end to ensure fairness. This is obviously not
                // ideal, but it shouldn't be a big deal in practice provided the critical section
                // is fairly small (so we won't get too many threads contending the mutex at once).
                // There's a *chance* we could get away with a LIFO queue for our use case, but I
                // don't wanna risk that.
                var parent: ?*Waiter = null;
                var waiter: *Waiter = last_state.waiter().?;
                while (waiter.next) |next| {
                    parent = waiter;
                    waiter = next;
                }
                // `waiter` is next in line for the lock. Remove them from the list.
                if (parent) |p| {
                    assert(p.next == waiter);
                    p.next = null;
                } else {
                    // We're waking the last waiter, so clear the list head.
                    if (m.state.cmpxchgWeak(
                        .fromWaiter(last_state.waiter().?),
                        .locked_once,
                        .acquire,
                        .acquire, // acquire any new `Waiter` memory
                    )) |new_state| {
                        continue :state new_state;
                    }
                }
                // Now we're ready to actually hand the lock over to them.
                const tid = waiter.tid; // load before the unpark below potentially invalidates `waiter`
                if (need_unpark_flag) setUnparkFlag(&waiter.unpark_flag);
                unpark(&.{tid}, m);
                return;
            },
        }
    }
}