src/client/open_wait.zig

Ref: main worktree-web-client-spec   Size: 6.0 KiB   History

//! Shared cancellation and absolute deadline for one opening/request operation.
const std = @import("std");
const interrupt = @import("interrupt.zig");
pub const Wait = struct {
    alloc: std.mem.Allocator,
    abort_fd: std.posix.fd_t = -1,
    carry: ?*std.ArrayList(u8) = null,
    deadline: ?i64 = null,
    watching: bool = true,
    pub fn check(self: *Wait) !void {
        if (self.abort_fd >= 0 and self.watching) {
            var fds = [_]std.posix.pollfd{.{ .fd = self.abort_fd, .events = std.posix.POLL.IN, .revents = 0 }};
            _ = try std.posix.poll(&fds, 0);
            if (fds[0].revents != 0) {
                var buf: [256]u8 = undefined;
                const n = std.posix.read(self.abort_fd, &buf) catch 0;
                if (n == 0) self.watching = false;
                if (std.mem.indexOfScalar(u8, buf[0..n], interrupt.detach_key) != null) return error.UserAbort;
                if (self.carry) |out| try out.appendSlice(self.alloc, buf[0..n]);
            }
        }
        if (self.deadline) |end| if (std.time.milliTimestamp() >= end) return error.Timeout;
    }
    pub fn remaining(self: *Wait, cap: u32) !u32 {
        try self.check();
        return if (self.deadline) |end| @intCast(@min(@as(i64, cap), @max(0, end - std.time.milliTimestamp()))) else cap;
    }
    pub fn poll(self: *Wait, fds: []std.posix.pollfd, cap: u32) !void {
        std.debug.assert(fds.len < 8);
        var all: [8]std.posix.pollfd = undefined;
        @memcpy(all[0..fds.len], fds);
        all[fds.len] = .{ .fd = if (self.watching) self.abort_fd else -1, .events = std.posix.POLL.IN, .revents = 0 };
        const ms = try self.remaining(cap);
        _ = try std.posix.poll(all[0 .. fds.len + 1], @intCast(@min(ms, std.math.maxInt(i32))));
        try self.check();
        @memcpy(fds, all[0..fds.len]);
    }
};

/// A full Unix listen queue gives EAGAIN on Linux, without initiating a
/// connection. Retry connect after a cancellable delay; writable/SO_ERROR=0
/// does not prove such a socket connected. Other pending connects use SO_ERROR.
pub fn connectUnix(path: []const u8, wait: *Wait) !std.net.Stream {
    const addr = try std.net.Address.initUnix(path);
    const fd = try std.posix.socket(std.posix.AF.UNIX, std.posix.SOCK.STREAM | std.posix.SOCK.CLOEXEC | std.posix.SOCK.NONBLOCK, 0);
    errdefer std.posix.close(fd);
    while (true) {
        try wait.check();
        switch (std.posix.errno(std.posix.system.connect(fd, &addr.any, addr.getOsSockLen()))) {
            .SUCCESS, .ISCONN => {},
            .INTR => continue,
            .AGAIN => {
                var none: [0]std.posix.pollfd = .{};
                try wait.poll(&none, 10);
                continue;
            },
            .INPROGRESS, .ALREADY => {
                var fds = [_]std.posix.pollfd{.{ .fd = fd, .events = std.posix.POLL.OUT, .revents = 0 }};
                while (fds[0].revents == 0) try wait.poll(&fds, 1000);
                try std.posix.getsockoptError(fd);
            },
            .NOENT => return error.FileNotFound,
            .ACCES => return error.AccessDenied,
            .PERM => return error.PermissionDenied,
            .NOMEM, .NOBUFS => return error.SystemResources,
            else => return error.ConnectionRefused,
        }
        break;
    }
    try wait.check();
    const flags = try std.posix.fcntl(fd, std.posix.F.GETFL, 0);
    const bits: u32 = @bitCast(std.posix.O{ .NONBLOCK = true });
    _ = try std.posix.fcntl(fd, std.posix.F.SETFL, flags & ~@as(usize, bits));
    return .{ .handle = fd };
}

test "full Unix listen backlog remains cancellable and restores blocking only after connect" {
    if (@import("builtin").os.tag != .linux) return error.SkipZigTest;
    const a = std.testing.allocator;
    var tmp = try @import("testtmp").TmpDir.make();
    defer tmp.cleanup();
    const path = try std.fmt.allocPrint(a, "{s}/full.sock", .{tmp.path()});
    defer a.free(path);
    const addr = try std.net.Address.initUnix(path);
    var server = try addr.listen(.{ .kernel_backlog = 1 });
    defer server.deinit();
    var queued: std.ArrayList(std.posix.fd_t) = .empty;
    defer {
        for (queued.items) |fd| std.posix.close(fd);
        queued.deinit(a);
    }
    var full = false;
    for (0..16) |_| {
        const fd = try std.posix.socket(std.posix.AF.UNIX, std.posix.SOCK.STREAM | std.posix.SOCK.NONBLOCK | std.posix.SOCK.CLOEXEC, 0);
        const err = std.posix.errno(std.posix.system.connect(fd, &addr.any, addr.getOsSockLen()));
        if (err == .AGAIN) {
            std.posix.close(fd);
            full = true;
            break;
        }
        if (err != .SUCCESS) {
            std.posix.close(fd);
            return error.UnexpectedConnect;
        }
        try queued.append(a, fd);
    }
    try std.testing.expect(full);
    var wait: Wait = .{ .alloc = a, .deadline = std.time.milliTimestamp() + 30 };
    try std.testing.expectError(error.Timeout, connectUnix(path, &wait));
    const pipe = try std.posix.pipe2(.{ .CLOEXEC = true });
    defer for (pipe) |fd| std.posix.close(fd);
    const Cancel = struct {
        fn fire(fd: std.posix.fd_t) void {
            std.Thread.sleep(20 * std.time.ns_per_ms);
            _ = std.posix.write(fd, &.{interrupt.detach_key}) catch {};
        }
    };
    const thread = try std.Thread.spawn(.{}, Cancel.fire, .{pipe[1]});
    defer thread.join();
    wait = .{ .alloc = a, .abort_fd = pipe[0], .deadline = std.time.milliTimestamp() + 2000 };
    const start = std.time.milliTimestamp();
    try std.testing.expectError(error.UserAbort, connectUnix(path, &wait));
    try std.testing.expect(std.time.milliTimestamp() - start < 500);
    const accepted = try server.accept();
    accepted.stream.close();
    wait = .{ .alloc = a, .deadline = std.time.milliTimestamp() + 500 };
    const stream = try connectUnix(path, &wait);
    defer stream.close();
    const flags = try std.posix.fcntl(stream.handle, std.posix.F.GETFL, 0);
    const bits: u32 = @bitCast(std.posix.O{ .NONBLOCK = true });
    try std.testing.expect(flags & bits == 0);
}